-- This file is part of the Wire Server implementation.
--
-- Copyright (C) 2025 Wire Swiss GmbH <opensource@wire.com>
--
-- This program is free software: you can redistribute it and/or modify it under
-- the terms of the GNU Affero General Public License as published by the Free
-- Software Foundation, either version 3 of the License, or (at your option) any
-- later version.
--
-- This program is distributed in the hope that it will be useful, but WITHOUT
-- ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
-- FOR A PARTICULAR PURPOSE. See the GNU Affero General Public License for more
-- details.
--
-- You should have received a copy of the GNU Affero General Public License along
-- with this program. If not, see <https://www.gnu.org/licenses/>.

module Test.Auth where

import API.Brig
import API.BrigInternal
import API.Common
import API.GalleyInternal
import qualified API.Nginz as Nginz
import qualified Data.ByteString.Char8 as BSChar8
import SetupHelpers
import Testlib.Prelude
import Text.Read
import UnliftIO.Async
import UnliftIO.Concurrent

-- Happy flow: login yields a valid zauth token.
--
-- See also: 'testBearerToken'
testBearerToken2 :: (HasCallStack) => App ()
testBearerToken2 :: HasCallStack => App ()
testBearerToken2 = do
  alice <- Domain -> CreateUser -> App Value
forall domain.
(HasCallStack, MakesValue domain) =>
domain -> CreateUser -> App Value
randomUser Domain
OwnDomain CreateUser
forall a. Default a => a
def
  email <- asString $ alice %. "email"
  loginResp <- login alice email defPassword >>= getJSON 200
  token <- asString $ loginResp %. "access_token"

  req <-
    rawBaseRequest alice Nginz Versioned "/self"
      <&> addHeader "Authorization" ("Bearer " <> token)
  submit "GET" req `bindResponse` \Response
resp -> do
    Response
resp.status Int -> Int -> App ()
forall a. (MakesValue a, HasCallStack) => a -> Int -> App ()
`shouldMatchInt` Int
200
    Response
resp.json App Value -> String -> App Value
forall a. (HasCallStack, MakesValue a) => a -> String -> App Value
%. String
"email" App Value -> String -> App ()
forall a b.
(MakesValue a, MakesValue b, HasCallStack) =>
a -> b -> App ()
`shouldMatch` String
email

-- Happy flow (zauth token encoded in AWS4_HMAC_SHA256)
--
-- See also: 'testAWS4_HMAC_SHA256_token'
testAWS4_HMAC_SHA256_token2 :: (HasCallStack) => App ()
testAWS4_HMAC_SHA256_token2 :: HasCallStack => App ()
testAWS4_HMAC_SHA256_token2 = do
  alice <- Domain -> CreateUser -> App Value
forall domain.
(HasCallStack, MakesValue domain) =>
domain -> CreateUser -> App Value
randomUser Domain
OwnDomain CreateUser
forall a. Default a => a
def
  email <- asString $ alice %. "email"
  loginResp <- login alice email defPassword >>= getJSON 200
  token <- asString $ loginResp %. "access_token"

  let testCases =
        [ (Bool
True, String
"AWS4-HMAC-SHA256 Credential=" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
token String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
", foo=bar"),
          (Bool
True, String
"AWS4-HMAC-SHA256 Credential=" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
token),
          (Bool
True, String
"AWS4-HMAC-SHA256 foo=bar, Credential=" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
token),
          (Bool
True, String
"AWS4-HMAC-SHA256 foo=bar, Credential=" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
token String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
", baz=qux"),
          (Bool
True, String
"AWS4-HMAC-SHA256 foo=bar,Credential=" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
token String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
",baz=qux"),
          (Bool
False, String
"AWS4-HMAC-SHA256 Credential=badtoken")
        ]

  for_ testCases $ \(Bool
good, String
header) -> do
    req <-
      Value -> Service -> Versioned -> String -> App Request
forall domain.
(HasCallStack, MakesValue domain) =>
domain -> Service -> Versioned -> String -> App Request
rawBaseRequest Value
alice Service
Nginz Versioned
Versioned String
"/self"
        App Request -> (Request -> Request) -> App Request
forall (f :: * -> *) a b. Functor f => f a -> (a -> b) -> f b
<&> String -> String -> Request -> Request
addHeader String
"Authorization" String
header
    submit "GET" req `bindResponse` \Response
resp -> do
      if Bool
good
        then do
          Response
resp.status Int -> Int -> App ()
forall a. (MakesValue a, HasCallStack) => a -> Int -> App ()
`shouldMatchInt` Int
200
          Response
