split

src/main.rs

@@ -56,16 +56,28 @@ 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 (mut router_readers, 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());
+    let router_writers3: Vec<(Arc<Socket>, HashMap<u8, RouterWriter>)> = router_writers
+        .into_iter()
+        .grouping_by(|(k, v)| v.key())
+        .into_iter()
+        .map(|(k, v)| {
+            (
+                Arc::clone(sockets.get_mut(&k).unwrap()),
+                v.into_iter().collect(),
+            )
+        })
+        .collect();
 
     thread::scope(|s| {
         for router in router_readers.values_mut() {
@@ -109,8 +121,8 @@ fn main() -> Result<(), Box<dyn Error>> {
             });
         }
 
-        for socket in sockets.values() {
-            s.spawn(|_| {
+        for (socket, mut router_writers) in router_writers3 {
+            s.spawn(move |_| {
                 let mut recv_buf = [MaybeUninit::uninit(); 1500];
                 loop {
                     match socket.recv_from(&mut recv_buf) {
@@ -133,7 +145,7 @@ fn main() -> Result<(), Box<dyn Error>> {
 
                                 if let Some(router) = router_writers.get_mut(&meta.src_id) {
                                     if meta.dst_id == config.local_id && meta.reversed == 0 {
-                                        *router_readers.get(&meta.src_id).unwrap().endpoint.write().unwrap() = Some(addr);
+                                        *router.endpoint.write().unwrap() = Some(addr);
                                         router.decrypt(payload);
                                         router.tun_writer.write_all(payload).unwrap();
                                     }

src/router.rs

@@ -15,7 +15,7 @@ pub struct RouterReader {
     pub secret: [u8; SECRET_LENGTH],
     pub tun_reader: Reader,
     pub socket: Arc<Socket>,
-    pub endpoint: RwLock<Option<SockAddr>>,
+    pub endpoint: Arc<RwLock<Option<SockAddr>>>,
 }
 
 impl RouterReader {
@@ -30,6 +30,7 @@ impl RouterReader {
 pub struct RouterWriter {
     pub config: &'static ConfigRouter,
     pub tun_writer: Writer,
+    pub endpoint: Arc<RwLock<Option<SockAddr>>>,
 }
 
 impl RouterWriter {
@@ -50,7 +51,7 @@ pub struct Router {
     pub tun_reader: Reader,
     pub tun_writer: Writer,
     pub socket: Arc<Socket>,
-    pub endpoint: RwLock<Option<SockAddr>>,
+    pub endpoint: Arc<RwLock<Option<SockAddr>>>,
 }
 
 impl Router {
@@ -104,9 +105,9 @@ impl Router {
 
     fn create_endpoint(
         config: &ConfigRouter,
-    ) -> Result<RwLock<Option<SockAddr>>, Box<dyn std::error::Error>> {
+    ) -> Result<Arc<RwLock<Option<SockAddr>>>, Box<dyn std::error::Error>> {
         let parsed = config.endpoint.to_socket_addrs()?.next().unwrap();
-        Ok(RwLock::new(Some(parsed.into())))
+        Ok(Arc::new(RwLock::new(Some(parsed.into()))))
     }
 
     pub fn new(
@@ -134,7 +135,7 @@ impl Router {
     pub fn split(self) -> (RouterReader, RouterWriter) {
         let writer = RouterWriter {
             config: self.config,
-            // endpoint: self.endpoint.read().unwrap().clone(),
+            endpoint: Arc::clone(&self.endpoint),
             tun_writer: self.tun_writer,
         };