From f1f3a730d0d6d4af16ba545fa149b5ec9d2a1c31 Mon Sep 17 00:00:00 2001
From: ssrlive <30760636+ssrlive@users.noreply.github.com>
Date: Sun, 24 Sep 2023 00:30:25 +0800
Subject: [PATCH] pipe_client_event

---
 src/tun2proxy.rs       | 26 ++++++++++++++++--------
 src/wintuninterface.rs | 46 ++++++++++++++++++++++++++++++++++++++++--
 2 files changed, 62 insertions(+), 10 deletions(-)

diff --git a/src/tun2proxy.rs b/src/tun2proxy.rs
index d746f13..9abf7c6 100644
--- a/src/tun2proxy.rs
+++ b/src/tun2proxy.rs
@@ -11,12 +11,10 @@ use smoltcp::phy::RawSocket;
 // #[cfg(any(target_os = "linux", target_os = "android"))]
 // use smoltcp::phy::TunTapInterface;
 #[cfg(target_os = "windows")]
-use crate::wintuninterface::WinTunInterface;
-#[cfg(target_family = "unix")]
-use smoltcp::phy::{RxToken, TxToken};
+use crate::wintuninterface::{NamedPipeSource, WinTunInterface};
 use smoltcp::{
     iface::{Config, Interface, SocketHandle, SocketSet},
-    phy::{Device, Medium},
+    phy::{Device, Medium, RxToken, TxToken},
     socket::{tcp, tcp::State, udp, udp::UdpMetadata},
     time::Instant,
     wire::{IpCidr, IpProtocol, Ipv4Packet, Ipv6Packet, TcpPacket, UdpPacket, UDP_HEADER_LEN},
@@ -211,7 +209,8 @@ pub(crate) trait ConnectionManager {
 }
 
 const TUN_TOKEN: Token = Token(0);
-const EXIT_TOKEN: Token = Token(1);
+const PIPE_TOKEN: Token = Token(1);
+const EXIT_TOKEN: Token = Token(2);
 
 pub struct TunToProxy<'a> {
     #[cfg(any(target_os = "linux", target_os = "android"))]
@@ -261,7 +260,11 @@ impl<'a> TunToProxy<'a> {
             .register(&mut SourceFd(&tun.as_raw_fd()), TUN_TOKEN, Interest::READABLE)?;
 
         #[cfg(target_os = "windows")]
-        poll.registry().register(&mut tun, TUN_TOKEN, Interest::READABLE)?;
+        {
+            poll.registry().register(&mut tun, TUN_TOKEN, Interest::READABLE)?;
+            let mut pipe = NamedPipeSource(tun.pipe_client());
+            poll.registry().register(&mut pipe, PIPE_TOKEN, Interest::READABLE)?;
+        }
 
         #[cfg(target_family = "unix")]
         let (exit_sender, mut exit_receiver) = mio::unix::pipe::new()?;
@@ -324,7 +327,6 @@ impl<'a> TunToProxy<'a> {
             let _slice = vec.as_slice();
 
             // TODO: Actual write. Replace.
-            #[cfg(target_family = "unix")]
             self.tun
                 .transmit(Instant::now())
                 .ok_or("tx token not available")?
@@ -931,7 +933,6 @@ impl<'a> TunToProxy<'a> {
 
     fn tun_event(&mut self, event: &Event) -> Result<(), Error> {
         if event.is_readable() {
-            #[cfg(target_family = "unix")]
             while let Some((rx_token, _)) = self.tun.receive(Instant::now()) {
                 rx_token.consume(|frame| self.receive_tun(frame))?;
             }
@@ -939,6 +940,14 @@ impl<'a> TunToProxy<'a> {
         Ok(())
     }
 
+    fn pipe_event(&mut self, event: &Event) -> Result<(), Error> {
+        if event.is_readable() {
+            #[cfg(target_os = "windows")]
+            self.tun.pipe_client_event()?;
+        }
+        Ok(())
+    }
+
     fn send_to_smoltcp(&mut self) -> Result<(), Error> {
         for token in self.write_sockets.clone().into_iter() {
             if let Some(connection) = self.find_info_by_token(token) {
@@ -1132,6 +1141,7 @@ impl<'a> TunToProxy<'a> {
                         return Ok(());
                     }
                     TUN_TOKEN => self.tun_event(event)?,
+                    PIPE_TOKEN => self.pipe_event(event)?,
                     _ => self.mio_socket_event(event)?,
                 }
             }
diff --git a/src/wintuninterface.rs b/src/wintuninterface.rs
index 57ce229..28b35dd 100644
--- a/src/wintuninterface.rs
+++ b/src/wintuninterface.rs
@@ -41,7 +41,7 @@ pub struct WinTunInterface {
     mtu: usize,
     medium: Medium,
     pipe_server: Rc<RefCell<NamedPipe>>,
-    _pipe_client: Rc<RefCell<NamedPipe>>,
+    pipe_client: Rc<RefCell<NamedPipe>>,
 }
 
 impl event::Source for WinTunInterface {
@@ -93,9 +93,35 @@ impl WinTunInterface {
             mtu,
             medium,
             pipe_server: Rc::new(RefCell::new(pipe_server)),
-            _pipe_client: Rc::new(RefCell::new(pipe_client)),
+            pipe_client: Rc::new(RefCell::new(pipe_client)),
         })
     }
+
+    pub fn pipe_client(&self) -> Rc<RefCell<NamedPipe>> {
+        self.pipe_client.clone()
+    }
+
+    // pub fn pipe_server(&self) -> Rc<RefCell<NamedPipe>> {
+    //     self.pipe_server.clone()
+    // }
+
+    pub fn pipe_client_event(&self) -> Result<(), io::Error> {
+        let mut buffer = vec![0; self.mtu];
+        match self.pipe_client.borrow_mut().read(&mut buffer) {
+            Ok(len) => {
+                let write_pack = self.inner.allocate_send_packet(len as u16);
+                if let Ok(mut write_pack) = write_pack {
+                    write_pack.bytes_mut().copy_from_slice(&buffer[..len]);
+                    self.inner.send_packet(write_pack);
+                } else if let Err(err) = write_pack {
+                    log::error!("phy: failed to allocate send packet: {}", err);
+                }
+            }
+            Err(err) if err.kind() == io::ErrorKind::WouldBlock => {}
+            Err(err) => panic!("{}", err),
+        }
+        Ok(())
+    }
 }
 
 impl Drop for WinTunInterface {
@@ -204,3 +230,19 @@ impl phy::TxToken for TxToken {
         result
     }
 }
+
+pub struct NamedPipeSource(pub Rc<RefCell<NamedPipe>>);
+
+impl event::Source for NamedPipeSource {
+    fn register(&mut self, registry: &Registry, token: Token, interests: Interest) -> io::Result<()> {
+        self.0.borrow_mut().register(registry, token, interests)
+    }
+
+    fn reregister(&mut self, registry: &Registry, token: Token, interests: Interest) -> io::Result<()> {
+        self.0.borrow_mut().reregister(registry, token, interests)
+    }
+
+    fn deregister(&mut self, registry: &Registry) -> io::Result<()> {
+        self.0.borrow_mut().deregister(registry)
+    }
+}