-- 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.PasswordReset (testPasswordResetMigration) where

import API.Brig
import API.BrigInternal (getPasswordResetCode)
import qualified API.BrigInternal as BI
import API.Common (defPassword)
import Control.Monad.Codensity
import Control.Monad.Reader
import SetupHelpers
import Test.Migration.Util (waitForMigration)
import Testlib.Prelude
import Testlib.ResourcePool

-- | Drives the password-reset store through the full cutover lifecycle
-- (cassandra -> migration-to-postgresql -> postgresql). Reset keys are
-- deterministic per user and a second reset for the same user is a silent
-- no-op, so every phase writes with a fresh user and reads (and verifies) the
-- write of the previous phase. This exercises the writes of all three
-- interpreters (Cassandra, dual-write, Postgres) as well as their reads.
--
-- The rows created before the cutover (u1: pure Cassandra write, u2/u3:
-- dual-write) must be served by the pure-Postgres interpreter afterwards — u1
-- proves the background worker backfilled a row that only ever existed in
-- Cassandra.
--
-- u4 also covers the retry edge cases: a wrong code decrements the remaining
-- retries (and refreshes the row's expiry), and exhausting the retries deletes
-- the row, after which even the correct code is rejected and the password is
-- unchanged.
testPasswordResetMigration :: (HasCallStack) => App ()
testPasswordResetMigration :: HasCallStack => App ()
testPasswordResetMigration = 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: write via the pure Cassandra interpreter
    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 -> Bool -> ServiceOverrides
forall {b} {b}. (ToJSON b, ToJSON b) => b -> b -> ServiceOverrides
conf String
"cassandra" Bool
False)) ((String -> App ResetUser) -> App ResetUser)
-> (String -> App ResetUser) -> App ResetUser
forall a b. (a -> b) -> a -> b
$ \String
_ ->
        HasCallStack => String -> App ResetUser
String -> App ResetUser
initiateReset String
domain

    -- P2 migration-to-postgresql (worker off): reads still come from
    -- Cassandra, writes go to both stores
    u2 <-
      runCodensity (startDynamicBackend backend (conf "migration-to-postgresql" False)) $ \String
_ -> do
        HasCallStack => String -> ResetUser -> App ()
String -> ResetUser -> App ()
checkCode String
domain ResetUser
u1
        HasCallStack => String -> App ResetUser
String -> App ResetUser
initiateReset String
domain

    -- P3 migration-to-postgresql (worker on): dual-write while backfilling
    u3 <-
      runCodensity (startDynamicBackend backend (conf "migration-to-postgresql" True)) $ \String
_ -> do
        HasCallStack => String -> ResetUser -> App ()
String -> ResetUser -> App ()
checkCode String
domain ResetUser
u2
        -- A wrong code decrements the retries and re-inserts the row (with a
        -- refreshed expiry); the API rejects it with a 400.
        String -> String -> String -> String -> App Response
