split

Cargo.lock

@@ -301,6 +301,12 @@ version = "0.3.2"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "a8d1add55171497b4705a648c6b583acafb01d58050a51727785f0b2c8e0a2b2"
 
+[[package]]
+name = "grouping_by"
+version = "0.2.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8ad3dbb820113291dd2f6aea7fb4fa6594d80381b287d4ec03ee56bb692870a9"
+
 [[package]]
 name = "ipnet"
 version = "2.10.1"
@@ -664,6 +670,7 @@ dependencies = [
  "base64",
  "crossbeam",
  "crossbeam-utils",
+ "grouping_by",
  "lazy_static",
  "pnet",
  "pnet_macros",

Cargo.toml

@@ -14,3 +14,4 @@ base64 = "0.22.1"
 lazy_static = "1.5.0"
 crossbeam = "0.8.4"
 crossbeam-utils = "0.8.20"
+grouping_by = "0.2.2"

src/main.rs

 mod router;
 
-use crate::router::{Router, SECRET_LENGTH};
+use crate::router::{Router, RouterReader, RouterWriter, SECRET_LENGTH};
 use std::collections::HashMap;
 use std::env;
 use std::error::Error;
@@ -38,6 +38,7 @@ pub struct Config {
     pub routers: Vec<ConfigRouter>,
 }
 use crossbeam_utils::thread;
+use grouping_by::GroupingBy;
 use lazy_static::lazy_static;
 use pnet::packet::ipv4::Ipv4Packet;
 use socket2::Socket;
@@ -55,9 +56,19 @@ fn main() -> Result<(), Box<dyn Error>> {
         .iter()
         .map(|c| Router::new(c, &mut sockets).map(|router| (c.remote_id, router)))
         .collect::<Result<_, _>>()?;
+    let  (mut router_readers, mut router_writers): (HashMap<u8, RouterReader>, HashMap<u8, RouterWriter>) =
+        routers
+            .into_iter()
+            .map(|(id, router)| {
+                let (reader, writer) = router.split();
+                ((id, reader), (id, writer))
+            })
+            .unzip();
+
+    let router_writers2 = router_writers.into_values().grouping_by(|k| k.key());
 
     thread::scope(|s| {
-        for router in routers.values_mut() {
+        for router in router_readers.values_mut() {
             s.spawn(|_| {
                 let mut buffer = [0u8; 1500 - 20]; // minus typical IP header space
                 let meta_size = size_of::<Meta>();
@@ -78,7 +89,7 @@ fn main() -> Result<(), Box<dyn Error>> {
                     match router.tun_reader.read(&mut buffer[meta_size..]) {
                         Ok(n) if n > 0 => {
                             // If we have a known remote address, encrypt and send
-                            if let Some(ref addr) = router.endpoint {
+                            if let Some(ref addr) = *router.endpoint.read().unwrap() {
                                 router.encrypt(&mut buffer[meta_size..meta_size + n]);
                                 #[cfg(target_os = "linux")]
                                 router.socket.set_mark(router.config.mark)?;
@@ -120,9 +131,9 @@ fn main() -> Result<(), Box<dyn Error>> {
                                     reversed: u16::from_le_bytes([meta_bytes[2], meta_bytes[3]]),
                                 };
 
-                                if let Some(router) = routers.get_mut(&meta.src_id) {
+                                if let Some(router) = router_writers.get_mut(&meta.src_id) {
                                     if meta.dst_id == config.local_id && meta.reversed == 0 {
-                                        router.endpoint = Some(addr);
+                                        *router_readers.get(&meta.src_id).unwrap().endpoint.write().unwrap() = Some(addr);
                                         router.decrypt(payload);
                                         router.tun_writer.write_all(payload).unwrap();
                                     }

src/router.rs

@@ -3,19 +3,54 @@ use std::collections::hash_map::Entry;
 use std::collections::HashMap;
 use std::net::ToSocketAddrs;
 use std::process::Command;
-use std::sync::Arc;
+use std::sync::{Arc, RwLock};
 use tun::{Reader, Writer};
 pub const SECRET_LENGTH: usize = 32;
-use crate::ConfigRouter;
+use crate::{local_secret, ConfigRouter};
 use base64::prelude::*;
 
+// tun -> raw
+pub struct RouterReader {
+    pub config: &'static ConfigRouter,
+    pub secret: [u8; SECRET_LENGTH],
+    pub tun_reader: Reader,
+    pub socket: Arc<Socket>,
+    pub endpoint: RwLock<Option<SockAddr>>,
+}
+
+impl RouterReader {
+    pub(crate) fn encrypt(&self, data: &mut [u8]) {
+        for (i, b) in data.iter_mut().enumerate() {
+            *b ^= self.secret[i % SECRET_LENGTH];
+        }
+    }
+}
+
+// raw -> tun
+pub struct RouterWriter {
+    pub config: &'static ConfigRouter,
+    pub tun_writer: Writer,
+}
+
+impl RouterWriter {
+    pub(crate) fn decrypt(&self, data: &mut [u8]) {
+        for (i, b) in data.iter_mut().enumerate() {
+            *b ^= local_secret[i % SECRET_LENGTH];
+        }
+    }
+
+    pub(crate) fn key(&self) -> u16 {
+        Router::key(self.config)
+    }
+}
+
 pub struct Router {
     pub config: &'static ConfigRouter,
     pub secret: [u8; SECRET_LENGTH],
-    pub endpoint: Option<SockAddr>,
     pub tun_reader: Reader,
     pub tun_writer: Writer,
     pub socket: Arc<Socket>,
+    pub endpoint: RwLock<Option<SockAddr>>,
 }
 
 impl Router {
@@ -28,11 +63,15 @@ impl Router {
         secret[..len].copy_from_slice(&decoded[..len]);
         Ok(secret)
     }
+
+    fn key(config: &ConfigRouter) -> u16 {
+        (config.family as u16) << 8 | config.proto as u16
+    }
     fn create_raw_socket(
         config: &ConfigRouter,
         sockets: &mut HashMap<u16, Arc<Socket>>,
     ) -> Result<Arc<Socket>, Box<dyn std::error::Error>> {
-        let key = (config.family as u16) << 8 | config.proto as u16;
+        let key = Router::key(config);
         let result = match sockets.entry(key) {
             Entry::Occupied(entry) => entry.get().clone(),
             Entry::Vacant(entry) => entry
@@ -65,9 +104,9 @@ impl Router {
 
     fn create_endpoint(
         config: &ConfigRouter,
-    ) -> Result<Option<SockAddr>, Box<dyn std::error::Error>> {
+    ) -> Result<RwLock<Option<SockAddr>>, Box<dyn std::error::Error>> {
         let parsed = config.endpoint.to_socket_addrs()?.next().unwrap();
-        Ok(Some(parsed.into()))
+        Ok(RwLock::new(Some(parsed.into())))
     }
 
     pub fn new(
@@ -92,12 +131,21 @@ impl Router {
         Ok(router)
     }
 
-    pub(crate) fn encrypt(&self, data: &mut [u8]) {
-        for (i, b) in data.iter_mut().enumerate() {
-            *b ^= self.secret[i % SECRET_LENGTH];
-        }
-    }
-    pub(crate) fn decrypt(&self, data: &mut [u8]) {
-        self.encrypt(data);
+    pub fn split(self) -> (RouterReader, RouterWriter) {
+        let writer = RouterWriter {
+            config: self.config,
+            // endpoint: self.endpoint.read().unwrap().clone(),
+            tun_writer: self.tun_writer,
+        };
+
+        let reader = RouterReader {
+            config: self.config,
+            secret: self.secret,
+            endpoint: self.endpoint,
+            tun_reader: self.tun_reader,
+            socket: self.socket,
+        };
+
+        (reader, writer)
     }
 }