resp.json App Value -> String -> App Value
forall a. (HasCallStack, MakesValue a) => a -> String -> App Value
%. String
"email" App Value -> String -> App ()
forall a b.
(MakesValue a, MakesValue b, HasCallStack) =>
a -> b -> App ()
`shouldMatch` String
email
        else do
          Response
resp.status Int -> Int -> App ()
forall a. (MakesValue a, HasCallStack) => a -> Int -> App ()
`shouldMatchInt` Int
401

-- The testLimitRetries test conforms to the following testing standards:
-- @SF.Channel @TSFI.RESTfulAPI @TSFI.NTP @S2
--
-- The following test tests the login retries. It checks that a user can make
-- only a prespecified number of attempts to log in with an invalid password,
-- after which the user is unable to try again for a configured amount of time.
-- After the configured amount of time has passed, the test asserts the user can
-- successfully log in again. Furthermore, the test asserts that another
-- unrelated user can successfully log-in in parallel to the failed attempts of
-- the aforementioned user.
testLimitRetries :: (HasCallStack) => App ()
testLimitRetries :: HasCallStack => App ()
testLimitRetries = do
  let retryLimit :: Int
retryLimit = Int
5
      timeout :: Int
timeout = Int
5
  ServiceOverrides -> (HasCallStack => String -> App ()) -> App ()
forall a.
HasCallStack =>
ServiceOverrides -> (HasCallStack => String -> App a) -> App a
withModifiedBackend
    ServiceOverrides
forall a. Default a => a
def
      { brigCfg =
          -- Set a small timeout to make this test fast
          setField @_ @Int "optSettings.setLimitFailedLogins.timeout" timeout
            >=> setField @_ @Int "optSettings.setLimitFailedLogins.retryLimit" retryLimit
            -- Disable password hashing rate limiting, so we can login many times without making this test slow
            >=> setField @_ @Int "optSettings.setPasswordHashingRateLimit.userLimit.inverseRate" 0
      }
    ((HasCallStack => String -> App ()) -> App ())
-> (HasCallStack => String -> App ()) -> App ()
forall a b. (a -> b) -> a -> b
$ \String
domain -> do
      alice <- String -> CreateUser -> App Value
forall domain.
(HasCallStack, MakesValue domain) =>
domain -> CreateUser -> App Value
randomUser String
domain CreateUser
forall a. Default a => a
def
      aliceEmail <- asString $ alice %. "email"

      bob <- randomUser domain def
      bobEmail <- asString $ bob %. "email"

      -- Alice tries to login 5 times with wrong password
      forM_ [1 .. retryLimit] $ \Int
_ ->
        String -> String -> String -> App Response
forall domain.
(HasCallStack, MakesValue domain) =>
domain -> String -> String -> App Response
login String
domain String
aliceEmail String
"wrong-password" App Response -> (Response -> App ()) -> App ()
forall a.
HasCallStack =>
App Response -> (Response -> App a) -> App a
`bindResponse` \Response
resp -> do
          Response
resp.status Int -> Int -> App ()
forall a. (MakesValue a, HasCallStack) => a -> Int -> App ()
`shouldMatchInt` Int
403
          Response
resp.json App Value -> String -> App Value
forall a. (HasCallStack, MakesValue a) => a -> String -> App Value
%. String
"label" App Value -> String -> App ()
forall a b.
(MakesValue a, MakesValue b, HasCallStack) =>
a -> b -> App ()
`shouldMatch` String
"invalid-credentials"

      -- Now alice cannot login even with correct password
      retryAfter <-
        login domain aliceEmail defPassword `bindResponse` \Response
resp -> do
          Response
resp.status Int -> Int -> App ()
forall a. (MakesValue a, HasCallStack) => a -> Int -> App ()
`shouldMatchInt` Int
403
          Response
resp.json App Value -> String -> App Value
forall a. (HasCallStack, MakesValue a) => a -> String -> App Value
%. String
"label" App Value -> String -> App ()
forall a b.
(MakesValue a, MakesValue b, HasCallStack) =>
a -> b -> App ()
`shouldMatch` String
"client-error"
          let Just Int
retryAfter = String -> Maybe Int
forall a. Read a => String -> Maybe a
readMaybe (String -> Maybe Int)
-> (ByteString -> String) -> ByteString -> Maybe Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> String
BSChar8.unpack (ByteString -> Maybe Int) -> Maybe ByteString -> Maybe Int
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< HeaderName -> [(HeaderName, ByteString)] -> Maybe ByteString
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup (String -> HeaderName
forall a. IsString a => String -> a
fromString String
"Retry-After") Response
resp.headers
          (Int
retryAfter Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= Int
timeout) Bool -> Bool -> App ()
forall a b.
(MakesValue a, MakesValue b, HasCallStack) =>
a -> b -> App ()
`shouldMatch` Bool
True
          Int -> App Int
