Firelink/src-tauri/src/extension_server.rs

use axum::{
    body::{Body, Bytes},
    extract::State,
    http::{HeaderMap, HeaderValue, Method, Request, StatusCode},
    middleware::{self, Next},
    response::Response,
    routing::{get, post},
    Router,
};
use hmac::{Hmac, KeyInit, Mac};
use reqwest::Url;
use serde::{Deserialize, Serialize};
use sha2::Sha256;
use std::collections::{HashMap, HashSet};
use std::path::Path;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex, RwLock};
use std::time::{SystemTime, UNIX_EPOCH};
use tauri::{AppHandle, Emitter, Manager};
use tokio::sync::watch;
use tower_http::cors::{Any, CorsLayer};
use ts_rs::TS;

pub const EXTENSION_SERVER_PORT: u16 = 6412;
pub const EXTENSION_SERVER_PORT_RANGE: std::ops::RangeInclusive<u16> = EXTENSION_SERVER_PORT..=6422;
const MAX_URL_COUNT: usize = 200;
const SIGNATURE_MAX_AGE_MS: u64 = 60_000;
const SERVER_HEADER: &str = "x-firelink-server";
const PROTOCOL_VERSION_HEADER: &str = "x-firelink-protocol-version";
const CLIENT_NONCE_HEADER: &str = "x-firelink-client-nonce";
const SERVER_PROOF_HEADER: &str = "x-firelink-server-proof";
const SERVER_PORT_HEADER: &str = "x-firelink-server-port";
const SMOKE_PROCESS_ID_HEADER: &str = "x-firelink-smoke-process-id";
const SERVER_PROOF_PREFIX: &[u8] = b"firelink-server-proof\n";
const PROTOCOL_VERSION: &str = "4";

type HmacSha256 = Hmac<Sha256>;
pub type SharedExtensionToken = Arc<RwLock<String>>;
pub type SharedFrontendReady = Arc<AtomicBool>;
pub type SharedServerPort = Arc<RwLock<Option<u16>>>;
type ReplayCache = Arc<Mutex<HashMap<String, u64>>>;

#[derive(Clone)]
pub struct ServerState {
    pub app_handle: AppHandle,
    pub pairing_token: SharedExtensionToken,
    pub frontend_ready: SharedFrontendReady,
    pub replay_cache: ReplayCache,
    pub bound_port: u16,
}

#[derive(Deserialize)]
struct ExtensionRequest {
    urls: Vec<String>,
    #[serde(default)]
    referer: Option<String>,
    #[serde(default)]
    silent: bool,
    #[serde(default)]
    filename: Option<String>,
    #[serde(default)]
    headers: Option<String>,
    #[serde(default)]
    cookies: Option<String>,
    #[serde(default)]
    media: bool,
}

#[derive(Clone, Serialize, TS)]
#[ts(export, export_to = "../../src/bindings/")]
pub struct ExtensionDownload {
    urls: Vec<String>,
    referer: Option<String>,
    silent: bool,
    filename: Option<String>,
    headers: Option<String>,
    cookies: Option<String>,
    media: bool,
}

pub async fn start_server(
    app_handle: AppHandle,
    pairing_token: SharedExtensionToken,
    frontend_ready: SharedFrontendReady,
    server_port: SharedServerPort,
    mut shutdown_rx: watch::Receiver<bool>,
) -> Result<(), String> {
    let (port, listener) = bind_extension_listener().await?;
    let state = ServerState {
        app_handle,
        pairing_token,
        frontend_ready,
        replay_cache: Arc::new(Mutex::new(HashMap::new())),
        bound_port: port,
    };

    let cors = CorsLayer::new()
        .allow_origin(tower_http::cors::AllowOrigin::predicate(|origin, _| {
            is_allowed_origin(origin.to_str().unwrap_or(""))
        }))
        .allow_methods([Method::GET, Method::POST, Method::OPTIONS])
        .allow_headers(Any)
        .expose_headers(Any);

    let app = Router::new()
        .route("/ping", get(ping_handler))
        .route("/download", post(download_handler))
        .layer(cors)
        .layer(middleware::from_fn(add_server_identity))
        .with_state(state);

    if let Ok(mut current_port) = server_port.write() {
        *current_port = Some(port);
    }

    log::info!("Browser extension server bound to 127.0.0.1:{port}");

    let server_result = axum::serve(listener, app)
        .with_graceful_shutdown(async move {
            if *shutdown_rx.borrow() {
                return;
            }
            let _ = shutdown_rx.changed().await;
        })
        .await
        .map_err(|e| format!("Server error: {}", e));

    if let Ok(mut current_port) = server_port.write() {
        *current_port = None;
    }

    server_result
}

