TCP/IP Analysis support (#35)

Author: UwUDev

Cargo.toml

@@ -56,10 +56,16 @@ rpmalloc = { version = "0.2.2", optional = true }
 jemallocator = { package = "tikv-jemallocator", version = "0.6", optional = true }
 mimalloc = { version = "0.1.39", default-features = false, optional = true }
 
+[target.'cfg(target_os = "linux")'.dependencies]
+pcap = "2.0"
+pnet = "0.35"
+serde_json = "1.0"
+pktparse = { version = "0.7.1", features = ["serde"] }
+
 [target.'cfg(target_family = "unix")'.dependencies]
 daemonize = "0.5.0"
 nix = { version = "0.30.1", features = ["user", "signal"] }
-sysinfo = { version = "0.37", default-features = false, features = ["system"] }
+sysinfo = "0.30"
 
 [dev-dependencies]
 serde_json = "1.0"

src/main.rs

@@ -1,6 +1,6 @@
 #![cfg_attr(not(test), warn(unused_crate_dependencies))]
 
-pub mod alloc;
+mod alloc;
 #[cfg(target_family = "unix")]
 mod daemon;
 mod error;
@@ -22,24 +22,38 @@ pub struct Opt {
 #[derive(clap::Args, Clone)]
 pub struct Args {
     /// Debug mode
-    #[clap(long, default_value = "info", env = "PINGLY_LOG")]
+    #[arg(long, default_value = "info", env = "PINGLY_LOG")]
     pub log: String,
 
     /// Bind address
-    #[clap(short, long, default_value = "0.0.0.0:8181")]
+    #[arg(short, long, default_value = "0.0.0.0:8181")]
     pub bind: SocketAddr,
 
     /// Concurrent connections
-    #[clap(short, long, default_value = "1024")]
+    #[arg(short, long, default_value = "1024")]
     pub concurrent: usize,
 
+    /// Keep alive timeout (seconds)
+    #[arg(short, long, default_value = "60")]
+    pub keep_alive_timeout: u64,
+
     /// TLS certificate file path
-    #[clap(short = 'C', long)]
+    #[arg(short = 'C', long)]
     pub tls_cert: Option<PathBuf>,
 
     /// TLS private key file path (EC/PKCS8/RSA)
-    #[clap(short = 'K', long)]
+    #[arg(short = 'K', long)]
     pub tls_key: Option<PathBuf>,
+
+    /// Enable packet capture for TCP/IP analysis (requires root privileges)
+    #[cfg(target_os = "linux")]
+    #[arg(long, short = 'T')]
+    pub tcp_capture_packet: bool,
+
+    /// Network interface to capture packets from (default: auto-detect)
+    #[cfg(target_os = "linux")]
+    #[arg(long, short = 'I')]
+    pub tcp_capture_interface: Option<String>,
 }
 
 #[derive(Subcommand)]
@@ -70,6 +84,7 @@ pub enum Commands {
 
 fn main() -> Result<()> {
     let opt = Opt::parse();
+    #[cfg(target_family = "unix")]
     let daemon = daemon::Daemon::default();
     match opt.commands {
         Commands::Run(config) => server::run(config),

src/server/mod.rs

@@ -15,7 +15,8 @@ use axum::{
 use axum_extra::response::ErasedJson;
 use axum_server::Handle;
 use hyper_util::rt::TokioTimer;
-use tower::limit::ConcurrencyLimitLayer;
+use tower::{limit::ConcurrencyLimitLayer, ServiceBuilder};
+
 use tower_http::{
     cors::{AllowHeaders, AllowMethods, AllowOrigin, CorsLayer},
     trace::{DefaultMakeSpan, DefaultOnFailure, DefaultOnResponse, TraceLayer},
@@ -27,6 +28,9 @@ use tracker::{
     info::{ConnectionTrack, Track, TrackInfo},
 };
 
+#[cfg(target_os = "linux")]
+use tracker::capture::TcpCaptureTrack;
+
 use crate::{error::Error, Args, Result};
 
 #[tokio::main]
@@ -44,8 +48,8 @@ pub async fn run(args: Args) -> Result<()> {
     tracing::info!("Concurrent limit: {}", args.concurrent);
     tracing::info!("Bind address: {}", args.bind);
 
-    // init global layer provider
-    let global_layer = tower::ServiceBuilder::new()
+    // Init global layer
+    let global_layer = ServiceBuilder::new()
         .layer(
             TraceLayer::new_for_http()
                 .make_span_with(DefaultMakeSpan::new().level(Level::INFO))
@@ -61,17 +65,50 @@ pub async fn run(args: Args) -> Result<()> {
         )
         .layer(ConcurrencyLimitLayer::new(args.concurrent));
 
-    let router = Router::new()
+    // Create the router with the tracking endpoints
+    #[cfg_attr(not(target_os = "linux"), allow(unused_mut))]
+    let mut router = Router::new()
         .route("/api/all", any(track))
         .route("/api/tls", any(tls_track))
         .route("/api/http1", any(http1_track))
-        .route("/api/http2", any(http2_track))
-        .layer(global_layer);
+        .route("/api/http2", any(http2_track));
 
     // Signal the server to shutdown using Handle.
     let handle = Handle::new();
 
+    // Add TCP tracking layer
+    #[cfg(target_os = "linux")]
+    {
+        let mut tcp_capture_track: Option<TcpCaptureTrack> = None;
+        if args.tcp_capture_packet {
+            tracing::info!("Enabling TCP/IP packet capture (requires root)");
+            let capture = TcpCaptureTrack::new(128, args.bind.port());
+            if let Err(err) = capture.start_capture(args.tcp_capture_interface.clone()) {
+                tracing::error!("Failed to start TCP/IP packet capture: {err}");
+            } else {
+                if let Some(iface) = args.tcp_capture_interface {
+                    tracing::info!(
+                        "TCP/IP packet capture started successfully on interface {iface}"
+                    );
+                }
+                tcp_capture_track = Some(capture);
+            }
+        }
+
+        if let Some(capture) = tcp_capture_track.clone() {
+            router = router
+                .route("/api/tcp", any(tcp_track))
+                .layer(Extension(capture));
+        }
+
+        tokio::spawn(signal::graceful_shutdown(
+            handle.clone(),
+            tcp_capture_track.clone(),
+        ));
+    }
+
     // Spawn a task to gracefully shutdown server.
+    #[cfg(not(target_os = "linux"))]
     tokio::spawn(signal::graceful_shutdown(handle.clone()));
 
     // Load TLS configuration with HTTP/2 ALPN preference
@@ -99,7 +136,11 @@ pub async fn run(args: Args) -> Result<()> {
     server
         .handle(handle)
         .map(TrackAcceptor::new)
-        .serve(router.into_make_service_with_connect_info::<SocketAddr>())
+        .serve(
+            router
+                .layer(global_layer)
+                .into_make_service_with_connect_info::<SocketAddr>(),
+        )
         .await
         .map_err(Into::into)
 }
@@ -115,12 +156,42 @@ impl IntoResponse for Error {
 pub async fn track(
     Extension(ConnectInfo(addr)): Extension<ConnectInfo<SocketAddr>>,
     Extension(track): Extension<ConnectionTrack>,
+    #[cfg(target_os = "linux")] tcp_capture: Option<Extension<TcpCaptureTrack>>,
     req: Request<Body>,
 ) -> Result<ErasedJson> {
-    tokio::task::spawn_blocking(move || TrackInfo::new(Track::All, addr, req, track))
+    // get TCP packets if capture is available
+    #[cfg(target_os = "linux")]
+    let tcp_packets = if let Some(Extension(capture)) = tcp_capture {
+        // small delay to capture packets
+        tokio::time::sleep(Duration::from_millis(100)).await;
+
+        let client_ip = addr.ip().to_string();
+        let client_port = addr.port();
+
+        let packets = capture.get_packets_for_client(&client_ip, client_port);
+        capture.clear_packets_for_client(&client_ip, client_port);
+        packets
+    } else {
+        Vec::new()
+    };
+
+    #[cfg(target_os = "linux")]
+    {
+        tokio::task::spawn_blocking(move || {
+            TrackInfo::new_with_tcp(Track::All, addr, req, track, tcp_packets)
+        })
         .await
         .map(ErasedJson::pretty)
         .map_err(Error::from)
+    }
+
+    #[cfg(not(target_os = "linux"))]
+    {
+        tokio::task::spawn_blocking(move || TrackInfo::new(Track::All, addr, req, track))
+            .await
+            .map(ErasedJson::pretty)
+            .map_err(Error::from)
+    }
 }
 
 #[inline]
@@ -158,3 +229,21 @@ pub async fn http2_track(
         .map(ErasedJson::pretty)
         .map_err(Error::from)
 }
+
+#[inline]
+#[cfg(target_os = "linux")]
+pub async fn tcp_track(
+    Extension(ConnectInfo(addr)): Extension<ConnectInfo<SocketAddr>>,
+    Extension(capture): Extension<TcpCaptureTrack>,
+) -> Result<ErasedJson> {
+    tokio::time::sleep(Duration::from_millis(100)).await;
+
+    let client_ip = addr.ip().to_string();
+    let client_port = addr.port();
+
+    let packets = capture.get_packets_for_client(&client_ip, client_port);
+
+    capture.clear_packets_for_client(&client_ip, client_port);
+
+    Ok(ErasedJson::pretty(&packets))
+}

src/server/signal.rs

@@ -4,19 +4,37 @@ use axum_server::Handle;
 use tokio::time::sleep;
 use tracing::info;
 
-pub(super) async fn graceful_shutdown(handle: Handle) {
+#[cfg(target_os = "linux")]
+use super::tracker::capture::TcpCaptureTrack;
+
+pub(super) async fn graceful_shutdown(
+    handle: Handle,
+    #[cfg(target_os = "linux")] capture: Option<TcpCaptureTrack>,
+) {
     tokio::signal::ctrl_c()
         .await
         .expect("Ctrl+C signal hanlde error");
 
     info!("Ctrl+C signal received: starting graceful shutdown");
 
+    #[cfg(target_os = "linux")]
+    if let Some(capture) = capture {
+        capture.shutdown();
+    }
+
     // Signal the server to shutdown using Handle.
     handle.graceful_shutdown(Some(Duration::from_secs(1)));
 
     // Print alive connection count every second.
     loop {
         sleep(Duration::from_secs(1)).await;
-        info!("Alive connections: {}", handle.connection_count());
+        let connections = handle.connection_count();
+        info!("Alive connections: {}", connections);
+
+        // Exit the loop when all connections are closed
+        if connections == 0 {
+            info!("All connections closed, exiting gracefully");
+            break;
+        }
     }
 }