itsi-server 0.1.1 → 0.1.13

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

@@ -0,0 +1,216 @@
+use std::{
+    collections::HashMap,
+    convert::Infallible,
+    net::SocketAddr,
+    sync::{Arc, OnceLock},
+    time::Duration,
+};
+
+use crate::server::{
+    bind::{Bind, BindAddress},
+    itsi_service::RequestContext,
+    types::{HttpRequest, HttpResponse},
+};
+
+use super::{string_rewrite::StringRewrite, ErrorResponse, FromValue, MiddlewareLayer};
+
+use async_trait::async_trait;
+use either::Either;
+use futures::TryStreamExt;
+use http::Response;
+use http_body_util::{combinators::BoxBody, BodyExt, StreamBody};
+use hyper::body::Frame;
+use magnus::error::Result;
+use reqwest::{dns::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 headers: HashMap<String, Option<ProxiedHeader>>,
+    pub verify_ssl: bool,
+    pub timeout: u64,
+    pub tls_sni: bool,
+    #[serde(skip_deserializing)]
+    pub client: OnceLock<Client>,
+    pub error_response: ErrorResponse,
+}
+
+#[derive(Debug, Clone, Deserialize)]
+pub enum ProxiedHeader {
+    #[serde(rename(deserialize = "value"))]
+    String(String),
+    #[serde(rename(deserialize = "rewrite"))]
+    StringRewrite(StringRewrite),
+}
+
+impl ProxiedHeader {
+    pub fn to_string(&self, req: &HttpRequest, context: &RequestContext) -> String {
+        match self {
+            ProxiedHeader::String(value) => value.clone(),
+            ProxiedHeader::StringRewrite(rewrite) => rewrite.rewrite_request(req, context),
+        }
+    }
+}
+
+#[derive(Debug, Clone)]
+pub struct Resolver {
+    backends: Arc<Vec<SocketAddr>>,
+}
+
+/// An iterator that owns an Arc to the backend list and iterates over it.
+pub struct ResolverIter {
+    backends: Arc<Vec<SocketAddr>>,
+    index: usize,
+}
+
+impl Iterator for ResolverIter {
+    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)
+        } else {
+            None
+        }
+    }
+}
+
+impl Resolve for Resolver {
+    fn resolve(&self, _name: reqwest::dns::Name) -> reqwest::dns::Resolving {
+        let backends = self.backends.clone();
+        let fut = async move {
+            let iter = ResolverIter { backends, index: 0 };
+            Ok(Box::new(iter) as Box<dyn Iterator<Item = SocketAddr> + Send>)
+        };
+        Box::pin(fut)
+    }
+}
+
+#[async_trait]
+impl MiddlewareLayer for Proxy {
+    async fn initialize(&self) -> Result<()> {
+        let backends = self
+            .backends
+            .iter()
+            .filter_map(|be| {
+                let bind: Bind = be.parse().ok()?;
+                match (bind.address, bind.port) {
+                    (BindAddress::Ip(ip_addr), port) => {
+                        Some(SocketAddr::new(ip_addr, port.unwrap()))
+                    }
+                    (BindAddress::UnixSocket(_), _) => None,
+                }
+            })
+            .collect::<Vec<_>>();
+
+        self.client
+            .set(
+                Client::builder()
+                    .timeout(Duration::from_secs(self.timeout))
+                    .danger_accept_invalid_certs(!self.verify_ssl)
+                    .danger_accept_invalid_hostnames(!self.verify_ssl)
+                    .dns_resolver(Arc::new(Resolver {
+                        backends: Arc::new(backends),
+                    }))
+                    .tls_sni(self.tls_sni)
+                    .build()
+                    .map_err(|e| {
+                        magnus::Error::new(
+                            magnus::exception::runtime_error(),
+                            format!("Failed to build Reqwest client: {}", e),
+                        )
+                    })?,
+            )
+            .map_err(|_e| {
+                magnus::Error::new(
+                    magnus::exception::exception(),
+                    "Failed to save resolver backends",
+                )
+            })?;
+        Ok(())
+    }
+
+    async fn before(
+        &self,
+        req: HttpRequest,
+        context: &mut RequestContext,
+    ) -> Result<Either<HttpRequest, HttpResponse>> {
+        let url = self.to.rewrite_request(&req, context);
+        let error_response = self.error_response.to_http_response(&req).await;
+
+        let destination = match Url::parse(&url) {
+            Ok(dest) => dest,
+            Err(_) => return Ok(Either::Right(error_response)),
+        };
+
+        let host_str = destination.host_str().unwrap_or_else(|| {
+            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);
+        }
+
+        // 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);
+        }
+
+        // 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));
+            }
+        }
+
+        let reqwest_builder = reqwest_builder.body(Body::wrap_stream(req.into_data_stream()));
+        let reqwest_response = reqwest_builder.send().await;
+
+        let response = match reqwest_response {
+            Ok(response) => {
+                let status = response.status();
+                let mut builder = Response::builder().status(status);
+                for (hn, hv) in response.headers() {
+                    builder = builder.header(hn, hv);
+                }
+                let response = builder.body(BoxBody::new(StreamBody::new(
+                    response
+                        .bytes_stream()
+                        .map_ok(Frame::data)
+                        .map_err(|_| -> Infallible { unreachable!("We handle IO errors above") }),
+                )));
+                if let Ok(response) = response {
+                    response
+                } else {
+                    error_response
+                }
+            }
+            Err(e) => {
+                error!("Error sending request: {}", e);
+                error_response
+            }
+        };
+        Ok(Either::Right(response))
+    }
+}
+impl FromValue for Proxy {}

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

