spiral-auth 0.1.12

use std::future::IntoFuture;
use std::sync::Arc;
use std::time::Duration;

use axum::extract::{Query, State};
use axum::http::{HeaderMap, StatusCode, header};
use axum::response::{Html, IntoResponse, Response};
use axum::routing::get;
use axum::{Router, serve};
use base64::Engine;
use base64::engine::general_purpose::URL_SAFE_NO_PAD;
use parking_lot::Mutex;
use serde::Deserialize;
use sha2::{Digest as _, Sha256};
use tokio::net::TcpListener;
use tokio::select;
use tokio::sync::oneshot;
use tokio::time::timeout;
use url::Url;
use webbrowser::open;

use super::{DeviceCodeAuth, TokensModel, tokens_from_response};
use crate::tls::crypto_provider;
use crate::{CredentialResult, OAuth2Grant, OAuth2TokenRequest, OAuth2TokenResponse};

const CALLBACK_PATH: &str = "/oauth/callback";
const AUTHORIZATION_TIMEOUT: Duration = Duration::from_secs(300);
const CALLBACK_PAGE: &str = r#"<!doctype html>
<html lang="en">
<meta charset="utf-8">
<meta name="viewport" content="width=device-width, initial-scale=1">
<title>Return to Spiral Studio</title>
<style>
  :root { color-scheme: light dark; font: 16px system-ui, sans-serif; }
  body { margin: 0; min-height: 100vh; display: grid; place-items: center; }
  main { max-width: 28rem; padding: 2rem; text-align: center; }
  h1 { font-size: 1.5rem; font-weight: 600; text-wrap: balance; }
  p { margin: 0; line-height: 1.6; text-wrap: balance; }
</style>
<main>
  <h1>Return to Spiral Studio</h1>
  <p>Continue in the app to finish signing in.</p>
  <p>You can close this tab.</p>
</main>
</html>"#;

type AuthorizationResult = Result<String, &'static str>;

impl DeviceCodeAuth {
    /// Sign in through the system browser using PKCE S256 and an ephemeral loopback callback.
    ///
    /// The client must register `http://127.0.0.1/oauth/callback` at the authorization server.
    /// Dropping this future cancels the listener. Tokens are returned without changing cached
    /// or persisted credentials; the host validates them before calling [`Self::store_tokens`].
    pub async fn authorize_browser(
        &self,
        organization_id: Option<&str>,
    ) -> CredentialResult<TokensModel> {
        self.authorize_browser_with(organization_id, |url| {
            open(url).map_err(|_| invalid!("Could not open browser for sign-in"))
        })
        .await
    }

    async fn authorize_browser_with(
        &self,
        organization_id: Option<&str>,
        open_browser: impl FnOnce(&str) -> CredentialResult<()>,
    ) -> CredentialResult<TokensModel> {
        timeout(AUTHORIZATION_TIMEOUT, async {
            let authorization = BrowserAuthorization::start(
                self.endpoint("/oauth/authorize")?,
                &self.client_id,
                organization_id,
            )
            .await?;
            open_browser(authorization.url.as_str())?;
            let redirect_uri = authorization.redirect_uri.to_string();
            let code_verifier = authorization.verifier.clone();
            let code = authorization.wait().await?;
            let response = self
                .http
                .post(self.endpoint("/oauth/token")?)
                .form(&OAuth2TokenRequest::public(
                    &self.client_id,
                    OAuth2Grant::AuthorizationCode {
                        code,
                        redirect_uri,
                        code_verifier,
                    },
                ))
                .timeout(Duration::from_secs(30))
                .send()
                .await?;
            if !response.status().is_success() {
                bail!("Could not complete browser sign-in");
            }
            tokens_from_response(response.json::<OAuth2TokenResponse>().await?, None)
        })
        .await
        .map_err(|_| invalid!("Timeout waiting for browser sign-in"))?
    }
}

struct BrowserAuthorization {
    listener: TcpListener,
    url: Url,
    redirect_uri: Url,
    verifier: String,
    callback: Arc<CallbackState>,
    result: oneshot::Receiver<AuthorizationResult>,
}

