ai_usagebar/kiro/
oauth.rs

1//! AWS SSO OIDC token refresh — `POST https://oidc.<region>.amazonaws.com/token`.
2//! Public, documented API (`CreateToken`:
3//! <https://docs.aws.amazon.com/singlesignon/latest/OIDCAPIReference/API_CreateToken.html>),
4//! unlike `fetch.rs`'s reverse-engineered `GetUsageLimits` call.
5//!
6//! kiro-cli's own cached access token lives about an hour (see `db.rs`); this
7//! refreshes it using the `refresh_token` + `client_id`/`client_secret`
8//! kiro-cli already registered for itself at `kiro-cli login` time — nothing
9//! new to authenticate, just the same local credentials `db.rs` already read.
10//! The refreshed token is **not** written back to kiro-cli's own database
11//! (mirroring `db.rs`'s read-only treatment of that live file); `fetch.rs`
12//! persists it in ai-usagebar's own account-scoped credential sidecar.
13
14use serde::{Deserialize, Serialize};
15
16use crate::error::{AppError, Result};
17
18/// Refresh happens this far ahead of the cached expiry so a slow round-trip
19/// never races the token's actual death. Mirrors `openai::oauth::REFRESH_BUFFER_SECS`.
20pub const REFRESH_BUFFER_SECS: i64 = 300;
21
22pub fn validate_region(region: &str) -> Result<()> {
23    let parts: Vec<_> = region.split('-').collect();
24    let valid = (3..=5).contains(&parts.len())
25        && region.len() <= 32
26        && parts[0].chars().all(|c| c.is_ascii_lowercase())
27        && parts[parts.len() - 1].chars().all(|c| c.is_ascii_digit())
28        && parts.iter().all(|part| {
29            !part.is_empty()
30                && part
31                    .chars()
32                    .all(|c| c.is_ascii_lowercase() || c.is_ascii_digit())
33        });
34    if valid {
35        Ok(())
36    } else {
37        Err(AppError::Credentials(
38            "Kiro CLI token contains an invalid AWS region. Run `kiro-cli login` again.".into(),
39        ))
40    }
41}
42
43pub fn token_endpoint(region: &str) -> Result<String> {
44    validate_region(region)?;
45    Ok(format!("https://oidc.{region}.amazonaws.com/token"))
46}
47
48#[derive(Debug, Serialize)]
49struct RefreshRequest<'a> {
50    #[serde(rename = "clientId")]
51    client_id: &'a str,
52    #[serde(rename = "clientSecret")]
53    client_secret: &'a str,
54    #[serde(rename = "grantType")]
55    grant_type: &'a str,
56    #[serde(rename = "refreshToken")]
57    refresh_token: &'a str,
58}
59
60#[derive(Debug, Deserialize)]
61pub struct RefreshResponse {
62    #[serde(rename = "accessToken", deserialize_with = "de_nonempty_string")]
63    pub access_token: String,
64    #[serde(
65        rename = "refreshToken",
66        default,
67        deserialize_with = "de_opt_nonempty_string"
68    )]
69    pub refresh_token: Option<String>,
70    #[serde(rename = "expiresIn", deserialize_with = "de_positive_u64")]
71    pub expires_in: u64,
72}
73
74fn de_nonempty_string<'de, D>(d: D) -> std::result::Result<String, D::Error>
75where
76    D: serde::Deserializer<'de>,
77{
78    let value = String::deserialize(d)?;
79    if value.trim().is_empty() {
80        Err(serde::de::Error::custom("accessToken cannot be empty"))
81    } else {
82        Ok(value)
83    }
84}
85
86fn de_opt_nonempty_string<'de, D>(d: D) -> std::result::Result<Option<String>, D::Error>
87where
88    D: serde::Deserializer<'de>,
89{
90    Option::<String>::deserialize(d)?
91        .map(|value| {
92            if value.trim().is_empty() {
93                Err(serde::de::Error::custom("refreshToken cannot be empty"))
94            } else {
95                Ok(value)
96            }
97        })
98        .transpose()
99}
100
101fn de_positive_u64<'de, D>(d: D) -> std::result::Result<u64, D::Error>
102where
103    D: serde::Deserializer<'de>,
104{
105    let v = serde_json::Value::deserialize(d)?;
106    match v {
107        serde_json::Value::Number(n) => {
108            const MAX_SAFE: u64 = (i64::MAX as u64) / 2;
109            if let Some(value) = n.as_u64().filter(|value| (1..=MAX_SAFE).contains(value)) {
110                Ok(value)
111            } else {
112                Err(serde::de::Error::custom(
113                    "expiresIn must be a positive integer in range",
114                ))
115            }
116        }
117        _ => Err(serde::de::Error::custom("expiresIn must be a number")),
118    }
119}
120
121/// Refresh the access token against `endpoint` (build it with
122/// [`token_endpoint`] for production; tests point it at mockito instead).
123/// Never echoes the upstream error body verbatim — an `invalid_grant`
124/// response from an OAuth token endpoint is not guaranteed not to include
125/// account-identifying detail, same reasoning as `cursor::fetch::error_to_pair`.
126pub async fn refresh(
127    client: &reqwest::Client,
128    endpoint: &str,
129    client_id: &str,
130    client_secret: &str,
131    refresh_token: &str,
132) -> Result<RefreshResponse> {
133    let req = RefreshRequest {
134        client_id,
135        client_secret,
136        grant_type: "refresh_token",
137        refresh_token,
138    };
139
140    let resp = client
141        .post(endpoint)
142        .header("Content-Type", "application/json")
143        .json(&req)
144        .send()
145        .await?;
146
147    let status = resp.status();
148    let body = crate::vendor::read_body_capped(resp, crate::vendor::MAX_BODY_BYTES).await?;
149    if !status.is_success() {
150        return Err(AppError::Http {
151            status: status.as_u16(),
152            body: "Kiro CLI token refresh failed".into(),
153        });
154    }
155    serde_json::from_slice(&body)
156        .map_err(|e| AppError::Schema(format!("kiro token refresh response: {e}")))
157}
158
159pub fn needs_refresh(expires_at_secs: i64, now_secs: i64) -> bool {
160    expires_at_secs < now_secs + REFRESH_BUFFER_SECS
161}
162
163#[cfg(test)]
164mod tests {
165    use super::*;
166
167    #[test]
168    fn token_endpoint_is_region_scoped() {
169        assert_eq!(
170            token_endpoint("us-east-1").unwrap(),
171            "https://oidc.us-east-1.amazonaws.com/token"
172        );
173        assert_eq!(
174            token_endpoint("us-gov-west-1").unwrap(),
175            "https://oidc.us-gov-west-1.amazonaws.com/token"
176        );
177    }
178
179    #[test]
180    fn unsafe_regions_are_rejected_before_url_construction() {
181        for region in [
182            "evil.example/#",
183            "us-east-1@evil.example",
184            "US-EAST-1",
185            "us--1",
186        ] {
187            assert!(token_endpoint(region).is_err(), "{region}");
188        }
189    }
190
191    #[test]
192    fn needs_refresh_threshold() {
193        let now = 1_000_000;
194        assert!(needs_refresh(now + 100, now));
195        assert!(!needs_refresh(now + 1000, now));
196    }
197
198    #[test]
199    fn empty_access_token_is_rejected() {
200        let body = r#"{"accessToken":"","expiresIn":3600}"#;
201        assert!(serde_json::from_str::<RefreshResponse>(body).is_err());
202    }
203
204    #[test]
205    fn malformed_expires_in_is_rejected_not_dropped() {
206        for value in [r#""3600""#, "-1", "0", "null", "true"] {
207            let body = format!(r#"{{"accessToken":"new","expiresIn":{value}}}"#);
208            assert!(
209                serde_json::from_str::<RefreshResponse>(&body).is_err(),
210                "{body}"
211            );
212        }
213        assert!(serde_json::from_str::<RefreshResponse>(r#"{"accessToken":"new"}"#).is_err());
214    }
215
216    #[test]
217    fn empty_rotated_refresh_token_is_rejected() {
218        let body = r#"{"accessToken":"new","refreshToken":" ","expiresIn":3600}"#;
219        assert!(serde_json::from_str::<RefreshResponse>(body).is_err());
220    }
221
222    #[tokio::test]
223    async fn refresh_success_parses_the_new_token() {
224        let mut server = mockito::Server::new_async().await;
225        let m = server
226            .mock("POST", "/token")
227            .with_status(200)
228            .with_body(r#"{"accessToken":"new-at","tokenType":"Bearer","expiresIn":3600}"#)
229            .create_async()
230            .await;
231        let client = reqwest::Client::new();
232        let r = refresh(
233            &client,
234            &format!("{}/token", server.url()),
235            "cid",
236            "csecret",
237            "old-rt",
238        )
239        .await
240        .unwrap();
241        assert_eq!(r.access_token, "new-at");
242        assert_eq!(r.expires_in, 3600);
243        assert_eq!(r.refresh_token, None);
244        m.assert_async().await;
245    }
246
247    #[tokio::test]
248    async fn refresh_sends_the_expected_json_body() {
249        let mut server = mockito::Server::new_async().await;
250        let m = server
251            .mock("POST", "/token")
252            .match_body(mockito::Matcher::Json(serde_json::json!({
253                "clientId": "cid",
254                "clientSecret": "csecret",
255                "grantType": "refresh_token",
256                "refreshToken": "old-rt",
257            })))
258            .with_status(200)
259            .with_body(r#"{"accessToken":"new-at","expiresIn":3600}"#)
260            .create_async()
261            .await;
262        let client = reqwest::Client::new();
263        refresh(
264            &client,
265            &format!("{}/token", server.url()),
266            "cid",
267            "csecret",
268            "old-rt",
269        )
270        .await
271        .unwrap();
272        m.assert_async().await;
273    }
274
275    #[tokio::test]
276    async fn refresh_400_does_not_echo_the_body() {
277        let mut server = mockito::Server::new_async().await;
278        server
279            .mock("POST", "/token")
280            .with_status(400)
281            .with_body(r#"{"error":"invalid_grant","error_description":"sensitive detail"}"#)
282            .create_async()
283            .await;
284        let client = reqwest::Client::new();
285        let err = refresh(
286            &client,
287            &format!("{}/token", server.url()),
288            "cid",
289            "csecret",
290            "old-rt",
291        )
292        .await
293        .unwrap_err();
294        match err {
295            AppError::Http { status, body } => {
296                assert_eq!(status, 400);
297                assert!(!body.contains("sensitive detail"));
298            }
299            other => panic!("expected Http error, got {other:?}"),
300        }
301    }
302}
ai_usagebar/kiro/oauth.rs

ai_usagebar/kiro/
oauth.rs