@@ -0,0 +1,124 @@
+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 async_trait::async_trait;
+use either::Either;
+use magnus::error::Result;
+use serde::Deserialize;
+use std::sync::{Arc, OnceLock};
+use std::time::Duration;
+
+#[derive(Debug, Clone, Deserialize)]
+pub struct RateLimit {
+    pub requests: u64,
+    pub seconds: u64,
+    pub key: RateLimitKey,
+    #[serde(skip_deserializing)]
+    pub rate_limiter: OnceLock<Arc<dyn RateLimiter>>,
+    pub store_config: RateLimiterConfig,
+    pub error_response: ErrorResponse,
+}
+
+#[derive(Debug, Clone, Deserialize)]
+pub enum RateLimitKey {
+    #[serde(rename(deserialize = "address"))]
+    SocketAddress,
+    #[serde(rename(deserialize = "parameter"))]
+    Parameter(TokenSource),
+}
+
+#[async_trait]
+impl MiddlewareLayer for RateLimit {
+    async fn initialize(&self) -> Result<()> {
+        // Instantiate our rate limiter based on the rate limit config here.
+        // This will automatically fall back to in-memory if Redis fails
+        if let Ok(limiter) = get_rate_limiter(&self.store_config).await {
+            let _ = self.rate_limiter.set(limiter);
+        }
+        Ok(())
+    }
+
+    async fn before(
+        &self,
+        req: HttpRequest,
+        context: &mut RequestContext,
+    ) -> Result<Either<HttpRequest, HttpResponse>> {
+        // Get the key to rate limit on
+        let key_value = match &self.key {
+            RateLimitKey::SocketAddress => {
+                // Use the socket address from the context
+                &context.addr
+            }
+            RateLimitKey::Parameter(token_source) => {
+                match token_source {
+                    TokenSource::Header { name, prefix } => {
+                        if let Some(header) = req.header(name) {
+                            if let Some(prefix) = prefix {
+                                header.strip_prefix(prefix).unwrap_or("").trim_ascii()
+                            } else {
+                                header.trim_ascii()
+                            }
+                        } else {
+                            // If no token is found, skip rate limiting
+                            tracing::warn!("No token found in header for rate limiting");
+                            return Ok(Either::Left(req));
+                        }
+                    }
+                    TokenSource::Query(query_name) => {
+                        if let Some(value) = req.query_param(query_name) {
+                            value
+                        } else {
+                            // If no token is found, skip rate limiting
+                            tracing::warn!("No token found in query for rate limiting");
+                            return Ok(Either::Left(req));
+                        }
+                    }
+                }
+            }
+        };
+
+        // Create a rate limit key
+        let rate_limit_key = create_rate_limit_key(key_value, req.uri().path());
+
+        // Get the rate limiter
+        if let Some(limiter) = self.rate_limiter.get() {
+            // Check if rate limit is exceeded
+            let timeout = Duration::from_secs(self.seconds);
+            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,
+                    ))
+                }
+                Err(e) => {
+                    // Other error, log and allow request (fail open)
+                    tracing::error!("Rate limiter error: {:?}", e);
+                    Ok(Either::Left(req))
+                }
+            }
+        } else {
+            // If rate limiter is not initialized, allow request
+            tracing::warn!("Rate limiter not initialized");
+            Ok(Either::Left(req))
+        }
+    }
+}
+impl FromValue for RateLimit {}

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

