primary tokens (#356)

* primary tokens

* support repeat token registration (TODO fix tests)

* fix notifications tests

* fix/test repeat/new registrations of the same token

* re-register token when subsequent ntf command fails with AUTH error (e.g. when server is re-started)

* cancel periodic notifications when token is deleted on the server

* debug failing test on CI

* fix notification test in CI

* debug CI test

* add delay in notificaitons test after server restart

src/Simplex/Messaging/Notifications/Server/Store.hs

@@ -2,30 +2,38 @@
 {-# LANGUAGE DuplicateRecordFields #-}
 {-# LANGUAGE GADTs #-}
 {-# LANGUAGE KindSignatures #-}
+{-# LANGUAGE LambdaCase #-}
 {-# LANGUAGE NamedFieldPuns #-}
+{-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+{-# LANGUAGE TupleSections #-}
 
 module Simplex.Messaging.Notifications.Server.Store where
 
 import Control.Concurrent.STM
+import Control.Monad
+import Data.ByteString.Char8 (ByteString)
+import qualified Data.Map.Strict as M
 import qualified Simplex.Messaging.Crypto as C
 import Simplex.Messaging.Notifications.Protocol
 import Simplex.Messaging.TMap (TMap)
 import qualified Simplex.Messaging.TMap as TM
-import Simplex.Messaging.Util ((<$$>))
+import Simplex.Messaging.Util (whenM, ($>>=))
 
 data NtfStore = NtfStore
   { tokens :: TMap NtfTokenId NtfTknData,
-    tokenIds :: TMap DeviceToken NtfTokenId
+    tokenRegistrations :: TMap DeviceToken (TMap ByteString NtfTokenId)
   }
 
 newNtfStore :: STM NtfStore
 newNtfStore = do
   tokens <- TM.empty
-  tokenIds <- TM.empty
-  pure NtfStore {tokens, tokenIds}
+  tokenRegistrations <- TM.empty
+  pure NtfStore {tokens, tokenRegistrations}
 
 data NtfTknData = NtfTknData
-  { token :: DeviceToken,
+  { ntfTknId :: NtfTokenId,
+    token :: DeviceToken,
     tknStatus :: TVar NtfTknStatus,
     tknVerifyKey :: C.APublicVerifyKey,
     tknDhKeys :: C.KeyPair 'C.X25519,
@@ -33,10 +41,10 @@ data NtfTknData = NtfTknData
     tknRegCode :: NtfRegCode
   }
 
-mkNtfTknData :: NewNtfEntity 'Token -> C.KeyPair 'C.X25519 -> C.DhSecretX25519 -> NtfRegCode -> STM NtfTknData
-mkNtfTknData (NewNtfTkn token tknVerifyKey _) tknDhKeys tknDhSecret tknRegCode = do
+mkNtfTknData :: NtfTokenId -> NewNtfEntity 'Token -> C.KeyPair 'C.X25519 -> C.DhSecretX25519 -> NtfRegCode -> STM NtfTknData
+mkNtfTknData ntfTknId (NewNtfTkn token tknVerifyKey _) tknDhKeys tknDhSecret tknRegCode = do
   tknStatus <- newTVar NTRegistered
-  pure NtfTknData {token, tknStatus, tknVerifyKey, tknDhKeys, tknDhSecret, tknRegCode}
+  pure NtfTknData {ntfTknId, token, tknStatus, tknVerifyKey, tknDhKeys, tknDhSecret, tknRegCode}
 
 -- data NtfSubscriptionsStore = NtfSubscriptionsStore
 
@@ -58,19 +66,57 @@ data NtfEntityRec (e :: NtfEntity) where
   NtfTkn :: NtfTknData -> NtfEntityRec 'Token
   NtfSub :: NtfSubData -> NtfEntityRec 'Subscription
 
-data ANtfEntityRec = forall e. NtfEntityI e => NER (SNtfEntity e) (NtfEntityRec e)
-
-getNtfToken :: NtfStore -> NtfTokenId -> STM (Maybe (NtfEntityRec 'Token))
-getNtfToken st tknId = NtfTkn <$$> TM.lookup tknId (tokens st)
+getNtfToken :: NtfStore -> NtfTokenId -> STM (Maybe NtfTknData)
+getNtfToken st tknId = TM.lookup tknId (tokens st)
 
 addNtfToken :: NtfStore -> NtfTokenId -> NtfTknData -> STM ()
-addNtfToken st tknId tkn@NtfTknData {token} = do
+addNtfToken st tknId tkn@NtfTknData {token, tknVerifyKey} = do
   TM.insert tknId tkn $ tokens st
-  TM.insert token tknId $ tokenIds st
+  TM.lookup token regs >>= \case
+    Just tIds -> TM.insert regKey tknId tIds
+    _ -> do
+      tIds <- TM.singleton regKey tknId
+      TM.insert token tIds regs
+  where
+    regs = tokenRegistrations st
+    regKey = C.toPubKey C.pubKeyBytes tknVerifyKey
+
+getNtfTokenRegistration :: NtfStore -> NewNtfEntity 'Token -> STM (Maybe NtfTknData)
+getNtfTokenRegistration st (NewNtfTkn token tknVerifyKey _) =
+  TM.lookup token (tokenRegistrations st)
+    $>>= TM.lookup regKey
+    $>>= (`TM.lookup` tokens st)
+  where
+    regKey = C.toPubKey C.pubKeyBytes tknVerifyKey
+
+removeInactiveTokenRegistrations :: NtfStore -> NtfTknData -> STM [NtfTokenId]
+removeInactiveTokenRegistrations st NtfTknData {ntfTknId = tId, token} =
+  TM.lookup token (tokenRegistrations st)
+    >>= maybe (pure []) removeRegs
+  where
+    removeRegs :: TMap ByteString NtfTokenId -> STM [NtfTokenId]
+    removeRegs tknRegs = do
+      tIds <- filter ((/= tId) . snd) . M.assocs <$> readTVar tknRegs
+      forM_ tIds $ \(regKey, tId') -> do
+        TM.delete regKey tknRegs
+        TM.delete tId' $ tokens st
+      pure $ map snd tIds
 
 deleteNtfToken :: NtfStore -> NtfTokenId -> STM ()
 deleteNtfToken st tknId = do
-  TM.lookupDelete tknId (tokens st) >>= mapM_ (\NtfTknData {token} -> TM.delete token $ tokenIds st)
+  TM.lookupDelete tknId (tokens st)
+    >>= mapM_
+      ( \NtfTknData {token, tknVerifyKey} ->
+          TM.lookup token regs
+            >>= mapM_
+              ( \tIds -> do
+                  TM.delete (regKey tknVerifyKey) tIds
+                  whenM (TM.null tIds) $ TM.delete token regs
+              )
+      )
+  where
+    regs = tokenRegistrations st
+    regKey = C.toPubKey C.pubKeyBytes
 
 -- getNtfRec :: NtfStore -> SNtfEntity e -> NtfEntityId -> STM (Maybe (NtfEntityRec e))
 -- getNtfRec st ent entId = case ent of