OAuth2 & OpenID Connect

[dependencies]
oauth2 = "5.0.0"   # default features: reqwest 0.12 client with rustls
axum-extra = { version = "0.12.6", features = ["cookie-private"] }

use axum::{
    Router,
    extract::{FromRef, Query, State},
    http::StatusCode,
    response::{IntoResponse, Redirect, Response},
    routing::get,
};
use axum_extra::extract::{
    PrivateCookieJar,
    cookie::{Cookie, Key, SameSite},
};
use oauth2::{
    AuthUrl, AuthorizationCode, ClientId, ClientSecret, CsrfToken, EndpointNotSet, EndpointSet, PkceCodeChallenge,
    PkceCodeVerifier, RedirectUrl, Scope, TokenResponse, TokenUrl, basic::BasicClient,
};
use serde::{Deserialize, Serialize};

type GitHubClient = BasicClient<EndpointSet, EndpointNotSet, EndpointNotSet, EndpointNotSet, EndpointSet>;

#[derive(Clone)]
struct AppState {
    oauth: GitHubClient,
    http: oauth2::reqwest::Client,
    user_api: String,
    cookie_key: Key,
}

impl FromRef<AppState> for Key {
    fn from_ref(state: &AppState) -> Self {
        state.cookie_key.clone()
    }
}

/// What we must remember between the redirect and the callback.
#[derive(Serialize, Deserialize)]
struct Flow {
    state: String,
    pkce_verifier: String,
}

const FLOW_COOKIE: &str = "oauth_flow";

async fn start(State(app): State<AppState>, jar: PrivateCookieJar) -> (PrivateCookieJar, Redirect) {
    let (challenge, verifier) = PkceCodeChallenge::new_random_sha256();
    let (url, csrf_state) = app
        .oauth
        .authorize_url(CsrfToken::new_random)
        .add_scope(Scope::new("read:user".into()))
        .set_pkce_challenge(challenge)
        .url();

    let flow = Flow { state: csrf_state.secret().clone(), pkce_verifier: verifier.secret().clone() };
    let cookie = Cookie::build((FLOW_COOKIE, serde_json::to_string(&flow).expect("serialisable")))
        .http_only(true)
        .secure(true)
        .same_site(SameSite::Lax) // must survive the top-level redirect back from GitHub
        .path("/auth")
        .max_age(time::Duration::minutes(10))
        .build();
    (jar.add(cookie), Redirect::to(url.as_str()))
}

#[derive(Deserialize)]
struct Callback {
    code: String,
    state: String,
}

#[derive(Deserialize, Debug)]
struct GitHubUser {
    id: u64,
    login: String,
}

