move registry implementations to poll files

Arian8j2 · Arian8j2 · commit 828fdae669c6 · 2024-10-19T00:02:14.000+03:30
diff --git a/forwarder/src/poll.rs b/forwarder/src/poll.rs
@@ -1,17 +1,16 @@
 use crate::{
     peer::{Peer, PeerManager},
+    socket::NonBlockingSocket,
     uri::Protocol,
 };
 use parking_lot::RwLock;
 use std::sync::Arc;
 
-mod registry;
-pub(crate) use registry::Registry;
-
 type OnPeerRecvCallback = dyn Fn(&Peer, &mut [u8]);
 
+/// trait to be able to listen on multiple sockets asynchronously
 pub trait Poll: Send {
-    /// blocks the current thread and listens on multiple registered `NonBlockingSocket`s
+    /// blocks the current thread and listens on multiple registered `NonBlockingSocket`'s
     /// at the same time and calls `on_peer_recv` on new packets from peer
     fn poll(
         &mut self,
@@ -23,6 +22,13 @@ pub trait Poll: Send {
     fn get_registry(&self) -> anyhow::Result<Box<dyn Registry>>;
 }
 
+/// trait that allows others to register socket to `Poll`
+pub trait Registry: Send + Sync {
+    // need Sync because parking_lot::RwLock needs inner to be Sync
+    fn register(&self, socket: &NonBlockingSocket) -> anyhow::Result<()>;
+    fn deregister(&self, socket: &NonBlockingSocket) -> anyhow::Result<()>;
+}
+
 mod icmp;
 mod udp;
 
diff --git a/forwarder/src/poll/icmp.rs b/forwarder/src/poll/icmp.rs
@@ -1,10 +1,7 @@
-use super::{
-    registry::{IcmpRegistry, Registry},
-    Poll,
-};
+use super::{Poll, Registry};
 use crate::{
     peer::{Peer, PeerManager},
-    socket::icmp::IcmpSocket,
+    socket::{icmp::IcmpSocket, NonBlockingSocket},
     MAX_PACKET_SIZE,
 };
 use parking_lot::RwLock;
@@ -49,3 +46,15 @@ impl Poll for IcmpPoll {
         }
     }
 }
+
+#[derive(Debug)]
+pub struct IcmpRegistry;
+// icmp doesn't need a registry because we manage it's poll ourself
+impl Registry for IcmpRegistry {
+    fn register(&self, _socket: &NonBlockingSocket) -> anyhow::Result<()> {
+        Ok(())
+    }
+    fn deregister(&self, _socket: &NonBlockingSocket) -> anyhow::Result<()> {
+        Ok(())
+    }
+}
diff --git a/forwarder/src/poll/registry.rs b/forwarder/src/poll/registry.rs
diff --git a/forwarder/src/poll/udp.rs b/forwarder/src/poll/udp.rs
@@ -1,12 +1,10 @@
-use super::{
-    registry::{Registry, UdpRegistry},
-    Poll,
-};
+use super::{Poll, Registry};
 use crate::{
     peer::{Peer, PeerManager},
+    socket::{NonBlockingSocket, NonBlockingSocketTrait},
     MAX_PACKET_SIZE,
 };
-use mio::Events;
+use mio::{unix::SourceFd, Events, Interest, Token};
 use parking_lot::RwLock;
 use std::sync::Arc;
 
@@ -48,3 +46,26 @@ impl Poll for UdpPoll {
         }
     }
 }
+
+#[derive(Debug)]
+pub struct UdpRegistry(pub mio::Registry);
+impl Registry for UdpRegistry {
+    fn register(&self, socket: &NonBlockingSocket) -> anyhow::Result<()> {
+        let socket = socket.as_udp().unwrap();
+        let local_port = socket.local_addr()?.port();
+        self.0.register(
+            &mut SourceFd(&socket.as_raw_fd()),
+            Token(local_port.into()),
+            Interest::READABLE,
+        )?;
+        Ok(())
+    }
+
+    fn deregister(&self, socket: &NonBlockingSocket) -> anyhow::Result<()> {
+        let socket = socket.as_udp().unwrap();
+        let raw_fd = socket.as_raw_fd();
+        let source = &mut SourceFd(&raw_fd);
+        self.0.deregister(source)?;
+        Ok(())
+    }
+}
