Track client addresses and add banned IPs

Attach client SocketAddr to connections and requests, enable per-server banned IP list, and improve worker wait/shutdown handling.

- engine.rs: store banned_ips in Hteapot, add add_banned_ip(), initialize banned_ips in constructors, accept listener (stream, addr), drop connections from banned addresses, store (TcpStream, SocketAddr) in the worker pool, persist addr in SocketData, and use HttpRequestBuilder::new_with_addr() when creating a SocketStatus. Also refine cvar.wait_while logic to honor the shutdown signal while waiting.
- request.rs: add addr: SocketAddr to HttpRequest, add new_with_addr() to HttpRequestBuilder to populate the request addr, and initialize default addr values in builders.

These changes let the server know the remote address for each request (so handlers can inspect or log it) and allow banning specific client IPs before they enter the worker pool. The wait logic change ensures worker threads properly observe shutdown while waiting for new connections.

Author: Az107

Cargo.lock

@@ -1,6 +1,6 @@
 [package]
 name = "hteapot"
-version = "0.6.5"
+version = "0.6.6"
 edition = "2024"
 authors = ["Alb Ruiz G. <me@albruiz.dev>"]
 description = "HTeaPot is a lightweight HTTP server library designed to be easy to use and extend."

Cargo.toml

@@ -1,6 +1,6 @@
 use std::collections::VecDeque;
 use std::io::{self, Read, Write};
-use std::net::{Shutdown, TcpListener, TcpStream};
+use std::net::{Shutdown, SocketAddr, TcpListener, TcpStream};
 
 use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::{Arc, Condvar, Mutex};
@@ -29,6 +29,7 @@ pub struct Hteapot {
     port: u16,
     address: String,
     threads: u16,
+    banned_ips: Vec<SocketAddr>,
     shutdown_signal: Option<Arc<AtomicBool>>,
     shutdown_hooks: Vec<Arc<dyn Fn() + Send + Sync + 'static>>,
 }
@@ -52,6 +53,7 @@ struct SocketStatus {
 /// Wraps a TCP stream and its associated state.
 struct SocketData {
     stream: TcpStream,
+    _addr: SocketAddr,
     status: Option<SocketStatus>,
 }
 
@@ -64,6 +66,10 @@ impl Hteapot {
         self.shutdown_signal.clone()
     }
 
+    pub fn add_banned_ip(&mut self, ip: SocketAddr) {
+        self.banned_ips.push(ip);
+    }
+
     pub fn add_shutdown_hook<F>(&mut self, hook: F)
     where
         F: Fn() + Send + Sync + 'static,
@@ -81,6 +87,7 @@ impl Hteapot {
             port,
             address: address.to_string(),
             threads: 1,
+            banned_ips: Vec::new(),
             shutdown_signal: None,
             shutdown_hooks: Vec::new(),
         }
@@ -91,6 +98,7 @@ impl Hteapot {
             port,
             address: address.to_string(),
             threads: if threads == 0 { 1 } else { threads },
+            banned_ips: Vec::new(),
             shutdown_signal: None,
             shutdown_hooks: Vec::new(),
         }
@@ -110,7 +118,7 @@ impl Hteapot {
             }
         };
 
-        let pool: Arc<(Mutex<VecDeque<TcpStream>>, Condvar)> =
+        let pool: Arc<(Mutex<VecDeque<(TcpStream, SocketAddr)>>, Condvar)> =
             Arc::new((Mutex::new(VecDeque::new()), Condvar::new()));
         let priority_list: Arc<Mutex<Vec<usize>>> =
             Arc::new(Mutex::new(vec![0; self.threads as usize]));
