1#[cfg(all(feature = "rust_crypto", feature = "aws_lc_rs"))]
23compile_error!("features `rust_crypto` and `aws_lc_rs` are mutually exclusive; enable exactly one");
24
25#[cfg(not(any(feature = "rust_crypto", feature = "aws_lc_rs")))]
26compile_error!("exactly one JWT crypto backend must be enabled: `rust_crypto` or `aws_lc_rs`");
27
28pub mod auth;
29pub mod config;
30mod embed;
31pub mod hooks;
32pub mod oidc;
33pub mod openapi;
34pub mod shield;
35pub mod transcode;
36
37use axum::extract::State;
38use axum::http::{Request, StatusCode};
39use axum::middleware::Next;
40use axum::response::{IntoResponse, Response};
41use axum::routing::get;
42use axum::{Json, Router};
43use prost_reflect::DescriptorPool;
44use std::net::SocketAddr;
45use tower_http::cors::{AllowOrigin, CorsLayer};
46use tower_http::trace::TraceLayer;
47
48use std::sync::Arc;
49
50use config::{DescriptorSource, ProxyConfig};
51use hooks::{AuthDecider, ExtraRoute, OidcBackend};
52
53#[derive(Clone, Debug)]
55pub struct ProxyState {
56 pub service_name: String,
58 pub grpc_upstream: String,
60 pub grpc_channel: tonic::transport::Channel,
62 pub maintenance_mode: bool,
64 pub maintenance_exempt: Vec<String>,
66 pub maintenance_message: String,
68 pub forwarded_headers: Vec<String>,
70 pub metrics_namespace: String,
72 pub metrics_classes: Vec<config::MetricsClassConfig>,
74 pub sse_keep_alive_secs: u64,
76}
77
78pub struct ProxyServer {
80 config: ProxyConfig,
81 descriptor_pool: Option<DescriptorPool>,
83 auth_decider: Option<Arc<dyn AuthDecider>>,
85 oidc_backend: Option<Arc<dyn OidcBackend>>,
87 extra_routes: Vec<ExtraRoute>,
89 verify_path: Option<String>,
91}
92
93impl ProxyServer {
94 pub fn from_config(config: ProxyConfig) -> Self {
96 Self {
97 config,
98 descriptor_pool: None,
99 auth_decider: None,
100 oidc_backend: None,
101 extra_routes: Vec::new(),
102 verify_path: None,
103 }
104 }
105
106 pub fn with_descriptors(mut self, pool: DescriptorPool) -> Self {
108 self.descriptor_pool = Some(pool);
109 self
110 }
111
112 pub fn with_auth_decider(mut self, decider: Arc<dyn AuthDecider>) -> Self {
118 self.auth_decider = Some(decider);
119 self
120 }
121
122 pub fn with_oidc_backend(mut self, backend: Arc<dyn OidcBackend>) -> Self {
128 self.oidc_backend = Some(backend);
129 self
130 }
131
132 pub fn with_extra_routes(mut self, routes: impl IntoIterator<Item = ExtraRoute>) -> Self {
135 self.extra_routes.extend(routes);
136 self
137 }
138
139 pub fn with_verify_path(mut self, path: impl Into<String>) -> Self {
146 self.verify_path = Some(path.into());
147 self
148 }
149
150 fn load_descriptors(&self) -> anyhow::Result<DescriptorPool> {
155 if let Some(pool) = &self.descriptor_pool {
156 return Ok(pool.clone());
157 }
158
159 let mut pool = DescriptorPool::new();
160
161 for source in &self.config.descriptors {
162 match source {
163 DescriptorSource::File { file } => {
164 let bytes = std::fs::read(file).map_err(|e| {
165 anyhow::anyhow!("Failed to read descriptor file {:?}: {}", file, e)
166 })?;
167 pool.decode_file_descriptor_set(bytes.as_slice())
168 .map_err(|e| {
169 anyhow::anyhow!("Failed to decode descriptor file {:?}: {}", file, e)
170 })?;
171 tracing::info!("Loaded descriptor from {:?}", file);
172 }
173 DescriptorSource::Reflection { reflection } => {
174 tracing::warn!(
175 "gRPC reflection client not supported — use descriptor files instead (reflection endpoint: {})",
176 reflection
177 );
178 }
179 DescriptorSource::Embedded { bytes } => {
180 pool.decode_file_descriptor_set(*bytes).map_err(|e| {
181 anyhow::anyhow!("Failed to decode embedded descriptors: {}", e)
182 })?;
183 }
184 }
185 }
186
187 Ok(pool)
188 }
189
190 fn decider_verify_path(&self) -> String {
195 self.verify_path.clone().unwrap_or_else(|| {
196 self.config
197 .auth
198 .as_ref()
199 .and_then(|a| a.forward_auth.as_ref())
200 .map(|fa| fa.path.clone())
201 .unwrap_or_else(|| "/auth/verify".to_string())
202 })
203 }
204
205 fn mounted_verify_path(&self) -> Option<String> {
215 if self.auth_decider.is_some() {
216 return Some(self.decider_verify_path());
217 }
218 self.config.auth.as_ref().and_then(|a| {
219 if a.mode != "jwt" {
220 return None;
221 }
222 a.forward_auth
223 .as_ref()
224 .filter(|fa| fa.enabled)
225 .map(|fa| fa.path.clone())
226 })
227 }
228
229 fn reserved_routes(&self, pool: &DescriptorPool) -> anyhow::Result<Vec<(String, String)>> {
240 let mut routes = Vec::new();
241 let mut get = |path: String| routes.push(("GET".to_string(), path));
242 if self.config.health.enabled {
243 get(self.config.health.path.clone());
244 get(self.config.health.live_path.clone());
245 get(self.config.health.ready_path.clone());
246 get(self.config.health.startup_path.clone());
247 }
248 if self.config.metrics.enabled {
249 get(self.config.metrics.path.clone());
250 }
251 if let Some(openapi) = self.config.openapi.as_ref().filter(|o| o.enabled) {
252 get(openapi.path.clone());
253 get(openapi.docs_path.clone());
254 }
255 if let Some(backend) = &self.oidc_backend {
257 for doc in backend.metadata_documents() {
258 get(doc.path);
259 }
260 get(backend.jwks().path);
261 get(backend.userinfo_path());
262 } else if let Some(cfg) = &self.config.oidc_discovery {
263 if let Some(oidc) = oidc::Oidc::build(cfg)
264 .map_err(|e| anyhow::anyhow!("invalid oidc_discovery config: {e}"))?
265 {
266 for path in oidc.paths() {
267 get(path);
268 }
269 }
270 }
271 for route in &self.extra_routes {
272 routes.push((route.method.as_str().to_string(), route.path.clone()));
273 }
274 routes.extend(transcode::route_paths(pool, &self.config.aliases));
275 Ok(routes)
276 }
277
278 pub fn router(&self) -> anyhow::Result<Router> {
280 self.config.validate()?;
283 let pool = self.load_descriptors()?;
284
285 let grpc_upstream = self.config.upstream.default.clone();
286 let grpc_channel = tonic::transport::Channel::from_shared(grpc_upstream.clone())
287 .map_err(|e| anyhow::anyhow!("invalid gRPC upstream URL: {}", e))?
288 .connect_timeout(std::time::Duration::from_secs(5))
289 .timeout(std::time::Duration::from_secs(5))
290 .connect_lazy();
291
292 let service_name = self.config.service.name.clone();
293 let metrics_namespace = service_name.replace('-', "_");
294
295 let verify_path = self.mounted_verify_path();
297
298 let mut mounted = self.reserved_routes(&pool)?;
307 if let Some(vp) = &verify_path {
308 mounted.push(("*".to_string(), vp.clone()));
309 }
310 let mut methods_by_shape: std::collections::HashMap<
314 String,
315 std::collections::HashSet<&str>,
316 > = std::collections::HashMap::new();
317 for (method, path) in &mounted {
318 if !path.starts_with('/') {
319 anyhow::bail!("route path {path:?} must start with '/'");
320 }
321 let methods = methods_by_shape
322 .entry(normalize_route_shape(path))
323 .or_default();
324 let conflict = if method == "*" {
327 !methods.is_empty()
328 } else {
329 methods.contains("*") || methods.contains(method.as_str())
330 };
331 if conflict {
332 anyhow::bail!("route path {path:?} is registered by more than one endpoint");
333 }
334 methods.insert(method.as_str());
335 }
336
337 let mut maintenance_exempt = self.config.maintenance.exempt_paths.clone();
343 if self.config.health.enabled {
344 maintenance_exempt.push(self.config.health.path.clone());
345 maintenance_exempt.push(self.config.health.live_path.clone());
346 maintenance_exempt.push(self.config.health.ready_path.clone());
347 maintenance_exempt.push(self.config.health.startup_path.clone());
348 }
349 if self.config.metrics.enabled {
350 maintenance_exempt.push(self.config.metrics.path.clone());
351 }
352 if let Some(vp) = &verify_path {
353 maintenance_exempt.push(vp.clone());
354 }
355
356 let state = ProxyState {
357 service_name: service_name.clone(),
358 grpc_upstream,
359 grpc_channel,
360 maintenance_mode: self.config.maintenance.enabled,
361 maintenance_exempt,
362 maintenance_message: self.config.maintenance.message.clone(),
363 forwarded_headers: self.config.forwarded_headers.clone(),
364 metrics_namespace,
365 metrics_classes: self.config.metrics_classes.clone(),
366 sse_keep_alive_secs: self.config.streaming.sse_keep_alive_secs,
367 };
368
369 let cors = self.build_cors();
370
371 let mut transcode_routes = transcode::routes(&pool, &self.config.aliases);
373
374 let authz = match self.config.auth.as_ref().and_then(|a| a.authz.as_ref()) {
379 Some(cfg) => auth::authz::Authz::build(cfg)
380 .map_err(|e| anyhow::anyhow!("invalid authz config: {e}"))?,
381 None => None,
382 };
383
384 if let Some(decider) = &self.auth_decider {
390 transcode_routes = transcode_routes.layer(axum::middleware::from_fn_with_state(
391 decider.clone(),
392 embed::auth_decider_gate,
393 ));
394 }
395 if let Some(authz) = authz {
396 transcode_routes = transcode_routes.layer(axum::middleware::from_fn_with_state(
397 authz,
398 auth::authz::middleware,
399 ));
400 }
401
402 let health_routes = if self.config.health.enabled {
404 let health = &self.config.health;
405 let health_service_name = service_name.clone();
406 Router::new()
407 .route(
408 &health.path,
409 get({
410 let name = health_service_name.clone();
411 move || async move {
412 Json(serde_json::json!({
413 "status": "ok",
414 "service": name,
415 }))
416 }
417 }),
418 )
419 .route(&health.live_path, get(|| async { StatusCode::OK }))
420 .route(
421 &health.ready_path,
422 get(|State(state): State<ProxyState>| async move {
423 let mut client =
424 tonic_health::pb::health_client::HealthClient::new(state.grpc_channel);
425 match client
426 .check(tonic_health::pb::HealthCheckRequest {
427 service: String::new(),
428 })
429 .await
430 {
431 Ok(resp) => {
432 let status = resp.into_inner().status;
433 if status
434 == tonic_health::pb::health_check_response::ServingStatus::Serving
435 as i32
436 {
437 StatusCode::OK
438 } else {
439 StatusCode::SERVICE_UNAVAILABLE
440 }
441 }
442 Err(_) => StatusCode::SERVICE_UNAVAILABLE,
443 }
444 }),
445 )
446 .route(&health.startup_path, get(|| async { StatusCode::OK }))
447 } else {
448 Router::new()
449 };
450
451 let metrics_routes = if self.config.metrics.enabled {
453 Router::new().route(
454 &self.config.metrics.path,
455 get(|| async {
456 let encoder = prometheus::TextEncoder::new();
457 let metric_families = prometheus::default_registry().gather();
458 match encoder.encode_to_string(&metric_families) {
459 Ok(text) => (
460 StatusCode::OK,
461 [(
462 axum::http::header::CONTENT_TYPE,
463 "text/plain; version=0.0.4; charset=utf-8",
464 )],
465 text,
466 )
467 .into_response(),
468 Err(_) => StatusCode::INTERNAL_SERVER_ERROR.into_response(),
469 }
470 }),
471 )
472 } else {
473 Router::new()
474 };
475
476 let openapi_routes = self.build_openapi_routes(&pool);
478
479 let oidc_routes = match &self.oidc_backend {
483 Some(backend) => embed::oidc_backend_routes(backend.clone()),
484 None => match &self.config.oidc_discovery {
485 Some(cfg) => oidc::Oidc::build(cfg)
486 .map_err(|e| anyhow::anyhow!("invalid oidc_discovery config: {e}"))?
487 .map(|o| o.routes())
488 .unwrap_or_default(),
489 None => Router::new(),
490 },
491 };
492
493 let shield = match &self.config.shield {
495 Some(cfg) => shield::Shield::build(cfg)
496 .map_err(|e| anyhow::anyhow!("invalid shield config: {e}"))?,
497 None => None,
498 };
499
500 let auth = match &self.config.auth {
502 Some(cfg) => {
503 auth::Auth::build(cfg).map_err(|e| anyhow::anyhow!("invalid auth config: {e}"))?
504 }
505 None => None,
506 };
507
508 let mut router = Router::new()
509 .merge(health_routes)
510 .merge(metrics_routes)
511 .merge(openapi_routes)
512 .merge(oidc_routes)
513 .merge(embed::extra_routes_router(&self.extra_routes))
514 .merge(transcode_routes);
515 let forward_auth = auth.as_ref().and_then(|built| {
523 auth::forward::ForwardAuth::build(self.config.auth.as_ref()?, built.clone())
524 });
525
526 if let Some(shield) = &shield {
535 router = router.layer(axum::middleware::from_fn_with_state(
536 shield.clone(),
537 shield::post_auth_middleware,
538 ));
539 }
540
541 if let Some(auth) = auth {
542 router = router.layer(axum::middleware::from_fn_with_state(auth, auth::middleware));
543 }
544
545 if let Some(decider) = &self.auth_decider {
549 let decider = decider.clone();
551 let path = self.decider_verify_path();
552 router = router.route(
553 &path,
554 axum::routing::any(move |req: axum::extract::Request| {
555 let decider = decider.clone();
556 async move { embed::verify_via_decider(decider, req).await }
557 }),
558 );
559 } else if let Some(forward_auth) = &forward_auth {
560 router = router.merge(forward_auth.routes());
561 }
562
563 if let Some(shield) = &shield {
569 router = router.layer(axum::middleware::from_fn_with_state(
570 shield.clone(),
571 shield::pre_auth_middleware,
572 ));
573 }
574
575 let router = router
576 .layer(axum::middleware::from_fn_with_state(
577 state.clone(),
578 maintenance_middleware,
579 ))
580 .layer(TraceLayer::new_for_http())
581 .layer(cors)
584 .with_state(state);
585
586 Ok(router)
587 }
588
589 fn build_openapi_routes(&self, pool: &DescriptorPool) -> Router<ProxyState> {
590 let openapi_config = match &self.config.openapi {
591 Some(cfg) if cfg.enabled => cfg,
592 _ => return Router::new(),
593 };
594
595 let spec = openapi::generate(pool, openapi_config, &self.config.aliases);
596 let spec_json = serde_json::to_string_pretty(&spec).unwrap_or_default();
597 let openapi_path = openapi_config.path.clone();
598 let docs_path = openapi_config.docs_path.clone();
599 let title = openapi_config
600 .title
601 .clone()
602 .unwrap_or_else(|| self.config.service.name.clone());
603 let openapi_path_for_docs = openapi_path.clone();
604
605 tracing::info!("OpenAPI spec at {}, docs at {}", openapi_path, docs_path,);
606
607 Router::new()
608 .route(
609 &openapi_path,
610 get(move || async move {
611 (
612 StatusCode::OK,
613 [(
614 axum::http::header::CONTENT_TYPE,
615 "application/json; charset=utf-8",
616 )],
617 spec_json,
618 )
619 }),
620 )
621 .route(
622 &docs_path,
623 get(move || async move {
624 let html = openapi::docs_html(&openapi_path_for_docs, &title);
625 (
626 StatusCode::OK,
627 [(axum::http::header::CONTENT_TYPE, "text/html; charset=utf-8")],
628 html,
629 )
630 }),
631 )
632 }
633
634 fn build_cors(&self) -> CorsLayer {
635 if self.config.cors.origins.is_empty() {
636 tracing::warn!("CORS origins not set — using permissive CORS (dev mode)");
637 CorsLayer::permissive()
638 } else {
639 let origins: Vec<_> = self
640 .config
641 .cors
642 .origins
643 .iter()
644 .filter_map(|o| o.parse().ok())
645 .collect();
646 CorsLayer::new()
647 .allow_origin(AllowOrigin::list(origins))
648 .allow_methods(tower_http::cors::Any)
649 .allow_headers(tower_http::cors::Any)
650 .allow_credentials(true)
651 .expose_headers([
652 "grpc-status".parse().unwrap(),
653 "grpc-message".parse().unwrap(),
654 "ratelimit-limit".parse().unwrap(),
656 "ratelimit-remaining".parse().unwrap(),
657 "ratelimit-reset".parse().unwrap(),
658 "retry-after".parse().unwrap(),
659 ])
660 }
661 }
662
663 pub async fn serve(&self) -> anyhow::Result<()> {
665 let router = self.router()?;
666 let app = router.into_make_service_with_connect_info::<SocketAddr>();
667 let addr: SocketAddr = self.config.listen.http.parse()?;
668 let listener = tokio::net::TcpListener::bind(addr).await?;
669
670 tracing::info!("{} listening on {}", self.config.service.name, addr);
671 axum::serve(listener, app).await?;
672 Ok(())
673 }
674}
675
676fn normalize_route_shape(path: &str) -> String {
682 path.split('/')
683 .map(|seg| {
684 if seg.starts_with("{*") && seg.ends_with('}') {
685 "{*}"
686 } else if seg.starts_with('{') && seg.ends_with('}') {
687 "{}"
688 } else {
689 seg
690 }
691 })
692 .collect::<Vec<_>>()
693 .join("/")
694}
695
696async fn maintenance_middleware(
698 State(state): State<ProxyState>,
699 request: Request<axum::body::Body>,
700 next: Next,
701) -> Response {
702 if state.maintenance_mode {
703 let path = request.uri().path();
704 let exempt = state.maintenance_exempt.iter().any(|pattern| {
705 if pattern.ends_with("/**") {
706 let prefix = &pattern[..pattern.len() - 3];
707 path.starts_with(prefix)
708 } else {
709 path == pattern
710 }
711 });
712 if !exempt {
713 return (
714 StatusCode::SERVICE_UNAVAILABLE,
715 [("retry-after", "300")],
716 state.maintenance_message.clone(),
717 )
718 .into_response();
719 }
720 }
721 next.run(request).await
722}
723
724#[cfg(test)]
726pub(crate) fn test_channel() -> tonic::transport::Channel {
727 tonic::transport::Channel::from_static("http://127.0.0.1:1")
728 .connect_timeout(std::time::Duration::from_millis(100))
729 .connect_lazy()
730}
731
732#[cfg(test)]
735pub(crate) fn test_state() -> ProxyState {
736 ProxyState {
737 service_name: "test".into(),
738 grpc_upstream: "http://127.0.0.1:1".into(),
739 grpc_channel: test_channel(),
740 maintenance_mode: false,
741 maintenance_exempt: vec![],
742 maintenance_message: String::new(),
743 forwarded_headers: vec![],
744 metrics_namespace: "test".into(),
745 metrics_classes: vec![],
746 sse_keep_alive_secs: 15,
747 }
748}
749
750#[cfg(test)]
751mod tests {
752 use super::*;
753
754 #[test]
755 fn normalize_route_shape_collapses_param_names() {
756 assert_eq!(
758 normalize_route_shape("/v1/x/{profile_id}"),
759 normalize_route_shape("/v1/x/{id}")
760 );
761 assert_eq!(normalize_route_shape("/a/{p}/b"), "/a/{}/b");
763 assert_eq!(normalize_route_shape("/a/{*rest}"), "/a/{*}");
764 assert_ne!(
765 normalize_route_shape("/a/{p}"),
766 normalize_route_shape("/a/b")
767 );
768 }
769
770 #[test]
771 fn test_minimal_config_server() {
772 let yaml = r#"
773upstream:
774 default: "http://127.0.0.1:50051"
775"#;
776 let config: ProxyConfig = serde_yaml::from_str(yaml).unwrap();
777 let server = ProxyServer::from_config(config);
778 assert!(server.descriptor_pool.is_none());
779 }
780
781 #[tokio::test]
782 async fn test_maintenance_exempt_matching() {
783 let state = ProxyState {
784 service_name: "test".into(),
785 grpc_upstream: "http://localhost:50051".into(),
786 grpc_channel: test_channel(),
787 maintenance_mode: true,
788 maintenance_exempt: vec![
789 "/health/**".into(),
790 "/.well-known/**".into(),
791 "/metrics".into(),
792 ],
793 maintenance_message: "Down".into(),
794 forwarded_headers: vec![],
795 metrics_namespace: "test".into(),
796 metrics_classes: vec![],
797 sse_keep_alive_secs: 15,
798 };
799
800 let check = |path: &str| -> bool {
801 state.maintenance_exempt.iter().any(|pattern| {
802 if pattern.ends_with("/**") {
803 let prefix = &pattern[..pattern.len() - 3];
804 path.starts_with(prefix)
805 } else {
806 path == pattern
807 }
808 })
809 };
810
811 assert!(check("/health"));
812 assert!(check("/health/ready"));
813 assert!(check("/.well-known/openid-configuration"));
814 assert!(check("/metrics"));
815 assert!(!check("/v1/auth/login"));
816 assert!(!check("/oauth2/token"));
817 }
818}