async fn add_server_identity(request: Request<Body>, next: Next) -> Response {
    let mut response = next.run(request).await;
    response
        .headers_mut()
        .insert(SERVER_HEADER, HeaderValue::from_static("1"));
    response.headers_mut().insert(
        PROTOCOL_VERSION_HEADER,
        HeaderValue::from_static(PROTOCOL_VERSION),
    );
    if std::env::var_os("FIRELINK_SMOKE_TEST").is_some() {
        if let Ok(process_id) = HeaderValue::from_str(&std::process::id().to_string()) {
            response
                .headers_mut()
                .insert(SMOKE_PROCESS_ID_HEADER, process_id);
        }
    }
    response
}

async fn bind_extension_listener() -> Result<(u16, tokio::net::TcpListener), String> {
    let mut errors = Vec::new();
    for port in EXTENSION_SERVER_PORT_RANGE {
        match tokio::net::TcpListener::bind(("127.0.0.1", port)).await {
            Ok(listener) => return Ok((port, listener)),
            Err(error) => {
                errors.push(format!("{port}: {error}"));
            }
        }
    }
    Err(format!(
        "Failed to bind extension server in port range {}-{} ({})",
        EXTENSION_SERVER_PORT,
        *EXTENSION_SERVER_PORT_RANGE.end(),
        errors.join("; ")
    ))
}

async fn ping_handler(
    State(state): State<ServerState>,
    headers: HeaderMap,
    body: Bytes,
) -> Result<Response, StatusCode> {
    let signature = match headers
        .get("x-firelink-signature")
        .and_then(|v| v.to_str().ok())
    {
        Some(v) => v,
        None => return Err(StatusCode::FORBIDDEN),
    };

    let timestamp_str = match headers
        .get("x-firelink-timestamp")
        .and_then(|v| v.to_str().ok())
    {
        Some(v) => v,
        None => return Err(StatusCode::FORBIDDEN),
    };

    let nonce = match headers
        .get(CLIENT_NONCE_HEADER)
        .and_then(|v| v.to_str().ok())
        .filter(|value| is_valid_client_nonce(value))
    {
        Some(v) => v,
        None => return Err(StatusCode::FORBIDDEN),
    };

    if verify_signature(signature, timestamp_str, &body, &state.pairing_token).is_err() {
        return Err(StatusCode::FORBIDDEN);
    }

    let proof = sign_server_proof(timestamp_str, nonce, state.bound_port, &state.pairing_token)
        .map_err(|_| StatusCode::FORBIDDEN)?;

    let mut response = Response::new(Body::empty());
    response.headers_mut().insert(
        SERVER_PROOF_HEADER,
        HeaderValue::from_str(&proof).map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?,
    );
    response.headers_mut().insert(
        SERVER_PORT_HEADER,
        HeaderValue::from_str(&state.bound_port.to_string())
            .map_err(|_| StatusCode::INTERNAL_SERVER_ERROR)?,
    );
    Ok(response)
}

