use std::time::Duration;
use serde::Deserialize;
use crate::error::{Error, Result};
#[derive(Debug, Clone)]
pub struct OAuthEndpoints {
pub device_authorization_endpoint: String,
pub token_endpoint: String,
pub client_id: String,
pub scope: Option<String>,
}
#[derive(Debug, Clone, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
pub struct McpOAuthTokens {
pub access_token: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub refresh_token: Option<String>,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub expires_at_secs: Option<u64>,
}
impl McpOAuthTokens {
pub fn is_expired_at(&self, now_secs: u64, skew_secs: u64) -> bool {
self.expires_at_secs
.map(|exp| now_secs.saturating_add(skew_secs) >= exp)
.unwrap_or(false)
}
}
#[derive(Debug, Clone)]
pub struct DeviceAuthorization {
pub device_code: String,
pub user_code: String,
pub verification_uri: String,
pub verification_uri_complete: Option<String>,
pub interval_secs: u64,
pub expires_in_secs: u64,
}
#[derive(Deserialize)]
struct DeviceAuthResponse {
device_code: String,
user_code: String,
verification_uri: String,
#[serde(default)]
verification_uri_complete: Option<String>,
#[serde(default = "default_interval")]
interval: u64,
#[serde(default = "default_expires_in")]
expires_in: u64,
}
fn default_interval() -> u64 {
5
}
fn default_expires_in() -> u64 {
600
}
#[derive(Deserialize)]
struct TokenResponse {
access_token: String,
#[serde(default)]
refresh_token: Option<String>,
#[serde(default)]
expires_in: Option<u64>,
}
#[derive(Deserialize)]
struct TokenErrorResponse {
error: String,
}
pub async fn start_device_authorization(
client: &reqwest::Client,
ep: &OAuthEndpoints,
) -> Result<DeviceAuthorization> {
let mut form = vec![("client_id", ep.client_id.as_str())];
if let Some(scope) = &ep.scope {
form.push(("scope", scope.as_str()));
}
let resp = client
.post(&ep.device_authorization_endpoint)
.form(&form)
.send()
.await
.map_err(|e| Error::tool("mcp_oauth", format!("device authorization request: {e}")))?;
if !resp.status().is_success() {
return Err(Error::tool(
"mcp_oauth",
format!("device authorization: http status {}", resp.status()),
));
}
let body: DeviceAuthResponse = resp.json().await.map_err(|e| {
Error::tool(
"mcp_oauth",
format!("decoding device authorization response: {e}"),
)
})?;
Ok(DeviceAuthorization {
device_code: body.device_code,
user_code: body.user_code,
verification_uri: body.verification_uri,
verification_uri_complete: body.verification_uri_complete,
interval_secs: body.interval,
expires_in_secs: body.expires_in,
})
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum DevicePollOutcome {
Pending,
SlowDown,
Authorized(McpOAuthTokens),
Denied,
Expired,
}
pub async fn poll_device_token(
client: &reqwest::Client,
ep: &OAuthEndpoints,
device_code: &str,
) -> Result<DevicePollOutcome> {
let form = [
("grant_type", "urn:ietf:params:oauth:grant-type:device_code"),
("device_code", device_code),
("client_id", ep.client_id.as_str()),
];
let resp = client
.post(&ep.token_endpoint)
.form(&form)
.send()
.await
.map_err(|e| Error::tool("mcp_oauth", format!("token poll: {e}")))?;
if resp.status().is_success() {
let body: TokenResponse = resp
.json()
.await
.map_err(|e| Error::tool("mcp_oauth", format!("decoding token response: {e}")))?;
return Ok(DevicePollOutcome::Authorized(McpOAuthTokens {
access_token: body.access_token,
refresh_token: body.refresh_token,
expires_at_secs: body.expires_in.map(|secs| now_secs() + secs),
}));
}
let body: TokenErrorResponse = resp.json().await.unwrap_or(TokenErrorResponse {
error: "unknown_error".to_string(),
});
match body.error.as_str() {
"authorization_pending" => Ok(DevicePollOutcome::Pending),
"slow_down" => Ok(DevicePollOutcome::SlowDown),
"expired_token" => Ok(DevicePollOutcome::Expired),
_ => Ok(DevicePollOutcome::Denied),
}
}
pub async fn run_device_flow(
client: &reqwest::Client,
ep: &OAuthEndpoints,
on_prompt: impl FnOnce(&DeviceAuthorization),
) -> Result<McpOAuthTokens> {
let auth = start_device_authorization(client, ep).await?;
on_prompt(&auth);
let deadline = now_secs() + auth.expires_in_secs;
let mut interval = auth.interval_secs.max(1);
loop {
tokio::time::sleep(Duration::from_secs(interval)).await;
match poll_device_token(client, ep, &auth.device_code).await? {
DevicePollOutcome::Authorized(tokens) => return Ok(tokens),
DevicePollOutcome::Pending => {
if now_secs() >= deadline {
return Err(Error::tool(
"mcp_oauth",
"device code expired while polling",
));
}
}
DevicePollOutcome::SlowDown => {
interval += 5;
if now_secs() >= deadline {
return Err(Error::tool(
"mcp_oauth",
"device code expired while polling",
));
}
}
DevicePollOutcome::Denied => {
return Err(Error::tool("mcp_oauth", "authorization was denied"));
}
DevicePollOutcome::Expired => {
return Err(Error::tool("mcp_oauth", "device code expired"));
}
}
}
}
pub async fn refresh_token(
client: &reqwest::Client,
ep: &OAuthEndpoints,
refresh_token: &str,
) -> Result<McpOAuthTokens> {
let form = [
("grant_type", "refresh_token"),
("refresh_token", refresh_token),
("client_id", ep.client_id.as_str()),
];
let resp = client
.post(&ep.token_endpoint)
.form(&form)
.send()
.await
.map_err(|e| Error::tool("mcp_oauth", format!("refresh request: {e}")))?;
if !resp.status().is_success() {
return Err(Error::tool(
"mcp_oauth",
format!("refresh: http status {}", resp.status()),
));
}
let body: TokenResponse = resp
.json()
.await
.map_err(|e| Error::tool("mcp_oauth", format!("decoding refresh response: {e}")))?;
Ok(McpOAuthTokens {
access_token: body.access_token,
refresh_token: body.refresh_token,
expires_at_secs: body.expires_in.map(|secs| now_secs() + secs),
})
}
fn now_secs() -> u64 {
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.map(|d| d.as_secs())
.unwrap_or(0)
}
pub fn bearer_header(tokens: &McpOAuthTokens) -> (String, String) {
(
"Authorization".to_string(),
format!("Bearer {}", tokens.access_token),
)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn is_expired_at_treats_unknown_lifetime_as_not_expired() {
let t = McpOAuthTokens {
access_token: "x".into(),
refresh_token: None,
expires_at_secs: None,
};
assert!(!t.is_expired_at(u64::MAX / 2, 0));
}
#[test]
fn is_expired_at_honors_skew() {
let t = McpOAuthTokens {
access_token: "x".into(),
refresh_token: None,
expires_at_secs: Some(1000),
};
assert!(!t.is_expired_at(900, 30));
assert!(t.is_expired_at(980, 30)); assert!(t.is_expired_at(1000, 0));
}
#[test]
fn bearer_header_has_the_expected_shape() {
let t = McpOAuthTokens {
access_token: "secret123".into(),
refresh_token: None,
expires_at_secs: None,
};
let (name, value) = bearer_header(&t);
assert_eq!(name, "Authorization");
assert_eq!(value, "Bearer secret123");
}
}