forall a. a -> App a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
retryAfter

      -- Bob can still login
      login domain bobEmail defPassword `bindResponse` \Response
resp -> do
        Response
resp.status Int -> Int -> App ()
forall a. (MakesValue a, HasCallStack) => a -> Int -> App ()
`shouldMatchInt` Int
200

      -- Waiting 2s less than retryAfter should still cause a failure
      threadDelay ((retryAfter - 2) * 1_000_000)
      login domain aliceEmail defPassword `bindResponse` \Response
resp -> do
        Response
resp.status Int -> Int -> App ()
forall a. (MakesValue a, HasCallStack) => a -> Int -> App ()
`shouldMatchInt` Int
403
        Response
resp.json App Value -> String -> App Value
forall a. (HasCallStack, MakesValue a) => a -> String -> App Value
%. String
"label" App Value -> String -> App ()
forall a b.
(MakesValue a, MakesValue b, HasCallStack) =>
a -> b -> App ()
`shouldMatch` String
"client-error"
        let Just Int
retryAfter2 = String -> Maybe Int
forall a. Read a => String -> Maybe a
readMaybe (String -> Maybe Int)
-> (ByteString -> String) -> ByteString -> Maybe Int
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> String
BSChar8.unpack (ByteString -> Maybe Int) -> Maybe ByteString -> Maybe Int
forall (m :: * -> *) a b. Monad m => (a -> m b) -> m a -> m b
=<< HeaderName -> [(HeaderName, ByteString)] -> Maybe ByteString
forall a b. Eq a => a -> [(a, b)] -> Maybe b
lookup (String -> HeaderName
forall a. IsString a => String -> a
fromString String
"Retry-After") Response
resp.headers
        -- This should be about 2 seconds or slightly less because we didn't
        -- wait long enough. This also asserts that the throttling doesn't get
        -- reset by making another call
        (Int
retryAfter2 Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
<= (Int
2 :: Int)) Bool -> Bool -> App ()
forall a b.
(MakesValue a, MakesValue b, HasCallStack) =>
a -> b -> App ()
`shouldMatch` Bool
True

      -- Waiting 2 more seconds should make the login succeed
      threadDelay 2_000_000
      login domain aliceEmail defPassword `bindResponse` \Response
resp -> do
        Response
resp.status Int -> Int -> App ()
forall a. (MakesValue a, HasCallStack) => a -> Int -> App ()
`shouldMatchInt` Int
200

-- @END

-- The testTooManyCookies test conforms to the following testing standards:
-- @SF.Provisioning @TSFI.RESTfulAPI @S2
--
-- The test asserts that there is an upper limit for the number of user cookies
-- per cookie type. It does that by concurrently attempting to create more
-- persistent and session cookies than the configured maximum.
-- Creation of new cookies beyond the limit causes deletion of the
-- oldest cookies.
testTooManyCookies :: (HasCallStack) => App ()
testTooManyCookies :: HasCallStack => App ()
testTooManyCookies = do
  let cookieLimit :: Int
cookieLimit = Int
5
  ServiceOverrides -> (HasCallStack => String -> App ()) -> App ()
forall a.
HasCallStack =>
ServiceOverrides -> (HasCallStack => String -> App a) -> App a
withModifiedBackend
    ServiceOverrides
forall a. Default a => a
def
      { brigCfg =
          -- Disable password hashing rate limiting, so we can login many times without making this test slow
          setField @_ @Int "optSettings.setPasswordHashingRateLimit.userLimit.inverseRate" 0
            -- Disable cookie throttling so this test is not slow
            >=> setField @_ @Int "optSettings.setUserCookieThrottle.retryAfter" 0
            >=> setField @_ @Int "optSettings.setUserCookieLimit" cookieLimit
      }
    ((HasCallStack => String -> App ()) -> App ())
-> (HasCallStack => String -> App ()) -> App ()
forall a b. (a -> b) -> a -> b
$ \String
domain -> do
      alice <- String -> CreateUser -> App Value
forall domain.
(HasCallStack, MakesValue domain) =>
domain -> CreateUser -> App Value
randomUser String
domain CreateUser
forall a. Default a => a
def
      aliceEmail <- asString $ alice %. "email"

      let testCookieLimit String
