1 files changed, 28 insertions, 1 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 ->