init

src/http/mod.rs

@@ -0,0 +1,185 @@
+use crate::{
+    auth::{client_ip::ClientIpPolicy, AuthorizationService},
+    config::Config,
+    logs::{follower::LiveMessage, index::HistoryStore},
+};
+use axum::{
+    extract::{
+        ws::{Message, WebSocket, WebSocketUpgrade},
+        ConnectInfo, Query, State,
+    },
+    http::{header, HeaderMap, StatusCode},
+    response::{IntoResponse, Response},
+    routing::get,
+    Json, Router,
+};
+use serde::Deserialize;
+use std::{
+    net::{IpAddr, SocketAddr},
+    sync::{
+        atomic::{AtomicUsize, Ordering},
+        Arc,
+    },
+};
+use tokio::sync::broadcast;
+use tower_http::{
+    catch_panic::CatchPanicLayer, limit::RequestBodyLimitLayer, set_header::SetResponseHeaderLayer,
+    trace::TraceLayer,
+};
+
+#[derive(Clone)]
+pub struct AppState {
+    pub config: Arc<Config>,
+    pub auth: AuthorizationService,
+    pub history: HistoryStore,
+    pub live_tx: broadcast::Sender<LiveMessage>,
+    pub ip_policy: ClientIpPolicy,
+    pub ws_count: Arc<AtomicUsize>,
+}
+
+pub fn router(state: AppState) -> Router {
+    Router::new().route("/api/status",get(status)).route("/api/logs/recent",get(recent)).route("/api/logs/history",get(history)).route("/api/live",get(websocket)).fallback(crate::frontend::serve).with_state(state).layer(RequestBodyLimitLayer::new(16*1024)).layer(SetResponseHeaderLayer::if_not_present(header::X_CONTENT_TYPE_OPTIONS,header::HeaderValue::from_static("nosniff"))).layer(SetResponseHeaderLayer::if_not_present(header::REFERRER_POLICY,header::HeaderValue::from_static("no-referrer"))).layer(SetResponseHeaderLayer::if_not_present(header::CONTENT_SECURITY_POLICY,header::HeaderValue::from_static("default-src 'self'; connect-src 'self' ws: wss:; style-src 'self'; script-src 'self'; frame-ancestors 'none'"))).layer(CatchPanicLayer::new()).layer(TraceLayer::new_for_http())
+}
+
+async fn status() -> impl IntoResponse {
+    Json(serde_json::json!({"status":"ok"}))
+}
+#[derive(Deserialize)]
+struct PageQuery {
+    before: Option<String>,
+    limit: Option<usize>,
+}
+async fn recent(
+    State(state): State<AppState>,
+    ConnectInfo(peer): ConnectInfo<SocketAddr>,
+    headers: HeaderMap,
+    Query(q): Query<PageQuery>,
+) -> Response {
+    logs_response(state, peer, headers, None, q.limit).await
+}
+async fn history(
+    State(state): State<AppState>,
+    ConnectInfo(peer): ConnectInfo<SocketAddr>,
+    headers: HeaderMap,
+    Query(q): Query<PageQuery>,
+) -> Response {
+    logs_response(state, peer, headers, q.before.as_deref(), q.limit).await
+}
+async fn logs_response(
+    state: AppState,
+    peer: SocketAddr,
+    headers: HeaderMap,
+    before: Option<&str>,
+    limit: Option<usize>,
+) -> Response {
+    if let Err(r) = authorize(&state, peer, &headers).await {
+        return r;
+    }
+    let limit = limit
+        .unwrap_or(state.config.initial_log_lines)
+        .min(state.config.max_history_lines)
+        .max(1);
+    match state.history.page(before, limit).await {
+        Ok(page) => Json(page).into_response(),
+        Err(code) if code == "invalid_cursor" => error(StatusCode::BAD_REQUEST, "invalid_cursor"),
+        Err(code) if code == "stale_cursor" => error(StatusCode::CONFLICT, "stale_cursor"),
+        Err(_) => error(StatusCode::UNPROCESSABLE_ENTITY, "history_unavailable"),
+    }
+}
+
+async fn websocket(
+    ws: WebSocketUpgrade,
+    State(state): State<AppState>,
+    ConnectInfo(peer): ConnectInfo<SocketAddr>,
+    headers: HeaderMap,
+) -> Response {
+    if let Err(r) = authorize(&state, peer, &headers).await {
+        return r;
+    }
+    if state.ws_count.load(Ordering::Relaxed) >= state.config.max_ws_connections {
+        return error(StatusCode::SERVICE_UNAVAILABLE, "connection_limit");
+    }
+    state.ws_count.fetch_add(1, Ordering::Relaxed);
+    ws.on_upgrade(move |socket| serve_socket(socket, state))
+        .into_response()
+}
+async fn serve_socket(mut socket: WebSocket, state: AppState) {
+    let _guard = WsGuard(state.ws_count.clone());
+    let mut rx = state.live_tx.subscribe();
+    if let Ok(hello) = serde_json::to_string(&LiveMessage::Hello) {
+        if socket.send(Message::Text(hello.into())).await.is_err() {
+            return;
+        }
+    }
+    loop {
+        match rx.recv().await {
+            Ok(event) => {
+                if let Ok(json) = serde_json::to_string(&event) {
+                    if socket.send(Message::Text(json.into())).await.is_err() {
+                        break;
+                    }
+                }
+            }
+            Err(broadcast::error::RecvError::Lagged(_)) => {
+                let _ = socket
+                    .send(Message::Text(r#"{"type":"resync_required"}"#.into()))
+                    .await;
+                break;
+            }
+            Err(_) => break,
+        }
+    }
+}
+struct WsGuard(Arc<AtomicUsize>);
+impl Drop for WsGuard {
+    fn drop(&mut self) {
+        self.0.fetch_sub(1, Ordering::Relaxed);
+    }
+}
+
+async fn authorize(
+    state: &AppState,
+    peer: SocketAddr,
+    headers: &HeaderMap,
+) -> Result<IpAddr, Response> {
+    validate_origin(&state.config.public_origin, headers)
+        .map_err(|_| error(StatusCode::FORBIDDEN, "origin_denied"))?;
+    let ip = state
+        .ip_policy
+        .extract(peer.ip(), headers)
+        .map_err(|_| error(StatusCode::BAD_REQUEST, "invalid_client_ip"))?;
+    if state.auth.is_allowed(ip).await {
+        Ok(ip)
+    } else {
+        Err(error(StatusCode::FORBIDDEN, "access_denied"))
+    }
+}
+fn validate_origin(expected: &str, headers: &HeaderMap) -> Result<(), ()> {
+    if let Some(origin) = headers.get(header::ORIGIN) {
+        if origin.to_str().ok() != Some(expected) {
+            return Err(());
+        }
+    }
+    Ok(())
+}
+fn error(status: StatusCode, code: &str) -> Response {
+    (status, Json(serde_json::json!({"error":code}))).into_response()
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    #[test]
+    fn rejects_cross_origin() {
+        let mut h = HeaderMap::new();
+        h.insert(header::ORIGIN, "https://evil.example".parse().unwrap());
+        assert!(validate_origin("https://logs.example", &h).is_err());
+    }
+    #[test]
+    fn accepts_matching_or_non_browser_origin() {
+        let mut h = HeaderMap::new();
+        assert!(validate_origin("https://logs.example", &h).is_ok());
+        h.insert(header::ORIGIN, "https://logs.example".parse().unwrap());
+        assert!(validate_origin("https://logs.example", &h).is_ok());
+    }
+}