label = do
            let loginFn :: String -> String -> String -> App Response
loginFn = if String
label String -> String -> Bool
forall a. Eq a => a -> a -> Bool
== String
"persistent" then String -> String -> String -> App Response
forall domain.
(HasCallStack, MakesValue domain) =>
domain -> String -> String -> App Response
login else String -> String -> String -> App Response
forall domain.
(HasCallStack, MakesValue domain) =>
domain -> String -> String -> App Response
loginWithSessionCookie
            (deletedCookie1 : deletedCookie2 : validCookies) <-
              Int -> App String -> App [String]
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM (Int
cookieLimit Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
2)
                (App String -> App [String]) -> App String -> App [String]
forall a b. (a -> b) -> a -> b
$ do
                  -- This threadDelay is required to get around problems caused
                  -- by: https://wearezeta.atlassian.net/browse/WPB-15446
                  Int -> App ()
forall (m :: * -> *). MonadIO m => Int -> m ()
threadDelay Int
1_000_000
                  String -> String -> String -> App Response
loginFn String
domain String
aliceEmail String
defPassword
                    App Response -> (Response -> App String) -> App String
forall a.
HasCallStack =>
App Response -> (Response -> App a) -> App a
`bindResponse` \Response
resp -> do
                      Response
resp.status Int -> Int -> App ()
forall a. (MakesValue a, HasCallStack) => a -> Int -> App ()
`shouldMatchInt` Int
200
                      String -> App String
forall a. a -> App a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (String -> App String)
-> (Maybe String -> String) -> Maybe String -> App String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Maybe String -> String
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe String -> App String) -> Maybe String -> App String
forall a b. (a -> b) -> a -> b
$ String -> Response -> Maybe String
getCookie String
"zuid" Response
resp
            addFailureContext ("deletedCookie1: " <> deletedCookie1 <> "\ndeletedCookie2: " <> deletedCookie2 <> "\nvalidCookies:\n" <> unlines validCookies) $ do
              forM_ [deletedCookie1, deletedCookie2] $ \String
deletedCookie -> do
                Value -> String -> App Response
forall uid.
(HasCallStack, MakesValue uid) =>
uid -> String -> App Response
renewToken Value
alice String
deletedCookie App Response -> (Response -> App ()) -> App ()
forall a.
HasCallStack =>
App Response -> (Response -> App a) -> App a
`bindResponse` \Response
resp ->
                  Response
resp.status Int -> Int -> App ()
forall a. (MakesValue a, HasCallStack) => a -> Int -> App ()
`shouldMatchInt` Int
403
              forM_ validCookies $ \String
validCookie ->
                Value -> String -> App Response
forall uid.
(HasCallStack, MakesValue uid) =>
uid -> String -> App Response
renewToken Value
alice String
validCookie App Response -> (Response -> App ()) -> App ()
forall a.
HasCallStack =>
App Response -> (Response -> App a) -> App a
`bindResponse` \Response
resp ->
                  Response
resp.status Int -> Int -> App ()
forall a. (MakesValue a, HasCallStack) => a -> Int -> App ()
`shouldMatchInt` Int
200
      concurrently_ (testCookieLimit "persistent") (testCookieLimit "session")

-- @END

-- The testInvalidCookie test conforms to the following testing standards:
-- @SF.Provisioning @TSFI.RESTfulAPI @TSFI.NTP @S2
--
-- Test that invalid and expired tokens do not work.
testInvalidCookie :: (HasCallStack) => App ()
testInvalidCookie :: HasCallStack => App ()
testInvalidCookie = do
  let cookieTimeout :: Int
cookieTimeout = Int
2
  ServiceOverrides -> (HasCallStack => String -> App ()) -> App ()
forall a.
HasCallStack =>
ServiceOverrides -> (HasCallStack => String -> App a) -> App a
withModifiedBackend
    ServiceOverrides
forall a. Default a => a
def
      { brigCfg =
          setField @_ @Int "zauth.authSettings.userTokenTimeout" cookieTimeout
            >=> setField @_ @Int "zauth.authSettings.legalHoldUserTokenTimeout" cookieTimeout
      }
    ((HasCallStack => String -> App ()) -> App ())
-> (HasCallStack => String -> App ()) -> App ()
forall a b. (a -> b) -> a -> b
$ \String
domain -> do
      String -> String -> App Response
forall domain cookie.
(HasCallStack, MakesValue domain, MakesValue cookie) =>
domain -> cookie -> App Response
Nginz.access String
domain String
"zuid=xxx" App Response -> (Response -> App ()) -> App ()
forall a.
HasCallStack =>
App Response -> (Response -> App a) -> App a
`bindResponse` \Response
resp -> do
        Response
