1 files changed, 95 insertions, 1 deletions

diff --git a/src/Data/Tox/Relay.hs b/src/Data/Tox/Relay.hs
index e35f49f0..a5366f85 100644
--- a/src/Data/Tox/Relay.hs
+++ b/src/Data/Tox/Relay.hs
@@ -1,10 +1,19 @@
 {-# LANGUAGE DeriveDataTypeable #-}
+{-# LANGUAGE FlexibleInstances  #-}
+{-# LANGUAGE KindSignatures     #-}
+{-# LANGUAGE MultiParamTypeClasses #-}
+{-# LANGUAGE ConstraintKinds #-}
 module Data.Tox.Relay where
 
-import Data.ByteString
+import Crypto.Error
+import Data.ByteArray as BA
+import Data.ByteString as B
 import Data.Data
+import Data.Functor.Contravariant
+import Data.Monoid
 import Data.Serialize
 import Data.Word
+import qualified Rank2
 
 import Crypto.Tox
 import Network.Tox.Onion.Transport
@@ -29,6 +38,24 @@ packetNumber :: RelayPacket -> Word8
 packetNumber (RelayData conid _) = conid -- 0 to 15 not allowed.
 packetNumber rp                  = fromIntegral $ pred $ constrIndex $ toConstr rp
 
+instance Sized RelayPacket where
+    size = mappend (ConstSize 1) $ VarSize $ \x -> case x of
+        RoutingRequest k             -> 32
+        RoutingResponse rpid k       -> 33
+        ConnectNotification conid    -> 1
+        DisconnectNotification conid -> 1
+        RelayPing pingid             -> 8
+        RelayPong pingid             -> 8
+        OOBSend k bs                 -> 32 + B.length bs
+        OOBRecv k bs                 -> 32 + B.length bs
+        OnionPacket query            -> case contramap (`asTypeOf` query) size of
+                                            ConstSize n -> n
+                                            VarSize f   -> f query
+        OnionPacketResponse answer   -> case contramap (`asTypeOf` answer) size of
+                                            ConstSize n -> n
+                                            VarSize f   -> f answer
+        RelayData _ bs               -> B.length bs
+
 instance Serialize RelayPacket where
 
     get = do
@@ -60,3 +87,70 @@ instance Serialize RelayPacket where
             OnionPacket query            -> put query
             OnionPacketResponse answer   -> put answer
             RelayData _ bs               -> putByteString bs
+
+-- | Initial client-to-server handshake message.
+data Hello (f :: * -> *) = Hello
+    { helloFrom  :: PublicKey
+    , helloNonce :: Nonce24
+    , helloData  :: f HelloData
+    }
+
+instance Rank2.Functor Hello where
+    f <$> Hello k n dta = Hello k n (f dta)
+
+instance Payload Serialize Hello where
+    mapPayload _ f (Hello k n dta) = Hello k n (f dta)
+
+instance Rank2.Foldable Hello where
+    foldMap f (Hello k n dta) = f dta
+
+instance Rank2.Traversable Hello where
+    traverse f (Hello k n dta) = Hello k n <$> f dta
+
+instance Sized (Hello Encrypted) where
+    size = ConstSize 56 <> contramap helloData size
+
+instance Serialize (Hello Encrypted) where
+    get = do CryptoPassed k <- publicKey <$> getBytes 32
+             n <- get
+             dta <- get
+             return $ Hello k n dta
+    put (Hello k n dta) = do mapM_ putWord8 $ BA.unpack k
+                             put n
+                             put dta
+
+data HelloData = HelloData
+    { sessionPublicKey :: PublicKey
+    , sessionBaseNonce :: Nonce24
+    }
+
+instance Sized HelloData where size = ConstSize 56
+
+instance Serialize HelloData where
+    get = HelloData <$> (id2key <$> get) <*> get
+    put (HelloData k n) = put (key2id k) >> put n
+
+-- | Handshake server-to-client response packet.
+data Welcome (f :: * -> *) = Welcome
+    { welcomeNonce :: Nonce24
+    , welcomeData  :: f HelloData
+    }
+
+instance Rank2.Functor Welcome where
+    f <$> Welcome n dta = Welcome n (f dta)
+
+instance Payload Serialize Welcome where
+    mapPayload _ f (Welcome n dta) = Welcome n (f dta)
+
+instance Rank2.Foldable Welcome where
+    foldMap f (Welcome _ dta) = f dta
+
+instance Rank2.Traversable Welcome where
+    traverse f (Welcome n dta) = Welcome n <$> f dta
+
+instance Sized (Welcome Encrypted) where
+    size = ConstSize 24 <> contramap welcomeData size
+
+instance Serialize (Welcome Encrypted) where
+    get = Welcome <$> get <*> get
+    put (Welcome n dta) = put n >> put dta