async fn download_handler(
    State(state): State<ServerState>,
    headers: HeaderMap,
    body: Bytes,
) -> Result<StatusCode, StatusCode> {
    let signature = match headers
        .get("x-firelink-signature")
        .and_then(|v| v.to_str().ok())
    {
        Some(v) => v,
        None => return Err(StatusCode::FORBIDDEN),
    };

    let timestamp_str = match headers
        .get("x-firelink-timestamp")
        .and_then(|v| v.to_str().ok())
    {
        Some(v) => v,
        None => return Err(StatusCode::FORBIDDEN),
    };

    let timestamp = match verify_signature(signature, timestamp_str, &body, &state.pairing_token) {
        Ok(v) => v,
        Err(_) => return Err(StatusCode::FORBIDDEN),
    };

    if !claim_request(signature, timestamp, &state.replay_cache) {
        return Err(StatusCode::FORBIDDEN);
    }

    let payload: ExtensionRequest = match serde_json::from_slice(&body) {
        Ok(v) => v,
        Err(_) => return Err(StatusCode::BAD_REQUEST),
    };

    let download = match normalize_download(payload) {
        Some(v) => v,
        None => return Err(StatusCode::BAD_REQUEST),
    };

    if let Some(window) = state.app_handle.get_webview_window("main") {
        let is_visible = window.is_visible().unwrap_or(true);
        if !is_visible {
            let _ = window.show();
            let _ = window.set_focus();
            // Sleep briefly to let the webview wake up from macOS App Nap
            // otherwise the IPC event emitted immediately after is dropped.
            tokio::time::sleep(std::time::Duration::from_millis(300)).await;
        } else {
            let _ = window.show();
            let _ = window.set_focus();
        }
    }

    if !wait_for_frontend(&state.frontend_ready).await {
        return Err(StatusCode::SERVICE_UNAVAILABLE);
    }

    if state
        .app_handle
        .emit("extension-add-download", download)
        .is_err()
    {
        return Err(StatusCode::INTERNAL_SERVER_ERROR);
    }

    Ok(StatusCode::OK)
}

async fn wait_for_frontend(frontend_ready: &SharedFrontendReady) -> bool {
    for _ in 0..40 {
        if frontend_ready.load(Ordering::Acquire) {
            return true;
        }
        tokio::time::sleep(std::time::Duration::from_millis(50)).await;
    }
    false
}

fn normalize_download(payload: ExtensionRequest) -> Option<ExtensionDownload> {
    let mut seen = HashSet::new();
    let urls = payload
        .urls
        .into_iter()
        .take(MAX_URL_COUNT)
        .filter_map(|raw_url| normalize_url(&raw_url))
        .filter(|url| seen.insert(url.clone()))
        .collect::<Vec<_>>();
    if urls.is_empty() {
        return None;
    }

    let referer = payload.referer.and_then(|value| {
        let url = Url::parse(value.trim()).ok()?;
        matches!(url.scheme(), "http" | "https").then(|| url.to_string())
    });
    let filename = payload.filename.and_then(|value| sanitize_filename(&value));
    let headers = normalize_headers(payload.headers, payload.media);

    Some(ExtensionDownload {
        urls,
        referer,
        silent: payload.silent,
        filename,
        headers,
        // Explicit media is resolved by yt-dlp, which must use Firelink's
        // configured browser-cookie source. Forwarding a browser's complete
        // Cookie header can exceed upstream limits and makes old extension
        // builds pay for a doomed metadata request before retrying. Ordinary
        // captured downloads still need their exact request cookies.
        cookies: (!payload.media)
            .then_some(payload.cookies)
            .flatten()
            .filter(|value| !value.trim().is_empty()),
        media: payload.media,
    })
}

fn normalize_headers(headers: Option<String>, media: bool) -> Option<String> {
    let headers = headers?;
    if !media {
        return (!headers.trim().is_empty()).then_some(headers);
    }

    let filtered = headers
        .lines()
        .filter(|line| {
            line.split_once(':')
                .map(|(name, _)| !name.trim().eq_ignore_ascii_case("cookie"))
                .unwrap_or(true)
        })
        .collect::<Vec<_>>()
        .join("\n");
    (!filtered.trim().is_empty()).then_some(filtered)
}

