weaver / crates / weaver-app / src / middleware.rs

//! Axum middleware for host-based routing.

use std::sync::Arc;

use axum::{
    extract::Request,
    http::header::HOST,
    middleware::Next,
    response::{IntoResponse, Redirect, Response},
};
use jacquard::IntoStatic;
use jacquard::client::AgentSessionExt;
use jacquard::smol_str::SmolStr;
use jacquard::smol_str::ToSmolStr;

use crate::{
    env::WEAVER_APP_DOMAIN,
    fetch::Fetcher,
    host_mode::{CustomDomainContext, HostContext, SubdomainContext},
    subdomain_app::{extract_subdomain, lookup_subdomain_context},
};

/// Reserved subdomains that should not be used for notebooks.
const RESERVED_SUBDOMAINS: &[&str] = &[
    "www", "api", "admin", "app", "auth", "cdn", "alpha", "beta", "staging", "index",
];

/// Prefix for encoded publication hostnames (nb-{encoded}.weaver.sh).
const ENCODED_PUBLICATION_PREFIX: &str = "nb-";

/// Middleware that resolves host context and inserts it into request extensions.
///
/// This handles:
/// - Encoded publication hostnames (nb-xxx.pub.weaver.sh) → redirect to real domain
/// - Subdomain routing (notebook.weaver.sh) → SubdomainContext
/// - Custom domain routing (myblog.com) → CustomDomainContext
/// - Main domain (weaver.sh) → MainDomain
pub async fn host_context_middleware(
    mut req: Request,
    next: Next,
    fetcher: Arc<Fetcher>,
) -> Response {
    // Extract everything we need from the request BEFORE any async calls.
    // The body type isn't Sync, so we can't hold &req across await points.
    let host = extract_host(&req);
    let path = req.uri().path().to_string();
    let query = req
        .uri()
        .query()
        .map(|q| format!("?{}", q))
        .unwrap_or_default();

    let Some(host) = host else {
        tracing::warn!("No Host header in request");
        req.extensions_mut().insert(HostContext::MainDomain);
        return next.run(req).await;
    };

    // Strip port if present.
    let host_str = host.split(':').next().unwrap_or(&host).to_string();

    // Check if this is a weaver domain subdomain.
    if let Some(subdomain) = extract_subdomain(&host_str, WEAVER_APP_DOMAIN) {
        // Check for encoded publication hostname (nb-xxx.pub.weaver.sh).
        if let Some(redirect_url) =
            handle_encoded_publication(&subdomain, &path, &query, fetcher.clone()).await
        {
            return Redirect::permanent(&redirect_url).into_response();
        }

        // Check for reserved subdomains.
        if RESERVED_SUBDOMAINS.contains(&subdomain.as_str()) {
            tracing::debug!(subdomain, "Reserved subdomain, using main domain");
            req.extensions_mut().insert(HostContext::MainDomain);
            return next.run(req).await;
        }

        // Try to look up notebook for this subdomain.
        if let Some(ctx) = lookup_subdomain_context(&fetcher, &subdomain).await {
            tracing::debug!(subdomain, "Resolved subdomain context");
            req.extensions_mut().insert(HostContext::Subdomain(ctx));
            return next.run(req).await;
        }

        // Subdomain not found, fall through to main domain.
        tracing::warn!(subdomain, "Subdomain not found, using main domain");
        req.extensions_mut().insert(HostContext::MainDomain);
        return next.run(req).await;
    }

    // Not a weaver subdomain - check for custom domain.
    if let Some(ctx) = lookup_custom_domain(fetcher.clone(), &host_str).await {
        tracing::debug!(domain = host_str, "Resolved custom domain context");
        req.extensions_mut().insert(HostContext::CustomDomain(ctx));
        return next.run(req).await;
    }

    // Unknown domain, use main domain routing.
    tracing::debug!(host = host_str, "Unknown host, using main domain");
    req.extensions_mut().insert(HostContext::MainDomain);
    next.run(req).await
}

/// Extract the Host header from the request.
fn extract_host(req: &Request) -> Option<String> {
    req.headers()
        .get(HOST)?
        .to_str()
        .ok()
        .map(|s| s.to_string())
}

/// Handle encoded publication hostname (nb-{encoded}.weaver.sh).
///
/// Returns Some(redirect_url) if this is an encoded hostname that should redirect.
async fn handle_encoded_publication(
    subdomain: &str,
    path: &str,
    query: &str,
    fetcher: Arc<Fetcher>,
) -> Option<String> {
    // Check if subdomain matches pattern: nb-{encoded}
    if !subdomain.starts_with(ENCODED_PUBLICATION_PREFIX) {
        return None;
    }

    // Decode the publication subdomain.
    let Some((did, rkey)) = weaver_common::domain_encoding::decode_publication_subdomain(subdomain)
    else {
        tracing::warn!(subdomain, "Failed to decode publication subdomain");
        return None;
    };

    // Look up the publication to get the real domain.
    let domain = lookup_publication_domain(fetcher, did.as_str(), rkey.as_str()).await?;

    Some(format!("https://{}{}{}", domain, path, query))
}

/// Look up publication domain by DID and rkey.
async fn lookup_publication_domain(fetcher: Arc<Fetcher>, did: &str, rkey: &str) -> Option<String> {
    use jacquard::types::string::Uri;
    use weaver_api::site_standard::publication::Publication;

    let uri_str = format!("at://{}/site.standard.publication/{}", did, rkey);
    let uri = Publication::uri(&uri_str).ok()?;

    let record = fetcher.fetch_record(&uri).await.ok()?;

    // Extract domain from the already-parsed publication URL.
    match &record.value.url {
        Uri::Https(url) | Uri::Wss(url) => url.host_str().map(|h| h.to_string()),
        _ => None,
    }
}

/// Look up custom domain context.
async fn lookup_custom_domain(fetcher: Arc<Fetcher>, domain: &str) -> Option<CustomDomainContext> {
    use jacquard::prelude::XrpcClient;
    use weaver_api::sh_weaver::domain::resolve_by_domain::ResolveByDomain;

    let output = fetcher
        .send(ResolveByDomain::new().domain(domain).build())
        .await
        .ok()?
        .into_output()
        .ok()?;

    let pub_view = output.publication;

    Some(CustomDomainContext {
        domain: SmolStr::new(domain),
        owner: jacquard::types::string::AtIdentifier::Did(
            jacquard::types::string::Did::new(&pub_view.did)
                .ok()?
                .into_static(),
        ),
        publication_rkey: pub_view.rkey.to_smolstr(),
        publication_name: pub_view.name.to_smolstr(),
        notebook_uri: pub_view.notebook_uri.map(|u| u.to_smolstr()),
    })
}