WIP

package.yaml

@@ -71,6 +71,7 @@ dependencies:
   - transformers == 0.6.*
   - unliftio == 0.2.*
   - unliftio-core == 0.2.*
+  - unordered-containers
   - websockets == 0.12.*
   - yaml == 0.11.*
   - zstd == 0.1.3.*

simplexmq.cabal

@@ -230,6 +230,7 @@ library
     , transformers ==0.6.*
     , unliftio ==0.2.*
     , unliftio-core ==0.2.*
+    , unordered-containers
     , websockets ==0.12.*
     , yaml ==0.11.*
     , zstd ==0.1.3.*
@@ -305,6 +306,7 @@ executable ntf-server
     , transformers ==0.6.*
     , unliftio ==0.2.*
     , unliftio-core ==0.2.*
+    , unordered-containers
     , websockets ==0.12.*
     , yaml ==0.11.*
     , zstd ==0.1.3.*
@@ -380,6 +382,7 @@ executable smp-agent
     , transformers ==0.6.*
     , unliftio ==0.2.*
     , unliftio-core ==0.2.*
+    , unordered-containers
     , websockets ==0.12.*
     , yaml ==0.11.*
     , zstd ==0.1.3.*
@@ -455,6 +458,7 @@ executable smp-server
     , transformers ==0.6.*
     , unliftio ==0.2.*
     , unliftio-core ==0.2.*
+    , unordered-containers
     , websockets ==0.12.*
     , yaml ==0.11.*
     , zstd ==0.1.3.*
@@ -530,6 +534,7 @@ executable xftp
     , transformers ==0.6.*
     , unliftio ==0.2.*
     , unliftio-core ==0.2.*
+    , unordered-containers
     , websockets ==0.12.*
     , yaml ==0.11.*
     , zstd ==0.1.3.*
@@ -605,6 +610,7 @@ executable xftp-server
     , transformers ==0.6.*
     , unliftio ==0.2.*
     , unliftio-core ==0.2.*
+    , unordered-containers
     , websockets ==0.12.*
     , yaml ==0.11.*
     , zstd ==0.1.3.*
@@ -721,6 +727,7 @@ test-suite simplexmq-test
     , transformers ==0.6.*
     , unliftio ==0.2.*
     , unliftio-core ==0.2.*
+    , unordered-containers
     , websockets ==0.12.*
     , yaml ==0.11.*
     , zstd ==0.1.3.*

src/Simplex/Messaging/Server.hs

@@ -51,6 +51,7 @@ import qualified Data.ByteString.Char8 as B
 import qualified Data.ByteString.Lazy.Char8 as LB
 import Data.Either (fromRight, partitionEithers)
 import Data.Functor (($>))
+import qualified Data.HashMap.Strict as HM
 import Data.Int (Int64)
 import qualified Data.IntMap.Strict as IM
 import Data.List (intercalate, mapAccumR)
