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(
319 crate::audit::AuditEvent::new("denied")
320 .with_route(format!("{method} {matched}"))
321 .with_denial(401, reason),
322 );
323 return Err(StatusCode::UNAUTHORIZED);
324 }
325 };
326
327 if let Some(identity) = identity.clone() {
329 request.extensions_mut().insert(identity);
330 }
331
332 match required {
333 RequiredCapability::Public => Ok(next.run(request).await),
335 RequiredCapability::Authenticated => Ok(next.run(request).await),
337 RequiredCapability::Capability(cap) => {
339 if capabilities.satisfies(cap) {
340 Ok(next.run(request).await)
341 } else {
342 audit.record(
343 crate::audit::AuditEvent::new("denied")
344 .with_identity(identity)
345 .with_route(format!("{method} {matched}"))
346 .with_denial(403, format!("missing-capability:{cap}")),
347 );
348 tracing::debug!(
349 %method,
350 path = %matched,
351 required = cap,
352 "authorization denied (403): insufficient capability"
353 );
354 Err(StatusCode::FORBIDDEN)
355 }
356 }
357 }
358}
359
360fn host_is_allowed(header: Option<&str>, allowed: &[String]) -> bool {
366 let Some(value) = header else {
367 return false;
370 };
371
372 let host = value
373 .rsplit_once(':')
374 .map_or(value, |(host, port)| {
375 if port.chars().all(|c| c.is_ascii_digit()) {
377 host
378 } else {
379 value
380 }
381 })
382 .trim_matches(|c| c == '[' || c == ']');
383
384 allowed
385 .iter()
386 .any(|candidate| candidate.eq_ignore_ascii_case(host))
387}
388
389async fn host_check_middleware(
391 State(allowed): State<Arc<Vec<String>>>,
392 request: Request,
393 next: Next,
394) -> Result<Response, (StatusCode, String)> {
395 let header = request
396 .headers()
397 .get(axum::http::header::HOST)
398 .and_then(|value| value.to_str().ok());
399
400 if host_is_allowed(header, &allowed) {
401 return Ok(next.run(request).await);
402 }
403
404 let seen = header.unwrap_or("(none)").to_string();
407 tracing::debug!(host = %seen, "request refused: host not allowed");
408 Err((
409 StatusCode::FORBIDDEN,
410 format!(
411 "host {seen} is not allowed; pass --allow-host {seen} to permit it
412"
413 ),
414 ))
415}
416
417pub fn create_secure_router(
419 state: AppState,
420 security: SecurityConfig,
421) -> (Router, Arc<ApiKeyStore>, Arc<RateLimiter>) {
422 let auth_store = Arc::new(ApiKeyStore::new(security.auth));
424 let rate_limiter = Arc::new(RateLimiter::new(security.rate_limit));
425
426 for key in &security.api_keys {
428 register_key(&auth_store, key, &security.capabilities);
429 }
430
431 let session_routes = Router::new()
433 .route("/", get(list_sessions).post(create_session))
434 .route("/{id}", get(get_session).delete(delete_session))
435 .route("/{id}/execute", post(execute_command))
436 .route("/{id}/ws", any(ws_handler));
437
438 let api_v1 = Router::new()
440 .route("/", get(api_info))
441 .route("/execute", post(execute_oneshot))
442 .route("/ws", any(ws_oneshot_handler))
443 .nest("/fs", fs_routes())
444 .nest("/sessions", session_routes);
445
446 let allowed_hosts = security.allowed_hosts.clone();
447
448 let mut router = Router::new()
450 .route("/health", get(health))
451 .nest("/api/v1", api_v1)
452 .layer(middleware::from_fn_with_state(
453 (Arc::clone(&auth_store), Arc::clone(&state.audit)),
454 capability_auth_middleware,
455 ))
456 .layer(middleware::from_fn_with_state(
457 Arc::clone(&rate_limiter),
458 rate_limit_middleware,
459 ))
460 .layer(TraceLayer::new_for_http());
461
462 if let Some(hosts) = allowed_hosts {
465 router = router.layer(middleware::from_fn_with_state(
466 Arc::new(hosts),
467 host_check_middleware,
468 ));
469 }
470
471 if let Some(cors) = cors_layer(&security.cors) {
473 router = router.layer(cors);
474 }
475
476 let router = router.with_state(state);
477
478 (router, auth_store, rate_limiter)
479}
480
481#[derive(Debug, Clone)]
483pub struct ServerConfig {
484 pub host: String,
486 pub port: u16,
488 pub security: SecurityConfig,
490 pub graceful_shutdown: bool,
492}
493
494impl ServerConfig {
495 pub fn new(host: impl Into<String>, port: u16) -> Self {
496 Self {
497 host: host.into(),
498 port,
499 security: SecurityConfig::default(),
500 graceful_shutdown: true,
501 }
502 }
503
504 pub fn bind_address(&self) -> String {
505 format!("{}:{}", self.host, self.port)
506 }
507
508 pub fn with_security(mut self, security: SecurityConfig) -> Self {
510 self.security = security;
511 self
512 }
513
514 pub fn without_graceful_shutdown(mut self) -> Self {
516 self.graceful_shutdown = false;
517 self
518 }
519}
520
521impl Default for ServerConfig {
522 fn default() -> Self {
523 Self {
524 host: "127.0.0.1".to_string(),
525 port: 3000,
526 security: SecurityConfig::default(),
527 graceful_shutdown: true,
528 }
529 }
530}
531
532pub async fn serve(config: ServerConfig) -> crate::Result<()> {
534 serve_with_state(config, AppState::new()).await
535}
536
537pub async fn bind(config: &ServerConfig) -> crate::Result<tokio::net::TcpListener> {
545 tokio::net::TcpListener::bind(config.bind_address())
546 .await
547 .map_err(crate::error::ShellTunnelError::Io)
548}
549
550pub async fn serve_with_state(config: ServerConfig, state: AppState) -> crate::Result<()> {
552 let listener = bind(&config).await?;
553 serve_on(listener, config, state).await
554}
555
556pub async fn serve_on(
558 listener: tokio::net::TcpListener,
559 config: ServerConfig,
560 state: AppState,
561) -> crate::Result<()> {
562 let addr = config.bind_address();
563
564 let (router, auth_store, _rate_limiter) = create_secure_router(state, config.security.clone());
566
567 if auth_store.is_enabled() {
569 if auth_store.count() == 0 {
570 let key = crate::security::generate_api_key();
573 register_key(&auth_store, &key, &config.security.capabilities);
574 tracing::info!("Generated API key: {}", key);
575 }
576 tracing::info!(
577 "Authentication enabled with {} API key(s)",
578 auth_store.count()
579 );
580 } else {
581 tracing::warn!("Authentication is DISABLED - server is open to all requests");
582 }
583
584 let _ = addr;
585 tracing::info!(
586 "Starting shell-tunnel API server on {}",
587 listener
588 .local_addr()
589 .map(|a| a.to_string())
590 .unwrap_or_else(|_| config.bind_address())
591 );
592
593 let service: IntoMakeServiceWithConnectInfo<Router, SocketAddr> =
595 router.into_make_service_with_connect_info::<SocketAddr>();
596
597 if config.graceful_shutdown {
598 axum::serve(listener, service)
600 .with_graceful_shutdown(shutdown_signal())
601 .await
602 .map_err(|e| {
603 crate::error::ShellTunnelError::Io(std::io::Error::other(e.to_string()))
604 })?;
605
606 tracing::info!("Server shutdown complete");
607 } else {
608 axum::serve(listener, service).await.map_err(|e| {
610 crate::error::ShellTunnelError::Io(std::io::Error::other(e.to_string()))
611 })?;
612 }
613
614 Ok(())
615}
616
617async fn shutdown_signal() {
619 let ctrl_c = async {
620 tokio::signal::ctrl_c()
621 .await
622 .expect("Failed to install Ctrl+C handler");
623 };
624
625 #[cfg(unix)]
626 let terminate = async {
627 tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())
628 .expect("Failed to install SIGTERM handler")
629 .recv()
630 .await;
631 };
632
633 #[cfg(not(unix))]
634 let terminate = std::future::pending::<()>();
635
636 tokio::select! {
637 _ = ctrl_c => {
638 tracing::info!("Received Ctrl+C, initiating graceful shutdown...");
639 }
640 _ = terminate => {
641 tracing::info!("Received SIGTERM, initiating graceful shutdown...");
642 }
643 }
644}
645
646fn fs_routes() -> Router<AppState> {
651 let upload_session_routes = Router::new()
667 .route(
668 "/uploads/{id}",
669 get(super::fs::upload_status)
670 .patch(super::fs::append_chunk)
671 .delete(super::fs::cancel_upload),
672 )
673 .route_layer(DefaultBodyLimit::max(crate::fs::MAX_CHUNK_SIZE));
674
675 Router::new()
676 .route("/list", get(super::fs::list))
677 .route("/stat", get(super::fs::stat))
678 .route(
679 "/file",
680 get(super::fs::download).delete(super::fs::delete_file),
681 )
682 .route("/uploads", post(super::fs::create_upload))
683 .merge(upload_session_routes)
684 .route("/uploads/{id}/complete", post(super::fs::complete_upload))
685}
686
687#[cfg(test)]
688mod tests {
689 use super::*;
690
691 #[test]
692 fn test_server_config_default() {
693 let config = ServerConfig::default();
694 assert_eq!(config.host, "127.0.0.1");
695 assert_eq!(config.port, 3000);
696 assert_eq!(config.bind_address(), "127.0.0.1:3000");
697 assert!(config.graceful_shutdown);
698 }
699
700 #[test]
701 fn test_server_config_custom() {
702 let config = ServerConfig::new("0.0.0.0", 8080);
703 assert_eq!(config.bind_address(), "0.0.0.0:8080");
704 }
705
706 #[test]
707 fn test_server_config_with_security() {
708 let config = ServerConfig::new("0.0.0.0", 8080)
709 .with_security(SecurityConfig::secure().with_api_key("test-key"));
710
711 assert!(config.security.auth.enabled);
712 assert_eq!(config.security.api_keys.len(), 1);
713 }
714
715 #[test]
716 fn test_security_config_default() {
717 let config = SecurityConfig::default();
718 assert!(!config.auth.enabled); assert!(config.rate_limit.enabled);
720 }
721
722 #[test]
723 fn test_security_config_secure() {
724 let config = SecurityConfig::secure();
725 assert!(config.auth.enabled);
726 assert!(config.rate_limit.enabled);
727 }
728
729 #[test]
730 fn test_cors_restrictive_by_default() {
731 assert!(!SecurityConfig::default().cors.allow_any);
732 assert!(!SecurityConfig::secure().cors.allow_any);
733 assert!(cors_layer(&CorsConfig::default()).is_none());
734 }
735
736 #[test]
737 fn test_cors_allow_any_opt_in() {
738 let config = SecurityConfig::development().with_cors_allow_any();
739 assert!(config.cors.allow_any);
740 assert!(cors_layer(&config.cors).is_some());
741 }
742
743 #[test]
744 fn test_security_config_development() {
745 let config = SecurityConfig::development();
746 assert!(!config.auth.enabled);
747 assert!(config.rate_limit.enabled);
748 }
749
750 #[test]
751 fn test_router_creation() {
752 let _router = create_router();
753 }
755
756 #[test]
757 fn test_required_capability_mapping() {
758 use RequiredCapability::{Authenticated, Capability, Public};
759
760 assert_eq!(required_capability(&Method::GET, "/health"), Public);
762 assert_eq!(required_capability(&Method::GET, "/api/v1"), Authenticated);
763
764 assert_eq!(
766 required_capability(&Method::POST, "/api/v1/execute"),
767 Capability("exec")
768 );
769 assert_eq!(
770 required_capability(&Method::GET, "/api/v1/ws"),
771 Capability("exec")
772 );
773 assert_eq!(
774 required_capability(&Method::POST, "/api/v1/sessions/{id}/execute"),
775 Capability("exec")
776 );
777 assert_eq!(
778 required_capability(&Method::GET, "/api/v1/sessions/{id}/ws"),
779 Capability("exec")
780 );
781
782 assert_eq!(
784 required_capability(&Method::GET, "/api/v1/sessions"),
785 Capability("session.read")
786 );
787 assert_eq!(
788 required_capability(&Method::POST, "/api/v1/sessions"),
789 Capability("session.manage")
790 );
791 assert_eq!(
792 required_capability(&Method::GET, "/api/v1/sessions/{id}"),
793 Capability("session.read")
794 );
795 assert_eq!(
796 required_capability(&Method::DELETE, "/api/v1/sessions/{id}"),
797 Capability("session.manage")
798 );
799 }
800
801 #[test]
802 fn test_required_capability_unknown_fails_closed() {
803 assert_eq!(
805 required_capability(&Method::GET, "/api/v1/unknown"),
806 RequiredCapability::Authenticated
807 );
808 }
809
810 #[test]
818 fn test_secure_router_creation() {
819 let state = AppState::new();
820 let security = SecurityConfig::secure().with_api_key("test-key");
821 let (router, auth_store, rate_limiter) = create_secure_router(state, security);
822
823 assert_eq!(auth_store.count(), 1);
824 assert!(auth_store.is_valid("test-key"));
825 assert!(rate_limiter.is_enabled());
826
827 drop(router);
829 }
830}