Use seqence number context when serializing CryptoMessage.

3 files changed, 55 insertions, 22 deletions

diff --git a/examples/atox.hs b/examples/atox.hs
index 7845d911..3bae5203 100644
--- a/examples/atox.hs
+++ b/examples/atox.hs
@@ -92,7 +92,7 @@ data SetCmd = SetME
     deriving (Eq,Bounded,Ord,Enum,Show)
 
 forkToxInputThread myRead = forkIO $ do
-    let myconduit = Conduit.sourceHandle myRead  .| conduitGet2 (S.get :: Get CryptoMessage) -- :: ConduitT i CryptoMessage IO ()
+    let myconduit = Conduit.sourceHandle myRead  .| conduitGet2 (getCryptoMessage 0 :: Get CryptoMessage) -- :: ConduitT i CryptoMessage IO ()
     Conduit.runConduit (myconduit .| Conduit.awaitForever handle)
  where
     handle (UpToN IPC (B.uncons -> Just (ord -> toEnum -> i,arg))) = updateState i arg
@@ -113,7 +113,7 @@ forkToxInputThread myRead = forkIO $ do
                                     modifyTVar' sMap (Map.insert key view)
 
     updateState AppendMsg arg
-             = case S.decode arg of
+             = case getCryptoMessage 0 `S.runGet` arg of
                  Left str  -> puts (packUtf8 str)
                  Right msg -> liftIO . atomically $ do
                                 me <- readTVar sMe
@@ -133,7 +133,7 @@ doit myReadFd myWriteFd = do
     terminalInputLoop myWrite
 
 hSend :: MonadIO m => Handle -> CryptoMessage -> m ()
-hSend h msg = liftIO $ B.hPutStrLn h (S.encode msg)
+hSend h msg = liftIO $ B.hPutStrLn h (S.runPut $ putCryptoMessage 0 msg)
 
 terminalInputLoop myWriteH = fix $ \loop -> do
     line <- B.getLine
diff --git a/examples/dhtd.hs b/examples/dhtd.hs
index 6161fbb4..5bc2b87a 100644
--- a/examples/dhtd.hs
+++ b/examples/dhtd.hs
@@ -99,7 +99,7 @@ import qualified Network.Tox.DHT.Transport as Tox
 import qualified Network.Tox.DHT.Handlers as Tox
 import qualified Network.Tox.Onion.Transport as Tox
 import qualified Network.Tox.Onion.Handlers as Tox
-import qualified Network.Tox.Crypto.Transport as Tox (CryptoMessage(..))
+import qualified Network.Tox.Crypto.Transport as Tox (CryptoMessage(..),putCryptoMessage,getCryptoMessage)
 import qualified Network.Tox.Crypto.Handlers as Tox
 import Data.Typeable
 import Network.Tox.ContactInfo as Tox
@@ -1407,7 +1407,7 @@ netcrypto (Just (DHT {..})) (Just mypubkey) h roster (Just tox) exes paramStr =
                                              when (them /= their_pub) $ do
                                                 atomically $ writeTVar whoAreThey their_pub
                                                 B.hPutStr myWrite ("\NUL\SOH" `B.append` S.encode (Tox.key2id their_pub))
-                                             B.hPutStr myWrite (S.encode msg)
+                                             B.hPutStr myWrite (S.runPut $ Tox.putCryptoMessage 0 msg)
                                              return (Just id)
                                 addHooks currentHooks typs = forM_ typs $ \typ -> modifyTVar (Tox.ncHooks session) (Map.insert typ (currentHooks typ ++ [makeHook session typ]))
                             case ranges of
@@ -1426,7 +1426,7 @@ netcrypto (Just (DHT {..})) (Just mypubkey) h roster (Just tox) exes paramStr =
                                 tid <- myThreadId
                                 let sidStr = printf "(%x)" (Tox.ncSessionId session)
                                 labelThread tid (exekey ++ ".forward" ++ sidStr)
-                                let myconduit = Conduit.sourceHandle myRead  .| conduitGet2 S.get -- :: ConduitT i CryptoMessage IO ()
+                                let myconduit = Conduit.sourceHandle myRead  .| conduitGet2 (Tox.getCryptoMessage 0) -- :: ConduitT i CryptoMessage IO ()
                                 Conduit.runConduit (myconduit .| awaitForever (\msg -> do
                                     let typ = toWord64 (getMessageType msg)
                                     mbSendIt <- liftIO $ atomically (lookupInRangeMap typ (Tox.ncOutHooks session))
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