impl BrowserAuthorization {
    async fn start(
        mut url: Url,
        client_id: &str,
        organization_id: Option<&str>,
    ) -> CredentialResult<Self> {
        let listener = TcpListener::bind(("127.0.0.1", 0)).await?;
        let address = listener.local_addr()?;
        let redirect_uri = Url::parse(&format!("http://{address}{CALLBACK_PATH}"))?;
        let verifier = random_token()?;
        let state = random_token()?;
        let (sender, result) = oneshot::channel();
        let callback = Arc::new(CallbackState {
            state: state.clone(),
            issuer: url.origin().ascii_serialization(),
            authority: address.to_string(),
            sender: Mutex::new(Some(sender)),
        });
        url.query_pairs_mut()
            .append_pair("response_type", "code")
            .append_pair("client_id", client_id)
            .append_pair("redirect_uri", redirect_uri.as_str())
            .append_pair("state", &state)
            .append_pair("code_challenge", &pkce_challenge(&verifier))
            .append_pair("code_challenge_method", "S256");
        if let Some(organization_id) = organization_id {
            url.query_pairs_mut()
                .append_pair("organization_id", organization_id);
        }
        Ok(Self {
            listener,
            url,
            redirect_uri,
            verifier,
            callback,
            result,
        })
    }

    async fn wait(self) -> CredentialResult<String> {
        let router = Router::new()
            .route(CALLBACK_PATH, get(callback))
            .with_state(self.callback);
        // Serving stays inside the authorization future so cancellation closes the listener.
        select! {
            result = serve(self.listener, router).into_future() => {
                result?;
                bail!("Browser sign-in listener stopped")
            }
            result = self.result => result
                .map_err(|_| invalid!("Browser sign-in was cancelled"))?
                .map_err(|message| invalid!("{message}")),
        }
    }
}

/// 32 random bytes from the TLS crypto provider, base64url-encoded.
fn random_token() -> CredentialResult<String> {
    let mut bytes = [0_u8; 32];
    crypto_provider()?
        .secure_random
        .fill(&mut bytes)
        .map_err(|_| invalid!("the crypto provider has no random source"))?;
    Ok(URL_SAFE_NO_PAD.encode(bytes))
}

/// RFC 7636 §4.2 S256 code challenge.
fn pkce_challenge(verifier: &str) -> String {
    URL_SAFE_NO_PAD.encode(Sha256::digest(verifier.as_bytes()))
}

struct CallbackState {
    state: String,
    issuer: String,
    authority: String,
    sender: Mutex<Option<oneshot::Sender<AuthorizationResult>>>,
}

#[derive(Deserialize)]
struct AuthorizationResponse {
    state: String,
    iss: String,
    code: Option<String>,
    error: Option<String>,
}

async fn callback(
    State(pending): State<Arc<CallbackState>>,
    headers: HeaderMap,
    Query(response): Query<AuthorizationResponse>,
) -> Response {
    if headers
        .get(header::HOST)
        .and_then(|value| value.to_str().ok())
        != Some(pending.authority.as_str())
        || response.state != pending.state
        || response.iss != pending.issuer
    {
        return (
            StatusCode::BAD_REQUEST,
            "This sign-in request does not match.",
        )
            .into_response();
    }
    let result = match (response.code, response.error) {
        (Some(code), None) if !code.is_empty() => Ok(code),
        (None, Some(error)) if error == "access_denied" => Err("Authorization request was denied"),
        (None, Some(error)) if !error.is_empty() => Err("Browser sign-in failed"),
        _ => return (StatusCode::BAD_REQUEST, "Invalid sign-in response.").into_response(),
    };
    let Some(sender) = pending.sender.lock().take() else {
        return (
            StatusCode::CONFLICT,
            "This sign-in request is already complete.",
        )
            .into_response();
    };
    drop(sender.send(result));
    (
        [
            (header::CACHE_CONTROL, "no-store"),
            (header::REFERRER_POLICY, "no-referrer"),
            (
                header::CONTENT_SECURITY_POLICY,
                "default-src 'none'; style-src 'unsafe-inline'; frame-ancestors 'none'",
            ),
        ],
        Html(CALLBACK_PAGE),
    )
        .into_response()
}

#[cfg(test)]
mod tests {
    use std::collections::HashMap;
    use std::env;

    use serde_json::json;
    use serde_urlencoded::from_bytes;
    use tokio::join;
    use wiremock::matchers::{method, path};
    use wiremock::{Mock, MockServer, Request, ResponseTemplate};

    use super::*;
    use crate::http::test_client;

