Menu
akurai-tasks
publicLatest change e0eb0fd4f7da2c0a23ef16030b8f6f34518268b3 - Accept scoped IDP tokens on Tasks MCP by Ólafur Búi Ólafsson
use akurai_json::{parse, Value};
use sha2::{Digest, Sha256};
use std::collections::BTreeMap;
use std::fs::File;
use std::io::{Read, Write};
use std::net::TcpStream;
use std::time::{Duration, SystemTime, UNIX_EPOCH};
#[derive(Clone)]
pub struct OidcConfig {
pub issuer: String,
pub internal_url: String,
pub client_id: String,
pub client_secret: String,
pub redirect_uri: String,
pub allowed_emails: Vec<String>,
}
#[derive(Clone)]
pub(crate) struct Pending {
pub verifier: String,
pub nonce: String,
pub expires_at: i64,
pub completed: Option<Identity>,
}
#[derive(Clone)]
pub(crate) struct Identity {
pub sub: String,
pub email: String,
}
pub(crate) fn verify_access_token(
token: &str,
config: &OidcConfig,
audience: &str,
) -> Result<Identity, String> {
let payload = token
.split('.')
.nth(1)
.ok_or_else(|| "access token is not a JWT".to_string())
.and_then(base64url_decode)?;
let claims =
parse(&String::from_utf8(payload).map_err(|_| "access token payload is not UTF-8")?)
.map_err(|error| format!("access token payload: {error}"))?;
if claims.get("iss").and_then(Value::as_str) != Some(config.issuer.trim_end_matches('/')) {
return Err("access token issuer mismatch".into());
}
if claims.get("aud").and_then(Value::as_str) != Some(audience) {
return Err("access token audience mismatch".into());
}
let response = request(
&config.internal_url,
"GET",
"/userinfo",
&[("Authorization", &format!("Bearer {token}"))],
None,
)?;
let user = parse(&response).map_err(|error| format!("userinfo response: {error}"))?;
let identity = Identity {
sub: user
.get("sub")
.and_then(Value::as_str)
.ok_or("userinfo omitted subject")?
.to_string(),
email: user
.get("email")
.and_then(Value::as_str)
.unwrap_or("")
.to_string(),
};
if !config.allowed_emails.is_empty()
&& !config
.allowed_emails
.iter()
.any(|allowed| allowed.eq_ignore_ascii_case(&identity.email))
{
return Err("IDP identity is not allowed".into());
}
Ok(identity)
}
pub(crate) fn begin(
config: &OidcConfig,
pending: &mut BTreeMap<String, Pending>,
) -> Result<String, String> {
if config.client_id.is_empty() || config.client_secret.is_empty() {
return Err("OIDC client is not configured".into());
}
let state = random_token(24)?;
let verifier = random_token(32)?;
let nonce = random_token(24)?;
let challenge = base64url(&Sha256::digest(verifier.as_bytes()));
pending.retain(|_, value| value.expires_at > now());
pending.insert(
state.clone(),
Pending {
verifier,
nonce: nonce.clone(),
expires_at: now() + 300,
completed: None,
},
);
Ok(format!(
"{}/authorize?client_id={}&redirect_uri={}&response_type=code&scope=openid%20profile%20email&state={}&nonce={}&code_challenge={}&code_challenge_method=S256",
config.issuer.trim_end_matches('/'), encode(&config.client_id), encode(&config.redirect_uri), encode(&state), encode(&nonce), encode(&challenge)
))
}
pub(crate) fn callback(
config: &OidcConfig,
query: Option<&str>,
pending: &mut BTreeMap<String, Pending>,
) -> Result<Identity, String> {
let params = form(query.unwrap_or(""));
let state = params.get("state").ok_or("OIDC state missing")?;
let code = params
.get("code")
.ok_or("OIDC authorization code missing")?;
let flow = pending.get_mut(state).ok_or("OIDC state is invalid")?;
if flow.expires_at <= now() {
return Err("OIDC state expired".into());
}
if let Some(identity) = &flow.completed {
return Ok(identity.clone());
}
let token_body = format!("grant_type=authorization_code&code={}&redirect_uri={}&client_id={}&client_secret={}&code_verifier={}", encode(code), encode(&config.redirect_uri), encode(&config.client_id), encode(&config.client_secret), encode(&flow.verifier));
let token_response = request(
&config.internal_url,
"POST",
"/token",
&[("Content-Type", "application/x-www-form-urlencoded")],
Some(&token_body),
)?;
let token = parse(&token_response).map_err(|error| format!("OIDC token response: {error}"))?;
let access = token
.get("access_token")
.and_then(Value::as_str)
.ok_or("OIDC access token missing")?;
let id_token = token
.get("id_token")
.and_then(Value::as_str)
.ok_or("OIDC ID token missing")?;
validate_nonce(id_token, &flow.nonce)?;
let userinfo = request(
&config.internal_url,
"GET",
"/userinfo",
&[("Authorization", &format!("Bearer {access}"))],
None,
)?;
let user = parse(&userinfo).map_err(|error| format!("OIDC userinfo response: {error}"))?;
let sub = user
.get("sub")
.and_then(Value::as_str)
.ok_or("OIDC subject missing")?
.to_string();
let email = user
.get("email")
.and_then(Value::as_str)
.unwrap_or("")
.to_string();
if !config.allowed_emails.is_empty()
&& !config
.allowed_emails
.iter()
.any(|allowed| allowed.eq_ignore_ascii_case(&email))
{
return Err("OIDC identity is not allowed".into());
}
let identity = Identity { sub, email };
flow.completed = Some(identity.clone());
Ok(identity)
}
fn validate_nonce(token: &str, expected: &str) -> Result<(), String> {
let payload = token
.split('.')
.nth(1)
.ok_or("OIDC ID token is malformed")?;
let bytes = base64url_decode(payload)?;
let json = parse(&String::from_utf8(bytes).map_err(|_| "OIDC ID token payload is not UTF-8")?)
.map_err(|error| format!("OIDC ID token payload: {error}"))?;
if json.get("nonce").and_then(Value::as_str) != Some(expected) {
return Err("OIDC nonce mismatch".into());
}
Ok(())
}
fn request(
base: &str,
method: &str,
path: &str,
headers: &[(&str, &str)],
body: Option<&str>,
) -> Result<String, String> {
let without_scheme = base
.strip_prefix("http://")
.ok_or("OIDC internal URL must use loopback HTTP")?;
let (host, port) = without_scheme
.split_once(':')
.unwrap_or((without_scheme, "80"));
if !matches!(host, "127.0.0.1" | "localhost") {
return Err("OIDC internal URL must be loopback".into());
}
let port: u16 = port.parse().map_err(|_| "OIDC internal port is invalid")?;
let mut stream = TcpStream::connect_timeout(
&format!("{host}:{port}")
.parse()
.map_err(|_| "OIDC internal address invalid")?,
Duration::from_secs(5),
)
.map_err(|error| format!("OIDC connect: {error}"))?;
stream
.set_read_timeout(Some(Duration::from_secs(10)))
.map_err(|error| format!("OIDC timeout: {error}"))?;
let body = body.unwrap_or("");
let mut wire = format!("{method} {path} HTTP/1.1\r\nHost: {host}:{port}\r\nConnection: close\r\nContent-Length: {}\r\n", body.len());
for (name, value) in headers {
wire.push_str(&format!("{name}: {value}\r\n"));
}
wire.push_str("\r\n");
wire.push_str(body);
stream
.write_all(wire.as_bytes())
.map_err(|error| format!("OIDC write: {error}"))?;
let mut bytes = Vec::new();
stream
.read_to_end(&mut bytes)
.map_err(|error| format!("OIDC read: {error}"))?;
let response = String::from_utf8_lossy(&bytes);
let status = response
.lines()
.next()
.and_then(|line| line.split_whitespace().nth(1))
.and_then(|value| value.parse::<u16>().ok())
.unwrap_or(0);
let content = response
.split_once("\r\n\r\n")
.map(|(_, body)| body)
.unwrap_or("");
if !(200..300).contains(&status) {
return Err(format!("OIDC endpoint returned HTTP {status}"));
}
Ok(content.to_string())
}
fn form(query: &str) -> BTreeMap<String, String> {
query
.split('&')
.filter_map(|pair| pair.split_once('='))
.map(|(key, value)| (decode(key), decode(value)))
.collect()
}
fn encode(value: &str) -> String {
let mut out = String::new();
for byte in value.bytes() {
if byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'_' | b'.' | b'~') {
out.push(byte as char);
} else {
out.push_str(&format!("%{byte:02X}"));
}
}
out
}
fn decode(value: &str) -> String {
let bytes = value.as_bytes();
let mut out = Vec::with_capacity(bytes.len());
let mut index = 0;
while index < bytes.len() {
if bytes[index] == b'%' && index + 2 < bytes.len() {
if let Ok(byte) = u8::from_str_radix(&value[index + 1..index + 3], 16) {
out.push(byte);
index += 3;
continue;
}
}
out.push(if bytes[index] == b'+' {
b' '
} else {
bytes[index]
});
index += 1;
}
String::from_utf8_lossy(&out).into_owned()
}
fn random_token(bytes: usize) -> Result<String, String> {
let mut data = vec![0u8; bytes];
File::open("/dev/urandom")
.and_then(|mut file| file.read_exact(&mut data))
.map_err(|error| format!("secure random: {error}"))?;
Ok(base64url(&data))
}
fn base64url(bytes: &[u8]) -> String {
const TABLE: &[u8] = b"ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_";
let mut out = String::new();
for chunk in bytes.chunks(3) {
let n = ((chunk[0] as u32) << 16)
| ((chunk.get(1).copied().unwrap_or(0) as u32) << 8)
| chunk.get(2).copied().unwrap_or(0) as u32;
out.push(TABLE[((n >> 18) & 63) as usize] as char);
out.push(TABLE[((n >> 12) & 63) as usize] as char);
if chunk.len() > 1 {
out.push(TABLE[((n >> 6) & 63) as usize] as char);
}
if chunk.len() > 2 {
out.push(TABLE[(n & 63) as usize] as char);
}
}
out
}
fn base64url_decode(value: &str) -> Result<Vec<u8>, String> {
fn digit(byte: u8) -> Option<u8> {
match byte {
b'A'..=b'Z' => Some(byte - b'A'),
b'a'..=b'z' => Some(byte - b'a' + 26),
b'0'..=b'9' => Some(byte - b'0' + 52),
b'-' => Some(62),
b'_' => Some(63),
_ => None,
}
}
let mut out = Vec::new();
for chunk in value.as_bytes().chunks(4) {
if chunk.len() < 2 {
return Err("invalid base64url".into());
}
let a = digit(chunk[0]).ok_or("invalid base64url")? as u32;
let b = digit(chunk[1]).ok_or("invalid base64url")? as u32;
let c = chunk.get(2).and_then(|value| digit(*value)).unwrap_or(0) as u32;
let d = chunk.get(3).and_then(|value| digit(*value)).unwrap_or(0) as u32;
let n = (a << 18) | (b << 12) | (c << 6) | d;
out.push((n >> 16) as u8);
if chunk.len() > 2 {
out.push((n >> 8) as u8);
}
if chunk.len() > 3 {
out.push(n as u8);
}
}
Ok(out)
}
fn now() -> i64 {
SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_secs() as i64
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn completed_callback_replays_without_reusing_authorization_code() {
let config = OidcConfig {
issuer: "https://issuer.invalid".into(),
internal_url: "http://127.0.0.1:1".into(),
client_id: "client".into(),
client_secret: "secret".into(),
redirect_uri: "https://tasks.invalid/auth/callback".into(),
allowed_emails: Vec::new(),
};
let mut pending = BTreeMap::from([(
"state".into(),
Pending {
verifier: "verifier".into(),
nonce: "nonce".into(),
expires_at: now() + 300,
completed: Some(Identity {
sub: "subject".into(),
email: "agent@example.com".into(),
}),
},
)]);
let replay = callback(&config, Some("state=state&code=already-used"), &mut pending)
.expect("completed callback should replay");
assert_eq!(replay.sub, "subject");
assert_eq!(replay.email, "agent@example.com");
}
}