From facb6f4561bc60e6e141f72877e0967e00a53cdb Mon Sep 17 00:00:00 2001
From: JRoberts <8711996+jr-simplex@users.noreply.github.com>
Date: Thu, 2 Jun 2022 15:36:16 +0400
Subject: [PATCH] ntf: separate notifierKey from SMPQueueNtf (#382)

---
 src/Simplex/Messaging/Notifications/Protocol.hs | 17 ++++++++---------
 src/Simplex/Messaging/Notifications/Server.hs   |  6 +++---
 .../Messaging/Notifications/Server/Store.hs     | 17 +++++++++--------
 3 files changed, 20 insertions(+), 20 deletions(-)

diff --git a/src/Simplex/Messaging/Notifications/Protocol.hs b/src/Simplex/Messaging/Notifications/Protocol.hs
index 3640004e0..8830ff457 100644
--- a/src/Simplex/Messaging/Notifications/Protocol.hs
+++ b/src/Simplex/Messaging/Notifications/Protocol.hs
@@ -120,7 +120,7 @@ instance ToJSON NtfRegCode where
 
 data NewNtfEntity (e :: NtfEntity) where
   NewNtfTkn :: DeviceToken -> C.APublicVerifyKey -> C.PublicKeyX25519 -> NewNtfEntity 'Token
-  NewNtfSub :: NtfTokenId -> SMPQueueNtf -> NewNtfEntity 'Subscription
+  NewNtfSub :: NtfTokenId -> SMPQueueNtf -> NtfPrivateSignKey -> NewNtfEntity 'Subscription
 
 deriving instance Show (NewNtfEntity e)
 
@@ -129,7 +129,7 @@ data ANewNtfEntity = forall e. NtfEntityI e => ANE (SNtfEntity e) (NewNtfEntity
 instance NtfEntityI e => Encoding (NewNtfEntity e) where
   smpEncode = \case
     NewNtfTkn tkn verifyKey dhPubKey -> smpEncode ('T', tkn, verifyKey, dhPubKey)
-    NewNtfSub tknId smpQueue -> smpEncode ('S', tknId, smpQueue)
+    NewNtfSub tknId smpQueue notifierKey -> smpEncode ('S', tknId, smpQueue, notifierKey)
   smpP = (\(ANE _ c) -> checkEntity c) <$?> smpP
 
 instance Encoding ANewNtfEntity where
@@ -137,7 +137,7 @@ instance Encoding ANewNtfEntity where
   smpP =
     A.anyChar >>= \case
       'T' -> ANE SToken <$> (NewNtfTkn <$> smpP <*> smpP <*> smpP)
-      'S' -> ANE SSubscription <$> (NewNtfSub <$> smpP <*> smpP)
+      'S' -> ANE SSubscription <$> (NewNtfSub <$> smpP <*> smpP <*> smpP)
       _ -> fail "bad ANewNtfEntity"
 
 instance Protocol NtfResponse where
@@ -307,16 +307,15 @@ instance ProtocolEncoding NtfResponse where
 
 data SMPQueueNtf = SMPQueueNtf
   { smpServer :: ProtocolServer,
-    notifierId :: NotifierId,
-    notifierKey :: NtfPrivateSignKey
+    notifierId :: NotifierId
   }
-  deriving (Show)
+  deriving (Eq, Ord, Show)
 
 instance Encoding SMPQueueNtf where
-  smpEncode SMPQueueNtf {smpServer, notifierId, notifierKey} = smpEncode (smpServer, notifierId, notifierKey)
+  smpEncode SMPQueueNtf {smpServer, notifierId} = smpEncode (smpServer, notifierId)
   smpP = do
-    (smpServer, notifierId, notifierKey) <- smpP
-    pure $ SMPQueueNtf smpServer notifierId notifierKey
+    (smpServer, notifierId) <- smpP
+    pure $ SMPQueueNtf smpServer notifierId
 
 data PushProvider = PPApns
   deriving (Eq, Ord, Show)
diff --git a/src/Simplex/Messaging/Notifications/Server.hs b/src/Simplex/Messaging/Notifications/Server.hs
index 03b291f17..18448c19e 100644
--- a/src/Simplex/Messaging/Notifications/Server.hs
+++ b/src/Simplex/Messaging/Notifications/Server.hs
@@ -94,8 +94,8 @@ ntfSubscriber NtfSubscriber {smpSubscribers, newSubQ, smpAgent = ca@SMPClientAge
     runSMPSubscriber SMPSubscriber {newSubQ = subscriberSubQ} =
       forever $
         atomically (peekTBQueue subscriberSubQ) >>= \case
-          NtfSub NtfSubData {smpQueue} -> do
-            let SMPQueueNtf {smpServer, notifierId, notifierKey} = smpQueue
+          NtfSub NtfSubData {smpQueue, notifierKey} -> do
+            let SMPQueueNtf {smpServer, notifierId} = smpQueue
             liftIO (runExceptT $ subscribeQueue ca smpServer ((SPNotifier, notifierId), notifierKey)) >>= \case
               Right _ -> do
                 -- update subscription status
@@ -215,7 +215,7 @@ verifyNtfTransmission (sig_, signed, (corrId, entId, _)) cmd = do
           | verifyCmdSignature sig_ signed tknVerifyKey -> verifiedTknCmd t c
           | otherwise -> VRFailed
         _ -> maybe False (dummyVerifyCmd signed) sig_ `seq` VRFailed
-    NtfCmd SSubscription c@(SNEW sub@(NewNtfSub tknId _)) -> do
+    NtfCmd SSubscription c@(SNEW sub@(NewNtfSub tknId _ _)) -> do
       -- TODO move active token check here to differentiate error
       r_ <- atomically $ findNtfSubscription st sub
       pure $ case r_ of
diff --git a/src/Simplex/Messaging/Notifications/Server/Store.hs b/src/Simplex/Messaging/Notifications/Server/Store.hs
index c15428f63..37a0b4f37 100644
--- a/src/Simplex/Messaging/Notifications/Server/Store.hs
+++ b/src/Simplex/Messaging/Notifications/Server/Store.hs
@@ -17,7 +17,7 @@ import qualified Data.Map.Strict as M
 import Data.Set (Set, insert)
 import qualified Simplex.Messaging.Crypto as C
 import Simplex.Messaging.Notifications.Protocol
-import Simplex.Messaging.Protocol (NotifierId, ProtocolServer)
+import Simplex.Messaging.Protocol (NtfPrivateSignKey)
 import Simplex.Messaging.TMap (TMap)
 import qualified Simplex.Messaging.TMap as TM
 import Simplex.Messaging.Util (whenM, ($>>=), (<$$>))
@@ -27,7 +27,7 @@ data NtfStore = NtfStore
     tokenRegistrations :: TMap DeviceToken (TMap ByteString NtfTokenId),
     subscriptions :: TMap NtfSubscriptionId NtfSubData,
     tokenSubscriptions :: TMap NtfTokenId (TVar (Set NtfSubscriptionId)),
-    subscriptionLookup :: TMap (NtfTokenId, ProtocolServer, NotifierId) NtfSubscriptionId
+    subscriptionLookup :: TMap (NtfTokenId, SMPQueueNtf) NtfSubscriptionId
   }
 
 newNtfStore :: STM NtfStore
@@ -66,6 +66,7 @@ mkNtfTknData ntfTknId (NewNtfTkn token tknVerifyKey _) tknDhKeys tknDhSecret tkn
 
 data NtfSubData = NtfSubData
   { smpQueue :: SMPQueueNtf,
+    notifierKey :: NtfPrivateSignKey,
     tokenId :: NtfTokenId,
     subStatus :: TVar NtfSubStatus
   }
@@ -133,12 +134,12 @@ getNtfSubscription st subId =
       (sub,) <$$> getActiveNtfToken st tokenId
 
 findNtfSubscription :: NtfStore -> NewNtfEntity 'Subscription -> STM (Maybe (NtfTknData, Maybe NtfSubData))
-findNtfSubscription st (NewNtfSub tknId SMPQueueNtf {smpServer, notifierId}) = do
+findNtfSubscription st (NewNtfSub tknId smpQueue _) = do
   getActiveNtfToken st tknId >>= mapM (\tkn -> (tkn,) <$> getSub)
   where
     getSub :: STM (Maybe NtfSubData)
     getSub =
-      TM.lookup (tknId, smpServer, notifierId) (subscriptionLookup st)
+      TM.lookup (tknId, smpQueue) (subscriptionLookup st)
         $>>= (`TM.lookup` subscriptions st)
 
 getActiveNtfToken :: NtfStore -> NtfTokenId -> STM (Maybe NtfTknData)
@@ -148,18 +149,18 @@ getActiveNtfToken st tknId =
     pure $ if tStatus == NTActive then Just tkn else Nothing
 
 mkNtfSubData :: NewNtfEntity 'Subscription -> STM NtfSubData
-mkNtfSubData (NewNtfSub tokenId smpQueue) = do
+mkNtfSubData (NewNtfSub tokenId smpQueue notifierKey) = do
   subStatus <- newTVar NSNew
-  pure NtfSubData {smpQueue, tokenId, subStatus}
+  pure NtfSubData {smpQueue, tokenId, subStatus, notifierKey}
 
 addNtfSubscription :: NtfStore -> NtfSubscriptionId -> NtfSubData -> STM (Maybe ())
-addNtfSubscription st subId sub@NtfSubData {smpQueue = SMPQueueNtf {smpServer, notifierId}, tokenId} =
+addNtfSubscription st subId sub@NtfSubData {smpQueue, tokenId} =
   TM.lookup tokenId (tokenSubscriptions st) >>= mapM insertSub
   where
     insertSub ts = do
       modifyTVar' ts $ insert subId
       TM.insert subId sub $ subscriptions st
-      TM.insert (tokenId, smpServer, notifierId) subId (subscriptionLookup st)
+      TM.insert (tokenId, smpQueue) subId (subscriptionLookup st)
 
 -- getNtfRec :: NtfStore -> SNtfEntity e -> NtfEntityId -> STM (Maybe (NtfEntityRec e))
 -- getNtfRec st ent entId = case ent of