perfect-postcode/server-rs/src/routes/pb_proxy.rs

use std::sync::{Arc, LazyLock};
use std::time::Duration;

use axum::body::Body;
use axum::extract::{Request, State};
use axum::http::{HeaderName, StatusCode};
use axum::response::{IntoResponse, Response};
use tracing::warn;

use crate::state::SharedState;

/// PocketBase API paths the frontend is allowed to reach via /pb/*.
/// Everything else (admins API, settings, logs, backups, collection schema,
/// arbitrary collection records like checkout_sessions/invites/short_urls)
/// is rejected at the proxy layer as defense-in-depth on top of PB's own
/// collection rules.
fn is_allowed_pb_path(path: &str) -> bool {
    // Exact paths
    if matches!(
        path,
        "/api/health" | "/api/oauth2-redirect" | "/api/realtime"
    ) {
        return true;
    }
    // Prefix-allowed paths. The trailing slash is intentional. Without it,
    // `/api/collections/users` (the schema endpoint) would match.
    const ALLOWED_PREFIXES: &[&str] = &[
        "/api/collections/users/",
        "/api/collections/saved_searches/",
        "/api/files/",
    ];
    ALLOWED_PREFIXES
        .iter()
        .any(|prefix| path.starts_with(prefix))
}

/// Dedicated HTTP client for proxying: does not follow redirects so 3xx
/// responses are passed through to the browser (needed for OAuth flows).
/// No client-wide timeout because SSE (Server-Sent Events) connections used
/// by PocketBase realtime/OAuth2 are long-lived streams; non-realtime
/// requests get a per-request timeout instead (see PROXY_REQUEST_TIMEOUT).
static PROXY_CLIENT: LazyLock<reqwest::Client> = LazyLock::new(|| {
    reqwest::Client::builder()
        .redirect(reqwest::redirect::Policy::none())
        .connect_timeout(Duration::from_secs(5))
        .referer(false)
        .build()
        .expect("Failed to build proxy HTTP client")
});

/// Timeout for proxied requests other than the realtime SSE stream, so a hung
/// PocketBase cannot pile up handlers indefinitely. Generous enough for file
/// uploads/downloads over slow links.
const PROXY_REQUEST_TIMEOUT: Duration = Duration::from_secs(30);

pub async fn proxy_to_pocketbase(
    State(shared): State<Arc<SharedState>>,
    req: Request,
) -> impl IntoResponse {
    let state = shared.load_state();
    let pb_url = state.pocketbase_url.trim_end_matches('/');

    let path = req.uri().path();
    let target_path = path.strip_prefix("/pb").unwrap_or(path);
    if !is_allowed_pb_path(target_path) {
        warn!(path = %target_path, "Rejected PocketBase proxy request to disallowed path");
        return StatusCode::NOT_FOUND.into_response();
    }
    let query = req
        .uri()
        .query()
        .map(|qs| format!("?{qs}"))
        .unwrap_or_default();
    let url = format!("{pb_url}{target_path}{query}");

    let method = req.method().clone();
    let mut builder = PROXY_CLIENT.request(method, &url);

    // The realtime SSE stream is intentionally unbounded; everything else
    // must complete within the timeout.
    if target_path != "/api/realtime" {
        builder = builder.timeout(PROXY_REQUEST_TIMEOUT);
    }

    // Forward only safe headers (allowlist)
    const ALLOWED_HEADERS: &[&str] = &[
        "content-type",
        "accept",
        "authorization",
        "cookie",
        "accept-language",
    ];
    for (name, value) in req.headers() {
        if ALLOWED_HEADERS.contains(&name.as_str()) {
            builder = builder.header(name.clone(), value.clone());
        }
    }

    // Do not forward client-supplied X-Forwarded-For/X-Real-IP. PocketBase
    // may use trusted proxy headers for rate limits, so accepting public
    // values here lets callers choose their own source IP.

    // Forward body
    let body_bytes = match axum::body::to_bytes(req.into_body(), 10 * 1024 * 1024).await {
        Ok(bytes) => bytes,
        Err(err) => {
            warn!("Failed to read request body: {err}");
            return (StatusCode::BAD_REQUEST, "Failed to read request body").into_response();
        }
    };
    builder = builder.body(body_bytes);

    match builder.send().await {
        Ok(upstream) => {
            let status = upstream.status();
            let mut response = Response::builder().status(status);

            for (name, value) in upstream.headers() {
                // Skip hop-by-hop headers
                if name == "transfer-encoding" {
                    continue;
                }
                match HeaderName::from_bytes(name.as_ref()) {
                    Ok(header_name) => {
                        response = response.header(header_name, value.clone());
                    }
                    Err(err) => {
                        warn!(header = ?name, error = %err, "Skipping unparseable upstream header");
                    }
                }
            }

            // Stream the response body instead of buffering it entirely.
            // This is critical for SSE (Server-Sent Events) used by PocketBase's
            // realtime system and OAuth2 flow: buffering would hang forever
            // since SSE responses never complete.
            let body = Body::from_stream(upstream.bytes_stream());
            response.body(body).unwrap_or_else(|err| {
                warn!("Failed to build proxied response: {err}");
                (StatusCode::BAD_GATEWAY, "Invalid upstream response").into_response()
            })
        }
        Err(err) => {
            warn!("PocketBase proxy error: {err}");
            (StatusCode::BAD_GATEWAY, "PocketBase unavailable").into_response()
        }
    }
}