update example

narrowlink · Jan 22, 2024 · b5c9ecc · b5c9ecc
1 parent 9be0637
commit b5c9ecc
Show file tree

Hide file tree

Showing 2 changed files with 50 additions and 31 deletions.
diff --git a/README.md b/README.md
@@ -6,14 +6,8 @@ Unstable, under development.
 use std::net::{IpAddr, Ipv4Addr, SocketAddr};
 use udp_stream::UdpStream;
 use tokio::io{AsyncRead, AsyncWrite};
-async fn copy_from_lhs_to_rhs(lhs: impl AsyncRead + AsyncWrite, rhs: impl AsyncRead + AsyncWrite) {
-    let (mut lhs_reader, mut lhs_writer) = tokio::io::split(lhs);
-    let (mut rhs_reader, mut rhs_writer) = tokio::io::split(rhs);
-    let _r = tokio::join! {
-        tokio::io::copy(&mut lhs_reader, &mut rhs_writer) ,
-        tokio::io::copy(&mut rhs_reader, &mut lhs_writer),
-    };
-}
+use etherparse::{IcmpEchoHeader, Icmpv4Header};
+
 #[tokio::main]
 async fn main(){
     const MTU: u16 = 1500;
@@ -30,18 +24,44 @@ async fn main(){
 
     while let Ok(stream) = ip_stack.accept().await {
         match stream {
-            IpStackStream::Tcp(tcp) => {
-                let rhs = TcpStream::connect("1.1.1.1:80").await.unwrap();
+            IpStackStream::Tcp(mut tcp) => {
+                let mut rhs = TcpStream::connect("1.1.1.1:80").await.unwrap();
                 tokio::spawn(async move {
-                    copy_from_lhs_to_rhs(tcp, rhs).await;
+                    let _ = tokio::io::copy_bidirectional(& mut tcp, & mut rhs).await;
                 });
             }
-            IpStackStream::Udp(udp) => {
+            IpStackStream::Udp(mut udp) => {
                 let addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(1, 1, 1, 1)), 53);
-                let rhs = UdpStream::connect(addr).await.unwrap();
+                let mut rhs = UdpStream::connect(addr).await.unwrap();
                 tokio::spawn(async move {
-                    copy_from_lhs_to_rhs(udp, rhs).await;
+                    let _ = tokio::io::copy_bidirectional(& mut udp, & mut rhs).await;
                 });
+            }
+			IpStackStream::UnknownTransport(u) => {
+                if u.src_addr().is_ipv4() && u.ip_protocol() == 1 {
+                    let (icmp_header, req_payload) = Icmpv4Header::from_slice(u.payload())?;
+                    if let etherparse::Icmpv4Type::EchoRequest(req) = icmp_header.icmp_type {
+                        println!("ICMPv4 echo");
+                        let echo = IcmpEchoHeader {
+                            id: req.id,
+                            seq: req.seq,
+                        };
+                        let mut resp = Icmpv4Header::new(etherparse::Icmpv4Type::EchoReply(echo));
+                        resp.update_checksum(req_payload);
+                        let mut payload = resp.to_bytes().to_vec();
+                        payload.extend_from_slice(req_payload);
+                        u.send(payload).await?;
+                    } else {
+                        println!("ICMPv4");
+                    }
+                    continue;
+                }
+                println!("unknown transport - Ip Protocol {}", u.ip_protocol());
+                continue;
+            }
+            IpStackStream::UnknownNetwork(pkt) => {
+                println!("unknown transport - {} bytes", pkt.len());
+                continue;
             }
         }
     }

diff --git a/examples/tun.rs b/examples/tun.rs
@@ -29,7 +29,7 @@ use clap::Parser;
 use etherparse::{IcmpEchoHeader, Icmpv4Header};
 use ipstack::stream::IpStackStream;
 use std::net::{Ipv4Addr, SocketAddr};
-use tokio::{join, net::TcpStream};
+use tokio::net::TcpStream;
 use udp_stream::UdpStream;
 
 // const MTU: u16 = 1500;
@@ -72,41 +72,40 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
 
     let mut ip_stack = ipstack::IpStack::new(ipstack_config, tun2::create_as_async(&config)?);
 
+    #[cfg(target_os = "macos")]
+    std::process::Command::new("sh")
+        .arg("-c")
+        .arg("sudo route -n add -net 10.0.0.0/24 10.0.0.33")
+        .output()
+        .unwrap();
+
     let server_addr = args.server_addr;
 
     loop {
         match ip_stack.accept().await? {
-            IpStackStream::Tcp(tcp) => {
+            IpStackStream::Tcp(mut tcp) => {
                 let s = TcpStream::connect(server_addr).await;
                 if let Err(ref err) = s {
                     println!("connect TCP server failed \"{}\"", err);
                     continue;
                 }
                 println!("==== New TCP connection ====");
-                let (mut t_rx, mut t_tx) = tokio::io::split(tcp);
-                let (mut s_rx, mut s_tx) = tokio::io::split(s?);
+                let mut s = s?;
                 tokio::spawn(async move {
-                    let _r = join! {
-                         tokio::io::copy(&mut t_rx, &mut s_tx) ,
-                         tokio::io::copy(&mut s_rx, &mut t_tx),
-                    };
+                    let _ = tokio::io::copy_bidirectional(&mut tcp, &mut s).await;
                     println!("====== end tcp connection ======");
                 });
             }
-            IpStackStream::Udp(udp) => {
+            IpStackStream::Udp(mut udp) => {
                 let s = UdpStream::connect(server_addr).await;
                 if let Err(ref err) = s {
                     println!("connect UDP server failed \"{}\"", err);
                     continue;
                 }
                 println!("==== New UDP connection ====");
-                let (mut t_rx, mut t_tx) = tokio::io::split(udp);
-                let (mut s_rx, mut s_tx) = tokio::io::split(s?);
+                let mut s = s?;
                 tokio::spawn(async move {
-                    let _r = join! {
-                         tokio::io::copy(&mut t_rx, &mut s_tx) ,
-                         tokio::io::copy(&mut s_rx, &mut t_tx),
-                    };
+                    let _ = tokio::io::copy_bidirectional(&mut udp, &mut s).await;
                     println!("==== end UDP connection ====");
                 });
             }
@@ -132,8 +131,8 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
                 println!("unknown transport - Ip Protocol {}", u.ip_protocol());
                 continue;
             }
-            IpStackStream::UnknownNetwork(payload) => {
-                println!("unknown transport - {} bytes", payload.len());
+            IpStackStream::UnknownNetwork(pkt) => {
+                println!("unknown transport - {} bytes", pkt.len());
                 continue;
             }
         };