2 files changed, 20 insertions, 14 deletions

diff --git a/src/Network/Tox/Crypto/Handlers.hs b/src/Network/Tox/Crypto/Handlers.hs
index 12818b2e..c5476371 100644
--- a/src/Network/Tox/Crypto/Handlers.hs
+++ b/src/Network/Tox/Crypto/Handlers.hs
@@ -28,7 +28,7 @@ data NetCryptoSessionStatus = Unaccepted | Accepted | Confirmed
 
 
 type IOHook addr x = addr -> x -> IO (Maybe (x -> x))
-type NetCryptoHook = IOHook SockAddr CryptoMessage
+type NetCryptoHook = IOHook SockAddr CryptoData
 
 
 data NetCryptoSession = NCrypto { ncState         :: TVar NetCryptoSessionStatus
@@ -160,7 +160,7 @@ cryptoNetHandler sessions addr (NetCrypto (CryptoPacket nonce16 encrypted)) = do
                         decodePlain =<< decrypt (computeSharedSecret ncSessionSecret pubkey tempNonce) encrypted
             case lr of
                 Left _ -> return Nothing -- decryption failed, ignore packet
-                Right (CryptoData {bufferStart, bufferEnd, bufferData=cm}) -> do -- decryption succeeded,
+                Right cd@(CryptoData {bufferStart, bufferEnd, bufferData=cm}) -> do -- decryption succeeded,
                     -- TODO: Why do I need bufferStart & bufferEnd?
                     --
                     --   buffer_start = highest packet number handled + 1
@@ -177,22 +177,23 @@ cryptoNetHandler sessions addr (NetCrypto (CryptoPacket nonce16 encrypted)) = do
                     -- then set session confirmed,
                     atomically $ writeTVar ncState Confirmed
                     hookmap <- atomically $ readTVar ncHooks
-                    -- if lossy, just run hook
-                    if lossyness (msgID cm) == Lossy
-                      then
-                            case Map.lookup (cm ^. messageType) hookmap of
+                    -- run hook
+                    flip fix cd $ \lookupAgain cd -> do
+                        let msgTyp = cd ^. messageType
+                        case Map.lookup msgTyp hookmap of
                                 Nothing -> return Nothing -- discarding, because no hooks
-                                Just hooks -> flip fix (hooks,cm) $ \loop (hooks,msg) -> do
-                                                let _ = cm :: CryptoMessage
-                                                case (hooks,cm) of
+                                Just hooks -> flip fix (hooks,cd,msgTyp) $ \loop (hooks,cd,typ) -> do
+                                                let _ = cd :: CryptoData
+                                                case (hooks,cd) of
                                                     ([],_) ->  return Nothing
-                                                    (hook:more,cm) -> do
-                                                        r <- hook addr cm :: IO (Maybe (CryptoMessage -> CryptoMessage))
+                                                    (hook:more,cd) -> do
+                                                        r <- hook addr cd :: IO (Maybe (CryptoData -> CryptoData))
                                                         case r of
-                                                            Just f -> loop (more,f cm)
+                                                            Just f -> let newcd = f cd
+                                                                          newtyp = newcd ^. messageType
+                                                                          in if newtyp == typ then loop (more,newcd,newtyp)
+                                                                                              else lookupAgain newcd
                                                             Nothing -> return Nothing -- message consumed
-                      else -- Lossless message, so try to restore sequence
-                        error "todo try to restore sequence of lossless messages"
     where
         last2Bytes :: Nonce24 -> Word
         last2Bytes (Nonce24 bs) = case S.decode (B.drop 22 bs) of
diff --git a/src/Network/Tox/Crypto/Transport.hs b/src/Network/Tox/Crypto/Transport.hs
index 7bc6e67f..8739c853 100644
--- a/src/Network/Tox/Crypto/Transport.hs
+++ b/src/Network/Tox/Crypto/Transport.hs
@@ -1,4 +1,5 @@
 {-# LANGUAGE KindSignatures #-}
+{-# LANGUAGE NamedFieldPuns #-}
 {-# LANGUAGE ViewPatterns #-}
 {-# LANGUAGE FlexibleInstances #-}
 {-# LANGUAGE LambdaCase #-}
@@ -434,6 +435,10 @@ instance HasMessageType CryptoMessage where
     setMessageType (TwoByte mid0 x) (Msg mid) = UpToN mid (B.singleton x)
     setMessageType (UpToN mid0 x)   (Msg mid) = UpToN mid x
 
+instance HasMessageType CryptoData where
+    getMessageType (CryptoData { bufferData })     = getMessageType bufferData
+    setMessageType cd@(CryptoData { bufferData=bd }) typ = cd { bufferData=setMessageType bd typ }
+
 -- | This lens should always succeed on CryptoMessage
 messageType :: (Functor f, HasMessageType x) => (MessageType -> f MessageType) -> (x -> f x)
 messageType = lens getMessageType setMessageType