From 7658e2cb427ad804fbc26ce52da2a568f302dde5 Mon Sep 17 00:00:00 2001
From: John Smith <jsmithveilid@protonmail.com>
Date: Mon, 17 Jul 2023 17:51:05 -0400
Subject: [PATCH] fix punishment

---
 veilid-core/src/network_manager/network_connection.rs | 6 ++++++
 veilid-core/src/network_manager/tasks/mod.rs          | 3 +++
 veilid-tools/src/network_result.rs                    | 3 +++
 3 files changed, 12 insertions(+)

diff --git a/veilid-core/src/network_manager/network_connection.rs b/veilid-core/src/network_manager/network_connection.rs
index 9d7b7eb6..dbda6da2 100644
--- a/veilid-core/src/network_manager/network_connection.rs
+++ b/veilid-core/src/network_manager/network_connection.rs
@@ -315,6 +315,12 @@ impl NetworkConnection {
                                         return RecvLoopAction::Finish;
                                     }
 
+                                    // Punish invalid messages
+                                    if v.is_invalid_message() {
+                                        address_filter.punish_ip_addr(peer_address.to_socket_addr().ip());
+                                        return RecvLoopAction::Finish;
+                                    }
+
                                     // Log other network results
                                     let mut message = network_result_value_or_log!(v => [ format!(": protocol_connection={:?}", protocol_connection) ] {
                                         return RecvLoopAction::Finish;
diff --git a/veilid-core/src/network_manager/tasks/mod.rs b/veilid-core/src/network_manager/tasks/mod.rs
index d7ff6fb0..a57b0d3c 100644
--- a/veilid-core/src/network_manager/tasks/mod.rs
+++ b/veilid-core/src/network_manager/tasks/mod.rs
@@ -66,6 +66,9 @@ impl NetworkManager {
         // Run the rolling transfers task
         self.unlocked_inner.rolling_transfers_task.tick().await?;
 
+        // Run the address filter task
+        self.unlocked_inner.address_filter_task.tick().await?;
+
         // Run the routing table tick
         routing_table.tick().await?;
 
diff --git a/veilid-tools/src/network_result.rs b/veilid-tools/src/network_result.rs
index ae93fe97..760cd460 100644
--- a/veilid-tools/src/network_result.rs
+++ b/veilid-tools/src/network_result.rs
@@ -202,6 +202,9 @@ impl<T> NetworkResult<T> {
     pub fn is_already_exists(&self) -> bool {
         matches!(self, Self::AlreadyExists(_))
     }
+    pub fn is_invalid_message(&self) -> bool {
+        matches!(self, Self::InvalidMessage(_))
+    }
     pub fn is_value(&self) -> bool {
         matches!(self, Self::Value(_))
     }