nuthatch/test/Nuthatch/Probe/Oidc.hs

module Nuthatch.Probe.Oidc
  ( ProbeContext (..)
  , Session (..)
  , login
  , newProbeContext
  , verifyPkceBoundary
  ) where

import           Data.Aeson
  ( FromJSON (..)
  , eitherDecodeStrict'
  , withObject
  , (.:)
  )
import qualified Data.ByteString as BS
import qualified Data.ByteString.Char8 as B8
import           Data.Char (toLower)
import           Data.List (find)
import           Data.Maybe (fromMaybe)
import           Nuthatch.Probe.Assert
import           Nuthatch.Probe.Curl
import           System.Environment (lookupEnv)
import           System.FilePath ((</>))

data ProbeContext = ProbeContext
  { probeBase             :: String
  , probeDirectory        :: FilePath
  , probeOidcIssuer       :: String
  , probeOidcProvider     :: String
  , probeOidcPassword     :: String
  , probeOidcLoginPat     :: Maybe String
  , probeOidcClientId     :: String
  , probeOidcClientSecret :: String
  , probeOidcRedirectUri  :: String
  }

data Session = Session
  { sessionCookies :: FilePath
  , sessionCsrf    :: BS.ByteString
  }

data ZitadelSession = ZitadelSession
  { zitadelSessionId    :: String
  , zitadelSessionToken :: String
  }

instance FromJSON ZitadelSession where
  parseJSON = withObject "ZitadelSession" $ \object ->
    ZitadelSession
      <$> object .: "sessionId"
      <*> object .: "sessionToken"

newtype ZitadelCallback = ZitadelCallback String

instance FromJSON ZitadelCallback where
  parseJSON = withObject "ZitadelCallback" $ \object ->
    ZitadelCallback <$> object .: "callbackUrl"

newtype OidcDiscovery = OidcDiscovery String

instance FromJSON OidcDiscovery where
  parseJSON = withObject "OidcDiscovery" $ \object ->
    OidcDiscovery <$> object .: "token_endpoint"

newProbeContext :: FilePath -> IO ProbeContext
newProbeContext directory = do
  lane <- requireEnv "NUTHATCH_FIXTURE_LANE"
  issuer <- requireEnv "NUTHATCH_OIDC_ISSUER"
  provider <- requireEnv "NUTHATCH_OIDC_PROVIDER"
  password <- requireEnv "NUTHATCH_OIDC_USER_PASSWORD"
  loginPat <- optionalEnv "NUTHATCH_OIDC_LOGIN_PAT"
  clientId <- requireEnv "NUTHATCH_OIDC_CLIENT_ID"
  clientSecret <- requireEnv "NUTHATCH_OIDC_CLIENT_SECRET"
  redirectUri <- requireEnv "NUTHATCH_OIDC_REDIRECT_URI"
  pure ProbeContext
    { probeBase = "https://nuthatch-" <> lane <> ":8080"
    , probeDirectory = directory
    , probeOidcIssuer = issuer
    , probeOidcProvider = provider
    , probeOidcPassword = password
    , probeOidcLoginPat = loginPat
    , probeOidcClientId = clientId
    , probeOidcClientSecret = clientSecret
    , probeOidcRedirectUri = redirectUri
    }

login :: ProbeContext -> Bool -> String -> String -> IO Session
login context verifyBoundary user suffix = do
  let cookies = probeDirectory context </> "cookies-" <> suffix
      idpCookies = probeDirectory context </> "idp-" <> suffix
  started <- curl
    [ "--cookie-jar", cookies
    , probeBase context <> "/login"
    ]
  assertEqual "login redirect status" 303 (curlStatus started)
  authorizationUrl <-
    maybe (fail "login response omitted Location") (pure . B8.unpack) $
      headerValue "location" started
  requireCookieAttributes
    "OIDC transaction"
    ["secure", "httponly", "samesite=lax", "path=/oidc/callback"]
    (findCookie "nuthatch_oidc" started)
  state <- queryValue "state" authorizationUrl
  if verifyBoundary
    then verifyTransactionBoundary context cookies authorizationUrl state
    else pure ()
  callback <- completeAuthorization context idpCookies user suffix authorizationUrl
  completed <- curl
    [ "--cookie", cookies
    , "--cookie-jar", cookies
    , callback
    ]
  assertEqual "OIDC callback status" 303 (curlStatus completed)
  assertEqual "OIDC callback location"
    (Just "/app")
    (B8.unpack <$> headerValue "location" completed)
  requireCookieAttributes
    "session"
    ["secure", "httponly", "samesite=lax", "path=/;"]
    (findCookie "nuthatch_session" completed)
  requireCookieAttributes
    "CSRF"
    ["secure", "samesite=lax", "path=/;"]
    (findCookie "nuthatch_csrf" completed)
  assert "CSRF cookie must be readable by the browser" $
    not ("httponly" `containsLower` findCookie "nuthatch_csrf" completed)
  requireCookieAttributes
    "expired OIDC transaction"
    ["max-age=0", "path=/oidc/callback;"]
    (findCookie "nuthatch_oidc" completed)
  let responseCookies = headerValues "set-cookie" completed
  assert "application cookies must be host-only" $
    all (not . containsLower "domain=") responseCookies
  csrf <- cookieValue cookies "nuthatch_csrf"
  pure Session
    { sessionCookies = cookies
    , sessionCsrf = csrf
    }

