Security Testing Auth Flows: PKCE, State Parameter, and Token Leakage

1. Client generates: code_verifier (random, 43-128 chars)
2. Client computes: code_challenge = BASE64URL(SHA256(code_verifier))
3. Authorization request includes: code_challenge + code_challenge_method=S256
4. Token request includes: code_verifier
5. Server verifies: SHA256(code_verifier) == code_challenge

import hashlib
import base64
import secrets

def generate_pkce():
    code_verifier = base64.urlsafe_b64encode(secrets.token_bytes(32)).rstrip(b"=").decode()
    code_challenge = base64.urlsafe_b64encode(
        hashlib.sha256(code_verifier.encode()).digest()
    ).rstrip(b"=").decode()
    return code_verifier, code_challenge

def test_pkce_happy_path():
    code_verifier, code_challenge = generate_pkce()

    # Authorization request with code_challenge
    code = get_authorization_code(
        code_challenge=code_challenge,
        code_challenge_method="S256"
    )

    # Token request with code_verifier
    response = exchange_code(code, code_verifier=code_verifier)
    assert response.status_code == 200
    assert "access_token" in response.json()

def test_pkce_wrong_verifier_rejected():
    code_verifier, code_challenge = generate_pkce()
    wrong_verifier, _ = generate_pkce()  # Different verifier

    code = get_authorization_code(
        code_challenge=code_challenge,
        code_challenge_method="S256"
    )

    # Send wrong verifier—server must reject
    response = exchange_code(code, code_verifier=wrong_verifier)
    assert response.status_code == 400
    assert response.json()["error"] == "invalid_grant"

def test_pkce_missing_verifier_rejected():
    _, code_challenge = generate_pkce()

    code = get_authorization_code(
        code_challenge=code_challenge,
        code_challenge_method="S256"
    )

    # No verifier in token request
    response = exchange_code(code, code_verifier=None)
    assert response.status_code == 400

def test_pkce_plain_method_rejected():
    # S256 is required; plain is deprecated and weaker
    # If your server is configured to reject plain, test it
    code_verifier = secrets.token_urlsafe(32)

    response = get_authorization_code(
        code_challenge=code_verifier,  # plain: challenge == verifier
        code_challenge_method="plain"
    )
    # Server should reject if plain is disabled
    if plain_pkce_disabled():
        assert response.status_code in [400, 302]  # Error or redirect with error

def test_public_client_requires_pkce():
    # Public client (no client_secret)—PKCE should be required
    response = get_authorization_code_without_pkce(
        client_id="public-client"
    )
    # Server should require PKCE for public clients
    assert_pkce_required(response)

def test_authorization_code_interception_prevented():
    code_verifier, code_challenge = generate_pkce()

    # Legitimate client starts flow with PKCE
    code = get_authorization_code(
        client_id="public-client",
        code_challenge=code_challenge,
        code_challenge_method="S256"
    )

    # Attacker intercepts the code—but doesn't have code_verifier
    attacker_response = exchange_code(code, code_verifier=None)
    assert attacker_response.status_code == 400

    # With wrong verifier
    _, attacker_challenge = generate_pkce()
    attacker_response2 = exchange_code(code, code_verifier=attacker_challenge)
    assert attacker_response2.status_code == 400

def test_state_parameter_validated():
    # Start a legitimate flow with state
    session = create_session()
    state = "random-state-12345"
    store_state(session, state)

    # Simulate callback with different state
    response = client.get(
        f"/callback?code=valid-code&state=different-state",
        cookies=session.cookies
    )

    # Must reject—state doesn't match
    assert response.status_code in [400, 403]
    assert not is_authenticated(session)

def test_missing_state_in_callback_rejected():
    session = create_session()
    state = "random-state-12345"
    store_state(session, state)

    # No state in callback at all
    response = client.get(
        "/callback?code=valid-code",
        cookies=session.cookies
    )
    assert response.status_code in [400, 403]

def test_empty_state_rejected():
    session = create_session()
    state = "random-state-12345"
    store_state(session, state)

    response = client.get(
        "/callback?code=valid-code&state=",
        cookies=session.cookies
    )
    assert response.status_code in [400, 403]

def test_state_is_random_and_unpredictable():
    # State should be cryptographically random
    states = []
    for _ in range(10):
        session = create_session()
        auth_url = start_oauth_flow(session)
        state = extract_state_from_url(auth_url)
        states.append(state)

    # All states should be unique
    assert len(set(states)) == len(states)

    # States should be long enough to be unpredictable (at least 16 bytes)
    for state in states:
        assert len(state) >= 22  # Base64url(16 bytes) = 22 chars minimum

def test_csrf_attack_fails():
    # Attacker sets up their own OAuth flow
    attacker_session = create_session()
    attacker_auth_url = start_oauth_flow(attacker_session)
    attacker_code = get_code_for_attacker(attacker_auth_url)
    attacker_state = extract_state_from_url(attacker_auth_url)

    # Victim has their own session (different state)
    victim_session = create_session()
    victim_auth_url = start_oauth_flow(victim_session)
    victim_state = extract_state_from_url(victim_auth_url)

    assert attacker_state != victim_state  # Different sessions, different states

    # Attacker tries to inject their code into victim's callback
    response = client.get(
        f"/callback?code={attacker_code}&state={attacker_state}",
        cookies=victim_session.cookies  # Victim's session
    )

    # State doesn't match victim's session—rejected
    assert response.status_code in [400, 403]
    assert not is_authenticated(victim_session)

