rustauth-plugins 0.2.0

use super::*;
use time::Duration;

use std::sync::atomic::{AtomicBool, Ordering};

use rustauth_core::context::request_state;
use rustauth_core::options::SessionAdditionalField;
use rustauth_core::plugin::{AuthPlugin, PluginAfterHookAction};

#[tokio::test]
async fn token_route_returns_authorization_pending_before_approval(
) -> Result<(), Box<dyn std::error::Error>> {
    let adapter = Arc::new(TestAdapter::default());
    let router = router(adapter, DeviceAuthorizationOptions::default())?;
    let code = create_device_code(&router, "test-client", None).await?;

    let response = poll_token(&router, string_field(&code, "device_code"), "test-client").await?;

    assert_eq!(response.status(), StatusCode::BAD_REQUEST);
    assert_token_cache_headers(&response);
    let body: Value = serde_json::from_slice(response.body())?;
    assert_eq!(body["error"], "authorization_pending");
    assert_eq!(body["error_description"], "Authorization pending");
    Ok(())
}

#[tokio::test]
async fn token_route_enforces_polling_interval() -> Result<(), Box<dyn std::error::Error>> {
    let adapter = Arc::new(TestAdapter::default());
    let router = router(
        adapter,
        DeviceAuthorizationOptions::new().interval(Duration::seconds(30)),
    )?;
    let code = create_device_code(&router, "test-client", None).await?;
    let device_code = string_field(&code, "device_code");
    let _ = poll_token(&router, device_code, "test-client").await?;

    let response = poll_token(&router, device_code, "test-client").await?;

    assert_eq!(response.status(), StatusCode::BAD_REQUEST);
    assert_token_cache_headers(&response);
    let body: Value = serde_json::from_slice(response.body())?;
    assert_eq!(body["error"], "slow_down");
    assert_eq!(body["error_description"], "Polling too frequently");
    Ok(())
}

#[tokio::test]
async fn token_route_rejects_invalid_device_code_and_grant_type(
) -> Result<(), Box<dyn std::error::Error>> {
    let adapter = Arc::new(TestAdapter::default());
    let router = router(adapter, DeviceAuthorizationOptions::default())?;

    let bad_grant = router
        .handle_async(json_request(
            Method::POST,
            "/api/auth/device/token",
            r#"{"grant_type":"bad","device_code":"missing","client_id":"test-client"}"#,
            None,
        )?)
        .await?;
    assert_eq!(bad_grant.status(), StatusCode::BAD_REQUEST);
    assert_token_cache_headers(&bad_grant);
    let body: Value = serde_json::from_slice(bad_grant.body())?;
    assert_eq!(body["error"], "invalid_request");

    let invalid_code = poll_token(&router, "missing", "test-client").await?;
    assert_eq!(invalid_code.status(), StatusCode::BAD_REQUEST);
    assert_token_cache_headers(&invalid_code);
    let body: Value = serde_json::from_slice(invalid_code.body())?;
    assert_eq!(body["error"], "invalid_grant");
    assert_eq!(body["error_description"], "Invalid device code");
    Ok(())
}

#[tokio::test]
async fn token_route_rejects_invalid_and_mismatched_client(
) -> Result<(), Box<dyn std::error::Error>> {
    let adapter = Arc::new(TestAdapter::default());
    let router = router(
        adapter,
        DeviceAuthorizationOptions::new().validate_client(|client_id| async move {
            Ok(client_id == "valid-client-1" || client_id == "valid-client-2")
        }),
    )?;
    let code = create_device_code(&router, "valid-client-1", None).await?;
    let device_code = string_field(&code, "device_code");

    let invalid_client = poll_token(&router, device_code, "invalid-client").await?;
    assert_token_cache_headers(&invalid_client);
    let body: Value = serde_json::from_slice(invalid_client.body())?;
    assert_eq!(body["error"], "invalid_grant");
    assert_eq!(body["error_description"], "Invalid client ID");

    let mismatched_client = poll_token(&router, device_code, "valid-client-2").await?;
    assert_token_cache_headers(&mismatched_client);
    let body: Value = serde_json::from_slice(mismatched_client.body())?;
    assert_eq!(body["error"], "invalid_grant");
    assert_eq!(body["error_description"], "Client ID mismatch");
    Ok(())
}