verifyPkceBoundary :: ProbeContext -> String -> IO ()
verifyPkceBoundary context user = do
  let cookies = probeDirectory context </> "cookies-pkce"
      idpCookies = probeDirectory context </> "idp-pkce"
  started <- curl
    [ "--cookie-jar", cookies
    , probeBase context <> "/login"
    ]
  authorizationUrl <-
    maybe (fail "PKCE login omitted Location") (pure . B8.unpack) $
      headerValue "location" started
  callback <- completeAuthorization
    context
    idpCookies
    user
    "pkce"
    authorizationUrl
  code <- queryValue "code" callback
  discoveryResponse <- curl
    [probeOidcIssuer context <> "/.well-known/openid-configuration"]
  assertEqual "OIDC discovery status" 200 (curlStatus discoveryResponse)
  OidcDiscovery tokenEndpoint <-
    decodeBody "OIDC discovery" discoveryResponse
  rejected <- curl
    [ "--user"
    , probeOidcClientId context <> ":" <> probeOidcClientSecret context
    , "--data-urlencode", "grant_type=authorization_code"
    , "--data-urlencode", "code=" <> code
    , "--data-urlencode", "redirect_uri=" <> probeOidcRedirectUri context
    , "--data-urlencode"
    , "code_verifier=aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"
    , tokenEndpoint
    ]
  assert "provider accepted an invalid PKCE verifier" $
    curlStatus rejected >= 400

verifyTransactionBoundary
  :: ProbeContext
  -> FilePath
  -> String
  -> String
  -> IO ()
verifyTransactionBoundary context cookies authorizationUrl state = do
  wrongRedirect <- curl [replaceRedirectUri authorizationUrl]
  assert "provider accepted an unregistered redirect URI" $
    curlStatus wrongRedirect >= 400
  invalidState <- curl
    [ "--cookie", cookies
    , probeBase context
      <> "/oidc/callback?code=unused&state=invalid-"
      <> state
    ]
  assert "application accepted an invalid OIDC state" $
    curlStatus invalidState >= 400
  transaction <- cookieValue cookies "nuthatch_oidc"
  malformed <- curl
    [ "--header", "Cookie: nuthatch_oidc=malformed"
    , probeBase context
      <> "/oidc/callback?code=unused&state="
      <> state
    ]
  assertEqual "malformed transaction cookie" 400 (curlStatus malformed)
  let tampered = BS.cons 120 (BS.drop 1 transaction)
  tamperedResponse <- curl
    [ "--header", "Cookie: nuthatch_oidc=" <> B8.unpack tampered
    , probeBase context
      <> "/oidc/callback?code=unused&state="
      <> state
    ]
  assertEqual "tampered transaction cookie" 400 (curlStatus tamperedResponse)

completeAuthorization
  :: ProbeContext
  -> FilePath
  -> String
  -> String
  -> String
  -> IO String
completeAuthorization context idpCookies user suffix authorizationUrl =
  case probeOidcProvider context of
    "keycloak" ->
      completeKeycloak context idpCookies user authorizationUrl
    "zitadel" ->
      completeZitadel context user suffix authorizationUrl
    provider ->
      fail ("unsupported OIDC provider: " <> provider)

completeKeycloak
  :: ProbeContext
  -> FilePath
  -> String
  -> String
  -> IO String
completeKeycloak context idpCookies user authorizationUrl = do
  form <- curl
    [ "--cookie-jar", idpCookies
    , authorizationUrl
    ]
  assertEqual "Keycloak authorization form" 200 (curlStatus form)
  action <- htmlAttribute "action" (B8.unpack (curlBody form))
  submitted <- curl
    [ "--cookie", idpCookies
    , "--cookie-jar", idpCookies
    , "--data-urlencode", "username=" <> user
    , "--data-urlencode", "password=" <> probeOidcPassword context
    , "--data-urlencode", "credentialId="
    , action
    ]
  maybe
    (fail "Keycloak login omitted callback Location")
    (pure . B8.unpack)
    (headerValue "location" submitted)

completeZitadel
  :: ProbeContext
  -> String
  -> String
  -> String
  -> IO String
