[WPB-11101] Refactor invitation store

changelog.d/5-internal/WPB-11101-internal-types

@@ -0,0 +1 @@
+Improve abstraction in the invitation store and hide DB interaction-specific internal types from the application code.

libs/wire-subsystems/src/Wire/InvitationCodeStore.hs → libs/wire-subsystems/src/Wire/InvitationStore.hs

@@ -18,25 +18,21 @@
 {-# LANGUAGE StrictData #-}
 {-# LANGUAGE TemplateHaskell #-}
 
-module Wire.InvitationCodeStore where
+module Wire.InvitationStore where
 
-import Control.Monad.Trans.Maybe (MaybeT (MaybeT, runMaybeT))
 import Data.Id (InvitationId, TeamId, UserId)
 import Data.Json.Util (UTCTimeMillis)
 import Data.Range (Range)
 import Database.CQL.Protocol (Record (..), TupleType, recordInstance)
 import Imports
 import Polysemy
-import Polysemy.TinyLog (TinyLog)
-import System.Logger.Message qualified as Log
 import URI.ByteString
 import Util.Timeout
 import Wire.API.Team.Invitation (Invitation (inviteeEmail))
 import Wire.API.Team.Invitation qualified as Public
 import Wire.API.Team.Role (Role, defaultRole)
 import Wire.API.User (EmailAddress, InvitationCode, Name)
 import Wire.Arbitrary (Arbitrary, GenericUniform (..))
-import Wire.Sem.Logger qualified as Log
 
 data StoredInvitation = MkStoredInvitation
   { teamId :: TeamId,
@@ -53,16 +49,6 @@ data StoredInvitation = MkStoredInvitation
 
 recordInstance ''StoredInvitation
 
-data StoredInvitationInfo = MkStoredInvitationInfo
-  { teamId :: TeamId,
-    invitationId :: InvitationId,
-    code :: InvitationCode
-  }
-  deriving (Show, Eq, Generic)
-  deriving (Arbitrary) via (GenericUniform StoredInvitationInfo)
-
-recordInstance ''StoredInvitationInfo
-
 data InsertInvitation = MkInsertInvitation
   { invitationId :: InvitationId,
     teamId :: TeamId,
@@ -84,45 +70,21 @@ data PaginatedResult a
 
 ----------------------------
 
-data InvitationCodeStore :: Effect where
-  InsertInvitation :: InsertInvitation -> Timeout -> InvitationCodeStore m StoredInvitation
-  LookupInvitation :: TeamId -> InvitationId -> InvitationCodeStore m (Maybe StoredInvitation)
-  LookupInvitationInfo :: InvitationCode -> InvitationCodeStore m (Maybe StoredInvitationInfo)
-  LookupInvitationCodesByEmail :: EmailAddress -> InvitationCodeStore m [StoredInvitationInfo]
+data InvitationStore :: Effect where
+  InsertInvitation :: InsertInvitation -> Timeout -> InvitationStore m StoredInvitation
+  LookupInvitation :: TeamId -> InvitationId -> InvitationStore m (Maybe StoredInvitation)
+  LookupInvitationByCode :: InvitationCode -> InvitationStore m (Maybe StoredInvitation)
+  LookupInvitationsByEmail :: EmailAddress -> InvitationStore m [StoredInvitation]
   -- | Range is page size, it defaults to 100
-  LookupInvitationsPaginated :: Maybe (Range 1 500 Int32) -> TeamId -> Maybe InvitationId -> InvitationCodeStore m (PaginatedResult [StoredInvitation])
-  CountInvitations :: TeamId -> InvitationCodeStore m Int64
-  DeleteInvitation :: TeamId -> InvitationId -> InvitationCodeStore m ()
-  DeleteAllTeamInvitations :: TeamId -> InvitationCodeStore m ()
+  LookupInvitationsPaginated :: Maybe (Range 1 500 Int32) -> TeamId -> Maybe InvitationId -> InvitationStore m (PaginatedResult [StoredInvitation])
+  CountInvitations :: TeamId -> InvitationStore m Int64
+  DeleteInvitation :: TeamId -> InvitationId -> InvitationStore m ()
+  DeleteAllTeamInvitations :: TeamId -> InvitationStore m ()
 
-makeSem ''InvitationCodeStore
+makeSem ''InvitationStore
 
 ----------------------------
 
-lookupInvitationByEmail :: (Member InvitationCodeStore r, Member TinyLog r) => EmailAddress -> Sem r (Maybe StoredInvitation)
-lookupInvitationByEmail email = runMaybeT do
-  MkStoredInvitationInfo {teamId, invitationId} <- MaybeT $ lookupSingleInvitationCodeByEmail email
-  MaybeT $ lookupInvitation teamId invitationId
-
-lookupInvitationByCode :: (Member InvitationCodeStore r) => InvitationCode -> Sem r (Maybe StoredInvitation)
-lookupInvitationByCode code = runMaybeT do
-  info <- MaybeT $ lookupInvitationInfo code
-  MaybeT $ lookupInvitation info.teamId info.invitationId
-
-lookupSingleInvitationCodeByEmail :: (Member TinyLog r, Member InvitationCodeStore r) => EmailAddress -> Sem r (Maybe StoredInvitationInfo)
-lookupSingleInvitationCodeByEmail email = do
-  invs <- lookupInvitationCodesByEmail email
-  case invs of
-    [] -> pure Nothing
-    [inv] -> pure $ Just inv
-    (_ : _ : _) -> do
-      -- edge case: more than one pending invite from different teams
-      Log.info $
-        Log.msg (Log.val "team_invidation_email: multiple pending invites from different teams for the same email")
-          . Log.field "email" (show email)
-
-      pure Nothing
-
 invitationFromStored :: Maybe (URIRef Absolute) -> StoredInvitation -> Public.Invitation
 invitationFromStored maybeUrl MkStoredInvitation {..} =
   Public.Invitation

libs/wire-subsystems/src/Wire/InvitationCodeStore/Cassandra.hs → libs/wire-subsystems/src/Wire/InvitationStore/Cassandra.hs

@@ -1,31 +1,33 @@
-module Wire.InvitationCodeStore.Cassandra where
+module Wire.InvitationStore.Cassandra
+  ( interpretInvitationStoreToCassandra,
+  )
+where
 
 import Cassandra
+import Control.Monad.Trans.Maybe
 import Data.Conduit (runConduit, (.|))
 import Data.Conduit.List qualified as Conduit
 import Data.Id
 import Data.Json.Util (UTCTimeMillis, toUTCTimeMillis)
 import Data.Range (Range, fromRange)
-import Data.Text.Ascii (encodeBase64Url)
-import Database.CQL.Protocol (TupleType, asRecord)
+import Database.CQL.Protocol (Record (..), TupleType, asRecord)
 import Imports
-import OpenSSL.Random (randBytes)
 import Polysemy
 import Polysemy.Embed
 import UnliftIO.Async (pooledMapConcurrentlyN_)
 import Util.Timeout
 import Wire.API.Team.Role (Role)
 import Wire.API.User
-import Wire.InvitationCodeStore
+import Wire.InvitationStore
 
-interpretInvitationCodeStoreToCassandra :: (Member (Embed IO) r) => ClientState -> InterpreterFor InvitationCodeStore r
-interpretInvitationCodeStoreToCassandra casClient =
+interpretInvitationStoreToCassandra :: (Member (Embed IO) r) => ClientState -> InterpreterFor InvitationStore r
+interpretInvitationStoreToCassandra casClient =
   interpret $
     runEmbedded (runClient casClient) . \case
       InsertInvitation newInv timeout -> embed $ insertInvitationImpl newInv timeout
       LookupInvitation tid iid -> embed $ lookupInvitationImpl tid iid
-      LookupInvitationCodesByEmail email -> embed $ lookupInvitationCodesByEmailImpl email
-      LookupInvitationInfo code -> embed $ lookupInvitationInfoImpl code
+      LookupInvitationsByEmail email -> embed $ lookupInvitationsByEmailImpl email
+      LookupInvitationByCode code -> embed $ lookupInvitationByCodeImpl code
       LookupInvitationsPaginated mSize tid miid -> embed $ lookupInvitationsPaginatedImpl mSize tid miid
       CountInvitations tid -> embed $ countInvitationsImpl tid
       DeleteInvitation tid invId -> embed $ deleteInvitationImpl tid invId
@@ -108,24 +110,41 @@ countInvitationsImpl t =
     cql :: PrepQuery R (Identity TeamId) (Identity Int64)
     cql = [sql| SELECT count(*) FROM team_invitation WHERE team = ?|]
 
-lookupInvitationInfoImpl :: InvitationCode -> Client (Maybe StoredInvitationInfo)
-lookupInvitationInfoImpl code =
-  fmap asRecord <$> retry x1 (query1 cql (params LocalQuorum (Identity code)))
+lookupInvitationByCodeImpl :: InvitationCode -> Client (Maybe StoredInvitation)
+lookupInvitationByCodeImpl code = runMaybeT do
+  (teamId, invId, _) <-
+    MaybeT $
+      retry x1 (query1 cqlInfo (params LocalQuorum (Identity code)))
+  MaybeT $ fmap asRecord <$> retry x1 (query1 cqlMain (params LocalQuorum (teamId, invId)))
   where
-    cql :: PrepQuery R (Identity InvitationCode) (TupleType StoredInvitationInfo)
-    cql =
+    cqlInfo :: PrepQuery R (Identity InvitationCode) (TeamId, InvitationId, InvitationCode)
+    cqlInfo =
       [sql|
-      SELECT  team, id, code FROM team_invitation_info WHERE code = ?
+      SELECT team, id, code FROM team_invitation_info WHERE code = ?
+      |]
+    cqlMain :: PrepQuery R (TeamId, InvitationId) (TupleType StoredInvitation)
+    cqlMain =
+      [sql|
+      SELECT team, role, id, created_at, created_by, email, name, code FROM team_invitation WHERE team = ? AND id = ?
       |]
 
-lookupInvitationCodesByEmailImpl :: EmailAddress -> Client [StoredInvitationInfo]
-lookupInvitationCodesByEmailImpl email = map asRecord <$> retry x1 (query cql (params LocalQuorum (Identity email)))
+lookupInvitationsByEmailImpl :: EmailAddress -> Client [StoredInvitation]
+lookupInvitationsByEmailImpl email = do
+  infoList <-
+    retry x1 (query cqlInfo (params LocalQuorum (Identity email)))
+  fmap catMaybes $ forM infoList $ \(tid, invId, _invCode) ->
+    fmap asRecord <$> retry x1 (query1 cqlMain (params LocalQuorum (tid, invId)))
   where
-    cql :: PrepQuery R (Identity EmailAddress) (TeamId, InvitationId, InvitationCode)
-    cql =
+    cqlInfo :: PrepQuery R (Identity EmailAddress) (TeamId, InvitationId, InvitationCode)
+    cqlInfo =
       [sql|
       SELECT team, invitation, code FROM team_invitation_email WHERE email = ?
       |]
+    cqlMain :: PrepQuery R (TeamId, InvitationId) (TupleType StoredInvitation)
+    cqlMain =
+      [sql|
+      SELECT team, role, id, created_at, created_by, email, name, code FROM team_invitation WHERE team = ? AND id = ?
+      |]
 
 lookupInvitationImpl :: TeamId -> InvitationId -> Client (Maybe StoredInvitation)
 lookupInvitationImpl tid iid =
@@ -186,8 +205,3 @@ deleteInvitationsImpl teamId =
   where
     cqlSelect :: PrepQuery R (Identity TeamId) (Identity InvitationId)
     cqlSelect = "SELECT id FROM team_invitation WHERE team = ? ORDER BY id ASC"
-
--- | This function doesn't really belong here, and may want to have return type `Sem (Random :
--- ...)` instead of `IO`.  Meh.
-mkInvitationCode :: IO InvitationCode
-mkInvitationCode = InvitationCode . encodeBase64Url <$> randBytes 24

libs/wire-subsystems/src/Wire/TeamInvitationSubsystem/Interpreter.hs

@@ -27,8 +27,8 @@ import Wire.Arbitrary
 import Wire.EmailSubsystem
 import Wire.GalleyAPIAccess hiding (AddTeamMember)
 import Wire.GalleyAPIAccess qualified as GalleyAPIAccess
-import Wire.InvitationCodeStore (InvitationCodeStore, StoredInvitation)
-import Wire.InvitationCodeStore qualified as Store
+import Wire.InvitationStore (InvitationStore, StoredInvitation)
+import Wire.InvitationStore qualified as Store
 import Wire.Sem.Logger qualified as Log
 import Wire.Sem.Now (Now)
 import Wire.Sem.Now qualified as Now
@@ -52,7 +52,7 @@ runTeamInvitationSubsystem ::
     Member GalleyAPIAccess r,
     Member UserSubsystem r,
     Member Random r,
-    Member InvitationCodeStore r,
+    Member InvitationStore r,
     Member Now r,
     Member EmailSubsystem r
   ) =>
@@ -69,7 +69,7 @@ inviteUserImpl ::
     Member UserSubsystem r,
     Member TinyLog r,
     Member Random r,
-    Member InvitationCodeStore r,
+    Member InvitationStore r,
     Member (Input TeamInvitationSubsystemConfig) r,
     Member Now r,
     Member EmailSubsystem r
@@ -106,7 +106,7 @@ inviteUserImpl luid tid request = do
 createInvitation' ::
   ( Member GalleyAPIAccess r,
     Member UserSubsystem r,
-    Member InvitationCodeStore r,
+    Member InvitationStore r,
     Member TinyLog r,
     Member (Error TeamInvitationSubsystemError) r,
     Member Random r,

libs/wire-subsystems/src/Wire/UserSubsystem.hs

@@ -27,7 +27,7 @@ import Wire.API.User.Search
 import Wire.Arbitrary
 import Wire.GalleyAPIAccess (GalleyAPIAccess)
 import Wire.GalleyAPIAccess qualified as GalleyAPIAccess
-import Wire.InvitationCodeStore
+import Wire.InvitationStore
 import Wire.UserKeyStore (EmailKey, emailKeyOrig)
 import Wire.UserSearch.Types
 import Wire.UserSubsystem.Error (UserSubsystemError (..))

libs/wire-subsystems/src/Wire/UserSubsystem/Interpreter.hs

@@ -55,7 +55,7 @@ import Wire.GalleyAPIAccess qualified as GalleyAPIAccess
 import Wire.IndexedUserStore (IndexedUserStore)
 import Wire.IndexedUserStore qualified as IndexedUserStore
 import Wire.IndexedUserStore.Bulk.ElasticSearch (teamSearchVisibilityInbound)
-import Wire.InvitationCodeStore
+import Wire.InvitationStore
 import Wire.Sem.Concurrency
 import Wire.Sem.Metrics
 import Wire.Sem.Metrics qualified as Metrics
@@ -99,7 +99,7 @@ runUserSubsystem ::
     Member IndexedUserStore r,
     Member FederationConfigStore r,
     Member Metrics r,
-    Member InvitationCodeStore r,
+    Member InvitationStore r,
     Member TinyLog r
   ) =>
   UserSubsystemConfig ->
@@ -172,7 +172,7 @@ runUserSubsystem cfg authInterpreter =
           internalFindTeamInvitationImpl mEmailKey code
 
 internalFindTeamInvitationImpl ::
-  ( Member InvitationCodeStore r,
+  ( Member InvitationStore r,
     Member (Error UserSubsystemError) r,
     Member (Input UserSubsystemConfig) r,
     Member (GalleyAPIAccess) r,
@@ -183,15 +183,11 @@ internalFindTeamInvitationImpl ::
   Sem r StoredInvitation
 internalFindTeamInvitationImpl Nothing _ = throw UserSubsystemMissingIdentity
 internalFindTeamInvitationImpl (Just e) c =
-  lookupInvitationInfo c >>= \case
-    Just invitationInfo -> do
-      inv <- lookupInvitation invitationInfo.teamId invitationInfo.invitationId
-      case (inv, (.email) <$> inv) of
-        (Just invite, Just em)
-          | e == mkEmailKey em -> do
-              ensureMemberCanJoin invitationInfo.teamId
-              pure invite
-        _ -> throw UserSubsystemInvalidInvitationCode
+  lookupInvitationByCode c >>= \case
+    Just inv -> do
+      if e == mkEmailKey (inv.email)
+        then ensureMemberCanJoin inv.teamId $> inv
+        else throw UserSubsystemInvalidInvitationCode
     Nothing -> throw UserSubsystemInvalidInvitationCode
   where
     ensureMemberCanJoin tid = do
@@ -849,8 +845,7 @@ getAccountsByImpl ::
   ( Member UserStore r,
     Member DeleteQueue r,
     Member (Input UserSubsystemConfig) r,
-    Member InvitationCodeStore r,
-    Member TinyLog r
+    Member InvitationStore r
   ) =>
   Local GetBy ->
   Sem r [User]
@@ -883,7 +878,7 @@ getAccountsByImpl (tSplit -> (domain, MkGetBy {includePendingInvitations, getByH
                 -- validated one cannot be found.  that's probably wrong?  split up into
                 -- validEmailIdentity, anyEmailIdentity?
                 Just email -> do
-                  hasInvitation <- isJust <$> lookupInvitationByEmail email
+                  hasInvitation <- isJust . listToMaybe <$> lookupInvitationsByEmail email
                   gcHack hasInvitation (User.userId user)
                   pure hasInvitation
                 Nothing -> error "getExtendedAccountsByImpl: should never happen, user invited via scim always has an email"
@@ -916,7 +911,7 @@ acceptTeamInvitationImpl ::
     Member UserStore r,
     Member GalleyAPIAccess r,
     Member (Error UserSubsystemError) r,
-    Member InvitationCodeStore r,
+    Member InvitationStore r,
     Member IndexedUserStore r,
     Member Metrics r,
     Member Events r,

libs/wire-subsystems/test/unit/Wire/MiniBackend.hs

@@ -72,10 +72,10 @@ import Wire.GalleyAPIAccess
 import Wire.HashPassword (HashPassword)
 import Wire.IndexedUserStore
 import Wire.InternalEvent hiding (DeleteUser)
-import Wire.InvitationCodeStore
+import Wire.InvitationStore
 import Wire.MockInterpreters
 import Wire.MockInterpreters.ActivationCodeStore (inMemoryActivationCodeStoreInterpreter)
-import Wire.MockInterpreters.InvitationCodeStore (inMemoryInvitationCodeStoreInterpreter)
+import Wire.MockInterpreters.InvitationStore (inMemoryInvitationStoreInterpreter)
 import Wire.PasswordResetCodeStore
 import Wire.PasswordStore
 import Wire.Sem.Concurrency
@@ -137,10 +137,10 @@ type MiniBackendEffects = UserSubsystem ': MiniBackendLowerEffects
 type MiniBackendLowerEffects =
   [ EmailSubsystem,
     GalleyAPIAccess,
-    InvitationCodeStore,
+    InvitationStore,
     PasswordStore,
     State (Map (TeamId, InvitationId) StoredInvitation),
-    State (Map InvitationCode StoredInvitationInfo),
+    State (Map InvitationCode StoredInvitation),
     ActivationCodeStore,
     State (Map EmailKey (Maybe UserId, ActivationCode)),
     BlockListStore,
@@ -178,7 +178,7 @@ data MiniBackend = MkMiniBackend
     passwordResetCodes :: Map PasswordResetKey (PRQueryData Identity),
     blockList :: [EmailKey],
     activationCodes :: Map EmailKey (Maybe UserId, ActivationCode),
-    invitationInfos :: Map InvitationCode StoredInvitationInfo,
+    invitationInfos :: Map InvitationCode StoredInvitation,
     invitations :: Map (TeamId, InvitationId) StoredInvitation
   }
   deriving stock (Eq, Show, Generic)
@@ -428,20 +428,20 @@ interpretMaybeFederationStackState maybeFederationAPIAccess localBackend teamMem
         . liftActivationCodeStoreState
         . inMemoryActivationCodeStoreInterpreter
         . liftInvitationInfoStoreState
-        . liftInvitationCodeStoreState
+        . liftInvitationStoreState
         . runInMemoryPasswordStoreInterpreter
-        . inMemoryInvitationCodeStoreInterpreter
+        . inMemoryInvitationStoreInterpreter
         . miniGalleyAPIAccess teamMember galleyConfigs
         . noopEmailSubsystemInterpreter
         . userSubsystemInterpreter
 
-liftInvitationInfoStoreState :: (Member (State MiniBackend) r) => Sem (State (Map InvitationCode StoredInvitationInfo) : r) a -> Sem r a
+liftInvitationInfoStoreState :: (Member (State MiniBackend) r) => Sem (State (Map InvitationCode StoredInvitation) : r) a -> Sem r a
 liftInvitationInfoStoreState = interpret \case
   Polysemy.State.Get -> gets (.invitationInfos)
   Put newAcs -> modify $ \b -> b {invitationInfos = newAcs}
 
-liftInvitationCodeStoreState :: (Member (State MiniBackend) r) => Sem (State (Map (TeamId, InvitationId) StoredInvitation) : r) a -> Sem r a
-liftInvitationCodeStoreState = interpret \case
+liftInvitationStoreState :: (Member (State MiniBackend) r) => Sem (State (Map (TeamId, InvitationId) StoredInvitation) : r) a -> Sem r a
+liftInvitationStoreState = interpret \case
   Polysemy.State.Get -> gets (.invitations)
   Put newInvs -> modify $ \b -> b {invitations = newInvs}