#[tokio::test]
async fn token_route_returns_expired_token_and_deletes_record(
) -> Result<(), Box<dyn std::error::Error>> {
    let adapter = Arc::new(TestAdapter::default());
    let router = router(
        adapter.clone(),
        DeviceAuthorizationOptions::new().expires_in(Duration::milliseconds(1)),
    )?;
    let code = create_device_code(&router, "test-client", None).await?;
    tokio::time::sleep(std::time::Duration::from_millis(10)).await;

    let response = poll_token(&router, string_field(&code, "device_code"), "test-client").await?;

    assert_eq!(response.status(), StatusCode::BAD_REQUEST);
    assert_token_cache_headers(&response);
    let body: Value = serde_json::from_slice(response.body())?;
    assert_eq!(body["error"], "expired_token");
    assert_eq!(body["error_description"], "Device code has expired");
    assert_eq!(adapter.len("device_code").await, 0);
    Ok(())
}

#[tokio::test]
async fn token_route_exchanges_approved_code_for_bearer_token_and_scope(
) -> Result<(), Box<dyn std::error::Error>> {
    let adapter = Arc::new(TestAdapter::default());
    let router = router(adapter.clone(), DeviceAuthorizationOptions::default())?;
    let code = create_device_code(&router, "test-client", Some("read write")).await?;
    let (_user_id, cookie) = create_user_session(&adapter).await?;
    approve(&router, string_field(&code, "user_code"), &cookie).await?;

    let response = poll_token(&router, string_field(&code, "device_code"), "test-client").await?;

    assert_eq!(response.status(), StatusCode::OK);
    assert_token_cache_headers(&response);
    let body: Value = serde_json::from_slice(response.body())?;
    assert!(!string_field(&body, "access_token").is_empty());
    assert_eq!(body["token_type"], "Bearer");
    assert_eq!(body["scope"], "read write");
    assert_eq!(adapter.len("device_code").await, 0);
    assert_eq!(adapter.len("session").await, 2);
    Ok(())
}

#[tokio::test]
async fn token_route_rejects_reused_approved_device_code() -> Result<(), Box<dyn std::error::Error>>
{
    let adapter = Arc::new(TestAdapter::default());
    let router = router(adapter.clone(), DeviceAuthorizationOptions::default())?;
    let code = create_device_code(&router, "test-client", None).await?;
    let (_user_id, cookie) = create_user_session(&adapter).await?;
    approve(&router, string_field(&code, "user_code"), &cookie).await?;
    let device_code = string_field(&code, "device_code");

    let first = poll_token(&router, device_code, "test-client").await?;
    assert_eq!(first.status(), StatusCode::OK);

    let second = poll_token(&router, device_code, "test-client").await?;
    assert_eq!(second.status(), StatusCode::BAD_REQUEST);
    assert_token_cache_headers(&second);
    let body: Value = serde_json::from_slice(second.body())?;
    assert_eq!(body["error"], "invalid_grant");
    assert_eq!(body["error_description"], "Invalid device code");
    assert_eq!(adapter.len("device_code").await, 0);
    assert_eq!(adapter.len("session").await, 2);
    Ok(())
}

#[tokio::test]
async fn token_route_allows_only_one_concurrent_approved_exchange(
) -> Result<(), Box<dyn std::error::Error>> {
    let adapter = Arc::new(TestAdapter::default());
    let router = Arc::new(router(
        adapter.clone(),
        DeviceAuthorizationOptions::default(),
    )?);
    let code = create_device_code(&router, "test-client", None).await?;
    let (_user_id, cookie) = create_user_session(&adapter).await?;
    approve(&router, string_field(&code, "user_code"), &cookie).await?;
    let device_code = string_field(&code, "device_code").to_owned();
    let client_id = "test-client".to_owned();

    let router_a = Arc::clone(&router);
    let router_b = Arc::clone(&router);
    let device_code_a = device_code.clone();
    let device_code_b = device_code;
    let client_id_a = client_id.clone();
    let client_id_b = client_id;
    let (response_a, response_b) = tokio::join!(
        poll_token(&router_a, &device_code_a, &client_id_a),
        poll_token(&router_b, &device_code_b, &client_id_b),
    );

    let responses = [response_a?, response_b?];
    let success_count = responses
        .iter()
        .filter(|response| response.status() == StatusCode::OK)
        .count();
    let rejection_count = responses
        .iter()
        .filter(|response| response.status() == StatusCode::BAD_REQUEST)
        .count();
    assert_eq!(success_count, 1);
    assert_eq!(rejection_count, 1);

    let rejected = responses
        .iter()
        .find(|response| response.status() == StatusCode::BAD_REQUEST)
        .ok_or("one rejected response")?;
    assert_token_cache_headers(rejected);
    let body: Value = serde_json::from_slice(rejected.body())?;
    assert_eq!(body["error"], "invalid_grant");
    assert!(
        body["error_description"] == "Device code has already been used"
            || body["error_description"] == "Invalid device code",
        "unexpected rejection description: {}",
        body["error_description"]
    );
    assert_eq!(adapter.len("device_code").await, 0);
    assert_eq!(adapter.len("session").await, 2);
    Ok(())
}

