From 0cae4fd05cdfda4e4860cef4ace904566885835a Mon Sep 17 00:00:00 2001 From: Hare Date: Sun, 30 Aug 2026 12:11:26 +0900 Subject: [PATCH] fix: retry device login code collisions --- crates/workspace-server/src/auth.rs | 33 +++++++- crates/workspace-server/src/server.rs | 104 ++++++++++++++++++++------ crates/workspace-server/src/store.rs | 73 ++++++++++++++++-- 3 files changed, 178 insertions(+), 32 deletions(-) diff --git a/crates/workspace-server/src/auth.rs b/crates/workspace-server/src/auth.rs index ae43d57a..c014e37a 100644 --- a/crates/workspace-server/src/auth.rs +++ b/crates/workspace-server/src/auth.rs @@ -105,8 +105,13 @@ pub fn token_hash(token: &str) -> String { } pub fn new_user_code() -> String { - let hex = Uuid::now_v7().simple().to_string().to_ascii_uppercase(); - format!("{}-{}", &hex[0..4], &hex[4..8]) + user_code_from_uuid(Uuid::now_v7()) +} + +fn user_code_from_uuid(id: Uuid) -> String { + let hex = id.simple().to_string().to_ascii_uppercase(); + let random_tail = &hex[24..32]; + format!("{}-{}", &random_tail[..4], &random_tail[4..]) } pub fn parse_bearer(headers: &HeaderMap) -> Option { @@ -207,3 +212,27 @@ pub fn auth_error(code: &str, message: &str) -> Error { message: message.to_string(), } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn user_code_uses_the_uuid_random_tail() { + let id = Uuid::parse_str("01234567-89ab-7cde-8123-456789abcdef").expect("UUID"); + assert_eq!(user_code_from_uuid(id), "89AB-CDEF"); + } + + #[test] + fn user_code_preserves_the_human_readable_format() { + let code = new_user_code(); + assert_eq!(code.len(), 9); + assert_eq!(code.as_bytes()[4], b'-'); + assert!( + code.chars() + .enumerate() + .all(|(index, value)| index == 4 || value.is_ascii_hexdigit()) + ); + assert_eq!(code, code.to_ascii_uppercase()); + } +} diff --git a/crates/workspace-server/src/server.rs b/crates/workspace-server/src/server.rs index 33e34c47..2759e670 100644 --- a/crates/workspace-server/src/server.rs +++ b/crates/workspace-server/src/server.rs @@ -10576,35 +10576,48 @@ async fn post_device_login_start( State(api): State, Json(request): Json, ) -> ApiResult> { + const MAX_CODE_ALLOCATION_ATTEMPTS: usize = 8; + let auth = auth_public_config(&api.config); - let device_code = mint_secret("yoi_device"); - let user_code = new_user_code(); let verification_uri = format!( "{}/login/device", auth.public_base_url.trim_end_matches('/') ); - let verification_uri_complete = format!("{verification_uri}?user_code={user_code}"); - api.store.create_device_login_flow(&DeviceLoginFlowRecord { - device_code: device_code.clone(), - user_code: user_code.clone(), - verification_uri: verification_uri.clone(), - client_name: request.client_name, - user_id: None, - api_token_id: None, - issued_access_token: None, - created_at: crate::auth::now_rfc3339(), - expires_at: rfc3339_after(Duration::minutes(10)), - approved_at: None, - consumed_at: None, - })?; - Ok(Json(DeviceLoginStartResponse { - device_code, - user_code, - verification_uri, - verification_uri_complete, - expires_in: 600, - interval: 5, - })) + for _ in 0..MAX_CODE_ALLOCATION_ATTEMPTS { + let device_code = mint_secret("yoi_device"); + let user_code = new_user_code(); + let verification_uri_complete = format!("{verification_uri}?user_code={user_code}"); + let flow = DeviceLoginFlowRecord { + device_code: device_code.clone(), + user_code: user_code.clone(), + verification_uri: verification_uri.clone(), + client_name: request.client_name.clone(), + user_id: None, + api_token_id: None, + issued_access_token: None, + created_at: crate::auth::now_rfc3339(), + expires_at: rfc3339_after(Duration::minutes(10)), + approved_at: None, + consumed_at: None, + }; + if !api.store.try_create_device_login_flow(&flow)? { + continue; + } + return Ok(Json(DeviceLoginStartResponse { + device_code, + user_code, + verification_uri, + verification_uri_complete, + expires_in: 600, + interval: 5, + })); + } + + Err(auth_error( + "device_login_code_allocation_failed", + "could not allocate a unique device login user code", + ) + .into()) } async fn post_device_login_approve( @@ -17762,6 +17775,49 @@ mod tests { assert_ne!(csrf_accepted.status(), StatusCode::FORBIDDEN); } + #[tokio::test] + async fn device_login_start_can_retry_while_previous_flow_is_pending() { + let workspace = tempfile::tempdir().unwrap(); + let api = test_api(workspace.path()).await; + let app = build_router(api.clone()); + let mut starts = Vec::new(); + + for _ in 0..2 { + let response = app + .clone() + .oneshot( + Request::builder() + .method(Method::POST) + .uri("/api/auth/device-login/start") + .header(CONTENT_TYPE, "application/json") + .body(Body::from(r#"{"client_name":"retry-test"}"#)) + .unwrap(), + ) + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + let body = to_bytes(response.into_body(), usize::MAX).await.unwrap(); + starts.push( + serde_json::from_slice::(&body) + .expect("device login start response"), + ); + } + + assert_ne!(starts[0]["device_code"], starts[1]["device_code"]); + assert_ne!(starts[0]["user_code"], starts[1]["user_code"]); + for start in starts { + let device_code = start["device_code"].as_str().expect("device code"); + let user_code = start["user_code"].as_str().expect("user code"); + assert_eq!(user_code.len(), 9); + assert!( + api.store + .get_device_login_flow_by_device_code(device_code) + .expect("stored device login flow") + .is_some() + ); + } + } + #[tokio::test] async fn direct_workspace_router_enforces_origin_on_browser_auth_mutations() { let workspace = tempfile::tempdir().unwrap(); diff --git a/crates/workspace-server/src/store.rs b/crates/workspace-server/src/store.rs index 09f836a0..87717309 100644 --- a/crates/workspace-server/src/store.rs +++ b/crates/workspace-server/src/store.rs @@ -996,7 +996,7 @@ pub trait ControlPlaneStore: Send + Sync { fn create_api_token(&self, record: &ApiTokenRecord) -> Result<()>; fn resolve_api_token(&self, token_hash: &str) -> Result>; fn mark_api_token_used(&self, token_hash: &str, used_at: &str) -> Result<()>; - fn create_device_login_flow(&self, record: &DeviceLoginFlowRecord) -> Result<()>; + fn try_create_device_login_flow(&self, record: &DeviceLoginFlowRecord) -> Result; fn get_device_login_flow_by_user_code( &self, user_code: &str, @@ -3123,14 +3123,15 @@ impl ControlPlaneStore for SqliteWorkspaceStore { }) } - fn create_device_login_flow(&self, record: &DeviceLoginFlowRecord) -> Result<()> { + fn try_create_device_login_flow(&self, record: &DeviceLoginFlowRecord) -> Result { self.with_conn(|conn| { - conn.execute( + let changed = conn.execute( r#"INSERT INTO device_login_flows (device_code, user_code, verification_uri, client_name, user_id, api_token_id, issued_access_token, created_at, expires_at, approved_at, consumed_at) - VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11)"#, + VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11) + ON CONFLICT(user_code) DO NOTHING"#, params![record.device_code, record.user_code, record.verification_uri, record.client_name, record.user_id, record.api_token_id, record.issued_access_token, record.created_at, record.expires_at, record.approved_at, record.consumed_at], )?; - Ok(()) + Ok(changed == 1) }) } @@ -13524,7 +13525,67 @@ CREATE TABLE ticket_assignment_operations ( approved_at: None, consumed_at: None, }; - store.create_device_login_flow(&flow).unwrap(); + assert!(store.try_create_device_login_flow(&flow).unwrap()); + let conflicting_flow = DeviceLoginFlowRecord { + device_code: "device-2".into(), + user_code: flow.user_code.clone(), + verification_uri: flow.verification_uri.clone(), + client_name: Some("other-cli".into()), + user_id: None, + api_token_id: None, + issued_access_token: None, + created_at: "2026-08-22T00:00:01Z".into(), + expires_at: "2026-08-22T00:10:01Z".into(), + approved_at: None, + consumed_at: None, + }; + assert!( + !store + .try_create_device_login_flow(&conflicting_flow) + .expect("user code collision") + ); + assert!( + store + .get_device_login_flow_by_device_code("device-2") + .expect("read conflicting flow") + .is_none() + ); + let fresh_flow = DeviceLoginFlowRecord { + device_code: "device-3".into(), + user_code: "DCBA-5678".into(), + verification_uri: flow.verification_uri.clone(), + client_name: Some("retry-cli".into()), + user_id: None, + api_token_id: None, + issued_access_token: None, + created_at: "2026-08-22T00:00:02Z".into(), + expires_at: "2026-08-22T00:10:02Z".into(), + approved_at: None, + consumed_at: None, + }; + assert!( + store + .try_create_device_login_flow(&fresh_flow) + .expect("fresh user code") + ); + let device_code_conflict = DeviceLoginFlowRecord { + device_code: flow.device_code.clone(), + user_code: "ABCD-9999".into(), + verification_uri: flow.verification_uri.clone(), + client_name: Some("invalid-cli".into()), + user_id: None, + api_token_id: None, + issued_access_token: None, + created_at: "2026-08-22T00:00:03Z".into(), + expires_at: "2026-08-22T00:10:03Z".into(), + approved_at: None, + consumed_at: None, + }; + assert!( + store + .try_create_device_login_flow(&device_code_conflict) + .is_err() + ); store .approve_device_login_flow( "device",