itsi-server 0.1.1 → 0.1.20

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

@@ -0,0 +1,338 @@
+use super::{error_response::ErrorResponse, token_source::TokenSource, FromValue, MiddlewareLayer};
+use crate::{
+    server::http_message_types::{HttpRequest, HttpResponse, RequestExt},
+    services::itsi_http_service::HttpRequestContext,
+};
+
+use async_trait::async_trait;
+use base64::{engine::general_purpose, Engine};
+use derive_more::Debug;
+use either::Either;
+use itsi_error::ItsiError;
+use jsonwebtoken::{
+    decode, decode_header, Algorithm as JwtAlg, DecodingKey, TokenData, Validation,
+};
+use magnus::error::Result;
+use serde::Deserialize;
+use std::{
+    collections::{HashMap, HashSet},
+    sync::OnceLock,
+};
+use tracing::{debug, error};
+
+#[derive(Debug, Clone, Deserialize)]
+pub struct AuthJwt {
+    pub token_source: TokenSource,
+    // The verifiers map still holds base64-encoded key strings keyed by algorithm.
+    pub verifiers: HashMap<JwtAlgorithm, Vec<String>>,
+    // We now store jsonwebtoken’s DecodingKey in our OnceLock.
+    #[serde(skip_deserializing)]
+    #[debug(skip)]
+    pub keys: OnceLock<HashMap<JwtAlgorithm, Vec<DecodingKey>>>,
+    pub audiences: Option<HashSet<String>>,
+    pub subjects: Option<HashSet<String>>,
+    pub issuers: Option<HashSet<String>>,
+    pub leeway: Option<u64>,
+    #[serde(default = "unauthorized_error_response")]
+    pub error_response: ErrorResponse,
+}
+
+fn unauthorized_error_response() -> ErrorResponse {
+    ErrorResponse::unauthorized()
+}
+
+#[derive(Debug, Clone, Deserialize, PartialEq, Eq, Hash)]
+pub enum JwtAlgorithm {
+    #[serde(rename(deserialize = "hs256"))]
+    Hs256,
+    #[serde(rename(deserialize = "hs384"))]
+    Hs384,
+    #[serde(rename(deserialize = "hs512"))]
+    Hs512,
+    #[serde(rename(deserialize = "rs256"))]
+    Rs256,
+    #[serde(rename(deserialize = "rs384"))]
+    Rs384,
+    #[serde(rename(deserialize = "rs512"))]
+    Rs512,
+    #[serde(rename(deserialize = "es256"))]
+    Es256,
+    #[serde(rename(deserialize = "es384"))]
+    Es384,
+    #[serde(rename(deserialize = "ps256"))]
+    Ps256,
+    #[serde(rename(deserialize = "ps384"))]
+    Ps384,
+    #[serde(rename(deserialize = "ps512"))]
+    Ps512,
+}
+
+// Allow conversion from jsonwebtoken’s Algorithm to our JwtAlgorithm.
+impl From<JwtAlg> for JwtAlgorithm {
+    fn from(alg: JwtAlg) -> Self {
+        match alg {
+            JwtAlg::HS256 => JwtAlgorithm::Hs256,
+            JwtAlg::HS384 => JwtAlgorithm::Hs384,
+            JwtAlg::HS512 => JwtAlgorithm::Hs512,
+            JwtAlg::RS256 => JwtAlgorithm::Rs256,
+            JwtAlg::RS384 => JwtAlgorithm::Rs384,
+            JwtAlg::RS512 => JwtAlgorithm::Rs512,
+            JwtAlg::ES256 => JwtAlgorithm::Es256,
+            JwtAlg::ES384 => JwtAlgorithm::Es384,
+            JwtAlg::PS256 => JwtAlgorithm::Ps256,
+            JwtAlg::PS384 => JwtAlgorithm::Ps384,
+            JwtAlg::PS512 => JwtAlgorithm::Ps512,
+            _ => panic!("Unsupported algorithm"),
+        }
+    }
+}
+
+impl JwtAlgorithm {
+    /// Given a base64-encoded key string, decode and construct a jsonwebtoken::DecodingKey.
+    pub fn key_from(&self, base64: &str) -> itsi_error::Result<DecodingKey> {
+        match self {
+            // For HMAC algorithms, use the secret directly.
+            JwtAlgorithm::Hs256 | JwtAlgorithm::Hs384 | JwtAlgorithm::Hs512 => {
+                Ok(DecodingKey::from_secret(
+                    &general_purpose::STANDARD
+                        .decode(base64)
+                        .map_err(ItsiError::new)?,
+                ))
+            }
+            // For RSA (and PS) algorithms, expect a PEM-formatted key.
+            JwtAlgorithm::Rs256
+            | JwtAlgorithm::Rs384
+            | JwtAlgorithm::Rs512
+            | JwtAlgorithm::Ps256
+            | JwtAlgorithm::Ps384
+            | JwtAlgorithm::Ps512 => DecodingKey::from_rsa_pem(base64.trim_ascii().as_bytes())
+                .map_err(|e| ItsiError::new(e.to_string())),
+            // For ECDSA algorithms, expect a PEM-formatted key.
+            JwtAlgorithm::Es256 | JwtAlgorithm::Es384 => {
+                DecodingKey::from_ec_pem(base64.trim_ascii().as_bytes())
+                    .map_err(|e| ItsiError::new(e.to_string()))
+            }
+        }
+    }
+}
+
+#[derive(Debug, Deserialize)]
+#[serde(untagged)]
+enum Audience {
+    Single(String),
+    Multiple(Vec<String>),
+}
+
+#[derive(Debug, Deserialize)]
+struct Claims {
+    // Here we assume the token includes an expiration.
+    #[allow(dead_code)]
+    exp: usize,
+    // The audience claim may be a single string or an array.
+    aud: Option<Audience>,
+    sub: Option<String>,
+    iss: Option<String>,
+}
+
+#[async_trait]
+impl MiddlewareLayer for AuthJwt {
+    async fn initialize(&self) -> Result<()> {
+        debug!(
+            target: "middleware::auth_jwt",
+            "Instantiating auth_jwt with {} verifiers", self.verifiers.len()
+        );
+
+        let keys: HashMap<JwtAlgorithm, Vec<DecodingKey>> = self
+            .verifiers
+            .iter()
+            .map(|(algorithm, key_strings)| {
+                let algo = algorithm.clone();
+                let keys: itsi_error::Result<Vec<DecodingKey>> = key_strings
+                    .iter()
+                    .map(|key_string| algorithm.key_from(key_string))
+                    .inspect(|key_result| {
+                        if key_result.is_err() {
+                            debug!(
+                                target: "middleware::auth_jwt",
+                                "Failed to load key for algorithm {:?}", algorithm
+                            )
+                        } else {
+                            debug!(
+                                target: "middleware::auth_jwt",
+                                "Loaded key for algorithm {:?}", algorithm
+                            )
+                        }
+                    })
+                    .collect();
+                keys.map(|keys| (algo, keys))
+            })
+            .collect::<itsi_error::Result<HashMap<JwtAlgorithm, Vec<DecodingKey>>>>()?;
+
+        self.keys
+            .set(keys)
+            .map_err(|_| ItsiError::new("Failed to set keys"))?;
+        Ok(())
+    }
+
+    async fn before(
+        &self,
+        req: HttpRequest,
+        _: &mut HttpRequestContext,
+    ) -> Result<Either<HttpRequest, HttpResponse>> {
+        // Retrieve the JWT token from either a header or a query parameter.
+        let token_str = match &self.token_source {
+            TokenSource::Header { name, prefix } => {
+                debug!(
+                    target: "middleware::auth_jwt",
+                    "Extracting JWT from header: {}, prefix: {:?}",
+                    name, prefix
+                );
+                if let Some(header) = req.header(name) {
+                    if let Some(prefix) = prefix {
+                        Some(header.strip_prefix(prefix).unwrap_or("").trim_ascii())
+                    } else {
+                        Some(header.trim_ascii())
+                    }
+                } else {
+                    None
+                }
+            }
+            TokenSource::Query(query_name) => {
+                debug!(
+                  target: "middleware::auth_jwt",
+                    "Extracting JWT from query parameter: {}",
+                    query_name
+                );
+                req.query_param(query_name)
+            }
+        };
+
+        if token_str.is_none() {
+            debug!(
+              target: "middleware::auth_jwt",
+                "No JWT found in headers or query parameters"
+            );
+            return Ok(Either::Right(
+                self.error_response
+                    .to_http_response(req.accept().into())
+                    .await,
+            ));
+        }
+        let token_str = token_str.unwrap();
+        let header =
+            decode_header(token_str).map_err(|_| ItsiError::new("Invalid token header"))?;
+
+        let alg: JwtAlgorithm = header.alg.into();
+
+        debug!(
+          target: "middleware::auth_jwt",
+            "Matched algorithm {:?}", alg
+        );
+        if !self.verifiers.contains_key(&alg) {
+            return Ok(Either::Right(
+                self.error_response
+                    .to_http_response(req.accept().into())
+                    .await,
+            ));
+        }
+        let keys = self.keys.get().unwrap().get(&alg).unwrap();
+
+        // Build validation based on the algorithm and optional leeway.
+        let mut validation = Validation::new(match alg {
+            JwtAlgorithm::Hs256 => JwtAlg::HS256,
+            JwtAlgorithm::Hs384 => JwtAlg::HS384,
+            JwtAlgorithm::Hs512 => JwtAlg::HS512,
+            JwtAlgorithm::Rs256 => JwtAlg::RS256,
+            JwtAlgorithm::Rs384 => JwtAlg::RS384,
+            JwtAlgorithm::Rs512 => JwtAlg::RS512,
+            JwtAlgorithm::Es256 => JwtAlg::ES256,
+            JwtAlgorithm::Es384 => JwtAlg::ES384,
+            JwtAlgorithm::Ps256 => JwtAlg::PS256,
+            JwtAlgorithm::Ps384 => JwtAlg::PS384,
+            JwtAlgorithm::Ps512 => JwtAlg::PS512,
+        });
+
+        if let Some(leeway) = self.leeway {
+            validation.leeway = leeway;
+        }
+
+        let token_data: Option<TokenData<Claims>> =
+            keys.iter()
+                .find_map(|key| match decode::<Claims>(token_str, key, &validation) {
+                    Ok(data) => Some(data),
+                    Err(e) => {
+                        error!("Token validation failed: {:?}", e);
+                        None
+                    }
+                });
+
+        let token_data = if let Some(data) = token_data {
+            data
+        } else {
+            return Ok(Either::Right(
+                self.error_response
+                    .to_http_response(req.accept().into())
+                    .await,
+            ));
+        };
+
+        let claims = token_data.claims;
+
+        if let Some(expected_audiences) = &self.audiences {
+            if let Some(aud) = &claims.aud {
+                let token_auds: HashSet<String> = match aud {
+                    Audience::Single(s) => [s.clone()].into_iter().collect(),
+                    Audience::Multiple(v) => v.iter().cloned().collect(),
+                };
+                if expected_audiences.is_disjoint(&token_auds) {
+                    debug!(
+                        "AUD check failed, token_auds: {:?}, expected_audiences: {:?}",
+                        token_auds, expected_audiences
+                    );
+                    return Ok(Either::Right(
+                        self.error_response
+                            .to_http_response(req.accept().into())
+                            .await,
+                    ));
+                }
+            }
+        }
+
+        if let Some(expected_subjects) = &self.subjects {
+            if let Some(sub) = &claims.sub {
+                if !expected_subjects.contains(sub) {
+                    debug!(
+                        "SUB check failed, token_sub: {:?}, expected_subjects: {:?}",
+                        sub, expected_subjects
+                    );
+                    return Ok(Either::Right(
+                        self.error_response
+                            .to_http_response(req.accept().into())
+                            .await,
+                    ));
+                }
+            }
+        }
+
+        // Verify expected issuer.
+        if let Some(expected_issuers) = &self.issuers {
+            if let Some(iss) = &claims.iss {
+                if !expected_issuers.contains(iss) {
+                    debug!(
+                        "ISS check failed, token_iss: {:?}, expected_issuers: {:?}",
+                        iss, expected_issuers
+                    );
+                    return Ok(Either::Right(
+                        self.error_response
+                            .to_http_response(req.accept().into())
+                            .await,
+                    ));
+                }
+            }
+        }
+
+        Ok(Either::Left(req))
+    }
+}
+
+impl FromValue for AuthJwt {}

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