#[tokio::test]
async fn approved_token_session_uses_secondary_storage_when_configured(
) -> Result<(), Box<dyn std::error::Error>> {
    let adapter = Arc::new(TestAdapter::default());
    let storage = Arc::new(TestSecondaryStorage::default());
    let mut auth_options = RustAuthOptions::default().secondary_storage(storage.clone());
    auth_options.session.store_session_in_database = true;
    let router = router_with_rustauth_options(
        adapter.clone(),
        DeviceAuthorizationOptions::default(),
        auth_options,
    )?;
    let code = create_device_code(&router, "test-client", Some("read")).await?;
    let (_user_id, cookie) = create_user_session(&adapter).await?;
    approve(&router, string_field(&code, "user_code"), &cookie).await?;

    let response = poll_token(&router, string_field(&code, "device_code"), "test-client").await?;
    let body: Value = serde_json::from_slice(response.body())?;
    let token = string_field(&body, "access_token");

    assert_eq!(response.status(), StatusCode::OK);
    assert!(storage.value(&format!("session:{token}"))?.is_some());
    assert_eq!(adapter.len("session").await, 2);
    Ok(())
}

#[tokio::test]
async fn token_route_returns_access_denied_after_denial() -> Result<(), Box<dyn std::error::Error>>
{
    let adapter = Arc::new(TestAdapter::default());
    let router = router(adapter.clone(), DeviceAuthorizationOptions::default())?;
    let code = create_device_code(&router, "test-client", None).await?;
    let (_user_id, cookie) = create_user_session(&adapter).await?;
    deny(&router, string_field(&code, "user_code"), &cookie).await?;

    let response = poll_token(&router, string_field(&code, "device_code"), "test-client").await?;

    assert_eq!(response.status(), StatusCode::BAD_REQUEST);
    assert_token_cache_headers(&response);
    let body: Value = serde_json::from_slice(response.body())?;
    assert_eq!(body["error"], "access_denied");
    assert_eq!(body["error_description"], "Access denied");
    assert_eq!(adapter.len("device_code").await, 0);
    Ok(())
}

#[tokio::test]
async fn token_route_accepts_form_body() -> Result<(), Box<dyn std::error::Error>> {
    let adapter = Arc::new(TestAdapter::default());
    let router = router(adapter.clone(), DeviceAuthorizationOptions::default())?;
    let code = create_device_code(&router, "test-client", None).await?;
    let (_user_id, cookie) = create_user_session(&adapter).await?;
    approve(&router, string_field(&code, "user_code"), &cookie).await?;

    let response =
        poll_token_form(&router, string_field(&code, "device_code"), "test-client").await?;

    assert_eq!(response.status(), StatusCode::OK);
    assert_token_cache_headers(&response);
    let body: Value = serde_json::from_slice(response.body())?;
    assert_eq!(body["token_type"], "Bearer");
    Ok(())
}

#[tokio::test]
async fn approved_token_session_records_request_metadata_and_additional_fields(
) -> Result<(), Box<dyn std::error::Error>> {
    let adapter = Arc::new(TestAdapter::default());
    let mut auth_options = RustAuthOptions::default();
    auth_options.session.additional_fields.insert(
        "device_origin".to_owned(),
        SessionAdditionalField::new(DbFieldType::String)
            .generated()
            .default_value(DbValue::String("device-flow".to_owned())),
    );
    let router = router_with_rustauth_options(
        adapter.clone(),
        DeviceAuthorizationOptions::default(),
        auth_options,
    )?;
    let code = create_device_code(&router, "test-client", None).await?;
    let (_user_id, cookie) = create_user_session(&adapter).await?;
    approve(&router, string_field(&code, "user_code"), &cookie).await?;

    let response = poll_token_with_headers(
        &router,
        string_field(&code, "device_code"),
        "test-client",
        "203.0.113.7, 10.0.0.1",
        "DeviceTest/1.0",
    )
    .await?;

    assert_eq!(response.status(), StatusCode::OK);
    let body: Value = serde_json::from_slice(response.body())?;
    let token = string_field(&body, "access_token");
    let session = adapter
        .records("session")
        .await
        .into_iter()
        .find(|record| record.get("token") == Some(&DbValue::String(token.to_owned())))
        .ok_or("missing token-created session")?;
    assert_eq!(
        session.get("ip_address"),
        Some(&DbValue::String("203.0.113.7".to_owned()))
    );
    assert_eq!(
        session.get("user_agent"),
        Some(&DbValue::String("DeviceTest/1.0".to_owned()))
    );
    assert_eq!(
        session.get("device_origin"),
        Some(&DbValue::String("device-flow".to_owned()))
    );
    Ok(())
}