@@ -0,0 +1,76 @@
+use crate::server::{
+    itsi_service::RequestContext,
+    types::{HttpRequest, HttpResponse},
+};
+
+use super::{string_rewrite::StringRewrite, FromValue, MiddlewareLayer};
+
+use async_trait::async_trait;
+use either::Either;
+use http::{Response, StatusCode};
+use http_body_util::{combinators::BoxBody, Empty};
+use magnus::error::Result;
+use serde::Deserialize;
+
+/// A simple API key filter.
+/// The API key can be given inside the header or a query string
+/// Keys are validated against a list of allowed key values (Changing these requires a restart)
+///
+#[derive(Debug, Clone, Deserialize)]
+pub struct Redirect {
+    pub to: StringRewrite,
+    #[serde(default)]
+    #[serde(rename(deserialize = "type"))]
+    pub redirect_type: RedirectType,
+}
+
+#[derive(Debug, Clone, Deserialize, Default)]
+pub enum RedirectType {
+    #[serde(rename(deserialize = "permanent"))]
+    #[default]
+    Permanent,
+    #[serde(rename(deserialize = "temporary"))]
+    Temporary,
+    #[serde(rename(deserialize = "found"))]
+    Found,
+    #[serde(rename(deserialize = "moved_permanently"))]
+    MovedPermanently,
+}
+
+#[async_trait]
+impl MiddlewareLayer for Redirect {
+    async fn before(
+        &self,
+        req: HttpRequest,
+        context: &mut RequestContext,
+    ) -> Result<Either<HttpRequest, HttpResponse>> {
+        Ok(Either::Right(self.redirect_response(&req, context)?))
+    }
+}
+
+impl Redirect {
+    pub fn redirect_response(
+        &self,
+        req: &HttpRequest,
+        context: &mut RequestContext,
+    ) -> Result<HttpResponse> {
+        let mut response = Response::new(BoxBody::new(Empty::new()));
+        *response.status_mut() = match self.redirect_type {
+            RedirectType::Permanent => StatusCode::PERMANENT_REDIRECT,
+            RedirectType::Temporary => StatusCode::TEMPORARY_REDIRECT,
+            RedirectType::MovedPermanently => StatusCode::MOVED_PERMANENTLY,
+            RedirectType::Found => StatusCode::FOUND,
+        };
+        response.headers_mut().append(
+            "Location",
+            self.to.rewrite_request(req, context).parse().map_err(|e| {
+                magnus::Error::new(
+                    magnus::exception::exception(),
+                    format!("Invalid Rewrite String: {:?}: {}", self.to, e),
+                )
+            })?,
+        );
+        Ok(response)
+    }
+}
+impl FromValue for Redirect {}

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

@@ -0,0 +1,43 @@
+use std::collections::HashMap;
+
+use super::{FromValue, MiddlewareLayer};
+use crate::server::{
+    itsi_service::RequestContext,
+    types::{HttpRequest, HttpResponse},
+};
+use async_trait::async_trait;
+use either::Either;
+use http::HeaderName;
+use magnus::error::Result;
+use serde::Deserialize;
+
+#[derive(Debug, Clone, Deserialize)]
+pub struct RequestHeaders {
+    pub additions: HashMap<String, Vec<String>>,
+    pub removals: Vec<String>,
+}
+
+#[async_trait]
+impl MiddlewareLayer for RequestHeaders {
+    async fn before(
+        &self,
+        mut req: HttpRequest,
+        _: &mut RequestContext,
+    ) -> Result<Either<HttpRequest, HttpResponse>> {
+        let headers = req.headers_mut();
+        for removal in &self.removals {
+            headers.remove(removal);
+        }
+        for (header_name, header_values) in &self.additions {
+            for header_value in header_values {
+                if let Ok(parsed_header_name) = header_name.parse::<HeaderName>() {
+                    if let Ok(parsed_header_value) = header_value.parse() {
+                        headers.append(parsed_header_name, parsed_header_value);
+                    }
+                }
+            }
+        }
+        Ok(Either::Left(req))
+    }
+}
+impl FromValue for RequestHeaders {}

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

