Skip to main content

trustee_api/
lib.rs

1//! Trustee API — REST + WebSocket server for the Trustee agent.
2//!
3//! Wraps a [`trustee_core::session::Session`] and exposes it over HTTP.
4//! Static frontend files are served from [`trustee_web`].
5//!
6//! Authentication is optional. When `[oidc]` or `[dev]` sections are present
7//! in the config TOML, all `/api/v1/*` endpoints require a valid JWT or dev
8//! token. Otherwise, all endpoints are open.
9
10pub mod auth;
11pub mod tls;
12mod routes;
13mod state;
14
15use std::net::SocketAddr;
16use std::sync::Arc;
17
18use anyhow::Result;
19use axum::routing::{get, post};
20use tower_http::cors::CorsLayer;
21
22pub use auth::{AuthConfig, AuthState};
23pub use state::ServerState;
24
25/// Run the API server.
26///
27/// Creates a `Session` with the given config, starts a background task to
28/// drain workflow messages and broadcast them to WebSocket clients, then
29/// serves the REST + WebSocket + static files on `addr`.
30///
31/// If `[oidc]` or `[dev]` sections are found in the config TOML, auth is
32/// enabled — all `/api/v1/*` endpoints (except health) require a valid token.
33///
34/// By default serves over HTTPS using a self-signed certificate from
35/// `~/.trustee/certs/`. If `use_tls` is false, serves plain HTTP.
36pub async fn run(
37    config_toml: String,
38    secrets: std::collections::HashMap<String, String>,
39    build_info: trustee_core::types::BuildInfo,
40    addr: SocketAddr,
41    use_tls: bool,
42) -> Result<()> {
43    // Parse auth config from TOML (returns None if no [oidc] or [dev] sections)
44    let auth_state = AuthConfig::from_toml(&config_toml).map(|cfg| {
45        let is_dev = cfg.dev_config.local_dev_mode;
46        tracing::info!(
47            "Auth enabled: {} mode, issuer={}",
48            if is_dev { "development" } else { "production" },
49            cfg.issuer_url
50        );
51        Arc::new(AuthState::new(cfg))
52    });
53
54    // Build the session
55    let (mut session, workflow_rx) = trustee_core::session::Session::new();
56    session.config_toml = Some(config_toml);
57    session.secrets = Some(secrets);
58    session.build_info = Some(build_info);
59    session.parse_auto_handoff_config();
60
61    // Create the broadcast channel for WebSocket fan-out
62    let (ws_tx, _ws_rx) = tokio::sync::broadcast::channel::<String>(256);
63
64    // Wrap session in shared state
65    let state = ServerState::new(session, ws_tx, auth_state);
66
67    // Start background message drain task (owns workflow_rx directly — no deadlock)
68    state.clone().spawn_drain_task(workflow_rx);
69
70    // Build router
71    //
72    // Auth middleware approach: since axum 0.8's from_fn_with_state has
73    // trait bound issues with nested routers, we apply auth checking at
74    // the handler level via a helper. Each protected route's handler
75    // calls auth::check_auth() first. This is simpler and avoids type
76    // complexity.
77    let app = axum::Router::new()
78        // Public routes
79        .route("/api/v1/health", get(routes::health))
80        .nest("/auth", auth::auth_routes())
81        // Protected API routes
82        .route("/api/v1/session", get(routes::get_session))
83        .route("/api/v1/session/command", post(routes::post_command))
84        .route("/api/v1/session/cancel", post(routes::post_cancel))
85        .route("/api/v1/session/handoff", post(routes::post_handoff))
86        .route("/api/v1/session/stream", get(routes::ws_handler))
87        // Static files from trustee-web
88        .route("/", get(routes::serve_index))
89        .route("/{file}", get(routes::serve_static))
90        .layer(CorsLayer::permissive())
91        .layer(axum::extract::DefaultBodyLimit::max(10 * 1024 * 1024))
92        .with_state(state);
93
94    // Start server
95    let listener = tokio::net::TcpListener::bind(addr).await?;
96
97    if use_tls {
98        // Install ring as the process-level crypto provider (required when
99        // rustls is built with default-features=false to avoid ambiguity
100        // with aws-lc-rs pulled in transitively by other crates).
101        let _ = rustls::crypto::ring::default_provider().install_default();
102
103        // Ensure self-signed certs exist
104        let cert_dir = tls::default_cert_dir();
105        let (cert_path, key_path) = tls::ensure_certs(&cert_dir)?;
106
107        // Load TLS config
108        let tls_config = tls::load_tls_config(&cert_path, &key_path)?;
109        let acceptor = tokio_rustls::TlsAcceptor::from(std::sync::Arc::new(tls_config));
110
111        tracing::info!("Trustee API listening on https://{}", addr);
112
113        // Manual accept loop — spawn hyper-util auto connection per TLS stream
114        loop {
115            let (tcp_stream, peer_addr) = match listener.accept().await {
116                Ok(stream) => stream,
117                Err(e) => {
118                    tracing::warn!("TCP accept failed: {}", e);
119                    continue;
120                }
121            };
122
123            let acceptor = acceptor.clone();
124            let app = app.clone();
125
126            tokio::spawn(async move {
127                let tls_stream = match acceptor.accept(tcp_stream).await {
128                    Ok(s) => s,
129                    Err(e) => {
130                        tracing::warn!("TLS accept failed from {}: {}", peer_addr, e);
131                        return;
132                    }
133                };
134
135                // Use hyper-util auto builder with the tower service from axum.
136                // serve_connection_with_upgrades is required for WebSocket support.
137                let io = hyper_util::rt::TokioIo::new(tls_stream);
138                let svc = hyper_util::service::TowerToHyperService::new(app);
139
140                let _ = hyper_util::server::conn::auto::Builder::new(hyper_util::rt::TokioExecutor::new())
141                    .serve_connection_with_upgrades(io, svc)
142                    .await;
143            });
144        }
145    } else {
146        tracing::info!("Trustee API listening on http://{}", addr);
147        axum::serve(listener, app).await?;
148    }
149
150    Ok(())
151}