feat(controller): serve the gateway JWKS off the management listener

README.md

@@ -697,7 +697,7 @@ routes:
       issuer: agentdesktop-controller
       audiences: [agentgateway]
       jwks:
-        url: https://<controller>/.well-known/jwks.json
+        url: https://<controller>/.well-known/jwks.json  # fleet listener or gatewayJwt.jwksListener
     urlRewrite:
       authority:
         full: api.business.githubcopilot.com
@@ -815,6 +815,13 @@ cross-site browser mutations and non-loopback `Host` headers. Reach it remotely
 through an SSH tunnel. Remote management (`adminAccess.mode: authenticated`)
 requires a controller build with a compiled-in authentication provider; the
 standard executable rejects it at startup. Device enrollment is configured separately.
+LLM gateways need the controller's JWKS (`/.well-known/jwks.json`) to verify
+the JWTs it issues. The management listener does not serve it. Gateways fetch
+it from the fleet listener, or from an optional `gatewayJwt.jwksListener` that
+serves nothing else. With the dedicated listener, a NetworkPolicy or firewall
+can allow the gateway that one port and block the management API, so a
+compromised gateway cannot read the device inventory or delete devices. A
+gateway that exports usage to the controller also needs the fleet port.
 See the [controller configuration schema](schema/controller-config.md).
 
 ### Configuration status

crates/controller/Cargo.toml

@@ -27,7 +27,6 @@ rand.workspace = true
 rcgen.workspace = true
 reqwest.workspace = true
 rust-embed = "8"
-rustls.workspace = true
 serde.workspace = true
 serde_json.workspace = true
 sha2.workspace = true

crates/controller/src/application.rs

@@ -65,7 +65,18 @@ pub async fn run(
         tracing::info!("LLM gateway JWT issuance enabled");
     }
     let gateway_jwks = gateway_jwt_issuer.as_ref().map(GatewayJwtIssuer::jwks);
-    let admin_gateway_jwks = gateway_jwks.clone();
+    let jwks_listener = config
+        .gateway_jwt
+        .as_ref()
+        .and_then(|gateway| gateway.jwks_listener.as_ref())
+        .zip(gateway_jwks.clone())
+        .map(|(listener, jwks)| {
+            let tls = listener.tls.then(|| gateway_jwt::JwksTls {
+                certificate: tls.certificate.clone(),
+                key: tls.key.clone(),
+            });
+            gateway_jwt::serve(listener.listen, jwks, tls)
+        });
     let usage_token = config
         .usage
         .as_ref()
@@ -145,13 +156,13 @@ pub async fn run(
             .await
             .context("serve fleet gRPC API")
     };
-    let admin = management::serve(
-        admin_listen,
-        admin_state,
-        admin_gateway_jwks,
-        config.admin_access,
-        extensions,
-    );
-    tokio::try_join!(fleet, admin)?;
+    let admin = management::serve(admin_listen, admin_state, config.admin_access, extensions);
+    let jwks = async move {
+        match jwks_listener {
+            Some(serve) => serve.await,
+            None => std::future::pending().await,
+        }
+    };
+    tokio::try_join!(fleet, admin, jwks)?;
     Ok(())
 }

crates/controller/src/gateway_jwt.rs

@@ -1,7 +1,8 @@
 use std::{
     collections::BTreeMap,
     fs,
-    path::Path,
+    net::SocketAddr,
+    path::{Path, PathBuf},
     sync::Arc,
     time::{Duration, SystemTime, UNIX_EPOCH},
 };
@@ -10,11 +11,13 @@ use anyhow::Context;
 use aws_lc_rs::{rsa::KeyPair as RsaKeyPair, signature::KeyPair as _};
 use axum::{
     Json, Router,
-    extract::State,
+    extract::{Request, State},
     http::{StatusCode, header},
+    middleware::{self, Next},
     response::{IntoResponse, Response},
     routing::get,
 };
+use axum_server::tls_rustls::RustlsConfig;
 use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD};
 use jsonwebtoken::{Algorithm, EncodingKey, Header, encode};
 use serde::Serialize;
