itsi 0.1.14 → 0.1.18

data/crates/itsi_server/src/server/middleware_stack/middlewares/proxy.rs

@@ -1,43 +1,68 @@
 use std::{
     collections::HashMap,
     convert::Infallible,
+    error::Error,
     net::SocketAddr,
-    sync::{Arc, OnceLock},
+    sync::{
+        atomic::{AtomicUsize, Ordering},
+        Arc, LazyLock, OnceLock,
+    },
     time::Duration,
 };
 
-use crate::server::{
-    bind::{Bind, BindAddress},
-    itsi_service::RequestContext,
-    types::{HttpRequest, HttpResponse},
-};
-
 use super::{string_rewrite::StringRewrite, ErrorResponse, FromValue, MiddlewareLayer};
-
+use crate::{
+    server::{
+        binds::bind::{Bind, BindAddress},
+        http_message_types::{HttpRequest, HttpResponse, RequestExt, ResponseFormat},
+        size_limited_incoming::MaxBodySizeReached,
+    },
+    services::itsi_http_service::HttpRequestContext,
+};
 use async_trait::async_trait;
+use bytes::{Bytes, BytesMut};
 use either::Either;
 use futures::TryStreamExt;
-use http::Response;
-use http_body_util::{combinators::BoxBody, BodyExt, StreamBody};
+use http::{HeaderMap, Method, Response, StatusCode};
+use http_body_util::{combinators::BoxBody, BodyExt, Empty, StreamBody};
 use hyper::body::Frame;
 use magnus::error::Result;
-use reqwest::{dns::Resolve, Body, Client, Url};
+use rand::Rng;
+use reqwest::{
+    dns::{Name, Resolve},
+    Body, Client, Url,
+};
 use serde::Deserialize;
-use tracing::error;
 
 #[derive(Debug, Clone, Deserialize)]
 pub struct Proxy {
     pub to: StringRewrite,
     pub backends: Vec<String>,
+    pub backend_priority: BackendPriority,
     pub headers: HashMap<String, Option<ProxiedHeader>>,
     pub verify_ssl: bool,
     pub timeout: u64,
     pub tls_sni: bool,
     #[serde(skip_deserializing)]
     pub client: OnceLock<Client>,
+    #[serde(default = "bad_gateway_error_response")]
     pub error_response: ErrorResponse,
 }
 
