cube/apps/llm-proxy/src/main.rs

//! llm.famzheng.me — gemma-4-31b-it 反向代理 + 简单 token 鉴权。
//!
//! - `GET  /`                       → `/chat` 跳转
//! - `GET  /chat`                   → 静态 web UI
//! - `POST /v1/chat/completions`    → OpenAI 兼容透传 (要 Authorization: token <PROXY_AUTH_TOKEN>)
//! - `GET  /healthz`                → 不带 auth, 给 k8s probe

mod proxy;

use std::sync::Arc;

use axum::{
    extract::State,
    http::{header, StatusCode},
    middleware::{self, Next},
    response::{Html, IntoResponse, Redirect, Response},
    routing::{get, post},
    Router,
};
use tower_http::trace::TraceLayer;

#[tokio::main]
async fn main() -> std::io::Result<()> {
    cube_core::init_tracing();
    let cfg = Arc::new(proxy::Config::from_env());
    let port: u16 = std::env::var("PORT")
        .ok()
        .and_then(|s| s.parse().ok())
        .unwrap_or(8080);

    let chat_api = Router::new()
        .route("/v1/chat/completions", post(proxy::handle))
        .route_layer(middleware::from_fn_with_state(cfg.clone(), require_token))
        .with_state(cfg);

    let app = Router::new()
        .route("/healthz", get(|| async { "ok" }))
        .route("/", get(|| async { Redirect::permanent("/chat") }))
        .route("/chat", get(chat_ui))
        .route("/favicon.svg", get(favicon))
        .route("/favicon.ico", get(favicon)) // 浏览器默认会请求 .ico，让它共享同一 SVG
        .merge(chat_api)
        .layer(TraceLayer::new_for_http());

    let addr = format!("0.0.0.0:{port}");
    let listener = tokio::net::TcpListener::bind(&addr).await?;
    tracing::info!(%addr, "llm-proxy listening");
    axum::serve(listener, app).await
}

const CHAT_HTML: &str = include_str!("../web/chat.html");
const FAVICON_SVG: &str = include_str!("../web/favicon.svg");

async fn chat_ui() -> Html<&'static str> {
    Html(CHAT_HTML)
}

async fn favicon() -> impl IntoResponse {
    (
        [
            (axum::http::header::CONTENT_TYPE, "image/svg+xml"),
            (axum::http::header::CACHE_CONTROL, "public, max-age=604800"),
        ],
        FAVICON_SVG,
    )
}

/// 验 `Authorization: token <PROXY_AUTH_TOKEN>`，错的直接 401。
async fn require_token(
    State(cfg): State<Arc<proxy::Config>>,
    req: axum::extract::Request,
    next: Next,
) -> Response {
    let header_val = req
        .headers()
        .get(header::AUTHORIZATION)
        .and_then(|v| v.to_str().ok())
        .map(str::trim);

    match header_val {
        Some(v) if check_token(v, &cfg.proxy_auth_token) => next.run(req).await,
        _ => (
            StatusCode::UNAUTHORIZED,
            "缺少或不匹配 `Authorization: token <your-token>`",
        )
            .into_response(),
    }
}

/// 接受 `token <T>` 或 `Bearer <T>`（OpenAI client 习惯发 Bearer，宽容点）。
pub fn check_token(header_value: &str, expected: &str) -> bool {
    if expected.is_empty() {
        return false;
    }
    let trimmed = header_value.trim();
    if let Some(rest) = trimmed.strip_prefix("token ") {
        return constant_time_eq(rest.trim().as_bytes(), expected.as_bytes());
    }
    if let Some(rest) = trimmed.strip_prefix("Bearer ") {
        return constant_time_eq(rest.trim().as_bytes(), expected.as_bytes());
    }
    false
}

/// 常时间比较，防 timing attack（虽然这场景影响小，做了不亏）。
fn constant_time_eq(a: &[u8], b: &[u8]) -> bool {
    if a.len() != b.len() {
        return false;
    }
    let mut diff: u8 = 0;
    for (x, y) in a.iter().zip(b.iter()) {
        diff |= x ^ y;
    }
    diff == 0
}

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

    #[test]
    fn check_token_accepts_token_scheme() {
        assert!(check_token("token famzheng-llm-2026", "famzheng-llm-2026"));
    }

    #[test]
    fn check_token_accepts_bearer_scheme() {
        assert!(check_token("Bearer famzheng-llm-2026", "famzheng-llm-2026"));
    }

    #[test]
    fn check_token_rejects_wrong_value() {
        assert!(!check_token("token wrong", "famzheng-llm-2026"));
    }

    #[test]
    fn check_token_rejects_unknown_scheme() {
        assert!(!check_token("Basic famzheng-llm-2026", "famzheng-llm-2026"));
        assert!(!check_token("famzheng-llm-2026", "famzheng-llm-2026"));
    }

    #[test]
    fn check_token_rejects_empty_expected() {
        // 防 misconfigured：空 expected 不应该让任何人通过
        assert!(!check_token("token any", ""));
        assert!(!check_token("Bearer ", ""));
    }

    #[test]
    fn check_token_strips_extra_whitespace() {
        assert!(check_token("  token   famzheng-llm-2026  ", "famzheng-llm-2026"));
    }

    #[test]
    fn check_token_rejects_prefix_match() {
        // 防止"famzheng-llm-2026-extra" 通过
        assert!(!check_token("token famzheng-llm-2026-extra", "famzheng-llm-2026"));
        assert!(!check_token("token famzheng-llm", "famzheng-llm-2026"));
    }
}