1 files changed, 57 insertions, 4 deletions

diff --git a/src/Network/Tox/Crypto/Handlers.hs b/src/Network/Tox/Crypto/Handlers.hs
index ac3d1ef0..0cb2d4db 100644
--- a/src/Network/Tox/Crypto/Handlers.hs
+++ b/src/Network/Tox/Crypto/Handlers.hs
@@ -19,6 +19,7 @@ import Data.Serialize             as S
 import Data.Word
 import GHC.Conc (unsafeIOToSTM)
 import qualified Data.Set         as Set
+import qualified Data.Array.Unboxed as A
 
 -- util, todo: move to another module
 maybeToEither :: Maybe b -> Either String b
@@ -31,7 +32,7 @@ data NetCryptoSessionStatus = Unaccepted | Accepted | Confirmed
 
 type IOHook addr x = addr -> x -> IO (Maybe (x -> x))
 type NetCryptoHook = IOHook NetCryptoSession CryptoData
-
+type MsgTypeArray = A.UArray Word8 Word16
 
 data NetCryptoSession = NCrypto { ncState         :: TVar NetCryptoSessionStatus
                                 , ncTheirBaseNonce :: TVar Nonce24 -- base nonce + packet number
@@ -44,6 +45,12 @@ data NetCryptoSession = NCrypto { ncState         :: TVar NetCryptoSessionStatus
                                 , ncSockAddr      :: SockAddr
                                 , ncHooks         :: TVar (Map.Map MessageType [NetCryptoHook])
                                 , ncUnrecognizedHook :: TVar (MessageType -> NetCryptoHook)
+                                , ncMessageTypes :: TVar MsgTypeArray -- ^ supported messages, 0 for unsupported,
+                                                                      --   otherwise the messageType, some message types
+                                                                      --   may not be in ncHooks yet, but they should appear
+                                                                      --   here if ncUnrecognizedHook will add them to ncHooks
+                                                                      --   on an as-need basis. On 0 entries, this Transport
+                                                                      --   will return id in case they are handled by another Transport.
                                 , ncAllSessions :: NetCryptoSessions -- ^ may be needed if one net-crypto session
                                                                      --   needs to possibly start another, as is
                                                                      --   the case in group chats
@@ -54,6 +61,7 @@ data NetCryptoSessions = NCSessions { netCryptoSessions :: TVar (Map.Map SockAdd
                                     , transportCrypto :: TransportCrypto
                                     , defaultHooks :: Map.Map MessageType [NetCryptoHook]
                                     , defaultUnrecognizedHook :: MessageType -> NetCryptoHook
+                                    , msgTypeArray :: MsgTypeArray
                                     }
 
 newSessionsState :: TransportCrypto -> (MessageType -> NetCryptoHook) -> Map.Map MessageType [NetCryptoHook] -> IO NetCryptoSessions
@@ -63,6 +71,7 @@ newSessionsState crypto unrechook hooks = do
                       , transportCrypto = crypto
                       , defaultHooks = hooks
                       , defaultUnrecognizedHook = unrechook
+                      , msgTypeArray = allMsgTypes -- todo make this a parameter
                       }
 
 data HandshakeParams
@@ -109,6 +118,7 @@ freshCryptoSession  sessions
     ncHooks0 <- atomically $ newTVar (defaultHooks sessions)
     ncUnrecognizedHook0 <- atomically $ newTVar (defaultUnrecognizedHook sessions)
     ncGroups0 <- atomically $ newTVar (Map.empty)
+    ncMessageTypes0 <- atomically $ newTVar (msgTypeArray sessions)
     let netCryptoSession =
             NCrypto { ncState         = ncState0
                     , ncTheirBaseNonce= ncTheirBaseNonce0
@@ -123,6 +133,7 @@ freshCryptoSession  sessions
                     , ncUnrecognizedHook = ncUnrecognizedHook0
                     , ncAllSessions = sessions
                     , ncGroups = ncGroups0
+                    , ncMessageTypes = ncMessageTypes0
                     }
     atomically $ modifyTVar allsessions (Map.insert addr netCryptoSession)
 
@@ -239,18 +250,24 @@ cryptoNetHandler sessions addr (NetCrypto (CryptoPacket nonce16 encrypted)) = do
                     hookmap <- atomically $ readTVar ncHooks
                     -- run hook
                     flip fix (cd,hookmap) $ \lookupAgain (cd,hookmap) -> do
+                        msgTyps <- atomically $ readTVar ncMessageTypes
                         let msgTyp = cd ^. messageType
-                        case Map.lookup msgTyp hookmap of
+                            msgTypMapped16 = msgTypes ! msgId cd
+                            msgTypMapped = fromIntegral msgTypMapped16
+                        if msgTypMapped16 == 0
+                         then return id
+                         else
+                           case Map.lookup msgTypMapped hookmap of
                                 Nothing -> do -- no recognizing hook, run ncUnrecognizedHook0, loopAgain on result
                                              unrecognize <- atomically $ readTVar (ncUnrecognizedHook session)
-                                             mbConsume <- unrecognize msgTyp session cd
+                                             mbConsume <- unrecognize msgTypMapped session cd
                                              case mbConsume of
                                                 Just f -> do
                                                            -- ncUnrecognizedHook0 may have updated the hookmap
                                                            hookmap' <- atomically $ readTVar ncHooks
                                                            lookupAgain (f cd,hookmap')
                                                 Nothing -> return Nothing
-                                Just hooks -> flip fix (hooks,cd,msgTyp) $ \loop (hooks,cd,typ) -> do
+                                Just hooks -> flip fix (hooks,cd,msgTypMapped) $ \loop (hooks,cd,typ) -> do
                                                 let _ = cd :: CryptoData
                                                 case (hooks,cd) of
                                                     ([],_) ->  return Nothing
@@ -269,6 +286,42 @@ cryptoNetHandler sessions addr (NetCrypto (CryptoPacket nonce16 encrypted)) = do
                                     _ -> error "unreachable-last2Bytes"
         dATA_NUM_THRESHOLD = 21845 -- = 65535 / 3
 
+-- | construct a 'MsgTypeArray' for specified types, using their known common positions
+--   in the MessageId space if they have such a thing.
+mkMsgTypes :: [MessageType] -> MsgTypeArray
+mkMsgTypes msgs = let zeros = A.listArray (0,255) (replicate 256 0)
+                      in zeros A.// map (\x -> (toIndex x,toWord16 x)) msgs
+    where
+        toIndex (Msg mid) = fromIntegral . fromEnum  $ mid
+        toIndex (GrpMsg KnownLossless nam) = 0x63 -- fromEnum MESSAGE_GROUPCHAT
+        toIndex (GrpMsg KnownLossy    nam) = 0xC7 -- fromEnum LOSSY_GROUPCHAT
+
+-- | Handle all Tox messages that this code base is aware of.
+allMsgTypes :: MsgTypeArray
+allMsgTypes = A.listArray (minBound,maxBound) (0:knownMsgs)
+  where
+    knownMsgs :: [Word16]
+    knownMsgs =
+      concat [ map (fromIntegral . fromEnum) [ PacketRequest .. KillPacket ]
+             , map (const 0)                 [ 3 .. 15 ]  -- UnspecifiedPacket
+             , map (const 0)                 [ 16 .. 23 ] -- MessengerLoseless
+             , map (fromIntegral . fromEnum) [ ONLINE .. OFFLINE ]
+             , map (const 0)                 [ 26 .. 47 ] -- MessengerLoseless
+             , map (fromIntegral . fromEnum) [ NICKNAME .. TYPING ]
+             , map (const 0)                 [ 52 .. 63 ] -- MessengerLoseless
+             , map (fromIntegral . fromEnum) [ MESSAGE .. ACTION ]
+             , map (const 0)                 [ 66 .. 68 ] -- MessengerLoseless
+             , map (fromIntegral . fromEnum) [ MSI ]
+             , map (const 0)                 [ 70 .. 79 ] -- MessengerLoseless
+             , map (fromIntegral . fromEnum) [ FILE_SENDREQUEST .. FILE_DATA ]
+             , map (const 0)                 [ 83 .. 95 ] -- MessengerLoseless
+             , map (fromIntegral . fromEnum) [ INVITE_GROUPCHAT .. MESSAGE_GROUPCHAT ]
+             , map (const 0)                 [ 100 .. 191 ] -- MessengerLoseless
+             , map (const 0)                 [ 192 .. 198 ] -- MessengerLossy
+             , map (fromIntegral . fromEnum) [ LOSSY_GROUPCHAT ]
+             , map (const 0)                 [ 200 .. 255 ] -- All lossy, exept the last
+             ]
+
 -- | handles nothing
 defaultCryptoDataHooks :: Map.Map MessageType [NetCryptoHook]
 defaultCryptoDataHooks = Map.empty