FastAPI Auth: OpenID Connect and Single Sign-On

{
  "iss": "https://accounts.google.com",
  "sub": "110248495921238986420",
  "aud": "your-client-id.apps.googleusercontent.com",
  "exp": 1700000000,
  "iat": 1699999000,
  "nonce": "random-nonce-from-your-app",
  "email": "user@gmail.com",
  "email_verified": true,
  "name": "Jane Doe",
  "picture": "https://lh3.googleusercontent.com/..."
}

import jwt
from jwt import PyJWKClient

# Cache the JWKS client — it fetches Google's public keys
jwks_client = PyJWKClient("https://www.googleapis.com/oauth2/v3/certs")


def validate_id_token(id_token: str, expected_nonce: str) -> dict:
    """Validate a Google ID token and return the claims."""
    # Fetch the signing key from Google's JWKS
    signing_key = jwks_client.get_signing_key_from_jwt(id_token)

    # Decode and validate
    claims = jwt.decode(
        id_token,
        signing_key.key,
        algorithms=["RS256"],
        audience=settings.google_client_id,  # Validates 'aud' claim
        issuer="https://accounts.google.com",  # Validates 'iss' claim
        options={"require": ["sub", "iss", "aud", "exp", "iat", "nonce"]},
    )

    # Verify nonce (replay protection)
    if claims.get("nonce") != expected_nonce:
        raise ValueError("Nonce mismatch — possible replay attack")

    # Verify email is verified
    if not claims.get("email_verified", False):
        raise ValueError("Email not verified by Google")

    return claims

@router.get("/callback")
async def google_callback(
    code: str,
    state: str,
    request: Request,
    response: Response,
    db: AsyncSession = Depends(get_db),
):
    # Verify state
    stored_state = request.cookies.get("oauth_state")
    stored_nonce = request.cookies.get("oauth_nonce")
    if not stored_state or not secrets.compare_digest(stored_state, state):
        raise HTTPException(status_code=400, detail="Invalid OAuth state")

    config = await get_google_config()

    # Exchange code for tokens
    async with httpx.AsyncClient() as client:
        token_response = await client.post(
            config["token_endpoint"],
            data={
                "code": code,
                "client_id": settings.google_client_id,
                "client_secret": settings.google_client_secret,
                "redirect_uri": settings.google_redirect_uri,
                "grant_type": "authorization_code",
            },
        )
        tokens = token_response.json()

    # Validate the ID token (no extra API call needed!)
    id_token = tokens.get("id_token")
    if not id_token:
        raise HTTPException(status_code=400, detail="No ID token in response")

    claims = validate_id_token(id_token, expected_nonce=stored_nonce)

    # Find or create user using the stable 'sub' claim
    user = await find_or_create_oauth_user(
        db,
        provider="google",
        provider_id=claims["sub"],      # Stable — never changes
        email=claims["email"],           # May change — not a primary key
        display_name=claims.get("name", claims["email"]),
        avatar_url=claims.get("picture"),
    )

    # Issue your own session/JWT
    jwt_token = create_access_token(user_id=str(user.id), roles=user.roles)

    # Clean up OAuth cookies
    response.delete_cookie("oauth_state")
    response.delete_cookie("oauth_nonce")

    return RedirectResponse(
        f"{settings.frontend_url}/auth/callback?token={jwt_token}"
    )

async def find_or_create_oauth_user(
    db: AsyncSession,
    provider: str,
    provider_id: str,
    email: str,
    display_name: str,
    avatar_url: str | None = None,
) -> User:
    # Check if this OAuth account exists
    result = await db.execute(
        select(OAuthAccount)
        .where(OAuthAccount.provider == provider, OAuthAccount.provider_user_id == provider_id)
        .options(joinedload(OAuthAccount.user))
    )
    oauth_account = result.scalar_one_or_none()

    if oauth_account:
        return oauth_account.user

    # Check if a user with this email exists (link accounts)
    result = await db.execute(select(User).where(User.email == email))
    user = result.scalar_one_or_none()

    if not user:
        user = User(
            email=email,
            display_name=display_name,
            hashed_password="",     # No password — OAuth-only user
            is_active=True,
            is_verified=True,       # Email verified by provider
        )
        db.add(user)
        await db.flush()

    # Link the OAuth account
    link = OAuthAccount(
        user_id=user.id,
        provider=provider,
        provider_user_id=provider_id,
        provider_email=email,
    )
    db.add(link)
    await db.commit()

    return user