resp.status Int -> Int -> App ()
forall a. (MakesValue a, HasCallStack) => a -> Int -> App ()
`shouldMatchInt` Int
403
        Response
resp.json App Value -> String -> App Value
forall a. (HasCallStack, MakesValue a) => a -> String -> App Value
%. String
"label" App Value -> String -> App ()
forall a b.
(MakesValue a, MakesValue b, HasCallStack) =>
a -> b -> App ()
`shouldMatch` String
"client-error"
        msg <- App Value -> App String
forall a. (HasCallStack, MakesValue a) => a -> App String
asString (App Value -> App String) -> App Value -> App String
forall a b. (a -> b) -> a -> b
$ Response
resp.json App Value -> String -> App Value
forall a. (HasCallStack, MakesValue a) => a -> String -> App Value
%. String
"message"
        msg `shouldContain` "Invalid zauth token"

      (owner, tid, [alice]) <- String -> Int -> App (Value, String, [Value])
forall domain.
(HasCallStack, MakesValue domain) =>
domain -> Int -> App (Value, String, [Value])
createTeam String
domain Int
2
      aliceEmail <- asString $ alice %. "email"
      aliceId <- asString $ alice %. "qualified_id.id"
      userCookie <-
        login domain aliceEmail defPassword `bindResponse` \Response
resp -> do
          Response
resp.status Int -> Int -> App ()
forall a. (MakesValue a, HasCallStack) => a -> Int -> App ()
`shouldMatchInt` Int
200
          String -> App String
forall a. a -> App a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (String -> App String)
-> (Maybe String -> String) -> Maybe String -> App String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Maybe String -> String
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe String -> App String) -> Maybe String -> App String
forall a b. (a -> b) -> a -> b
$ String -> Response -> Maybe String
getCookie String
"zuid" Response
resp

      legalholdWhitelistTeam tid owner >>= assertSuccess

      lhCookie <-
        legalholdLogin domain aliceId defPassword `bindResponse` \Response
resp -> do
          Response
resp.status Int -> Int -> App ()
forall a. (MakesValue a, HasCallStack) => a -> Int -> App ()
`shouldMatchInt` Int
200
          String -> App String
forall a. a -> App a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (String -> App String)
-> (Maybe String -> String) -> Maybe String -> App String
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Maybe String -> String
forall a. HasCallStack => Maybe a -> a
fromJust (Maybe String -> App String) -> Maybe String -> App String
forall a b. (a -> b) -> a -> b
$ String -> Response -> Maybe String
getCookie String
"zuid" Response
resp

      -- Wait for both cookies to expire
      -- In order to reduce flakiness, the delay has been increased by 2 seconds
      -- after an investigation revealed that check for cookie expiration is somewhat lenient due to
      -- the way the time is calculated in the backend.
      -- See the interpreter of Now which is implemented using `Control.AutoUpdate` which defaults to an update frequency of 1 sec.
      -- Furthermore the timestamp is then again rounded down to the nearest second before it is compared to the cookie expiration time.
      threadDelay $ (cookieTimeout + 2) * 1_000_000
      -- Assert that the cookies are considered expired
      for_ [userCookie, lhCookie] $ \String
cookie ->
        String -> String -> App Response
forall domain cookie.
(HasCallStack, MakesValue domain, MakesValue cookie) =>
domain -> cookie -> App Response
Nginz.access String
domain (String
"zuid=" String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
cookie) App Response -> (Response -> App ()) -> App ()
forall a.
HasCallStack =>
App Response -> (Response -> App a) -> App a
`bindResponse` \Response
resp -> do
          Response
resp.status Int -> Int -> App ()
forall a. (MakesValue a, HasCallStack) => a -> Int -> App ()
`shouldMatchInt` Int
403
          Response
resp.json App Value -> String -> App Value
forall a. (HasCallStack, MakesValue a) => a -> String -> App Value
%. String
"label" App Value -> String -> App ()
forall a b.
(MakesValue a, MakesValue b, HasCallStack) =>
a -> b -> App ()
`shouldMatch` String
"invalid-credentials"
          Response
resp.json App Value -> String -> App Value
forall a. (HasCallStack, MakesValue a) => a -> String -> App Value
%. String
"message" App Value -> String -> App ()
forall a b.
(MakesValue a, MakesValue b, HasCallStack) =>
a -> b -> App ()
`shouldMatch` String
"Zauth token expired"

-- @END