1use 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#[derive(Debug, Clone, Default)]
42pub struct CorsConfig {
43 pub allow_any: bool,
46}
47
48#[derive(Debug, Clone)]
50pub struct SecurityConfig {
51 pub auth: AuthConfig,
53 pub rate_limit: RateLimitConfig,
55 pub api_keys: Vec<String>,
57 pub capabilities: Option<CapabilitySet>,
64 pub cors: CorsConfig,
66 pub allowed_hosts: Option<Vec<String>>,
75}
76
77impl Default for SecurityConfig {
78 fn default() -> Self {
79 Self {
80 auth: AuthConfig::disabled(), rate_limit: RateLimitConfig::default(),
82 api_keys: Vec::new(),
83 capabilities: None, cors: CorsConfig::default(), allowed_hosts: None,
86 }
87 }
88}
89
90fn 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 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 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 pub fn with_api_key(mut self, key: impl Into<String>) -> Self {
129 self.api_keys.push(key.into());
130 self
131 }
132
133 pub fn with_capabilities(mut self, capabilities: CapabilitySet) -> Self {
138 self.capabilities = Some(capabilities);
139 self
140 }
141
142 pub fn with_allowed_hosts(mut self, hosts: Vec<String>) -> Self {
144 self.allowed_hosts = Some(hosts);
145 self
146 }
147
148 pub fn with_cors_allow_any(mut self) -> Self {
150 self.cors.allow_any = true;
151 self
152 }
153}
154
155fn 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
164pub fn create_router() -> Router {
166 create_router_with_state(AppState::new())
167}
168
169pub fn create_router_with_state(state: AppState) -> Router {
171 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 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 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#[derive(Debug, Clone, Copy, PartialEq, Eq)]
201pub enum RequiredCapability {
202 Public,
204 Authenticated,
207 Capability(&'static str),
209}
210
211pub fn required_capability(method: &Method, matched_path: &str) -> RequiredCapability {
221 use RequiredCapability::{Authenticated, Capability, Public};
222
223 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
259async fn capability_auth_middleware(
268 State((store, audit)): State<(
269 std::sync::Arc<ApiKeyStore>,
270 std::sync::Arc<crate::audit::AuditSink>,
271 )>,
272 mut request: Request,
273 next: Next,
274) -> Result<Response, StatusCode> {
275 if !store.is_enabled() {
277 return Ok(next.run(request).await);
278 }
279
280 let method = request.method().clone();
281 let matched = request
283 .extensions()
284 .get::<MatchedPath>()
285 .map(|m| m.as_str().to_owned())
286 .unwrap_or_default();
287 let required = required_capability(&method, &matched);
288
289 if required == RequiredCapability::Public {
291 return Ok(next.run(request).await);
292 }
293
294 let token = request
297 .headers()
298 .get(AUTHORIZATION)
299 .and_then(|v| v.to_str().ok())
300 .and_then(|header| store.extract_key(header));
301
302 let identity = token.as_deref().and_then(|t| store.identity(t));
303
304 let capabilities = match token.as_deref().and_then(|t| store.capabilities(t)) {
305 Some(caps) => caps,
306 None => {
307 let reason = if token.is_none() {
311 "missing-token"
312 } else {
313 "invalid-token"
314 };
315 tracing::debug!(%method, path = %matched, reason, "auth rejected (401)");
316 audit.record(
326 crate::audit::AuditEvent::new("denied")
327 .with_route(format!("{method} {matched}"))
328 .with_denial(401, reason),
329 );
330 return Err(StatusCode::UNAUTHORIZED);
331 }
332 };
333
334 if let Some(identity) = identity.clone() {
336 request.extensions_mut().insert(identity);
337 }
338
339 match required {
340 RequiredCapability::Public => Ok(next.run(request).await),
342 RequiredCapability::Authenticated => Ok(next.run(request).await),
344 RequiredCapability::Capability(cap) => {
346 if capabilities.satisfies(cap) {
347 Ok(next.run(request).await)
348 } else {
349 audit.record(
350 crate::audit::AuditEvent::new("denied")
351 .with_identity(identity)
352 .with_route(format!("{method} {matched}"))
353 .with_denial(403, format!("missing-capability:{cap}")),
354 );
355 tracing::debug!(
356 %method,
357 path = %matched,
358 required = cap,
359 "authorization denied (403): insufficient capability"
360 );
361 Err(StatusCode::FORBIDDEN)
362 }
363 }
364 }
365}
366
367fn host_is_allowed(header: Option<&str>, allowed: &[String]) -> bool {
373 let Some(value) = header else {
374 return false;
377 };
378
379 let host = value
380 .rsplit_once(':')
381 .map_or(value, |(host, port)| {
382 if port.chars().all(|c| c.is_ascii_digit()) {
384 host
385 } else {
386 value
387 }
388 })
389 .trim_matches(|c| c == '[' || c == ']');
390
391 allowed
392 .iter()
393 .any(|candidate| candidate.eq_ignore_ascii_case(host))
394}
395
396async fn host_check_middleware(
398 State(allowed): State<Arc<Vec<String>>>,
399 request: Request,
400 next: Next,
401) -> Result<Response, (StatusCode, String)> {
402 let header = request
403 .headers()
404 .get(axum::http::header::HOST)
405 .and_then(|value| value.to_str().ok());
406
407 if host_is_allowed(header, &allowed) {
408 return Ok(next.run(request).await);
409 }
410
411 let seen = header.unwrap_or("(none)").to_string();
414 tracing::debug!(host = %seen, "request refused: host not allowed");
415 Err((
416 StatusCode::FORBIDDEN,
417 format!(
418 "host {seen} is not allowed; pass --allow-host {seen} to permit it
419"
420 ),
421 ))
422}
423
424pub fn create_secure_router(
426 state: AppState,
427 security: SecurityConfig,
428) -> (Router, Arc<ApiKeyStore>, Arc<RateLimiter>) {
429 let auth_store = Arc::new(ApiKeyStore::new(security.auth));
431 let rate_limiter = Arc::new(RateLimiter::new(security.rate_limit));
432
433 for key in &security.api_keys {
435 register_key(&auth_store, key, &security.capabilities);
436 }
437
438 let session_routes = Router::new()
440 .route("/", get(list_sessions).post(create_session))
441 .route("/{id}", get(get_session).delete(delete_session))
442 .route("/{id}/execute", post(execute_command))
443 .route("/{id}/ws", any(ws_handler));
444
445 let api_v1 = Router::new()
447 .route("/", get(api_info))
448 .route("/execute", post(execute_oneshot))
449 .route("/ws", any(ws_oneshot_handler))
450 .nest("/fs", fs_routes())
451 .nest("/sessions", session_routes);
452
453 let allowed_hosts = security.allowed_hosts.clone();
454
455 let mut router = Router::new()
457 .route("/health", get(health))
458 .nest("/api/v1", api_v1)
459 .layer(middleware::from_fn_with_state(
460 (Arc::clone(&auth_store), Arc::clone(&state.audit)),
461 capability_auth_middleware,
462 ))
463 .layer(middleware::from_fn_with_state(
464 Arc::clone(&rate_limiter),
465 rate_limit_middleware,
466 ))
467 .layer(TraceLayer::new_for_http());
468
469 if let Some(hosts) = allowed_hosts {
472 router = router.layer(middleware::from_fn_with_state(
473 Arc::new(hosts),
474 host_check_middleware,
475 ));
476 }
477
478 if let Some(cors) = cors_layer(&security.cors) {
480 router = router.layer(cors);
481 }
482
483 let router = router.with_state(state);
484
485 (router, auth_store, rate_limiter)
486}
487
488#[derive(Debug, Clone)]
490pub struct ServerConfig {
491 pub host: String,
493 pub port: u16,
495 pub security: SecurityConfig,
497 pub graceful_shutdown: bool,
499}
500
501impl ServerConfig {
502 pub fn new(host: impl Into<String>, port: u16) -> Self {
503 Self {
504 host: host.into(),
505 port,
506 security: SecurityConfig::default(),
507 graceful_shutdown: true,
508 }
509 }
510
511 pub fn bind_address(&self) -> String {
512 format!("{}:{}", self.host, self.port)
513 }
514
515 pub fn with_security(mut self, security: SecurityConfig) -> Self {
517 self.security = security;
518 self
519 }
520
521 pub fn without_graceful_shutdown(mut self) -> Self {
523 self.graceful_shutdown = false;
524 self
525 }
526}
527
528impl Default for ServerConfig {
529 fn default() -> Self {
530 Self {
531 host: "127.0.0.1".to_string(),
532 port: 3000,
533 security: SecurityConfig::default(),
534 graceful_shutdown: true,
535 }
536 }
537}
538
539pub async fn serve(config: ServerConfig) -> crate::Result<()> {
541 serve_with_state(config, AppState::new()).await
542}
543
544pub async fn bind(config: &ServerConfig) -> crate::Result<tokio::net::TcpListener> {
552 tokio::net::TcpListener::bind(config.bind_address())
553 .await
554 .map_err(crate::error::ShellTunnelError::Io)
555}
556
557pub async fn serve_with_state(config: ServerConfig, state: AppState) -> crate::Result<()> {
559 let listener = bind(&config).await?;
560 serve_on(listener, config, state).await
561}
562
563pub async fn serve_on(
565 listener: tokio::net::TcpListener,
566 config: ServerConfig,
567 state: AppState,
568) -> crate::Result<()> {
569 let addr = config.bind_address();
570
571 let (router, auth_store, _rate_limiter) = create_secure_router(state, config.security.clone());
573
574 if auth_store.is_enabled() {
576 if auth_store.count() == 0 {
577 let key = crate::security::generate_api_key();
580 register_key(&auth_store, &key, &config.security.capabilities);
581 tracing::info!("Generated API key: {}", key);
582 }
583 tracing::info!(
584 "Authentication enabled with {} API key(s)",
585 auth_store.count()
586 );
587 } else {
588 tracing::warn!("Authentication is DISABLED - server is open to all requests");
589 }
590
591 let _ = addr;
592 tracing::info!(
593 "Starting shell-tunnel API server on {}",
594 listener
595 .local_addr()
596 .map(|a| a.to_string())
597 .unwrap_or_else(|_| config.bind_address())
598 );
599
600 let service: IntoMakeServiceWithConnectInfo<Router, SocketAddr> =
602 router.into_make_service_with_connect_info::<SocketAddr>();
603
604 if config.graceful_shutdown {
605 axum::serve(listener, service)
607 .with_graceful_shutdown(shutdown_signal())
608 .await
609 .map_err(|e| {
610 crate::error::ShellTunnelError::Io(std::io::Error::other(e.to_string()))
611 })?;
612
613 tracing::info!("Server shutdown complete");
614 } else {
615 axum::serve(listener, service).await.map_err(|e| {
617 crate::error::ShellTunnelError::Io(std::io::Error::other(e.to_string()))
618 })?;
619 }
620
621 Ok(())
622}
623
624async fn shutdown_signal() {
626 let ctrl_c = async {
627 tokio::signal::ctrl_c()
628 .await
629 .expect("Failed to install Ctrl+C handler");
630 };
631
632 #[cfg(unix)]
633 let terminate = async {
634 tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
635 .expect("Failed to install SIGTERM handler")
636 .recv()
637 .await;
638 };
639
640 #[cfg(not(unix))]
641 let terminate = std::future::pending::<()>();
642
643 tokio::select! {
644 _ = ctrl_c => {
645 tracing::info!("Received Ctrl+C, initiating graceful shutdown...");
646 }
647 _ = terminate => {
648 tracing::info!("Received SIGTERM, initiating graceful shutdown...");
649 }
650 }
651}
652
653fn fs_routes() -> Router<AppState> {
658 let upload_session_routes = Router::new()
674 .route(
675 "/uploads/{id}",
676 get(super::fs::upload_status)
677 .patch(super::fs::append_chunk)
678 .delete(super::fs::cancel_upload),
679 )
680 .route_layer(DefaultBodyLimit::max(crate::fs::MAX_CHUNK_SIZE));
681
682 Router::new()
683 .route("/list", get(super::fs::list))
684 .route("/stat", get(super::fs::stat))
685 .route(
686 "/file",
687 get(super::fs::download).delete(super::fs::delete_file),
688 )
689 .route("/uploads", post(super::fs::create_upload))
690 .merge(upload_session_routes)
691 .route("/uploads/{id}/complete", post(super::fs::complete_upload))
692}
693
694#[cfg(test)]
695mod tests {
696 use super::*;
697
698 #[test]
699 fn test_server_config_default() {
700 let config = ServerConfig::default();
701 assert_eq!(config.host, "127.0.0.1");
702 assert_eq!(config.port, 3000);
703 assert_eq!(config.bind_address(), "127.0.0.1:3000");
704 assert!(config.graceful_shutdown);
705 }
706
707 #[test]
708 fn test_server_config_custom() {
709 let config = ServerConfig::new("0.0.0.0", 8080);
710 assert_eq!(config.bind_address(), "0.0.0.0:8080");
711 }
712
713 #[test]
714 fn test_server_config_with_security() {
715 let config = ServerConfig::new("0.0.0.0", 8080)
716 .with_security(SecurityConfig::secure().with_api_key("test-key"));
717
718 assert!(config.security.auth.enabled);
719 assert_eq!(config.security.api_keys.len(), 1);
720 }
721
722 #[test]
723 fn test_security_config_default() {
724 let config = SecurityConfig::default();
725 assert!(!config.auth.enabled); assert!(config.rate_limit.enabled);
727 }
728
729 #[test]
730 fn test_security_config_secure() {
731 let config = SecurityConfig::secure();
732 assert!(config.auth.enabled);
733 assert!(config.rate_limit.enabled);
734 }
735
736 #[test]
737 fn test_cors_restrictive_by_default() {
738 assert!(!SecurityConfig::default().cors.allow_any);
739 assert!(!SecurityConfig::secure().cors.allow_any);
740 assert!(cors_layer(&CorsConfig::default()).is_none());
741 }
742
743 #[test]
744 fn test_cors_allow_any_opt_in() {
745 let config = SecurityConfig::development().with_cors_allow_any();
746 assert!(config.cors.allow_any);
747 assert!(cors_layer(&config.cors).is_some());
748 }
749
750 #[test]
751 fn test_security_config_development() {
752 let config = SecurityConfig::development();
753 assert!(!config.auth.enabled);
754 assert!(config.rate_limit.enabled);
755 }
756
757 #[test]
758 fn test_router_creation() {
759 let _router = create_router();
760 }
762
763 #[test]
764 fn test_required_capability_mapping() {
765 use RequiredCapability::{Authenticated, Capability, Public};
766
767 assert_eq!(required_capability(&Method::GET, "/health"), Public);
769 assert_eq!(required_capability(&Method::GET, "/api/v1"), Authenticated);
770
771 assert_eq!(
773 required_capability(&Method::POST, "/api/v1/execute"),
774 Capability("exec")
775 );
776 assert_eq!(
777 required_capability(&Method::GET, "/api/v1/ws"),
778 Capability("exec")
779 );
780 assert_eq!(
781 required_capability(&Method::POST, "/api/v1/sessions/{id}/execute"),
782 Capability("exec")
783 );
784 assert_eq!(
785 required_capability(&Method::GET, "/api/v1/sessions/{id}/ws"),
786 Capability("exec")
787 );
788
789 assert_eq!(
791 required_capability(&Method::GET, "/api/v1/sessions"),
792 Capability("session.read")
793 );
794 assert_eq!(
795 required_capability(&Method::POST, "/api/v1/sessions"),
796 Capability("session.manage")
797 );
798 assert_eq!(
799 required_capability(&Method::GET, "/api/v1/sessions/{id}"),
800 Capability("session.read")
801 );
802 assert_eq!(
803 required_capability(&Method::DELETE, "/api/v1/sessions/{id}"),
804 Capability("session.manage")
805 );
806 }
807
808 #[test]
809 fn test_required_capability_unknown_fails_closed() {
810 assert_eq!(
812 required_capability(&Method::GET, "/api/v1/unknown"),
813 RequiredCapability::Authenticated
814 );
815 }
816
817 #[test]
825 fn test_secure_router_creation() {
826 let state = AppState::new();
827 let security = SecurityConfig::secure().with_api_key("test-key");
828 let (router, auth_store, rate_limiter) = create_secure_router(state, security);
829
830 assert_eq!(auth_store.count(), 1);
831 assert!(auth_store.is_valid("test-key"));
832 assert!(rate_limiter.is_enabled());
833
834 drop(router);
836 }
837}