    #[tokio::test]
    async fn browser_sign_in_binds_the_callback_and_exchanges_pkce_without_saving_tokens() {
        let server = MockServer::start().await;
        let auth_file =
            env::temp_dir().join(format!("spiral-browser-{}.json", random_token().unwrap()));
        let auth = DeviceCodeAuth::new_empty(
            test_client(),
            auth_file.clone(),
            Url::parse(&server.uri()).unwrap(),
            "desktop".to_owned(),
        );
        let (opened, location) = oneshot::channel();
        let authorization = auth.authorize_browser_with(Some("org_test"), |url| {
            opened.send(url.to_owned()).unwrap();
            Ok(())
        });
        let browser = async {
            let url = Url::parse(&location.await.unwrap()).unwrap();
            let query: HashMap<_, _> = url.query_pairs().into_owned().collect();
            assert_eq!(url.path(), "/oauth/authorize");
            assert_eq!(query["client_id"], "desktop");
            assert_eq!(query["response_type"], "code");
            assert_eq!(query["organization_id"], "org_test");
            assert_eq!(query["code_challenge_method"], "S256");
            assert_eq!(query["state"].len(), 43);
            let mut redirect = Url::parse(&query["redirect_uri"]).unwrap();
            assert_eq!(redirect.host_str(), Some("127.0.0.1"));
            assert_eq!(redirect.path(), CALLBACK_PATH);
            assert!(redirect.port().is_some());

            let expected_redirect = query["redirect_uri"].clone();
            let challenge = query["code_challenge"].clone();
            Mock::given(method("POST"))
                .and(path("/oauth/token"))
                .respond_with(move |request: &Request| {
                    let form: HashMap<String, String> = from_bytes(&request.body).unwrap();
                    assert_eq!(form["grant_type"], "authorization_code");
                    assert_eq!(form["client_id"], "desktop");
                    assert_eq!(form["code"], "approved");
                    assert_eq!(form["redirect_uri"], expected_redirect);
                    assert_eq!(pkce_challenge(&form["code_verifier"]), challenge);
                    ResponseTemplate::new(200).set_body_json(json!({
                        "access_token": "access", "refresh_token": "refresh",
                        "token_type": "Bearer", "resource": "https://api.example"
                    }))
                })
                .expect(1)
                .mount(&server)
                .await;

            let client = test_client();
            for (state, issuer) in [
                ("wrong-state", server.uri()),
                (
                    query["state"].as_str(),
                    "https://wrong-issuer.example".to_owned(),
                ),
            ] {
                redirect.set_query(None);
                redirect
                    .query_pairs_mut()
                    .append_pair("state", state)
                    .append_pair("iss", &issuer)
                    .append_pair("code", "untrusted");
                assert_eq!(
                    client.get(redirect.clone()).send().await.unwrap().status(),
                    StatusCode::BAD_REQUEST
                );
            }
            redirect.set_query(None);
            redirect
                .query_pairs_mut()
                .append_pair("state", &query["state"])
                .append_pair("iss", &server.uri())
                .append_pair("code", "approved");
            let mut duplicate = redirect.clone();
            duplicate.query_pairs_mut().append_pair("code", "untrusted");
            assert_eq!(
                client.get(duplicate).send().await.unwrap().status(),
                StatusCode::BAD_REQUEST
            );
            let response = client.get(redirect).send().await.unwrap();
            assert_eq!(response.status(), StatusCode::OK);
            assert_eq!(response.headers()[header::CACHE_CONTROL], "no-store");
            assert!(
                response
                    .text()
                    .await
                    .unwrap()
                    .contains("Return to Spiral Studio")
            );
        };
        let (result, ()) = join!(authorization, browser);
        let tokens = result.unwrap();
        assert_eq!(tokens.access_token.expose_secret(), "access");
        assert_eq!(tokens.refresh_token.expose_secret(), "refresh");
        assert_eq!(tokens.resource.as_deref(), Some("https://api.example"));
        assert!(!auth_file.exists());
    }

    #[tokio::test]
    async fn cancellation_closes_the_loopback_listener() {
        let authorization = BrowserAuthorization::start(
            Url::parse("https://auth.example/oauth/authorize").unwrap(),
            "desktop",
            None,
        )
        .await
        .unwrap();
        let address = authorization.listener.local_addr().unwrap();
        assert!(
            timeout(Duration::from_millis(20), authorization.wait())
                .await
                .is_err()
        );
        let listener = TcpListener::bind(address).await.unwrap();
        drop(listener);
    }

    #[tokio::test]
    async fn browser_denial_returns_actionable_error() {
        let authorization = BrowserAuthorization::start(
            Url::parse("https://auth.example/oauth/authorize").unwrap(),
            "desktop",
            None,
        )
        .await
        .unwrap();
        let mut redirect = authorization.redirect_uri.clone();
        redirect
            .query_pairs_mut()
            .append_pair("state", &authorization.callback.state)
            .append_pair("iss", &authorization.callback.issuer)
            .append_pair("error", "access_denied");
        let client = test_client();
        let (result, response) = join!(authorization.wait(), client.get(redirect).send());
        assert!(result.unwrap_err().to_string().contains("denied"));
        assert_eq!(response.unwrap().status(), StatusCode::OK);
    }
}