braid-core 0.1.4

//! Axum middleware for Braid protocol support.
//!
//! Provides an Axum layer that extracts Braid protocol information from incoming
//! requests and makes it available to handlers.

use super::resource_state::ResourceStateManager;
use crate::core::protocol_mod as protocol;
use crate::core::protocol_mod::constants::headers;
use crate::core::Version;
use axum::{extract::Request, middleware::Next, response::Response};
use futures::StreamExt;
use std::collections::BTreeMap;
use std::sync::Arc;

pub use crate::core::protocol_mod::BraidState;

/// Newtype wrapper indicating Firefox browser detection.
#[derive(Clone, Copy, Debug)]
pub struct IsFirefox(pub bool);

/// Middleware handler function for use with `from_fn_with_state`.
async fn braid_middleware_handler(
    axum::extract::State(state): axum::extract::State<BraidLayer>,
    req: Request,
    next: Next,
) -> Response {
    state.handle_middleware(req, next).await
}

/// Axum middleware layer for Braid protocol support.
#[derive(Clone)]
pub struct BraidLayer {
    config: super::config::ServerConfig,
    pub resource_manager: Arc<ResourceStateManager>,
    pub multiplexer_registry: Arc<super::multiplex::MultiplexerRegistry>,
}

impl BraidLayer {
    /// Create a new Braid layer with default configuration.
    #[must_use]
    pub fn new() -> Self {
        Self {
            config: super::config::ServerConfig::default(),
            resource_manager: Arc::new(ResourceStateManager::new()),
            multiplexer_registry: Arc::new(super::multiplex::MultiplexerRegistry::new()),
        }
    }

    /// Create a new Braid layer with custom configuration.
    #[must_use]
    pub fn with_config(config: super::config::ServerConfig) -> Self {
        Self {
            config,
            resource_manager: Arc::new(ResourceStateManager::new()),
            multiplexer_registry: Arc::new(super::multiplex::MultiplexerRegistry::new()),
        }
    }

    /// Get a reference to the layer's configuration.
    #[inline]
    #[must_use]
    pub fn config(&self) -> &super::config::ServerConfig {
        &self.config
    }

    /// Create the middleware function for use with Axum.
    #[must_use]
    pub fn middleware(
        &self,
    ) -> impl tower::Layer<
        axum::routing::Route,
        Service = impl tower::Service<
            Request,
            Response = Response,
            Error = std::convert::Infallible,
            Future = impl Send + 'static,
        > + Clone
                      + Send
                      + Sync
                      + 'static,
    > + Clone {
        axum::middleware::from_fn_with_state(self.clone(), braid_middleware_handler)
    }

    async fn handle_middleware(&self, mut req: Request, next: Next) -> Response {
        let resource_manager = self.resource_manager.clone();
        let multiplexer_registry = self.multiplexer_registry.clone();

        if req.method().as_str() == "MULTIPLEX" {
            let version = req
                .headers()
                .get(headers::MULTIPLEX_VERSION)
                .and_then(|v| v.to_str().ok())
                .unwrap_or("1.0");

            if version == "1.0" {
                let (tx, mut rx) = tokio::sync::mpsc::channel(1024);
                let id = format!("{:x}", rand::random::<u64>());

                multiplexer_registry.add(id.clone(), tx).await;

                let stream = async_stream::stream! {
                    while let Some(data) = rx.recv().await {
                        yield Ok::<_, std::io::Error>(axum::body::Bytes::from(data));
                    }
                };

                let body = axum::body::Body::from_stream(stream);
                return Response::builder()
                    .status(200)
                    .header(headers::MULTIPLEX_VERSION, "1.0")
                    .body(body)
                    .unwrap();
            }
        }

        let braid_state = BraidState::from_headers(req.headers());
        let multiplex_through = braid_state.multiplex_through.clone();
        let m_registry = multiplexer_registry.clone();

        let is_firefox = req
            .headers()
            .get("user-agent")
            .and_then(|v| v.to_str().ok())
            .map(|ua| ua.to_lowercase().contains("firefox"))
            .unwrap_or(false);

        req.extensions_mut().insert(Arc::new(braid_state));
        req.extensions_mut().insert(resource_manager);
        req.extensions_mut().insert(multiplexer_registry);
        req.extensions_mut().insert(IsFirefox(is_firefox));

        let mut response = next.run(req).await;

        let headers = response.headers_mut();
        headers.insert(
            axum::http::header::HeaderName::from_static("range-request-allow-methods"),
            axum::http::header::HeaderValue::from_static("PATCH, PUT"),
        );
        headers.insert(
            axum::http::header::HeaderName::from_static("range-request-allow-units"),
            axum::http::header::HeaderValue::from_static("json"),
        );

        if let Some(through) = multiplex_through {
            let parts: Vec<&str> = through.split('/').collect();
            if parts.len() >= 5 && parts[1] == ".well-known" && parts[2] == "multiplexer" {
                let m_id = parts[3];
                let r_id = parts[4];

                if let Some(conn) = m_registry.get(m_id).await {
                    let sender = conn.sender.clone();
                    let r_id = r_id.to_string();

                    let mut cors_headers = axum::http::HeaderMap::new();
                    for (k, v) in response.headers() {
                        let k_str = k.as_str();
                        if k_str.starts_with("access-control-") {
                            cors_headers.insert(k.clone(), v.clone());
                        }
                    }

                    tokio::spawn(async move {
                        let mut header_block =
                            format!(":status: {}\r\n", response.status().as_u16());
                        for (name, value) in response.headers() {
                            header_block.push_str(&format!(
                                "{}: {}\r\n",
                                name,
                                value.to_str().unwrap_or("")
                            ));
                        }
                        header_block.push_str("\r\n");

                        let start_evt = protocol::multiplex::MultiplexEvent::Data(
                            r_id.clone(),
                            header_block.clone().into_bytes(),
                        );
                        let _ = sender.send(start_evt.to_string().into_bytes()).await;
                        let _ = sender.send(header_block.into_bytes()).await;

                        let mut body_stream = response.into_body().into_data_stream();
                        while let Some(Ok(chunk)) = body_stream.next().await {
                            let data_evt = protocol::multiplex::MultiplexEvent::Data(
                                r_id.clone(),
                                chunk.to_vec(),
                            );
                            let _ = sender.send(data_evt.to_string().into_bytes()).await;
                            let _ = sender.send(chunk.to_vec()).await;
                        }

                        let close_evt = protocol::multiplex::MultiplexEvent::CloseResponse(r_id);
                        let _ = sender.send(close_evt.to_string().into_bytes()).await;
                    });

                    let mut builder = Response::builder()
                        .status(293)
                        .header(headers::MULTIPLEX_VERSION, "1.0");

                    if let Some(headers) = builder.headers_mut() {
                        headers.extend(cors_headers);
                    }

                    return builder.body(axum::body::Body::empty()).unwrap();
                }
            }
        }

        response
    }
}

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

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_parse_version_header() {
        let result = protocol::parse_version_header("\"v1\", \"v2\", \"v3\"");
        assert!(result.is_ok());
        assert_eq!(result.unwrap().len(), 3);
    }
}