De-flake testIdpUpdate: atomic issuer deletion in Cassandra batch

changelog.d/3-bug-fixes/fix-idp-update-issuer-tombstone

@@ -0,0 +1 @@
+Fix flaky IdP updates caused by Cassandra tombstone shadowing when an issuer is reused across updates. Old issuer mappings are now deleted atomically within the same Cassandra batch as the new config insertion.

integration/test/Test/Spar.hs

@@ -1290,17 +1290,19 @@ testIdpUpdate = do
   -- update the IdP
   idp2 <- do
     (resp, meta) <- updateTestIdpWithMetaWithPrivateCreds owner idpId
+    resp.status `shouldMatchInt` 200
     (,meta) <$> asString (resp.json %. "id")
   -- the SCIM users can login
   for_ uids $ \(_, email) -> do
-    void $ loginWithSamlEmail True tid email idp2
+    eventually $ void $ loginWithSamlEmail True tid email idp2
   -- update the IdP again and use the original metadata
   idp3 <- do
     resp <- updateIdp owner idpId idpmeta
+    resp.status `shouldMatchInt` 200
     (,(idpmeta, pCreds)) <$> asString (resp.json %. "id")
   -- the SCIM users can still login
   for_ uids $ \(_, email) -> do
-    void $ loginWithSamlEmail True tid email idp3
+    eventually $ void $ loginWithSamlEmail True tid email idp3
 
 -- @SF.Provisioning @TSFI.RESTfulAPI @S2
 --

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

