jokoway 0.1.0-rc.1

Jokoway is a high-performance API Gateway built on Pingora (Rust) with dead-simple YAML configs.
Documentation
use jokoway_transformer::{RequestTransformer, ResponseTransformer};

use bytes::BytesMut;
use flate2::Decompress;
pub use jokoway_core::{AppContext, Context, RequestContext};
use pingora::protocols::http::bridge::grpc_web::GrpcWebCtx;
use std::any::Any;
use std::sync::Arc;

#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum GrpcMode {
    #[default]
    None,
    Native,
    Web,
}

pub struct ProxyContext {
    pub upstream_name: Option<Arc<str>>,
    pub response_transformer: Option<Arc<dyn ResponseTransformer>>,
    pub req_transformer: Option<Arc<dyn RequestTransformer>>,
    pub is_upgrade: bool,
    pub grpc_mode: GrpcMode,
    pub grpc_web: GrpcWebCtx,
    pub ws_client_buf: BytesMut,
    pub ws_upstream_buf: BytesMut,
    pub grpc_client_buf: BytesMut,
    pub grpc_upstream_buf: BytesMut,
    pub rewrite_host: Option<String>,
    pub max_retries: u32,
    pub retries_attempted: u32,

    pub ws_client_decompressor: Option<Decompress>,
    pub ws_upstream_decompressor: Option<Decompress>,

    pub middleware_ctx: Vec<Box<dyn Any + Send + Sync>>,
    pub request_ctx: RequestContext,
}

impl ProxyContext {
    /// Create a new ProxyContext with optimized buffer sizes
    pub fn new() -> Self {
        Self {
            upstream_name: None,
            response_transformer: None,
            req_transformer: None,
            is_upgrade: false,
            grpc_mode: GrpcMode::None,
            grpc_web: GrpcWebCtx::default(),
            // Pre-allocate reasonable buffer sizes for WebSocket frames
            ws_client_buf: BytesMut::with_capacity(4096),
            ws_upstream_buf: BytesMut::with_capacity(4096),
            grpc_client_buf: BytesMut::with_capacity(4096),
            grpc_upstream_buf: BytesMut::with_capacity(4096),
            rewrite_host: None,
            max_retries: 1,
            retries_attempted: 0,

            ws_client_decompressor: None,
            ws_upstream_decompressor: None,

            middleware_ctx: Vec::new(),
            request_ctx: RequestContext::new(),
        }
    }

    /// Clear websocket buffers for reuse
    pub fn clear_ws_buffers(&mut self) {
        self.ws_client_buf.clear();
        self.ws_upstream_buf.clear();
    }

    /// Clear grpc buffers for reuse
    pub fn clear_grpc_buffers(&mut self) {
        self.grpc_client_buf.clear();
        self.grpc_upstream_buf.clear();
    }
}

impl Default for ProxyContext {
    fn default() -> Self {
        Self::new()
    }
}

#[cfg(test)]
mod tests {
    use super::{AppContext, Context};

    #[test]
    fn app_ctx_insert_get_remove() {
        let ctx = AppContext::new();

        assert!(ctx.get::<usize>().is_none());
        ctx.insert(12usize);
        assert_eq!(*ctx.get::<usize>().unwrap(), 12);

        ctx.insert(24usize);
        assert_eq!(*ctx.get::<usize>().unwrap(), 24);

        let removed = ctx.remove::<usize>().unwrap();
        assert_eq!(*removed, 24);
        assert!(ctx.get::<usize>().is_none());
    }

    #[test]
    fn app_ctx_handles_multiple_types() {
        let ctx = AppContext::new();

        ctx.insert(10usize);
        ctx.insert("jokoway".to_string());

        assert_eq!(*ctx.get::<usize>().unwrap(), 10);
        assert_eq!(&*ctx.get::<String>().unwrap(), "jokoway");
    }

    #[test]
    fn app_ctx_remove_missing_returns_none() {
        let ctx = AppContext::new();
        assert!(ctx.remove::<u64>().is_none());
    }

    #[derive(Debug, PartialEq, Eq)]
    struct CustomData {
        id: u32,
        label: String,
    }

    #[test]
    fn app_ctx_store_custom_struct() {
        let ctx = AppContext::new();
        let data = CustomData {
            id: 7,
            label: "alpha".to_string(),
        };

        ctx.insert(data);

        let stored = ctx.get::<CustomData>().unwrap();
        assert_eq!(stored.id, 7);
    }
}