Merge pull request #22 from ssrlive/main

refine code

Cargo.toml

@@ -29,6 +29,7 @@ tracing = { version = "0.1", default-features = false, features = [
 
 [dev-dependencies]
 clap = { version = "4.5", features = ["derive"] }
+env_logger = "0.11"
 udp-stream = { version = "0.0", default-features = false }
 tokio = { version = "1.36", features = [
     "rt-multi-thread",
@@ -54,3 +55,7 @@ debug-assertions = false # Remove assertions from the binary.
 incremental = false      # Disable incremental compilation.
 overflow-checks = false  # Disable overflow checks.
 strip = true             # Automatically strip symbols from the binary.
+
+[[example]]
+name = "tun2"
+required-features = ["log"]

examples/tun2.rs

@@ -26,7 +26,7 @@
 //!
 
 use clap::Parser;
-use etherparse::{IcmpEchoHeader, Icmpv4Header};
+use etherparse::{IcmpEchoHeader, Icmpv4Header, IpNumber};
 use ipstack::stream::IpStackStream;
 use std::net::{Ipv4Addr, SocketAddr};
 use tokio::net::TcpStream;
@@ -35,18 +35,37 @@ use udp_stream::UdpStream;
 // const MTU: u16 = 1500;
 const MTU: u16 = u16::MAX;
 
+#[repr(C)]
+#[derive(Default, Debug, Copy, Clone, PartialEq, Eq, PartialOrd, Ord, clap::ValueEnum)]
+pub enum ArgVerbosity {
+    Off = 0,
+    Error,
+    Warn,
+    #[default]
+    Info,
+    Debug,
+    Trace,
+}
+
 #[derive(Parser)]
 #[command(author, version, about = "Testing app for tun.", long_about = None)]
 struct Args {
     /// echo server address, likes `127.0.0.1:8080`
     #[arg(short, long, value_name = "IP:port")]
     server_addr: SocketAddr,
+
+    /// Verbosity level
+    #[arg(short, long, value_name = "level", value_enum, default_value = "info")]
+    pub verbosity: ArgVerbosity,
 }
 
 #[tokio::main]
 async fn main() -> Result<(), Box<dyn std::error::Error>> {
     let args = Args::parse();
 
+    let default = format!("{:?}", args.verbosity);
+    env_logger::Builder::from_env(env_logger::Env::default().default_filter_or(default)).init();
+
     let ipv4 = Ipv4Addr::new(10, 0, 0, 33);
     let netmask = Ipv4Addr::new(255, 255, 255, 0);
     let gateway = Ipv4Addr::new(10, 0, 0, 1);
@@ -103,7 +122,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
                 });
             }
             IpStackStream::UnknownTransport(u) => {
-                if u.src_addr().is_ipv4() && u.ip_protocol() == 1.into() {
+                if u.src_addr().is_ipv4() && u.ip_protocol() == IpNumber::ICMP {
                     let (icmp_header, req_payload) = Icmpv4Header::from_slice(u.payload())?;
                     if let etherparse::Icmpv4Type::EchoRequest(req) = icmp_header.icmp_type {
                         println!("ICMPv4 echo");

examples/tun_wintun.rs

@@ -1,7 +1,7 @@
 use std::net::{Ipv4Addr, SocketAddr};
 
 use clap::Parser;
-use etherparse::{IcmpEchoHeader, Icmpv4Header};
+use etherparse::{IcmpEchoHeader, Icmpv4Header, IpNumber};
 use ipstack::stream::IpStackStream;
 use tokio::net::TcpStream;
 use udp_stream::UdpStream;
@@ -82,7 +82,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
                 });
             }
             IpStackStream::UnknownTransport(u) => {
-                if u.src_addr().is_ipv4() && u.ip_protocol() == 1.into() {
+                if u.src_addr().is_ipv4() && u.ip_protocol() == IpNumber::ICMP {
                     let (icmp_header, req_payload) = Icmpv4Header::from_slice(u.payload())?;
                     if let etherparse::Icmpv4Type::EchoRequest(req) = icmp_header.icmp_type {
                         println!("ICMPv4 echo");

src/error.rs

@@ -10,9 +10,6 @@ pub enum IpStackError {
     #[error("ValueTooBigError<u16> {0}")]
     ValueTooBigErrorU16(#[from] etherparse::err::ValueTooBigError<u16>),
 
-    #[error("From<ValueTooBigError<u32>> {0}")]
-    ValueTooBigErrorU32(#[from] etherparse::err::ValueTooBigError<u32>),
-
     #[error("ValueTooBigError<usize> {0}")]
     ValueTooBigErrorUsize(#[from] etherparse::err::ValueTooBigError<usize>),
 

src/lib.rs

@@ -109,22 +109,9 @@ impl IpStack {
 
                         match streams.entry(packet.network_tuple()){
                             Occupied(entry) =>{
-                                // let t = packet.transport_protocol();
                                 if let Err(_x) = entry.get().send(packet){
                                     #[cfg(feature = "log")]
                                     trace!("{}", _x);
-                                    // match t{
-                                    //     IpStackPacketProtocol::Tcp(_t) => {
-                                    //         // dbg!(t.flags());
-                                    //     }
-                                    //     IpStackPacketProtocol::Udp => {
-                                    //         // dbg!("udp");
-                                    //     }
-                                    //     IpStackPacketProtocol::Unknown => {
-                                    //         // dbg!("unknown");
-                                    //     }
-                                    // }
-
                                 }
                             }
                             Vacant(entry) => {
@@ -183,11 +170,11 @@ impl IpStack {
 
         IpStack { accept_receiver }
     }
+
     pub async fn accept(&mut self) -> Result<IpStackStream, IpStackError> {
-        if let Some(s) = self.accept_receiver.recv().await {
-            Ok(s)
-        } else {
-            Err(IpStackError::AcceptError)
-        }
+        self.accept_receiver
+            .recv()
+            .await
+            .ok_or(IpStackError::AcceptError)
     }
 }

src/packet.rs

@@ -4,7 +4,7 @@ use etherparse::{NetHeaders, PacketHeaders, TcpHeader, UdpHeader};
 
 use crate::error::IpStackError;
 
-#[derive(Eq, Hash, PartialEq, Debug)]
+#[derive(Eq, Hash, PartialEq, Debug, Clone, Copy)]
 pub struct NetworkTuple {
     pub src: SocketAddr,
     pub dst: SocketAddr,
@@ -21,18 +21,21 @@ pub mod tcp_flags {
     pub const FIN: u8 = 0b00000001;
 }
 
+#[derive(Debug, Clone)]
 pub(crate) enum IpStackPacketProtocol {
     Tcp(TcpPacket),
     Unknown,
     Udp,
 }
 
+#[derive(Debug, Clone)]
 pub(crate) enum TransportHeader {
     Tcp(TcpHeader),
     Udp(UdpHeader),
     Unknown,
 }
 
+#[derive(Debug, Clone)]
 pub struct NetworkPacket {
     pub(crate) ip: NetHeaders,
     pub(crate) transport: TransportHeader,
@@ -130,6 +133,7 @@ impl NetworkPacket {
     }
 }
 
+#[derive(Debug, Clone)]
 pub(super) struct TcpPacket {
     header: TcpHeader,
 }

src/stream/tcb.rs

@@ -1,8 +1,4 @@
-use std::{
-    collections::BTreeMap,
-    pin::Pin,
-    time::{Duration, SystemTime},
-};
+use std::{collections::BTreeMap, pin::Pin, time::Duration};
 
 use tokio::time::Sleep;
 
@@ -29,6 +25,7 @@ pub(super) enum PacketStatus {
     KeepAlive,
 }
 
+#[derive(Debug)]
 pub(super) struct Tcb {
     pub(super) seq: u32,
     pub(super) retransmission: Option<u32>,
@@ -41,21 +38,20 @@ pub(super) struct Tcb {
     state: TcpState,
     pub(super) avg_send_window: (u64, u64),
     pub(super) inflight_packets: Vec<InflightPacket>,
-    pub(super) unordered_packets: BTreeMap<u32, UnorderedPacket>,
+    unordered_packets: BTreeMap<u32, UnorderedPacket>,
 }
 
 impl Tcb {
     pub(super) fn new(ack: u32, tcp_timeout: Duration) -> Tcb {
         let seq = 100;
+        let deadline = tokio::time::Instant::now() + tcp_timeout;
         Tcb {
             seq,
             retransmission: None,
             ack,
             last_ack: seq,
             tcp_timeout,
-            timeout: Box::pin(tokio::time::sleep_until(
-                tokio::time::Instant::now() + tcp_timeout,
-            )),
+            timeout: Box::pin(tokio::time::sleep_until(deadline)),
             send_window: u16::MAX,
             recv_window: 0,
             state: TcpState::SynReceived(false),
@@ -176,9 +172,6 @@ impl Tcb {
         }
     }
     pub(super) fn change_last_ack(&mut self, ack: u32) {
-        self.timeout
-            .as_mut()
-            .reset(tokio::time::Instant::now() + self.tcp_timeout);
         let distance = ack.wrapping_sub(self.last_ack);
 
         if matches!(self.state, TcpState::Established) {
@@ -198,37 +191,44 @@ impl Tcb {
     pub fn is_send_buffer_full(&self) -> bool {
         self.seq.wrapping_sub(self.last_ack) >= MAX_UNACK
     }
+
+    pub(crate) fn reset_timeout(&mut self) {
+        let deadline = tokio::time::Instant::now() + self.tcp_timeout;
+        self.timeout.as_mut().reset(deadline);
+    }
 }
 
+#[derive(Debug, Clone)]
 pub struct InflightPacket {
     pub seq: u32,
     pub payload: Vec<u8>,
-    pub send_time: SystemTime,
+    // pub send_time: SystemTime, // todo
 }
 
 impl InflightPacket {
     fn new(seq: u32, payload: Vec<u8>) -> Self {
         Self {
             seq,
             payload,
-            send_time: SystemTime::now(),
+            // send_time: SystemTime::now(), // todo
         }
     }
     pub(crate) fn contains(&self, seq: u32) -> bool {
         self.seq < seq && self.seq + self.payload.len() as u32 >= seq
     }
 }
 
-pub struct UnorderedPacket {
-    pub payload: Vec<u8>,
-    pub recv_time: SystemTime,
+#[derive(Debug, Clone)]
+struct UnorderedPacket {
+    payload: Vec<u8>,
+    // pub recv_time: SystemTime, // todo
 }
 
 impl UnorderedPacket {
     pub(crate) fn new(payload: Vec<u8>) -> Self {
         Self {
             payload,
-            recv_time: SystemTime::now(),
+            // recv_time: SystemTime::now(), // todo
         }
     }
 }