async fn callback(State(app): State<AppState>, jar: PrivateCookieJar, Query(cb): Query<Callback>) -> Response {
    // 1. the flow cookie must exist and its state must match: otherwise this is a forged or replayed callback
    let Some(flow) = jar.get(FLOW_COOKIE).and_then(|c| serde_json::from_str::<Flow>(c.value()).ok()) else {
        return (StatusCode::BAD_REQUEST, "login expired, please try again\n").into_response();
    };
    let jar = jar.remove(Cookie::build(FLOW_COOKIE).path("/auth"));
    if flow.state != cb.state {
        return (jar, (StatusCode::BAD_REQUEST, "state mismatch\n")).into_response();
    }

    // 2. exchange the code, proving possession of the PKCE verifier
    let token = match app
        .oauth
        .exchange_code(AuthorizationCode::new(cb.code))
        .set_pkce_verifier(PkceCodeVerifier::new(flow.pkce_verifier))
        .request_async(&app.http)
        .await
    {
        Ok(token) => token,
        Err(err) => {
            tracing::warn!(%err, "code exchange failed");
            return (jar, StatusCode::BAD_GATEWAY).into_response();
        }
    };

    // 3. who is this? (GitHub is OAuth2-only, so we call its API instead of reading an ID token)
    let user: GitHubUser = match app
        .http
        .get(&app.user_api)
        .bearer_auth(token.access_token().secret())
        .header("user-agent", "bookshelf")
        .send()
        .await
        .and_then(|r| r.error_for_status())
    {
        // oauth2's reqwest is built without the `json` feature: parse the bytes with serde_json
        Ok(resp) => match resp.bytes().await.ok().and_then(|b| serde_json::from_slice(&b).ok()) {
            Some(user) => user,
            None => return (jar, StatusCode::BAD_GATEWAY).into_response(),
        },
        Err(_) => return (jar, StatusCode::BAD_GATEWAY).into_response(),
    };

    // 4. link to a local account by the provider's immutable id (not by e-mail or login name),
    //    then start a session with session.cycle_id() as on the Sessions page
    tracing::info!(github_id = user.id, login = %user.login, "signed in with GitHub");
    (jar, Redirect::to("/account")).into_response()
}

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    // non-secret endpoints may default to GitHub's; the client credentials must be configured
    let env = |k: &str, default: &str| std::env::var(k).unwrap_or_else(|_| default.to_owned());
    let client_id = std::env::var("GITHUB_CLIENT_ID").expect("GITHUB_CLIENT_ID must be set");
    let client_secret = std::env::var("GITHUB_CLIENT_SECRET").expect("GITHUB_CLIENT_SECRET must be set");
    let oauth = BasicClient::new(ClientId::new(client_id))
        .set_client_secret(ClientSecret::new(client_secret))
        .set_auth_uri(AuthUrl::new(env("OAUTH_AUTH_URL", "https://github.com/login/oauth/authorize"))?)
        .set_token_uri(TokenUrl::new(env("OAUTH_TOKEN_URL", "https://github.com/login/oauth/access_token"))?)
        .set_redirect_uri(RedirectUrl::new(env("OAUTH_REDIRECT_URL", "http://localhost:3000/auth/github/callback"))?);

    let http = oauth2::reqwest::ClientBuilder::new()
        .redirect(oauth2::reqwest::redirect::Policy::none()) // never follow redirects on token requests
        .build()?;

    let state = AppState {
        oauth,
        http,
        user_api: env("GITHUB_USER_API", "https://api.github.com/user"),
        cookie_key: Key::generate(), // load a persistent 64-byte key from a secret store in production
    };
    let app = Router::new()
        .route("/auth/github", get(start))
        .route("/auth/github/callback", get(callback))
        .with_state(state);
    let listener = tokio::net::TcpListener::bind("127.0.0.1:3000").await?;
    axum::serve(listener, app).await?;
    Ok(())
}

$ curl -si localhost:3000/auth/github | grep -i -E '^HTTP|^location|^set-cookie' | cut -c1-120
HTTP/1.1 303 See Other
location: https://github.com/login/oauth/authorize?response_type=code&client_id=...&state=...&code_challenge=...
set-cookie: oauth_flow=6ADLyrbu...; HttpOnly; SameSite=Lax; Secure; Path=/auth; Max-Age=600

use openidconnect::{
    AuthenticationFlow, ClientId, ClientSecret, CsrfToken, IssuerUrl, Nonce, PkceCodeChallenge, RedirectUrl, Scope,
    core::{CoreClient, CoreProviderMetadata, CoreResponseType},
};

#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let http = openidconnect::reqwest::ClientBuilder::new()
        .redirect(openidconnect::reqwest::redirect::Policy::none())
        .build()?;

    // Discovery: fetches https://accounts.example.com/.well-known/openid-configuration and the JWKS
    let issuer = IssuerUrl::new(std::env::var("OIDC_ISSUER").unwrap_or_else(|_| "https://accounts.example.com".into()))?;
    let metadata = CoreProviderMetadata::discover_async(issuer, &http).await?;

    let client = CoreClient::from_provider_metadata(
        metadata,
        ClientId::new(std::env::var("OIDC_CLIENT_ID")?),
        Some(ClientSecret::new(std::env::var("OIDC_CLIENT_SECRET")?)),
    )
    .set_redirect_uri(RedirectUrl::new("https://bookshelf.example/auth/oidc/callback".into())?);

    let (challenge, _verifier) = PkceCodeChallenge::new_random_sha256();
    let (url, _state, _nonce) = client
        .authorize_url(AuthenticationFlow::<CoreResponseType>::AuthorizationCode, CsrfToken::new_random, Nonce::new_random)
        .add_scope(Scope::new("email".into()))
        .set_pkce_challenge(challenge)
        .url();
    println!("redirect the browser to {url}");

    // In the callback, after exchange_code(...).set_pkce_verifier(...).request_async(&http):
    //   let id_token = token.id_token().ok_or("no ID token")?;
    //   let claims = id_token.claims(&client.id_token_verifier(), &nonce)?;  // signature, iss, aud, exp, nonce
    //   let subject = claims.subject();                                      // the stable user id
    Ok(())
}