#[tokio::test]
async fn approved_token_records_new_session_for_after_hooks(
) -> Result<(), Box<dyn std::error::Error>> {
    let adapter = Arc::new(TestAdapter::default());
    let observed = Arc::new(AtomicBool::new(false));
    let observed_hook = Arc::clone(&observed);
    let observer = AuthPlugin::new("device-session-observer").with_after_hook(
        "/device/token",
        move |_context, _request, response| {
            if request_state::current_new_session()?.is_some() {
                observed_hook.store(true, Ordering::SeqCst);
            }
            Ok(PluginAfterHookAction::Continue(response))
        },
    );
    let auth_options = RustAuthOptions {
        plugins: vec![
            device_authorization(DeviceAuthorizationOptions::default())?,
            observer,
        ],
        secret: Some(secret().to_owned()),
        base_url: Some("http://localhost:3000".to_owned()),
        advanced: AdvancedOptions {
            disable_csrf_check: true,
            disable_origin_check: true,
            ..AdvancedOptions::default()
        },
        ..RustAuthOptions::default()
    };
    let context = create_auth_context_with_adapter(auth_options, adapter.clone())?;
    let router = AuthRouter::with_async_endpoints(context, Vec::new(), Vec::new())?;
    let code = create_device_code(&router, "test-client", None).await?;
    let (_user_id, cookie) = create_user_session(&adapter).await?;
    approve(&router, string_field(&code, "user_code"), &cookie).await?;

    let response = poll_token(&router, string_field(&code, "device_code"), "test-client").await?;

    assert_eq!(response.status(), StatusCode::OK);
    assert!(observed.load(Ordering::SeqCst));
    Ok(())
}

async fn poll_token(
    router: &AuthRouter,
    device_code: &str,
    client_id: &str,
) -> Result<http::Response<Vec<u8>>, Box<dyn std::error::Error>> {
    let body = format!(
        r#"{{"grant_type":"urn:ietf:params:oauth:grant-type:device_code","device_code":"{device_code}","client_id":"{client_id}"}}"#
    );
    Ok(router
        .handle_async(json_request(
            Method::POST,
            "/api/auth/device/token",
            &body,
            None,
        )?)
        .await?)
}

async fn poll_token_form(
    router: &AuthRouter,
    device_code: &str,
    client_id: &str,
) -> Result<http::Response<Vec<u8>>, Box<dyn std::error::Error>> {
    let body = format!(
        "grant_type=urn%3Aietf%3Aparams%3Aoauth%3Agrant-type%3Adevice_code&device_code={device_code}&client_id={client_id}"
    );
    Ok(router
        .handle_async(form_request(
            Method::POST,
            "/api/auth/device/token",
            &body,
            None,
        )?)
        .await?)
}

async fn poll_token_with_headers(
    router: &AuthRouter,
    device_code: &str,
    client_id: &str,
    ip_address: &str,
    user_agent: &str,
) -> Result<http::Response<Vec<u8>>, Box<dyn std::error::Error>> {
    let body = format!(
        r#"{{"grant_type":"urn:ietf:params:oauth:grant-type:device_code","device_code":"{device_code}","client_id":"{client_id}"}}"#
    );
    Ok(router
        .handle_async(
            Request::builder()
                .method(Method::POST)
                .uri("http://localhost:3000/api/auth/device/token")
                .header(header::CONTENT_TYPE, "application/json")
                .header("x-forwarded-for", ip_address)
                .header(header::USER_AGENT, user_agent)
                .body(body.into_bytes())?,
        )
        .await?)
}

fn assert_token_cache_headers(response: &http::Response<Vec<u8>>) {
    assert_eq!(
        response.headers().get(header::CACHE_CONTROL),
        Some(&http::HeaderValue::from_static("no-store"))
    );
    assert_eq!(
        response.headers().get(header::PRAGMA),
        Some(&http::HeaderValue::from_static("no-cache"))
    );
}

async fn approve(
    router: &AuthRouter,
    user_code: &str,
    cookie: &str,
) -> Result<(), Box<dyn std::error::Error>> {
    let response = router
        .handle_async(json_request(
            Method::POST,
            "/api/auth/device/approve",
            &format!(r#"{{"userCode":"{user_code}"}}"#),
            Some(cookie),
        )?)
        .await?;
    assert_eq!(response.status(), StatusCode::OK);
    Ok(())
}

async fn deny(
    router: &AuthRouter,
    user_code: &str,
    cookie: &str,
) -> Result<(), Box<dyn std::error::Error>> {
    let response = router
        .handle_async(json_request(
            Method::POST,
            "/api/auth/device/deny",
            &format!(r#"{{"userCode":"{user_code}"}}"#),
            Some(cookie),
        )?)
        .await?;
    assert_eq!(response.status(), StatusCode::OK);
    Ok(())
}