itsi 0.1.20 → 0.2.2

data/crates/itsi_server/src/server/serve_strategy/cluster_mode.rs

@@ -2,9 +2,7 @@ use crate::ruby_types::itsi_server::itsi_server_config::ItsiServerConfig;
 use crate::server::signal::SIGNAL_HANDLER_CHANNEL;
 use crate::server::{lifecycle_event::LifecycleEvent, process_worker::ProcessWorker};
 use itsi_error::{ItsiError, Result};
-use itsi_rb_helpers::{
-    call_proc_and_log_errors, call_with_gvl, call_without_gvl, create_ruby_thread,
-};
+use itsi_rb_helpers::{call_with_gvl, call_without_gvl, create_ruby_thread};
 use itsi_tracing::{error, info, warn};
 use magnus::Value;
 use nix::{libc::exit, unistd::Pid};
@@ -56,6 +54,12 @@ impl ClusterMode {
             .expect("Failed to build Tokio runtime")
     }
 
+    pub fn invoke_hook(&self, hook_name: &str) {
+        if let Some(hook) = self.server_config.server_params.read().hooks.get(hook_name) {
+            call_with_gvl(|_| hook.call::<_, Value>(()).ok());
+        }
+    }
+
     #[allow(clippy::await_holding_lock)]
     pub async fn handle_lifecycle_event(
         self: Arc<Self>,
@@ -70,10 +74,12 @@ impl ClusterMode {
             LifecycleEvent::Shutdown => {
                 self.server_config.stop_watcher()?;
                 self.shutdown().await?;
+                self.invoke_hook("before_shutdown");
                 Ok(())
             }
             LifecycleEvent::Restart => {
-                if self.server_config.check_config() {
+                if self.server_config.check_config().await {
+                    self.invoke_hook("before_restart");
                     self.server_config.dup_fds()?;
                     self.shutdown().await.ok();
                     info!("Shutdown complete. Calling reload exec");
@@ -82,7 +88,7 @@ impl ClusterMode {
                 Ok(())
             }
             LifecycleEvent::Reload => {
-                if !self.server_config.check_config() {
+                if !self.server_config.check_config().await {
                     return Ok(());
                 }
                 let should_reexec = self.server_config.clone().reload(true)?;
@@ -268,15 +274,7 @@ impl ClusterMode {
     #[instrument(skip(self), fields(mode = "cluster", pid=format!("{:?}", Pid::this())))]
     pub fn run(self: Arc<Self>) -> Result<()> {
         info!("Starting in Cluster mode");
-        if let Some(proc) = self
-            .server_config
-            .server_params
-            .read()
-            .hooks
-            .get("before_fork")
-        {
-            call_with_gvl(|_| call_proc_and_log_errors(proc.clone()))
-        }
+        self.invoke_hook("before_fork");
         self.process_workers
             .lock()
             .iter()
@@ -290,7 +288,15 @@ impl ClusterMode {
 
         self.build_runtime().block_on(async {
           let self_ref = self_ref.clone();
-          let mut memory_check_interval = time::interval(time::Duration::from_secs(2));
+          let memory_check_duration = if self_ref.server_config.server_params.read().worker_memory_limit.is_some(){
+            time::Duration::from_secs(15)
+          } else {
+            time::Duration::from_secs(60 * 60 * 24 * 365 * 100)
+          };
+
+          let mut memory_check_interval = time::interval(memory_check_duration);
+
+          self.invoke_hook("after_start");
 
           loop {
             tokio::select! {
@@ -314,7 +320,7 @@ impl ClusterMode {
                     if let Some(current_mem_usage) = largest_worker.memory_usage(){
                       if current_mem_usage > memory_limit {
                         largest_worker.reboot(self_ref.clone()).await.ok();
-                        if let Some(hook) = self_ref.server_config.server_params.read().hooks.get("after_memory_threshold_reached") {
+                        if let Some(hook) = self_ref.server_config.server_params.read().hooks.get("after_memory_limit_reached") {
                           call_with_gvl(|_|  hook.call::<_, Value>((largest_worker.pid(),)).ok() );
                         }
                       }

data/crates/itsi_server/src/server/serve_strategy/single_mode.rs

@@ -19,7 +19,7 @@ use itsi_rb_helpers::{
     call_with_gvl, call_without_gvl, create_ruby_thread, funcall_no_ret, print_rb_backtrace,
 };
 use itsi_tracing::{debug, error, info};
-use magnus::value::ReprValue;
+use magnus::{value::ReprValue, Value};
 use nix::unistd::Pid;
 use parking_lot::RwLock;
 use std::{
@@ -83,8 +83,7 @@ impl SingleMode {
         builder
             .thread_name("itsi-server-accept-loop")
             .thread_stack_size(3 * 1024 * 1024)
-            .enable_io()
-            .enable_time()
+            .enable_all()
             .build()
             .expect("Failed to build Tokio runtime")
     }
@@ -160,12 +159,12 @@ impl SingleMode {
     pub fn start_monitors(
         self: Arc<Self>,
         thread_workers: Arc<Vec<Arc<ThreadWorker>>>,
-    ) -> magnus::Thread {
+    ) -> Option<magnus::Thread> {
         call_with_gvl(move |_| {
             create_ruby_thread(move || {
                 call_without_gvl(move || {
                     let monitor_runtime = RuntimeBuilder::new_current_thread()
-                        .enable_time()
+                        .enable_all()
                         .build()
                         .unwrap();
                     let receiver = self.clone();
@@ -190,10 +189,10 @@ impl SingleMode {
                                   lifecycle_event = lifecycle_rx.recv() => {
                                       match lifecycle_event {
                                           Ok(LifecycleEvent::Restart) => {
-                                              receiver.restart().ok();
+                                              receiver.restart().await.ok();
                                           }
                                           Ok(LifecycleEvent::Reload) => {
-                                              receiver.reload().ok();
+                                              receiver.reload().await.ok();
                                           }
                                           Ok(LifecycleEvent::Shutdown) => {
                                             break;
@@ -233,6 +232,11 @@ impl SingleMode {
 
         let (shutdown_sender, _) = watch::channel(RunningPhase::Running);
         let monitor_thread = self.clone().start_monitors(thread_workers.clone());
+        if monitor_thread.is_none() {
+            error!("Failed to start monitor thread");
+            return Err(ItsiError::new("Failed to start monitor thread"));
+        }
+        let monitor_thread = monitor_thread.unwrap();
         if SHUTDOWN_REQUESTED.load(Ordering::SeqCst) {
             return Ok(());
         }
@@ -290,6 +294,7 @@ impl SingleMode {
                             }
                             lifecycle_event = lifecycle_rx.recv() => match lifecycle_event{
                               Ok(LifecycleEvent::Shutdown) => {
+                                debug!("Received lifecycle event: {:?}", lifecycle_event);
                                 shutdown_sender.send(RunningPhase::ShutdownPending).unwrap();
                                 tokio::time::sleep(Duration::from_millis(25)).await;
                                 for _i in 0..workers_clone.len() {
@@ -303,11 +308,23 @@ impl SingleMode {
                           }
                         }
                     }
-                    while let Some(_res) = acceptor_task_set.join_next().await {}
+
+                    let deadline = Instant::now()
+                        + Duration::from_secs_f64(self_ref.server_config.server_params.read().shutdown_timeout);
+                    tokio::select! {
+                        _ = async {
+                            while let Some(_res) = acceptor_task_set.join_next().await {}
+                        } => {},
+                        _ = tokio::time::sleep_until(tokio::time::Instant::from_std(deadline)) => {},
+                    }
                 });
 
               }
 
+              if self.is_single_mode() {
+                self.invoke_hook("after_start");
+              }
+
               while let Some(_res) = listener_task_set.join_next().await {}
 
               // Explicitly drop all listeners to ensure file descriptors are released
@@ -346,15 +363,24 @@ impl SingleMode {
             sleep(Duration::from_millis(50));
         }
 
+        if self.is_single_mode() {
+            self.invoke_hook("before_shutdown");
+        }
+
         if self.restart_requested.load(Ordering::SeqCst) {
             self.restart_requested.store(false, Ordering::SeqCst);
             info!("Worker restarting");
             self.run()?;
         }
+
         debug!("Runtime has shut down");
         result
     }
 
+    pub fn is_single_mode(&self) -> bool {
+        self.server_config.server_params.read().workers == 1
+    }
+
     pub(crate) async fn serve_connection(
         &self,
         stream: IoStream,
@@ -417,12 +443,15 @@ impl SingleMode {
     /// Attempts to reload the config "live"
     /// Not that when running in single mode this will not unload
     /// old code. If you need a clean restart, use the `restart` (SIGHUP) method instead
-    pub fn reload(&self) -> Result<()> {
-        if !self.server_config.check_config() {
+    pub async fn reload(&self) -> Result<()> {
+        if !self.server_config.check_config().await {
             return Ok(());
         }
         let should_reexec = self.server_config.clone().reload(false)?;
         if should_reexec {
+            if self.is_single_mode() {
+                self.invoke_hook("before_restart");
+            }
             self.server_config.dup_fds()?;
             self.server_config.reload_exec()?;
         }
@@ -432,11 +461,19 @@ impl SingleMode {
         Ok(())
     }
 
+    pub fn invoke_hook(&self, hook_name: &str) {
+        if let Some(hook) = self.server_config.server_params.read().hooks.get(hook_name) {
+            call_with_gvl(|_| hook.call::<_, Value>(()).ok());
+        }
+    }
     /// Restart the server while keeping connections open.
-    pub fn restart(&self) -> Result<()> {
-        if !self.server_config.check_config() {
+    pub async fn restart(&self) -> Result<()> {
+        if !self.server_config.check_config().await {
             return Ok(());
         }
+        if self.is_single_mode() {
+            self.invoke_hook("before_restart");
+        }
         self.server_config.dup_fds()?;
         self.server_config.reload_exec()?;
         Ok(())

data/crates/itsi_server/src/server/signal.rs

@@ -49,6 +49,7 @@ fn receive_signal(signum: i32, _: sighandler_t) {
 pub fn reset_signal_handlers() -> bool {
     SIGINT_COUNT.store(0, std::sync::atomic::Ordering::SeqCst);
     SHUTDOWN_REQUESTED.store(false, std::sync::atomic::Ordering::SeqCst);
+
     unsafe {
         libc::signal(libc::SIGTERM, receive_signal as usize);
         libc::signal(libc::SIGINT, receive_signal as usize);

data/crates/itsi_server/src/server/size_limited_incoming.rs

@@ -10,6 +10,7 @@ use std::sync::atomic::AtomicUsize;
 use std::sync::atomic::Ordering;
 use std::task::Context;
 use std::task::Poll;
+use tracing::debug;
 
 /// Custom error to indicate that the maximum body size was exceeded.
 #[derive(Debug)]
@@ -71,6 +72,11 @@ where
                         Ok(data) => {
                             let len = data.remaining();
                             self.current += len;
+                            debug!(
+                              target: "option::max_body",
+                              "current: {}, limit: {}",
+                              self.current, self.limit.load(Ordering::Relaxed)
+                            );
                             if self.current > self.limit.load(Ordering::Relaxed) {
                                 Poll::Ready(Some(Err(Box::new(MaxBodySizeReached))))
                             } else {

data/crates/itsi_server/src/server/thread_worker.rs

@@ -1,4 +1,5 @@
 use async_channel::Sender;
+use itsi_error::ItsiError;
 use itsi_rb_helpers::{
     call_with_gvl, call_without_gvl, create_ruby_thread, kill_threads, HeapValue,
 };
@@ -159,7 +160,7 @@ impl ThreadWorker {
     pub fn poll_shutdown(&self, deadline: Instant) -> bool {
         if let Some(thread) = self.thread.read().deref() {
             if Instant::now() > deadline {
-                warn!("Worker shutdown timed out. Killing thread {:?}", thread);
+                debug!("Worker shutdown timed out. Killing thread {:?}", thread);
                 self.terminated.store(true, Ordering::SeqCst);
                 kill_threads(vec![thread.as_value()]);
             }
@@ -198,6 +199,9 @@ impl ThreadWorker {
                         self_ref.accept_loop(params, name, receiver, terminated);
                     }
                 })
+                .ok_or_else(|| {
+                    ItsiError::InternalServerError("Failed to create Ruby thread".to_owned())
+                })?
                 .into(),
             );
             Ok::<(), magnus::Error>(())

data/crates/itsi_server/src/services/itsi_http_service.rs

@@ -1,5 +1,5 @@
 use crate::default_responses::{NOT_FOUND_RESPONSE, TIMEOUT_RESPONSE};
-use crate::ruby_types::itsi_server::itsi_server_config::ServerParams;
+use crate::ruby_types::itsi_server::itsi_server_config::{ItsiServerTokenPreference, ServerParams};
 use crate::server::binds::listener::ListenerInfo;
 use crate::server::http_message_types::{ConversionExt, HttpResponse, RequestExt, ResponseFormat};
 use crate::server::lifecycle_event::LifecycleEvent;
@@ -18,6 +18,7 @@ use itsi_error::ItsiError;
 use regex::Regex;
 use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::OnceLock;
+use tracing::error;
 
 use std::{future::Future, ops::Deref, pin::Pin, sync::Arc};
 use tokio::sync::watch::{self};
@@ -153,6 +154,10 @@ impl HttpRequestContext {
     }
 }
 
+const SERVER_TOKEN_VERSION: HeaderValue =
+    HeaderValue::from_static(concat!("Itsi/", env!("CARGO_PKG_VERSION")));
+const SERVER_TOKEN_NAME: HeaderValue = HeaderValue::from_static("Itsi");
+
 impl Service<Request<Incoming>> for ItsiHttpService {
     type Response = HttpResponse;
     type Error = ItsiError;
@@ -195,7 +200,10 @@ impl Service<Request<Incoming>> for ItsiHttpService {
                         depth = index;
                         break;
                     }
-                    Err(e) => return Err(e.into()),
+                    Err(e) => {
+                        error!("Middleware error: {}", e);
+                        break;
+                    }
                 }
             }
 
@@ -208,6 +216,16 @@ impl Service<Request<Incoming>> for ItsiHttpService {
                 resp = elm.after(resp, &mut context).await;
             }
 
+            match params.itsi_server_token_preference {
+                ItsiServerTokenPreference::Version => {
+                    resp.headers_mut().insert("Server", SERVER_TOKEN_VERSION);
+                }
+                ItsiServerTokenPreference::Name => {
+                    resp.headers_mut().insert("Server", SERVER_TOKEN_NAME);
+                }
+                ItsiServerTokenPreference::None => {}
+            }
+
             Ok(resp)
         };
 

data/crates/itsi_server/src/services/rate_limiter.rs

@@ -5,10 +5,12 @@ use redis::{Client, RedisError, Script};
 use serde::Deserialize;
 use std::any::Any;
 use std::collections::{HashMap, HashSet};
+use std::result::Result;
 use std::sync::{Arc, LazyLock, Mutex};
 use std::time::{Duration, Instant};
 use tokio::sync::{Mutex as AsyncMutex, RwLock};
 use tokio::time::timeout;
+use tracing::warn;
 use url::Url;
 
 #[derive(Debug)]
@@ -109,10 +111,19 @@ impl RedisRateLimiter {
         // Create the Lua script once when initializing the rate limiter
         let increment_script = Script::new(
             r#"
+            -- Increment the counter
             local current = redis.call('INCR', KEYS[1])
-            if redis.call('TTL', KEYS[1]) < 0 then
-                redis.call('EXPIRE', KEYS[1], ARGV[1])
+
+            -- Fetch existing TTL
+            local ttl = redis.call('TTL', KEYS[1])
+            if ttl < 0 then
+                -- If no TTL is set, apply the window interval
+                local window = tonumber(ARGV[1])
+                redis.call('EXPIRE', KEYS[1], window)
+                ttl = window
             end
+
+            -- Return both the current count and remaining TTL
             return { current, ttl }
             "#,
         );
@@ -135,14 +146,14 @@ impl RedisRateLimiter {
         let mut connection = (*self.connection).clone();
 
         // Set the ban with the reason as the value
-        let _: () = redis::cmd("SET")
+        let _: Result<(), RedisError> = redis::cmd("SET")
             .arg(&ban_key)
             .arg(reason)
             .arg("EX")
             .arg(timeout_secs)
             .query_async(&mut connection)
             .await
-            .map_err(RateLimitError::RedisError)?;
+            .inspect_err(|e| warn!("Exception banning IP {:?}", e));
 
         Ok(())
     }

data/crates/itsi_server/src/services/static_file_server.rs

@@ -1,9 +1,10 @@
 use crate::{
-    default_responses::{INTERNAL_SERVER_ERROR_RESPONSE, NOT_FOUND_RESPONSE},
+    default_responses::NOT_FOUND_RESPONSE,
     prelude::*,
     server::{
         http_message_types::{HttpRequest, HttpResponse, RequestExt, ResponseFormat},
         middleware_stack::ErrorResponse,
+        redirect_type::RedirectType,
     },
 };
 use base64::{engine::general_purpose, Engine};
@@ -46,6 +47,7 @@ pub static ROOT_STATIC_FILE_SERVER: LazyLock<StaticFileServer> = LazyLock::new(|
         recheck_interval: Duration::from_secs(1),
         try_html_extension: true,
         auto_index: true,
+        headers: None,
         not_found_behavior: NotFoundBehavior::Error(ErrorResponse::not_found()),
         serve_hidden_files: false,
         allowed_extensions: vec!["html".to_string(), "css".to_string(), "js".to_string()],
@@ -56,6 +58,7 @@ pub static ROOT_STATIC_FILE_SERVER: LazyLock<StaticFileServer> = LazyLock::new(|
 #[derive(Debug, Clone, Deserialize)]
 pub struct Redirect {
     pub to: String,
+    pub r#type: RedirectType,
 }
 
 #[derive(Debug, Clone, Deserialize)]
@@ -68,8 +71,6 @@ pub enum NotFoundBehavior {
     IndexFile(PathBuf),
     #[serde(rename = "redirect")]
     Redirect(Redirect),
-    #[serde(rename = "internal_server_error")]
-    InternalServerError,
 }
 
 #[derive(Debug, Clone)]
@@ -81,6 +82,7 @@ pub struct StaticFileServerConfig {
     pub try_html_extension: bool,
     pub auto_index: bool,
     pub not_found_behavior: NotFoundBehavior,
+    pub headers: Option<HashMap<String, String>>,
     pub serve_hidden_files: bool,
     pub allowed_extensions: Vec<String>,
 }
@@ -113,23 +115,26 @@ struct CacheEntry {
 }
 
 impl CacheEntry {
-    pub fn suggest_content_for(&self, supported_encodings: &[HeaderValue]) -> (Arc<Bytes>, &str) {
+    pub fn suggest_content_for(
+        &self,
+        supported_encodings: &[HeaderValue],
+    ) -> (Arc<Bytes>, Option<&str>) {
         for encoding_header in supported_encodings {
             if let Ok(header_value) = encoding_header.to_str() {
                 for header_value in header_value.split(",").map(|hv| hv.trim()) {
                     for algo in header_value.split(";").map(|hv| hv.trim()) {
                         match algo {
                             "zstd" if self.zstd_encoded.is_some() => {
-                                return (self.zstd_encoded.clone().unwrap(), "zstd")
+                                return (self.zstd_encoded.clone().unwrap(), Some("zstd"))
                             }
                             "gzip" if self.gzip_encoded.is_some() => {
-                                return (self.gzip_encoded.clone().unwrap(), "gzip")
+                                return (self.gzip_encoded.clone().unwrap(), Some("gzip"))
                             }
                             "br" if self.br_encoded.is_some() => {
-                                return (self.br_encoded.clone().unwrap(), "br")
+                                return (self.br_encoded.clone().unwrap(), Some("br"))
                             }
                             "deflate" if self.deflate_encoded.is_some() => {
-                                return (self.deflate_encoded.clone().unwrap(), "deflate")
+                                return (self.deflate_encoded.clone().unwrap(), Some("deflate"))
                             }
                             _ => {}
                         }
@@ -137,7 +142,7 @@ impl CacheEntry {
                 }
             }
         }
-        (self.content.clone(), "identity")
+        (self.content.clone(), None)
     }
 }
 
@@ -304,15 +309,10 @@ impl StaticFileServer {
                         .await
                 }
                 NotFoundBehavior::Redirect(redirect) => Response::builder()
-                    .status(StatusCode::MOVED_PERMANENTLY)
+                    .status(redirect.r#type.status_code())
                     .header(header::LOCATION, redirect.to)
                     .body(BoxBody::new(Full::new(Bytes::new())))
                     .unwrap(),
-                NotFoundBehavior::InternalServerError => {
-                    INTERNAL_SERVER_ERROR_RESPONSE
-                        .to_http_response(request.accept().into())
-                        .await
-                }
             },
         })
     }
@@ -434,8 +434,8 @@ impl StaticFileServer {
 
         // No valid cached entry, resolve the key to a file path
         let decoded_key = percent_decode_str(key).decode_utf8_lossy();
-        let normalized_path =
-            normalize_path(decoded_key).ok_or(NotFoundBehavior::InternalServerError)?;
+        let normalized_path = normalize_path(decoded_key)
+            .ok_or(NotFoundBehavior::Error(NOT_FOUND_RESPONSE.clone()))?;
 
         if !self.config.serve_hidden_files
             && normalized_path
@@ -449,6 +449,7 @@ impl StaticFileServer {
 
         let mut full_path = self.config.root_dir.clone();
         full_path.push(normalized_path);
+        debug!("Resolving path {:?}", full_path);
         // Check if path exists and is a file
         match tokio::fs::metadata(&full_path).await {
             Ok(metadata) => {
@@ -492,7 +493,9 @@ impl StaticFileServer {
                         // Check for case insensitive index.html
                         let entries = match tokio::fs::read_dir(&full_path).await {
                             Ok(entries) => entries,
-                            Err(_) => return Err(NotFoundBehavior::InternalServerError),
+                            Err(_) => {
+                                return Err(NotFoundBehavior::Error(NOT_FOUND_RESPONSE.clone()))
+                            }
                         };
 
                         tokio::pin!(entries);
@@ -558,6 +561,7 @@ impl StaticFileServer {
             }
             Err(_) => {
                 // Path doesn't exist, try with .html extension if configured
+                debug!("Path doesn't exist");
                 if self.config.try_html_extension {
                     let mut html_path = full_path.clone();
                     html_path.set_extension("html");
@@ -754,6 +758,7 @@ impl StaticFileServer {
                 (end_idx - start) as usize,
                 last_modified,
                 content_range,
+                &self.headers,
                 self.stream_file_range(file, start, end_idx).await.unwrap(),
             )
         } else {
@@ -765,6 +770,7 @@ impl StaticFileServer {
                 content_length as usize,
                 last_modified,
                 content_range,
+                &self.headers,
                 self.stream_file(file).await.unwrap(),
             )
         }
@@ -863,6 +869,7 @@ impl StaticFileServer {
                 range_bytes.len(),
                 cache_entry.last_modified,
                 content_range,
+                &self.headers,
                 BoxBody::new(Full::new(range_bytes)),
             )
         } else {
@@ -871,12 +878,13 @@ impl StaticFileServer {
             let body = build_ok_body(content);
             build_file_response(
                 status,
-                Some(encoding),
+                encoding,
                 Some(&cache_entry.etag),
                 get_mime_type(path),
                 content_length as usize,
                 cache_entry.last_modified,
                 content_range,
+                &self.headers,
                 body,
             )
         }
@@ -947,6 +955,7 @@ fn build_file_response(
     content_length: usize,
     last_modified: SystemTime,
     range_header: Option<String>,
+    headers: &Option<HashMap<String, String>>,
     body: BoxBody<Bytes, Infallible>,
 ) -> http::Response<BoxBody<Bytes, Infallible>> {
     let mut builder = Response::builder()
@@ -966,6 +975,11 @@ fn build_file_response(
     if let Some(range) = range_header {
         builder = builder.header(CONTENT_RANGE, range);
     }
+    if let Some(headers) = headers {
+        for (key, value) in headers {
+            builder = builder.header(key, value);
+        }
+    }
 
     builder.body(body).unwrap()
 }