From 70be08ed910680fa08e6a6c3ca3f8c96e0fff1a0 Mon Sep 17 00:00:00 2001
From: Kurt Siegfried <kurt@helius.xyz>
Date: Fri, 17 May 2024 14:55:52 +0000
Subject: [PATCH] Add deny list to config

---
 Cargo.toml                           |  1 +
 validator-firewall/src/ip_service.rs | 18 ++++++++-----
 validator-firewall/src/main.rs       | 38 ++++++++++++++++++++++------
 3 files changed, 42 insertions(+), 15 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 939d092..a73f11a 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,2 +1,3 @@
 [workspace]
 members = ["xtask", "validator-firewall", "validator-firewall-common"]
+resolver = "2"
\ No newline at end of file
diff --git a/validator-firewall/src/ip_service.rs b/validator-firewall/src/ip_service.rs
index 0de027f..e437266 100644
--- a/validator-firewall/src/ip_service.rs
+++ b/validator-firewall/src/ip_service.rs
@@ -11,7 +11,7 @@ use tokio::time::sleep;
 pub struct GossipWatcher {
     exit_flag: Arc<AtomicBool>,
     rpc_client: Arc<RpcClient>,
-    static_overrides: Arc<HashSet<u32>>,
+    static_overrides: Arc<(HashSet<u32>, HashSet<u32>)>,
     // allow_list: Map, //Annoying type-erased BPF Map <u32, u8>
 }
 
@@ -19,7 +19,7 @@ impl GossipWatcher {
     pub fn new(
         exit_flag: Arc<AtomicBool>,
         rpc_client: Arc<RpcClient>,
-        static_overrides: Arc<HashSet<u32>>,
+        static_overrides: Arc<(HashSet<u32>, HashSet<u32>)>,
         // allow_list: Map,
     ) -> Self {
         Self {
@@ -31,7 +31,11 @@ impl GossipWatcher {
     }
     pub async fn run(&self, allow_list: Map) {
         let mut allow_list: HashMap<_, u32, u8> = HashMap::try_from(allow_list).unwrap();
-        for ip in self.static_overrides.iter() {
+
+        let allow_overrides = &self.static_overrides.clone().0;
+        let deny_overrides = &self.static_overrides.clone().1;
+
+        for ip in allow_overrides.iter() {
             allow_list.insert(ip, 0, 0).unwrap();
         }
         let mut gossip_set = HashSet::new();
@@ -45,7 +49,9 @@ impl GossipWatcher {
                             let ip_numeric: u32 = (*sock.ip())
                                 .try_into()
                                 .expect("Received invalid ip address");
-                            gossip_set.insert(ip_numeric);
+                            if !deny_overrides.contains(&ip_numeric) {
+                                gossip_set.insert(ip_numeric);
+                            }
                         }
                         SocketAddr::V6(_) => {}
                     }
@@ -55,9 +61,7 @@ impl GossipWatcher {
                     allow_list
                         .iter()
                         .filter_map(|r| r.ok())
-                        .filter(|(x, _)| {
-                            !gossip_set.contains(x) && !self.static_overrides.contains(x)
-                        })
+                        .filter(|(x, _)| !gossip_set.contains(x) && !allow_overrides.contains(x))
                         .map(|x| x.0)
                         .collect::<Vec<u32>>()
                 };
diff --git a/validator-firewall/src/main.rs b/validator-firewall/src/main.rs
index 7a1efd4..8178852 100644
--- a/validator-firewall/src/main.rs
+++ b/validator-firewall/src/main.rs
@@ -11,7 +11,7 @@ use aya::{
 };
 use aya_log::BpfLogger;
 use clap::Parser;
-use log::{debug, info, warn};
+use log::{debug, error, info, warn};
 use serde::Deserialize;
 use solana_rpc_client::nonblocking::rpc_client::RpcClient;
 use std::collections::HashSet;
@@ -29,10 +29,10 @@ struct HVFConfig {
     static_overrides: Option<PathBuf>,
     #[clap(short, long, default_value = "https://api.mainnet-beta.solana.com")]
     rpc_endpoint: String,
-    #[arg(short, long, value_name = "PORT", value_parser = clap::value_parser!(u16), num_args = 1..)]
+    #[arg(short, long, value_name = "PORT", value_parser = clap::value_parser!(u16), num_args = 0.., default_value="8009 8010")]
     protected_ports: Vec<u16>,
 }
-
+#[allow(dead_code)] //Used in Debug
 #[derive(Deserialize, Debug)]
 struct NameAddressPair {
     name: String,
@@ -41,7 +41,8 @@ struct NameAddressPair {
 
 #[derive(Deserialize, Debug)]
 struct StaticOverrides {
-    nodes: Vec<NameAddressPair>,
+    allow: Vec<NameAddressPair>,
+    deny: Vec<NameAddressPair>,
 }
 
 const ALLOW_LIST_MAP: &str = "hvf_allow_list";
@@ -55,15 +56,36 @@ async fn main() -> Result<(), anyhow::Error> {
 
     env_logger::init();
     let static_overrides = {
-        let mut local_overrides = HashSet::new();
+        let mut local_allow = HashSet::new();
+        let mut local_deny = HashSet::new();
+
+        // Load static overrides if provided
         if let Some(path) = config.static_overrides {
             let overrides = load_static_overrides(path)?;
-            for node in overrides.nodes.iter() {
-                local_overrides.insert(u32::from(node.ip));
+            let denied: HashSet<Ipv4Addr> = overrides.deny.iter().map(|x| x.ip.clone()).collect();
+            let intersection: Vec<&NameAddressPair> = overrides
+                .allow
+                .iter()
+                .filter(|x| denied.contains(&x.ip))
+                .collect();
+
+            if !intersection.is_empty() {
+                error!(
+                    "Static overrides contain overlapping entries for deny and allow: {:?}",
+                    intersection
+                );
+                std::process::exit(1);
+            }
+            for node in overrides.allow.iter() {
+                local_allow.insert(u32::from(node.ip));
             }
+            for node in overrides.deny.iter() {
+                local_deny.insert(u32::from(node.ip));
+            }
+
             info!("Loaded static overrides: {:?}", overrides);
         };
-        Arc::new(local_overrides)
+        Arc::new((local_allow, local_deny))
     };
 
     // Bump the memlock rlimit. This is needed for older kernels that don't use the