@@ -541,7 +542,11 @@ verifyTransmission :: Maybe (THandleAuth 'TServer, C.CbNonce) -> Maybe Transmiss
 verifyTransmission auth_ tAuth authorized queueId cmd =
   case cmd of
     Cmd SRecipient (NEW k _ _ _) -> pure $ Nothing `verifiedWith` k
-    Cmd SRecipient _ -> verifyQueue (\q -> Just q `verifiedWith` recipientKey q) <$> get SRecipient
+    Cmd SRecipient _ -> do
+      QueueStore {recipientKeys} <- asks queueStore
+      (queueId `HM.lookup`) <$> readTVarIO recipientKeys >>= \case
+        Nothing -> pure $! dummyVerify
+        Just rKey -> if verify rKey then slowIO >> get SRecipient >>= either (\_notGonnaHappenUnlessDeletedWhileCheckingAuth -> pure VRFailed) (pure . VRVerified . Just) else pure VRFailed
     -- SEND will be accepted without authorization before the queue is secured with KEY command
     Cmd SSender SEND {} -> verifyQueue (\q -> Just q `verified` maybe (isNothing tAuth) verify (senderKey q)) <$> get SSender
     Cmd SSender PING -> pure $ VRVerified Nothing
@@ -554,6 +559,8 @@ verifyTransmission auth_ tAuth authorized queueId cmd =
     verifyQueue = either (\_ -> dummyVerify)
     verified q cond = if cond then VRVerified q else VRFailed
     verifiedWith q k = q `verified` verify k
+    slowIO :: M ()
+    slowIO = liftIO $ threadDelay 100000 -- XXX: Some thumb twiddling to prove that fetching queue data doesn't allow to distinguish bad key and non-existent queue.
     get :: SParty p -> M (Either ErrorType QueueRec)
     get party = do
       st <- asks queueStore

src/Simplex/Messaging/Server/QueueStore.hs

@@ -8,6 +8,9 @@ module Simplex.Messaging.Server.QueueStore where
 import Simplex.Messaging.Encoding.String
 import Simplex.Messaging.Protocol
 
+-- normalized/storage form
+-- can be split into per-party pieces
+-- id/public key form party index, dhsecrets can use dense storage or loaded on demand
 data QueueRec = QueueRec
   { recipientId :: !RecipientId,
     recipientKey :: !RcvPublicAuthKey,

src/Simplex/Messaging/Server/QueueStore/STM.hs

@@ -24,6 +24,8 @@ where
 
 import Control.Monad
 import Data.Functor (($>))
+import Data.HashMap.Strict (HashMap)
+import qualified Data.HashMap.Strict as HM
 import Simplex.Messaging.Protocol
 import Simplex.Messaging.Server.QueueStore
 import Simplex.Messaging.TMap (TMap)
@@ -32,24 +34,32 @@ import Simplex.Messaging.Util (ifM, ($>>=))
 import UnliftIO.STM
 
 data QueueStore = QueueStore
-  { queues :: TMap RecipientId (TVar QueueRec),
-    senders :: TMap SenderId RecipientId,
-    notifiers :: TMap NotifierId RecipientId
+  { recipientKeys :: TVar (HashMap RecipientId RcvPublicAuthKey), -- O(1) index for const-time auth checks
+    senderKeys :: TVar (HashMap SenderId SndPublicAuthKey), -- O(1) index for const-time auth checks
+    notifierKeys :: TVar (HashMap NotifierId NtfPublicAuthKey), -- O(1) index for const-time auth checks
+    queues :: TMap RecipientId (TVar QueueRec), -- can be stored cold; and indexed with ints (using annotated keys below)
+    senders :: TMap SenderId RecipientId, -- not needed with annotated keys
+    notifiers :: TMap NotifierId RecipientId -- not needed with annotated keys
   }
 
 newQueueStore :: STM QueueStore
 newQueueStore = do
+  recipientKeys <- newTVar mempty
+  senderKeys <- newTVar mempty
+  notifierKeys <- newTVar mempty
   queues <- TM.empty
   senders <- TM.empty
   notifiers <- TM.empty
-  pure QueueStore {queues, senders, notifiers}
+  pure QueueStore {recipientKeys, senderKeys, notifierKeys, queues, senders, notifiers}
 
 addQueue :: QueueStore -> QueueRec -> STM (Either ErrorType ())
-addQueue QueueStore {queues, senders} q@QueueRec {recipientId = rId, senderId = sId} = do
+addQueue QueueStore {recipientKeys, senderKeys, queues, senders} q@QueueRec {recipientId = rId, recipientKey, senderId = sId, senderKey} = do
   ifM hasId (pure $ Left DUPLICATE_) $ do
     qVar <- newTVar q
     TM.insert rId qVar queues
+    modifyTVar' recipientKeys $ HM.insert rId recipientKey
     TM.insert sId rId senders
+    forM_ senderKey $ modifyTVar' senderKeys . HM.insert rId -- XXX: should not exist here yet, unless testing?
     pure $ Right ()
   where
     hasId = (||) <$> TM.member rId queues <*> TM.member sId senders
@@ -64,13 +74,14 @@ getQueue QueueStore {queues, senders, notifiers} party qId =
       SNotifier -> TM.lookup qId notifiers $>>= (`TM.lookup` queues)
 
 secureQueue :: QueueStore -> RecipientId -> SndPublicAuthKey -> STM (Either ErrorType QueueRec)
-secureQueue QueueStore {queues} rId sKey =
+secureQueue QueueStore {senderKeys, queues} rId sKey =
   withQueue rId queues $ \qVar ->
     readTVar qVar >>= \q -> case senderKey q of
       Just k -> pure $ if sKey == k then Just q else Nothing
-      _ ->
+      _ -> do
         let q' = q {senderKey = Just sKey}
-         in writeTVar qVar q' $> Just q'
+        modifyTVar' senderKeys $ HM.insert rId sKey
+        writeTVar qVar q' $> Just q'
 
 addQueueNotifier :: QueueStore -> RecipientId -> NtfCreds -> STM (Either ErrorType QueueRec)
 addQueueNotifier QueueStore {queues, notifiers} rId ntfCreds@NtfCreds {notifierId = nId} = do

tests/ServerTests.hs

@@ -62,7 +62,7 @@ serverTests t@(ATransport t') = do
   describe "Store log" $ testWithStoreLog t
   describe "Restore messages" $ testRestoreMessages t
   describe "Restore messages (old / v2)" $ testRestoreExpireMessages t
-  describe "Timing of AUTH error" $ testTiming t
+  fdescribe "Timing of AUTH error" $ testTiming t
   describe "Message notifications" $ testMessageNotifications t
   describe "Message expiration" $ do
     testMsgExpireOnSend t'