@@ -120,10 +128,10 @@ impl Hteapot {
         let shutdown_signal = self.shutdown_signal.clone();
         let shutdown_hooks = Arc::new(self.shutdown_hooks.clone());
 
-        for thread_index in 0..self.threads {
+        for _thread_index in 0..self.threads {
             let pool_clone = pool.clone();
             let action_clone = arc_action.clone();
-            let priority_list_clone = priority_list.clone();
+            let _priority_list_clone = priority_list.clone();
             let shutdown_signal_clone = shutdown_signal.clone();
 
             thread::spawn(move || {
@@ -132,28 +140,33 @@ impl Hteapot {
                     {
                         let (lock, cvar) = &*pool_clone;
                         let mut pool = lock.lock().expect("Error locking pool");
+
                         if streams_to_handle.is_empty() {
                             // Store the returned guard back into pool
-                            pool = cvar
-                                .wait_while(pool, |pool| pool.is_empty())
-                                .expect("Error waiting on cvar");
-                        }
-                        //TODO: move this to allow process the last request
-                        if let Some(signal) = &shutdown_signal_clone {
-                            if !signal.load(Ordering::SeqCst) {
-                                break; // Exit the server loop
-                            }
+                            pool = if let Some(signal) = &shutdown_signal_clone {
+                                if !signal.load(Ordering::SeqCst) {
+                                    break;
+                                }
+                                cvar.wait_while(pool, |pool| {
+                                    pool.is_empty() && signal.load(Ordering::SeqCst)
+                                })
+                                .expect("Error waiting on cvar")
+                            } else {
+                                cvar.wait_while(pool, |pool| pool.is_empty())
+                                    .expect("Error waiting on cvar")
+                            };
                         }
 
-                        while let Some(stream) = pool.pop_back() {
+                        while let Some((stream, addr)) = pool.pop_back() {
                             let socket_status = SocketStatus {
                                 ttl: Instant::now(),
                                 status: Status::Read,
                                 response: Box::new(EmptyHttpResponse {}),
-                                request: HttpRequestBuilder::new(),
+                                request: HttpRequestBuilder::new_with_addr(addr),
                                 index_writed: 0,
                             };
                             let socket_data = SocketData {
+                                _addr: addr,
                                 stream,
                                 status: Some(socket_status),
                             };
@@ -190,10 +203,15 @@ impl Hteapot {
                     break;
                 }
             }
-            let stream = match listener.accept() {
-                Ok((stream, _)) => stream,
+
+            let (stream, addr) = match listener.accept() {
+                Ok((stream, addr)) => (stream, addr),
                 Err(_) => continue,
             };
+            if self.banned_ips.contains(&addr) {
+                let _ = stream.shutdown(Shutdown::Both);
+                continue;
+            }
 
             if stream.set_nonblocking(true).is_err() {
                 eprintln!("Error setting non-blocking mode on stream");
@@ -209,7 +227,7 @@ impl Hteapot {
                 let mut pool = lock.lock().expect("Error locking pool");
 
                 // Add the connection to the pool for the least-loaded thread
-                pool.push_front(stream);
+                pool.push_front((stream, addr));
                 cvar.notify_one();
             }
         }

src/hteapot/engine.rs

@@ -10,6 +10,7 @@
 use super::HttpHeaders;
 use super::HttpMethod;
 use std::hash::Hash;
+use std::net::SocketAddr;
 use std::{cmp::min, collections::HashMap, net::TcpStream, str};
 
 const MAX_HEADER_SIZE: usize = 1024 * 16;
@@ -25,6 +26,7 @@ pub struct HttpRequest {
     pub args: HashMap<String, String>,
     pub headers: HttpHeaders,
     pub body: Vec<u8>,
+    pub addr: SocketAddr,
     stream: Option<TcpStream>,
 }
 
@@ -60,6 +62,8 @@ impl HttpRequest {
             args: HashMap::new(),
             headers: HttpHeaders::new(),
             body: Vec::new(),
+            addr: SocketAddr::from(([0, 0, 0, 0], 0)),
+
             stream: None,
         };
     }
@@ -71,6 +75,8 @@ impl HttpRequest {
             path: String::new(),
             args: HashMap::new(),
             headers: HttpHeaders::new(),
+            addr: SocketAddr::from(([0, 0, 0, 0], 0)),
+
             body: Vec::new(),
             stream: None,
         }
@@ -83,6 +89,7 @@ impl HttpRequest {
             path: self.path.clone(),
             args: self.args.clone(),
             headers: self.headers.clone(),
+            addr: self.addr.clone(),
             body: self.body.clone(),
             stream: None,
         };
@@ -135,6 +142,25 @@ impl HttpRequestBuilder {
                 args: HashMap::new(),
                 headers: HttpHeaders::new(),
                 body: Vec::new(),
+                addr: SocketAddr::from(([0, 0, 0, 0], 0)),
+                stream: None,
+            },
+            chunked: false,
+            state: State::Init,
+            body_size: 0,
+            buffer: Vec::new(),
+        };
+    }
+
+    pub fn new_with_addr(addr: SocketAddr) -> Self {
+        return HttpRequestBuilder {
+            request: HttpRequest {
+                method: HttpMethod::GET,
+                path: String::new(),
+                args: HashMap::new(),
+                headers: HttpHeaders::new(),
+                body: Vec::new(),
+                addr,
                 stream: None,
             },
             chunked: false,
@@ -152,30 +178,6 @@ impl HttpRequestBuilder {
         }
     }
 
-    /// Reads bytes into the request body based on `Content-Length`.
-    fn read_body_len(&mut self) -> Option<()> {
-        let body_left = self.body_size.saturating_sub(self.request.body.len());
-
-        let body_left = self.body_size.saturating_sub(self.request.body.len());
-
-        if body_left > 0 {
-            return None;
-        } else {
-            return Some(());
-        }
-    }
-
-    /// Placeholder for future support of chunked body parsing.
-    fn _read_body_chunk(&mut self) -> Option<()> {
-        //TODO: this will support chunked body in the future
-        todo!()
-    }
-
-    /// Main entry point for reading the request body.
-    fn read_body(&mut self) -> Option<()> {
-        return self.read_body_len();
-    }
-
     pub fn done(&self) -> bool {
         self.state == State::Finish
     }

src/hteapot/request.rs

@@ -218,6 +218,7 @@ fn main() {
 
         let response = handlers.get_handler(&ctx);
         if response.is_none() {
+            logger.error("No handler found for request".to_string());
             return HttpResponse::new(HttpStatus::InternalServerError, "content", None);
         }
         let response = response.unwrap().run(&mut ctx);