hephaestus/crates/hephd/tests/oauth.rs

//! Device-code flow + token store (tech-spec §13, slice 10b), offline.
//!
//! A mock OAuth provider serves discovery, the device-authorization endpoint,
//! and the token endpoint (which reports `authorization_pending` once before
//! issuing tokens). We drive `DeviceFlow` against it with an injected no-op
//! sleep, so the polling loop is exercised deterministically and instantly.

use std::collections::HashMap;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{mpsc, Arc};
use std::thread;
use std::time::Duration;

use axum::extract::{Form, State};
use axum::http::StatusCode;
use axum::response::{IntoResponse, Response};
use axum::routing::{get, post};
use axum::{Json, Router};
use serde_json::{json, Value};

use hephd::oauth::{DeviceFlow, MemoryTokenStore, StoredToken, TokenStore};

#[derive(Clone)]
struct IdpState {
    base: String,
    /// How many times the token endpoint has been polled for the device code.
    polls: Arc<AtomicUsize>,
}

/// Start a mock OIDC provider; return its base URL.
fn start_idp() -> String {
    let (tx, rx) = mpsc::channel();
    thread::spawn(move || {
        let rt = tokio::runtime::Builder::new_current_thread()
            .enable_all()
            .build()
            .unwrap();
        rt.block_on(async move {
            let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
            let base = format!("http://{}", listener.local_addr().unwrap());
            tx.send(base.clone()).unwrap();
            let state = IdpState {
                base,
                polls: Arc::new(AtomicUsize::new(0)),
            };
            let app = Router::new()
                .route("/.well-known/openid-configuration", get(discovery))
                .route("/device", post(device_authorization))
                .route("/token", post(token))
                .with_state(state);
            axum::serve(listener, app).await.unwrap();
        });
    });
    rx.recv_timeout(Duration::from_secs(5)).unwrap()
}

async fn discovery(State(s): State<IdpState>) -> Json<Value> {
    Json(json!({
        "issuer": s.base,
        "device_authorization_endpoint": format!("{}/device", s.base),
        "token_endpoint": format!("{}/token", s.base),
    }))
}

async fn device_authorization(State(s): State<IdpState>) -> Json<Value> {
    Json(json!({
        "device_code": "dev-code-xyz",
        "user_code": "WDJB-MJHT",
        "verification_uri": format!("{}/activate", s.base),
        "interval": 1,
        "expires_in": 300,
    }))
}

async fn token(State(s): State<IdpState>, Form(form): Form<HashMap<String, String>>) -> Response {
    match form.get("grant_type").map(String::as_str) {
        Some("urn:ietf:params:oauth:grant-type:device_code") => {
            // Report pending on the first poll, then issue tokens.
            if s.polls.fetch_add(1, Ordering::SeqCst) == 0 {
                return (
                    StatusCode::BAD_REQUEST,
                    Json(json!({ "error": "authorization_pending" })),
                )
                    .into_response();
            }
            Json(json!({
                "access_token": "access-1",
                "refresh_token": "refresh-1",
                "expires_in": 3600,
            }))
            .into_response()
        }
        Some("refresh_token") => Json(json!({
            "access_token": "access-2",
            "expires_in": 3600,
        }))
        .into_response(),
        _ => (
            StatusCode::BAD_REQUEST,
            Json(json!({ "error": "unsupported_grant_type" })),
        )
            .into_response(),
    }
}

#[test]
fn device_flow_polls_pending_then_issues_a_token() {
    let issuer = start_idp();
    let flow = DeviceFlow::discover(&issuer, "heph-cli").unwrap();

    let auth = flow.start("openid").unwrap();
    assert_eq!(auth.user_code, "WDJB-MJHT");
    assert!(auth.verification_uri.contains("/activate"));

    // No real waiting — the injected sleep is a no-op.
    let token = flow.poll(&auth, |_| {}).unwrap();
    assert_eq!(token.access_token, "access-1");
    assert_eq!(token.refresh_token.as_deref(), Some("refresh-1"));
    assert!(token.expires_at > 0);
}

#[test]
fn refresh_keeps_the_old_refresh_token_when_omitted() {
    let issuer = start_idp();
    let flow = DeviceFlow::discover(&issuer, "heph-cli").unwrap();
    let refreshed = flow.refresh("refresh-1").unwrap();
    assert_eq!(refreshed.access_token, "access-2");
    // The provider omitted a new refresh token, so the old one is retained.
    assert_eq!(refreshed.refresh_token.as_deref(), Some("refresh-1"));
}

#[test]
fn memory_token_store_round_trips_and_reports_expiry() {
    let store = MemoryTokenStore::default();
    assert!(store.load().is_none());

    let token = StoredToken {
        access_token: "a".into(),
        refresh_token: Some("r".into()),
        expires_at: 10_000,
    };
    store.save(&token).unwrap();
    assert_eq!(store.load(), Some(token.clone()));

    assert!(!token.is_expired(5_000), "still valid well before expiry");
    assert!(
        token.is_expired(10_000),
        "expired at the boundary (with skew)"
    );

    store.clear().unwrap();
    assert!(store.load().is_none());
}