1 files changed, 18 insertions, 0 deletions

diff --git a/src/Network/Tox/Crypto/Transport.hs b/src/Network/Tox/Crypto/Transport.hs
index 8b7a7222..0011fac0 100644
--- a/src/Network/Tox/Crypto/Transport.hs
+++ b/src/Network/Tox/Crypto/Transport.hs
@@ -1,5 +1,7 @@
 {-# LANGUAGE KindSignatures #-}
 {-# LANGUAGE ViewPatterns #-}
+{-# LANGUAGE FlexibleInstances #-}
+{-# LANGUAGE LambdaCase #-}
 module Network.Tox.Crypto.Transport
     ( parseNetCrypto
     , encodeNetCrypto
@@ -40,6 +42,7 @@ import Network.Tox.DHT.Transport (Cookie)
 
 import Network.Socket
 import Data.ByteString as B
+import Data.Monoid
 import Data.Word
 import Crypto.Hash
 import Control.Lens
@@ -76,6 +79,15 @@ data CryptoPacket (f :: * -> *) = CryptoPacket
     , pktData :: f CryptoData
     }
 
+instance Sized CryptoData where
+    size = contramap bufferStart size
+            <> contramap bufferEnd size
+            <> contramap bufferData size
+
+instance Serialize (CryptoPacket Encrypted) where
+    get = CryptoPacket <$> get <*> get
+    put (CryptoPacket n16 dta) = put n16 >> put dta
+
 data CryptoData = CryptoData
     { -- | [our recvbuffers buffer_start, (highest packet number handled + 1), (big endian)]
       bufferStart :: Word32
@@ -108,6 +120,12 @@ data CryptoMessage
         | UpToN    { msgID :: MessageID, msgBytes :: ByteString } -- length < N
     deriving (Eq,Show)
 
+instance Sized CryptoMessage where
+    size = VarSize $ \case
+        OneByte {}              -> 1
+        TwoByte {}              -> 2
+        UpToN { msgBytes = bs } -> 1 + B.length bs
+
 erCompat lens = error $ "Use of '" ++ lens ++ "'  lens on incompatible CryptoMessage type"
 
 typingStatus :: Functor f => (UserStatus -> f UserStatus)-> (CryptoMessage -> f CryptoMessage)