+fn bad_gateway_error_response() -> ErrorResponse {
+    ErrorResponse::bad_gateway()
+}
+
+#[derive(Debug, Clone, Deserialize)]
+pub enum BackendPriority {
+    #[serde(rename(deserialize = "round_robin"))]
+    RoundRobin,
+    #[serde(rename(deserialize = "ordered"))]
+    Ordered,
+    #[serde(rename(deserialize = "random"))]
+    Random,
+}
+
 #[derive(Debug, Clone, Deserialize)]
 pub enum ProxiedHeader {
     #[serde(rename(deserialize = "value"))]
@@ -47,7 +72,7 @@ pub enum ProxiedHeader {
 }
 
 impl ProxiedHeader {
-    pub fn to_string(&self, req: &HttpRequest, context: &RequestContext) -> String {
+    pub fn to_string(&self, req: &HttpRequest, context: &HttpRequestContext) -> String {
         match self {
             ProxiedHeader::String(value) => value.clone(),
             ProxiedHeader::StringRewrite(rewrite) => rewrite.rewrite_request(req, context),
@@ -58,22 +83,24 @@ impl ProxiedHeader {
 #[derive(Debug, Clone)]
 pub struct Resolver {
     backends: Arc<Vec<SocketAddr>>,
+    counter: Arc<AtomicUsize>,
+    backend_priority: BackendPriority,
 }
 
-/// An iterator that owns an Arc to the backend list and iterates over it.
-pub struct ResolverIter {
+pub struct StatefulResolverIter {
     backends: Arc<Vec<SocketAddr>>,
-    index: usize,
+    start_index: usize,
+    current: usize,
 }
 
-impl Iterator for ResolverIter {
+impl Iterator for StatefulResolverIter {
     type Item = SocketAddr;
 
     fn next(&mut self) -> Option<Self::Item> {
-        if self.index < self.backends.len() {
-            let addr = self.backends[self.index];
-            self.index += 1;
-            Some(addr)
+        if self.current < self.backends.len() {
+            let index = (self.start_index + self.current) % self.backends.len();
+            self.current += 1;
+            Some(self.backends[index])
         } else {
             None
         }
@@ -81,16 +108,160 @@ impl Iterator for ResolverIter {
 }
 
 impl Resolve for Resolver {
-    fn resolve(&self, _name: reqwest::dns::Name) -> reqwest::dns::Resolving {
+    fn resolve(&self, _name: Name) -> reqwest::dns::Resolving {
         let backends = self.backends.clone();
+        let len = backends.len();
+
+        let start_index = match self.backend_priority {
+            BackendPriority::Ordered => 0,
+            BackendPriority::Random => rand::rng().random_range(0..len),
+            BackendPriority::RoundRobin => self.counter.fetch_add(1, Ordering::Relaxed) % len,
+        };
+
         let fut = async move {
-            let iter = ResolverIter { backends, index: 0 };
+            let iter = StatefulResolverIter {
+                backends,
+                start_index,
+                current: 0,
+            };
             Ok(Box::new(iter) as Box<dyn Iterator<Item = SocketAddr> + Send>)
         };
+
         Box::pin(fut)
     }
 }
 
+static BAD_GATEWAY_RESPONSE: LazyLock<ErrorResponse> = LazyLock::new(ErrorResponse::bad_gateway);
+static GATEWAY_TIMEOUT_RESPONSE: LazyLock<ErrorResponse> =
+    LazyLock::new(ErrorResponse::gateway_timeout);
+static SERVICE_UNAVAILABLE_RESPONSE: LazyLock<ErrorResponse> =
+    LazyLock::new(ErrorResponse::service_unavailable);
+static INTERNAL_SERVER_ERROR_RESPONSE: LazyLock<ErrorResponse> =
+    LazyLock::new(ErrorResponse::internal_server_error);
+
+fn is_idempotent(method: &Method) -> bool {
+    matches!(
+        *method,
+        Method::GET | Method::HEAD | Method::PUT | Method::DELETE | Method::OPTIONS
+    )
+}
+
+/// A helper that stores the immutable parts of the incoming request.
+struct RequestInfo {
+    method: Method,
+    headers: HeaderMap,
+}
+
+impl Proxy {
+    /// Build a header map of overriding headers based on the configured values.
+    /// This uses the full HttpRequest and context to compute each header value.
+    fn build_overriding_headers(
+        &self,
+        req: &HttpRequest,
+        context: &mut HttpRequestContext,
+    ) -> http::HeaderMap {
+        let mut headers = http::HeaderMap::new();
+        for (name, header_opt) in self.headers.iter() {
+            if let Some(header_value) = header_opt {
+                // Compute the header value using the full HttpRequest.
+                let value_str = header_value.to_string(req, context);
+                if let Ok(header_val) = http::HeaderValue::from_str(&value_str) {
+                    if let Ok(header_name) = name.parse::<http::header::HeaderName>() {
+                        headers.insert(header_name, header_val);
+                    }
+                }
+            }
+        }
+        headers
+    }
+
+    /// Build a reqwest::RequestBuilder by merging the original request headers
+    /// (unless overridden) with the precomputed overriding headers.
+    fn build_reqwest_request_info(
+        &self,
+        req_info: &RequestInfo,
+        url: &str,
+        host_str: &str,
+        body: Body,
+        overriding_headers: &http::HeaderMap,
+    ) -> reqwest::RequestBuilder {
+        let mut builder = self
+            .client
+            .get()
+            .unwrap()
+            .request(req_info.method.clone(), url);
+
+        // Forward headers from the original request unless they are overridden.
+        for (name, value) in req_info.headers.iter() {
+            if overriding_headers.contains_key(name) {
+                continue;
+            }
+            builder = builder.header(name, value);
+        }
+
+        // Add a Host header if not overridden.
+        if !overriding_headers.contains_key("host") && !host_str.is_empty() {
+            builder = builder.header("Host", host_str);
+        }
+
+        for (name, value) in overriding_headers.iter() {
+            builder = builder.header(name, value);
+        }
+
+        builder.body(body)
+    }
+
+    /// Sends an idempotent request using a replayable (buffered) body.
+    async fn send_request_idempotent(
+        &self,
+        req_info: &RequestInfo,
+        url: &str,
+        host_str: &str,
+        max_attempts: usize,
+        replayable_bytes: Bytes,
+        overriding_headers: &http::HeaderMap,
+    ) -> std::result::Result<reqwest::Response, reqwest::Error> {
+        let mut last_err = None;
+        for attempt in 0..max_attempts {
+            let body = Body::from(replayable_bytes.clone());
+            let builder =
+                self.build_reqwest_request_info(req_info, url, host_str, body, overriding_headers);
+            match builder.send().await {
+                Ok(response) => return Ok(response),
+                Err(e) => {
+                    // Retry for connectivity-related errors.
+                    if e.is_connect() {
+                        last_err = Some(e);
+                        if attempt + 1 < max_attempts {
+                            continue;
+                        } else {
+                            break;
+                        }
+                    } else {
+                        return Err(e);
+                    }
+                }
+            }
+        }
+        Err(last_err.expect("At least one attempt should have set last_err"))
+    }
+
+    /// Sends a non-idempotent request once using its streaming body.
+    async fn send_request_non_idempotent(
+        &self,
+        req: HttpRequest,
+        req_info: &RequestInfo,
+        url: &str,
+        host_str: &str,
+        overriding_headers: &http::HeaderMap,
+    ) -> std::result::Result<reqwest::Response, reqwest::Error> {
+        let body = Body::wrap_stream(req.into_data_stream());
+        let builder =
+            self.build_reqwest_request_info(req_info, url, host_str, body, overriding_headers);
+        builder.send().await
+    }
+}
+
 #[async_trait]
 impl MiddlewareLayer for Proxy {
     async fn initialize(&self) -> Result<()> {
@@ -116,6 +287,8 @@ impl MiddlewareLayer for Proxy {
                     .danger_accept_invalid_hostnames(!self.verify_ssl)
                     .dns_resolver(Arc::new(Resolver {
                         backends: Arc::new(backends),
+                        counter: Arc::new(AtomicUsize::new(0)),
+                        backend_priority: self.backend_priority.clone(),
                     }))
                     .tls_sni(self.tls_sni)
                     .build()
@@ -128,7 +301,7 @@ impl MiddlewareLayer for Proxy {
             )
             .map_err(|_e| {
                 magnus::Error::new(
-                    magnus::exception::exception(),
+                    magnus::exception::standard_error(),
                     "Failed to save resolver backends",
                 )
             })?;
@@ -138,55 +311,67 @@ impl MiddlewareLayer for Proxy {
     async fn before(
         &self,
         req: HttpRequest,
-        context: &mut RequestContext,
+        context: &mut HttpRequestContext,
     ) -> Result<Either<HttpRequest, HttpResponse>> {
         let url = self.to.rewrite_request(&req, context);
-        let error_response = self.error_response.to_http_response(&req).await;
+        let accept: ResponseFormat = req.accept().into();
+        let error_response = self.error_response.to_http_response(accept.clone()).await;
 
         let destination = match Url::parse(&url) {
             Ok(dest) => dest,
             Err(_) => return Ok(Either::Right(error_response)),
         };
 
+        // Clone the headers before consuming the request.
+        let req_headers = req.headers().clone();
         let host_str = destination.host_str().unwrap_or_else(|| {
-            req.headers()
+            req_headers
                 .get("Host")
                 .and_then(|h| h.to_str().ok())
                 .unwrap_or("")
         });
 
-        let mut reqwest_builder = self
-            .client
-            .get()
-            .unwrap()
-            .request(req.method().clone(), url);
-
-        // Forward incoming headers unless they're in remove_headers or overridden.
-        for (name, value) in req.headers().iter() {
-            let name_str = name.as_str();
-            if self.headers.contains_key(name_str) {
-                continue;
-            }
-            reqwest_builder = reqwest_builder.header(name, value);
-        }
+        let req_info = RequestInfo {
+            method: req.method().clone(),
+            headers: req_headers.clone(),
+        };
 
-        // Add the host header if it's not overridden and host_str is non-empty.
-        if !self.headers.contains_key("host") && !host_str.is_empty() {
-            reqwest_builder = reqwest_builder.header("Host", host_str);
-        }
+        // Precompute the overriding headers from the full request.
+        let overriding_headers = self.build_overriding_headers(&req, context);
 
-        // Add overriding headers.
-        for (name, header_value) in self.headers.iter() {
-            if let Some(header_value) = header_value {
-                reqwest_builder =
-                    reqwest_builder.header(name, header_value.to_string(&req, context));
-            }
-        }
+        // Determine max_attempts based on the number of backends.
+        let max_attempts = self.backends.len();
 
-        let reqwest_builder = reqwest_builder.body(Body::wrap_stream(req.into_data_stream()));
-        let reqwest_response = reqwest_builder.send().await;
+        let reqwest_response_result = if is_idempotent(&req_info.method) {
+            let (_parts, body) = req.into_parts();
+            let replayable_bytes = match body.into_data_stream().try_collect::<Vec<Bytes>>().await {
+                Ok(chunks) => {
+                    let mut buf = BytesMut::new();
+                    for chunk in chunks {
+                        buf.extend_from_slice(&chunk);
+                    }
+                    buf.freeze()
+                }
+                Err(e) => {
+                    tracing::error!("Error buffering request body: {}", e);
+                    return Ok(Either::Right(error_response));
+                }
+            };
+            self.send_request_idempotent(
+                &req_info,
+                &url,
+                host_str,
+                max_attempts,
+                replayable_bytes,
+                &overriding_headers,
+            )
+            .await
+        } else {
+            self.send_request_non_idempotent(req, &req_info, &url, host_str, &overriding_headers)
+                .await
+        };
 
-        let response = match reqwest_response {
+        let response = match reqwest_response_result {
             Ok(response) => {
                 let status = response.status();
                 let mut builder = Response::builder().status(status);
@@ -199,17 +384,30 @@ impl MiddlewareLayer for Proxy {
                         .map_ok(Frame::data)
                         .map_err(|_| -> Infallible { unreachable!("We handle IO errors above") }),
                 )));
-                if let Ok(response) = response {
-                    response
-                } else {
-                    error_response
-                }
+                response.unwrap_or(error_response)
             }
             Err(e) => {
-                error!("Error sending request: {}", e);
-                error_response
+                if let Some(inner) = e.source() {
+                    if inner.downcast_ref::<MaxBodySizeReached>().is_some() {
+                        let mut max_body_response = Response::new(BoxBody::new(Empty::new()));
+                        *max_body_response.status_mut() = StatusCode::PAYLOAD_TOO_LARGE;
+                        return Ok(Either::Right(max_body_response));
+                    }
+                }
+                if e.is_timeout() {
+                    GATEWAY_TIMEOUT_RESPONSE.to_http_response(accept).await
+                } else if e.is_connect() {
+                    BAD_GATEWAY_RESPONSE.to_http_response(accept).await
+                } else if e.is_status() {
+                    SERVICE_UNAVAILABLE_RESPONSE.to_http_response(accept).await
+                } else {
+                    INTERNAL_SERVER_ERROR_RESPONSE
+                        .to_http_response(accept)
+                        .await
+                }
             }
         };
+
         Ok(Either::Right(response))
     }
 }

data/crates/itsi_server/src/server/middleware_stack/middlewares/rate_limit.rs

@@ -1,10 +1,8 @@
 use super::{token_source::TokenSource, ErrorResponse, FromValue, MiddlewareLayer};
-use crate::server::{
-    itsi_service::RequestContext,
-    rate_limiter::{
-        create_rate_limit_key, get_rate_limiter, RateLimitError, RateLimiter, RateLimiterConfig,
-    },
-    types::{HttpRequest, HttpResponse, RequestExt},
+use crate::server::http_message_types::{HttpRequest, HttpResponse, RequestExt};
+use crate::services::itsi_http_service::HttpRequestContext;
+use crate::services::rate_limiter::{
+    create_rate_limit_key, get_rate_limiter, RateLimitError, RateLimiter, RateLimiterConfig,
 };
 use async_trait::async_trait;
 use either::Either;
@@ -21,9 +19,14 @@ pub struct RateLimit {
     #[serde(skip_deserializing)]
     pub rate_limiter: OnceLock<Arc<dyn RateLimiter>>,
     pub store_config: RateLimiterConfig,
+    #[serde(default = "too_many_requests_error_response")]
     pub error_response: ErrorResponse,
 }
 
+fn too_many_requests_error_response() -> ErrorResponse {
+    ErrorResponse::too_many_requests()
+}
+
 #[derive(Debug, Clone, Deserialize)]
 pub enum RateLimitKey {
     #[serde(rename(deserialize = "address"))]
@@ -46,7 +49,7 @@ impl MiddlewareLayer for RateLimit {
     async fn before(
         &self,
         req: HttpRequest,
-        context: &mut RequestContext,
+        context: &mut HttpRequestContext,
     ) -> Result<Either<HttpRequest, HttpResponse>> {
         // Get the key to rate limit on
         let key_value = match &self.key {
@@ -92,21 +95,25 @@ impl MiddlewareLayer for RateLimit {
             let limit = self.requests;
 
             match limiter.check_limit(&rate_limit_key, limit, timeout).await {
-                Ok(_) => {
-                    // Rate limit not exceeded, allow request
-                    Ok(Either::Left(req))
-                }
-                Err(RateLimitError::RateLimitExceeded { limit, count }) => {
-                    // Rate limit exceeded, return error response
-                    tracing::info!(
-                        "Rate limit exceeded for key '{}': {}/{} requests",
-                        rate_limit_key,
-                        count,
-                        limit
-                    );
-                    Ok(Either::Right(
-                        self.error_response.to_http_response(&req).await,
-                    ))
+                Ok(_) => Ok(Either::Left(req)),
+                Err(RateLimitError::RateLimitExceeded { limit, ttl, .. }) => {
+                    let mut response = self
+                        .error_response
+                        .to_http_response(req.accept().into())
+                        .await;
+                    response
+                        .headers_mut()
+                        .insert("X-RateLimit-Limit", limit.to_string().parse().unwrap());
+                    response
+                        .headers_mut()
+                        .insert("X-RateLimit-Remaining", "0".parse().unwrap());
+                    response
+                        .headers_mut()
+                        .insert("X-RateLimit-Reset", ttl.to_string().parse().unwrap());
+                    response
+                        .headers_mut()
+                        .insert("Retry-After", ttl.to_string().parse().unwrap());
+                    Ok(Either::Right(response))
                 }
                 Err(e) => {
                     // Other error, log and allow request (fail open)

data/crates/itsi_server/src/server/middleware_stack/middlewares/redirect.rs

@@ -1,6 +1,6 @@
-use crate::server::{
-    itsi_service::RequestContext,
-    types::{HttpRequest, HttpResponse},
+use crate::{
+    server::http_message_types::{HttpRequest, HttpResponse},
+    services::itsi_http_service::HttpRequestContext,
 };
 
 use super::{string_rewrite::StringRewrite, FromValue, MiddlewareLayer};
@@ -42,7 +42,7 @@ impl MiddlewareLayer for Redirect {
     async fn before(
         &self,
         req: HttpRequest,
-        context: &mut RequestContext,
+        context: &mut HttpRequestContext,
     ) -> Result<Either<HttpRequest, HttpResponse>> {
         Ok(Either::Right(self.redirect_response(&req, context)?))
     }
@@ -52,7 +52,7 @@ impl Redirect {
     pub fn redirect_response(
         &self,
         req: &HttpRequest,
-        context: &mut RequestContext,
+        context: &mut HttpRequestContext,
     ) -> Result<HttpResponse> {
         let mut response = Response::new(BoxBody::new(Empty::new()));
         *response.status_mut() = match self.redirect_type {
@@ -65,7 +65,7 @@ impl Redirect {
             "Location",
             self.to.rewrite_request(req, context).parse().map_err(|e| {
                 magnus::Error::new(
-                    magnus::exception::exception(),
+                    magnus::exception::standard_error(),
                     format!("Invalid Rewrite String: {:?}: {}", self.to, e),
                 )
             })?,

data/crates/itsi_server/src/server/middleware_stack/middlewares/request_headers.rs

@@ -1,10 +1,11 @@
 use std::collections::HashMap;
 
-use super::{FromValue, MiddlewareLayer};
-use crate::server::{
-    itsi_service::RequestContext,
-    types::{HttpRequest, HttpResponse},
+use crate::{
+    server::http_message_types::{HttpRequest, HttpResponse},
+    services::itsi_http_service::HttpRequestContext,
 };
+
+use super::{FromValue, MiddlewareLayer};
 use async_trait::async_trait;
 use either::Either;
 use http::HeaderName;
@@ -22,7 +23,7 @@ impl MiddlewareLayer for RequestHeaders {
     async fn before(
         &self,
         mut req: HttpRequest,
-        _: &mut RequestContext,
+        _: &mut HttpRequestContext,
     ) -> Result<Either<HttpRequest, HttpResponse>> {
         let headers = req.headers_mut();
         for removal in &self.removals {

data/crates/itsi_server/src/server/middleware_stack/middlewares/response_headers.rs

@@ -1,7 +1,9 @@
 use std::collections::HashMap;
 
 use super::{FromValue, MiddlewareLayer};
-use crate::server::{itsi_service::RequestContext, types::HttpResponse};
+use crate::{
+    server::http_message_types::HttpResponse, services::itsi_http_service::HttpRequestContext,
+};
 use async_trait::async_trait;
 use http::HeaderName;
 use serde::Deserialize;
@@ -14,7 +16,7 @@ pub struct ResponseHeaders {
 
 #[async_trait]
 impl MiddlewareLayer for ResponseHeaders {
-    async fn after(&self, mut resp: HttpResponse, _: &mut RequestContext) -> HttpResponse {
+    async fn after(&self, mut resp: HttpResponse, _: &mut HttpRequestContext) -> HttpResponse {
         let headers = resp.headers_mut();
         for removal in &self.removals {
             headers.remove(removal);