@@ -0,0 +1,142 @@
+use crate::{
+    server::http_message_types::HttpResponse, services::itsi_http_service::HttpRequestContext,
+};
+
+use super::{FromValue, MiddlewareLayer};
+use async_trait::async_trait;
+use http::{HeaderName, HeaderValue};
+use magnus::error::Result;
+use serde::Deserialize;
+use std::{collections::HashMap, sync::OnceLock};
+
+#[derive(Debug, Clone, Deserialize)]
+pub struct CacheControl {
+    #[serde(default)]
+    pub max_age: Option<u64>,
+    #[serde(default)]
+    pub s_max_age: Option<u64>,
+    #[serde(default)]
+    pub stale_while_revalidate: Option<u64>,
+    #[serde(default)]
+    pub stale_if_error: Option<u64>,
+    #[serde(default)]
+    pub public: bool,
+    #[serde(default)]
+    pub private: bool,
+    #[serde(default)]
+    pub no_cache: bool,
+    #[serde(default)]
+    pub no_store: bool,
+    #[serde(default)]
+    pub must_revalidate: bool,
+    #[serde(default)]
+    pub proxy_revalidate: bool,
+    #[serde(default)]
+    pub immutable: bool,
+    #[serde(default)]
+    pub vary: Vec<String>,
+    #[serde(default)]
+    pub additional_headers: HashMap<String, String>,
+    #[serde(skip_deserializing)]
+    pub cache_control_str: OnceLock<String>,
+}
+
+#[async_trait]
+impl MiddlewareLayer for CacheControl {
+    async fn initialize(&self) -> Result<()> {
+        let mut directives = Vec::new();
+
+        if self.public && !self.private {
+            directives.push("public".to_owned());
+        } else if self.private && !self.public {
+            directives.push("private".to_owned());
+        }
+        if self.no_cache {
+            directives.push("no-cache".to_owned());
+        }
+        if self.no_store {
+            directives.push("no-store".to_owned());
+        }
+        if self.must_revalidate {
+            directives.push("must-revalidate".to_owned());
+        }
+        if self.proxy_revalidate {
+            directives.push("proxy-revalidate".to_owned());
+        }
+        if self.immutable {
+            directives.push("immutable".to_owned());
+        }
+
+        // Add age parameters
+        if let Some(max_age) = self.max_age {
+            directives.push(format!("max-age={}", max_age));
+        }
+
+        if let Some(s_max_age) = self.s_max_age {
+            directives.push(format!("s-maxage={}", s_max_age));
+        }
+
+        if let Some(stale_while_revalidate) = self.stale_while_revalidate {
+            directives.push(format!("stale-while-revalidate={}", stale_while_revalidate));
+        }
+
+        if let Some(stale_if_error) = self.stale_if_error {
+            directives.push(format!("stale-if-error={}", stale_if_error));
+        }
+
+        // Set the Cache-Control header if we have directives
+        if !directives.is_empty() {
+            let cache_control_value = directives.join(", ");
+            self.cache_control_str.set(cache_control_value).unwrap();
+        }
+
+        Ok(())
+    }
+
+    async fn after(&self, mut resp: HttpResponse, _: &mut HttpRequestContext) -> HttpResponse {
+        // Skip for statuses where caching doesn't make sense
+        let status = resp.status().as_u16();
+        if matches!(status, 401 | 403 | 500..=599) {
+            return resp;
+        }
+
+        // Set the Cache-Control header if we have directives
+        if let Some(cache_control_value) = self.cache_control_str.get() {
+            if let Ok(value) = HeaderValue::from_str(cache_control_value) {
+                resp.headers_mut().insert("Cache-Control", value);
+            }
+        }
+
+        // Set Expires header based on max-age if present
+        if let Some(max_age) = self.max_age {
+            // Set the Expires header based on max-age
+            // Use a helper to format the HTTP date correctly
+            let expires = chrono::Utc::now() + chrono::Duration::seconds(max_age as i64);
+            let expires_str = expires.format("%a, %d %b %Y %H:%M:%S GMT").to_string();
+            if let Ok(value) = HeaderValue::from_str(&expires_str) {
+                resp.headers_mut().insert("Expires", value);
+            }
+        }
+
+        // Set Vary header
+        if !self.vary.is_empty() {
+            let vary_value = self.vary.join(", ");
+            if let Ok(value) = HeaderValue::from_str(&vary_value) {
+                resp.headers_mut().insert("Vary", value);
+            }
+        }
+
+        // Set additional custom headers
+        for (name, value) in &self.additional_headers {
+            if let Ok(header_value) = HeaderValue::from_str(value) {
+                if let Ok(header_name) = name.parse::<HeaderName>() {
+                    resp.headers_mut().insert(header_name, header_value);
+                }
+            }
+        }
+
+        resp
+    }
+}
+
+impl FromValue for CacheControl {}