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
259const MAX_AUDITED_PATH: usize = 256;
268
269fn 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 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
309async 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 if !store.is_enabled() {
327 return Ok(next.run(request).await);
328 }
329
330 let method = request.method().clone();
331 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 if required == RequiredCapability::Public {
347 return Ok(next.run(request).await);
348 }
349
350 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 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 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 if let Some(identity) = identity.clone() {
393 request.extensions_mut().insert(identity);
394 }
395
396 match required {
397 RequiredCapability::Public => Ok(next.run(request).await),
399 RequiredCapability::Authenticated => Ok(next.run(request).await),
401 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
424fn host_is_allowed(header: Option<&str>, allowed: &[String]) -> bool {
430 let Some(value) = header else {
431 return false;
434 };
435
436 let host = value
437 .rsplit_once(':')
438 .map_or(value, |(host, port)| {
439 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
453async 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 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
481pub fn create_secure_router(
483 state: AppState,
484 security: SecurityConfig,
485) -> (Router, Arc<ApiKeyStore>, Arc<RateLimiter>) {
486 let auth_store = Arc::new(ApiKeyStore::new(security.auth));
488 let rate_limiter = Arc::new(RateLimiter::new(security.rate_limit));
489
490 for key in &security.api_keys {
492 register_key(&auth_store, key, &security.capabilities);
493 }
494
495 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 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 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 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 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#[derive(Debug, Clone)]
547pub struct ServerConfig {
548 pub host: String,
550 pub port: u16,
552 pub security: SecurityConfig,
554 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 pub fn with_security(mut self, security: SecurityConfig) -> Self {
574 self.security = security;
575 self
576 }
577
578 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
596pub async fn serve(config: ServerConfig) -> crate::Result<()> {
598 serve_with_state(config, AppState::new()).await
599}
600
601pub 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
614pub 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
620pub 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 let (router, auth_store, _rate_limiter) = create_secure_router(state, config.security.clone());
630
631 if auth_store.is_enabled() {
633 if auth_store.count() == 0 {
634 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 let service: IntoMakeServiceWithConnectInfo<Router, SocketAddr> =
659 router.into_make_service_with_connect_info::<SocketAddr>();
660
661 if config.graceful_shutdown {
662 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 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
681async 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
710fn fs_routes() -> Router<AppState> {
715 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); 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 }
819
820 #[test]
821 fn test_required_capability_mapping() {
822 use RequiredCapability::{Authenticated, Capability, Public};
823
824 assert_eq!(required_capability(&Method::GET, "/health"), Public);
826 assert_eq!(required_capability(&Method::GET, "/api/v1"), Authenticated);
827
828 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 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 assert_eq!(
869 required_capability(&Method::GET, "/api/v1/unknown"),
870 RequiredCapability::Authenticated
871 );
872 }
873
874 #[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 drop(router);
893 }
894}