@@ -0,0 +1,34 @@
+use std::collections::HashMap;
+
+use super::{FromValue, MiddlewareLayer};
+use crate::server::{itsi_service::RequestContext, types::HttpResponse};
+use async_trait::async_trait;
+use http::HeaderName;
+use serde::Deserialize;
+
+#[derive(Debug, Clone, Deserialize)]
+pub struct ResponseHeaders {
+    pub additions: HashMap<String, Vec<String>>,
+    pub removals: Vec<String>,
+}
+
+#[async_trait]
+impl MiddlewareLayer for ResponseHeaders {
+    async fn after(&self, mut resp: HttpResponse, _: &mut RequestContext) -> HttpResponse {
+        let headers = resp.headers_mut();
+        for removal in &self.removals {
+            headers.remove(removal);
+        }
+        for (header_name, header_values) in &self.additions {
+            for header_value in header_values {
+                if let Ok(parsed_header_name) = header_name.parse::<HeaderName>() {
+                    if let Ok(parsed_header_value) = header_value.parse() {
+                        headers.append(parsed_header_name, parsed_header_value);
+                    }
+                }
+            }
+        }
+        resp
+    }
+}
+impl FromValue for ResponseHeaders {}

data/ext/itsi_server/src/server/middleware_stack/middlewares/ruby_app.rs

@@ -0,0 +1,93 @@
+use super::MiddlewareLayer;
+use crate::ruby_types::itsi_grpc_request::ItsiGrpcRequest;
+use crate::server::static_file_server::ROOT_STATIC_FILE_SERVER;
+use crate::{
+    ruby_types::itsi_http_request::ItsiHttpRequest,
+    server::{
+        itsi_service::RequestContext,
+        types::{HttpRequest, HttpResponse},
+    },
+};
+use async_trait::async_trait;
+use derive_more::Debug;
+use either::Either;
+use itsi_rb_helpers::{HeapVal, HeapValue};
+use magnus::{block::Proc, error::Result, value::ReprValue, Symbol};
+use std::str::FromStr;
+use std::sync::Arc;
+
+#[derive(Debug)]
+pub struct RubyApp {
+    app: Arc<HeapValue<Proc>>,
+    request_type: RequestType,
+    sendfile: bool,
+}
+
+#[derive(Debug)]
+pub enum RequestType {
+    Http,
+    Grpc,
+}
+
+impl FromStr for RequestType {
+    type Err = &'static str;
+
+    fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
+        match s {
+            "http" => Ok(RequestType::Http),
+            "grpc" => Ok(RequestType::Grpc),
+            _ => Err("Invalid request type"),
+        }
+    }
+}
+
+impl RubyApp {
+    pub fn from_value(params: HeapVal) -> magnus::error::Result<Self> {
+        let app = params.funcall::<_, _, Proc>(Symbol::new("[]"), ("app_proc",))?;
+        let sendfile = params
+            .funcall::<_, _, bool>(Symbol::new("[]"), ("sendfile",))
+            .unwrap_or(true);
+        let request_type: RequestType = params
+            .funcall::<_, _, String>(Symbol::new("[]"), ("request_type",))
+            .unwrap_or("http".to_string())
+            .parse()
+            .unwrap_or(RequestType::Http);
+
+        Ok(RubyApp {
+            app: Arc::new(app.into()),
+            sendfile,
+            request_type,
+        })
+    }
+}
+
+#[async_trait]
+impl MiddlewareLayer for RubyApp {
+    async fn before(
+        &self,
+        req: HttpRequest,
+        context: &mut RequestContext,
+    ) -> Result<Either<HttpRequest, HttpResponse>> {
+        match self.request_type {
+            RequestType::Http => ItsiHttpRequest::process_request(self.app.clone(), req, context)
+                .await
+                .map_err(|e| e.into())
+                .map(Either::Right),
+            RequestType::Grpc => ItsiGrpcRequest::process_request(self.app.clone(), req, context)
+                .await
+                .map_err(|e| e.into())
+                .map(Either::Right),
+        }
+    }
+
+    async fn after(&self, resp: HttpResponse, _context: &mut RequestContext) -> HttpResponse {
+        if self.sendfile {
+            if let Some(sendfile_header) = resp.headers().get("X-Sendfile") {
+                return ROOT_STATIC_FILE_SERVER
+                    .serve_single(sendfile_header.to_str().unwrap())
+                    .await;
+            }
+        }
+        resp
+    }
+}