From c811f2d0a1d5b8121bacf0952ef7ef7dd6efee11 Mon Sep 17 00:00:00 2001
From: Yureka <yuka@yuka.dev>
Date: Wed, 15 Jan 2025 10:44:39 +0100
Subject: [PATCH] add forbidden_nets config option

---
 Cargo.lock  | 12 +++++++++++-
 Cargo.toml  |  1 +
 src/main.rs | 42 ++++++++++++++++++++++++++++++++++++++++--
 3 files changed, 52 insertions(+), 3 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 955511a..92d3082 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1,6 +1,6 @@
 # This file is automatically @generated by Cargo.
 # It is not intended for manual editing.
-version = 3
+version = 4
 
 [[package]]
 name = "aead"
@@ -288,6 +288,15 @@ dependencies = [
  "generic-array",
 ]
 
+[[package]]
+name = "ipnet"
+version = "2.10.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ddc24109865250148c2e0f3d25d4f0f479571723792d3802153c60922a4fb708"
+dependencies = [
+ "serde",
+]
+
 [[package]]
 name = "is-terminal"
 version = "0.4.12"
@@ -621,6 +630,7 @@ dependencies = [
  "blake3",
  "get_if_addrs",
  "igd",
+ "ipnet",
  "log",
  "pretty_env_logger",
  "serde",
diff --git a/Cargo.toml b/Cargo.toml
index e036b4a..a7b0bc1 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -19,3 +19,4 @@ blake3 = "1.5"
 
 igd = { version = "0.12", default-features = false }
 get_if_addrs = "0.5"
+ipnet = { version = "2.10.1", features = ["serde"] }
diff --git a/src/main.rs b/src/main.rs
index fdf7104..dc6e0fb 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -50,6 +50,9 @@ struct Config {
     /// The list of peers we try to connect to
     #[serde(default)]
     peers: Vec<Peer>,
+
+    #[serde(default)]
+    forbidden_nets: Vec<ipnet::IpNet>,
 }
 
 #[derive(Deserialize)]
@@ -527,7 +530,20 @@ impl State {
         let mut peer_vec = self
             .peers
             .iter()
-            .filter(|(_, info)| now < info.last_seen + TIMEOUT.as_secs() && info.endpoint.is_some())
+            .filter(|(_, info)| {
+                let seen = now < info.last_seen + TIMEOUT.as_secs();
+                let endpoint_valid = info
+                    .endpoint
+                    .map(|ep| {
+                        !daemon
+                            .config
+                            .forbidden_nets
+                            .iter()
+                            .any(|net| net.contains(&ep))
+                    })
+                    .unwrap_or(false);
+                seen && endpoint_valid
+            })
             .map(|(_, info)| (info.gossip_ip, info.gossip_prio))
             .collect::<Vec<_>>();
         peer_vec.sort_by_key(|(_, prio)| *prio);
@@ -628,10 +644,22 @@ impl State {
                     (Some((addr1, _)), Some(addr2)) => addr1 != addr2,
                     _ => false,
                 };
+                // If the current endpoint is in a forbidden net, reconfigure the peer even if it has a connection
+                let forbidden_endpoint = peer
+                    .endpoint
+                    .map(|ep| {
+                        daemon
+                            .config
+                            .forbidden_nets
+                            .iter()
+                            .any(|net| net.contains(&ep))
+                    })
+                    .unwrap_or(false);
 
                 // if peer is connected and endpoint is the correct one,
                 // set higher keepalive and then skip reconfiguring it
-                if !bad_endpoint && now < peer.last_seen + TIMEOUT.as_secs() {
+                if !bad_endpoint && !forbidden_endpoint && now < peer.last_seen + TIMEOUT.as_secs()
+                {
                     Command::new("wg")
                         .args([
                             "set",
@@ -671,6 +699,16 @@ impl State {
                         }
                     }
                     endpoints.sort();
+                    endpoints = endpoints
+                        .into_iter()
+                        .filter(|(ep, _)| {
+                            !daemon
+                                .config
+                                .forbidden_nets
+                                .iter()
+                                .any(|net| net.contains(ep))
+                        })
+                        .collect();
                     endpoints
                 }
             };