notification server TRPL command (#420)

* notification server TRPL command

* test

* client methods

* only remove current token registration

src/Simplex/Messaging/Agent/Client.hs

@@ -32,6 +32,7 @@ module Simplex.Messaging.Agent.Client
     agentNtfRegisterToken,
     agentNtfVerifyToken,
     agentNtfCheckToken,
+    agentNtfReplaceToken,
     agentNtfDeleteToken,
     agentNtfEnableCron,
     agentNtfCreateSubscription,
@@ -570,6 +571,10 @@ agentNtfCheckToken :: AgentMonad m => AgentClient -> NtfTokenId -> NtfToken -> m
 agentNtfCheckToken c tknId NtfToken {ntfServer, ntfPrivKey} =
   withLogClient c ntfServer tknId "TCHK" $ \ntf -> ntfCheckToken ntf ntfPrivKey tknId
 
+agentNtfReplaceToken :: AgentMonad m => AgentClient -> NtfTokenId -> NtfToken -> DeviceToken -> m ()
+agentNtfReplaceToken c tknId NtfToken {ntfServer, ntfPrivKey} token =
+  withLogClient c ntfServer tknId "TRPL" $ \ntf -> ntfReplaceToken ntf ntfPrivKey tknId token
+
 agentNtfDeleteToken :: AgentMonad m => AgentClient -> NtfTokenId -> NtfToken -> m ()
 agentNtfDeleteToken c tknId NtfToken {ntfServer, ntfPrivKey} =
   withLogClient c ntfServer tknId "TDEL" $ \ntf -> ntfDeleteToken ntf ntfPrivKey tknId

src/Simplex/Messaging/Notifications/Client.hs

@@ -42,6 +42,9 @@ ntfCheckToken c pKey tknId =
     NRTkn stat -> pure stat
     _ -> throwE PCEUnexpectedResponse
 
+ntfReplaceToken :: NtfClient -> C.APrivateSignKey -> NtfTokenId -> DeviceToken -> ExceptT ProtocolClientError IO ()
+ntfReplaceToken c pKey tknId token = okNtfCommand (TRPL token) c pKey tknId
+
 ntfDeleteToken :: NtfClient -> C.APrivateSignKey -> NtfTokenId -> ExceptT ProtocolClientError IO ()
 ntfDeleteToken = okNtfCommand TDEL
 

src/Simplex/Messaging/Notifications/Protocol.hs

@@ -55,6 +55,7 @@ data NtfCommandTag (e :: NtfEntity) where
   TNEW_ :: NtfCommandTag 'Token
   TVFY_ :: NtfCommandTag 'Token
   TCHK_ :: NtfCommandTag 'Token
+  TRPL_ :: NtfCommandTag 'Token
   TDEL_ :: NtfCommandTag 'Token
   TCRN_ :: NtfCommandTag 'Token
   SNEW_ :: NtfCommandTag 'Subscription
@@ -71,6 +72,7 @@ instance NtfEntityI e => Encoding (NtfCommandTag e) where
     TNEW_ -> "TNEW"
     TVFY_ -> "TVFY"
     TCHK_ -> "TCHK"
+    TRPL_ -> "TRPL"
     TDEL_ -> "TDEL"
     TCRN_ -> "TCRN"
     SNEW_ -> "SNEW"
@@ -88,6 +90,7 @@ instance ProtocolMsgTag NtfCmdTag where
     "TNEW" -> Just $ NCT SToken TNEW_
     "TVFY" -> Just $ NCT SToken TVFY_
     "TCHK" -> Just $ NCT SToken TCHK_
+    "TRPL" -> Just $ NCT SToken TRPL_
     "TDEL" -> Just $ NCT SToken TDEL_
     "TCRN" -> Just $ NCT SToken TCRN_
     "SNEW" -> Just $ NCT SSubscription SNEW_
@@ -157,6 +160,8 @@ data NtfCommand (e :: NtfEntity) where
   TVFY :: NtfRegCode -> NtfCommand 'Token
   -- | check token status
   TCHK :: NtfCommand 'Token
+  -- | replace device token (while keeping all existing subscriptions)
+  TRPL :: DeviceToken -> NtfCommand 'Token
   -- | delete token - all subscriptions will be removed and no more notifications will be sent
   TDEL :: NtfCommand 'Token
   -- | enable periodic background notification to fetch the new messages - interval is in minutes, minimum is 20, 0 to disable
@@ -182,6 +187,7 @@ instance NtfEntityI e => ProtocolEncoding (NtfCommand e) where
     TNEW newTkn -> e (TNEW_, ' ', newTkn)
     TVFY code -> e (TVFY_, ' ', code)
     TCHK -> e TCHK_
+    TRPL tkn -> e (TRPL_, ' ', tkn)
     TDEL -> e TDEL_
     TCRN int -> e (TCRN_, ' ', int)
     SNEW newSub -> e (SNEW_, ' ', newSub)
@@ -221,6 +227,7 @@ instance ProtocolEncoding NtfCmd where
         TNEW_ -> TNEW <$> _smpP
         TVFY_ -> TVFY <$> _smpP
         TCHK_ -> pure TCHK
+        TRPL_ -> TRPL <$> _smpP
         TDEL_ -> pure TDEL
         TCRN_ -> TCRN <$> _smpP
     NCT SSubscription tag ->
@@ -244,7 +251,7 @@ data NtfResponseTag
 
 instance Encoding NtfResponseTag where
   smpEncode = \case
-    NRTknId_ -> "IDTKN"
+    NRTknId_ -> "IDTKN" -- it should be "TID", "SID"
     NRSubId_ -> "IDSUB"
     NROk_ -> "OK"
     NRErr_ -> "ERR"

src/Simplex/Messaging/Notifications/Server.hs

@@ -296,6 +296,16 @@ client NtfServerClient {rcvQ, sndQ} NtfSubscriber {newSubQ} NtfPushServer {pushQ
               logDebug "TVFY - incorrect code or token status"
               pure $ NRErr AUTH
           TCHK -> pure $ NRTkn status
+          TRPL token' -> do
+            logDebug "TRPL - replace token"
+            st <- asks store
+            regCode <- getRegCode
+            atomically $ do
+              removeTokenRegistration st tkn
+              writeTVar tknStatus NTRegistered
+              addNtfToken st tknId tkn {token = token', tknRegCode = regCode}
+              writeTBQueue pushQ (tkn, PNVerification regCode)
+            pure NROk
           TDEL -> do
             logDebug "TDEL"
             st <- asks store

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

@@ -112,6 +112,15 @@ removeInactiveTokenRegistrations st NtfTknData {ntfTknId = tId, token} =
         TM.delete tId' $ tokens st
       pure $ map snd tIds
 
+removeTokenRegistration :: NtfStore -> NtfTknData -> STM ()
+removeTokenRegistration st NtfTknData {ntfTknId = tId, token, tknVerifyKey} =
+  TM.lookup token (tokenRegistrations st) >>= mapM_ removeReg
+  where
+    removeReg regs =
+      TM.lookup k regs
+        >>= mapM_ (\tId' -> when (tId == tId') $ TM.delete k regs)
+    k = C.toPubKey C.pubKeyBytes tknVerifyKey
+
 deleteNtfToken :: NtfStore -> NtfTokenId -> STM ()
 deleteNtfToken st tknId = do
   TM.lookupDelete tknId (tokens st)

tests/NtfServerTests.hs

@@ -138,3 +138,27 @@ testNotificationSubscription (ATransport t) =
           (decryptedMsg, Right "hello") #== "delivered from queue"
           Resp "6" _ OK <- signSendRecv rh rKey ("6", rId, ACK mId1)
           pure ()
+          -- replace token
+          let tkn' = DeviceToken PPApns "efgh"
+          RespNtf "7" tId' NROk <- signSendRecvNtf nh tknKey ("7", tId, TRPL tkn')
+          tId `shouldBe` tId'
+          APNSMockRequest {notification = APNSNotification {aps = APNSBackground _, notificationData = Just ntfData2}, sendApnsResponse = send2} <-
+            atomically $ readTBQueue apnsQ
+          send2 APNSRespOk
+          let Right verification2 = ntfData2 .-> "verification"
+              Right nonce2 = C.cbNonce <$> ntfData2 .-> "nonce"
+              Right code2 = NtfRegCode <$> C.cbDecrypt dhSecret nonce2 verification2
+          RespNtf "8" _ NROk <- signSendRecvNtf nh tknKey ("8", tId, TVFY code2)
+          RespNtf "8a" _ (NRTkn NTActive) <- signSendRecvNtf nh tknKey ("8a", tId, TCHK)
+          -- send message
+          Resp "9" _ OK <- signSendRecv sh sKey ("9", sId, _SEND' "hello 2")
+          APNSMockRequest {notification = notification3, sendApnsResponse = send3} <- atomically $ readTBQueue apnsQ
+          let APNSNotification {aps = APNSMutableContent {}, notificationData = Just ntfData3} = notification3
+              Right nonce3 = C.cbNonce <$> ntfData3 .-> "nonce"
+              Right message3 = ntfData3 .-> "message"
+              Right ntfDataDecrypted3 = C.cbDecrypt dhSecret nonce3 message3
+              Right APNS.PNMessageData {smpQueue = SMPQueueNtf {smpServer = smpServer3, notifierId = notifierId3}} =
+                parse strP (AP.INTERNAL "error parsing PNMessageData") ntfDataDecrypted3
+          smpServer3 `shouldBe` srv
+          notifierId3 `shouldBe` nId
+          send3 APNSRespOk