@@ -19,6 +19,7 @@
 
 module Wire.IdPConfigStore
   ( IdPConfigStore (..),
+    InsertMode (..),
     Replacing (..),
     Replaced (..),
     insertConfig,
@@ -49,8 +50,18 @@ newtype Replaced = Replaced SAML.IdPId
 newtype Replacing = Replacing SAML.IdPId
   deriving (Eq, Ord, Show)
 
+-- | Controls whether 'insertConfig' also cleans up old issuer mappings from the
+-- issuer lookup tables inside the same Cassandra batch.
+data InsertMode
+  = -- | Insert the config only; leave old issuer mappings intact.
+    InsertOnly
+  | -- | Insert the config and atomically delete old issuer mappings to prevent
+    -- Cassandra tombstone shadowing when an issuer is reused across updates.
+    InsertWithIssuerCleanup
+  deriving stock (Eq, Show)
+
 data IdPConfigStore m a where
-  InsertConfig :: IP.IdP -> IdPConfigStore m ()
+  InsertConfig :: InsertMode -> IP.IdP -> IdPConfigStore m ()
   NewHandle :: TeamId -> IdPConfigStore m IP.IdPHandle
   GetConfig :: SAML.IdPId -> IdPConfigStore m IP.IdP
   GetIdPByIssuerV1Maybe :: SAML.Issuer -> IdPConfigStore m (Maybe IP.IdP)

libs/wire-subsystems/src/Wire/IdPConfigStore/Cassandra.hs

@@ -39,7 +39,7 @@ import Wire.API.Routes.Public (ZHostValue)
 import Wire.API.User.IdentityProvider hiding (apiVersion, oldIssuers, replacedBy, team)
 import Wire.API.User.IdentityProvider qualified as IP
 import {- instance Cql SAML.IdPId -} Wire.DomainRegistrationStore.Cassandra ()
-import Wire.IdPConfigStore (IdPConfigStore (..), IdpDbError (..), Replaced (..), Replacing (..))
+import Wire.IdPConfigStore (IdPConfigStore (..), IdpDbError (..), InsertMode (..), Replaced (..), Replacing (..))
 import Wire.IdPConfigStore.Orphans ()
 
 idPToCassandra ::
@@ -50,7 +50,7 @@ idPToCassandra ::
 idPToCassandra =
   interpret $
     \case
-      InsertConfig iw -> embed @m (runExceptT $ insertIdPConfig iw) >>= either throw pure
+      InsertConfig mode iw -> embed @m (runExceptT $ insertIdPConfig mode iw) >>= either throw pure
       NewHandle tid -> embed @m (runExceptT $ newIdPHandleForTeam tid) >>= either throw pure
       GetConfig i -> embed @m (runExceptT $ getIdPConfig i) >>= either throw pure
       GetIdPByIssuerV1 i -> embed @m (runExceptT $ getIdPByIssuerV1 i) >>= either throw pure
@@ -72,10 +72,14 @@ type IdPConfigRow = (SAML.IdPId, SAML.Issuer, URI, SignedCertificate, [SignedCer
 insertIdPConfig ::
   forall m.
   (HasCallStack, MonadClient m, MonadError IdpDbError m) =>
+  InsertMode ->
   IdP ->
   m ()
-insertIdPConfig idp = do
+insertIdPConfig insertMode idp = do
   ensureDoNotMixApiVersions
+  let oldIssuers = idp ^. SAML.idpExtraInfo . IP.oldIssuers
+      tid = idp ^. SAML.idpExtraInfo . IP.team
+      thisVersion = fromMaybe defWireIdPAPIVersion $ idp ^. SAML.idpExtraInfo . IP.apiVersion
   retry x5 . batch $ do
     setType BatchLogged
     setConsistency LocalQuorum
@@ -105,6 +109,20 @@ insertIdPConfig idp = do
       ( idp ^. SAML.idpId,
         idp ^. SAML.idpExtraInfo . IP.team
       )
+    -- Delete old issuer mappings atomically with the new config insertion to
+    -- prevent Cassandra tombstone shadowing when an issuer is reused across
+    -- multiple updates (e.g., updating from issuer A to B and back to A).
+    -- Only done when the caller requests issuer cleanup (e.g. on IdP update,
+    -- but not on creation where replaced IdPs must keep their mappings).
+    case insertMode of
+      InsertOnly -> pure ()
+      InsertWithIssuerCleanup ->
+        forM_ oldIssuers $ \oldIssuer ->
+          case thisVersion of
+            WireIdPAPIV2 -> addPrepQuery delOldIssuerV2 (oldIssuer, tid)
+            WireIdPAPIV1 -> do
+              addPrepQuery delOldIssuerV1 (Identity oldIssuer)
+              addPrepQuery delOldIssuerV1' (Identity oldIssuer)
   where
     ensureDoNotMixApiVersions :: m ()
     ensureDoNotMixApiVersions = do
@@ -129,6 +147,15 @@ insertIdPConfig idp = do
     byTeam :: PrepQuery W (SAML.IdPId, TeamId) ()
     byTeam = "INSERT INTO team_idp (idp, team) VALUES (?, ?)"
 
+    delOldIssuerV2 :: PrepQuery W (SAML.Issuer, TeamId) ()
+    delOldIssuerV2 = "DELETE FROM issuer_idp_v2 WHERE issuer = ? AND team = ?"
+
+    delOldIssuerV1 :: PrepQuery W (Identity SAML.Issuer) ()
+    delOldIssuerV1 = "DELETE FROM issuer_idp WHERE issuer = ?"
+
+    delOldIssuerV1' :: PrepQuery W (Identity SAML.Issuer) ()
+    delOldIssuerV1' = "DELETE FROM issuer_idp_v2 WHERE issuer = ?"
+
 newIdPHandleForTeam ::
   forall m.
   (HasCallStack, MonadClient m, MonadError IdpDbError m) =>

libs/wire-subsystems/src/Wire/IdPConfigStore/Mem.hs

@@ -44,7 +44,7 @@ idPToMem = evState . evEff
 
     evEff :: Sem (IdPConfigStore ': r) a -> Sem (State TypedState ': r) a
     evEff = reinterpret @_ @(State TypedState) $ \case
-      InsertConfig iw ->
+      InsertConfig _mode iw ->
         modify' (insertConfig iw)
       NewHandle _tid ->
         -- Same handle for all IdPs is good enough, for now

libs/wire-subsystems/test/unit/Wire/IdPSubsystem/InterpreterSpec.hs

@@ -70,7 +70,7 @@ runAllEffects ::
   (Either IdPSubsystemError a, [RecordedLog])
 runAllEffects enableDiscovery idps teams brigAPIMockFn action = swap $ run $ runState [] $ do
   (_idpState, result) <- idPToMem $ do
-    forM_ idps insertConfig
+    forM_ idps (insertConfig InsertOnly)
     -- Run the action
     miniGalleyAPIAccess teams def
       . brigAPIAccessMock brigAPIMockFn

services/spar/src/Spar/API.hs

@@ -132,7 +132,7 @@ import Wire.API.User.Saml
 import Wire.BrigAPIAccess (BrigAPIAccess)
 import qualified Wire.BrigAPIAccess as BrigAPIAccess
 import Wire.GalleyAPIAccess (GalleyAPIAccess)
-import Wire.IdPConfigStore (IdPConfigStore, Replaced (..), Replacing (..))
+import Wire.IdPConfigStore (IdPConfigStore, InsertMode (..), Replaced (..), Replacing (..))
 import qualified Wire.IdPConfigStore as IdPConfigStore
 import Wire.IdPSubsystem (IdPSubsystem)
 import qualified Wire.IdPSubsystem as IdPSubsystem
@@ -731,7 +731,7 @@ idpCreate samlConfig tid zUser uncheckedMbHost (IdPMetadataValue rawIdpMetadata
     maybe (IdPConfigStore.newHandle tid) (pure . IdPHandle . fromRange) mHandle
       >>= validateNewIdP apiversion idpmeta tid mReplaces mbHost
   IdPRawMetadataStore.store (idp ^. SAML.idpId) rawIdpMetadata
-  IdPConfigStore.insertConfig idp
+  IdPConfigStore.insertConfig InsertOnly idp
   forM_ mReplaces $ \replaces ->
     IdPConfigStore.setReplacedBy (Replaced replaces) (Replacing (idp ^. SAML.idpId))
   when (SAML.isMultiIngressConfig samlConfig) $
@@ -966,13 +966,10 @@ idpUpdateXML samlConfig mbZUsr mDomain raw idpmeta idpid mHandle = withDebugLog
   -- (if raw metadata is stored and then spar goes out, raw metadata won't match the
   -- structured idp config.  since this will lead to a 5xx response, the client is expected to
   -- try again, which would clean up cassandra state.)
-  IdPConfigStore.insertConfig idp''
-  -- if the IdP issuer is updated, the old issuer must be removed explicitly.
-  -- if this step is ommitted (due to a crash) resending the update request should fix the inconsistent state.
-  let mbteamid = case fromMaybe defWireIdPAPIVersion $ idp'' ^. SAML.idpExtraInfo . apiVersion of
-        WireIdPAPIV1 -> Nothing
-        WireIdPAPIV2 -> Just teamid
-  forM_ (idp'' ^. SAML.idpExtraInfo . oldIssuers) (flip IdPConfigStore.deleteIssuer mbteamid)
+  IdPConfigStore.insertConfig InsertWithIssuerCleanup idp''
+  -- Old issuer mappings are now deleted atomically as part of the same
+  -- Cassandra batch inside 'insertConfig' to prevent tombstone shadowing
+  -- when an issuer is reused across updates.
   when (SAML.isMultiIngressConfig samlConfig) $
     BrigAPIAccess.sendSAMLIdPChangedEmail $
       IdPUpdated zUsr previousIdP idp''

services/spar/src/Spar/Sem/SAML2/Library.hs

@@ -124,7 +124,7 @@ instance
   type IdPConfigExtra (SPImpl r) = WireIdP
   type IdPConfigSPId (SPImpl r) = TeamId
 
-  storeIdPConfig = SPImpl . IdPConfigStore.insertConfig
+  storeIdPConfig = SPImpl . IdPConfigStore.insertConfig IdPConfigStore.InsertOnly
   getIdPConfig = SPImpl . IdPConfigStore.getConfig
   getIdPConfigByIssuerOptionalSPId issuer mbteam = SPImpl $ case mbteam of
     Nothing -> IdPConfigStore.getIdPByIssuerV1 issuer

services/spar/test-integration/Test/Spar/DataSpec.hs

@@ -190,26 +190,26 @@ spec = do
     describe "IdPConfig" $ do
       it "insertIdPConfig, getIdPConfig are \"inverses\"" $ do
         idp <- makeTestIdP
-        () <- runSpar $ IdPEffect.insertConfig idp
+        () <- runSpar $ IdPEffect.insertConfig IdPEffect.InsertOnly idp
         midp <- runSpar $ IdPEffect.getConfig (idp ^. idpId)
         liftIO $ midp `shouldBe` idp
       it "getIdPByIssuer works" $ do
         idp <- makeTestIdP
-        () <- runSpar $ IdPEffect.insertConfig idp
+        () <- runSpar $ IdPEffect.insertConfig IdPEffect.InsertOnly idp
         midp <- getIdPByIssuer (idp ^. idpMetadata . edIssuer) (idp ^. SAML.idpExtraInfo . team)
         liftIO $ midp `shouldBe` Just idp
       it "getIdPConfigsByTeam works" $ do
         skipIdPAPIVersions [WireIdPAPIV1]
         teamid <- nextWireId
         idp <- makeTestIdP <&> idpExtraInfo .~ WireIdP teamid Nothing [] Nothing (IdPHandle "IdP 1") Nothing
-        () <- runSpar $ IdPEffect.insertConfig idp
+        () <- runSpar $ IdPEffect.insertConfig IdPEffect.InsertOnly idp
         idps <- runSpar $ IdPEffect.getConfigsByTeam teamid
         liftIO $ idps `shouldBe` [idp]
       it "deleteIdPConfig works" $ do
         teamid <- nextWireId
         idpApiVersion <- asks (^. teWireIdPAPIVersion)
         idp <- makeTestIdP <&> idpExtraInfo .~ WireIdP teamid (Just idpApiVersion) [] Nothing (IdPHandle "IdP 1") Nothing
-        () <- runSpar $ IdPEffect.insertConfig idp
+        () <- runSpar $ IdPEffect.insertConfig IdPEffect.InsertOnly idp
         do
           midp <- runSpar $ IdPEffect.getConfig (idp ^. idpId)
           liftIO $ midp `shouldBe` idp

services/spar/test/Test/Spar/Saml/IdPSpec.hs

@@ -665,7 +665,7 @@ spec =
                   let allow = singletonAllowlist bogusFingerprint
                   (_logs2, _notifs2, res) <-
                     interpretWithLoggingMockOptsE (withAllow (Just allow)) Nothing $ do
-                      insertConfig idp
+                      insertConfig InsertOnly idp
                       idpUpdate singleIngressSamlConfig zUser host idpInfo (idp._idpId) Nothing
                   res `shouldBe` Left (SAML.CustomError (SparIdPCertNotAllowed fingerprint))
 
@@ -681,7 +681,7 @@ spec =
                 Right idp -> do
                   (_logs2, _notifs2, res) <-
                     interpretWithLoggingMockOptsE (withAllow (Just partialAllow)) Nothing $ do
-                      insertConfig idp
+                      insertConfig InsertOnly idp
                       idpUpdate singleIngressSamlConfig zUser host idpInfo (idp._idpId) Nothing
                   res `shouldBe` Left (SAML.CustomError (SparIdPCertNotAllowed secondCertFingerprint))