fn normalize_url(raw_url: &str) -> Option<String> {
    let url = Url::parse(raw_url.trim()).ok()?;
    matches!(url.scheme(), "http" | "https" | "ftp" | "sftp").then(|| url.to_string())
}

fn sanitize_filename(filename: &str) -> Option<String> {
    let normalized = filename.trim().replace('\\', "/");
    let basename = Path::new(&normalized).file_name()?.to_str()?.trim();
    if basename.is_empty() || basename == "." || basename == ".." || basename.len() > 255 {
        return None;
    }
    Some(basename.to_string())
}

fn verify_signature(
    signature_hex: &str,
    timestamp_text: &str,
    body: &[u8],
    pairing_token: &SharedExtensionToken,
) -> Result<u64, ()> {
    let signature = decode_hex(signature_hex)?;
    let timestamp = timestamp_text.parse::<u64>().map_err(|_| ())?;
    let now = current_time_millis().ok_or(())?;
    if now.abs_diff(timestamp) >= SIGNATURE_MAX_AGE_MS {
        return Err(());
    }

    let token = pairing_token.read().unwrap_or_else(|e| e.into_inner());
    if token.is_empty() {
        return Err(());
    }

    let mut mac = HmacSha256::new_from_slice(token.as_bytes()).map_err(|_| ())?;
    mac.update(timestamp_text.as_bytes());
    mac.update(body);
    mac.verify_slice(&signature).map_err(|_| ())?;
    Ok(timestamp)
}

fn is_valid_client_nonce(value: &str) -> bool {
    value.len() == 32 && value.bytes().all(|byte| byte.is_ascii_hexdigit())
}

fn sign_server_proof(
    timestamp_text: &str,
    nonce: &str,
    bound_port: u16,
    pairing_token: &SharedExtensionToken,
) -> Result<String, ()> {
    let token = pairing_token.read().unwrap_or_else(|e| e.into_inner());
    if token.is_empty() {
        return Err(());
    }

    let mut mac = HmacSha256::new_from_slice(token.as_bytes()).map_err(|_| ())?;
    mac.update(SERVER_PROOF_PREFIX);
    mac.update(timestamp_text.as_bytes());
    mac.update(b"\n");
    mac.update(nonce.as_bytes());
    mac.update(b"\n");
    mac.update(bound_port.to_string().as_bytes());
    let signature = mac.finalize().into_bytes();
    Ok(encode_hex(signature.as_slice()))
}

fn encode_hex(bytes: &[u8]) -> String {
    bytes
        .iter()
        .map(|byte| format!("{byte:02x}"))
        .collect::<String>()
}

fn claim_request(signature: &str, timestamp: u64, replay_cache: &ReplayCache) -> bool {
    let now = match current_time_millis() {
        Some(now) => now,
        None => return false,
    };
    let mut cache = match replay_cache.lock() {
        Ok(cache) => cache,
        Err(_) => return false,
    };
    cache.retain(|_, seen_at| now.saturating_sub(*seen_at) < SIGNATURE_MAX_AGE_MS);
    if cache.len() > 10_000 {
        return false;
    }
    let key = format!("{timestamp}:{}", signature.to_ascii_lowercase());
    cache.insert(key, now).is_none()
}

fn current_time_millis() -> Option<u64> {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .ok()
        .and_then(|duration| u64::try_from(duration.as_millis()).ok())
}

fn decode_hex(value: &str) -> Result<Vec<u8>, ()> {
    if value.len() != 64 || !value.is_ascii() {
        return Err(());
    }
    value
        .as_bytes()
        .chunks_exact(2)
        .map(|pair| {
            let high = hex_digit(pair[0]).ok_or(())?;
            let low = hex_digit(pair[1]).ok_or(())?;
            Ok((high << 4) | low)
        })
        .collect()
}

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

fn is_allowed_origin(origin: &str) -> bool {
    Url::parse(origin)
        .ok()
        .is_some_and(|url| matches!(url.scheme(), "moz-extension" | "chrome-extension"))
}

