1use std::time::{Duration, SystemTime, UNIX_EPOCH};
18
19use base64::Engine as _;
20use base64::engine::general_purpose::URL_SAFE_NO_PAD;
21use serde::{Deserialize, Serialize};
22use sha2::{Digest, Sha256};
23use zenwave::{Client as _, client};
24
25use crate::constant::{
26 AUTHORIZE_URL, CLIENT_ID, DEVICE_CODE_URL, DEVICE_TOKEN_URL, OAUTH_PORT, OAUTH_SCOPES,
27 REDIRECT_URI, TOKEN_URL,
28};
29use crate::error::CodexError;
30
31const CALLBACK_BUF: usize = 8192;
33
34#[derive(Debug)]
36pub struct Pkce {
37 verifier: String,
38 challenge: String,
39 state: String,
40}
41
42impl Pkce {
43 #[must_use]
49 pub fn generate() -> Self {
50 let mut verifier_bytes = [0u8; 32];
51 let mut state_bytes = [0u8; 16];
52 getrandom::fill(&mut verifier_bytes).expect("secure random source unavailable");
53 getrandom::fill(&mut state_bytes).expect("secure random source unavailable");
54
55 let verifier = URL_SAFE_NO_PAD.encode(verifier_bytes);
56 let challenge = URL_SAFE_NO_PAD.encode(Sha256::digest(verifier.as_bytes()));
57 let state = URL_SAFE_NO_PAD.encode(state_bytes);
58
59 Self {
60 verifier,
61 challenge,
62 state,
63 }
64 }
65
66 #[must_use]
68 pub fn verifier(&self) -> &str {
69 &self.verifier
70 }
71
72 #[must_use]
74 pub fn state(&self) -> &str {
75 &self.state
76 }
77
78 #[must_use]
85 pub fn authorize_url(&self) -> String {
86 let mut url = url::Url::parse(AUTHORIZE_URL).expect("constant URL parses");
87 url.query_pairs_mut()
88 .append_pair("response_type", "code")
89 .append_pair("client_id", CLIENT_ID)
90 .append_pair("redirect_uri", REDIRECT_URI)
91 .append_pair("scope", OAUTH_SCOPES)
92 .append_pair("code_challenge", &self.challenge)
93 .append_pair("code_challenge_method", "S256")
94 .append_pair("state", &self.state)
95 .append_pair("codex_cli_simplified_flow", "true");
96 url.into()
97 }
98}
99
100#[cfg(not(target_arch = "wasm32"))]
114pub async fn listen_for_callback(expected_state: &str) -> Result<String, CodexError> {
115 use futures_lite::{AsyncReadExt, AsyncWriteExt};
116
117 let addr = std::net::SocketAddr::from(([127, 0, 0, 1], OAUTH_PORT));
118 let listener = async_io::Async::<std::net::TcpListener>::bind(addr)?;
119 tracing::info!(port = OAUTH_PORT, "waiting for OAuth callback");
120
121 let (mut stream, _) = listener.accept().await?;
122 let mut buf = vec![0u8; CALLBACK_BUF];
123 let n = stream.read(&mut buf).await?;
124 let request = String::from_utf8_lossy(&buf[..n]);
125
126 let (code, error) = match parse_callback_request(&request, expected_state) {
127 Ok(code) => (Some(code), None),
128 Err(err) => (None, Some(err.to_string())),
129 };
130
131 let (status, body) = code.as_ref().map_or(
132 (
133 "400 Bad Request",
134 "<html><body><h1>Login failed</h1><p>Return to the terminal for details.</p></body></html>",
135 ),
136 |_| {
137 (
138 "200 OK",
139 "<html><body><h1>Logged in</h1><p>You may close this tab and return to May.</p></body></html>",
140 )
141 },
142 );
143 let response = format!(
144 "HTTP/1.1 {status}\r\ncontent-type: text/html\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{body}",
145 body.len()
146 );
147 let _ = stream.write_all(response.as_bytes()).await;
148 let _ = stream.close().await;
149
150 match (code, error) {
151 (Some(code), _) => Ok(code),
152 (None, Some(err)) => Err(CodexError::Callback(err)),
153 _ => Err(CodexError::Callback("empty callback".into())),
154 }
155}
156
157fn parse_callback_request(request: &str, expected_state: &str) -> Result<String, CodexError> {
159 let path = request
160 .split_whitespace()
161 .nth(1)
162 .ok_or_else(|| CodexError::Callback("malformed HTTP request".into()))?;
163 let url = url::Url::parse(&format!("http://localhost{path}"))
164 .map_err(|e| CodexError::Callback(format!("malformed callback path: {e}")))?;
165
166 let mut code = None;
167 let mut state = None;
168 let mut error = None;
169 for (k, v) in url.query_pairs() {
170 match k.as_ref() {
171 "code" => code = Some(v.into_owned()),
172 "state" => state = Some(v.into_owned()),
173 "error" => error = Some(v.into_owned()),
174 _ => {}
175 }
176 }
177
178 if let Some(error) = error {
179 return Err(CodexError::Callback(format!(
180 "authorization failed: {error}"
181 )));
182 }
183 if state.as_deref() != Some(expected_state) {
184 return Err(CodexError::StateMismatch);
185 }
186 code.ok_or_else(|| CodexError::Callback("callback missing code".into()))
187}
188
189pub async fn exchange_code(code: &str, verifier: &str) -> Result<CodexCredentials, CodexError> {
196 #[derive(Serialize)]
197 struct ExchangeRequest<'a> {
198 grant_type: &'a str,
199 client_id: &'a str,
200 code: &'a str,
201 code_verifier: &'a str,
202 redirect_uri: &'a str,
203 }
204
205 let mut backend = client();
206 let response: TokenResponse = backend
207 .post(TOKEN_URL)?
208 .json_body(&ExchangeRequest {
209 grant_type: "authorization_code",
210 client_id: CLIENT_ID,
211 code,
212 code_verifier: verifier,
213 redirect_uri: REDIRECT_URI,
214 })?
215 .json()
216 .await?;
217 Ok(response.into_credentials())
218}
219
220pub async fn refresh_tokens(refresh_token: &str) -> Result<CodexCredentials, CodexError> {
229 #[derive(Serialize)]
230 struct RefreshRequest<'a> {
231 grant_type: &'a str,
232 refresh_token: &'a str,
233 client_id: &'a str,
234 }
235
236 let mut backend = client();
237 let response: TokenResponse = backend
238 .post(TOKEN_URL)?
239 .json_body(&RefreshRequest {
240 grant_type: "refresh_token",
241 refresh_token,
242 client_id: CLIENT_ID,
243 })?
244 .json()
245 .await?;
246 Ok(response.into_credentials())
247}
248
249pub async fn request_device_code() -> Result<DeviceCodeResponse, CodexError> {
256 #[derive(Serialize)]
257 struct DeviceCodeRequest<'a> {
258 client_id: &'a str,
259 }
260
261 let mut backend = client();
262 let response: DeviceCodeResponse = backend
263 .post(DEVICE_CODE_URL)?
264 .json_body(&DeviceCodeRequest {
265 client_id: CLIENT_ID,
266 })?
267 .json()
268 .await?;
269 Ok(response)
270}
271
272pub async fn poll_device_code(
283 device_auth_id: &str,
284 user_code: &str,
285) -> Result<(String, String), CodexError> {
286 #[derive(Serialize)]
287 struct PollRequest<'a> {
288 device_auth_id: &'a str,
289 user_code: &'a str,
290 }
291 #[derive(Deserialize)]
292 struct PollResponse {
293 authorization_code: String,
294 code_verifier: String,
295 }
296
297 let mut backend = client();
299 let result: Result<PollResponse, zenwave::Error> = backend
300 .post(DEVICE_TOKEN_URL)?
301 .json_body(&PollRequest {
302 device_auth_id,
303 user_code,
304 })?
305 .json()
306 .await;
307
308 match result {
309 Ok(poll) => Ok((poll.authorization_code, poll.code_verifier)),
310 Err(zenwave::Error::Http { status, .. })
311 if status.as_u16() == 403 || status.as_u16() == 404 =>
312 {
313 Err(CodexError::AuthorizationPending)
314 }
315 Err(err) => Err(CodexError::Http(err)),
316 }
317}
318
319#[derive(Debug, Clone, Serialize, Deserialize)]
324pub struct CodexCredentials {
325 pub access_token: String,
327 pub refresh_token: String,
329 #[serde(default, skip_serializing_if = "Option::is_none")]
331 pub id_token: Option<String>,
332 #[serde(default, skip_serializing_if = "Option::is_none")]
334 pub account_id: Option<String>,
335 #[serde(default, skip_serializing_if = "Option::is_none")]
337 pub residency: Option<String>,
338 #[serde(default, skip_serializing_if = "Option::is_none")]
340 pub expires_at: Option<u64>,
341 #[serde(default, skip_serializing_if = "Option::is_none")]
343 pub earliest_refresh_at: Option<u64>,
344}
345
346impl CodexCredentials {
347 #[must_use]
350 pub fn new(access_token: impl Into<String>, refresh_token: impl Into<String>) -> Self {
351 let mut credentials = Self {
352 access_token: access_token.into(),
353 refresh_token: refresh_token.into(),
354 id_token: None,
355 account_id: None,
356 residency: None,
357 expires_at: None,
358 earliest_refresh_at: None,
359 };
360 credentials.fill_claims();
361 credentials
362 }
363
364 #[must_use]
366 pub fn with_id_token(mut self, id_token: impl Into<String>) -> Self {
367 self.id_token = Some(id_token.into());
368 self.fill_claims();
369 self
370 }
371
372 #[must_use]
374 pub const fn with_expires_at(mut self, expires_at: u64) -> Self {
375 self.expires_at = Some(expires_at);
376 self
377 }
378
379 #[must_use]
381 pub fn is_expired(&self, skew: Duration) -> bool {
382 self.expires_at.is_some_and(|at| {
383 let now = SystemTime::now()
384 .duration_since(UNIX_EPOCH)
385 .map_or(0, |d| d.as_secs());
386 now.saturating_add(skew.as_secs()) >= at
387 })
388 }
389}
390
391#[derive(Debug, Deserialize)]
393pub struct DeviceCodeResponse {
394 pub device_auth_id: String,
396 pub user_code: String,
398 #[serde(default)]
400 pub interval: Option<String>,
401}
402
403#[derive(Debug, Deserialize)]
405struct TokenResponse {
406 access_token: String,
407 refresh_token: String,
408 #[serde(default)]
409 id_token: Option<String>,
410 #[serde(default)]
411 expires_in: Option<u64>,
412 #[serde(default)]
414 earliest_refresh_at: Option<serde_json::Value>,
415}
416
417impl TokenResponse {
418 fn into_credentials(self) -> CodexCredentials {
419 let now = SystemTime::now()
420 .duration_since(UNIX_EPOCH)
421 .map_or(0, |d| d.as_secs());
422 let expires_at = self.expires_in.map(|secs| now + secs);
423 let earliest_refresh_at = self
424 .earliest_refresh_at
425 .as_ref()
426 .and_then(serde_json::Value::as_u64)
427 .or(expires_at);
428 let mut creds = CodexCredentials {
429 access_token: self.access_token,
430 refresh_token: self.refresh_token,
431 id_token: self.id_token,
432 account_id: None,
433 residency: None,
434 expires_at,
435 earliest_refresh_at,
436 };
437 creds.fill_claims();
438 creds
439 }
440}
441
442impl CodexCredentials {
443 fn fill_claims(&mut self) {
448 if self.account_id.is_none() {
449 self.account_id =
450 decode_claims(&self.access_token).and_then(JwtClaims::chatgpt_account_id);
451 }
452 if self.account_id.is_none() {
453 self.account_id = self
454 .id_token
455 .as_deref()
456 .and_then(decode_claims)
457 .and_then(JwtClaims::chatgpt_account_id);
458 }
459 if self.residency.is_none() {
460 self.residency =
461 decode_claims(&self.access_token).and_then(JwtClaims::compute_residency);
462 }
463 }
464}
465
466#[derive(Debug, Deserialize)]
468struct JwtClaims {
469 #[serde(default)]
471 chatgpt_account_id: Option<String>,
472 #[serde(default)]
474 chatgpt_compute_residency: Option<String>,
475 #[serde(default, rename = "https://api.openai.com/auth")]
477 auth: Option<AuthClaims>,
478}
479
480impl JwtClaims {
481 fn chatgpt_account_id(self) -> Option<String> {
482 let flat = self.chatgpt_account_id;
483 match self.auth {
484 Some(mut auth) => flat
485 .or_else(|| auth.chatgpt_account_id.take())
486 .or_else(|| auth.first_organization_id()),
487 None => flat,
488 }
489 }
490
491 fn compute_residency(self) -> Option<String> {
492 let flat = self.chatgpt_compute_residency;
493 match self.auth {
494 Some(auth) => flat.or(auth.chatgpt_compute_residency),
495 None => flat,
496 }
497 }
498}
499
500#[derive(Debug, Deserialize)]
501struct AuthClaims {
502 #[serde(default)]
503 chatgpt_account_id: Option<String>,
504 #[serde(default)]
505 chatgpt_compute_residency: Option<String>,
506 #[serde(default)]
507 organizations: Option<Vec<Organization>>,
508}
509
510impl AuthClaims {
511 fn first_organization_id(self) -> Option<String> {
512 self.organizations?.first().and_then(|org| org.id.clone())
513 }
514}
515
516#[derive(Debug, Deserialize)]
517struct Organization {
518 id: Option<String>,
519}
520
521fn decode_claims(jwt: &str) -> Option<JwtClaims> {
524 let payload = jwt.split('.').nth(1)?;
525 let bytes = URL_SAFE_NO_PAD.decode(payload).ok()?;
526 serde_json::from_slice(&bytes).ok()
527}
528
529#[cfg(test)]
530mod tests {
531 use super::*;
532
533 #[test]
534 fn pkce_challenge_is_s256_of_verifier() {
535 let pkce = Pkce::generate();
536 let expected = URL_SAFE_NO_PAD.encode(Sha256::digest(pkce.verifier().as_bytes()));
537 assert_eq!(pkce.challenge, expected);
538 assert_eq!(pkce.state().len(), 22); }
540
541 #[test]
542 fn authorize_url_carries_pkce_params() {
543 let pkce = Pkce::generate();
544 let url = pkce.authorize_url();
545 assert!(url.starts_with(AUTHORIZE_URL));
546 assert!(url.contains("code_challenge_method=S256"));
547 assert!(url.contains(&format!("state={}", pkce.state())));
548 assert!(url.contains("client_id=app_EMoamEEZ73f0CkXaXp7hrann"));
549 }
550
551 #[test]
552 fn callback_parses_code_and_state() {
553 let request =
554 "GET /auth/callback?code=abc123&state=xyz HTTP/1.1\r\nHost: localhost\r\n\r\n";
555 assert_eq!(parse_callback_request(request, "xyz").unwrap(), "abc123");
556 }
557
558 #[test]
559 fn callback_rejects_wrong_state() {
560 let request = "GET /auth/callback?code=abc&state=xyz HTTP/1.1\r\n\r\n";
561 assert!(matches!(
562 parse_callback_request(request, "other"),
563 Err(CodexError::StateMismatch)
564 ));
565 }
566
567 #[test]
568 fn callback_surfaces_error_param() {
569 let request = "GET /auth/callback?error=access_denied&state=xyz HTTP/1.1\r\n\r\n";
570 assert!(matches!(
571 parse_callback_request(request, "xyz"),
572 Err(CodexError::Callback(_))
573 ));
574 }
575
576 #[test]
577 fn nested_claims_decode() {
578 let payload = serde_json::json!({
581 "https://api.openai.com/auth": {
582 "chatgpt_account_id": "acct-1",
583 "chatgpt_compute_residency": "no_constraint"
584 }
585 });
586 let jwt = format!(
587 "h.{}.s",
588 URL_SAFE_NO_PAD.encode(serde_json::to_vec(&payload).unwrap())
589 );
590 let claims = decode_claims(&jwt).expect("claims decode");
591 assert_eq!(claims.chatgpt_account_id().as_deref(), Some("acct-1"));
592 let claims = decode_claims(&jwt).expect("claims decode");
593 assert_eq!(claims.compute_residency().as_deref(), Some("no_constraint"));
594 }
595
596 #[test]
597 fn credentials_roundtrip() {
598 let creds = CodexCredentials::new("at", "rt").with_expires_at(42);
599 let json = serde_json::to_string(&creds).unwrap();
600 let back: CodexCredentials = serde_json::from_str(&json).unwrap();
601 assert_eq!(back.access_token, "at");
602 assert_eq!(back.expires_at, Some(42));
603 }
604}