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/Protocol.hs

@@ -264,6 +264,7 @@ data NtfResponse
   | NRTkn NtfTknStatus
   | NRSub NtfSubStatus
   | NRPong
+  deriving (Show)
 
 instance ProtocolEncoding NtfResponse where
   type Tag NtfResponse = NtfResponseTag
@@ -361,7 +362,7 @@ data NtfSubStatus
     NSEnd
   | -- | SMP AUTH error
     NSSMPAuth
-  deriving (Eq)
+  deriving (Eq, Show)
 
 instance Encoding NtfSubStatus where
   smpEncode = \case

src/Simplex/Messaging/Notifications/Server.hs

@@ -117,7 +117,7 @@ ntfPush s@NtfPushServer {pushQ} = liftIO . forever . runExceptT $ do
     (_, PNVerification _) -> do
       -- TODO check token status
       deliverNotification pp tkn ntf
-      atomically $ writeTVar tknStatus NTConfirmed
+      atomically $ modifyTVar tknStatus $ \status' -> if status' == NTActive then NTActive else NTConfirmed
     (NTActive, PNCheckMessages) -> do
       deliverNotification pp tkn ntf
     _ -> do
@@ -166,26 +166,26 @@ verifyNtfTransmission ::
 verifyNtfTransmission (sig_, signed, (corrId, entId, _)) cmd = do
   st <- asks store
   case cmd of
-    NtfCmd SToken c@(TNEW n@(NewNtfTkn _ k _)) -> do
-      r_ <- atomically $ getNtfToken st entId
+    NtfCmd SToken c@(TNEW tkn@(NewNtfTkn _ k _)) -> do
+      r_ <- atomically $ getNtfTokenRegistration st tkn
       pure $
         if verifyCmdSignature sig_ signed k
           then case r_ of
-            Just r@(NtfTkn NtfTknData {tknVerifyKey})
-              | k == tknVerifyKey -> tknCmd r c
+            Just t@NtfTknData {tknVerifyKey}
+              | k == tknVerifyKey -> verifiedTknCmd t c
               | otherwise -> VRFailed
-            _ -> VRVerified (NtfReqNew corrId (ANE SToken n))
+            _ -> VRVerified (NtfReqNew corrId (ANE SToken tkn))
           else VRFailed
     NtfCmd SToken c -> do
-      r_ <- atomically $ getNtfToken st entId
-      pure $ case r_ of
-        Just r@(NtfTkn NtfTknData {tknVerifyKey})
-          | verifyCmdSignature sig_ signed tknVerifyKey -> tknCmd r c
+      t_ <- atomically $ getNtfToken st entId
+      pure $ case t_ of
+        Just t@NtfTknData {tknVerifyKey}
+          | verifyCmdSignature sig_ signed tknVerifyKey -> verifiedTknCmd t c
           | otherwise -> VRFailed
         _ -> maybe False (dummyVerifyCmd signed) sig_ `seq` VRFailed
     _ -> pure VRFailed
   where
-    tknCmd r c = VRVerified (NtfReqCmd SToken r (corrId, entId, c))
+    verifiedTknCmd t c = VRVerified (NtfReqCmd SToken (NtfTkn t) (corrId, entId, c))
 
 client :: forall m. (MonadUnliftIO m, MonadReader NtfEnv m) => NtfServerClient -> NtfSubscriber -> NtfPushServer -> m ()
 client NtfServerClient {rcvQ, sndQ} NtfSubscriber {subQ = _} NtfPushServer {pushQ, intervalNotifiers} =
@@ -204,11 +204,11 @@ client NtfServerClient {rcvQ, sndQ} NtfSubscriber {subQ = _} NtfPushServer {push
         tknId <- getId
         regCode <- getRegCode
         atomically $ do
-          tkn <- mkNtfTknData newTkn ks dhSecret regCode
+          tkn <- mkNtfTknData tknId newTkn ks dhSecret regCode
           addNtfToken st tknId tkn
           writeTBQueue pushQ (tkn, PNVerification regCode)
         pure (corrId, "", NRId tknId srvDhPubKey)
-      NtfReqCmd SToken (NtfTkn tkn@NtfTknData {tknStatus, tknRegCode, tknDhSecret, tknDhKeys = (srvDhPubKey, srvDhPrivKey)}) (corrId, tknId, cmd) -> do
+      NtfReqCmd SToken (NtfTkn tkn@NtfTknData {ntfTknId, tknStatus, tknRegCode, tknDhSecret, tknDhKeys = (srvDhPubKey, srvDhPrivKey)}) (corrId, tknId, cmd) -> do
         status <- readTVarIO tknStatus
         (corrId,tknId,) <$> case cmd of
           TNEW (NewNtfTkn _ _ dhPubKey) -> do
@@ -218,12 +218,15 @@ client NtfServerClient {rcvQ, sndQ} NtfSubscriber {subQ = _} NtfPushServer {push
             if tknDhSecret == dhSecret
               then do
                 atomically $ writeTBQueue pushQ (tkn, PNVerification tknRegCode)
-                pure $ NRId tknId srvDhPubKey
+                pure $ NRId ntfTknId srvDhPubKey
               else pure $ NRErr AUTH
           TVFY code -- this allows repeated verification for cases when client connection dropped before server response
             | (status == NTRegistered || status == NTConfirmed || status == NTActive) && tknRegCode == code -> do
               logDebug "TVFY - token verified"
+              st <- asks store
               atomically $ writeTVar tknStatus NTActive
+              tIds <- atomically $ removeInactiveTokenRegistrations st tkn
+              forM_ tIds cancelInvervalNotifications
               pure NROk
             | otherwise -> do
               logDebug "TVFY - incorrect code or token status"
@@ -233,12 +236,12 @@ client NtfServerClient {rcvQ, sndQ} NtfSubscriber {subQ = _} NtfPushServer {push
             logDebug "TDEL"
             st <- asks store
             atomically $ deleteNtfToken st tknId
+            cancelInvervalNotifications tknId
             pure NROk
-          TCRN 0 ->
+          TCRN 0 -> do
             logDebug "TCRN 0"
-              >> atomically (TM.lookupDelete tknId intervalNotifiers)
-              >>= mapM_ (uninterruptibleCancel . action)
-              >> pure NROk
+            cancelInvervalNotifications tknId
+            pure NROk
           TCRN int
             | int < 20 -> pure $ NRErr QUOTA
             | otherwise -> do
@@ -274,6 +277,10 @@ client NtfServerClient {rcvQ, sndQ} NtfSubscriber {subQ = _} NtfPushServer {push
     getRandomBytes n = do
       gVar <- asks idsDrg
       atomically (C.pseudoRandomBytes n gVar)
+    cancelInvervalNotifications :: NtfTokenId -> m ()
+    cancelInvervalNotifications tknId =
+      atomically (TM.lookupDelete tknId intervalNotifiers)
+        >>= mapM_ (uninterruptibleCancel . action)
 
 -- NReqCreate corrId tokenId smpQueue -> pure (corrId, "", NROk)
 -- do

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