OAuthCodeStoreSpec.hs

{-# LANGUAGE DataKinds #-}

-- | Pure tests for the in-memory 'Shomei.OAuth.AuthorizationCode.Store' interpreter and for
-- 'Shomei.OAuth.Authorize.Workflow.authorize', the policy the authorize endpoint enforces.
--
-- The store's contract is consume-once: a code is redeemable exactly once, never after it
-- expires, and a replay is indistinguishable from an unknown code. Those three misses are what
-- the token endpoint answers @invalid_grant@ for, and getting any of them wrong turns a
-- single-use credential into a reusable one. The same behavior is re-proven against real
-- PostgreSQL — including under a genuine race — by @shomei-postgres@'s integration test.
--
-- The workflow's contract is the PKCE and scope policy: a public client cannot skip PKCE, only
-- S256 is accepted, and a client cannot be granted a scope it was never registered for.
module Shomei.OAuthCodeStoreSpec (tests) where

import Data.IORef (IORef, newIORef, readIORef)
import Data.Map.Strict qualified as Map
import Data.Maybe (isJust, isNothing)
import Data.Set (Set)
import Data.Set qualified as Set
import Data.Text (Text)
import Data.Text qualified as Text
import Data.Time (UTCTime (..), addUTCTime, fromGregorian)
import Shomei.Audit.Event.Domain qualified as Event
import Shomei.Authorization.Claims.Domain (Audience (..), AuthClaims (..), Issuer (..), Scope (..))
import Shomei.Config (ShomeiConfig, defaultShomeiConfig)
import Shomei.Id (SessionId, UserId, genOAuthClientId, genSessionId, genUserId, idText)
import Shomei.OAuth.AuthorizationCode.Domain (AuthorizationCode (..), NewAuthorizationCode (..))
import Shomei.OAuth.AuthorizationCode.Store
  ( bindAuthorizationCodeSession,
    consumeAuthorizationCode,
    deleteExpiredAuthorizationCodes,
    findConsumedAuthorizationCode,
    putAuthorizationCode,
  )
import Shomei.OAuth.Authorize.Workflow
  ( AuthorizeError (..),
    AuthorizeParams (..),
    IssuedCode (..),
    authorize,
    isValidS256Challenge,
  )
import Shomei.OAuth.Client.Domain (ClientType (..), NewOAuthClient (..), OAuthClient (..))
import Shomei.OAuth.Client.Store (createOAuthClient)
import Shomei.ServiceAccount.Secret (sha256Hex)
import Shomei.Session.Domain (NewSession (..), Session (..), SessionKind (InteractiveSession))
import Shomei.Session.Store (createSession)
import Shomei.Test.InMemory (World (..), emptyWorld, runInMemory)
import Test.Tasty (TestTree, testGroup)
import Test.Tasty.HUnit (assertBool, assertFailure, testCase, (@?=))

tests :: TestTree
tests =
  testGroup
    "OAuthCodeStore and the authorize workflow"
    [ testGroup
        "OAuthCodeStore (in-memory)"
        [ testCase "a stored code is consumable exactly once" consumeOnce,
          testCase "a consumed code can be bound to and recover its minted session" bindAndFindConsumed,
          testCase "an expired code never consumes" expiredNeverConsumes,
          testCase "an unknown code hash consumes to Nothing" unknownConsumes,
          testCase "deleteExpired removes only what is past its expiry" deleteExpired
        ],
      testGroup
        "authorize (workflow policy)"
        [ testCase "a valid request mints a code, stores only its digest, and audits it" happyPath,
          testCase "a public client without a code_challenge is refused" publicClientNeedsPkce,
          testCase "a confidential client may omit PKCE" confidentialMayOmitPkce,
          testCase "code_challenge_method other than S256 is refused" onlyS256,
          testCase "a code_challenge present with no method is refused (no silent `plain`)" noImplicitPlain,
          testCase "a malformed code_challenge is refused at authorize" malformedChallenge,
          testCase "response_type other than code is unsupported_response_type" onlyCodeResponseType,
          testCase "an absent scope grants the client's whole allow-list" absentScopeGrantsAll,
          testCase "an absent scope strips privilege scopes from a hand-inserted client" absentScopeStripsPrivileges,
          testCase "a requested privilege scope is invalid_scope" requestedPrivilegeScopeFails,
          testCase "an absent scope is invalid when only privilege scopes remain" privilegeOnlyDefaultFails,
          testCase "a scope outside the allow-list is invalid_scope" scopeOutsideAllowList,
          testCase "an empty scope parameter is invalid_scope, not a request for nothing" emptyScope,
          testCase "auth_time is copied from the authorizing token, not its iat or now" authTimeIsCredentialTime,
          testCase "isValidS256Challenge accepts only 43 unpadded base64url chars" challengeShape
        ]
    ]

-- Fixtures -------------------------------------------------------------------

t0 :: UTCTime
t0 = UTCTime (fromGregorian 2026 7 10) 0

-- | The stock config: 'authorize' reads only @oauthConfig.authorizationCodeTTL@ from it (60s by
-- default). @oidcEnabled@ gates the /route/, not the workflow.
cfg :: ShomeiConfig
cfg = defaultShomeiConfig (Issuer "https://shomei.test") (Audience "shomei-clients")

newWorld :: IO (IORef World)
newWorld = newIORef (emptyWorld t0)

callbackUri :: Text
callbackUri = "https://app.example.com/callback"

-- | A well-formed S256 challenge: 43 unpadded base64url characters.
challenge :: Text
challenge = Text.replicate 43 "a"

openidScope, profileScope :: Scope
openidScope = Scope "openid"
profileScope = Scope "profile"

allowed :: Set Scope
allowed = Set.fromList [openidScope, profileScope]

baseParams :: AuthorizeParams
baseParams =
  AuthorizeParams
    { responseType = Just "code",
      redirectUri = callbackUri,
      scope = Nothing,
      state = Nothing,
      nonce = Nothing,
      codeChallenge = Just challenge,
      codeChallengeMethod = Just "S256"
    }

-- | Claims for a user whose token was refreshed at request time but who authenticated an hour
-- earlier, so the authorization-code test can distinguish @auth_time@ from @iat@ and "now".
claimsFor :: UserId -> SessionId -> AuthClaims
claimsFor uid sid =
  AuthClaims
    { subject = uid,
      sessionId = sid,
      issuer = Issuer "https://shomei.test",
      audience = Audience "shomei-clients",
      issuedAt = t0,
      expiresAt = addUTCTime 900 t0,
      authTime = addUTCTime (-3600) t0,
      scopes = Set.empty,
      roles = Set.empty,
      permissions = Set.empty,
      actor = Nothing,
      extraClaims = mempty
    }

newCode :: UserId -> UTCTime -> Text -> NewAuthorizationCode
newCode uid expiresAt codeHash =
  NewAuthorizationCode
    { codeHash,
      clientId = "oauthclient_x",
      redirectUri = callbackUri,
      userId = uid,
      scopes = Set.singleton openidScope,
      nonce = Nothing,
      codeChallenge = Just challenge,
      authTime = t0,
      createdAt = t0,
      expiresAt
    }

-- Store ----------------------------------------------------------------------

-- | The single most important property in this plan: a code is a one-shot credential.
consumeOnce :: IO ()
consumeOnce = do
  ref <- newWorld
  (first', second') <- runInMemory ref do
    uid <- genUserId
    putAuthorizationCode (newCode uid (addUTCTime 60 t0) "hash-1")
    a <- consumeAuthorizationCode "hash-1" t0
    b <- consumeAuthorizationCode "hash-1" t0
    pure (a, b)
  assertBool "the first consume returns the code" (isJust first')
  fmap (.consumedAt) first' @?= Just (Just t0)
  assertBool "the second consume returns nothing" (isNothing second')

bindAndFindConsumed :: IO ()
bindAndFindConsumed = do
  ref <- newWorld
  (sid, found, expired) <- runInMemory ref do
    uid <- genUserId
    sid <- genSessionId
    putAuthorizationCode (newCode uid (addUTCTime 60 t0) "hash-bound")
    _ <- consumeAuthorizationCode "hash-bound" t0
    bindAuthorizationCodeSession "hash-bound" sid
    found <- findConsumedAuthorizationCode "hash-bound" t0
    expired <- findConsumedAuthorizationCode "hash-bound" (addUTCTime 61 t0)
    pure (sid, found, expired)
  fmap (.sessionId) found @?= Just (Just sid)
  assertBool "an expired consumed row is no longer replay-actionable" (isNothing expired)

expiredNeverConsumes :: IO ()
expiredNeverConsumes = do
  ref <- newWorld
  result <- runInMemory ref do
    uid <- genUserId
    putAuthorizationCode (newCode uid (addUTCTime 60 t0) "hash-1")
    -- One second past the expiry.
    consumeAuthorizationCode "hash-1" (addUTCTime 61 t0)
  assertBool "an expired code must not consume" (isNothing result)

unknownConsumes :: IO ()
unknownConsumes = do
  ref <- newWorld
  result <- runInMemory ref (consumeAuthorizationCode "no-such-hash" t0)
  assertBool "an unknown code hash consumes to Nothing" (isNothing result)

deleteExpired :: IO ()
deleteExpired = do
  ref <- newWorld
  remaining <- runInMemory ref do
    uid <- genUserId
    putAuthorizationCode (newCode uid (addUTCTime 10 t0) "expired")
    putAuthorizationCode (newCode uid (addUTCTime 600 t0) "live")
    deleteExpiredAuthorizationCodes (addUTCTime 60 t0)
    (,) <$> consumeAuthorizationCode "expired" (addUTCTime 60 t0) <*> consumeAuthorizationCode "live" (addUTCTime 60 t0)
  assertBool "the expired code is gone" (isNothing (fst remaining))
  assertBool "the live code survives" (isJust (snd remaining))

-- Workflow -------------------------------------------------------------------

-- | Run 'authorize' against a freshly registered client of the given type.
runAuthorize :: ClientType -> AuthorizeParams -> IO (Either AuthorizeError IssuedCode, World)
runAuthorize = runAuthorizeWithAllowed allowed

runAuthorizeWithAllowed :: Set Scope -> ClientType -> AuthorizeParams -> IO (Either AuthorizeError IssuedCode, World)
runAuthorizeWithAllowed registeredScopes clientType params = do
  ref <- newWorld
  result <- runInMemory ref do
    uid <- genUserId
    session <-
      createSession
        NewSession
          { userId = uid,
            createdAt = t0,
            expiresAt = addUTCTime 3600 t0,
            actor = Nothing,
            oauthClientId = Nothing,
            kind = InteractiveSession,
            grantedScopes = Set.empty,
            authenticatedAt = t0
          }
    ocid <- genOAuthClientId
    client <-
      createOAuthClient
        NewOAuthClient
          { oauthClientId = ocid,
            clientId = idText ocid,
            secretHash = case clientType of
              ConfidentialClient -> Just "hash"
              PublicClient -> Nothing,
            clientType,
            displayName = "test",
            redirectUris = [callbackUri],
            allowedScopes = registeredScopes,
            createdAt = t0
          }
    authorize cfg client (claimsFor uid session.sessionId) params
  world <- readIORef ref
  pure (result, world)

expectLeft :: Either AuthorizeError IssuedCode -> IO AuthorizeError
expectLeft = either pure (const (assertFailure "expected the authorize request to be refused"))

expectRight :: Either AuthorizeError IssuedCode -> IO IssuedCode
expectRight = either (\e -> assertFailure ("expected success, got " <> show e)) pure

happyPath :: IO ()
happyPath = do
  (result, world) <- runAuthorize ConfidentialClient baseParams {state = Just "xyz", nonce = Just "n-0S6"}
  issued <- expectRight result
  issued.state @?= Just "xyz"
  issued.grantedScopes @?= allowed
  -- Only the digest is stored: the code itself lives in the redirect URL and nowhere else.
  case Map.elems (oauthCodes world) of
    [stored] -> do
      stored.codeHash @?= sha256Hex issued.code
      assertBool "the plaintext code is never a key" (Map.notMember issued.code (oauthCodes world))
      stored.nonce @?= Just "n-0S6"
      stored.consumedAt @?= Nothing
      stored.expiresAt @?= addUTCTime 60 t0
    other -> assertFailure ("expected exactly one stored code, got " <> show (length other))
  -- The audit trail records the authorization without naming the code.
  assertBool
    "an oauth_code_issued event is published"
    (any isCodeIssued (publishedEvents world))
  where
    isCodeIssued = \case
      Event.OAuthCodeIssued _ -> True
      _ -> False

publicClientNeedsPkce :: IO ()
publicClientNeedsPkce = do
  (result, _) <- runAuthorize PublicClient baseParams {codeChallenge = Nothing, codeChallengeMethod = Nothing}
  e <- expectLeft result
  case e of
    AuthorizeInvalidRequest _ -> pure ()
    other -> assertFailure ("expected invalid_request, got " <> show other)

confidentialMayOmitPkce :: IO ()
confidentialMayOmitPkce = do
  (result, _) <- runAuthorize ConfidentialClient baseParams {codeChallenge = Nothing, codeChallengeMethod = Nothing}
  _ <- expectRight result
  pure ()

onlyS256 :: IO ()
onlyS256 = do
  (result, _) <- runAuthorize ConfidentialClient baseParams {codeChallengeMethod = Just "plain"}
  e <- expectLeft result
  case e of
    AuthorizeInvalidRequest _ -> pure ()
    other -> assertFailure ("expected invalid_request, got " <> show other)

-- | RFC 7636 defaults an absent method to @plain@. Accepting that default would silently downgrade
-- a client that meant S256, so the method must be spelled out.
noImplicitPlain :: IO ()
noImplicitPlain = do
  (result, _) <- runAuthorize ConfidentialClient baseParams {codeChallengeMethod = Nothing}
  e <- expectLeft result
  case e of
    AuthorizeInvalidRequest _ -> pure ()
    other -> assertFailure ("expected invalid_request, got " <> show other)

malformedChallenge :: IO ()
malformedChallenge = do
  (result, _) <- runAuthorize ConfidentialClient baseParams {codeChallenge = Just "too-short"}
  e <- expectLeft result
  case e of
    AuthorizeInvalidRequest _ -> pure ()
    other -> assertFailure ("expected invalid_request, got " <> show other)

onlyCodeResponseType :: IO ()
onlyCodeResponseType = do
  (result, _) <- runAuthorize ConfidentialClient baseParams {responseType = Just "token"}
  e <- expectLeft result
  e @?= UnsupportedResponseType

absentScopeGrantsAll :: IO ()
absentScopeGrantsAll = do
  (result, _) <- runAuthorize ConfidentialClient baseParams {scope = Nothing}
  issued <- expectRight result
  issued.grantedScopes @?= allowed

absentScopeStripsPrivileges :: IO ()
absentScopeStripsPrivileges = do
  let registered = Set.insert (Scope "shomei:admin") allowed
  (result, _) <- runAuthorizeWithAllowed registered ConfidentialClient baseParams {scope = Nothing}
  issued <- expectRight result
  issued.grantedScopes @?= allowed

requestedPrivilegeScopeFails :: IO ()
requestedPrivilegeScopeFails = do
  let registered = Set.insert (Scope "shomei:admin") allowed
  (result, _) <-
    runAuthorizeWithAllowed registered ConfidentialClient baseParams {scope = Just "openid shomei:admin"}
  e <- expectLeft result
  e @?= AuthorizeInvalidScope

privilegeOnlyDefaultFails :: IO ()
privilegeOnlyDefaultFails = do
  (result, _) <-
    runAuthorizeWithAllowed (Set.singleton (Scope "shomei:admin")) ConfidentialClient baseParams {scope = Nothing}
  e <- expectLeft result
  e @?= AuthorizeInvalidScope

scopeOutsideAllowList :: IO ()
scopeOutsideAllowList = do
  (result, _) <- runAuthorize ConfidentialClient baseParams {scope = Just "openid admin:everything"}
  e <- expectLeft result
  e @?= AuthorizeInvalidScope

emptyScope :: IO ()
emptyScope = do
  (result, _) <- runAuthorize ConfidentialClient baseParams {scope = Just "   "}
  e <- expectLeft result
  e @?= AuthorizeInvalidScope

-- | OIDC's @auth_time@ means "when the user authenticated", which is the authorizing access
-- token's carried credential time — an hour ago here — not its refreshed @iat@ or request time.
authTimeIsCredentialTime :: IO ()
authTimeIsCredentialTime = do
  (result, world) <- runAuthorize ConfidentialClient baseParams
  _ <- expectRight result
  case Map.elems (oauthCodes world) of
    [stored] -> stored.authTime @?= addUTCTime (-3600) t0
    _ -> assertFailure "expected exactly one stored code"

challengeShape :: IO ()
challengeShape = do
  assertBool "43 base64url chars is valid" (isValidS256Challenge challenge)
  assertBool "42 chars is not" (not (isValidS256Challenge (Text.replicate 42 "a")))
  assertBool "44 chars is not" (not (isValidS256Challenge (Text.replicate 44 "a")))
  -- Standard base64 (+ /) and padding are exactly what a client that forgot base64url emits.
  assertBool "'+' is not base64url" (not (isValidS256Challenge (Text.replicate 42 "a" <> "+")))
  assertBool "'/' is not base64url" (not (isValidS256Challenge (Text.replicate 42 "a" <> "/")))
  assertBool "'=' padding is not accepted" (not (isValidS256Challenge (Text.replicate 42 "a" <> "=")))
  assertBool "'-' and '_' are base64url" (isValidS256Challenge (Text.replicate 41 "a" <> "-_"))