first pass at token-specific TTL (100% claude code)

Author: david-crespo

nexus/db-model/src/device_auth.rs

@@ -32,6 +32,7 @@ pub struct DeviceAuthRequest {
     pub user_code: String,
     pub time_created: DateTime<Utc>,
     pub time_expires: DateTime<Utc>,
+    pub requested_ttl_seconds: Option<i64>,
 }
 
 impl DeviceAuthRequest {
@@ -98,7 +99,7 @@ fn generate_user_code() -> String {
 }
 
 impl DeviceAuthRequest {
-    pub fn new(client_id: Uuid) -> Self {
+    pub fn new(client_id: Uuid, requested_ttl_seconds: Option<u32>) -> Self {
         let now = Utc::now();
         Self {
             client_id,
@@ -107,6 +108,7 @@ impl DeviceAuthRequest {
             time_created: now,
             time_expires: now
                 + Duration::seconds(CLIENT_AUTHENTICATION_TIMEOUT),
+            requested_ttl_seconds: requested_ttl_seconds.map(|ttl| ttl as i64),
         }
     }
 

nexus/db-schema/src/schema.rs

@@ -1360,6 +1360,7 @@ table! {
         device_code -> Text,
         time_created -> Timestamptz,
         time_expires -> Timestamptz,
+        requested_ttl_seconds -> Nullable<Int8>,
     }
 }
 

nexus/src/app/device_auth.rs

@@ -78,12 +78,14 @@ impl super::Nexus {
         &self,
         opctx: &OpContext,
         client_id: Uuid,
+        requested_ttl_seconds: Option<u32>,
     ) -> CreateResult<DeviceAuthRequest> {
         // TODO-correctness: the `user_code` generated for a new request
         // is used as a primary key, but may potentially collide with an
         // existing outstanding request. So we should retry some (small)
         // number of times if inserting the new request fails.
-        let auth_request = DeviceAuthRequest::new(client_id);
+        let auth_request =
+            DeviceAuthRequest::new(client_id, requested_ttl_seconds);
         self.db_datastore.device_auth_request_create(opctx, auth_request).await
     }
 
@@ -115,17 +117,43 @@ impl super::Nexus {
             .silo_auth_settings_view(opctx, &authz_silo)
             .await?;
 
+        // Validate the requested TTL against the silo's max TTL
+        if let Some(requested_ttl) = db_request.requested_ttl_seconds {
+            if let Some(max_ttl) = silo_auth_settings.device_token_max_ttl_seconds {
+                if requested_ttl > max_ttl {
+                    return Err(Error::invalid_request(&format!(
+                        "Requested TTL {} exceeds maximum allowed TTL {} for this silo",
+                        requested_ttl, max_ttl
+                    )));
+                }
+            }
+        }
+
         // Create an access token record.
+        let token_ttl = match db_request.requested_ttl_seconds {
+            Some(requested_ttl) => {
+                // Use the requested TTL, but cap it at the silo max if there is one
+                let effective_ttl =
+                    match silo_auth_settings.device_token_max_ttl_seconds {
+                        Some(max_ttl) => std::cmp::min(requested_ttl, max_ttl),
+                        None => requested_ttl,
+                    };
+                Some(Utc::now() + Duration::seconds(effective_ttl))
+            }
+            None => {
+                // Use the silo max TTL if no specific TTL was requested
+                silo_auth_settings
+                    .device_token_max_ttl_seconds
+                    .map(|ttl| Utc::now() + Duration::seconds(ttl))
+            }
+        };
+
         let token = DeviceAccessToken::new(
             db_request.client_id,
             db_request.device_code,
             db_request.time_created,
             silo_user_id,
-            // Token gets the max TTL for the silo (if there is one) until we
-            // build a way for the user to ask for a different TTL
-            silo_auth_settings
-                .device_token_max_ttl_seconds
-                .map(|ttl| Utc::now() + Duration::seconds(ttl)),
+            token_ttl,
         );
 
         if db_request.time_expires < Utc::now() {

nexus/src/external_api/http_entrypoints.rs

@@ -7855,7 +7855,7 @@ impl NexusExternalApi for NexusExternalApiImpl {
             };
 
             let model = nexus
-                .device_auth_request_create(&opctx, params.client_id)
+                .device_auth_request_create(&opctx, params.client_id, params.ttl_seconds)
                 .await?;
             nexus.build_oauth_response(
                 StatusCode::OK,

nexus/tests/integration_tests/device_auth.rs

@@ -55,7 +55,7 @@ async fn test_device_auth_flow(cptestctx: &ControlPlaneTestContext) {
         .expect("failed to reject device auth start without client_id");
 
     let client_id = Uuid::new_v4();
-    let authn_params = DeviceAuthRequest { client_id };
+    let authn_params = DeviceAuthRequest { client_id, ttl_seconds: None };
 
     // Using a JSON encoded body fails.
     RequestBuilder::new(testctx, Method::POST, "/device/auth")
@@ -235,7 +235,7 @@ async fn test_device_auth_flow(cptestctx: &ControlPlaneTestContext) {
 /// as a string
 async fn get_device_token(testctx: &ClientTestContext) -> String {
     let client_id = Uuid::new_v4();
-    let authn_params = DeviceAuthRequest { client_id };
+    let authn_params = DeviceAuthRequest { client_id, ttl_seconds: None };
 
     // Start a device authentication flow
     let auth_response: DeviceAuthResponse =
@@ -421,6 +421,133 @@ async fn test_device_token_expiration(cptestctx: &ControlPlaneTestContext) {
     assert_eq!(settings.device_token_max_ttl_seconds, None);
 }
 
+#[nexus_test]
+async fn test_device_token_request_ttl(cptestctx: &ControlPlaneTestContext) {
+    let testctx = &cptestctx.external_client;
+
+    // Set silo max TTL to 10 seconds
+    let _settings: views::SiloSettings = object_put(
+        testctx,
+        "/v1/settings",
+        &params::SiloSettingsUpdate {
+            device_token_max_ttl_seconds: NonZeroU32::new(10),
+        },
+    )
+    .await;
+
+    let client_id = Uuid::new_v4();
+
+    // Test 1: Request TTL above the max should fail at verification time
+    let authn_params_invalid = DeviceAuthRequest { 
+        client_id, 
+        ttl_seconds: Some(20) // Above the 10 second max
+    };
+
+    let auth_response: DeviceAuthResponse =
+        RequestBuilder::new(testctx, Method::POST, "/device/auth")
+            .allow_non_dropshot_errors()
+            .body_urlencoded(Some(&authn_params_invalid))
+            .expect_status(Some(StatusCode::OK))
+            .execute()
+            .await
+            .expect("failed to start client authentication flow")
+            .parsed_body()
+            .expect("client authentication response");
+
+    let confirm_params = DeviceAuthVerify { user_code: auth_response.user_code };
+
+    // Confirmation should fail because requested TTL exceeds max
+    let confirm_response = NexusRequest::new(
+        RequestBuilder::new(testctx, Method::POST, "/device/confirm")
+            .body(Some(&confirm_params))
+            .expect_status(Some(StatusCode::BAD_REQUEST)),
+    )
+    .authn_as(AuthnMode::PrivilegedUser)
+    .execute()
+    .await
+    .expect("confirmation should fail for TTL above max");
+
+    // Check that the error message mentions TTL
+    let error_body = String::from_utf8_lossy(&confirm_response.body);
+    assert!(error_body.contains("TTL") || error_body.contains("ttl"));
+
+    // Test 2: Request TTL below the max should succeed and be used
+    let authn_params_valid = DeviceAuthRequest { 
+        client_id: Uuid::new_v4(), // New client ID for new flow
+        ttl_seconds: Some(5) // Below the 10 second max
+    };
+
+    let auth_response: DeviceAuthResponse =
+        RequestBuilder::new(testctx, Method::POST, "/device/auth")
+            .allow_non_dropshot_errors()
+            .body_urlencoded(Some(&authn_params_valid))
+            .expect_status(Some(StatusCode::OK))
+            .execute()
+            .await
+            .expect("failed to start client authentication flow")
+            .parsed_body()
+            .expect("client authentication response");
+
+    let device_code = auth_response.device_code;
+    let user_code = auth_response.user_code;
+    let confirm_params = DeviceAuthVerify { user_code };
+
+    // Confirmation should succeed
+    NexusRequest::new(
+        RequestBuilder::new(testctx, Method::POST, "/device/confirm")
+            .body(Some(&confirm_params))
+            .expect_status(Some(StatusCode::NO_CONTENT)),
+    )
+    .authn_as(AuthnMode::PrivilegedUser)
+    .execute()
+    .await
+    .expect("failed to confirm");
+
+    let token_params = DeviceAccessTokenRequest {
+        grant_type: "urn:ietf:params:oauth:grant-type:device_code".to_string(),
+        device_code,
+        client_id: authn_params_valid.client_id,
+    };
+
+    // Get the token
+    let token: DeviceAccessTokenGrant = NexusRequest::new(
+        RequestBuilder::new(testctx, Method::POST, "/device/token")
+            .allow_non_dropshot_errors()
+            .body_urlencoded(Some(&token_params))
+            .expect_status(Some(StatusCode::OK)),
+    )
+    .authn_as(AuthnMode::PrivilegedUser)
+    .execute()
+    .await
+    .expect("failed to get token")
+    .parsed_body()
+    .expect("failed to deserialize token response");
+
+    // Verify the token has the correct expiration (5 seconds from now)
+    let tokens = get_tokens_priv(testctx).await;
+    let our_token = tokens.iter().find(|t| t.time_expires.is_some()).unwrap();
+    let expires_at = our_token.time_expires.unwrap();
+    let now = Utc::now();
+    
+    // Should expire approximately 5 seconds from now (allow some tolerance for test timing)
+    let expected_expiry = now + chrono::Duration::seconds(5);
+    let time_diff = (expires_at - expected_expiry).num_seconds().abs();
+    assert!(time_diff <= 2, "Token expiry should be close to requested TTL");
+
+    // Token should work initially
+    project_list(&testctx, &token.access_token, StatusCode::OK)
+        .await
+        .expect("token should work initially");
+
+    // Wait for token to expire
+    sleep(Duration::from_secs(6)).await;
+
+    // Token should be expired now
+    project_list(&testctx, &token.access_token, StatusCode::UNAUTHORIZED)
+        .await
+        .expect("token should be expired");
+}
+
 async fn get_tokens_priv(
     testctx: &ClientTestContext,
 ) -> Vec<views::DeviceAccessToken> {

nexus/types/src/external_api/params.rs

@@ -2382,6 +2382,8 @@ impl TryFrom<String> for RelativeUri {
 #[derive(Clone, Debug, Deserialize, Serialize, JsonSchema)]
 pub struct DeviceAuthRequest {
     pub client_id: Uuid,
+    /// Optional TTL for the access token in seconds. If not specified, the silo's max TTL will be used.
+    pub ttl_seconds: Option<u32>,
 }
 
 #[derive(Clone, Debug, Deserialize, Serialize, JsonSchema)]

schema/crdb/dbinit.sql

@@ -2813,7 +2813,9 @@ CREATE TABLE IF NOT EXISTS omicron.public.device_auth_request (
     client_id UUID NOT NULL,
     device_code STRING(40) NOT NULL,
     time_created TIMESTAMPTZ NOT NULL,
-    time_expires TIMESTAMPTZ NOT NULL
+    time_expires TIMESTAMPTZ NOT NULL,
+    -- requested TTL for the token in seconds (if specified by the user)
+    requested_ttl_seconds INT8 CHECK (requested_ttl_seconds > 0)
 );
 
 -- Access tokens granted in response to successful device authorization flows.

schema/crdb/device-auth-request-ttl/up01.sql

@@ -0,0 +1,2 @@
+ALTER TABLE omicron.public.device_auth_request
+  ADD COLUMN IF NOT EXISTS requested_ttl_seconds INT8 CHECK (requested_ttl_seconds > 0);