slave-framework 0.6.0

//! A tiny OpenID Connect provider for tests and local development.
//!
//! **Not secure, never use it in production:** anyone can sign in as anyone,
//! and its signing key is published in this crate. It exists so apps can be
//! tested, and tried locally, without an account at a real provider.
//!
//! It serves discovery, a JWKS, an authorization endpoint with a "sign in as"
//! form (Authorization Code flow, PKCE checked) and a token endpoint, and can
//! mint ID tokens directly for tests.

use std::collections::HashMap;
use std::net::TcpListener;
use std::sync::{Arc, Mutex};
use std::thread::JoinHandle;

use axum::extract::{Form, Query, State};
use axum::http::{HeaderValue, StatusCode, header};
use axum::response::{Html, IntoResponse, Redirect, Response};
use axum::routing::get;
use axum::{Json, Router};
use openidconnect::core::{
    CoreIdToken, CoreIdTokenClaims, CoreJsonWebKeySet, CoreJwsSigningAlgorithm,
    CoreRsaPrivateSigningKey,
};
use openidconnect::{
    Audience, EmptyAdditionalClaims, EndUserEmail, EndUserName, IssuerUrl, JsonWebKeyId,
    LocalizedClaim, Nonce, PkceCodeChallenge, PkceCodeVerifier, PrivateSigningKey, StandardClaims,
    SubjectIdentifier,
};
use serde::Deserialize;
use serde_json::json;

/// The client ID the test provider's tokens are issued for by default.
pub const TEST_CLIENT_ID: &str = "slave-test";

/// Test-only RSA key: it is public, so tokens it signs prove nothing.
const SIGNING_KEY_PEM: &str = include_str!("test_signing_key.pem");
const KEY_ID: &str = "slave-test-key";

/// What to put in a minted ID token.
#[derive(Debug, Clone)]
pub struct TokenSpec {
    pub subject: String,
    pub email: Option<String>,
    pub name: Option<String>,
    /// The `aud` claim (default [`TEST_CLIENT_ID`]).
    pub audience: String,
    /// Seconds until expiry; negative for an already-expired token.
    pub expires_in: i64,
    pub nonce: Option<String>,
}

impl TokenSpec {
    /// A valid token for `email`, whose subject is derived from it.
    pub fn for_email(email: &str) -> Self {
        Self {
            subject: format!("test|{email}"),
            email: Some(email.to_owned()),
            name: email.split('@').next().map(str::to_owned),
            audience: TEST_CLIENT_ID.to_owned(),
            expires_in: 3600,
            nonce: None,
        }
    }
}

struct Provider {
    issuer: String,
    key: CoreRsaPrivateSigningKey,
    /// Authorization codes waiting to be exchanged.
    codes: Mutex<HashMap<String, PendingCode>>,
}

struct PendingCode {
    spec: TokenSpec,
    client_id: String,
    redirect_uri: String,
    code_challenge: Option<String>,
}

/// A running provider, stopped when dropped.
pub struct TestIdentityProvider {
    provider: Arc<Provider>,
    shutdown: Option<tokio::sync::oneshot::Sender<()>>,
    thread: Option<JoinHandle<()>>,
}

impl TestIdentityProvider {
    /// Starts on a free local port.
    pub fn start() -> Self {
        Self::start_on("127.0.0.1:0")
    }

    /// Starts on `addr` (e.g. `127.0.0.1:9000` for a fixed development issuer).
    pub fn start_on(addr: &str) -> Self {
        let listener = TcpListener::bind(addr).expect("bind the test identity provider");
        listener
            .set_nonblocking(true)
            .expect("non-blocking listener");
        let port = listener.local_addr().expect("local address").port();
        let provider = Arc::new(Provider {
            issuer: format!("http://127.0.0.1:{port}"),
            key: CoreRsaPrivateSigningKey::from_pem(
                SIGNING_KEY_PEM,
                Some(JsonWebKeyId::new(KEY_ID.to_owned())),
            )
            .expect("valid test signing key"),
            codes: Mutex::new(HashMap::new()),
        });

        let (shutdown, stopped) = tokio::sync::oneshot::channel::<()>();
        let state = Arc::clone(&provider);
        let thread = std::thread::spawn(move || {
            let runtime = tokio::runtime::Builder::new_current_thread()
                .enable_all()
                .build()
                .expect("tokio runtime");
            runtime.block_on(async move {
                let listener = tokio::net::TcpListener::from_std(listener).expect("tokio listener");
                let _ = axum::serve(listener, router(state))
                    .with_graceful_shutdown(async {
                        let _ = stopped.await;
                    })
                    .await;
            });
        });
        Self {
            provider,
            shutdown: Some(shutdown),
            thread: Some(thread),
        }
    }

    /// The issuer URL (`OIDC_ISSUER_URL`).
    pub fn issuer(&self) -> &str {
        &self.provider.issuer
    }