@@ -34,13 +37,71 @@ pub struct GatewayJwks {
     keys: Vec<GatewayJwk>,
 }
 
-/// HTTP routes published on the controller's device-facing TLS listener.
+/// The JWKS route, mounted on the fleet listener and the dedicated JWKS listener.
 pub fn routes(jwks: Option<GatewayJwks>) -> Router {
     Router::<Option<GatewayJwks>>::new()
         .route("/.well-known/jwks.json", get(serve_jwks))
         .with_state(jwks)
 }
 
+const JWKS_REQUEST_TIMEOUT: Duration = Duration::from_secs(10);
+
+/// PEM certificate and key for an HTTPS JWKS listener.
+pub struct JwksTls {
+    pub certificate: PathBuf,
+    pub key: PathBuf,
+}
+
+/// Serves only the gateway JWKS, over HTTPS when `tls` is set.
+pub async fn serve(
+    address: SocketAddr,
+    jwks: GatewayJwks,
+    tls: Option<JwksTls>,
+) -> anyhow::Result<()> {
+    let listener = tokio::net::TcpListener::bind(address)
+        .await
+        .with_context(|| format!("bind gateway JWKS listener on {address}"))?;
+    serve_listener(listener, jwks, tls).await
+}
+
+pub async fn serve_listener(
+    listener: tokio::net::TcpListener,
+    jwks: GatewayJwks,
+    tls: Option<JwksTls>,
+) -> anyhow::Result<()> {
+    let address = listener.local_addr()?;
+    let app = routes(Some(jwks)).layer(middleware::from_fn(safeguard_request));
+    match tls {
+        Some(tls) => {
+            let config = RustlsConfig::from_config(Arc::new(
+                agentdesktop_core::tls::server_config(&tls.certificate, &tls.key, "gateway JWKS")?,
+            ));
+            tracing::info!(listen = %address, tls = true, "gateway JWKS listening");
+            axum_server::from_tcp_rustls(listener.into_std()?, config)?
+                .serve(app.into_make_service())
+                .await
+                .context("serve gateway JWKS over HTTPS")
+        }
+        None => {
+            tracing::info!(listen = %address, tls = false, "gateway JWKS listening");
+            axum::serve(listener, app)
+                .await
+                .context("serve gateway JWKS")
+        }
+    }
+}
+
+async fn safeguard_request(request: Request, next: Next) -> Response {
+    let mut response = tokio::time::timeout(JWKS_REQUEST_TIMEOUT, next.run(request))
+        .await
+        .unwrap_or_else(|_| StatusCode::GATEWAY_TIMEOUT.into_response());
+    response.headers_mut().insert(
+        header::X_CONTENT_TYPE_OPTIONS,
+        header::HeaderValue::from_static("nosniff"),
+    );
+    response
+}
+
 async fn serve_jwks(State(jwks): State<Option<GatewayJwks>>) -> Result<Response, StatusCode> {
     jwks.map(|jwks| {
         (
@@ -195,8 +256,7 @@ mod tests {
 
     use super::{Actor, Claims, GatewayJwtIssuer};
 
-    #[test]
-    fn published_jwks_verifies_issued_credentials() {
+    fn test_issuer() -> GatewayJwtIssuer {
         let key_pair = KeyPair::generate(KeySize::Rsa2048).expect("generate RSA key");
         let der: Pkcs8V1Der<'static> = key_pair.as_der().expect("encode private key");
         let pem = pem_rfc7468::encode_string("PRIVATE KEY", LineEnding::LF, der.as_ref())
@@ -207,14 +267,106 @@ mod tests {
             rand::random::<u64>()
         ));
         fs::write(&path, pem).expect("write test private key");
-
         let issuer = GatewayJwtIssuer::from_rsa_pem(
             &path,
             "agentdesktop-controller".to_owned(),
             "agentdesktop".to_owned(),
             std::time::Duration::from_secs(300),
         )
         .expect("create issuer");
+        let _ = fs::remove_file(path);
+        issuer
+    }
+
+    async fn get_with_retry(client: &reqwest::Client, url: &str) -> reqwest::Response {
+        for _ in 0..50 {
+            match client.get(url).send().await {
+                Ok(response) => return response,
+                Err(error) if error.is_connect() => {
+                    tokio::time::sleep(std::time::Duration::from_millis(10)).await;
+                }
+                Err(error) => panic!("JWKS probe failed: {error}"),
+            }
+        }
+        panic!("JWKS listener did not start");
+    }
+
+    async fn assert_serves_jwks(client: &reqwest::Client, origin: &str) {
+        let response = get_with_retry(client, &format!("{origin}/.well-known/jwks.json")).await;
+        assert_eq!(response.status(), reqwest::StatusCode::OK);
+        assert_eq!(
+            response.headers()["content-type"],
+            "application/jwk-set+json"
+        );
+        assert_eq!(response.headers()["cache-control"], "public, max-age=300");
+        assert_eq!(response.headers()["x-content-type-options"], "nosniff");
+        let jwks: JwkSet = response.json().await.expect("decode JWKS");
+        assert!(jwks.find("agentdesktop").is_some());
+
+        let other = client
+            .get(format!("{origin}/api/v1/devices"))
+            .send()
+            .await
+            .expect("probe other path");
+        assert_eq!(other.status(), reqwest::StatusCode::NOT_FOUND);
+    }
+
+    #[tokio::test]
+    async fn dedicated_listener_serves_only_jwks_over_http() {
+        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
+        let address = listener.local_addr().unwrap();
+        let server = tokio::spawn(super::serve_listener(listener, test_issuer().jwks(), None));
+        let client = reqwest::Client::builder()
+            .timeout(std::time::Duration::from_secs(5))
+            .build()
+            .unwrap();
+        assert_serves_jwks(&client, &format!("http://{address}")).await;
+        server.abort();
+    }
+
+    #[tokio::test]
+    async fn dedicated_listener_serves_jwks_over_https_when_configured() {
+        let directory = std::env::temp_dir().join(format!(
+            "agentdesktop-jwks-tls-{}-{}",
+            std::process::id(),
+            rand::random::<u64>()
+        ));
+        fs::create_dir(&directory).unwrap();
+        let rcgen::CertifiedKey { cert, signing_key } =
+            rcgen::generate_simple_self_signed(vec!["127.0.0.1".into()]).unwrap();
+        let certificate = directory.join("controller.pem");
+        let key = directory.join("controller-key.pem");
+        fs::write(&certificate, cert.pem()).unwrap();
+        fs::write(&key, signing_key.serialize_pem()).unwrap();
+
+        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
+        let address = listener.local_addr().unwrap();
+        let server = tokio::spawn(super::serve_listener(
+            listener,
+            test_issuer().jwks(),
+            Some(super::JwksTls { certificate, key }),
+        ));
+        let client = reqwest::Client::builder()
+            .add_root_certificate(reqwest::Certificate::from_pem(cert.pem().as_bytes()).unwrap())
+            .timeout(std::time::Duration::from_secs(5))
+            .build()
+            .unwrap();
+        assert_serves_jwks(&client, &format!("https://{address}")).await;
+        assert!(
+            client
+                .get(format!("http://{address}/.well-known/jwks.json"))
+                .send()
+                .await
+                .map_or(true, |response| !response.status().is_success()),
+            "plaintext must not reach the TLS listener"
+        );
+        server.abort();
+        fs::remove_dir_all(directory).unwrap();
+    }
+
+    #[test]
+    fn published_jwks_verifies_issued_credentials() {
+        let issuer = test_issuer();
         let (credential, _) = issuer
             .issue("user", "device", "codex", "agentgateway", None, false)
             .expect("issue credential");
@@ -231,8 +383,6 @@ mod tests {
         validation.set_audience(&["agentgateway"]);
         decode::<serde_json::Value>(&credential, &decoding_key, &validation)
             .expect("verify credential with published JWKS");
-
-        let _ = fs::remove_file(path);
     }
 
     #[test]