tests/oauth.rs at 6f959c6c1b6e9ee99e7f6e49c3fc78a4affe0491 · lewis.moe/bspds-sandbox

lewis.moe / bspds-sandbox
this repo has no description
bspds-sandbox / tests / oauth.rs
at 6f959c6c1b6e9ee99e7f6e49c3fc78a4affe0491 50 kB view raw
   1mod common;
   2mod helpers;
   3use base64::{Engine as _, engine::general_purpose::URL_SAFE_NO_PAD};
   4use chrono::Utc;
   5use common::{base_url, client, create_account_and_login};
   6use reqwest::{redirect, StatusCode};
   7use serde_json::{json, Value};
   8use sha2::{Digest, Sha256};
   9use wiremock::{Mock, MockServer, ResponseTemplate};
  10use wiremock::matchers::{method, path};
  11fn no_redirect_client() -> reqwest::Client {
  12    reqwest::Client::builder()
  13        .redirect(redirect::Policy::none())
  14        .build()
  15        .unwrap()
  16}
  17fn generate_pkce() -> (String, String) {
  18    let verifier_bytes: [u8; 32] = rand::random();
  19    let code_verifier = URL_SAFE_NO_PAD.encode(verifier_bytes);
  20    let mut hasher = Sha256::new();
  21    hasher.update(code_verifier.as_bytes());
  22    let hash = hasher.finalize();
  23    let code_challenge = URL_SAFE_NO_PAD.encode(&hash);
  24    (code_verifier, code_challenge)
  25}
  26async fn setup_mock_client_metadata(redirect_uri: &str) -> MockServer {
  27    let mock_server = MockServer::start().await;
  28    let client_id = mock_server.uri();
  29    let metadata = json!({
  30        "client_id": client_id,
  31        "client_name": "Test OAuth Client",
  32        "redirect_uris": [redirect_uri],
  33        "grant_types": ["authorization_code", "refresh_token"],
  34        "response_types": ["code"],
  35        "token_endpoint_auth_method": "none",
  36        "dpop_bound_access_tokens": false
  37    });
  38    Mock::given(method("GET"))
  39        .and(path("/"))
  40        .respond_with(ResponseTemplate::new(200).set_body_json(metadata))
  41        .mount(&mock_server)
  42        .await;
  43    mock_server
  44}
  45#[allow(dead_code)]
  46async fn setup_mock_dpop_client(redirect_uri: &str) -> MockServer {
  47    let mock_server = MockServer::start().await;
  48    let client_id = mock_server.uri();
  49    let metadata = json!({
  50        "client_id": client_id,
  51        "client_name": "DPoP Test Client",
  52        "redirect_uris": [redirect_uri],
  53        "grant_types": ["authorization_code", "refresh_token"],
  54        "response_types": ["code"],
  55        "token_endpoint_auth_method": "none",
  56        "dpop_bound_access_tokens": true
  57    });
  58    Mock::given(method("GET"))
  59        .and(path("/"))
  60        .respond_with(ResponseTemplate::new(200).set_body_json(metadata))
  61        .mount(&mock_server)
  62        .await;
  63    mock_server
  64}
  65#[tokio::test]
  66async fn test_oauth_protected_resource_metadata() {
  67    let url = base_url().await;
  68    let client = client();
  69    let res = client
  70        .get(format!("{}/.well-known/oauth-protected-resource", url))
  71        .send()
  72        .await
  73        .expect("Failed to fetch protected resource metadata");
  74    assert_eq!(res.status(), StatusCode::OK);
  75    let body: Value = res.json().await.expect("Invalid JSON");
  76    assert!(body["resource"].is_string());
  77    assert!(body["authorization_servers"].is_array());
  78    assert!(body["bearer_methods_supported"].is_array());
  79    let bearer_methods = body["bearer_methods_supported"].as_array().unwrap();
  80    assert!(bearer_methods.contains(&json!("header")));
  81}
  82#[tokio::test]
  83async fn test_oauth_authorization_server_metadata() {
  84    let url = base_url().await;
  85    let client = client();
  86    let res = client
  87        .get(format!("{}/.well-known/oauth-authorization-server", url))
  88        .send()
  89        .await
  90        .expect("Failed to fetch authorization server metadata");
  91    assert_eq!(res.status(), StatusCode::OK);
  92    let body: Value = res.json().await.expect("Invalid JSON");
  93    assert!(body["issuer"].is_string());
  94    assert!(body["authorization_endpoint"].is_string());
  95    assert!(body["token_endpoint"].is_string());
  96    assert!(body["jwks_uri"].is_string());
  97    let response_types = body["response_types_supported"].as_array().unwrap();
  98    assert!(response_types.contains(&json!("code")));
  99    let grant_types = body["grant_types_supported"].as_array().unwrap();
 100    assert!(grant_types.contains(&json!("authorization_code")));
 101    assert!(grant_types.contains(&json!("refresh_token")));
 102    let code_challenge_methods = body["code_challenge_methods_supported"].as_array().unwrap();
 103    assert!(code_challenge_methods.contains(&json!("S256")));
 104    assert_eq!(body["require_pushed_authorization_requests"], json!(true));
 105    let dpop_algs = body["dpop_signing_alg_values_supported"].as_array().unwrap();
 106    assert!(dpop_algs.contains(&json!("ES256")));
 107}
 108#[tokio::test]
 109async fn test_oauth_jwks_endpoint() {
 110    let url = base_url().await;
 111    let client = client();
 112    let res = client
 113        .get(format!("{}/oauth/jwks", url))
 114        .send()
 115        .await
 116        .expect("Failed to fetch JWKS");
 117    assert_eq!(res.status(), StatusCode::OK);
 118    let body: Value = res.json().await.expect("Invalid JSON");
 119    assert!(body["keys"].is_array());
 120}
 121#[tokio::test]
 122async fn test_par_success() {
 123    let url = base_url().await;
 124    let client = client();
 125    let redirect_uri = "https://example.com/callback";
 126    let mock_client = setup_mock_client_metadata(redirect_uri).await;
 127    let client_id = mock_client.uri();
 128    let (_code_verifier, code_challenge) = generate_pkce();
 129    let res = client
 130        .post(format!("{}/oauth/par", url))
 131        .form(&[
 132            ("response_type", "code"),
 133            ("client_id", &client_id),
 134            ("redirect_uri", redirect_uri),
 135            ("code_challenge", &code_challenge),
 136            ("code_challenge_method", "S256"),
 137            ("scope", "atproto"),
 138            ("state", "test-state-123"),
 139        ])
 140        .send()
 141        .await
 142        .expect("Failed to send PAR request");
 143    assert_eq!(res.status(), StatusCode::OK, "PAR should succeed: {:?}", res.text().await);
 144    let body: Value = client
 145        .post(format!("{}/oauth/par", url))
 146        .form(&[
 147            ("response_type", "code"),
 148            ("client_id", &client_id),
 149            ("redirect_uri", redirect_uri),
 150            ("code_challenge", &code_challenge),
 151            ("code_challenge_method", "S256"),
 152            ("scope", "atproto"),
 153            ("state", "test-state-123"),
 154        ])
 155        .send()
 156        .await
 157        .unwrap()
 158        .json()
 159        .await
 160        .expect("Invalid JSON");
 161    assert!(body["request_uri"].is_string());
 162    assert!(body["expires_in"].is_number());
 163    let request_uri = body["request_uri"].as_str().unwrap();
 164    assert!(request_uri.starts_with("urn:ietf:params:oauth:request_uri:"));
 165}
 166#[tokio::test]
 167async fn test_authorize_get_with_valid_request_uri() {
 168    let url = base_url().await;
 169    let client = client();
 170    let redirect_uri = "https://example.com/callback";
 171    let mock_client = setup_mock_client_metadata(redirect_uri).await;
 172    let client_id = mock_client.uri();
 173    let (_, code_challenge) = generate_pkce();
 174    let par_res = client
 175        .post(format!("{}/oauth/par", url))
 176        .form(&[
 177            ("response_type", "code"),
 178            ("client_id", &client_id),
 179            ("redirect_uri", redirect_uri),
 180            ("code_challenge", &code_challenge),
 181            ("code_challenge_method", "S256"),
 182            ("scope", "atproto"),
 183            ("state", "test-state"),
 184        ])
 185        .send()
 186        .await
 187        .expect("PAR failed");
 188    let par_body: Value = par_res.json().await.expect("Invalid PAR JSON");
 189    let request_uri = par_body["request_uri"].as_str().unwrap();
 190    let auth_res = client
 191        .get(format!("{}/oauth/authorize", url))
 192        .header("Accept", "application/json")
 193        .query(&[("request_uri", request_uri)])
 194        .send()
 195        .await
 196        .expect("Authorize GET failed");
 197    assert_eq!(auth_res.status(), StatusCode::OK);
 198    let auth_body: Value = auth_res.json().await.expect("Invalid auth JSON");
 199    assert_eq!(auth_body["client_id"], client_id);
 200    assert_eq!(auth_body["redirect_uri"], redirect_uri);
 201    assert_eq!(auth_body["scope"], "atproto");
 202    assert_eq!(auth_body["state"], "test-state");
 203}
 204#[tokio::test]
 205async fn test_authorize_rejects_invalid_request_uri() {
 206    let url = base_url().await;
 207    let client = client();
 208    let res = client
 209        .get(format!("{}/oauth/authorize", url))
 210        .header("Accept", "application/json")
 211        .query(&[("request_uri", "urn:ietf:params:oauth:request_uri:nonexistent")])
 212        .send()
 213        .await
 214        .expect("Request failed");
 215    assert_eq!(res.status(), StatusCode::BAD_REQUEST);
 216    let body: Value = res.json().await.expect("Invalid JSON");
 217    assert_eq!(body["error"], "invalid_request");
 218}
 219#[tokio::test]
 220async fn test_authorize_requires_request_uri() {
 221    let url = base_url().await;
 222    let client = client();
 223    let res = client
 224        .get(format!("{}/oauth/authorize", url))
 225        .send()
 226        .await
 227        .expect("Request failed");
 228    assert_eq!(res.status(), StatusCode::BAD_REQUEST);
 229}
 230#[tokio::test]
 231async fn test_full_oauth_flow_without_dpop() {
 232    let url = base_url().await;
 233    let http_client = client();
 234    let (_, _user_did) = create_account_and_login(&http_client).await;
 235    let ts = Utc::now().timestamp_millis();
 236    let handle = format!("oauth-test-{}", ts);
 237    let email = format!("oauth-test-{}@example.com", ts);
 238    let password = "oauth-test-password";
 239    let create_res = http_client
 240        .post(format!("{}/xrpc/com.atproto.server.createAccount", url))
 241        .json(&json!({
 242            "handle": handle,
 243            "email": email,
 244            "password": password
 245        }))
 246        .send()
 247        .await
 248        .expect("Account creation failed");
 249    assert_eq!(create_res.status(), StatusCode::OK);
 250    let account: Value = create_res.json().await.unwrap();
 251    let user_did = account["did"].as_str().unwrap();
 252    let redirect_uri = "https://example.com/oauth/callback";
 253    let mock_client = setup_mock_client_metadata(redirect_uri).await;
 254    let client_id = mock_client.uri();
 255    let (code_verifier, code_challenge) = generate_pkce();
 256    let state = format!("state-{}", ts);
 257    let par_res = http_client
 258        .post(format!("{}/oauth/par", url))
 259        .form(&[
 260            ("response_type", "code"),
 261            ("client_id", &client_id),
 262            ("redirect_uri", redirect_uri),
 263            ("code_challenge", &code_challenge),
 264            ("code_challenge_method", "S256"),
 265            ("scope", "atproto"),
 266            ("state", &state),
 267        ])
 268        .send()
 269        .await
 270        .expect("PAR failed");
 271    let par_status = par_res.status();
 272    let par_text = par_res.text().await.unwrap_or_default();
 273    if par_status != StatusCode::OK {
 274        panic!("PAR failed with status {}: {}", par_status, par_text);
 275    }
 276    let par_body: Value = serde_json::from_str(&par_text).unwrap();
 277    let request_uri = par_body["request_uri"].as_str().unwrap();
 278    let auth_client = no_redirect_client();
 279    let auth_res = auth_client
 280        .post(format!("{}/oauth/authorize", url))
 281        .form(&[
 282            ("request_uri", request_uri),
 283            ("username", &handle),
 284            ("password", password),
 285            ("remember_device", "false"),
 286        ])
 287        .send()
 288        .await
 289        .expect("Authorize POST failed");
 290    let auth_status = auth_res.status();
 291    if auth_status != StatusCode::TEMPORARY_REDIRECT
 292        && auth_status != StatusCode::SEE_OTHER
 293        && auth_status != StatusCode::FOUND
 294    {
 295        let auth_text = auth_res.text().await.unwrap_or_default();
 296        panic!(
 297            "Expected redirect, got {}: {}",
 298            auth_status, auth_text
 299        );
 300    }
 301    let location = auth_res.headers().get("location")
 302        .expect("No Location header")
 303        .to_str()
 304        .unwrap();
 305    assert!(location.starts_with(redirect_uri), "Redirect to wrong URI: {}", location);
 306    assert!(location.contains("code="), "No code in redirect: {}", location);
 307    assert!(location.contains(&format!("state={}", state)), "Wrong state in redirect");
 308    let code = location
 309        .split("code=")
 310        .nth(1)
 311        .unwrap()
 312        .split('&')
 313        .next()
 314        .unwrap();
 315    let token_res = http_client
 316        .post(format!("{}/oauth/token", url))
 317        .form(&[
 318            ("grant_type", "authorization_code"),
 319            ("code", code),
 320            ("redirect_uri", redirect_uri),
 321            ("code_verifier", &code_verifier),
 322            ("client_id", &client_id),
 323        ])
 324        .send()
 325        .await
 326        .expect("Token request failed");
 327    let token_status = token_res.status();
 328    let token_text = token_res.text().await.unwrap_or_default();
 329    if token_status != StatusCode::OK {
 330        panic!("Token request failed with status {}: {}", token_status, token_text);
 331    }
 332    let token_body: Value = serde_json::from_str(&token_text).unwrap();
 333    assert!(token_body["access_token"].is_string());
 334    assert!(token_body["refresh_token"].is_string());
 335    assert_eq!(token_body["token_type"], "Bearer");
 336    assert!(token_body["expires_in"].is_number());
 337    assert_eq!(token_body["sub"], user_did);
 338}
 339#[tokio::test]
 340async fn test_token_refresh_flow() {
 341    let url = base_url().await;
 342    let http_client = client();
 343    let ts = Utc::now().timestamp_millis();
 344    let handle = format!("refresh-test-{}", ts);
 345    let email = format!("refresh-test-{}@example.com", ts);
 346    let password = "refresh-test-password";
 347    http_client
 348        .post(format!("{}/xrpc/com.atproto.server.createAccount", url))
 349        .json(&json!({
 350            "handle": handle,
 351            "email": email,
 352            "password": password
 353        }))
 354        .send()
 355        .await
 356        .expect("Account creation failed");
 357    let redirect_uri = "https://example.com/refresh-callback";
 358    let mock_client = setup_mock_client_metadata(redirect_uri).await;
 359    let client_id = mock_client.uri();
 360    let (code_verifier, code_challenge) = generate_pkce();
 361    let par_body: Value = http_client
 362        .post(format!("{}/oauth/par", url))
 363        .form(&[
 364            ("response_type", "code"),
 365            ("client_id", &client_id),
 366            ("redirect_uri", redirect_uri),
 367            ("code_challenge", &code_challenge),
 368            ("code_challenge_method", "S256"),
 369        ])
 370        .send()
 371        .await
 372        .unwrap()
 373        .json()
 374        .await
 375        .unwrap();
 376    let request_uri = par_body["request_uri"].as_str().unwrap();
 377    let auth_client = no_redirect_client();
 378    let auth_res = auth_client
 379        .post(format!("{}/oauth/authorize", url))
 380        .form(&[
 381            ("request_uri", request_uri),
 382            ("username", &handle),
 383            ("password", password),
 384            ("remember_device", "false"),
 385        ])
 386        .send()
 387        .await
 388        .unwrap();
 389    let location = auth_res.headers().get("location").unwrap().to_str().unwrap();
 390    let code = location.split("code=").nth(1).unwrap().split('&').next().unwrap();
 391    let token_body: Value = http_client
 392        .post(format!("{}/oauth/token", url))
 393        .form(&[
 394            ("grant_type", "authorization_code"),
 395            ("code", code),
 396            ("redirect_uri", redirect_uri),
 397            ("code_verifier", &code_verifier),
 398            ("client_id", &client_id),
 399        ])
 400        .send()
 401        .await
 402        .unwrap()
 403        .json()
 404        .await
 405        .unwrap();
 406    let refresh_token = token_body["refresh_token"].as_str().unwrap();
 407    let original_access_token = token_body["access_token"].as_str().unwrap();
 408    let refresh_res = http_client
 409        .post(format!("{}/oauth/token", url))
 410        .form(&[
 411            ("grant_type", "refresh_token"),
 412            ("refresh_token", refresh_token),
 413            ("client_id", &client_id),
 414        ])
 415        .send()
 416        .await
 417        .expect("Refresh request failed");
 418    assert_eq!(refresh_res.status(), StatusCode::OK);
 419    let refresh_body: Value = refresh_res.json().await.unwrap();
 420    assert!(refresh_body["access_token"].is_string());
 421    assert!(refresh_body["refresh_token"].is_string());
 422    let new_access_token = refresh_body["access_token"].as_str().unwrap();
 423    let new_refresh_token = refresh_body["refresh_token"].as_str().unwrap();
 424    assert_ne!(new_access_token, original_access_token, "Access token should rotate");
 425    assert_ne!(new_refresh_token, refresh_token, "Refresh token should rotate");
 426}
 427#[tokio::test]
 428async fn test_wrong_credentials_denied() {
 429    let url = base_url().await;
 430    let http_client = client();
 431    let ts = Utc::now().timestamp_millis();
 432    let handle = format!("wrong-creds-{}", ts);
 433    let email = format!("wrong-creds-{}@example.com", ts);
 434    let password = "correct-password";
 435    http_client
 436        .post(format!("{}/xrpc/com.atproto.server.createAccount", url))
 437        .json(&json!({
 438            "handle": handle,
 439            "email": email,
 440            "password": password
 441        }))
 442        .send()
 443        .await
 444        .unwrap();
 445    let redirect_uri = "https://example.com/wrong-creds-callback";
 446    let mock_client = setup_mock_client_metadata(redirect_uri).await;
 447    let client_id = mock_client.uri();
 448    let (_, code_challenge) = generate_pkce();
 449    let par_body: Value = http_client
 450        .post(format!("{}/oauth/par", url))
 451        .form(&[
 452            ("response_type", "code"),
 453            ("client_id", &client_id),
 454            ("redirect_uri", redirect_uri),
 455            ("code_challenge", &code_challenge),
 456            ("code_challenge_method", "S256"),
 457        ])
 458        .send()
 459        .await
 460        .unwrap()
 461        .json()
 462        .await
 463        .unwrap();
 464    let request_uri = par_body["request_uri"].as_str().unwrap();
 465    let auth_res = http_client
 466        .post(format!("{}/oauth/authorize", url))
 467        .header("Accept", "application/json")
 468        .form(&[
 469            ("request_uri", request_uri),
 470            ("username", &handle),
 471            ("password", "wrong-password"),
 472            ("remember_device", "false"),
 473        ])
 474        .send()
 475        .await
 476        .unwrap();
 477    assert_eq!(auth_res.status(), StatusCode::FORBIDDEN);
 478    let error_body: Value = auth_res.json().await.unwrap();
 479    assert_eq!(error_body["error"], "access_denied");
 480}
 481#[tokio::test]
 482async fn test_token_revocation() {
 483    let url = base_url().await;
 484    let http_client = client();
 485    let ts = Utc::now().timestamp_millis();
 486    let handle = format!("revoke-test-{}", ts);
 487    let email = format!("revoke-test-{}@example.com", ts);
 488    let password = "revoke-test-password";
 489    http_client
 490        .post(format!("{}/xrpc/com.atproto.server.createAccount", url))
 491        .json(&json!({
 492            "handle": handle,
 493            "email": email,
 494            "password": password
 495        }))
 496        .send()
 497        .await
 498        .unwrap();
 499    let redirect_uri = "https://example.com/revoke-callback";
 500    let mock_client = setup_mock_client_metadata(redirect_uri).await;
 501    let client_id = mock_client.uri();
 502    let (code_verifier, code_challenge) = generate_pkce();
 503    let par_body: Value = http_client
 504        .post(format!("{}/oauth/par", url))
 505        .form(&[
 506            ("response_type", "code"),
 507            ("client_id", &client_id),
 508            ("redirect_uri", redirect_uri),
 509            ("code_challenge", &code_challenge),
 510            ("code_challenge_method", "S256"),
 511        ])
 512        .send()
 513        .await
 514        .unwrap()
 515        .json()
 516        .await
 517        .unwrap();
 518    let request_uri = par_body["request_uri"].as_str().unwrap();
 519    let auth_client = no_redirect_client();
 520    let auth_res = auth_client
 521        .post(format!("{}/oauth/authorize", url))
 522        .form(&[
 523            ("request_uri", request_uri),
 524            ("username", &handle),
 525            ("password", password),
 526            ("remember_device", "false"),
 527        ])
 528        .send()
 529        .await
 530        .unwrap();
 531    let location = auth_res.headers().get("location").unwrap().to_str().unwrap();
 532    let code = location.split("code=").nth(1).unwrap().split('&').next().unwrap();
 533    let token_body: Value = http_client
 534        .post(format!("{}/oauth/token", url))
 535        .form(&[
 536            ("grant_type", "authorization_code"),
 537            ("code", code),
 538            ("redirect_uri", redirect_uri),
 539            ("code_verifier", &code_verifier),
 540            ("client_id", &client_id),
 541        ])
 542        .send()
 543        .await
 544        .unwrap()
 545        .json()
 546        .await
 547        .unwrap();
 548    let refresh_token = token_body["refresh_token"].as_str().unwrap();
 549    let revoke_res = http_client
 550        .post(format!("{}/oauth/revoke", url))
 551        .form(&[("token", refresh_token)])
 552        .send()
 553        .await
 554        .unwrap();
 555    assert_eq!(revoke_res.status(), StatusCode::OK);
 556    let refresh_after_revoke = http_client
 557        .post(format!("{}/oauth/token", url))
 558        .form(&[
 559            ("grant_type", "refresh_token"),
 560            ("refresh_token", refresh_token),
 561            ("client_id", &client_id),
 562        ])
 563        .send()
 564        .await
 565        .unwrap();
 566    assert_eq!(refresh_after_revoke.status(), StatusCode::BAD_REQUEST);
 567}
 568#[tokio::test]
 569async fn test_unsupported_grant_type() {
 570    let url = base_url().await;
 571    let http_client = client();
 572    let res = http_client
 573        .post(format!("{}/oauth/token", url))
 574        .form(&[
 575            ("grant_type", "client_credentials"),
 576            ("client_id", "https://example.com"),
 577        ])
 578        .send()
 579        .await
 580        .unwrap();
 581    assert_eq!(res.status(), StatusCode::BAD_REQUEST);
 582    let body: Value = res.json().await.unwrap();
 583    assert_eq!(body["error"], "unsupported_grant_type");
 584}
 585#[tokio::test]
 586async fn test_invalid_refresh_token() {
 587    let url = base_url().await;
 588    let http_client = client();
 589    let res = http_client
 590        .post(format!("{}/oauth/token", url))
 591        .form(&[
 592            ("grant_type", "refresh_token"),
 593            ("refresh_token", "invalid-refresh-token"),
 594            ("client_id", "https://example.com"),
 595        ])
 596        .send()
 597        .await
 598        .unwrap();
 599    assert_eq!(res.status(), StatusCode::BAD_REQUEST);
 600    let body: Value = res.json().await.unwrap();
 601    assert_eq!(body["error"], "invalid_grant");
 602}
 603#[tokio::test]
 604async fn test_expired_authorization_request() {
 605    let url = base_url().await;
 606    let http_client = client();
 607    let res = http_client
 608        .get(format!("{}/oauth/authorize", url))
 609        .header("Accept", "application/json")
 610        .query(&[("request_uri", "urn:ietf:params:oauth:request_uri:expired-or-nonexistent")])
 611        .send()
 612        .await
 613        .unwrap();
 614    assert_eq!(res.status(), StatusCode::BAD_REQUEST);
 615    let body: Value = res.json().await.unwrap();
 616    assert_eq!(body["error"], "invalid_request");
 617}
 618#[tokio::test]
 619async fn test_token_introspection() {
 620    let url = base_url().await;
 621    let http_client = client();
 622    let ts = Utc::now().timestamp_millis();
 623    let handle = format!("introspect-{}", ts);
 624    let email = format!("introspect-{}@example.com", ts);
 625    let password = "introspect-password";
 626    http_client
 627        .post(format!("{}/xrpc/com.atproto.server.createAccount", url))
 628        .json(&json!({
 629            "handle": handle,
 630            "email": email,
 631            "password": password
 632        }))
 633        .send()
 634        .await
 635        .unwrap();
 636    let redirect_uri = "https://example.com/introspect-callback";
 637    let mock_client = setup_mock_client_metadata(redirect_uri).await;
 638    let client_id = mock_client.uri();
 639    let (code_verifier, code_challenge) = generate_pkce();
 640    let par_body: Value = http_client
 641        .post(format!("{}/oauth/par", url))
 642        .form(&[
 643            ("response_type", "code"),
 644            ("client_id", &client_id),
 645            ("redirect_uri", redirect_uri),
 646            ("code_challenge", &code_challenge),
 647            ("code_challenge_method", "S256"),
 648        ])
 649        .send()
 650        .await
 651        .unwrap()
 652        .json()
 653        .await
 654        .unwrap();
 655    let request_uri = par_body["request_uri"].as_str().unwrap();
 656    let auth_client = no_redirect_client();
 657    let auth_res = auth_client
 658        .post(format!("{}/oauth/authorize", url))
 659        .form(&[
 660            ("request_uri", request_uri),
 661            ("username", &handle),
 662            ("password", password),
 663            ("remember_device", "false"),
 664        ])
 665        .send()
 666        .await
 667        .unwrap();
 668    let location = auth_res.headers().get("location").unwrap().to_str().unwrap();
 669    let code = location.split("code=").nth(1).unwrap().split('&').next().unwrap();
 670    let token_body: Value = http_client
 671        .post(format!("{}/oauth/token", url))
 672        .form(&[
 673            ("grant_type", "authorization_code"),
 674            ("code", code),
 675            ("redirect_uri", redirect_uri),
 676            ("code_verifier", &code_verifier),
 677            ("client_id", &client_id),
 678        ])
 679        .send()
 680        .await
 681        .unwrap()
 682        .json()
 683        .await
 684        .unwrap();
 685    let access_token = token_body["access_token"].as_str().unwrap();
 686    let introspect_res = http_client
 687        .post(format!("{}/oauth/introspect", url))
 688        .form(&[("token", access_token)])
 689        .send()
 690        .await
 691        .unwrap();
 692    assert_eq!(introspect_res.status(), StatusCode::OK);
 693    let introspect_body: Value = introspect_res.json().await.unwrap();
 694    assert_eq!(introspect_body["active"], true);
 695    assert!(introspect_body["client_id"].is_string());
 696    assert!(introspect_body["exp"].is_number());
 697}
 698#[tokio::test]
 699async fn test_introspect_invalid_token() {
 700    let url = base_url().await;
 701    let http_client = client();
 702    let res = http_client
 703        .post(format!("{}/oauth/introspect", url))
 704        .form(&[("token", "invalid.token.here")])
 705        .send()
 706        .await
 707        .unwrap();
 708    assert_eq!(res.status(), StatusCode::OK);
 709    let body: Value = res.json().await.unwrap();
 710    assert_eq!(body["active"], false);
 711}
 712#[tokio::test]
 713async fn test_introspect_revoked_token() {
 714    let url = base_url().await;
 715    let http_client = client();
 716    let ts = Utc::now().timestamp_millis();
 717    let handle = format!("introspect-revoked-{}", ts);
 718    let email = format!("introspect-revoked-{}@example.com", ts);
 719    let password = "introspect-revoked-password";
 720    http_client
 721        .post(format!("{}/xrpc/com.atproto.server.createAccount", url))
 722        .json(&json!({
 723            "handle": handle,
 724            "email": email,
 725            "password": password
 726        }))
 727        .send()
 728        .await
 729        .unwrap();
 730    let redirect_uri = "https://example.com/introspect-revoked-callback";
 731    let mock_client = setup_mock_client_metadata(redirect_uri).await;
 732    let client_id = mock_client.uri();
 733    let (code_verifier, code_challenge) = generate_pkce();
 734    let par_body: Value = http_client
 735        .post(format!("{}/oauth/par", url))
 736        .form(&[
 737            ("response_type", "code"),
 738            ("client_id", &client_id),
 739            ("redirect_uri", redirect_uri),
 740            ("code_challenge", &code_challenge),
 741            ("code_challenge_method", "S256"),
 742        ])
 743        .send()
 744        .await
 745        .unwrap()
 746        .json()
 747        .await
 748        .unwrap();
 749    let request_uri = par_body["request_uri"].as_str().unwrap();
 750    let auth_client = no_redirect_client();
 751    let auth_res = auth_client
 752        .post(format!("{}/oauth/authorize", url))
 753        .form(&[
 754            ("request_uri", request_uri),
 755            ("username", &handle),
 756            ("password", password),
 757            ("remember_device", "false"),
 758        ])
 759        .send()
 760        .await
 761        .unwrap();
 762    let location = auth_res.headers().get("location").unwrap().to_str().unwrap();
 763    let code = location.split("code=").nth(1).unwrap().split('&').next().unwrap();
 764    let token_body: Value = http_client
 765        .post(format!("{}/oauth/token", url))
 766        .form(&[
 767            ("grant_type", "authorization_code"),
 768            ("code", code),
 769            ("redirect_uri", redirect_uri),
 770            ("code_verifier", &code_verifier),
 771            ("client_id", &client_id),
 772        ])
 773        .send()
 774        .await
 775        .unwrap()
 776        .json()
 777        .await
 778        .unwrap();
 779    let access_token = token_body["access_token"].as_str().unwrap();
 780    let refresh_token = token_body["refresh_token"].as_str().unwrap();
 781    http_client
 782        .post(format!("{}/oauth/revoke", url))
 783        .form(&[("token", refresh_token)])
 784        .send()
 785        .await
 786        .unwrap();
 787    let introspect_res = http_client
 788        .post(format!("{}/oauth/introspect", url))
 789        .form(&[("token", access_token)])
 790        .send()
 791        .await
 792        .unwrap();
 793    assert_eq!(introspect_res.status(), StatusCode::OK);
 794    let body: Value = introspect_res.json().await.unwrap();
 795    assert_eq!(body["active"], false, "Revoked token should be inactive");
 796}
 797#[tokio::test]
 798async fn test_state_with_special_chars() {
 799    let url = base_url().await;
 800    let http_client = client();
 801    let ts = Utc::now().timestamp_millis();
 802    let handle = format!("state-special-{}", ts);
 803    let email = format!("state-special-{}@example.com", ts);
 804    let password = "state-special-password";
 805    http_client
 806        .post(format!("{}/xrpc/com.atproto.server.createAccount", url))
 807        .json(&json!({
 808            "handle": handle,
 809            "email": email,
 810            "password": password
 811        }))
 812        .send()
 813        .await
 814        .unwrap();
 815    let redirect_uri = "https://example.com/state-special-callback";
 816    let mock_client = setup_mock_client_metadata(redirect_uri).await;
 817    let client_id = mock_client.uri();
 818    let (_code_verifier, code_challenge) = generate_pkce();
 819    let special_state = "state=with&special=chars&plus+more";
 820    let par_body: Value = http_client
 821        .post(format!("{}/oauth/par", url))
 822        .form(&[
 823            ("response_type", "code"),
 824            ("client_id", &client_id),
 825            ("redirect_uri", redirect_uri),
 826            ("code_challenge", &code_challenge),
 827            ("code_challenge_method", "S256"),
 828            ("state", special_state),
 829        ])
 830        .send()
 831        .await
 832        .unwrap()
 833        .json()
 834        .await
 835        .unwrap();
 836    let request_uri = par_body["request_uri"].as_str().unwrap();
 837    let auth_client = no_redirect_client();
 838    let auth_res = auth_client
 839        .post(format!("{}/oauth/authorize", url))
 840        .form(&[
 841            ("request_uri", request_uri),
 842            ("username", &handle),
 843            ("password", password),
 844            ("remember_device", "false"),
 845        ])
 846        .send()
 847        .await
 848        .unwrap();
 849    assert!(
 850        auth_res.status().is_redirection(),
 851        "Should redirect even with special chars in state"
 852    );
 853    let location = auth_res.headers().get("location").unwrap().to_str().unwrap();
 854    assert!(location.contains("state="), "State should be in redirect URL");
 855    let encoded_state = urlencoding::encode(special_state);
 856    assert!(
 857        location.contains(&format!("state={}", encoded_state)),
 858        "State should be URL-encoded. Got: {}",
 859        location
 860    );
 861}
 862#[tokio::test]
 863async fn test_2fa_required_when_enabled() {
 864    let url = base_url().await;
 865    let http_client = client();
 866    let ts = Utc::now().timestamp_millis();
 867    let handle = format!("2fa-required-{}", ts);
 868    let email = format!("2fa-required-{}@example.com", ts);
 869    let password = "2fa-test-password";
 870    let create_res = http_client
 871        .post(format!("{}/xrpc/com.atproto.server.createAccount", url))
 872        .json(&json!({
 873            "handle": handle,
 874            "email": email,
 875            "password": password
 876        }))
 877        .send()
 878        .await
 879        .unwrap();
 880    assert_eq!(create_res.status(), StatusCode::OK);
 881    let account: Value = create_res.json().await.unwrap();
 882    let user_did = account["did"].as_str().unwrap();
 883    let db_url = common::get_db_connection_string().await;
 884    let pool = sqlx::postgres::PgPoolOptions::new()
 885        .max_connections(1)
 886        .connect(&db_url)
 887        .await
 888        .expect("Failed to connect to database");
 889    sqlx::query("UPDATE users SET two_factor_enabled = true WHERE did = $1")
 890        .bind(user_did)
 891        .execute(&pool)
 892        .await
 893        .expect("Failed to enable 2FA");
 894    let redirect_uri = "https://example.com/2fa-callback";
 895    let mock_client = setup_mock_client_metadata(redirect_uri).await;
 896    let client_id = mock_client.uri();
 897    let (_, code_challenge) = generate_pkce();
 898    let par_body: Value = http_client
 899        .post(format!("{}/oauth/par", url))
 900        .form(&[
 901            ("response_type", "code"),
 902            ("client_id", &client_id),
 903            ("redirect_uri", redirect_uri),
 904            ("code_challenge", &code_challenge),
 905            ("code_challenge_method", "S256"),
 906        ])
 907        .send()
 908        .await
 909        .unwrap()
 910        .json()
 911        .await
 912        .unwrap();
 913    let request_uri = par_body["request_uri"].as_str().unwrap();
 914    let auth_client = no_redirect_client();
 915    let auth_res = auth_client
 916        .post(format!("{}/oauth/authorize", url))
 917        .form(&[
 918            ("request_uri", request_uri),
 919            ("username", &handle),
 920            ("password", password),
 921            ("remember_device", "false"),
 922        ])
 923        .send()
 924        .await
 925        .unwrap();
 926    assert!(
 927        auth_res.status().is_redirection(),
 928        "Should redirect to 2FA page, got status: {}",
 929        auth_res.status()
 930    );
 931    let location = auth_res.headers().get("location").unwrap().to_str().unwrap();
 932    assert!(
 933        location.contains("/oauth/authorize/2fa"),
 934        "Should redirect to 2FA page, got: {}",
 935        location
 936    );
 937    assert!(
 938        location.contains("request_uri="),
 939        "2FA redirect should include request_uri"
 940    );
 941}
 942#[tokio::test]
 943async fn test_2fa_invalid_code_rejected() {
 944    let url = base_url().await;
 945    let http_client = client();
 946    let ts = Utc::now().timestamp_millis();
 947    let handle = format!("2fa-invalid-{}", ts);
 948    let email = format!("2fa-invalid-{}@example.com", ts);
 949    let password = "2fa-test-password";
 950    let create_res = http_client
 951        .post(format!("{}/xrpc/com.atproto.server.createAccount", url))
 952        .json(&json!({
 953            "handle": handle,
 954            "email": email,
 955            "password": password
 956        }))
 957        .send()
 958        .await
 959        .unwrap();
 960    assert_eq!(create_res.status(), StatusCode::OK);
 961    let account: Value = create_res.json().await.unwrap();
 962    let user_did = account["did"].as_str().unwrap();
 963    let db_url = common::get_db_connection_string().await;
 964    let pool = sqlx::postgres::PgPoolOptions::new()
 965        .max_connections(1)
 966        .connect(&db_url)
 967        .await
 968        .expect("Failed to connect to database");
 969    sqlx::query("UPDATE users SET two_factor_enabled = true WHERE did = $1")
 970        .bind(user_did)
 971        .execute(&pool)
 972        .await
 973        .expect("Failed to enable 2FA");
 974    let redirect_uri = "https://example.com/2fa-invalid-callback";
 975    let mock_client = setup_mock_client_metadata(redirect_uri).await;
 976    let client_id = mock_client.uri();
 977    let (_, code_challenge) = generate_pkce();
 978    let par_body: Value = http_client
 979        .post(format!("{}/oauth/par", url))
 980        .form(&[
 981            ("response_type", "code"),
 982            ("client_id", &client_id),
 983            ("redirect_uri", redirect_uri),
 984            ("code_challenge", &code_challenge),
 985            ("code_challenge_method", "S256"),
 986        ])
 987        .send()
 988        .await
 989        .unwrap()
 990        .json()
 991        .await
 992        .unwrap();
 993    let request_uri = par_body["request_uri"].as_str().unwrap();
 994    let auth_client = no_redirect_client();
 995    let auth_res = auth_client
 996        .post(format!("{}/oauth/authorize", url))
 997        .form(&[
 998            ("request_uri", request_uri),
 999            ("username", &handle),
1000            ("password", password),
1001            ("remember_device", "false"),
1002        ])
1003        .send()
1004        .await
1005        .unwrap();
1006    assert!(auth_res.status().is_redirection());
1007    let location = auth_res.headers().get("location").unwrap().to_str().unwrap();
1008    assert!(location.contains("/oauth/authorize/2fa"));
1009    let twofa_res = http_client
1010        .post(format!("{}/oauth/authorize/2fa", url))
1011        .form(&[
1012            ("request_uri", request_uri),
1013            ("code", "000000"),
1014        ])
1015        .send()
1016        .await
1017        .unwrap();
1018    assert_eq!(twofa_res.status(), StatusCode::OK);
1019    let body = twofa_res.text().await.unwrap();
1020    assert!(
1021        body.contains("Invalid verification code") || body.contains("invalid"),
1022        "Should show error for invalid code"
1023    );
1024}
1025#[tokio::test]
1026async fn test_2fa_valid_code_completes_auth() {
1027    let url = base_url().await;
1028    let http_client = client();
1029    let ts = Utc::now().timestamp_millis();
1030    let handle = format!("2fa-valid-{}", ts);
1031    let email = format!("2fa-valid-{}@example.com", ts);
1032    let password = "2fa-test-password";
1033    let create_res = http_client
1034        .post(format!("{}/xrpc/com.atproto.server.createAccount", url))
1035        .json(&json!({
1036            "handle": handle,
1037            "email": email,
1038            "password": password
1039        }))
1040        .send()
1041        .await
1042        .unwrap();
1043    assert_eq!(create_res.status(), StatusCode::OK);
1044    let account: Value = create_res.json().await.unwrap();
1045    let user_did = account["did"].as_str().unwrap();
1046    let db_url = common::get_db_connection_string().await;
1047    let pool = sqlx::postgres::PgPoolOptions::new()
1048        .max_connections(1)
1049        .connect(&db_url)
1050        .await
1051        .expect("Failed to connect to database");
1052    sqlx::query("UPDATE users SET two_factor_enabled = true WHERE did = $1")
1053        .bind(user_did)
1054        .execute(&pool)
1055        .await
1056        .expect("Failed to enable 2FA");
1057    let redirect_uri = "https://example.com/2fa-valid-callback";
1058    let mock_client = setup_mock_client_metadata(redirect_uri).await;
1059    let client_id = mock_client.uri();
1060    let (code_verifier, code_challenge) = generate_pkce();
1061    let par_body: Value = http_client
1062        .post(format!("{}/oauth/par", url))
1063        .form(&[
1064            ("response_type", "code"),
1065            ("client_id", &client_id),
1066            ("redirect_uri", redirect_uri),
1067            ("code_challenge", &code_challenge),
1068            ("code_challenge_method", "S256"),
1069        ])
1070        .send()
1071        .await
1072        .unwrap()
1073        .json()
1074        .await
1075        .unwrap();
1076    let request_uri = par_body["request_uri"].as_str().unwrap();
1077    let auth_client = no_redirect_client();
1078    let auth_res = auth_client
1079        .post(format!("{}/oauth/authorize", url))
1080        .form(&[
1081            ("request_uri", request_uri),
1082            ("username", &handle),
1083            ("password", password),
1084            ("remember_device", "false"),
1085        ])
1086        .send()
1087        .await
1088        .unwrap();
1089    assert!(auth_res.status().is_redirection());
1090    let twofa_code: String = sqlx::query_scalar(
1091        "SELECT code FROM oauth_2fa_challenge WHERE request_uri = $1"
1092    )
1093    .bind(request_uri)
1094    .fetch_one(&pool)
1095    .await
1096    .expect("Failed to get 2FA code from database");
1097    let twofa_res = auth_client
1098        .post(format!("{}/oauth/authorize/2fa", url))
1099        .form(&[
1100            ("request_uri", request_uri),
1101            ("code", &twofa_code),
1102        ])
1103        .send()
1104        .await
1105        .unwrap();
1106    assert!(
1107        twofa_res.status().is_redirection(),
1108        "Valid 2FA code should redirect to success, got status: {}",
1109        twofa_res.status()
1110    );
1111    let location = twofa_res.headers().get("location").unwrap().to_str().unwrap();
1112    assert!(
1113        location.starts_with(redirect_uri),
1114        "Should redirect to client callback, got: {}",
1115        location
1116    );
1117    assert!(
1118        location.contains("code="),
1119        "Redirect should include authorization code"
1120    );
1121    let auth_code = location.split("code=").nth(1).unwrap().split('&').next().unwrap();
1122    let token_res = http_client
1123        .post(format!("{}/oauth/token", url))
1124        .form(&[
1125            ("grant_type", "authorization_code"),
1126            ("code", auth_code),
1127            ("redirect_uri", redirect_uri),
1128            ("code_verifier", &code_verifier),
1129            ("client_id", &client_id),
1130        ])
1131        .send()
1132        .await
1133        .unwrap();
1134    assert_eq!(token_res.status(), StatusCode::OK, "Token exchange should succeed");
1135    let token_body: Value = token_res.json().await.unwrap();
1136    assert!(token_body["access_token"].is_string());
1137    assert_eq!(token_body["sub"], user_did);
1138}
1139#[tokio::test]
1140async fn test_2fa_lockout_after_max_attempts() {
1141    let url = base_url().await;
1142    let http_client = client();
1143    let ts = Utc::now().timestamp_millis();
1144    let handle = format!("2fa-lockout-{}", ts);
1145    let email = format!("2fa-lockout-{}@example.com", ts);
1146    let password = "2fa-test-password";
1147    let create_res = http_client
1148        .post(format!("{}/xrpc/com.atproto.server.createAccount", url))
1149        .json(&json!({
1150            "handle": handle,
1151            "email": email,
1152            "password": password
1153        }))
1154        .send()
1155        .await
1156        .unwrap();
1157    assert_eq!(create_res.status(), StatusCode::OK);
1158    let account: Value = create_res.json().await.unwrap();
1159    let user_did = account["did"].as_str().unwrap();
1160    let db_url = common::get_db_connection_string().await;
1161    let pool = sqlx::postgres::PgPoolOptions::new()
1162        .max_connections(1)
1163        .connect(&db_url)
1164        .await
1165        .expect("Failed to connect to database");
1166    sqlx::query("UPDATE users SET two_factor_enabled = true WHERE did = $1")
1167        .bind(user_did)
1168        .execute(&pool)
1169        .await
1170        .expect("Failed to enable 2FA");
1171    let redirect_uri = "https://example.com/2fa-lockout-callback";
1172    let mock_client = setup_mock_client_metadata(redirect_uri).await;
1173    let client_id = mock_client.uri();
1174    let (_, code_challenge) = generate_pkce();
1175    let par_body: Value = http_client
1176        .post(format!("{}/oauth/par", url))
1177        .form(&[
1178            ("response_type", "code"),
1179            ("client_id", &client_id),
1180            ("redirect_uri", redirect_uri),
1181            ("code_challenge", &code_challenge),
1182            ("code_challenge_method", "S256"),
1183        ])
1184        .send()
1185        .await
1186        .unwrap()
1187        .json()
1188        .await
1189        .unwrap();
1190    let request_uri = par_body["request_uri"].as_str().unwrap();
1191    let auth_client = no_redirect_client();
1192    let auth_res = auth_client
1193        .post(format!("{}/oauth/authorize", url))
1194        .form(&[
1195            ("request_uri", request_uri),
1196            ("username", &handle),
1197            ("password", password),
1198            ("remember_device", "false"),
1199        ])
1200        .send()
1201        .await
1202        .unwrap();
1203    assert!(auth_res.status().is_redirection());
1204    for i in 0..5 {
1205        let res = http_client
1206            .post(format!("{}/oauth/authorize/2fa", url))
1207            .form(&[
1208                ("request_uri", request_uri),
1209                ("code", "999999"),
1210            ])
1211            .send()
1212            .await
1213            .unwrap();
1214        if i < 4 {
1215            assert_eq!(res.status(), StatusCode::OK, "Attempt {} should show error page", i + 1);
1216            let body = res.text().await.unwrap();
1217            assert!(
1218                body.contains("Invalid verification code"),
1219                "Should show invalid code error on attempt {}", i + 1
1220            );
1221        }
1222    }
1223    let lockout_res = http_client
1224        .post(format!("{}/oauth/authorize/2fa", url))
1225        .form(&[
1226            ("request_uri", request_uri),
1227            ("code", "999999"),
1228        ])
1229        .send()
1230        .await
1231        .unwrap();
1232    assert_eq!(lockout_res.status(), StatusCode::OK);
1233    let body = lockout_res.text().await.unwrap();
1234    assert!(
1235        body.contains("Too many failed attempts") || body.contains("No 2FA challenge found"),
1236        "Should be locked out after max attempts. Body: {}",
1237        &body[..body.len().min(500)]
1238    );
1239}
1240#[tokio::test]
1241async fn test_account_selector_with_2fa_requires_verification() {
1242    let url = base_url().await;
1243    let http_client = client();
1244    let ts = Utc::now().timestamp_millis();
1245    let handle = format!("selector-2fa-{}", ts);
1246    let email = format!("selector-2fa-{}@example.com", ts);
1247    let password = "selector-2fa-password";
1248    let create_res = http_client
1249        .post(format!("{}/xrpc/com.atproto.server.createAccount", url))
1250        .json(&json!({
1251            "handle": handle,
1252            "email": email,
1253            "password": password
1254        }))
1255        .send()
1256        .await
1257        .unwrap();
1258    assert_eq!(create_res.status(), StatusCode::OK);
1259    let account: Value = create_res.json().await.unwrap();
1260    let user_did = account["did"].as_str().unwrap().to_string();
1261    let redirect_uri = "https://example.com/selector-2fa-callback";
1262    let mock_client = setup_mock_client_metadata(redirect_uri).await;
1263    let client_id = mock_client.uri();
1264    let (code_verifier, code_challenge) = generate_pkce();
1265    let par_body: Value = http_client
1266        .post(format!("{}/oauth/par", url))
1267        .form(&[
1268            ("response_type", "code"),
1269            ("client_id", &client_id),
1270            ("redirect_uri", redirect_uri),
1271            ("code_challenge", &code_challenge),
1272            ("code_challenge_method", "S256"),
1273        ])
1274        .send()
1275        .await
1276        .unwrap()
1277        .json()
1278        .await
1279        .unwrap();
1280    let request_uri = par_body["request_uri"].as_str().unwrap();
1281    let auth_client = no_redirect_client();
1282    let auth_res = auth_client
1283        .post(format!("{}/oauth/authorize", url))
1284        .form(&[
1285            ("request_uri", request_uri),
1286            ("username", &handle),
1287            ("password", password),
1288            ("remember_device", "true"),
1289        ])
1290        .send()
1291        .await
1292        .unwrap();
1293    assert!(auth_res.status().is_redirection());
1294    let device_cookie = auth_res.headers()
1295        .get("set-cookie")
1296        .and_then(|v| v.to_str().ok())
1297        .map(|s| s.split(';').next().unwrap_or("").to_string())
1298        .expect("Should have received device cookie");
1299    let location = auth_res.headers().get("location").unwrap().to_str().unwrap();
1300    assert!(location.contains("code="), "First auth should succeed");
1301    let code = location.split("code=").nth(1).unwrap().split('&').next().unwrap();
1302    let _token_body: Value = http_client
1303        .post(format!("{}/oauth/token", url))
1304        .form(&[
1305            ("grant_type", "authorization_code"),
1306            ("code", code),
1307            ("redirect_uri", redirect_uri),
1308            ("code_verifier", &code_verifier),
1309            ("client_id", &client_id),
1310        ])
1311        .send()
1312        .await
1313        .unwrap()
1314        .json()
1315        .await
1316        .unwrap();
1317    let db_url = common::get_db_connection_string().await;
1318    let pool = sqlx::postgres::PgPoolOptions::new()
1319        .max_connections(1)
1320        .connect(&db_url)
1321        .await
1322        .expect("Failed to connect to database");
1323    sqlx::query("UPDATE users SET two_factor_enabled = true WHERE did = $1")
1324        .bind(&user_did)
1325        .execute(&pool)
1326        .await
1327        .expect("Failed to enable 2FA");
1328    let (code_verifier2, code_challenge2) = generate_pkce();
1329    let par_body2: Value = http_client
1330        .post(format!("{}/oauth/par", url))
1331        .form(&[
1332            ("response_type", "code"),
1333            ("client_id", &client_id),
1334            ("redirect_uri", redirect_uri),
1335            ("code_challenge", &code_challenge2),
1336            ("code_challenge_method", "S256"),
1337        ])
1338        .send()
1339        .await
1340        .unwrap()
1341        .json()
1342        .await
1343        .unwrap();
1344    let request_uri2 = par_body2["request_uri"].as_str().unwrap();
1345    let select_res = auth_client
1346        .post(format!("{}/oauth/authorize/select", url))
1347        .header("cookie", &device_cookie)
1348        .form(&[
1349            ("request_uri", request_uri2),
1350            ("did", &user_did),
1351        ])
1352        .send()
1353        .await
1354        .unwrap();
1355    assert!(
1356        select_res.status().is_redirection(),
1357        "Account selector should redirect, got status: {}",
1358        select_res.status()
1359    );
1360    let select_location = select_res.headers().get("location").unwrap().to_str().unwrap();
1361    assert!(
1362        select_location.contains("/oauth/authorize/2fa"),
1363        "Account selector with 2FA enabled should redirect to 2FA page, got: {}",
1364        select_location
1365    );
1366    let twofa_code: String = sqlx::query_scalar(
1367        "SELECT code FROM oauth_2fa_challenge WHERE request_uri = $1"
1368    )
1369    .bind(request_uri2)
1370    .fetch_one(&pool)
1371    .await
1372    .expect("Failed to get 2FA code");
1373    let twofa_res = auth_client
1374        .post(format!("{}/oauth/authorize/2fa", url))
1375        .header("cookie", &device_cookie)
1376        .form(&[
1377            ("request_uri", request_uri2),
1378            ("code", &twofa_code),
1379        ])
1380        .send()
1381        .await
1382        .unwrap();
1383    assert!(twofa_res.status().is_redirection());
1384    let final_location = twofa_res.headers().get("location").unwrap().to_str().unwrap();
1385    assert!(
1386        final_location.starts_with(redirect_uri) && final_location.contains("code="),
1387        "After 2FA, should redirect to client with code, got: {}",
1388        final_location
1389    );
1390    let final_code = final_location.split("code=").nth(1).unwrap().split('&').next().unwrap();
1391    let token_res = http_client
1392        .post(format!("{}/oauth/token", url))
1393        .form(&[
1394            ("grant_type", "authorization_code"),
1395            ("code", final_code),
1396            ("redirect_uri", redirect_uri),
1397            ("code_verifier", &code_verifier2),
1398            ("client_id", &client_id),
1399        ])
1400        .send()
1401        .await
1402        .unwrap();
1403    assert_eq!(token_res.status(), StatusCode::OK);
1404    let final_token: Value = token_res.json().await.unwrap();
1405    assert_eq!(final_token["sub"], user_did, "Token should be for the correct user");
1406}