def test_token_not_in_url_after_login(browser):
    browser.login("user", "password")

    # Token must not appear in any URL after login
    current_url = browser.current_url
    assert "access_token" not in current_url
    assert "token=" not in current_url
    assert "id_token" not in current_url

def test_callback_redirects_to_clean_url(browser):
    # OAuth callback URL contains the code
    # After processing, the app should redirect to a clean URL
    browser.login("user", "password")

    # After login, URL should not contain OAuth parameters
    final_url = browser.current_url
    assert "code=" not in final_url
    assert "state=" not in final_url
    assert "session_state=" not in final_url

def test_referrer_policy_set():
    response = requests.get("http://localhost:3000/dashboard")

    # Referrer-Policy should prevent token leakage
    referrer_policy = response.headers.get("Referrer-Policy", "")
    # Any of these prevent full URL leakage
    safe_policies = [
        "no-referrer",
        "same-origin",
        "strict-origin",
        "strict-origin-when-cross-origin"
    ]
    assert any(policy in referrer_policy for policy in safe_policies), \
        f"Unsafe Referrer-Policy: {referrer_policy}"

def test_access_token_not_in_logs():
    tokens = authenticate("user", "password")
    access_token = tokens["access_token"]

    # Make several authenticated requests
    for _ in range(5):
        client.get("/api/data", headers=auth(access_token))

    # Check logs
    log_content = read_application_logs()
    assert access_token not in log_content, \
        "Access token found in application logs"

def test_refresh_token_not_in_logs():
    tokens = authenticate("user", "password")
    refresh_token = tokens["refresh_token"]

    # Use the refresh token
    use_refresh_token(refresh_token)

    log_content = read_application_logs()
    assert refresh_token not in log_content, \
        "Refresh token found in application logs"

def test_authorization_code_not_in_logs():
    # The callback endpoint receives the code—make sure it's not logged
    with capture_logs() as log_stream:
        complete_oauth_flow("user", "password")

    log_content = log_stream.getvalue()
    # Code should be short-lived but still shouldn't appear in logs
    # Check that no JWT-like strings appear in logs
    assert not re.search(r'[A-Za-z0-9_-]{20,}\.[A-Za-z0-9_-]{20,}\.[A-Za-z0-9_-]{20,}', log_content), \
        "JWT-format string found in logs"

def test_spa_does_not_store_token_in_url_fragment(browser):
    # Some SPAs briefly put the token in the URL hash
    # It must not persist or be accessible to scripts
    browser.login("user", "password")

    # After login redirect processing, fragment should be cleared
    current_url = browser.current_url
    assert "#" not in current_url or "access_token" not in current_url

def test_token_not_in_browser_history(browser):
    browser.login("user", "password")

    # Navigate away and back
    browser.goto("http://example.com")
    browser.go_back()

    # History should not expose tokens
    history = browser.get_history()
    for entry in history:
        assert "access_token" not in entry.get("url", "")

def test_exact_redirect_uri_required():
    # Registered URI: https://app.example.com/callback
    # Must reject anything different

    malicious_uris = [
        "https://app.example.com.evil.com/callback",  # Subdomain attack
        "https://app.example.com/callback/../steal",   # Path traversal
        "https://app.example.com/callback?next=https://evil.com",  # Open redirect
        "https://evil.com",  # Completely different domain
        "https://app.example.com/callbacK",  # Case variation
        "http://app.example.com/callback",   # HTTP instead of HTTPS
        "https://app.example.com/callback/extra",  # Longer path
    ]

    for uri in malicious_uris:
        response = requests.get(
            "/oauth/authorize",
            params={
                "client_id": "my-app",
                "redirect_uri": uri,
                "response_type": "code",
                "scope": "openid"
            },
            allow_redirects=False
        )
        # Must NOT redirect to the malicious URI
        location = response.headers.get("Location", "")
        assert not location.startswith(uri), \
            f"Server redirected to malicious URI: {uri}"
        assert response.status_code in [400, 403], \
            f"Expected error for URI {uri}, got {response.status_code}"

def test_redirect_uri_not_in_error_response():
    # Some servers echo back the redirect URI in error responses
    # Avoid doing this for unregistered URIs (prevents reflected XSS)
    response = requests.get(
        "/oauth/authorize",
        params={
            "client_id": "my-app",
            "redirect_uri": "https://evil.com/<script>alert(1)</script>",
            "response_type": "code"
        }
    )
    assert "<script>" not in response.text

def test_implicit_flow_tokens_not_in_server_logs():
    # In implicit flow, tokens come back in the URL fragment
    # Server never sees them—but check middleware doesn't log them either
    # (This test depends on your specific setup)
    pass

def test_implicit_flow_disabled_for_new_clients():
    # Implicit flow is deprecated (OAuth 2.1 removes it)
    # New clients should not be able to use it
    response = requests.get(
        "/oauth/authorize",
        params={
            "client_id": "new-client",
            "redirect_uri": "https://app.example.com/callback",
            "response_type": "token",  # Implicit
            "scope": "openid"
        },
        allow_redirects=False
    )
    # Should be rejected for new clients
    assert response.status_code in [400, 403] or \
        "unsupported_response_type" in response.headers.get("Location", "")

def test_access_token_bound_to_client():
    # Token issued to client-A can't be used by client-B
    # (relevant when tokens carry client-specific claims)
    token_a = get_token_for_client("client-A", "user", "password")
    token_b = get_token_for_client("client-B", "user", "password")

    # client-A's token should work for client-A's resources
    assert client_a_resource(token_a["access_token"]).status_code == 200

    # client-A's token should NOT work for client-B's resources
    result = client_b_resource(token_a["access_token"])
    assert result.status_code in [401, 403]