itsi-scheduler 0.1.5 → 0.2.2

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

@@ -0,0 +1,151 @@
+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};
+use tracing::debug;
+
+#[derive(Debug, 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(", ");
+            debug!(target: "middleware::cache_control", "Built cache-control directive {}", cache_control_value);
+            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) {
+            debug!(target: "middleware::cache_control", "Skipping cache-control for status {}", status);
+            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) {
+                debug!(target: "middleware::cache_control", "Setting cache-control header to {}", cache_control_value);
+                resp.headers_mut().insert("Cache-Control", value);
+            } else {
+                debug!(target: "middleware::cache_control", "Failed to parse cache-control value {}", cache_control_value);
+            }
+        } else {
+            debug!(target: "middleware::cache_control", "No cache-control value provided");
+        }
+
+        // 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
+            debug!(target: "middleware::cache_control", "Setting expires header to {}", max_age);
+            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);
+            }
+        }
+
+        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>() {
+                    debug!(target: "middleware::cache_control", "Setting custom header {} to {:?}", header_name, header_value);
+                    resp.headers_mut().insert(header_name, header_value);
+                }
+            }
+        }
+
+        resp
+    }
+}
+
+impl FromValue for CacheControl {}

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

@@ -0,0 +1,316 @@
+use crate::{
+    server::http_message_types::HttpResponse, services::itsi_http_service::HttpRequestContext,
+};
+
+use super::{
+    header_interpretation::{find_first_supported, header_contains},
+    FromValue, MiddlewareLayer,
+};
+
+use async_compression::{
+    tokio::bufread::{BrotliEncoder, DeflateEncoder, GzipEncoder, ZstdEncoder},
+    Level,
+};
+use async_trait::async_trait;
+use bytes::{Bytes, BytesMut};
+use futures::TryStreamExt;
+use http::{
+    header::{GetAll, CONTENT_ENCODING, CONTENT_LENGTH, CONTENT_TYPE},
+    HeaderValue, Response,
+};
+use http_body_util::{combinators::BoxBody, BodyExt, Full, StreamBody};
+use hyper::body::{Body, Frame};
+use serde::{Deserialize, Serialize};
+use std::convert::Infallible;
+use tokio::io::{AsyncRead, AsyncReadExt, BufReader};
+use tokio_stream::StreamExt;
+use tokio_util::io::{ReaderStream, StreamReader};
+use tracing::debug;
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct Compression {
+    min_size: usize,
+    algorithms: Vec<CompressionAlgorithm>,
+    compress_streams: bool,
+    mime_types: Vec<MimeType>,
+    level: CompressionLevel,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+enum CompressionLevel {
+    #[serde(rename(deserialize = "fastest"))]
+    Fastest,
+    #[serde(rename(deserialize = "best"))]
+    Best,
+    #[serde(rename(deserialize = "balanced"))]
+    Balanced,
+    #[serde(rename(deserialize = "precise"))]
+    Precise(i32),
+}
+
+impl CompressionLevel {
+    fn to_async_compression_level(&self) -> Level {
+        match self {
+            CompressionLevel::Fastest => Level::Fastest,
+            CompressionLevel::Best => Level::Best,
+            CompressionLevel::Balanced => Level::Default,
+            CompressionLevel::Precise(level) => Level::Precise(*level),
+        }
+    }
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, PartialOrd, Eq, Ord)]
+pub enum CompressionAlgorithm {
+    #[serde(rename(deserialize = "gzip"))]
+    Gzip,
+    #[serde(rename(deserialize = "br"))]
+    Brotli,
+    #[serde(rename(deserialize = "deflate"))]
+    Deflate,
+    #[serde(rename(deserialize = "zstd"))]
+    Zstd,
+    #[serde(rename(deserialize = "identity"))]
+    Identity,
+}
+
+impl CompressionAlgorithm {
+    pub fn as_str(&self) -> &'static str {
+        match self {
+            CompressionAlgorithm::Gzip => "gzip",
+            CompressionAlgorithm::Brotli => "br",
+            CompressionAlgorithm::Deflate => "deflate",
+            CompressionAlgorithm::Zstd => "zstd",
+            CompressionAlgorithm::Identity => "identity",
+        }
+    }
+
+    pub fn header_value(&self) -> HeaderValue {
+        HeaderValue::from_str(self.as_str()).unwrap()
+    }
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+enum MimeType {
+    #[serde(rename(deserialize = "text"))]
+    Text,
+    #[serde(rename(deserialize = "image"))]
+    Image,
+    #[serde(rename(deserialize = "application"))]
+    Application,
+    #[serde(rename(deserialize = "audio"))]
+    Audio,
+    #[serde(rename(deserialize = "video"))]
+    Video,
+    #[serde(rename(deserialize = "font"))]
+    Font,
+    #[serde(rename(deserialize = "other"))]
+    Other(String),
+    #[serde(rename(deserialize = "all"))]
+    All,
+}
+
+impl MimeType {
+    pub fn matches(&self, content_encodings: &GetAll<HeaderValue>) -> bool {
+        match self {
+            MimeType::Text => header_contains(content_encodings, "text/*"),
+            MimeType::Image => header_contains(content_encodings, "image/*"),
+            MimeType::Application => header_contains(content_encodings, "application/*"),
+            MimeType::Audio => header_contains(content_encodings, "audio/*"),
+            MimeType::Video => header_contains(content_encodings, "video/*"),
+            MimeType::Font => header_contains(content_encodings, "font/*"),
+            MimeType::Other(v) => header_contains(content_encodings, v),
+            MimeType::All => header_contains(content_encodings, "*"),
+        }
+    }
+}
+
+fn stream_encode<R>(encoder: R) -> BoxBody<Bytes, Infallible>
+where
+    R: AsyncRead + Unpin + Sync + Send + 'static,
+{
+    let encoded_stream = ReaderStream::new(encoder).map(|res| {
+        res.map(Frame::data)
+            .map_err(|_| -> Infallible { unreachable!("We handle IO errors above") })
+    });
+    BoxBody::new(StreamBody::new(encoded_stream))
+}
+
+fn update_content_encoding(parts: &mut http::response::Parts, new_encoding: HeaderValue) {
+    if let Some(existing) = parts.headers.get(CONTENT_ENCODING) {
+        let mut encodings = existing.to_str().unwrap_or("").to_owned();
+        if !encodings.is_empty() {
+            encodings.push_str(", ");
+        }
+        encodings.push_str(new_encoding.to_str().unwrap());
+        parts
+            .headers
+            .insert(CONTENT_ENCODING, HeaderValue::from_str(&encodings).unwrap());
+    } else {
+        parts.headers.insert(CONTENT_ENCODING, new_encoding);
+    }
+}
+
+#[async_trait]
+impl MiddlewareLayer for Compression {
+    /// We'll apply compression on the response, where appropriate.
+    /// This is if:
+    /// * The response body is larger than the minimum size.
+    /// * The response content type is supported.
+    /// * The client supports the compression algorithm.
+    async fn after(&self, resp: HttpResponse, context: &mut HttpRequestContext) -> HttpResponse {
+        let body_size = resp.size_hint().exact();
+        let resp = resp;
+
+        // Don't compress if it's not an explicitly listed compressable type
+        if !self
+            .mime_types
+            .iter()
+            .any(|mt| mt.matches(&resp.headers().get_all(CONTENT_TYPE)))
+        {
+            debug!(
+                target: "middleware::compress",
+                "Mime type not supported for compression"
+            );
+            return resp;
+        }
+
+        // Don't compress streams unless compress streams is enabled.
+        if body_size.is_none() && !self.compress_streams {
+            debug!(
+                target: "middleware::compress",
+                "Stream compression disabled"
+            );
+            return resp;
+        }
+
+        // Don't compress too small bodies
+        if body_size.is_some_and(|s| s < self.min_size as u64) {
+            debug!(
+                target: "middleware::compress",
+                "Body size too small for compression"
+            );
+            return resp;
+        }
+
+        // Don't recompress if we're already compressed in a supported format
+        for existing_encoding in resp.headers().get_all(CONTENT_ENCODING) {
+            if let Ok(encodings) = existing_encoding.to_str() {
+                for encoding in encodings.split(',').map(str::trim) {
+                    let encoding = encoding.split(';').next().unwrap_or(encoding).trim();
+                    if self.algorithms.iter().any(|algo| algo.as_str() == encoding) {
+                        debug!(
+                            target: "middleware::compress",
+                            "Body already compressed with supported algorithm"
+                        );
+                        return resp;
+                    }
+                }
+            }
+        }
+
+        let compression_method = match find_first_supported(
+            &context.supported_encoding_set,
+            self.algorithms.iter().map(|algo| algo.as_str()),
+        ) {
+            Some("gzip") => CompressionAlgorithm::Gzip,
+            Some("br") => CompressionAlgorithm::Brotli,
+            Some("deflate") => CompressionAlgorithm::Deflate,
+            Some("zstd") => CompressionAlgorithm::Zstd,
+            _ => CompressionAlgorithm::Identity,
+        };
+
+        debug!(
+            target: "middleware::compress",
+            "Selected compression method: {:?}", compression_method
+        );
+        if matches!(compression_method, CompressionAlgorithm::Identity) {
+            return resp;
+        }
+
+        let (mut parts, body) = resp.into_parts();
+
+        let new_body = if let Some(_size) = body_size {
+            let full_bytes: Bytes = body
+                .into_data_stream()
+                .try_fold(BytesMut::new(), |mut acc, chunk| async move {
+                    acc.extend_from_slice(&chunk);
+                    Ok(acc)
+                })
+                .await
+                .unwrap()
+                .freeze();
+
+            let cursor = std::io::Cursor::new(full_bytes);
+            let reader = BufReader::new(cursor);
+            let compressed_bytes = match compression_method {
+                CompressionAlgorithm::Gzip => {
+                    let mut encoder =
+                        GzipEncoder::with_quality(reader, self.level.to_async_compression_level());
+                    let mut buf = Vec::new();
+                    encoder.read_to_end(&mut buf).await.unwrap();
+                    buf
+                }
+                CompressionAlgorithm::Brotli => {
+                    let mut encoder = BrotliEncoder::with_quality(
+                        reader,
+                        self.level.to_async_compression_level(),
+                    );
+                    let mut buf = Vec::new();
+                    encoder.read_to_end(&mut buf).await.unwrap();
+                    buf
+                }
+                CompressionAlgorithm::Deflate => {
+                    let mut encoder = DeflateEncoder::with_quality(
+                        reader,
+                        self.level.to_async_compression_level(),
+                    );
+                    let mut buf = Vec::new();
+                    encoder.read_to_end(&mut buf).await.unwrap();
+                    buf
+                }
+                CompressionAlgorithm::Zstd => {
+                    let mut encoder =
+                        ZstdEncoder::with_quality(reader, self.level.to_async_compression_level());
+                    let mut buf = Vec::new();
+                    encoder.read_to_end(&mut buf).await.unwrap();
+                    buf
+                }
+                CompressionAlgorithm::Identity => unreachable!(),
+            };
+            BoxBody::new(Full::new(Bytes::from(compressed_bytes)))
+        } else {
+            let stream = body
+                .into_data_stream()
+                .map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e));
+            let async_read_fut = StreamReader::new(stream);
+            let reader = BufReader::new(async_read_fut);
+            match compression_method {
+                CompressionAlgorithm::Gzip => stream_encode(GzipEncoder::with_quality(
+                    reader,
+                    self.level.to_async_compression_level(),
+                )),
+                CompressionAlgorithm::Brotli => stream_encode(BrotliEncoder::with_quality(
+                    reader,
+                    self.level.to_async_compression_level(),
+                )),
+                CompressionAlgorithm::Deflate => stream_encode(DeflateEncoder::with_quality(
+                    reader,
+                    self.level.to_async_compression_level(),
+                )),
+                CompressionAlgorithm::Zstd => stream_encode(ZstdEncoder::with_quality(
+                    reader,
+                    self.level.to_async_compression_level(),
+                )),
+                CompressionAlgorithm::Identity => unreachable!(),
+            }
+        };
+
+        update_content_encoding(&mut parts, compression_method.header_value());
+        parts.headers.remove(CONTENT_LENGTH);
+        debug!(
+            target: "middleware::compress",
+            "Response compressed"
+        );
+        Response::from_parts(parts, new_body)
+    }
+}
+impl FromValue for Compression {}