completeZitadel context user _ authorizationUrl = do
  pat <- maybe (fail "ZITADEL login PAT is missing") pure $
    probeOidcLoginPat context
  started <- curl [authorizationUrl]
  location <- maybe
    (fail "ZITADEL login omitted Location")
    (pure . B8.unpack)
    (headerValue "location" started)
  authRequest <- queryValue "authRequest" location
  created <- curl
    [ "--header", "Authorization: Bearer " <> pat
    , "--header", "Content-Type: application/json"
    , "--data"
    , "{\"checks\":{\"user\":{\"userId\":\""
      <> user
      <> "\"},\"password\":{\"password\":\""
      <> probeOidcPassword context
      <> "\"}},\"lifetime\":\"300s\"}"
    , probeOidcIssuer context <> "/v2/sessions"
    ]
  assert "ZITADEL session creation failed" $
    curlStatus created >= 200 && curlStatus created < 300
  ZitadelSession {..} <- decodeBody "ZITADEL session" created
  finalized <- curl
    [ "--header", "Authorization: Bearer " <> pat
    , "--header", "Content-Type: application/json"
    , "--data"
    , "{\"session\":{\"sessionId\":\""
      <> zitadelSessionId
      <> "\",\"sessionToken\":\""
      <> zitadelSessionToken
      <> "\"}}"
    , probeOidcIssuer context
      <> "/v2/oidc/auth_requests/"
      <> authRequest
    ]
  assertEqual "ZITADEL authorization finalization" 200 (curlStatus finalized)
  ZitadelCallback callback <- decodeBody "ZITADEL callback" finalized
  pure callback

decodeBody :: FromJSON value => String -> CurlResponse -> IO value
decodeBody label response =
  either (fail . ((label <> ": ") <>)) pure $
    eitherDecodeStrict' (curlBody response)

findCookie :: BS.ByteString -> CurlResponse -> BS.ByteString
findCookie name response =
  fromMaybe BS.empty $
    find (containsLower (lower name <> "=")) (headerValues "set-cookie" response)

requireCookieAttributes
  :: String
  -> [BS.ByteString]
  -> BS.ByteString
  -> IO ()
requireCookieAttributes label attributes cookie = do
  assert (label <> " cookie is missing") (not (BS.null cookie))
  mapM_
    (\attribute ->
      assert (label <> " cookie omitted " <> B8.unpack attribute) $
        attribute `containsLower` cookie)
    attributes

containsLower :: BS.ByteString -> BS.ByteString -> Bool
containsLower needle haystack =
  lower needle `BS.isInfixOf` lower haystack

lower :: BS.ByteString -> BS.ByteString
lower = B8.map toLower

queryValue :: String -> String -> IO String
queryValue name url =
  case drop 1 (dropWhile (/= '?') url) of
    "" -> fail ("URL has no query: " <> url)
    query ->
      case find ((== name) . fst) (map splitPair (splitOn '&' query)) of
        Just (_, value) | not (null value) -> pure value
        _ -> fail ("URL query omitted " <> name)

splitPair :: String -> (String, String)
splitPair value =
  case break (== '=') value of
    (name, '=' : rest) -> (name, rest)
    (name, _) -> (name, "")

splitOn :: Char -> String -> [String]
splitOn delimiter value =
  case break (== delimiter) value of
    (prefix, []) -> [prefix]
    (prefix, _ : suffix) -> prefix : splitOn delimiter suffix

replaceRedirectUri :: String -> String
replaceRedirectUri url =
  let marker = "redirect_uri="
      (prefix, rest) = breakOn marker url
  in case rest of
      "" -> url
      _ ->
        let afterName = drop (length marker) rest
            suffix = dropWhile (/= '&') afterName
        in prefix
          <> marker
          <> "https%3A%2F%2Finvalid.example%2Fcallback"
          <> suffix

htmlAttribute :: String -> String -> IO String
htmlAttribute name html =
  let marker = name <> "=\""
      (_, rest) = breakOn marker html
  in case rest of
      "" -> fail ("HTML omitted " <> name)
      _ ->
        pure
          ( replaceAll "&amp;" "&"
          $ takeWhile (/= '"')
          $ drop (length marker) rest
          )

breakOn :: String -> String -> (String, String)
breakOn needle = go []
  where
    go prefix remaining
      | needle `isPrefixOf` remaining = (reverse prefix, remaining)
    go prefix [] = (reverse prefix, "")
    go prefix (next : remaining) = go (next : prefix) remaining

isPrefixOf :: String -> String -> Bool
isPrefixOf prefix value =
  take (length prefix) value == prefix

replaceAll :: String -> String -> String -> String
replaceAll needle replacement value =
  case breakOn needle value of
    (prefix, "") -> prefix
    (prefix, rest) ->
      prefix
      <> replacement
      <> replaceAll needle replacement (drop (length needle) rest)

optionalEnv :: String -> IO (Maybe String)
optionalEnv = lookupEnv