    /// `OIDC_ISSUER_URL` and `OIDC_AUDIENCES` for processes under test.
    pub fn envs(&self) -> [(&'static str, String); 2] {
        [
            (crate::oidc::ISSUER_ENV, self.provider.issuer.clone()),
            (crate::oidc::AUDIENCES_ENV, TEST_CLIENT_ID.to_owned()),
        ]
    }

    /// A valid ID token for `email`.
    pub fn id_token(&self, email: &str) -> String {
        self.provider.mint(&TokenSpec::for_email(email))
    }

    /// An ID token with exactly these claims.
    pub fn id_token_with(&self, spec: &TokenSpec) -> String {
        self.provider.mint(spec)
    }
}

impl Drop for TestIdentityProvider {
    fn drop(&mut self) {
        if let Some(shutdown) = self.shutdown.take() {
            let _ = shutdown.send(());
        }
        if let Some(thread) = self.thread.take() {
            let _ = thread.join();
        }
    }
}

impl Provider {
    fn mint(&self, spec: &TokenSpec) -> String {
        let now = chrono::Utc::now();
        let mut standard = StandardClaims::new(SubjectIdentifier::new(spec.subject.clone()));
        if let Some(email) = &spec.email {
            standard = standard
                .set_email(Some(EndUserEmail::new(email.clone())))
                .set_email_verified(Some(true));
        }
        if let Some(name) = &spec.name {
            standard =
                standard.set_name(Some(LocalizedClaim::from(EndUserName::new(name.clone()))));
        }
        let claims = CoreIdTokenClaims::new(
            IssuerUrl::new(self.issuer.clone()).expect("issuer URL"),
            vec![Audience::new(spec.audience.clone())],
            now + chrono::Duration::seconds(spec.expires_in),
            now,
            standard,
            EmptyAdditionalClaims {},
        )
        .set_nonce(spec.nonce.clone().map(Nonce::new));
        CoreIdToken::new(
            claims,
            &self.key,
            CoreJwsSigningAlgorithm::RsaSsaPkcs1V15Sha256,
            None,
            None,
        )
        .expect("sign ID token")
        .to_string()
    }
}

fn router(provider: Arc<Provider>) -> Router {
    Router::new()
        .route("/.well-known/openid-configuration", get(discovery))
        .route("/jwks", get(jwks))
        .route("/authorize", get(authorize_form).post(authorize))
        .route("/token", axum::routing::post(token).options(preflight))
        .with_state(provider)
}

async fn discovery(State(provider): State<Arc<Provider>>) -> Response {
    let issuer = &provider.issuer;
    cors(Json(json!({
        "issuer": issuer,
        "authorization_endpoint": format!("{issuer}/authorize"),
        "token_endpoint": format!("{issuer}/token"),
        "jwks_uri": format!("{issuer}/jwks"),
        "response_types_supported": ["code"],
        "subject_types_supported": ["public"],
        "id_token_signing_alg_values_supported": ["RS256"],
        "scopes_supported": ["openid", "email", "profile"],
        "token_endpoint_auth_methods_supported": ["none", "client_secret_post", "client_secret_basic"],
        "code_challenge_methods_supported": ["S256"],
        "claims_supported": ["sub", "email", "email_verified", "name"],
    }))
    .into_response())
}

async fn jwks(State(provider): State<Arc<Provider>>) -> Response {
    cors(
        Json(CoreJsonWebKeySet::new(vec![
            provider.key.as_verification_key(),
        ]))
        .into_response(),
    )
}

#[derive(Deserialize)]
struct AuthorizeParams {
    client_id: String,
    redirect_uri: String,
    #[serde(default)]
    state: Option<String>,
    #[serde(default)]
    nonce: Option<String>,
    #[serde(default)]
    code_challenge: Option<String>,
    #[serde(default)]
    code_challenge_method: Option<String>,
}

#[derive(Deserialize)]
struct SignIn {
    client_id: String,
    redirect_uri: String,
    #[serde(default)]
    state: Option<String>,
    #[serde(default)]
    nonce: Option<String>,
    #[serde(default)]
    code_challenge: Option<String>,
    email: String,
    #[serde(default)]
    name: String,
}

fn escape(text: &str) -> String {
    text.replace('&', "&amp;")
        .replace('"', "&quot;")
        .replace('<', "&lt;")
        .replace('>', "&gt;")
}

/// The "sign in as" page: no password, anyone can be anyone.
async fn authorize_form(Query(params): Query<AuthorizeParams>) -> Response {
    if params
        .code_challenge_method
        .as_deref()
        .is_some_and(|method| method != "S256")
    {
        return (StatusCode::BAD_REQUEST, "only S256 PKCE is supported").into_response();
    }
    let hidden = |name: &str, value: Option<&str>| {
        value.map_or(String::new(), |value| {
            format!(
                r#"<input type="hidden" name="{name}" value="{}">"#,
                escape(value)
            )
        })
    };
    let fields = [
        hidden("client_id", Some(&params.client_id)),
        hidden("redirect_uri", Some(&params.redirect_uri)),
        hidden("state", params.state.as_deref()),
        hidden("nonce", params.nonce.as_deref()),
        hidden("code_challenge", params.code_challenge.as_deref()),
    ]
    .concat();
    Html(format!(
        r#"<!doctype html><html><head><meta charset="utf-8"><title>Development sign-in</title></head>
<body style="font-family: system-ui; max-width: 28rem; margin: 3rem auto">
<h1>Development sign-in</h1>
<p><strong>Not secure.</strong> slave's test identity provider: choose who to be.</p>
<form method="post" action="/authorize" style="display: grid; gap: .5rem">
{fields}
<input name="email" type="email" value="ada@example.com" required>
<input name="name" value="Ada">
<button>Sign in</button>
</form></body></html>"#
    ))
    .into_response()
}

async fn authorize(State(provider): State<Arc<Provider>>, Form(form): Form<SignIn>) -> Response {
    let mut spec = TokenSpec::for_email(&form.email);
    if !form.name.trim().is_empty() {
        spec.name = Some(form.name.trim().to_owned());
    }
    spec.audience = form.client_id.clone();
    spec.nonce = form.nonce;
    let code = openidconnect::CsrfToken::new_random().secret().clone();
    provider.codes.lock().expect("codes lock").insert(
        code.clone(),
        PendingCode {
            spec,
            client_id: form.client_id,
            redirect_uri: form.redirect_uri.clone(),
            code_challenge: form.code_challenge,
        },
    );
    let separator = if form.redirect_uri.contains('?') {
        '&'
    } else {
        '?'
    };
    let mut location = format!("{}{separator}code={code}", form.redirect_uri);
    if let Some(state) = form.state {
        location.push_str("&state=");
        location.push_str(&url_encode(&state));
    }
    Redirect::to(&location).into_response()
}

#[derive(Deserialize)]
struct TokenRequest {
    grant_type: String,
    code: String,
    redirect_uri: String,
    #[serde(default)]
    client_id: Option<String>,
    #[serde(default)]
    code_verifier: Option<String>,
}

async fn token(
    State(provider): State<Arc<Provider>>,
    Form(request): Form<TokenRequest>,
) -> Response {
    let error = |message: &str| {
        cors(
            (
                StatusCode::BAD_REQUEST,
                Json(json!({"error": "invalid_grant", "error_description": message})),
            )
                .into_response(),
        )
    };
    if request.grant_type != "authorization_code" {
        return error("only authorization_code is supported");
    }
    let Some(pending) = provider
        .codes
        .lock()
        .expect("codes lock")
        .remove(&request.code)
    else {
        return error("unknown or used code");
    };
    if pending.redirect_uri != request.redirect_uri
        || request
            .client_id
            .as_deref()
            .is_some_and(|client| client != pending.client_id)
    {
        return error("redirect_uri or client_id does not match");
    }
    if let Some(challenge) = &pending.code_challenge {
        let Some(verifier) = request.code_verifier else {
            return error("code_verifier required");
        };
        let expected =
            PkceCodeChallenge::from_code_verifier_sha256(&PkceCodeVerifier::new(verifier));
        if expected.as_str() != challenge {
            return error("PKCE verification failed");
        }
    }
    let id_token = provider.mint(&pending.spec);
    cors(
        Json(json!({
            "access_token": openidconnect::CsrfToken::new_random().secret(),
            "token_type": "Bearer",
            "expires_in": pending.spec.expires_in,
            "id_token": id_token,
            "scope": "openid email profile",
        }))
        .into_response(),
    )
}

/// Browsers' CORS preflight for the token endpoint (single-page apps).
async fn preflight() -> Response {
    cors(StatusCode::NO_CONTENT.into_response())
}

fn cors(mut response: Response) -> Response {
    let headers = response.headers_mut();
    headers.insert(
        header::ACCESS_CONTROL_ALLOW_ORIGIN,
        HeaderValue::from_static("*"),
    );
    headers.insert(
        header::ACCESS_CONTROL_ALLOW_HEADERS,
        HeaderValue::from_static("content-type, authorization"),
    );
    headers.insert(
        header::ACCESS_CONTROL_ALLOW_METHODS,
        HeaderValue::from_static("GET, POST, OPTIONS"),
    );
    response
}

fn url_encode(value: &str) -> String {
    value
        .bytes()
        .map(|byte| match byte {
            b'A'..=b'Z' | b'a'..=b'z' | b'0'..=b'9' | b'-' | b'_' | b'.' | b'~' => {
                (byte as char).to_string()
            }
            _ => format!("%{byte:02X}"),
        })
        .collect()
}