AHashMap and Refactor (#35)

Co-authored-by: Sajjad <sajjad@pourali.com>

Author: RoDmitry

Cargo.toml

@@ -10,6 +10,7 @@ repository = 'https://github.com/narrowlink/ipstack'
 readme = "README.md"
 
 [dependencies]
+ahash = "0.8"
 tokio = { version = "1.37", features = [
     "sync",
     "rt",

src/lib.rs

@@ -1,10 +1,12 @@
-pub use error::{IpStackError, Result};
+use crate::{
+    packet::IpStackPacketProtocol,
+    stream::{IpStackStream, IpStackTcpStream, IpStackUdpStream, IpStackUnknownTransport},
+};
+use ahash::AHashMap;
+use log::{error, trace};
 use packet::{NetworkPacket, NetworkTuple};
 use std::{
-    collections::{
-        hash_map::Entry::{Occupied, Vacant},
-        HashMap,
-    },
+    collections::hash_map::Entry::{Occupied, Vacant},
     time::Duration,
 };
 use tokio::{
@@ -13,16 +15,12 @@ use tokio::{
     sync::mpsc::{self, UnboundedReceiver, UnboundedSender},
 };
 
-use log::{error, trace};
-
-use crate::{
-    packet::IpStackPacketProtocol,
-    stream::{IpStackStream, IpStackTcpStream, IpStackUdpStream, IpStackUnknownTransport},
-};
 mod error;
 mod packet;
 pub mod stream;
 
+pub use self::error::{IpStackError, Result};
+
 const DROP_TTL: u8 = 0;
 
 #[cfg(unix)]
@@ -89,7 +87,8 @@ impl IpStack {
         let (accept_sender, accept_receiver) = mpsc::unbounded_channel::<IpStackStream>();
 
         tokio::spawn(async move {
-            let mut streams: HashMap<NetworkTuple, UnboundedSender<NetworkPacket>> = HashMap::new();
+            let mut streams: AHashMap<NetworkTuple, UnboundedSender<NetworkPacket>> =
+                AHashMap::new();
             let mut buffer = [0u8; u16::MAX as usize];
 
             let (pkt_sender, mut pkt_receiver) = mpsc::unbounded_channel::<NetworkPacket>();
@@ -102,7 +101,16 @@ impl IpStack {
                             continue;
                         };
                         if let IpStackPacketProtocol::Unknown = packet.transport_protocol() {
-                            accept_sender.send(IpStackStream::UnknownTransport(IpStackUnknownTransport::new(packet.src_addr().ip(),packet.dst_addr().ip(),packet.payload,&packet.ip,config.mtu,pkt_sender.clone())))?;
+                            accept_sender.send(
+                                IpStackStream::UnknownTransport(IpStackUnknownTransport::new(
+                                    packet.src_addr().ip(),
+                                    packet.dst_addr().ip(),
+                                    packet.payload,
+                                    &packet.ip,
+                                    config.mtu,
+                                    pkt_sender.clone()
+                                ))
+                            )?;
                             continue;
                         }
 
@@ -115,7 +123,14 @@ impl IpStack {
                             Vacant(entry) => {
                                 match packet.transport_protocol(){
                                     IpStackPacketProtocol::Tcp(h) => {
-                                        match IpStackTcpStream::new(packet.src_addr(),packet.dst_addr(),h, pkt_sender.clone(),config.mtu,config.tcp_timeout){
+                                        match IpStackTcpStream::new(
+                                            packet.src_addr(),
+                                            packet.dst_addr(),
+                                            h,
+                                            pkt_sender.clone(),
+                                            config.mtu,
+                                            config.tcp_timeout
+                                        ){
                                             Ok(stream) => {
                                                 entry.insert(stream.stream_sender());
                                                 accept_sender.send(IpStackStream::Tcp(stream))?;
@@ -130,7 +145,14 @@ impl IpStack {
                                         }
                                     }
                                     IpStackPacketProtocol::Udp => {
-                                        let stream = IpStackUdpStream::new(packet.src_addr(),packet.dst_addr(),packet.payload, pkt_sender.clone(),config.mtu,config.udp_timeout);
+                                        let stream = IpStackUdpStream::new(
+                                            packet.src_addr(),
+                                            packet.dst_addr(),
+                                            packet.payload,
+                                            pkt_sender.clone(),
+                                            config.mtu,
+                                            config.udp_timeout
+                                        );
                                         entry.insert(stream.stream_sender());
                                         accept_sender.send(IpStackStream::Udp(stream))?;
                                     }

src/packet.rs

@@ -1,8 +1,6 @@
-use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr};
-
-use etherparse::{Ipv4Header, Ipv6Header, NetSlice, SlicedPacket, TcpHeader, UdpHeader};
-
 use crate::error::IpStackError;
+use etherparse::{Ipv4Header, Ipv6Header, NetSlice, SlicedPacket, TcpHeader, UdpHeader};
+use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr};
 
 #[derive(Eq, Hash, PartialEq, Debug, Clone, Copy)]
 pub struct NetworkTuple {

src/stream/tcb.rs

@@ -1,9 +1,7 @@
+use crate::packet::TcpPacket;
 use std::{collections::BTreeMap, pin::Pin, time::Duration};
-
 use tokio::time::Sleep;
 
-use crate::packet::TcpPacket;
-
 const MAX_UNACK: u32 = 1024 * 16; // 16KB
 const READ_BUFFER_SIZE: usize = 1024 * 16; // 16KB
 

src/stream/tcp_wrapper.rs

@@ -1,18 +1,15 @@
+use super::tcp::IpStackTcpStream as IpStackTcpStreamInner;
+use crate::{
+    packet::{NetworkPacket, TcpPacket},
+    IpStackError,
+};
 use std::{net::SocketAddr, pin::Pin, time::Duration};
-
 use tokio::{
     io::AsyncWriteExt,
     sync::mpsc::{self, UnboundedSender},
     time::timeout,
 };
 
-use crate::{
-    packet::{NetworkPacket, TcpPacket},
-    IpStackError,
-};
-
-use super::tcp::IpStackTcpStream as IpStackTcpStreamInner;
-
 pub struct IpStackTcpStream {
     inner: Option<Box<IpStackTcpStreamInner>>,
     peer_addr: SocketAddr,

src/stream/unknown.rs

@@ -1,12 +1,10 @@
-use std::{io::Error, mem, net::IpAddr};
-
-use etherparse::{IpNumber, Ipv4Header, Ipv6FlowLabel, Ipv6Header};
-use tokio::sync::mpsc::UnboundedSender;
-
 use crate::{
     packet::{IpHeader, NetworkPacket, TransportHeader},
     TTL,
 };
+use etherparse::{IpNumber, Ipv4Header, Ipv6FlowLabel, Ipv6Header};
+use std::{io::Error, mem, net::IpAddr};
+use tokio::sync::mpsc::UnboundedSender;
 
 pub struct IpStackUnknownTransport {
     src_addr: IpAddr,