agentmux_srv\muxbus/
pkce.rs

1// Copyright 2026, AgentMux Corp.
2// SPDX-License-Identifier: Apache-2.0
3
4//! PKCE Authorization Code flow for Cognito desktop login.
5//!
6//! Flow:
7//!   1. Generate code_verifier + code_challenge (S256).
8//!   2. Bind a random local TCP port for the redirect_uri.
9//!   3. Open browser to Cognito hosted UI.
10//!   4. Await HTTP callback (code + state).
11//!   5. Exchange code for tokens via /oauth2/token.
12//!   6. Decode email from id_token claims (no re-verification needed —
13//!      we fetched the token directly from Cognito).
14
15use base64::Engine as _;
16use sha2::{Digest, Sha256};
17use tokio::io::{AsyncReadExt, AsyncWriteExt};
18use tokio::net::TcpListener;
19
20use crate::backend::storage::muxbus::MuxBusCredentials;
21
22pub const LOGIN_TIMEOUT_SECS: u64 = 300;
23
24#[derive(Debug)]
25pub struct PkceResult {
26    pub credentials: MuxBusCredentials,
27}
28
29pub async fn run_pkce_login(
30    cognito_domain: &str,
31    client_id: &str,
32    http_client: &reqwest::Client,
33) -> Result<PkceResult, String> {
34    // 1. Generate code_verifier (43–128 chars of URL-safe chars)
35    let v1 = uuid::Uuid::new_v4();
36    let v2 = uuid::Uuid::new_v4();
37    let mut raw = [0u8; 32];
38    raw[..16].copy_from_slice(v1.as_bytes());
39    raw[16..].copy_from_slice(v2.as_bytes());
40    let code_verifier = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(raw);
41
42    // 2. code_challenge = BASE64URL(SHA-256(code_verifier))
43    let mut hasher = Sha256::new();
44    hasher.update(code_verifier.as_bytes());
45    let digest = hasher.finalize();
46    let code_challenge = base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(digest);
47
48    // 3. State (CSRF)
49    let state = uuid::Uuid::new_v4().to_string();
50
51    // 4. Bind a random port
52    let listener = TcpListener::bind("127.0.0.1:0")
53        .await
54        .map_err(|e| format!("failed to bind callback listener: {e}"))?;
55    let port = listener
56        .local_addr()
57        .map_err(|e| format!("listener addr: {e}"))?
58        .port();
59
60    let redirect_uri = format!("http://127.0.0.1:{port}/callback");
61
62    // 5. Build auth URL
63    let scopes = "openid+email+profile+https%3A%2F%2Fmuxbus.agentmux.ai%2Fread+https%3A%2F%2Fmuxbus.agentmux.ai%2Fwrite";
64    let auth_url = format!(
65        "{cognito_domain}/oauth2/authorize\
66         ?response_type=code\
67         &client_id={client_id}\
68         &redirect_uri={redirect_uri_enc}\
69         &scope={scopes}\
70         &code_challenge={code_challenge}\
71         &code_challenge_method=S256\
72         &state={state}",
73        redirect_uri_enc = percent_encode(&redirect_uri),
74    );
75
76    // 6. Open browser
77    open_browser(&auth_url);
78
79    tracing::info!(
80        cognito_domain = cognito_domain,
81        port = port,
82        "muxbus: PKCE login started, awaiting browser callback"
83    );
84
85    // 7. Accept one connection (with timeout)
86    let (mut stream, _) = tokio::time::timeout(
87        std::time::Duration::from_secs(LOGIN_TIMEOUT_SECS),
88        listener.accept(),
89    )
90    .await
91    .map_err(|_| "login timed out (5 min) — please try again")?
92    .map_err(|e| format!("callback accept failed: {e}"))?;
93
94    // 8. Read HTTP request and extract code + state
95    let mut buf = [0u8; 4096];
96    let n = tokio::time::timeout(
97        std::time::Duration::from_secs(10),
98        stream.read(&mut buf),
99    )
100    .await
101    .map_err(|_| "callback read timed out")?
102    .map_err(|e| format!("callback read failed: {e}"))?;
103    let request = std::str::from_utf8(&buf[..n]).unwrap_or("");
104
105    // Respond to the browser immediately
106    let html = if request.contains("code=") {
107        b"HTTP/1.1 200 OK\r\nContent-Type: text/html\r\n\r\n\
108          <html><body><h2>Connected to AgentMux Cloud.</h2>\
109          <p>You can close this tab.</p></body></html>" as &[u8]
110    } else {
111        b"HTTP/1.1 400 Bad Request\r\nContent-Type: text/html\r\n\r\n\
112          <html><body><h2>Login failed.</h2><p>Please retry from AgentMux.</p></body></html>"
113    };
114    let _ = stream.write_all(html).await;
115
116    // Parse query from first line: GET /callback?code=xxx&state=yyy HTTP/1.1
117    let query = request
118        .lines()
119        .next()
120        .and_then(|line| line.split_whitespace().nth(1))
121        .and_then(|path| path.split_once('?').map(|(_, q)| q))
122        .unwrap_or("");
123
124    let code = query_param(query, "code")
125        .ok_or_else(|| "callback missing 'code' parameter")?;
126    let returned_state = query_param(query, "state").unwrap_or_default();
127
128    if returned_state != state {
129        return Err("CSRF mismatch — state parameter did not match".to_string());
130    }
131
132    // 9. Exchange code for tokens
133    let token_url = format!("{cognito_domain}/oauth2/token");
134    let params = [
135        ("grant_type", "authorization_code"),
136        ("client_id", client_id),
137        ("code", &code),
138        ("redirect_uri", &redirect_uri),
139        ("code_verifier", &code_verifier),
140    ];
141
142    let resp = http_client
143        .post(&token_url)
144        .form(&params)
145        .send()
146        .await
147        .map_err(|e| format!("token exchange request failed: {e}"))?;
148
149    if !resp.status().is_success() {
150        let body = resp.text().await.unwrap_or_default();
151        return Err(format!("token exchange failed: {body}"));
152    }
153
154    let token_json: serde_json::Value = resp
155        .json()
156        .await
157        .map_err(|e| format!("token response parse failed: {e}"))?;
158
159    let access_token = token_json["access_token"]
160        .as_str()
161        .unwrap_or("")
162        .to_string();
163    let refresh_token = token_json["refresh_token"]
164        .as_str()
165        .unwrap_or("")
166        .to_string();
167    let id_token = token_json["id_token"].as_str().unwrap_or("").to_string();
168    let expires_in = token_json["expires_in"].as_i64().unwrap_or(3600);
169    let expires_at = (std::time::SystemTime::now()
170        .duration_since(std::time::UNIX_EPOCH)
171        .unwrap_or_default()
172        .as_secs() as i64)
173        + expires_in;
174
175    // 10. Extract email + sub from id_token payload (no re-verification needed)
176    let (user_email, user_sub) = extract_jwt_claims(&id_token);
177
178    tracing::info!(
179        email = user_email,
180        "muxbus: PKCE login succeeded"
181    );
182
183    Ok(PkceResult {
184        credentials: MuxBusCredentials {
185            cognito_domain: cognito_domain.to_string(),
186            client_id: client_id.to_string(),
187            access_token,
188            refresh_token,
189            id_token,
190            expires_at,
191            user_email,
192            user_sub,
193        },
194    })
195}
196
197pub async fn refresh_token(
198    creds: &MuxBusCredentials,
199    http_client: &reqwest::Client,
200) -> Result<MuxBusCredentials, String> {
201    if creds.refresh_token.is_empty() {
202        return Err("no refresh token stored".to_string());
203    }
204    let token_url = format!("{}/oauth2/token", creds.cognito_domain);
205    let params = [
206        ("grant_type", "refresh_token"),
207        ("client_id", creds.client_id.as_str()),
208        ("refresh_token", creds.refresh_token.as_str()),
209    ];
210    let resp = http_client
211        .post(&token_url)
212        .form(&params)
213        .send()
214        .await
215        .map_err(|e| format!("refresh request failed: {e}"))?;
216    if !resp.status().is_success() {
217        let body = resp.text().await.unwrap_or_default();
218        return Err(format!("token refresh failed: {body}"));
219    }
220    let json: serde_json::Value = resp
221        .json()
222        .await
223        .map_err(|e| format!("refresh response parse failed: {e}"))?;
224
225    let access_token = json["access_token"].as_str().unwrap_or("").to_string();
226    let expires_in = json["expires_in"].as_i64().unwrap_or(3600);
227    let expires_at = (std::time::SystemTime::now()
228        .duration_since(std::time::UNIX_EPOCH)
229        .unwrap_or_default()
230        .as_secs() as i64)
231        + expires_in;
232    // Cognito doesn't rotate refresh_token on refresh — keep existing
233    let new_id_token = json["id_token"]
234        .as_str()
235        .map(|s| s.to_string())
236        .unwrap_or_else(|| creds.id_token.clone());
237
238    Ok(MuxBusCredentials {
239        cognito_domain: creds.cognito_domain.clone(),
240        client_id: creds.client_id.clone(),
241        access_token,
242        refresh_token: creds.refresh_token.clone(),
243        id_token: new_id_token,
244        expires_at,
245        user_email: creds.user_email.clone(),
246        user_sub: creds.user_sub.clone(),
247    })
248}
249
250fn query_param(query: &str, key: &str) -> Option<String> {
251    query.split('&').find_map(|pair| {
252        let (k, v) = pair.split_once('=')?;
253        if k == key {
254            Some(percent_decode(v))
255        } else {
256            None
257        }
258    })
259}
260
261fn percent_encode(s: &str) -> String {
262    let mut out = String::with_capacity(s.len() * 3);
263    for b in s.bytes() {
264        match b {
265            b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9'
266            | b'-' | b'_' | b'.' | b'~' => out.push(b as char),
267            _ => {
268                out.push('%');
269                out.push_str(&format!("{b:02X}"));
270            }
271        }
272    }
273    out
274}
275
276fn percent_decode(s: &str) -> String {
277    let mut out = Vec::with_capacity(s.len());
278    let bytes = s.as_bytes();
279    let mut i = 0;
280    while i < bytes.len() {
281        if bytes[i] == b'%' && i + 2 < bytes.len() {
282            if let Ok(hex) = u8::from_str_radix(
283                std::str::from_utf8(&bytes[i + 1..i + 3]).unwrap_or(""),
284                16,
285            ) {
286                out.push(hex);
287                i += 3;
288                continue;
289            }
290        } else if bytes[i] == b'+' {
291            out.push(b' ');
292            i += 1;
293            continue;
294        }
295        out.push(bytes[i]);
296        i += 1;
297    }
298    String::from_utf8_lossy(&out).into_owned()
299}
300
301fn extract_jwt_claims(token: &str) -> (String, String) {
302    let payload = token.splitn(3, '.').nth(1).unwrap_or("");
303    let decoded = base64::engine::general_purpose::URL_SAFE_NO_PAD
304        .decode(payload)
305        .or_else(|_| base64::engine::general_purpose::STANDARD.decode(payload))
306        .unwrap_or_default();
307    let json: serde_json::Value =
308        serde_json::from_slice(&decoded).unwrap_or_default();
309    let email = json["email"].as_str().unwrap_or("").to_string();
310    let sub = json["sub"].as_str().unwrap_or("").to_string();
311    (email, sub)
312}
313
314fn open_browser(url: &str) {
315    #[cfg(target_os = "windows")]
316    {
317        let _ = std::process::Command::new("cmd")
318            .args(["/C", "start", "", url])
319            .spawn();
320    }
321    #[cfg(target_os = "macos")]
322    {
323        let _ = std::process::Command::new("open").arg(url).spawn();
324    }
325    #[cfg(not(any(target_os = "windows", target_os = "macos")))]
326    {
327        let _ = std::process::Command::new("xdg-open").arg(url).spawn();
328    }
329}
agentmux_srv\muxbus/pkce.rs

agentmux_srv\muxbus/
pkce.rs
