add udp support for windows

Arian8j2 · Arian8j2 · commit 53acda59d669 · 2024-11-18T22:11:34.000+03:30
diff --git a/forwarder/Cargo.toml b/forwarder/Cargo.toml
@@ -8,5 +8,5 @@ anyhow = "1.0.71"
 log = "0.4.20"
 etherparse = "0.13.0"
 socket2 = { version = "0.5.5", features = ["all"] }
-mio = { version = "1.0.2", features = ["net", "os-ext", "os-poll"] }
+mio = { version = "1.0.2", features = ["net", "os-poll"] }
 parking_lot = "0.12.3"
diff --git a/forwarder/src/lib.rs b/forwarder/src/lib.rs
@@ -150,9 +150,9 @@ fn try_cleanup(peer_manager: &RwLock<PeerManager>) {
     for peer in peers.get_all() {
         let used = peer.reset_used();
         if !used {
-            let client_addr = peer.get_client_addr();
+            let client_addr = *peer.get_client_addr();
             log::info!("cleaning peer that handled '{client_addr}'");
-            if let Err(error) = peers.remove_peer(&peer) {
+            if let Err(error) = peers.remove_peer(peer) {
                 log::warn!("couldn't remove peer of '{client_addr}': {error:?}");
             }
         } else {
diff --git a/forwarder/src/peer.rs b/forwarder/src/peer.rs
@@ -76,13 +76,13 @@ impl PeerManager {
         }
     }
 
-    pub fn add_peer(&mut self, new_peer: Peer) -> anyhow::Result<Arc<Peer>> {
+    pub fn add_peer(&mut self, mut new_peer: Peer) -> anyhow::Result<Arc<Peer>> {
         let client_addr = new_peer.client_addr;
+        self.registry.register(&mut new_peer.socket)?;
         let peer = Arc::new(new_peer);
         self.client_addr_to_peers.insert(client_addr, peer.clone());
         let peer_port = peer.socket.local_addr()?.port();
         self.port_to_peers.insert(peer_port, peer.clone());
-        self.registry.register(&peer.socket)?;
         Ok(peer)
     }
 
@@ -100,10 +100,13 @@ impl PeerManager {
         self.client_addr_to_peers.values().cloned().collect()
     }
 
-    pub fn remove_peer(&mut self, peer: &Peer) -> anyhow::Result<()> {
-        self.registry.deregister(&peer.socket)?;
+    pub fn remove_peer(&mut self, peer: Arc<Peer>) -> anyhow::Result<()> {
         self.client_addr_to_peers.remove(&peer.client_addr);
         self.port_to_peers.remove(&peer.socket.local_addr()?.port());
+
+        let mut peer =
+            Arc::try_unwrap(peer).map_err(|_| anyhow::anyhow!("can't unwrap Arc<peer>"))?;
+        self.registry.deregister(&mut peer.socket)?;
         Ok(())
     }
 }
diff --git a/forwarder/src/poll.rs b/forwarder/src/poll.rs
@@ -25,8 +25,8 @@ pub trait Poll: Send {
 /// 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<()>;
+    fn register(&self, socket: &mut NonBlockingSocket) -> anyhow::Result<()>;
+    fn deregister(&self, socket: &mut NonBlockingSocket) -> anyhow::Result<()>;
 }
 
 mod icmp;
diff --git a/forwarder/src/poll/icmp.rs b/forwarder/src/poll/icmp.rs
@@ -51,10 +51,10 @@ impl Poll for IcmpPoll {
 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<()> {
+    fn register(&self, _socket: &mut NonBlockingSocket) -> anyhow::Result<()> {
         Ok(())
     }
-    fn deregister(&self, _socket: &NonBlockingSocket) -> anyhow::Result<()> {
+    fn deregister(&self, _socket: &mut NonBlockingSocket) -> anyhow::Result<()> {
         Ok(())
     }
 }
diff --git a/forwarder/src/poll/udp.rs b/forwarder/src/poll/udp.rs
@@ -4,7 +4,7 @@ use crate::{
     socket::{NonBlockingSocket, NonBlockingSocketTrait},
     MAX_PACKET_SIZE,
 };
-use mio::{unix::SourceFd, Events, Interest, Token};
+use mio::{Events, Interest, Token};
 use parking_lot::RwLock;
 use std::sync::Arc;
 
@@ -50,22 +50,20 @@ 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();
+    fn register(&self, socket: &mut NonBlockingSocket) -> anyhow::Result<()> {
+        let socket = socket.as_mut_udp().unwrap();
         let local_port = socket.local_addr()?.port();
         self.0.register(
-            &mut SourceFd(&socket.as_raw_fd()),
+            socket.as_inner(),
             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)?;
+    fn deregister(&self, socket: &mut NonBlockingSocket) -> anyhow::Result<()> {
+        let socket = socket.as_mut_udp().unwrap();
+        self.0.deregister(socket.as_inner())?;
         Ok(())
     }
 }
diff --git a/forwarder/src/socket.rs b/forwarder/src/socket.rs
@@ -68,7 +68,7 @@ impl NonBlockingSocket {
         Ok(socket)
     }
 
-    pub fn as_udp(&self) -> Option<&udp::NonBlockingUdpSocket> {
+    pub fn as_mut_udp(&mut self) -> Option<&mut udp::NonBlockingUdpSocket> {
         match self {
             Self::Udp(inner) => Some(inner),
             _ => None,
diff --git a/forwarder/src/socket/udp.rs b/forwarder/src/socket/udp.rs
@@ -1,9 +1,5 @@
 use super::{NonBlockingSocketTrait, SocketTrait};
-use std::{
-    io,
-    net::SocketAddr,
-    os::fd::{AsRawFd, RawFd},
-};
+use std::{io, net::SocketAddr};
 
 #[derive(Debug)]
 pub struct UdpSocket(std::net::UdpSocket);
@@ -30,17 +26,16 @@ impl SocketTrait for UdpSocket {
 }
 
 #[derive(Debug)]
-pub struct NonBlockingUdpSocket(std::net::UdpSocket);
+pub struct NonBlockingUdpSocket(mio::net::UdpSocket);
 
 impl NonBlockingUdpSocket {
     pub fn bind(address: &SocketAddr) -> io::Result<Self> {
-        let socket = std::net::UdpSocket::bind(address)?;
-        socket.set_nonblocking(true)?;
+        let socket = mio::net::UdpSocket::bind(*address)?;
         Ok(Self(socket))
     }
 
-    pub fn as_raw_fd(&self) -> RawFd {
-        self.0.as_raw_fd()
+    pub fn as_inner(&mut self) -> &mut mio::net::UdpSocket {
+        &mut self.0
     }
 }
 
@@ -50,7 +45,7 @@ impl NonBlockingSocketTrait for NonBlockingUdpSocket {
     }
 
     fn connect(&mut self, addr: &SocketAddr) -> io::Result<()> {
-        self.0.connect(addr)
+        self.0.connect(*addr)
     }
 
     fn recv(&self, buffer: &mut [u8]) -> io::Result<usize> {

Original file line number	Diff line number	Diff line change
`@@ -76,13 +76,13 @@ impl PeerManager {`
`76`	`76`	`}`
`77`	`77`	`}`
`78`	`78`
`79`		`- pub fn add_peer(&mut self, new_peer: Peer) -> anyhow::Result<Arc<Peer>> {`
	`79`	`+ pub fn add_peer(&mut self, mut new_peer: Peer) -> anyhow::Result<Arc<Peer>> {`
`80`	`80`	`let client_addr = new_peer.client_addr;`
	`81`	`+ self.registry.register(&mut new_peer.socket)?;`
`81`	`82`	`let peer = Arc::new(new_peer);`
`82`	`83`	`self.client_addr_to_peers.insert(client_addr, peer.clone());`
`83`	`84`	`let peer_port = peer.socket.local_addr()?.port();`
`84`	`85`	`self.port_to_peers.insert(peer_port, peer.clone());`
`85`		`- self.registry.register(&peer.socket)?;`
`86`	`86`	`Ok(peer)`
`87`	`87`	`}`
`88`	`88`
`@@ -100,10 +100,13 @@ impl PeerManager {`
`100`	`100`	`self.client_addr_to_peers.values().cloned().collect()`
`101`	`101`	`}`
`102`	`102`
`103`		`- pub fn remove_peer(&mut self, peer: &Peer) -> anyhow::Result<()> {`
`104`		`- self.registry.deregister(&peer.socket)?;`
	`103`	`+ pub fn remove_peer(&mut self, peer: Arc<Peer>) -> anyhow::Result<()> {`
`105`	`104`	`self.client_addr_to_peers.remove(&peer.client_addr);`
`106`	`105`	`self.port_to_peers.remove(&peer.socket.local_addr()?.port());`
	`106`	`+`
	`107`	`+ let mut peer =`
	`108`	`+ Arc::try_unwrap(peer).map_err(\|_\| anyhow::anyhow!("can't unwrap Arc<peer>"))?;`
	`109`	`+ self.registry.deregister(&mut peer.socket)?;`
`107`	`110`	`Ok(())`
`108`	`111`	`}`
`109`	`112`	`}`
Original file line number	Diff line number	Diff line change
`@@ -51,10 +51,10 @@ impl Poll for IcmpPoll {`
`51`	`51`	`pub struct IcmpRegistry;`
`52`	`52`	`// icmp doesn't need a registry because we manage it's poll ourself`
`53`	`53`	`impl Registry for IcmpRegistry {`
`54`		`- fn register(&self, _socket: &NonBlockingSocket) -> anyhow::Result<()> {`
	`54`	`+ fn register(&self, _socket: &mut NonBlockingSocket) -> anyhow::Result<()> {`
`55`	`55`	`Ok(())`
`56`	`56`	`}`
`57`		`- fn deregister(&self, _socket: &NonBlockingSocket) -> anyhow::Result<()> {`
	`57`	`+ fn deregister(&self, _socket: &mut NonBlockingSocket) -> anyhow::Result<()> {`
`58`	`58`	`Ok(())`
`59`	`59`	`}`
`60`	`60`	`}`
Original file line number	Diff line number	Diff line change
`@@ -68,7 +68,7 @@ impl NonBlockingSocket {`
`68`	`68`	`Ok(socket)`
`69`	`69`	`}`
`70`	`70`
`71`		`- pub fn as_udp(&self) -> Option<&udp::NonBlockingUdpSocket> {`
	`71`	`+ pub fn as_mut_udp(&mut self) -> Option<&mut udp::NonBlockingUdpSocket> {`
`72`	`72`	`match self {`
`73`	`73`	`Self::Udp(inner) => Some(inner),`
`74`	`74`	`_ => None,`
Original file line number	Diff line number	Diff line change
`@@ -1,9 +1,5 @@`
`1`	`1`	`use super::{NonBlockingSocketTrait, SocketTrait};`
`2`		`-use std::{`
`3`		`- io,`
`4`		`- net::SocketAddr,`
`5`		`- os::fd::{AsRawFd, RawFd},`
`6`		`-};`
	`2`	`+use std::{io, net::SocketAddr};`
`7`	`3`
`8`	`4`	`#[derive(Debug)]`
`9`	`5`	`pub struct UdpSocket(std::net::UdpSocket);`
`@@ -30,17 +26,16 @@ impl SocketTrait for UdpSocket {`
`30`	`26`	`}`
`31`	`27`
`32`	`28`	`#[derive(Debug)]`
`33`		`-pub struct NonBlockingUdpSocket(std::net::UdpSocket);`
	`29`	`+pub struct NonBlockingUdpSocket(mio::net::UdpSocket);`
`34`	`30`
`35`	`31`	`impl NonBlockingUdpSocket {`
`36`	`32`	`pub fn bind(address: &SocketAddr) -> io::Result<Self> {`
`37`		`- let socket = std::net::UdpSocket::bind(address)?;`
`38`		`- socket.set_nonblocking(true)?;`
	`33`	`+ let socket = mio::net::UdpSocket::bind(*address)?;`
`39`	`34`	`Ok(Self(socket))`
`40`	`35`	`}`
`41`	`36`
`42`		`- pub fn as_raw_fd(&self) -> RawFd {`
`43`		`- self.0.as_raw_fd()`
	`37`	`+ pub fn as_inner(&mut self) -> &mut mio::net::UdpSocket {`
	`38`	`+ &mut self.0`
`44`	`39`	`}`
`45`	`40`	`}`
`46`	`41`
`@@ -50,7 +45,7 @@ impl NonBlockingSocketTrait for NonBlockingUdpSocket {`
`50`	`45`	`}`
`51`	`46`
`52`	`47`	`fn connect(&mut self, addr: &SocketAddr) -> io::Result<()> {`
`53`		`- self.0.connect(addr)`
	`48`	`+ self.0.connect(*addr)`
`54`	`49`	`}`
`55`	`50`
`56`	`51`	`fn recv(&self, buffer: &mut [u8]) -> io::Result<usize> {`