Skip to main content

shell_tunnel/api/
router.rs

1//! API router configuration.
2
3use std::net::SocketAddr;
4use std::sync::Arc;
5
6use axum::{
7    extract::{
8        connect_info::IntoMakeServiceWithConnectInfo, DefaultBodyLimit, MatchedPath, Request, State,
9    },
10    http::{header::AUTHORIZATION, Method, StatusCode},
11    middleware::{self, Next},
12    response::Response,
13    routing::{any, get, post},
14    Router,
15};
16use tower_http::{
17    cors::{Any, CorsLayer},
18    trace::TraceLayer,
19};
20
21use super::handlers::{
22    api_info, create_session, delete_session, execute_command, execute_oneshot, get_session,
23    health, list_sessions, AppState,
24};
25use super::websocket::{ws_handler, ws_oneshot_handler};
26use crate::security::{
27    rate_limit_middleware, ApiKeyStore, AuthConfig, CapabilitySet, RateLimitConfig, RateLimiter,
28};
29
30/// Cross-Origin Resource Sharing (CORS) configuration.
31///
32/// CORS is a browser-enforced mechanism; non-browser clients (`curl`, SDKs)
33/// ignore it entirely. shell-tunnel therefore emits **no** permissive CORS headers
34/// by default: this blocks a malicious web page from reading responses cross-origin
35/// and, because the JSON execute endpoints require a preflight, from issuing the
36/// request at all — with zero impact on the intended non-browser consumers.
37///
38/// Note: CORS does **not** defend against DNS-rebinding (the attacker's host is
39/// rebound to `127.0.0.1`, making the request same-origin); that requires
40/// Host-header validation and is tracked separately.
41#[derive(Debug, Clone, Default)]
42pub struct CorsConfig {
43    /// Allow any origin/method/header (restores the permissive `Any` behavior).
44    /// Off by default; enable only for trusted browser-based UIs.
45    pub allow_any: bool,
46}
47
48/// Security configuration for the server.
49#[derive(Debug, Clone)]
50pub struct SecurityConfig {
51    /// Authentication configuration.
52    pub auth: AuthConfig,
53    /// Rate limiting configuration.
54    pub rate_limit: RateLimitConfig,
55    /// API keys to pre-register.
56    pub api_keys: Vec<String>,
57    /// Capabilities granted to the pre-registered keys and to the
58    /// auto-generated fallback key.
59    ///
60    /// `None` (the default) means **full-control** (wildcard) — the backward-
61    /// compatible behavior for a bare `--api-key` / `--require-auth` (spec §4).
62    /// `Some(set)` issues fine-grained tokens scoped to that set (spec §9).
63    pub capabilities: Option<CapabilitySet>,
64    /// CORS configuration.
65    pub cors: CorsConfig,
66    /// Host names this server answers to, when it is worth checking.
67    ///
68    /// `None` disables the check. It is meant for a loopback-bound server, where
69    /// DNS rebinding is the one attack CORS cannot stop: the attacker's name is
70    /// rebound to `127.0.0.1`, so the browser considers the request same-origin
71    /// and sends it. The `Host` header still carries the attacker's name, which
72    /// is what this compares. A published server is deliberately reachable under
73    /// a name we may not know, so the check does not apply there.
74    pub allowed_hosts: Option<Vec<String>>,
75}
76
77impl Default for SecurityConfig {
78    fn default() -> Self {
79        Self {
80            auth: AuthConfig::disabled(), // Disabled by default for ease of use
81            rate_limit: RateLimitConfig::default(),
82            api_keys: Vec::new(),
83            capabilities: None, // Full-control by default (legacy-compatible)
84            cors: CorsConfig::default(), // Restrictive by default (no permissive CORS)
85            allowed_hosts: None,
86        }
87    }
88}
89
90/// Build a permissive CORS layer when `allow_any` is set, otherwise `None`.
91///
92/// Returning `None` means no CORS headers are emitted — the secure default.
93fn cors_layer(cfg: &CorsConfig) -> Option<CorsLayer> {
94    cfg.allow_any.then(|| {
95        CorsLayer::new()
96            .allow_origin(Any)
97            .allow_methods(Any)
98            .allow_headers(Any)
99    })
100}
101
102impl SecurityConfig {
103    /// Create a secure configuration.
104    pub fn secure() -> Self {
105        Self {
106            auth: AuthConfig::default(),
107            rate_limit: RateLimitConfig::default(),
108            api_keys: Vec::new(),
109            capabilities: None,
110            cors: CorsConfig::default(),
111            allowed_hosts: None,
112        }
113    }
114
115    /// Create a development configuration (no auth, relaxed limits).
116    pub fn development() -> Self {
117        Self {
118            auth: AuthConfig::disabled(),
119            rate_limit: RateLimitConfig::relaxed(),
120            api_keys: Vec::new(),
121            capabilities: None,
122            cors: CorsConfig::default(),
123            allowed_hosts: None,
124        }
125    }
126
127    /// Add an API key.
128    pub fn with_api_key(mut self, key: impl Into<String>) -> Self {
129        self.api_keys.push(key.into());
130        self
131    }
132
133    /// Scope the issued tokens to a fine-grained capability set (spec §9).
134    ///
135    /// Applies to the pre-registered keys and to the auto-generated fallback
136    /// key. Without this, tokens are full-control (legacy-compatible).
137    pub fn with_capabilities(mut self, capabilities: CapabilitySet) -> Self {
138        self.capabilities = Some(capabilities);
139        self
140    }
141
142    /// Answer only to these host names.
143    pub fn with_allowed_hosts(mut self, hosts: Vec<String>) -> Self {
144        self.allowed_hosts = Some(hosts);
145        self
146    }
147
148    /// Enable permissive (`Any`) CORS. Opt-in; only for trusted browser UIs.
149    pub fn with_cors_allow_any(mut self) -> Self {
150        self.cors.allow_any = true;
151        self
152    }
153}
154
155/// Register `key` into `store` with `capabilities` (fine-grained), or as a
156/// legacy full-control key when `capabilities` is `None` (spec §4/§9).
157fn register_key(store: &ApiKeyStore, key: &str, capabilities: &Option<CapabilitySet>) {
158    match capabilities {
159        Some(caps) => store.add_key_with_capabilities(key, caps.clone(), "configured"),
160        None => store.add_key(key),
161    }
162}
163
164/// Create the API router with all routes configured.
165pub fn create_router() -> Router {
166    create_router_with_state(AppState::new())
167}
168
169/// Create the API router with custom state (no security).
170pub fn create_router_with_state(state: AppState) -> Router {
171    // Session routes
172    let session_routes = Router::new()
173        .route("/", get(list_sessions).post(create_session))
174        .route("/{id}", get(get_session).delete(delete_session))
175        .route("/{id}/execute", post(execute_command))
176        .route("/{id}/ws", any(ws_handler));
177
178    // API v1 routes
179    let api_v1 = Router::new()
180        .route("/", get(api_info))
181        .route("/execute", post(execute_oneshot))
182        .route("/ws", any(ws_oneshot_handler))
183        .nest("/fs", fs_routes())
184        .nest("/sessions", session_routes);
185
186    // Build main router. This "no security" convenience constructor uses the
187    // restrictive CORS default (no permissive CORS headers emitted).
188    Router::new()
189        .route("/health", get(health))
190        .nest("/api/v1", api_v1)
191        .layer(TraceLayer::new_for_http())
192        .with_state(state)
193}
194
195/// The capability a route requires (Phase A spec §3).
196///
197/// Declared here at the router layer — co-located with the route definitions,
198/// which are the single source of truth for the matched-path strings this maps
199/// against. Not a per-handler attribute.
200#[derive(Debug, Clone, Copy, PartialEq, Eq)]
201pub enum RequiredCapability {
202    /// No authentication at all (e.g. `/health`).
203    Public,
204    /// Authenticated only: any valid token passes, no specific capability
205    /// required (spec §3 "인증만" tier, e.g. `GET /api/v1`).
206    Authenticated,
207    /// Requires a specific capability string (set membership, or wildcard).
208    Capability(&'static str),
209}
210
211/// Map a matched route (`method` + axum [`MatchedPath`]) to its required
212/// capability (spec §3 table).
213///
214/// Keyed on the **full nested** `MatchedPath` (confirmed against the real router
215/// structure) plus the HTTP method, because one path can require different
216/// capabilities per method (e.g. `GET /api/v1/sessions` = read, `POST` = manage).
217///
218/// Unknown routes fail **closed** to [`RequiredCapability::Authenticated`]: an
219/// unmapped route still requires a valid token, never less than that.
220pub fn required_capability(method: &Method, matched_path: &str) -> RequiredCapability {
221    use RequiredCapability::{Authenticated, Capability, Public};
222
223    // HEAD is GET without a body, and axum's `get()` serves it automatically, so
224    // its authorization must equal GET's. Keyed separately below, every GET
225    // route would need a twin HEAD arm — and a forgotten twin falls through to
226    // the closed default `Authenticated`, which means any valid token, not the
227    // capability GET requires. Normalising here is the one place it cannot be
228    // forgotten. Matched as `&str` (rather than comparing `Method` values) so
229    // the table below stays exactly as it reads for every other method.
230    let method = match method.as_str() {
231        "HEAD" => "GET",
232        other => other,
233    };
234
235    match (method, matched_path) {
236        (_, "/health") => Public,
237        ("GET", "/api/v1") => Authenticated,
238        ("POST", "/api/v1/execute") => Capability("exec"),
239        (_, "/api/v1/ws") => Capability("exec"),
240        ("GET", "/api/v1/sessions") => Capability("session.read"),
241        ("POST", "/api/v1/sessions") => Capability("session.manage"),
242        ("GET", "/api/v1/sessions/{id}") => Capability("session.read"),
243        ("DELETE", "/api/v1/sessions/{id}") => Capability("session.manage"),
244        ("POST", "/api/v1/sessions/{id}/execute") => Capability("exec"),
245        (_, "/api/v1/sessions/{id}/ws") => Capability("exec"),
246        ("GET", "/api/v1/fs/list") => Capability("fs.read"),
247        ("GET", "/api/v1/fs/stat") => Capability("fs.read"),
248        ("GET", "/api/v1/fs/file") => Capability("fs.read"),
249        ("DELETE", "/api/v1/fs/file") => Capability("fs.write"),
250        ("POST", "/api/v1/fs/uploads") => Capability("fs.write"),
251        ("GET", "/api/v1/fs/uploads/{id}") => Capability("fs.write"),
252        ("PATCH", "/api/v1/fs/uploads/{id}") => Capability("fs.write"),
253        ("POST", "/api/v1/fs/uploads/{id}/complete") => Capability("fs.write"),
254        ("DELETE", "/api/v1/fs/uploads/{id}") => Capability("fs.write"),
255        _ => Authenticated,
256    }
257}
258
259/// The longest raw request path an audit entry will carry, in bytes.
260///
261/// Only an *unmatched* path is recorded raw, and an unmatched path is whatever
262/// the caller asked for — a probe can send four kilobytes of it and did, in the
263/// measurement that set this number. A matched route is a router template and
264/// is never truncated. The same log-flood worry already put the `tracing` line
265/// beside this one at `debug`; the trail cannot be silenced by a log level, so
266/// it needs the bound instead.
267const MAX_AUDITED_PATH: usize = 256;
268
269/// The `route` value for an audit entry.
270///
271/// A matched route is recorded as its template (`/api/v1/sessions/{id}`) so
272/// that entries group rather than exploding into one bucket per id. An
273/// unmatched path has no template, so the raw path is recorded: it is the only
274/// thing that says *what was probed*, and probing is precisely what the trail
275/// is asked about afterwards. Recording only the method left two different
276/// probes with byte-identical entries — five distinct requests produced five
277/// indistinguishable lines, confirmed against a running server, and a majority
278/// of the `denied` entries on an internet-facing deployment were of that shape.
279///
280/// The caller passes `uri().path()`, never `path_and_query()`, and that is a
281/// guarantee rather than a shortcut: a query string is caller-controlled too,
282/// and `USAGE.md` §4 promises the trail never carries a credential. A probe of
283/// `/nope?token=…` is recorded as `/nope` — confirmed by running it.
284///
285/// Truncation is marked with a suffix that begins with a space, which is not a
286/// byte a request path can contain: a space terminates the request target, so
287/// hyper answers `400` long before this function sees it. A caller therefore
288/// cannot forge the marker by asking for a path that ends in it. (Raw control
289/// bytes and invalid UTF-8 are refused at the same layer for the same kind of
290/// reason, so this function is not where they need handling — measured, not
291/// assumed, because "the parser surely rejects that" is the premise this
292/// repository has been wrong about before.)
293fn audited_route(method: &Method, matched: Option<&str>, raw_path: &str) -> String {
294    let Some(template) = matched else {
295        if raw_path.len() <= MAX_AUDITED_PATH {
296            return format!("{method} {raw_path}");
297        }
298        // Back off to a character boundary; index 0 is always one, so this
299        // terminates.
300        let mut end = MAX_AUDITED_PATH;
301        while !raw_path.is_char_boundary(end) {
302            end -= 1;
303        }
304        return format!("{method} {} (truncated)", &raw_path[..end]);
305    };
306    format!("{method} {template}")
307}
308
309/// Scope-aware authentication + authorization middleware (spec §5).
310///
311/// 1. Resolve the route's required capability from `method` + `MatchedPath`.
312/// 2. `Public` route or auth disabled → pass through.
313/// 3. Extract the bearer token; look up its `TokenRecord` in the store.
314///    Missing/invalid token → **401**.
315/// 4. `Authenticated` route → any valid token passes. `Capability(c)` route →
316///    the token's set must satisfy `c` (membership or wildcard), else **403**.
317async fn capability_auth_middleware(
318    State((store, audit)): State<(
319        std::sync::Arc<ApiKeyStore>,
320        std::sync::Arc<crate::audit::AuditSink>,
321    )>,
322    mut request: Request,
323    next: Next,
324) -> Result<Response, StatusCode> {
325    // Auth disabled → open server (existing behavior).
326    if !store.is_enabled() {
327        return Ok(next.run(request).await);
328    }
329
330    let method = request.method().clone();
331    // Owned so it can be used in the rejection logs after `request` is consumed.
332    //
333    // `None` and `Some("")` are not the same thing and are no longer flattened
334    // together: a path the router did not match has no template, and the raw
335    // path is the only description of it that exists. Authorization still sees
336    // the empty string for it (`required_capability` fails closed on it), but
337    // the audit entry does not.
338    let matched = request
339        .extensions()
340        .get::<MatchedPath>()
341        .map(|m| m.as_str().to_owned());
342    let route = audited_route(&method, matched.as_deref(), request.uri().path());
343    let required = required_capability(&method, matched.as_deref().unwrap_or_default());
344
345    // Public routes (e.g. /health) skip auth entirely.
346    if required == RequiredCapability::Public {
347        return Ok(next.run(request).await);
348    }
349
350    // Extract the bearer token and resolve its capabilities.
351    // Missing header, wrong prefix, or unregistered token → 401.
352    let token = request
353        .headers()
354        .get(AUTHORIZATION)
355        .and_then(|v| v.to_str().ok())
356        .and_then(|header| store.extract_key(header));
357
358    let identity = token.as_deref().and_then(|t| store.identity(t));
359
360    let capabilities = match token.as_deref().and_then(|t| store.capabilities(t)) {
361        Some(caps) => caps,
362        None => {
363            // Logged at debug to avoid a log-flood amplifier under probing; the
364            // token value itself is never logged. `missing-token` = no/malformed
365            // Authorization header, `invalid-token` = present but unregistered.
366            let reason = if token.is_none() {
367                "missing-token"
368            } else {
369                "invalid-token"
370            };
371            tracing::debug!(%method, path = %route, reason, "auth rejected (401)");
372            // Probing is exactly what an audit trail is asked about afterwards,
373            // so this layer's refusals are recorded as well as successes.
374            //
375            // "This layer's" is the whole claim, not modesty. Refusals that
376            // happen after this middleware has let a request through — an
377            // extractor turning away an unrecognised field, a malformed body, a
378            // path parameter that will not parse, a body over the size limit —
379            // record nothing, so the trail is not a complete list of every
380            // refusal the server issued. `USAGE.md` §4 names that gap and lists
381            // the measured cases; keep the two in step as the layer grows.
382            audit.record(
383                crate::audit::AuditEvent::new("denied")
384                    .with_route(route)
385                    .with_denial(401, reason),
386            );
387            return Err(StatusCode::UNAUTHORIZED);
388        }
389    };
390
391    // Handlers record what actually ran, and need to know who asked.
392    if let Some(identity) = identity.clone() {
393        request.extensions_mut().insert(identity);
394    }
395
396    match required {
397        // Already handled above, but keep the match exhaustive.
398        RequiredCapability::Public => Ok(next.run(request).await),
399        // Any valid token satisfies an authenticated-only route.
400        RequiredCapability::Authenticated => Ok(next.run(request).await),
401        // Specific capability: set membership (or wildcard) required, else 403.
402        RequiredCapability::Capability(cap) => {
403            if capabilities.satisfies(cap) {
404                Ok(next.run(request).await)
405            } else {
406                audit.record(
407                    crate::audit::AuditEvent::new("denied")
408                        .with_identity(identity)
409                        .with_route(route.clone())
410                        .with_denial(403, format!("missing-capability:{cap}")),
411                );
412                tracing::debug!(
413                    %method,
414                    path = %route,
415                    required = cap,
416                    "authorization denied (403): insufficient capability"
417                );
418                Err(StatusCode::FORBIDDEN)
419            }
420        }
421    }
422}
423
424/// Whether `header` names a host this server answers to.
425///
426/// Compared without the port, since the port is not what an attacker controls
427/// in a rebinding attack, and a legitimate caller may reach the same server
428/// through different ports.
429fn host_is_allowed(header: Option<&str>, allowed: &[String]) -> bool {
430    let Some(value) = header else {
431        // HTTP/1.1 requires a Host header; its absence is not a shape any
432        // ordinary client produces.
433        return false;
434    };
435
436    let host = value
437        .rsplit_once(':')
438        .map_or(value, |(host, port)| {
439            // Only strip a trailing port, not part of a bare IPv6 address.
440            if port.chars().all(|c| c.is_ascii_digit()) {
441                host
442            } else {
443                value
444            }
445        })
446        .trim_matches(|c| c == '[' || c == ']');
447
448    allowed
449        .iter()
450        .any(|candidate| candidate.eq_ignore_ascii_case(host))
451}
452
453/// Reject requests carrying a `Host` this server does not answer to.
454async fn host_check_middleware(
455    State(allowed): State<Arc<Vec<String>>>,
456    request: Request,
457    next: Next,
458) -> Result<Response, (StatusCode, String)> {
459    let header = request
460        .headers()
461        .get(axum::http::header::HOST)
462        .and_then(|value| value.to_str().ok());
463
464    if host_is_allowed(header, &allowed) {
465        return Ok(next.run(request).await);
466    }
467
468    // Named explicitly: an operator hitting this from a container or behind a
469    // proxy needs to know which name was refused and how to permit it.
470    let seen = header.unwrap_or("(none)").to_string();
471    tracing::debug!(host = %seen, "request refused: host not allowed");
472    Err((
473        StatusCode::FORBIDDEN,
474        format!(
475            "host {seen} is not allowed; pass --allow-host {seen} to permit it
476"
477        ),
478    ))
479}
480
481/// Create the API router with security enabled.
482pub fn create_secure_router(
483    state: AppState,
484    security: SecurityConfig,
485) -> (Router, Arc<ApiKeyStore>, Arc<RateLimiter>) {
486    // Create security components
487    let auth_store = Arc::new(ApiKeyStore::new(security.auth));
488    let rate_limiter = Arc::new(RateLimiter::new(security.rate_limit));
489
490    // Register API keys with their configured capabilities.
491    for key in &security.api_keys {
492        register_key(&auth_store, key, &security.capabilities);
493    }
494
495    // Session routes
496    let session_routes = Router::new()
497        .route("/", get(list_sessions).post(create_session))
498        .route("/{id}", get(get_session).delete(delete_session))
499        .route("/{id}/execute", post(execute_command))
500        .route("/{id}/ws", any(ws_handler));
501
502    // API v1 routes
503    let api_v1 = Router::new()
504        .route("/", get(api_info))
505        .route("/execute", post(execute_oneshot))
506        .route("/ws", any(ws_oneshot_handler))
507        .nest("/fs", fs_routes())
508        .nest("/sessions", session_routes);
509
510    let allowed_hosts = security.allowed_hosts.clone();
511
512    // Build main router with security layers
513    let mut router = Router::new()
514        .route("/health", get(health))
515        .nest("/api/v1", api_v1)
516        .layer(middleware::from_fn_with_state(
517            (Arc::clone(&auth_store), Arc::clone(&state.audit)),
518            capability_auth_middleware,
519        ))
520        .layer(middleware::from_fn_with_state(
521            Arc::clone(&rate_limiter),
522            rate_limit_middleware,
523        ))
524        .layer(TraceLayer::new_for_http());
525
526    // Outermost, so a rebound request is refused before it reaches the token
527    // store or the rate limiter's bookkeeping.
528    if let Some(hosts) = allowed_hosts {
529        router = router.layer(middleware::from_fn_with_state(
530            Arc::new(hosts),
531            host_check_middleware,
532        ));
533    }
534
535    // Permissive CORS only when explicitly opted in (default: restrictive).
536    if let Some(cors) = cors_layer(&security.cors) {
537        router = router.layer(cors);
538    }
539
540    let router = router.with_state(state);
541
542    (router, auth_store, rate_limiter)
543}
544
545/// Server configuration.
546#[derive(Debug, Clone)]
547pub struct ServerConfig {
548    /// Host address to bind to.
549    pub host: String,
550    /// Port to listen on.
551    pub port: u16,
552    /// Security configuration.
553    pub security: SecurityConfig,
554    /// Enable graceful shutdown on SIGTERM/SIGINT.
555    pub graceful_shutdown: bool,
556}
557
558impl ServerConfig {
559    pub fn new(host: impl Into<String>, port: u16) -> Self {
560        Self {
561            host: host.into(),
562            port,
563            security: SecurityConfig::default(),
564            graceful_shutdown: true,
565        }
566    }
567
568    pub fn bind_address(&self) -> String {
569        format!("{}:{}", self.host, self.port)
570    }
571
572    /// Enable security with the given configuration.
573    pub fn with_security(mut self, security: SecurityConfig) -> Self {
574        self.security = security;
575        self
576    }
577
578    /// Disable graceful shutdown.
579    pub fn without_graceful_shutdown(mut self) -> Self {
580        self.graceful_shutdown = false;
581        self
582    }
583}
584
585impl Default for ServerConfig {
586    fn default() -> Self {
587        Self {
588            host: "127.0.0.1".to_string(),
589            port: 3000,
590            security: SecurityConfig::default(),
591            graceful_shutdown: true,
592        }
593    }
594}
595
596/// Start the API server.
597pub async fn serve(config: ServerConfig) -> crate::Result<()> {
598    serve_with_state(config, AppState::new()).await
599}
600
601/// Bind the API server's port without starting to serve.
602///
603/// Callers that need to know the port before traffic flows — anything binding
604/// port 0, where the OS chooses — take the listener from here and hand it to
605/// [`serve_on`]. Splitting bind from serve is what makes an ephemeral port
606/// usable: the alternative is binding twice and racing whoever grabs it in
607/// between.
608pub async fn bind(config: &ServerConfig) -> crate::Result<tokio::net::TcpListener> {
609    tokio::net::TcpListener::bind(config.bind_address())
610        .await
611        .map_err(crate::error::ShellTunnelError::Io)
612}
613
614/// Start the API server with custom state.
615pub async fn serve_with_state(config: ServerConfig, state: AppState) -> crate::Result<()> {
616    let listener = bind(&config).await?;
617    serve_on(listener, config, state).await
618}
619
620/// Serve on an already-bound listener.
621pub async fn serve_on(
622    listener: tokio::net::TcpListener,
623    config: ServerConfig,
624    state: AppState,
625) -> crate::Result<()> {
626    let addr = config.bind_address();
627
628    // Create router with security
629    let (router, auth_store, _rate_limiter) = create_secure_router(state, config.security.clone());
630
631    // Log API key if auth is enabled and keys are registered
632    if auth_store.is_enabled() {
633        if auth_store.count() == 0 {
634            // Generate and register a key if none provided, scoped to the
635            // configured capabilities (full-control when unset).
636            let key = crate::security::generate_api_key();
637            register_key(&auth_store, &key, &config.security.capabilities);
638            tracing::info!("Generated API key: {}", key);
639        }
640        tracing::info!(
641            "Authentication enabled with {} API key(s)",
642            auth_store.count()
643        );
644    } else {
645        tracing::warn!("Authentication is DISABLED - server is open to all requests");
646    }
647
648    let _ = addr;
649    tracing::info!(
650        "Starting shell-tunnel API server on {}",
651        listener
652            .local_addr()
653            .map(|a| a.to_string())
654            .unwrap_or_else(|_| config.bind_address())
655    );
656
657    // Create service with connection info for rate limiting
658    let service: IntoMakeServiceWithConnectInfo<Router, SocketAddr> =
659        router.into_make_service_with_connect_info::<SocketAddr>();
660
661    if config.graceful_shutdown {
662        // Serve with graceful shutdown
663        axum::serve(listener, service)
664            .with_graceful_shutdown(shutdown_signal())
665            .await
666            .map_err(|e| {
667                crate::error::ShellTunnelError::Io(std::io::Error::other(e.to_string()))
668            })?;
669
670        tracing::info!("Server shutdown complete");
671    } else {
672        // Serve without graceful shutdown
673        axum::serve(listener, service).await.map_err(|e| {
674            crate::error::ShellTunnelError::Io(std::io::Error::other(e.to_string()))
675        })?;
676    }
677
678    Ok(())
679}
680
681/// Wait for shutdown signal (Ctrl+C or SIGTERM).
682async fn shutdown_signal() {
683    let ctrl_c = async {
684        tokio::signal::ctrl_c()
685            .await
686            .expect("Failed to install Ctrl+C handler");
687    };
688
689    #[cfg(unix)]
690    let terminate = async {
691        tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
692            .expect("Failed to install SIGTERM handler")
693            .recv()
694            .await;
695    };
696
697    #[cfg(not(unix))]
698    let terminate = std::future::pending::<()>();
699
700    tokio::select! {
701        _ = ctrl_c => {
702            tracing::info!("Received Ctrl+C, initiating graceful shutdown...");
703        }
704        _ = terminate => {
705            tracing::info!("Received SIGTERM, initiating graceful shutdown...");
706        }
707    }
708}
709
710/// Filesystem routes, shared by the plain and the secured router.
711///
712/// Built in one place so the two constructors cannot drift apart — a route
713/// present in only one of them is reachable in only one deployment shape.
714fn fs_routes() -> Router<AppState> {
715    // The chunk-upload route carries request bodies up to `MAX_CHUNK_SIZE` (8
716    // MiB) — well above axum-core's own default body limit (2 MiB, axum-core
717    // 0.5.6's `DEFAULT_LIMIT`), which this app otherwise never overrides. A
718    // client sending a chunk at the server's own advertised `chunk_size` (4
719    // MiB default) would have it rejected before `append_chunk` ever ran.
720    //
721    // Raised only for this one path, via a merged sub-router, rather than
722    // `.layer()` on the whole `fs_routes` router: the latter would also raise
723    // the limit for `/list`, `/stat`, and `/file`, none of which need an 8
724    // MiB body, and a limit set any higher up would reach `/api/v1/execute`
725    // too. Set at the hard ceiling (`MAX_CHUNK_SIZE`) rather than the
726    // *configured* `chunk_size`, so a chunk larger than configured but still
727    // under the ceiling reaches `append_chunk`'s own `TooLarge` check and
728    // gets a 413 with a machine-readable body — if axum's limit cut it off
729    // first, the caller would get a bodyless 413 with no way to tell why.
730    let upload_session_routes = Router::new()
731        .route(
732            "/uploads/{id}",
733            get(super::fs::upload_status)
734                .patch(super::fs::append_chunk)
735                .delete(super::fs::cancel_upload),
736        )
737        .route_layer(DefaultBodyLimit::max(crate::fs::MAX_CHUNK_SIZE));
738
739    Router::new()
740        .route("/list", get(super::fs::list))
741        .route("/stat", get(super::fs::stat))
742        .route(
743            "/file",
744            get(super::fs::download).delete(super::fs::delete_file),
745        )
746        .route("/uploads", post(super::fs::create_upload))
747        .merge(upload_session_routes)
748        .route("/uploads/{id}/complete", post(super::fs::complete_upload))
749}
750
751#[cfg(test)]
752mod tests {
753    use super::*;
754
755    #[test]
756    fn test_server_config_default() {
757        let config = ServerConfig::default();
758        assert_eq!(config.host, "127.0.0.1");
759        assert_eq!(config.port, 3000);
760        assert_eq!(config.bind_address(), "127.0.0.1:3000");
761        assert!(config.graceful_shutdown);
762    }
763
764    #[test]
765    fn test_server_config_custom() {
766        let config = ServerConfig::new("0.0.0.0", 8080);
767        assert_eq!(config.bind_address(), "0.0.0.0:8080");
768    }
769
770    #[test]
771    fn test_server_config_with_security() {
772        let config = ServerConfig::new("0.0.0.0", 8080)
773            .with_security(SecurityConfig::secure().with_api_key("test-key"));
774
775        assert!(config.security.auth.enabled);
776        assert_eq!(config.security.api_keys.len(), 1);
777    }
778
779    #[test]
780    fn test_security_config_default() {
781        let config = SecurityConfig::default();
782        assert!(!config.auth.enabled); // Disabled by default
783        assert!(config.rate_limit.enabled);
784    }
785
786    #[test]
787    fn test_security_config_secure() {
788        let config = SecurityConfig::secure();
789        assert!(config.auth.enabled);
790        assert!(config.rate_limit.enabled);
791    }
792
793    #[test]
794    fn test_cors_restrictive_by_default() {
795        assert!(!SecurityConfig::default().cors.allow_any);
796        assert!(!SecurityConfig::secure().cors.allow_any);
797        assert!(cors_layer(&CorsConfig::default()).is_none());
798    }
799
800    #[test]
801    fn test_cors_allow_any_opt_in() {
802        let config = SecurityConfig::development().with_cors_allow_any();
803        assert!(config.cors.allow_any);
804        assert!(cors_layer(&config.cors).is_some());
805    }
806
807    #[test]
808    fn test_security_config_development() {
809        let config = SecurityConfig::development();
810        assert!(!config.auth.enabled);
811        assert!(config.rate_limit.enabled);
812    }
813
814    #[test]
815    fn test_router_creation() {
816        let _router = create_router();
817        // Router created successfully
818    }
819
820    #[test]
821    fn test_required_capability_mapping() {
822        use RequiredCapability::{Authenticated, Capability, Public};
823
824        // Public + authenticated-only tiers.
825        assert_eq!(required_capability(&Method::GET, "/health"), Public);
826        assert_eq!(required_capability(&Method::GET, "/api/v1"), Authenticated);
827
828        // exec routes (oneshot + session-scoped + WS).
829        assert_eq!(
830            required_capability(&Method::POST, "/api/v1/execute"),
831            Capability("exec")
832        );
833        assert_eq!(
834            required_capability(&Method::GET, "/api/v1/ws"),
835            Capability("exec")
836        );
837        assert_eq!(
838            required_capability(&Method::POST, "/api/v1/sessions/{id}/execute"),
839            Capability("exec")
840        );
841        assert_eq!(
842            required_capability(&Method::GET, "/api/v1/sessions/{id}/ws"),
843            Capability("exec")
844        );
845
846        // read vs manage split on the same path, keyed by method.
847        assert_eq!(
848            required_capability(&Method::GET, "/api/v1/sessions"),
849            Capability("session.read")
850        );
851        assert_eq!(
852            required_capability(&Method::POST, "/api/v1/sessions"),
853            Capability("session.manage")
854        );
855        assert_eq!(
856            required_capability(&Method::GET, "/api/v1/sessions/{id}"),
857            Capability("session.read")
858        );
859        assert_eq!(
860            required_capability(&Method::DELETE, "/api/v1/sessions/{id}"),
861            Capability("session.manage")
862        );
863    }
864
865    #[test]
866    fn test_required_capability_unknown_fails_closed() {
867        // An unmapped route requires at least a valid token (never less).
868        assert_eq!(
869            required_capability(&Method::GET, "/api/v1/unknown"),
870            RequiredCapability::Authenticated
871        );
872    }
873
874    // A direct unit-level check of the `HEAD` normalisation lived here once,
875    // hardcoding three fs paths. It is superseded by
876    // `every_get_fs_route_authorizes_head_identically` in `tests/fs_api.rs`,
877    // which derives the same check from the one authoritative route table
878    // (shared with `every_fs_route_declares_a_capability`) instead of a
879    // second, hand-maintained list that could drift from it.
880
881    #[test]
882    fn test_secure_router_creation() {
883        let state = AppState::new();
884        let security = SecurityConfig::secure().with_api_key("test-key");
885        let (router, auth_store, rate_limiter) = create_secure_router(state, security);
886
887        assert_eq!(auth_store.count(), 1);
888        assert!(auth_store.is_valid("test-key"));
889        assert!(rate_limiter.is_enabled());
890
891        // Router should be created
892        drop(router);
893    }
894}