More TCP relay packet serialization.

3 files changed, 130 insertions, 2 deletions

diff --git a/src/Crypto/Tox.hs b/src/Crypto/Tox.hs
index ad246d4e..bd4cdfba 100644
--- a/src/Crypto/Tox.hs
+++ b/src/Crypto/Tox.hs
@@ -6,11 +6,14 @@
 {-# LANGUAGE DeriveFunctor #-}
 {-# LANGUAGE DeriveGeneric #-}
 {-# LANGUAGE DeriveTraversable #-}
-{-# LANGUAGE ExplicitNamespaces #-}
 {-# LANGUAGE TypeOperators #-}
 {-# LANGUAGE MagicHash, UnboxedTuples, BangPatterns #-}
+{-# LANGUAGE MultiParamTypeClasses #-}
+{-# LANGUAGE ConstraintKinds #-}
+{-# LANGUAGE Rank2Types #-}
 {-# LANGUAGE NamedFieldPuns #-}
 {-# LANGUAGE PatternSynonyms #-}
+{-# LANGUAGE FlexibleContexts #-}
 module Crypto.Tox
     ( PublicKey
     , publicKey
@@ -38,8 +41,11 @@ module Crypto.Tox
     , lookupSharedSecret
     , lookupNonceFunction
     , lookupNonceFunctionSTM
+    , Payload(..)
     , encrypt
     , decrypt
+    , decryptPayload
+    , encryptPayload
     , Nonce8(..)
     , Nonce24(..)
     , incrementNonce24
@@ -110,6 +116,8 @@ import Data.Time.Clock.POSIX
 import Data.Hashable
 import System.IO.Unsafe (unsafeDupablePerformIO)
 import Data.Functor.Compose
+import qualified Rank2
+import Data.Functor.Identity
 
 -- | A 16-byte mac and an arbitrary-length encrypted stream.
 newtype Encrypted a = Encrypted ByteString
@@ -253,6 +261,18 @@ decrypt (State hash crypt) ciphertext
     m = fst . XSalsa.combine crypt $ c
     a = Auth . Poly1305.finalize . Poly1305.update hash $ c
 
+class Rank2.Functor g => Payload c g where
+    mapPayload :: proxy c -> (forall a. c a => p a -> q a) -> g p -> g q
+
+decryptPayload :: ( Rank2.Traversable g
+                  , Payload Serialize g
+                  ) => State -> g Encrypted -> Either String (g Identity)
+decryptPayload st g = do
+    plain <- Rank2.traverse (decrypt st) g
+    Rank2.sequence $ mapPayload (Proxy :: Proxy Serialize)
+                                (Composed . fmap pure . decodePlain)
+                                plain
+
 -- Encrypt-then-Mac:  Encrypt the cleartext, then compute the MAC on the
 -- ciphertext, and prepend it to the ciphertext
 encrypt :: State -> Plain s a -> Encrypted a
@@ -261,6 +281,13 @@ encrypt (State hash crypt) (Plain m) = Encrypted $ B.append (encode a) c
     c = fst . XSalsa.combine crypt $ m
     a = Auth . Poly1305.finalize . Poly1305.update hash $ c
 
+encryptPayload :: Payload Serialize g => State -> g Identity -> g Encrypted
+encryptPayload st g =
+    encrypt st
+        Rank2.<$> mapPayload (Proxy :: Proxy Serialize)
+                             (encodePlain . runIdentity)
+                             g
+
 -- (Poly1305.State, XSalsa.State)
 computeSharedSecret :: SecretKey -> PublicKey -> Nonce24 -> State
 computeSharedSecret sk recipient = k `seq` \nonce ->
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
diff --git a/src/Network/Tox/Onion/Transport.hs b/src/Network/Tox/Onion/Transport.hs
index ef9121f2..160b99f7 100644
--- a/src/Network/Tox/Onion/Transport.hs
+++ b/src/Network/Tox/Onion/Transport.hs
@@ -347,6 +347,11 @@ deriving instance ( Show (Forwarding (ThreeMinus n) (OnionMessage Encrypted))
                   , KnownNat (PeanoNat n)
                   ) => Show (OnionRequest n)
 
+instance Sized (OnionRequest N0) where -- N1 and N2 are the same, N3 does not encode the nonce.
+    size = contramap onionNonce size
+        <> contramap onionForward size
+        <> contramap pathFromOwner size
+
 instance ( Serialize (Forwarding (ThreeMinus n) (OnionMessage Encrypted))
          , Sized (ReturnPath n)
          , Serialize (ReturnPath n)
@@ -383,6 +388,8 @@ instance ( Serialize (ReturnPath n) ) => Serialize (OnionResponse n) where
                                              . getOnionReply)
     put (OnionResponse p m) = put p >> putOnionMsg m
 
+instance (Sized (ReturnPath n)) => Sized (OnionResponse (S n)) where
+    size = contramap pathToOwner size <> contramap msgToOwner size
 
 data Addressed a = Addressed { sockAddr :: SockAddr, unaddressed :: a }
  deriving (Eq,Show)