JwtSpec.hs

{-# OPTIONS_GHC -Wno-deprecations #-}

-- | Regression tests for the verifier's trust boundary.
module Shomei.SigningKey.Verify.JwtSpec (tests) where

import Control.Lens ((%~), (&), (.~), (?~), (^.))
import Control.Monad (forM_)
import Crypto.JOSE.Compact (encodeCompact)
import Crypto.JOSE.Error (runJOSE)
import Crypto.JOSE.Header (newHeaderParamProtected)
import Crypto.JOSE.JWA.JWS (Alg (ES256, HS256, RS256))
import Crypto.JOSE.JWK (JWK, asPublicKey, fromOctets)
import Crypto.JOSE.JWS (newJWSHeaderProtected)
import Crypto.JOSE.JWS qualified as JWS
import Crypto.JWT
  ( Audience (Audience),
    ClaimsSet,
    JWTError,
    NumericDate (..),
    SignedJWT,
    StringOrURI,
    addClaim,
    claimAud,
    claimNbf,
    signClaims,
    unregisteredClaims,
  )
import Data.Aeson (Object, Result (Error, Success), Value (Number, String))
import Data.Aeson qualified as Aeson
import Data.Aeson.KeyMap qualified as KeyMap
import Data.ByteArray.Encoding (Base (Base64URLUnpadded), convertFromBase, convertToBase)
import Data.ByteString (ByteString)
import Data.ByteString.Lazy qualified as BSL
import Data.Map.Strict qualified as Map
import Data.Maybe (fromMaybe)
import Data.String (fromString)
import Data.Text (Text)
import Data.Text qualified as Text
import Data.Text.Encoding qualified as Text
import Data.Time (addUTCTime, getCurrentTime)
import Data.Time.Clock.POSIX (utcTimeToPOSIXSeconds)
import Shomei.Authorization.Claims.Domain (AuthClaims (..), mkExtraClaims)
import Shomei.Config (ShomeiConfig (..), SigningKeyConfig (..))
import Shomei.Error (TokenError (..))
import Shomei.Id (genUserId, idText)
import Shomei.Session.Token.Domain (AccessToken (AccessToken))
import Shomei.SigningKey.Domain qualified as Domain
import Shomei.SigningKey.Key.Jwt (generateSigningKey, generateSigningKeyFor, keyKid)
import Shomei.SigningKey.Sign.Jwt (claimsFromAuth, signAccessToken, signAccessTokenWith)
import Shomei.SigningKey.TestSupport (mkClaims, mkClaimsWith, publicJwks, testConfig)
import Shomei.SigningKey.Verify.Jwt (VerifierSettings (..), verifierSettingsFromConfig, verifyToken, verifyTokenWith)
import Test.Tasty (TestTree, testGroup)
import Test.Tasty.HUnit (assertBool, assertFailure, testCase, (@?=))

tests :: TestTree
tests =
  testGroup
    "Verify"
    [ testCase "actor object emission and bounded legacy compatibility preserve delegation" $ do
        jwk <- generateSigningKey
        now <- getCurrentTime
        base <- mkClaims testConfig now
        uid <- genUserId
        let claims = base {actor = Just uid}
            deadline = ceiling (utcTimeToPOSIXSeconds (addUTCTime 7200 now))
            legacyCfg = testConfig {signingKeyConfig = testConfig.signingKeyConfig {legacyActorUntilUnix = Just deadline, emitLegacyActor = True}}
            expiredCfg = legacyCfg {signingKeyConfig = legacyCfg.signingKeyConfig {legacyActorUntilUnix = Just (deadline - 7300)}}
            actorObject = Aeson.object ["sub" Aeson..= idText uid]
        objectWire <- signAccessOrFail jwk claims
        payload <- decodeSegment 1 objectWire
        KeyMap.lookup "act" payload @?= Just actorObject
        legacyToken <- either (assertFailure . show) pure =<< signAccessTokenWith legacyCfg jwk claims
        let AccessToken legacyWire = legacyToken
        legacyPayload <- decodeSegment 1 legacyWire
        KeyMap.lookup "act" legacyPayload @?= Just (String (idText uid))
        forM_ [objectWire, legacyWire] \wire -> do
          verified <- verifyToken (publicJwks jwk []) legacyCfg wire
          case verified of
            Right ac -> ac.actor @?= Just uid
            Left err -> assertFailure (show err)
        verifyToken (publicJwks jwk []) testConfig legacyWire >>= (@?= Left TokenMalformed)
        verifyToken (publicJwks jwk []) expiredCfg legacyWire >>= (@?= Left TokenMalformed)
        AccessToken automatic <- either (assertFailure . show) pure =<< signAccessTokenWith expiredCfg jwk claims
        automaticPayload <- decodeSegment 1 automatic
        KeyMap.lookup "act" automaticPayload @?= Just actorObject
        let tooShort = legacyCfg {signingKeyConfig = legacyCfg.signingKeyConfig {legacyActorUntilUnix = Just (deadline - 7140)}}
        AccessToken bounded <- either (assertFailure . show) pure =<< signAccessTokenWith tooShort jwk claims
        boundedPayload <- decodeSegment 1 bounded
        KeyMap.lookup "act" boundedPayload @?= Just actorObject,
      testCase "accepted lifetime is positive and finite independently of clock skew" $ do
        jwk <- generateSigningKey
        now <- getCurrentTime
        let settings = (verifierSettingsFromConfig testConfig) {maxTokenLifetime = 60, allowedClockSkew = 3600}
        forM_ [0, -1, 61] \ttl -> do
          claims <- mkClaimsWith testConfig now (addUTCTime ttl now)
          wire <- signAccessOrFail jwk claims
          verifyTokenWith settings (publicJwks jwk []) wire >>= (@?= Left TokenMalformed)
        claims <- mkClaimsWith testConfig now (addUTCTime 60 now)
        wire <- signAccessOrFail jwk claims
        result <- verifyTokenWith settings (publicJwks jwk []) wire
        case result of
          Right _ -> pure ()
          Left err -> assertFailure (show err)
        verifyTokenWith settings {maxTokenLifetime = 0} (publicJwks jwk []) wire >>= (@?= Left TokenMalformed),
      testCase "a direct extension bag cannot inject actor or jti" $ do
        jwk <- generateSigningKey
        now <- getCurrentTime
        base <- mkClaims testConfig now
        uid <- genUserId
        let claims = base {extraClaims = KeyMap.fromList [("act", Aeson.object ["sub" Aeson..= idText uid]), ("jti", String "forged")]}
        wire <- signAccessOrFail jwk claims
        payload <- decodeSegment 1 wire
        KeyMap.lookup "act" payload @?= Nothing
        KeyMap.lookup "jti" payload @?= Nothing
        assertBool "sid remains the correlation identifier" (KeyMap.member "sid" payload),
      testCase "a present malformed actor never verifies as an ordinary token" $ do
        jwk <- generateSigningKey
        now <- getCurrentTime
        claims <- mkClaims testConfig now
        forM_ [Aeson.object ["sub" Aeson..= String "invalid"], Aeson.toJSON (["actor"] :: [Text]), Number 1, Aeson.Null, String ""] \actorValue -> do
          wire <- signClaimsOrFail jwk (addClaim "act" actorValue (claimsFromAuth claims))
          result <- verifyToken (publicJwks jwk []) testConfig wire
          result @?= Left TokenMalformed,
      testCase "accepts an iat within the configured 30-second skew" $ do
        jwk <- generateSigningKey
        now <- getCurrentTime
        claims <- mkClaimsWith testConfig (addUTCTime 10 now) (addUTCTime 3600 now)
        wire <- signAccessOrFail jwk claims
        result <- verifyToken (publicJwks jwk []) testConfig wire
        case result of
          Right _ -> pure ()
          Left err -> assertFailure ("expected the token to verify, got " <> show err),
      testCase "rejects an iat beyond the configured skew" $ do
        jwk <- generateSigningKey
        now <- getCurrentTime
        claims <- mkClaimsWith testConfig (addUTCTime 120 now) (addUTCTime 3600 now)
        wire <- signAccessOrFail jwk claims
        result <- verifyToken (publicJwks jwk []) testConfig wire
        result @?= Left (TokenOtherError "iat in the future"),
      testCase "rejects a string-valued roles claim as malformed" $ do
        jwk <- generateSigningKey
        now <- getCurrentTime
        claims <- mkClaims testConfig now
        wire <- signClaimsOrFail jwk (claimsFromAuth claims & addClaim "roles" (String "admin"))
        result <- verifyToken (publicJwks jwk []) testConfig wire
        result @?= Left TokenMalformed,
      testCase "rejects a multi-element audience even when one value matches" $ do
        jwk <- generateSigningKey
        now <- getCurrentTime
        claims <- mkClaims testConfig now
        let audienceValues = [sou "shomei-clients", sou "other"]
        wire <- signClaimsOrFail jwk (claimsFromAuth claims & claimAud ?~ Audience audienceValues)
        result <- verifyToken (publicJwks jwk []) testConfig wire
        result @?= Left TokenAudienceInvalid,
      testCase "mints integral iat and exp values" $ do
        jwk <- generateSigningKey
        now <- getCurrentTime
        claims <- mkClaims testConfig now
        wire <- signAccessOrFail jwk claims
        payload <- decodeSegment 1 wire
        assertIntegralNumber "iat" payload
        assertIntegralNumber "exp" payload
        assertIntegralNumber "auth_time" payload,
      testCase "a legacy access token without auth_time falls back to iat" $ do
        jwk <- generateSigningKey
        now <- getCurrentTime
        claims <- mkClaims testConfig now
        let legacyClaims = claimsFromAuth claims & unregisteredClaims %~ Map.delete "auth_time"
        wire <- signClaimsOrFail jwk legacyClaims
        result <- verifyToken (publicJwks jwk []) testConfig wire
        case result of
          Right verified -> verified.authTime @?= verified.issuedAt
          Left err -> assertFailure ("expected the legacy token to verify, got " <> show err),
      testCase "signed future nbf is rejected while the configured skew is honored" $ do
        jwk <- generateSigningKey
        now <- getCurrentTime
        claims <- mkClaims testConfig now
        let wireAt offset = signClaimsWithHeaderOrFail jwk ES256 (Just (keyKid jwk)) (Just "at+jwt") (claimsFromAuth claims & claimNbf ?~ NumericDate (addUTCTime offset now))
        future <- wireAt 120
        verifyToken (publicJwks jwk []) testConfig future >>= (@?= Left (TokenOtherError "token not yet valid"))
        withinSkew <- wireAt 10
        accepted <- verifyToken (publicJwks jwk []) testConfig withinSkew
        case accepted of
          Right _ -> pure ()
          Left _ -> assertFailure "nbf inside configured skew was rejected",
      testCase "drops nbf and jti from the extension claim bag" $ do
        let extras = mkExtraClaims (KeyMap.fromList [("nbf", Number 1), ("jti", String "forged")])
        assertBool "nbf must be reserved" (not (KeyMap.member "nbf" extras))
        assertBool "jti must be reserved" (not (KeyMap.member "jti" extras)),
      testCase "reports an unknown kid without trying other keys" $ do
        signingKey <- generateSigningKey
        publishedKey <- generateSigningKey
        now <- getCurrentTime
        claims <- mkClaims testConfig now
        wire <- signAccessOrFail signingKey claims
        result <- verifyToken (publicJwks publishedKey []) testConfig wire
        result @?= Left (TokenKeyNotFound (Just (keyKid signingKey))),
      testCase "reports a missing kid even when the signature key is published" $ do
        jwk <- generateSigningKey
        now <- getCurrentTime
        claims <- mkClaims testConfig now
        wire <- signClaimsWithHeaderOrFail jwk ES256 Nothing (Just "at+jwt") (claimsFromAuth claims)
        result <- verifyToken (publicJwks jwk []) testConfig wire
        result @?= Left (TokenKeyNotFound Nothing),
      testCase "rejects typ JWT on an access token" $ do
        jwk <- generateSigningKey
        now <- getCurrentTime
        claims <- mkClaims testConfig now
        wire <- signClaimsWithHeaderOrFail jwk ES256 (Just (keyKid jwk)) (Just "JWT") (claimsFromAuth claims)
        result <- verifyToken (publicJwks jwk []) testConfig wire
        result @?= Left (TokenOtherError "typ JWT is not at+jwt"),
      testCase "temporarily accepts an access token with no typ" $ do
        jwk <- generateSigningKey
        now <- getCurrentTime
        claims <- mkClaims testConfig now
        wire <- signClaimsWithHeaderOrFail jwk ES256 (Just (keyKid jwk)) Nothing (claimsFromAuth claims)
        result <- verifyToken (publicJwks jwk []) testConfig wire
        case result of
          Right _ -> pure ()
          Left err -> assertFailure ("expected the typ-less compatibility token to verify, got " <> show err),
      testCase "mints access-token typ at+jwt" $ do
        jwk <- generateSigningKey
        now <- getCurrentTime
        claims <- mkClaims testConfig now
        wire <- signAccessOrFail jwk claims
        header <- decodeSegment 0 wire
        KeyMap.lookup "typ" header @?= Just (String "at+jwt"),
      testCase "rejects alg none" $ do
        jwk <- generateSigningKey
        now <- getCurrentTime
        claims <- mkClaims testConfig now
        realWire <- signAccessOrFail jwk claims
        let payload = Text.splitOn "." realWire !! 1
            header = Aeson.object ["alg" Aeson..= String "none", "kid" Aeson..= String (keyKid jwk)]
            headerSegment = Text.decodeUtf8 (convertToBase Base64URLUnpadded (BSL.toStrict (Aeson.encode header)) :: ByteString)
            wire = Text.intercalate "." [headerSegment, payload, ""]
        result <- verifyToken (publicJwks jwk []) testConfig wire
        result @?= Left TokenSignatureInvalid,
      testCase "rejects HS256 signed with public-key bytes" $ do
        jwk <- generateSigningKey
        now <- getCurrentTime
        claims <- mkClaims testConfig now
        let publicKey = fromMaybe jwk (jwk ^. asPublicKey)
            hmacKey = fromOctets (BSL.toStrict (Aeson.encode publicKey))
        wire <- signClaimsWithHeaderOrFail hmacKey HS256 (Just (keyKid jwk)) (Just "at+jwt") (claimsFromAuth claims)
        result <- verifyToken (publicJwks jwk []) testConfig wire
        result @?= Left TokenSignatureInvalid,
      testCase "rejects RS256 under an EC kid" $ do
        ecKey <- generateSigningKey
        rsaKey <- generateSigningKeyFor Domain.RS256
        now <- getCurrentTime
        claims <- mkClaims testConfig now
        wire <- signClaimsWithHeaderOrFail rsaKey RS256 (Just (keyKid ecKey)) (Just "at+jwt") (claimsFromAuth claims)
        result <- verifyToken (publicJwks ecKey []) testConfig wire
        result @?= Left TokenSignatureInvalid
    ]

sou :: Text -> StringOrURI
sou = fromString . Text.unpack

signAccessOrFail :: JWK -> AuthClaims -> IO Text
signAccessOrFail jwk claims = do
  result <- signAccessToken jwk claims
  case result of
    Left err -> assertFailure ("signing failed: " <> show err)
    Right (AccessToken wire) -> pure wire

signClaimsOrFail :: JWK -> ClaimsSet -> IO Text
signClaimsOrFail jwk = signClaimsWithHeaderOrFail jwk ES256 (Just (keyKid jwk)) Nothing

signClaimsWithHeaderOrFail :: JWK -> Alg -> Maybe Text -> Maybe Text -> ClaimsSet -> IO Text
signClaimsWithHeaderOrFail jwk algorithm headerKid headerType claims = do
  let header =
        newJWSHeaderProtected algorithm
          & JWS.kid
            .~ fmap newHeaderParamProtected headerKid
          & JWS.typ
            .~ fmap newHeaderParamProtected headerType
  result <- runJOSE @JWTError do
    signed <- signClaims jwk header claims
    pure (encodeCompact (signed :: SignedJWT))
  case result of
    Left err -> assertFailure ("signing claims failed: " <> show err)
    Right wire -> pure (Text.decodeUtf8 (BSL.toStrict wire))

decodeSegment :: Int -> Text -> IO Object
decodeSegment index wire = do
  let segment = Text.encodeUtf8 (Text.splitOn "." wire !! index)
  raw <-
    either (assertFailure . ("segment base64url decode failed: " <>)) pure $
      (convertFromBase Base64URLUnpadded segment :: Either String ByteString)
  maybe (assertFailure "segment is not a JSON object") pure (Aeson.decodeStrict raw)

assertIntegralNumber :: Text -> Object -> IO ()
assertIntegralNumber name payload = case KeyMap.lookup (fromString (Text.unpack name)) payload of
  Just value@(Number _) -> case Aeson.fromJSON value :: Result Integer of
    Success _ -> pure ()
    Error err -> assertFailure (Text.unpack name <> " is not integral: " <> err)
  other -> assertFailure (Text.unpack name <> " must be a JSON number, got " <> show other)