From f8c749a1b0c9f9616bf52447f7c624d01d309642 Mon Sep 17 00:00:00 2001
From: Kirill Fomichev <fanatid@ya.ru>
Date: Tue, 18 Oct 2022 15:49:54 -0300
Subject: [PATCH] Add `local-ip-range` argument

---
 client/src/main.rs  | 30 +++++++++++++++++++++++++++++-
 shared/src/types.rs |  6 ++++++
 2 files changed, 35 insertions(+), 1 deletion(-)

diff --git a/client/src/main.rs b/client/src/main.rs
index 4e7201f..e936163 100644
--- a/client/src/main.rs
+++ b/client/src/main.rs
@@ -567,7 +567,7 @@ fn fetch(
     let State { peers, cidrs } = api.http("GET", "/user/state")?;
 
     let device = Device::get(interface, opts.network.backend)?;
-    let modifications = device.diff(&peers);
+    let mut modifications = device.diff(&peers);
 
     let updates = modifications
         .iter()
@@ -621,6 +621,34 @@ fn fetch(
     if nat.no_nat_traversal {
         log::debug!("NAT traversal explicitly disabled, not attempting.");
     } else {
+        let mut peers = Vec::with_capacity(modifications.len());
+        if let Some(local_ip_range) = nat.local_ip_range {
+            for diff in &modifications {
+                peers.push(
+                    diff.new
+                        .cloned()
+                        .map(|mut peer| {
+                            let mut candidates = vec![];
+                            let mut rest = vec![];
+                            for endpoint in peer.candidates.drain(..) {
+                                let addr = endpoint.resolve().with_str(endpoint.to_string())?;
+                                if local_ip_range.contains(&addr.ip()) {
+                                    candidates.push(endpoint);
+                                } else {
+                                    rest.push(endpoint);
+                                }
+                            }
+                            candidates.append(&mut rest);
+                            peer.candidates = candidates;
+                            Ok::<_, anyhow::Error>(peer)
+                        })
+                        .transpose()?,
+                );
+            }
+            for (i, mut diff) in modifications.iter_mut().enumerate() {
+                diff.new = peers.get(i).expect("same length").as_ref();
+            }
+        }
         let mut nat_traverse = NatTraverse::new(interface, opts.network.backend, &modifications)?;
 
         // Give time for handshakes with recently changed endpoints to complete before attempting traversal.
diff --git a/shared/src/types.rs b/shared/src/types.rs
index a0012b1..71846e5 100644
--- a/shared/src/types.rs
+++ b/shared/src/types.rs
@@ -437,6 +437,11 @@ pub struct NatOpts {
     /// Don't report any candidates to coordinating server.
     /// Shorthand for --exclude-nat-candidates '0.0.0.0/0'.
     pub no_nat_candidates: bool,
+
+    #[clap(long)]
+    /// Priorities candidads from given IP range.
+    /// ex. --local-ip-range '192.168.10.0/24'
+    pub local_ip_range: Option<IpNet>,
 }
 
 impl NatOpts {
@@ -445,6 +450,7 @@ impl NatOpts {
             no_nat_traversal: true,
             exclude_nat_candidates: vec![],
             no_nat_candidates: true,
+            local_ip_range: None,
         }
     }