From 145f099b556b2d3fa6d5b7c9c2c0ae3eab079c6d Mon Sep 17 00:00:00 2001
From: YanLien <YJQ980314@outlook.com>
Date: Sat, 14 Sep 2024 10:06:25 +0800
Subject: [PATCH] support communication with host

---
 src/smoltcp_impl/addr.rs | 15 ++++++++
 src/smoltcp_impl/mod.rs  | 79 ++++++++++++++++++++++++++++++++++++++--
 2 files changed, 91 insertions(+), 3 deletions(-)

diff --git a/src/smoltcp_impl/addr.rs b/src/smoltcp_impl/addr.rs
index 2cd0b30..e6e5f56 100644
--- a/src/smoltcp_impl/addr.rs
+++ b/src/smoltcp_impl/addr.rs
@@ -34,3 +34,18 @@ pub fn is_unspecified(ip: IpAddress) -> bool {
 
 pub const UNSPECIFIED_IP: IpAddress = IpAddress::v4(0, 0, 0, 0);
 pub const UNSPECIFIED_ENDPOINT: IpEndpoint = IpEndpoint::new(UNSPECIFIED_IP, 0);
+
+pub fn is_loopback_route(ip: IpAddress) -> bool {
+    match ip {
+        IpAddress::Ipv4(addr) => is_ipv4_loopback(addr),
+        IpAddress::Ipv6(addr) => is_ipv6_loopback(addr),
+    }
+}
+
+fn is_ipv4_loopback(addr: Ipv4Address) -> bool {
+    addr.0[0] == 127
+}
+
+fn is_ipv6_loopback(addr: Ipv6Address) -> bool {
+    addr.0 == [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1]
+}
\ No newline at end of file
diff --git a/src/smoltcp_impl/mod.rs b/src/smoltcp_impl/mod.rs
index 19293f9..61d05d0 100644
--- a/src/smoltcp_impl/mod.rs
+++ b/src/smoltcp_impl/mod.rs
@@ -148,9 +148,55 @@ impl<'a> SocketSetWrapper<'a> {
     }
 
     pub fn poll_interfaces(&self) {
-        let timestamp =
-            Instant::from_micros_const((current_time_nanos() / NANOS_PER_MICROS) as i64);
-        #[cfg(feature = "monolithic")]
+        let timestamp = Instant::from_micros_const((current_time_nanos() / NANOS_PER_MICROS) as i64);
+        
+        let mut poll_actions: Vec<Box<dyn Fn()>> = Vec::new();
+        {
+            let socket_set = self.0.lock();
+            for (handle, socket) in socket_set.iter() {
+                debug!("socket is {:?}",socket);
+                match socket {
+                    Socket::Tcp(tcp_socket) => {
+                        if let Some(remote_endpoint) = tcp_socket.remote_endpoint() {
+                            let dest_ip = remote_endpoint.addr;
+                            debug!("tcp_socket can send to {:?}", dest_ip);
+                            let self_clone = self;
+                            poll_actions.push(Box::new(move || {
+                                self_clone.route_and_poll(dest_ip, handle, timestamp);
+                            }));
+                        }
+                    },
+                    Socket::Udp(udp_socket) => {
+                        // if udp_socket.can_send() {
+                            if let Some(dest_ip) = udp_socket.endpoint().addr {
+                                debug!("udp_socket can send to {:?}", dest_ip);
+                                let self_clone = self;
+                                poll_actions.push(Box::new(move || {
+                                    self_clone.route_and_poll(dest_ip, handle, timestamp);
+                                }));
+                            }
+                        // }
+                    },
+                    Socket::Raw(raw_socket) => {
+                        debug!("socket {:?} is raw", raw_socket);
+                    },
+                    Socket::Icmp(icmp_socket) => {
+                        debug!("socket {:?} is icmp", icmp_socket);
+                    },
+                    Socket::Dns(dns_socket) => {
+                        debug!("socket {:?} is dns", dns_socket);
+                        ETH0.poll(&self.0);
+                    },
+                }
+            }
+        }
+
+        // 执行所有收集到的操作
+        for action in poll_actions {
+            action();
+        }
+
+        // // #[cfg(feature = "monolithic")]
         LOOPBACK.lock().poll(
             timestamp,
             LOOPBACK_DEV.lock().deref_mut(),
@@ -160,12 +206,39 @@ impl<'a> SocketSetWrapper<'a> {
         ETH0.poll(&self.0);
     }
 
+    pub fn route_and_poll(&self, dest_ip: IpAddress, _handle: SocketHandle, timestamp: Instant) {
+        debug!("route_and_poll {:?}", dest_ip);
+        match route_packet(dest_ip) {
+            "loopback" => {
+                LOOPBACK.lock().poll(timestamp, LOOPBACK_DEV.lock().deref_mut(), &mut self.0.lock());
+            },
+            "eth0" => {
+                ETH0.poll(&self.0);
+            },
+            _ => unreachable!(),
+        }
+    }
+
     pub fn remove(&self, handle: SocketHandle) {
         self.0.lock().remove(handle);
         debug!("socket {}: destroyed", handle);
     }
 }
 
+fn route_packet(destination: IpAddress) -> &'static str {
+    if is_loopback_route(destination) {
+        "loopback"
+    } else {
+        // 检查是否匹配配置的回环地址
+        let loopback = LOOPBACK.lock();
+        if loopback.ip_addrs().iter().any(|cidr| cidr.contains_addr(&destination)) {
+            "loopback"
+        } else {
+            "eth0"
+        }
+    }
+}
+
 #[allow(unused)]
 impl InterfaceWrapper {
     fn new(name: &'static str, dev: AxNetDevice, ether_addr: EthernetAddress) -> Self {