From 3af7b4343b045f1e8b963c8fcf537c41e0cccc51 Mon Sep 17 00:00:00 2001 From: Alex Bates Date: Wed, 5 Aug 2026 21:30:36 +0100 Subject: [PATCH] add oauth callback --- src/main.rs | 21 +++- src/oauth/callback.rs | 229 ++++++++++++++++++++++++++++++++++++++++++ src/oauth/mod.rs | 2 + 3 files changed, 251 insertions(+), 1 deletion(-) create mode 100644 src/oauth/callback.rs diff --git a/src/main.rs b/src/main.rs index d07890e..1228ba7 100644 --- a/src/main.rs +++ b/src/main.rs @@ -4,6 +4,7 @@ mod oauth; mod state; use axum::{ + http::{header, HeaderMap}, response::Html, routing::{get, post}, Router, @@ -11,6 +12,7 @@ use axum::{ use maud::{html, Markup, DOCTYPE}; use config::Config; +use oauth::session::{decode_identity_cookie, extract_cookie_value, IDENTITY_COOKIE_NAME}; use state::AppState; #[tokio::main] @@ -49,9 +51,26 @@ pub(crate) fn layout(body: Markup) -> Html { ) } -async fn index() -> Html { +async fn index(headers: HeaderMap) -> Html { + let identity = headers + .get(header::COOKIE) + .and_then(|v| v.to_str().ok()) + .and_then(|cookie_header| extract_cookie_value(cookie_header, IDENTITY_COOKIE_NAME)) + .and_then(|value| decode_identity_cookie(value).ok()); + layout(html! { h1 { "starhaven" } + @match &identity { + Some(identity) => { + p { "logged in as " (identity.handle.as_deref().unwrap_or(&identity.did)) } + form method="post" action="/auth/logout" { + button type="submit" { "log out" } + } + } + None => { + a href="/login" { "log in" } + } + } button hx-post="/hello" hx-target="#greeting" hx-swap="innerHTML" { "Say hello" } diff --git a/src/oauth/callback.rs b/src/oauth/callback.rs new file mode 100644 index 0000000..badeeda --- /dev/null +++ b/src/oauth/callback.rs @@ -0,0 +1,229 @@ +//! OAuth callback: exchange the authorization code for DPoP-bound tokens and +//! set the encrypted session + readable identity cookies. + +use atproto_identity::key::{identify_key, to_public}; +use atproto_oauth::resources::pds_resources; +use atproto_oauth::workflow::{oauth_complete, OAuthClient, OAuthRequest}; +use axum::extract::{Query, State}; +use axum::http::{header, HeaderMap}; +use axum::response::{IntoResponse, Redirect, Response}; +use axum::routing::get; +use axum::Router; +use chrono::{Duration, Utc}; +use serde::Deserialize; + +use crate::error::AppError; +use crate::oauth::session::{ + build_identity_cookie_header, build_session_cookie_header, encode_identity_cookie, + encode_session_cookie, IdentityCookie, SessionCookie, +}; +use crate::state::AppState; + +/// Query parameters for `GET /callback`. +#[derive(Debug, Deserialize)] +pub struct CallbackQuery { + /// The opaque state value echoed back. + pub state: String, + /// The issuer (RFC 9207), if present. + pub iss: Option, + /// The authorization code. + pub code: Option, + /// An error code, if the authorization failed. + pub error: Option, +} + +/// `GET /callback` -- complete the OAuth flow and set session cookies. +pub async fn callback( + State(state): State, + Query(query): Query, +) -> Result { + if let Some(error) = query.error { + return Err(AppError::BadRequest(format!( + "authorization failed: {error}" + ))); + } + + let code = query + .code + .ok_or_else(|| AppError::BadRequest("missing authorization code".to_string()))?; + + // Pop + delete: single-use, and a missing entry means it can't be trusted. + let persisted = { + let mut requests = state + .oauth_requests + .lock() + .expect("oauth_requests lock poisoned"); + requests.remove(&query.state) + } + .ok_or_else(|| AppError::BadRequest("unknown or expired oauth state".to_string()))?; + + if persisted.expires_at < Utc::now() { + return Err(AppError::BadRequest("oauth state expired".to_string())); + } + + // RFC 9207: `iss` must match the authorization server we sent the user to. + if let Some(iss) = &query.iss { + if iss != &persisted.issuer { + return Err(AppError::BadRequest("issuer mismatch".to_string())); + } + } + + // Re-resolve from the original login hint -- catches a handle re-pointed + // to a different DID mid-flow. + let login_hint = persisted + .login_hint + .trim_start_matches("at://") + .trim_start_matches('@'); + + let document = if login_hint.starts_with("https://") { + None + } else { + Some( + state + .identity_resolver + .resolve(login_hint) + .await + .map_err(|e| AppError::BadRequest(format!("identity resolution failed: {e}")))?, + ) + }; + + if let (Some(pinned), Some(document)) = (&persisted.subject, &document) { + if pinned != &document.id { + return Err(AppError::BadRequest( + "identity changed during the authorization flow".to_string(), + )); + } + } + + let pds_endpoint = if login_hint.starts_with("https://") { + login_hint.to_string() + } else { + document + .as_ref() + .and_then(|d| d.pds_endpoints().first().map(|s| s.to_string())) + .ok_or_else(|| AppError::BadRequest("no PDS endpoint".to_string()))? + }; + + // Re-discover the authorization server for the token exchange. + let (_protected_resource, authorization_server) = + pds_resources(&state.http_client, &pds_endpoint) + .await + .map_err(|e| AppError::Internal(anyhow::anyhow!("resource discovery failed: {e}")))?; + + if authorization_server.issuer != persisted.issuer { + return Err(AppError::BadRequest( + "authorization server issuer changed".to_string(), + )); + } + + // Rebuild the DPoP key + OAuth client used during the login/PAR step. + let dpop_key = identify_key(&persisted.dpop_private_key) + .map_err(|e| AppError::Internal(anyhow::anyhow!("invalid stored DPoP key: {e}")))?; + + let signing_key = state + .config + .oauth_private_keys + .first() + .cloned() + .ok_or_else(|| AppError::Internal(anyhow::anyhow!("no OAuth signing key configured")))?; + + let signing_public_key = to_public(&signing_key) + .map(|k| k.to_string()) + .unwrap_or_default(); + + let oauth_client = OAuthClient { + redirect_uri: state.config.oauth_redirect_uri(), + client_id: state.config.oauth_client_id(), + private_signing_key_data: signing_key, + }; + + // 0.14.5's `OAuthRequest` has no `subject` field — the crate doesn't bind + // the callback to an expected DID for us. We verify the resulting token's + // `sub` against `persisted.subject` ourselves below. + let oauth_request = OAuthRequest { + oauth_state: persisted.state.clone(), + issuer: persisted.issuer.clone(), + authorization_server: authorization_server.issuer.clone(), + nonce: persisted.nonce.clone(), + pkce_verifier: persisted.pkce_verifier.clone(), + signing_public_key, + dpop_private_key: persisted.dpop_private_key.clone(), + created_at: persisted.created_at, + expires_at: persisted.expires_at, + }; + + // Exchange the code for DPoP-bound access + refresh tokens. + let token_response = oauth_complete( + &state.http_client, + &oauth_client, + &dpop_key, + &code, + &oauth_request, + &authorization_server, + ) + .await + .map_err(|e| AppError::Internal(anyhow::anyhow!("token exchange failed: {e}")))?; + + // Subject-binding check: `oauth_complete` does not validate `sub` against + // anything for us. Never trust a token-supplied subject that contradicts + // the DID we resolved+pinned before sending the user off to authorize — + // that's how a hostile/compromised authorization server could mint a + // session as someone else. + let did = match (&persisted.subject, &token_response.sub) { + (Some(expected), Some(sub)) if expected == sub => sub.clone(), + (Some(expected), Some(sub)) => { + return Err(AppError::BadRequest(format!( + "token subject mismatch: expected {expected}, got {sub}" + ))); + } + (Some(expected), None) => expected.clone(), + (None, Some(sub)) => sub.clone(), + (None, None) => { + return Err(AppError::BadRequest( + "could not determine account DID".to_string(), + )); + } + }; + + let expires_at = Utc::now() + Duration::seconds(i64::from(token_response.expires_in)); + + let session = SessionCookie { + did: did.clone(), + access_token: token_response.access_token.clone(), + refresh_token: token_response.refresh_token.clone(), + expires_at, + dpop_private_key: persisted.dpop_private_key.clone(), + }; + + let identity = IdentityCookie { + did: did.clone(), + handle: document + .as_ref() + .and_then(|d| d.handles().map(|h| h.to_string())), + pds_url: Some(pds_endpoint), + }; + + let session_value = encode_session_cookie(&state.config.cookie_secret, &session) + .map_err(|e| AppError::Internal(anyhow::anyhow!("session encode failed: {e}")))?; + let identity_value = encode_identity_cookie(&identity) + .map_err(|e| AppError::Internal(anyhow::anyhow!("identity encode failed: {e}")))?; + + let max_age = Duration::days(30).num_seconds(); + let mut headers = HeaderMap::new(); + if let Ok(h) = build_session_cookie_header(&state.config.external_base, &session_value, max_age) + { + headers.append(header::SET_COOKIE, h); + } + if let Ok(h) = + build_identity_cookie_header(&state.config.external_base, &identity_value, max_age) + { + headers.append(header::SET_COOKIE, h); + } + + Ok((headers, Redirect::to("/")).into_response()) +} + +/// Routes for the callback step. +pub fn router() -> Router { + Router::new().route("/callback", get(callback)) +} diff --git a/src/oauth/mod.rs b/src/oauth/mod.rs index c9dbdae..33a4b82 100644 --- a/src/oauth/mod.rs +++ b/src/oauth/mod.rs @@ -1,6 +1,7 @@ //! OAuth confidential-client flow: client metadata + JWKS discovery, login, //! callback, session refresh and logout. Owns every OAuth-related route. +pub mod callback; pub mod login; pub mod metadata; pub mod session; @@ -16,4 +17,5 @@ pub fn router() -> Router { .route("/client-metadata.json", get(metadata::client_metadata)) .route("/jwks.json", get(metadata::jwks)) .merge(login::router()) + .merge(callback::router()) } -- 2.51.2