forall domain.
(HasCallStack, MakesValue domain) =>
domain -> String -> String -> String -> App Response
completePasswordReset String
domain ResetUser
u2.key (ResetUser
u2.code String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
"X") String
"some-password" App Response -> (Response -> 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 => Int -> Response -> App ()
Int -> Response -> App ()
assertStatus Int
400
        u3 <- HasCallStack => String -> App ResetUser
String -> App ResetUser
initiateReset String
domain
        waitForMigration domain counterName
        pure u3

    -- P4 postgresql: reads are served exclusively from Postgres
    runCodensity (startDynamicBackend backend (conf "postgresql" False)) $ \String
_ -> do
      -- Rows written by every interpreter are visible to the pure Postgres
      -- interpreter.
      HasCallStack => String -> ResetUser -> App ()
String -> ResetUser -> App ()
checkCode String
domain ResetUser
u1
      HasCallStack => String -> ResetUser -> App ()
String -> ResetUser -> App ()
checkCode String
domain ResetUser
u2
      HasCallStack => String -> ResetUser -> App ()
String -> ResetUser -> App ()
checkCode String
domain ResetUser
u3

      -- The decremented row from P3 still completes the flow.
      String -> String -> String -> String -> App Response
forall domain.
(HasCallStack, MakesValue domain) =>
domain -> String -> String -> String -> App Response
completePasswordReset String
domain ResetUser
u2.key ResetUser
u2.code String
"shiny-new-password" App Response -> (Response -> 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 => Response -> App ()
Response -> App ()
assertSuccess
      String -> String -> String -> App Response
forall domain.
(HasCallStack, MakesValue domain) =>
domain -> String -> String -> App Response
login String
domain ResetUser
u2.email String
"shiny-new-password" App Response -> (Response -> 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 => Response -> App ()
Response -> App ()
assertSuccess

      -- Write and read via the pure Postgres interpreter.
      u4 <- HasCallStack => String -> App ResetUser
String -> App ResetUser
initiateReset String
domain
      checkCode domain u4

      -- Retry depletion: three wrong codes exhaust the retries (3 -> 2 -> 1
      -- -> deleted). Afterwards even the correct code is rejected and the
      -- password is unchanged.
      for_ [1 :: Int .. 3] $ \Int
_ ->
        String -> String -> String -> String -> App Response
forall domain.
(HasCallStack, MakesValue domain) =>
domain -> String -> String -> String -> App Response
completePasswordReset String
domain ResetUser
u4.key (ResetUser
u4.code String -> String -> String
forall a. Semigroup a => a -> a -> a
<> String
"X") String
"some-password" App Response -> (Response -> 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 => Int -> Response -> App ()
Int -> Response -> App ()
assertStatus Int
400
      getPasswordResetCode domain u4.email >>= assertStatus 400
      completePasswordReset domain u4.key u4.code "shiny-new-password" >>= assertStatus 400
      login domain u4.email defPassword >>= assertSuccess
  where
    conf :: b -> b -> ServiceOverrides
conf b
db b
runMigration =
      ServiceOverrides
forall a. Default a => a
def
        { brigCfg = setField "postgresMigration.passwordReset" db,
          backgroundWorkerCfg =
            setField "postgresMigration.passwordReset" db
              >=> setField "migratePasswordReset" runMigration
        }
    counterName :: String
counterName = String
"^wire_password_reset_migration_finished"

-- | A user together with the reset data of an initiated password reset.
data ResetUser = ResetUser
  { ResetUser -> String
email :: String,
    ResetUser -> String
key :: String,
    ResetUser -> String
code :: String
  }

-- | Create a fresh user (with a known password) and initiate a password reset
-- for it, returning the reset key and code. Reset keys are deterministic per
-- user and a second reset for the same user is a silent no-op, which is why
-- every phase-write uses a fresh user.
initiateReset :: (HasCallStack) => String -> App ResetUser
initiateReset :: HasCallStack => String -> App ResetUser
initiateReset String
domain = do
  user <- String -> CreateUser -> App Value
forall domain.
(HasCallStack, MakesValue domain) =>
domain -> CreateUser -> App Value
randomUser String
domain CreateUser
forall a. Default a => a
def {BI.password = Just defPassword}
  email <- user %. "email" & asString
  passwordReset domain email >>= assertSuccess
  (key, code) <- getResetData domain email
  pure ResetUser {email = email, key = key, code = code}

getResetData :: (HasCallStack) => String -> String -> App (String, String)
getResetData :: HasCallStack => String -> String -> App (String, String)
getResetData String
domain String
email =
  App Response
-> (Response -> App (String, String)) -> App (String, String)
forall a.
HasCallStack =>
App Response -> (Response -> App a) -> App a
bindResponse (String -> String -> App Response
forall domain.
(HasCallStack, MakesValue domain) =>
domain -> String -> App Response
getPasswordResetCode String
domain String
email) ((Response -> App (String, String)) -> App (String, String))
-> (Response -> App (String, String)) -> App (String, String)
forall a b. (a -> b) -> a -> b
$ \Response
resp -> do
    Response
resp.status Int -> Int -> App ()
forall a. (MakesValue a, HasCallStack) => a -> Int -> App ()
`shouldMatchInt` Int
200
    (,) (String -> String -> (String, String))
-> App String -> App (String -> (String, String))
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (Response
resp.json Maybe Value -> String -> App Value
forall a. (HasCallStack, MakesValue a) => a -> String -> App Value
%. String
"key" App Value -> (App Value -> App String) -> App String
forall a b. a -> (a -> b) -> b
& App Value -> App String
forall a. (HasCallStack, MakesValue a) => a -> App String
asString) App (String -> (String, String))
-> App String -> App (String, String)
forall a b. App (a -> b) -> App a -> App b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (Response
resp.json Maybe Value -> String -> App Value
forall a. (HasCallStack, MakesValue a) => a -> String -> App Value
%. String
"code" App Value -> (App Value -> App String) -> App String
forall a b. a -> (a -> b) -> b
& App Value -> App String
forall a. (HasCallStack, MakesValue a) => a -> App String
asString)

-- | The stored reset code is still readable and unchanged.
checkCode :: (HasCallStack) => String -> ResetUser -> App ()
checkCode :: HasCallStack => String -> ResetUser -> App ()
checkCode String
domain ResetUser
u = do
  (key, code) <- HasCallStack => String -> String -> App (String, String)
String -> String -> App (String, String)
getResetData String
domain ResetUser
u.email
  key `shouldMatch` u.key
  code `shouldMatch` u.code