#[cfg(test)]
mod tests {
    use super::{
        add_server_identity, is_valid_client_nonce, normalize_download, sign_server_proof,
        ExtensionRequest, PROTOCOL_VERSION_HEADER, SERVER_HEADER,
    };
    use axum::{http::StatusCode, middleware, routing::get, Router};
    use hmac::{Hmac, KeyInit, Mac};
    use sha2::Sha256;
    use std::sync::{Arc, RwLock};

    #[tokio::test]
    async fn identifies_every_extension_server_response() {
        let app = Router::new()
            .route("/ping", get(|| async { StatusCode::FORBIDDEN }))
            .layer(middleware::from_fn(add_server_identity));
        let listener = tokio::net::TcpListener::bind(("127.0.0.1", 0))
            .await
            .unwrap();
        let address = listener.local_addr().unwrap();
        let server = tokio::spawn(async move {
            axum::serve(listener, app).await.unwrap();
        });

        crate::ensure_reqwest_crypto_provider();
        let response = reqwest::get(format!("http://{address}/ping"))
            .await
            .unwrap();
        assert_eq!(response.status(), StatusCode::FORBIDDEN);
        assert_eq!(response.headers().get(SERVER_HEADER).unwrap(), "1");
        assert_eq!(
            response.headers().get(PROTOCOL_VERSION_HEADER).unwrap(),
            "4"
        );

        server.abort();
    }

    #[test]
    fn validates_client_nonce_shape() {
        assert!(is_valid_client_nonce("0123456789abcdef0123456789abcdef"));
        assert!(is_valid_client_nonce("ABCDEF0123456789abcdef0123456789"));
        assert!(!is_valid_client_nonce("0123456789abcdef0123456789abcde"));
        assert!(!is_valid_client_nonce("0123456789abcdef0123456789abcdeg"));
    }

    #[test]
    fn explicit_media_drops_the_extension_cookie_header() {
        let download = normalize_download(ExtensionRequest {
            urls: vec!["https://www.youtube.com/watch?v=example".to_string()],
            referer: None,
            silent: false,
            filename: None,
            headers: Some(format!(
                "Cookie: stale={};\nUser-Agent: Firefox",
                "x".repeat(64 * 1024)
            )),
            cookies: Some(format!("large={}", "x".repeat(64 * 1024))),
            media: true,
        })
        .expect("valid media handoff");

        assert!(download.media);
        assert!(download.cookies.is_none());
        assert_eq!(download.headers.as_deref(), Some("User-Agent: Firefox"));
    }

    #[test]
    fn regular_capture_preserves_the_extension_cookie_header() {
        let download = normalize_download(ExtensionRequest {
            urls: vec!["https://example.com/private.zip".to_string()],
            referer: None,
            silent: true,
            filename: None,
            headers: None,
            cookies: Some("session=browser-cookie-header".to_string()),
            media: false,
        })
        .expect("valid download handoff");

        assert!(!download.media);
        assert_eq!(
            download.cookies.as_deref(),
            Some("session=browser-cookie-header")
        );
    }

    #[test]
    fn signs_server_proof_with_timestamp_nonce_and_bound_port() {
        let token = Arc::new(RwLock::new("pairing-token".to_string()));
        let timestamp = "1710000000000";
        let nonce = "0123456789abcdef0123456789abcdef";
        let port = 6414;

        let mut mac = Hmac::<Sha256>::new_from_slice(b"pairing-token").unwrap();
        mac.update(b"firelink-server-proof\n");
        mac.update(timestamp.as_bytes());
        mac.update(b"\n");
        mac.update(nonce.as_bytes());
        mac.update(b"\n");
        mac.update(port.to_string().as_bytes());
        let expected = mac
            .finalize()
            .into_bytes()
            .iter()
            .map(|byte| format!("{byte:02x}"))
            .collect::<String>();

        assert_eq!(
            sign_server_proof(timestamp, nonce, port, &token).unwrap(),
            expected
        );
        assert_ne!(
            sign_server_proof(timestamp, nonce, port + 1, &token).unwrap(),
            expected
        );
    }
}