-- This file is part of the Wire Server implementation.
--
-- Copyright (C) 2026 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.Migration.Budget (testBudgetMigration) where

import API.Brig (login)
import API.Common (defPassword)
import Control.Monad.Codensity
import Control.Monad.Reader
import qualified Data.ByteString.Char8 as BSChar8
import SetupHelpers
import Testlib.Prelude
import Testlib.ResourcePool
import Text.Read (readMaybe)

-- | Drives the login rate-limit (budget) store through the cutover lifecycle
-- (cassandra -> migration-to-postgresql -> postgresql). Budget rows are
-- ephemeral TTL'd counters, so unlike the other store migrations there is no
-- backfill worker: Cassandra rows age out on their own while new writes flow
-- to Postgres. The 60s login-block timeout is what keeps rows alive across
-- backend restarts within this test.
--
-- P2 proves reads stay Cassandra-sourced during dual-write (u1's Cassandra row
-- still blocks). P3 proves a row dual-written under migration-to-postgresql is
-- visible to the pure Postgres interpreter (u2 still blocked), and that a
-- Cassandra-only row does NOT survive cutover (u1 unblocked) — the documented
-- no-backfill behavior.
testBudgetMigration :: (HasCallStack) => App ()
testBudgetMigration :: HasCallStack => App ()
testBudgetMigration = do
  resourcePool <- (Env -> ResourcePool BackendResource)
-> App (ResourcePool BackendResource)
forall r (m :: * -> *) a. MonadReader r m => (r -> a) -> m a
asks (.resourcePool)
  runCodensity (acquireResources 1 resourcePool) $ \[BackendResource
backend] -> do
    let domain :: String
domain = BackendResource
backend.berDomain

    -- P1 cassandra: budget rows are written only to Cassandra
    u1 <-
      Codensity App String -> forall b. (String -> App b) -> App b
forall k (m :: k -> *) a.
Codensity m a -> forall (b :: k). (a -> m b) -> m b
runCodensity (HasCallStack =>
BackendResource -> ServiceOverrides -> Codensity App String
BackendResource -> ServiceOverrides -> Codensity App String
startDynamicBackend BackendResource
backend (String -> ServiceOverrides
forall {b}. ToJSON b => b -> ServiceOverrides
conf String
"cassandra")) ((String -> App Value) -> App Value)
-> (String -> App Value) -> App Value
forall a b. (a -> b) -> a -> b
$ \String
_ -> do
        u1 <- String -> CreateUser -> App Value
forall domain.
(HasCallStack, MakesValue domain) =>
domain -> CreateUser -> App Value
randomUser String
domain CreateUser
forall a. Default a => a
def
        exhaustLogins domain u1
        pure u1

    -- P2 migration-to-postgresql: reads still come from Cassandra (source of
    -- truth), writes are mirrored to Postgres
    u2 <-
      runCodensity (startDynamicBackend backend (conf "migration-to-postgresql")) $ \String
_ -> do
        -- u1 is still blocked: its budget row lives in Cassandra and reads
        -- are served from there during migration.
        HasCallStack => String -> Value -> App ()
String -> Value -> App ()
assertBlocked String
domain Value
u1
        u2 <- String -> CreateUser -> App Value
forall domain.
(HasCallStack, MakesValue domain) =>
domain -> CreateUser -> App Value
randomUser String
domain CreateUser
forall a. Default a => a
def
        exhaustLogins domain u2
        pure u2

    -- P3 postgresql: reads and writes are served exclusively from Postgres
    runCodensity (startDynamicBackend backend (conf "postgresql")) $ \String
_ -> do
      -- u2 is still blocked: the dual-written row is visible to the
      -- Postgres-only interpreter.
      HasCallStack => String -> Value -> App ()
String -> Value -> App ()
assertBlocked String
domain Value
u2
      -- u1 is NOT blocked: its budget row only ever existed in Cassandra and
      -- was never backfilled (no backfill worker for ephemeral data).
      HasCallStack => String -> Value -> App ()
String -> Value -> App ()
assertUnblocked String
domain Value
u1
      -- Write and read via the pure Postgres interpreter.
      u3 <- String -> CreateUser -> App Value
forall domain.
(HasCallStack, MakesValue domain) =>
domain -> CreateUser -> App Value
randomUser String
domain CreateUser
forall a. Default a => a
def
      exhaustLogins domain u3
  where
    conf :: b -> ServiceOverrides
conf b
db =
      ServiceOverrides
forall a. Default a => a
def
        { brigCfg =
            setField @_ @Int "optSettings.setLimitFailedLogins.timeout" 60
              >=> setField @_ @Int "optSettings.setPasswordHashingRateLimit.userLimit.inverseRate" 0
              >=> setField "postgresMigration.budget" db
        }

-- | Exhaust the failed-login budget of a user: 5 wrong passwords get plain
-- login-failure responses, and the 6th attempt (even with the correct
-- password) is blocked with a @Retry-After@ header.
exhaustLogins :: (HasCallStack) => String -> Value -> App ()
exhaustLogins :: HasCallStack => String -> Value -> App ()
exhaustLogins String
domain Value
u = do
  email <- App Value -> App String
forall a. (HasCallStack, MakesValue a) => a -> App String
asString (Value
u Value -> String -> App Value
forall a. (HasCallStack, MakesValue a) => a -> String -> App Value
%. String
"email")
  for_ [1 :: Int .. 5] $ \Int
_ ->
    String -> String -> String -> App Response
forall domain.
(HasCallStack, MakesValue domain) =>
domain -> String -> String -> App Response
login String
domain String
email 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 Maybe 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"
  assertBlockedEmail domain email

assertBlocked :: (HasCallStack) => String -> Value -> App ()
assertBlocked :: HasCallStack => String -> Value -> App ()
assertBlocked String
domain Value
u = App Value -> App String
forall a. (HasCallStack, MakesValue a) => a -> App String
asString (Value
u Value -> String -> App Value
forall a. (HasCallStack, MakesValue a) => a -> String -> App Value
%. String
"email") App String -> (String -> App ()) -> App ()
forall a b. App a -> (a -> App b) -> App b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= HasCallStack => String -> String -> App ()
String -> String -> App ()
assertBlockedEmail String
domain

assertBlockedEmail :: (HasCallStack) => String -> String -> App ()
assertBlockedEmail :: HasCallStack => String -> String -> App ()
assertBlockedEmail String
domain String
email =
  String -> String -> String -> App Response
forall domain.
(HasCallStack, MakesValue domain) =>
domain -> String -> String -> App Response
login String
domain String
email String
defPassword 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 Maybe 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"
    -- The Retry-After header distinguishes LoginBlocked from a plain
    -- login-failure.
    let Just (Int
_retryAfter :: Int) = 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
    () -> App ()
forall a. a -> App a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()

assertUnblocked :: (HasCallStack) => String -> Value -> App ()
assertUnblocked :: HasCallStack => String -> Value -> App ()
assertUnblocked String
domain Value
u = do
  email <- App Value -> App String
forall a. (HasCallStack, MakesValue a) => a -> App String
asString (Value
u Value -> String -> App Value
forall a. (HasCallStack, MakesValue a) => a -> String -> App Value
%. String
"email")
  login domain email defPassword >>= assertSuccess