hemx/slhx-axum/src/lib.rs

use axum::async_trait;
use axum::body::{to_bytes, Body};
use axum::extract::{FromRequest, FromRequestParts, Multipart};
use axum::http::{header, request::Parts, HeaderMap, HeaderValue, Request, Response, StatusCode};
use axum::response::sse::{Event, Sse};
use axum::response::IntoResponse;
use futures_util::{Stream, StreamExt};
use slhx_core::{BuildFingerprint, EffectBatch, Handle, IntoEffect};
use std::collections::BTreeMap;
use std::convert::Infallible;

pub const SLHX_PARTIAL_HEADER: &str = "x-slhx-partial";
pub const SLHX_FINGERPRINT_HEADER: &str = "x-slhx-fingerprint";
pub const SLHX_TITLE_HEADER: &str = "x-slhx-title";
pub const SLHX_CONTENT_TYPE: &str = "application/slhx";
pub const SLHX_HANDLE_FIELD: &str = "__h";
pub const SLHX_RUNTIME_CONTENT_TYPE: &str = "application/javascript; charset=utf-8";
pub const SLHX_SSE_EVENT: &str = "slhx";

#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
pub struct RuntimeJs;

pub const fn runtime_js() -> RuntimeJs {
    RuntimeJs
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum PageMode {
    Full,
    Partial,
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct PageRequest {
    pub mode: PageMode,
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct PageResponse {
    pub mode: PageMode,
    pub html: String,
    pub title: Option<String>,
    pub fingerprint: Option<BuildFingerprint>,
}

impl PageRequest {
    pub fn from_headers(headers: &HeaderMap) -> Self {
        Self {
            mode: PageMode::from_headers(headers),
        }
    }

    pub const fn is_partial(self) -> bool {
        matches!(self.mode, PageMode::Partial)
    }

    pub fn page(self, partial_html: impl Into<String>, shell: impl FnOnce(String) -> String) -> PageResponse {
        let partial_html = partial_html.into();
        match self.mode {
            PageMode::Full => PageResponse::full(shell(partial_html)),
            PageMode::Partial => PageResponse::partial(partial_html),
        }
    }
}

#[async_trait]
impl<S> FromRequestParts<S> for PageRequest
where
    S: Send + Sync,
{
    type Rejection = Infallible;

    async fn from_request_parts(parts: &mut Parts, _state: &S) -> Result<Self, Self::Rejection> {
        Ok(Self::from_headers(&parts.headers))
    }
}

impl PageResponse {
    pub fn full(html: impl Into<String>) -> Self {
        Self {
            mode: PageMode::Full,
            html: html.into(),
            title: None,
            fingerprint: None,
        }
    }

    pub fn partial(html: impl Into<String>) -> Self {
        Self {
            mode: PageMode::Partial,
            html: html.into(),
            title: None,
            fingerprint: None,
        }
    }

    pub fn title(mut self, title: impl Into<String>) -> Self {
        self.title = Some(title.into());
        self
    }

    pub fn fingerprint(mut self, fingerprint: BuildFingerprint) -> Self {
        self.fingerprint = Some(fingerprint);
        self
    }
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct EffectResponse {
    pub batch: EffectBatch,
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct InteractionFile {
    pub name: String,
    pub file_name: Option<String>,
    pub content_type: Option<String>,
    pub bytes: Vec<u8>,
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct InteractionForm {
    pub handle_id: u32,
    fields: Vec<(String, String)>,
    files: Vec<InteractionFile>,
}

pub struct HandlerRegistry {
    fingerprint: BuildFingerprint,
    handlers: BTreeMap<u32, Box<dyn Fn(InteractionForm) -> EffectBatch + Send + Sync>>,
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub enum InteractionFormRejection {
    InvalidBody,
    MissingHandle,
    InvalidHandle,
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub enum DispatchRejection {
    UnknownHandle(u32),
}

impl EffectResponse {
    pub fn new(effects: impl IntoEffect, fingerprint: BuildFingerprint) -> Self {
        Self {
            batch: effects.into_batch(fingerprint),
        }
    }
}

impl InteractionForm {
    pub fn new(handle_id: u32, fields: impl IntoIterator<Item = (String, String)>) -> Self {
        Self {
            handle_id,
            fields: fields.into_iter().collect(),
            files: Vec::new(),
        }
    }

    pub fn parse_urlencoded(body: &[u8]) -> Result<Self, InteractionFormRejection> {
        Self::from_parts(parse_urlencoded_pairs(body)?, Vec::new())
    }

    // req: multipart/001, req: multipart/002
    pub async fn parse_multipart(mut multipart: Multipart) -> Result<Self, InteractionFormRejection> {
        let mut fields = Vec::new();
        let mut files = Vec::new();

        while let Some(field) = multipart
            .next_field()
            .await
            .map_err(|_| InteractionFormRejection::InvalidBody)?
        {
            let Some(name) = field.name().map(str::to_owned) else {
                continue;
            };
            let file_name = field.file_name().map(str::to_owned);
            let content_type = field.content_type().map(str::to_owned);
            let bytes = field
                .bytes()
                .await
                .map_err(|_| InteractionFormRejection::InvalidBody)?;

            if file_name.is_some() {
                files.push(InteractionFile {
                    name,
                    file_name,
                    content_type,
                    bytes: bytes.to_vec(),
                });
            } else {
                let value = String::from_utf8(bytes.to_vec())
                    .map_err(|_| InteractionFormRejection::InvalidBody)?;
                fields.push((name, value));
            }
        }

        Self::from_parts(fields, files)
    }

    fn from_parts(fields: Vec<(String, String)>, files: Vec<InteractionFile>) -> Result<Self, InteractionFormRejection> {
        let Some(handle) = fields
            .iter()
            .find_map(|(name, value)| (name == SLHX_HANDLE_FIELD).then_some(value))
        else {
            return Err(InteractionFormRejection::MissingHandle);
        };
        let handle_id = handle
            .parse::<u32>()
            .map_err(|_| InteractionFormRejection::InvalidHandle)?;
        Ok(Self { handle_id, fields, files })
    }

    pub fn value(&self, name: &str) -> Option<&str> {
        self.fields
            .iter()
            .find_map(|(field, value)| (field == name).then_some(value.as_str()))
    }

    pub fn values<'a>(&'a self, name: &'a str) -> impl Iterator<Item = &'a str> + 'a {
        self.fields
            .iter()
            .filter_map(move |(field, value)| (field == name).then_some(value.as_str()))
    }

    pub fn fields(&self) -> &[(String, String)] {
        &self.fields
    }

    pub fn files(&self) -> &[InteractionFile] {
        &self.files
    }

    pub fn file(&self, name: &str) -> Option<&InteractionFile> {
        self.files.iter().find(|file| file.name == name)
    }
}

impl HandlerRegistry {
    pub const fn new(fingerprint: BuildFingerprint) -> Self {
        Self {
            fingerprint,
            handlers: BTreeMap::new(),
        }
    }

    pub fn register<E>(
        mut self,
        handle_id: u32,
        handler: impl Fn(InteractionForm) -> E + Send + Sync + 'static,
    ) -> Self
    where
        E: IntoEffect,
    {
        let fingerprint = self.fingerprint;
        self.handlers.insert(
            handle_id,
            Box::new(move |form| handler(form).into_batch(fingerprint)),
        );
        self
    }

    pub fn register_handle<I, E>(
        self,
        handle: Handle<I>,
        handler: impl Fn(InteractionForm) -> E + Send + Sync + 'static,
    ) -> Self
    where
        E: IntoEffect,
    {
        self.register(handle.id().id, handler)
    }

    pub fn dispatch(&self, form: InteractionForm) -> Result<EffectResponse, DispatchRejection> {
        let handle_id = form.handle_id;
        let Some(handler) = self.handlers.get(&handle_id) else {
            return Err(DispatchRejection::UnknownHandle(handle_id));
        };
        Ok(EffectResponse {
            batch: handler(form),
        })
    }

    pub fn contains(&self, handle_id: u32) -> bool {
        self.handlers.contains_key(&handle_id)
    }
}

impl IntoResponse for InteractionFormRejection {
    fn into_response(self) -> axum::response::Response {
        let (status, message) = match self {
            Self::InvalidBody => (StatusCode::BAD_REQUEST, "invalid slhx form body"),
            Self::MissingHandle => (StatusCode::BAD_REQUEST, "missing __h slhx handle field"),
            Self::InvalidHandle => (StatusCode::BAD_REQUEST, "invalid __h slhx handle field"),
        };
        (status, message).into_response()
    }
}

#[async_trait]
impl<S> FromRequest<S> for InteractionForm
where
    S: Send + Sync,
{
    type Rejection = InteractionFormRejection;

    async fn from_request(req: Request<Body>, _state: &S) -> Result<Self, Self::Rejection> {
        if is_multipart(req.headers()) {
            let multipart = Multipart::from_request(req, _state)
                .await
                .map_err(|_| InteractionFormRejection::InvalidBody)?;
            return Self::parse_multipart(multipart).await;
        }

        let bytes = to_bytes(req.into_body(), 1024 * 1024)
            .await
            .map_err(|_| InteractionFormRejection::InvalidBody)?;
        Self::parse_urlencoded(&bytes)
    }
}

fn is_multipart(headers: &HeaderMap) -> bool {
    headers
        .get(header::CONTENT_TYPE)
        .and_then(|value| value.to_str().ok())
        .is_some_and(|content_type| {
            content_type
                .split(';')
                .next()
                .is_some_and(|mime| mime.trim().eq_ignore_ascii_case("multipart/form-data"))
        })
}

impl PageMode {
    pub fn from_headers(headers: &HeaderMap) -> Self {
        match headers.get(SLHX_PARTIAL_HEADER).and_then(|value| value.to_str().ok()) {
            Some("1" | "true") => Self::Partial,
            _ => Self::Full,
        }
    }
}

impl IntoResponse for PageResponse {
    fn into_response(self) -> axum::response::Response {
        let html = match self.fingerprint {
            Some(fingerprint) => html_with_root_fingerprint(self.html, fingerprint),
            None => self.html,
        };
        let mut response = Response::new(Body::from(html));
        response.headers_mut().insert(
            header::CONTENT_TYPE,
            HeaderValue::from_static("text/html; charset=utf-8"),
        );
        if self.mode == PageMode::Partial {
            response
                .headers_mut()
                .insert(SLHX_PARTIAL_HEADER, HeaderValue::from_static("true"));
        }
        if let Some(fingerprint) = self.fingerprint.and_then(fingerprint_header) {
            response
                .headers_mut()
                .insert(SLHX_FINGERPRINT_HEADER, fingerprint);
        }
        if let Some(title) = self.title.and_then(|title| HeaderValue::from_str(&title).ok()) {
            response.headers_mut().insert(SLHX_TITLE_HEADER, title);
        }
        response
    }
}

fn html_with_root_fingerprint(mut html: String, fingerprint: BuildFingerprint) -> String {
    if html.contains("data-slhx-fp=") {
        return html;
    }
    let Some(root_attr) = html.find("data-slhx-root") else {
        return html;
    };
    let Some(tag_start) = html[..root_attr].rfind('<') else {
        return html;
    };
    let Some(tag_end) = html[tag_start..].find('>') else {
        return html;
    };
    let insert_at = tag_start + tag_end;
    html.insert_str(insert_at, &format!(" data-slhx-fp=\"{}\"", fingerprint.0));
    html
}

fn fingerprint_header(fingerprint: BuildFingerprint) -> Option<HeaderValue> {
    HeaderValue::from_str(&fingerprint.0.to_string()).ok()
}

impl IntoResponse for EffectResponse {
    fn into_response(self) -> axum::response::Response {
        let bytes = self.batch.to_wire();
        let mut response = Response::new(Body::from(bytes));
        response.headers_mut().insert(
            header::CONTENT_TYPE,
            HeaderValue::from_static(SLHX_CONTENT_TYPE),
        );
        if let Some(fingerprint) = fingerprint_header(self.batch.fingerprint) {
            response
                .headers_mut()
                .insert(SLHX_FINGERPRINT_HEADER, fingerprint);
        }
        response
    }
}

impl IntoResponse for DispatchRejection {
    fn into_response(self) -> axum::response::Response {
        match self {
            Self::UnknownHandle(handle_id) => (
                StatusCode::NOT_FOUND,
                format!("unknown slhx handle id {handle_id}"),
            )
                .into_response(),
        }
    }
}

impl IntoResponse for RuntimeJs {
    fn into_response(self) -> axum::response::Response {
        let mut response = Response::new(Body::from(slhx_js::RUNTIME_JS));
        response.headers_mut().insert(
            header::CONTENT_TYPE,
            HeaderValue::from_static(SLHX_RUNTIME_CONTENT_TYPE),
        );
        response.headers_mut().insert(
            header::CACHE_CONTROL,
            HeaderValue::from_static("public, max-age=31536000, immutable"),
        );
        response
    }
}

pub fn runtime_js_source() -> &'static str {
    slhx_js::RUNTIME_JS
}

// req: push/001, req: push/003, req: push/004
pub fn sse<S, E>(batches: S) -> Sse<impl Stream<Item = Result<Event, E>> + Send>
where
    S: Stream<Item = Result<EffectBatch, E>> + Send + 'static,
    E: Into<axum::BoxError>,
{
    Sse::new(batches.map(|batch| batch.map(sse_event)))
}

pub fn sse_event(batch: EffectBatch) -> Event {
    Event::default()
        .event(SLHX_SSE_EVENT)
        .data(encode_sse_batch(&batch))
}

pub fn encode_sse_batch(batch: &EffectBatch) -> String {
    base64_url_no_pad(&batch.to_wire())
}

fn base64_url_no_pad(input: &[u8]) -> String {
    const ALPHABET: &[u8; 64] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_";
    let mut out = String::with_capacity((input.len() * 4).div_ceil(3));
    let mut chunks = input.chunks_exact(3);
    for chunk in &mut chunks {
        let n = ((chunk[0] as u32) << 16) | ((chunk[1] as u32) << 8) | chunk[2] as u32;
        out.push(ALPHABET[((n >> 18) & 0x3f) as usize] as char);
        out.push(ALPHABET[((n >> 12) & 0x3f) as usize] as char);
        out.push(ALPHABET[((n >> 6) & 0x3f) as usize] as char);
        out.push(ALPHABET[(n & 0x3f) as usize] as char);
    }
    match chunks.remainder() {
        [a] => {
            let n = (*a as u32) << 16;
            out.push(ALPHABET[((n >> 18) & 0x3f) as usize] as char);
            out.push(ALPHABET[((n >> 12) & 0x3f) as usize] as char);
        }
        [a, b] => {
            let n = ((*a as u32) << 16) | ((*b as u32) << 8);
            out.push(ALPHABET[((n >> 18) & 0x3f) as usize] as char);
            out.push(ALPHABET[((n >> 12) & 0x3f) as usize] as char);
            out.push(ALPHABET[((n >> 6) & 0x3f) as usize] as char);
        }
        [] => {}
        _ => unreachable!(),
    }
    out
}

fn parse_urlencoded_pairs(body: &[u8]) -> Result<Vec<(String, String)>, InteractionFormRejection> {
    if body.is_empty() {
        return Ok(Vec::new());
    }

    body.split(|byte| *byte == b'&')
        .map(|pair| {
            let equals = pair.iter().position(|byte| *byte == b'=');
            let (name, value) = match equals {
                Some(index) => (&pair[..index], &pair[index + 1..]),
                None => (pair, &[][..]),
            };
            Ok((percent_decode(name)?, percent_decode(value)?))
        })
        .collect()
}

fn percent_decode(input: &[u8]) -> Result<String, InteractionFormRejection> {
    let mut out = Vec::with_capacity(input.len());
    let mut i = 0;
    while i < input.len() {
        match input[i] {
            b'+' => {
                out.push(b' ');
                i += 1;
            }
            b'%' if i + 2 < input.len() => {
                let high = hex(input[i + 1]).ok_or(InteractionFormRejection::InvalidBody)?;
                let low = hex(input[i + 2]).ok_or(InteractionFormRejection::InvalidBody)?;
                out.push((high << 4) | low);
                i += 3;
            }
            b'%' => return Err(InteractionFormRejection::InvalidBody),
            byte => {
                out.push(byte);
                i += 1;
            }
        }
    }
    String::from_utf8(out).map_err(|_| InteractionFormRejection::InvalidBody)
}

fn hex(byte: u8) -> Option<u8> {
    match byte {
        b'0'..=b'9' => Some(byte - b'0'),
        b'a'..=b'f' => Some(byte - b'a' + 10),
        b'A'..=b'F' => Some(byte - b'A' + 10),
        _ => None,
    }
}

#[cfg(test)]
mod tests {
    use super::{encode_sse_batch, html_with_root_fingerprint, sse, BuildFingerprint, InteractionForm, SLHX_SSE_EVENT};
    use axum::{body::{to_bytes, Body}, extract::FromRequest, http::{header, Request}, response::IntoResponse};
    use futures_util::stream;
    use scraper::{Html, Selector};
    use slhx_core::{EffectBatch, EFFECT_BATCH_ABI_VERSION};
    use std::convert::Infallible;

    fn selector(value: &str) -> Selector {
        Selector::parse(value).expect("test selector parses")
    }

    #[test]
    fn root_fingerprint_is_added_to_initial_root() {
        // req: test/005
        let html = html_with_root_fingerprint(
            "<html><body><main data-slhx-root>Docs</main></body></html>".into(),
            BuildFingerprint(99),
        );

        let document = Html::parse_document(&html);
        let root = document
            .select(&selector("main[data-slhx-root]"))
            .next()
            .expect("root element is rendered");
        assert_eq!(root.value().attr("data-slhx-fp"), Some("99"));
        assert_eq!(root.text().collect::<String>(), "Docs");
    }

    #[test]
    fn existing_root_fingerprint_is_preserved() {
        // req: test/005
        let html = html_with_root_fingerprint(
            "<main data-slhx-root data-slhx-fp=\"1\">Docs</main>".into(),
            BuildFingerprint(99),
        );

        let document = Html::parse_fragment(&html);
        let root = document
            .select(&selector("main[data-slhx-root]"))
            .next()
            .expect("root element is rendered");
        assert_eq!(root.value().attr("data-slhx-fp"), Some("1"));
        assert_eq!(root.text().collect::<String>(), "Docs");
    }

    #[tokio::test]
    async fn sse_response_streams_base64_url_effect_batches() {
        let batch = EffectBatch {
            abi_version: EFFECT_BATCH_ABI_VERSION,
            fingerprint: BuildFingerprint(11),
            ops: Vec::new(),
        };
        let encoded = encode_sse_batch(&batch);

        let response = sse(stream::iter([Ok::<_, Infallible>(batch)])).into_response();
        assert_eq!(
            response.headers().get(header::CONTENT_TYPE).unwrap(),
            "text/event-stream"
        );
        let body = to_bytes(response.into_body(), 1024).await.unwrap();

        assert_eq!(
            String::from_utf8(body.to_vec()).unwrap(),
            format!("event: {SLHX_SSE_EVENT}\ndata: {encoded}\n\n")
        );
        assert!(!encoded.contains('='));
    }

    #[tokio::test]
    async fn interaction_form_extracts_multipart_fields_and_files() {
        let boundary = "slhx-test-boundary";
        let body = concat!(
            "--slhx-test-boundary\r\n",
            "Content-Disposition: form-data; name=\"__h\"\r\n\r\n",
            "7\r\n",
            "--slhx-test-boundary\r\n",
            "Content-Disposition: form-data; name=\"title\"\r\n\r\n",
            "Report\r\n",
            "--slhx-test-boundary\r\n",
            "Content-Disposition: form-data; name=\"upload\"; filename=\"a.txt\"\r\n",
            "Content-Type: text/plain\r\n\r\n",
            "hello\r\n",
            "--slhx-test-boundary--\r\n",
        );
        let request = Request::builder()
            .header(
                axum::http::header::CONTENT_TYPE,
                format!("multipart/form-data; boundary={boundary}"),
            )
            .body(Body::from(body))
            .unwrap();

        let form = InteractionForm::from_request(request, &()).await.unwrap();

        assert_eq!(form.handle_id, 7);
        assert_eq!(form.value("title"), Some("Report"));
        let file = form.file("upload").unwrap();
        assert_eq!(file.file_name.as_deref(), Some("a.txt"));
        assert_eq!(file.content_type.as_deref(), Some("text/plain"));
        assert_eq!(file.bytes, b"hello");
    }
}