1 files changed, 49 insertions, 16 deletions

diff --git a/src/Network/Tox/Crypto/Transport.hs b/src/Network/Tox/Crypto/Transport.hs
index 2a0633f0..84929e63 100644
--- a/src/Network/Tox/Crypto/Transport.hs
+++ b/src/Network/Tox/Crypto/Transport.hs
@@ -10,6 +10,7 @@ module Network.Tox.Crypto.Transport
     , parseCrypto
     , encodeCrypto
     , unpadCryptoMsg
+    , decodeRawCryptoMsg
     , createRequestPacket
     , parseHandshakes
     , encodeHandshakes
@@ -59,6 +60,8 @@ module Network.Tox.Crypto.Transport
     , msgSizeParam
     , NegotiationID(..)
     , NegotiationMsg(..)
+    , getCryptoMessage
+    , putCryptoMessage
     , module Data.Tox.Message
     ) where
 
@@ -70,21 +73,21 @@ import Network.Tox.NodeId
 import Network.Socket
 import Data.ByteArray
 
+import Control.Monad
 import Data.ByteString as B
+import Data.Function
 import Data.Maybe
 import Data.Monoid
 import Data.Word
 import Data.Bits
 import Crypto.Hash
 import Control.Lens
-import Control.Monad
 import Data.Text as T
 import Data.Text.Encoding as T
 import Data.Serialize as S
 import Control.Arrow
 import DPut
 import Data.PacketBuffer as PB
-import Data.Function
 
 showCryptoMsg :: Word32 -> CryptoMessage -> [Char]
 showCryptoMsg seqno (UpToN PacketRequest bytes) = "UpToN PacketRequest --> "
@@ -234,8 +237,15 @@ One effect of this is that short messages will be padded to at least 5 bytes.
 -}
 
 instance Serialize CryptoData where
-    get = CryptoData <$> get <*> get <*> get
-    put (CryptoData start end dta) = put start >> put end >> put dta
+    get = do
+        ack <- get
+        seqno <- get
+        cm <- getCryptoMessage ack
+        return $ CryptoData ack seqno cm
+    put (CryptoData ack seqno dta) = do
+        put ack
+        put seqno
+        putCryptoMessage ack dta
 
 -- The 'UserStatus' equivalent in Presence is:
 --
@@ -274,32 +284,55 @@ unpadCryptoMsg x@(UpToN mid0 (B.dropWhile (==0) -> B.uncons -> Just (toEnum8 ->
             _ -> UpToN mid bytes
 unpadCryptoMsg x = x
 
+decodeRawCryptoMsg :: CryptoData -> CryptoMessage
+decodeRawCryptoMsg (CryptoData ack seqno cm) =
+    let cm' = unpadCryptoMsg cm
+    in case msgID cm' of
+        PacketRequest -> RequestResend PacketRequest $ decompressSequenceNumbers ack $ msgByteList cm'
+        _             -> cm'
+
 data CryptoMessage
         = OneByte  { msgID :: MessageID }
         | TwoByte  { msgID :: MessageID, msgByte  :: Word8 }
         | UpToN    { msgID :: MessageID, msgBytes :: ByteString } -- length < N
+          -- | TODO: The msgID field is redundant in this case and can be removed
+          -- after all uses are audited.
+        | RequestResend { msgID :: MessageID, requested :: [Word32] }
     deriving (Eq,Show)
 
+msgByteList :: CryptoMessage -> [Word8]
+msgByteList (UpToN _ bs)  = B.unpack bs
+msgByteList (TwoByte _ b) = [b]
+msgByteList (OneByte _)   = []
+
 instance Sized CryptoMessage where
     size = VarSize $ \case
-        OneByte {}              -> 1
-        TwoByte {}              -> 2
-        UpToN { msgBytes = bs } -> 1 + B.length bs
+        OneByte {}                       -> 1
+        TwoByte {}                       -> 2
+        UpToN { msgBytes = bs }          -> 1 + B.length bs
+        RequestResend { requested = ws } -> 1 + Prelude.length ws
 
-instance Serialize CryptoMessage where
-    get = do
+getCryptoMessage :: Word32 -> Get CryptoMessage
+getCryptoMessage seqno = do
         i <- get :: Get MessageID
         n <- remaining
-        case msgSizeParam i of
+        pkt <- case msgSizeParam i of
             Just (True,0) -> return $ OneByte i
             Just (True,1) -> TwoByte i <$> get
             _ -> UpToN i <$> getByteString n
-
-    put (OneByte i)   =    putWord8 (fromIntegral . fromEnum $ i)
-    put (TwoByte i b) = do putWord8 (fromIntegral . fromEnum $ i)
-                           putWord8 b 
-    put (UpToN i x)   = do putWord8 (fromIntegral . fromEnum $ i)
-                           putByteString x
+        return $ if msgID pkt == PacketRequest
+                  then RequestResend PacketRequest $ decompressSequenceNumbers seqno $ msgByteList pkt
+                  else pkt
+
+putCryptoMessage :: Word32 -> CryptoMessage -> Put
+putCryptoMessage seqno (OneByte i)   = putWord8 (fromIntegral . fromEnum $ i)
+putCryptoMessage seqno (TwoByte i b) = do putWord8 (fromIntegral . fromEnum $ i)
+                                          putWord8 b
+putCryptoMessage seqno (UpToN i x)   = do putWord8 (fromIntegral . fromEnum $ i)
+                                          putByteString x
+putCryptoMessage seqno (RequestResend _ ws) = do
+    putWord8 (fromIntegral . fromEnum $ PacketRequest)
+    mapM_ putWord8 $ compressSequenceNumbers seqno ws
 
 instance Serialize MessageID where
     get = toEnum . fromIntegral <$> getWord8