1 files changed, 20 insertions, 17 deletions

diff --git a/src/Network/Tox/DHT/Transport.hs b/src/Network/Tox/DHT/Transport.hs
index 736e84d1..bd108276 100644
--- a/src/Network/Tox/DHT/Transport.hs
+++ b/src/Network/Tox/DHT/Transport.hs
@@ -426,34 +426,37 @@ forwardDHTRequests crypto closeLookup dht = dht { awaitMessage = await' }
                   await' pass
         m -> pass m
 
-encrypt :: TransportCrypto -> DHTMessage ((,) Nonce8) -> NodeInfo -> (DHTMessage Encrypted8, NodeInfo)
-encrypt crypto msg ni = ( transcode (encryptMessage crypto (id2key $ nodeId ni)) msg
-                        , ni )
+encrypt :: TransportCrypto -> DHTMessage ((,) Nonce8) -> NodeInfo -> IO (DHTMessage Encrypted8, NodeInfo)
+encrypt crypto msg ni = do
+    let cipher n plain = Composed $ encryptMessage crypto (id2key $ nodeId ni) n plain
+    m <- sequenceMessage $ transcode cipher msg
+    return (m, ni)
 
 encryptMessage :: Serialize a =>
                   TransportCrypto ->
                   PublicKey ->
-                  Nonce24 -> Either (Nonce8,a) (Asymm (Nonce8,a)) -> Encrypted8 a
-encryptMessage crypto destKey n arg = E8 $ ToxCrypto.encrypt secret plain
- where
-    secret = computeSharedSecret (transportSecret crypto) destKey n
-    plain  = encodePlain $ swap $ either id asymmData arg
+                  Nonce24 -> Either (Nonce8,a) (Asymm (Nonce8,a)) -> IO (Encrypted8 a)
+encryptMessage crypto destKey n arg = do
+    let plain  = encodePlain $ swap $ either id asymmData arg
+    secret <- lookupSharedSecret crypto (transportSecret crypto) destKey n
+    return $ E8 $ ToxCrypto.encrypt secret plain
 
-decrypt :: TransportCrypto -> DHTMessage Encrypted8 -> NodeInfo -> Either String (DHTMessage ((,) Nonce8), NodeInfo)
+decrypt :: TransportCrypto -> DHTMessage Encrypted8 -> NodeInfo -> IO (Either String (DHTMessage ((,) Nonce8), NodeInfo))
 decrypt crypto msg ni = do
-    msg' <- sequenceMessage $ transcode (\n -> decryptMessage crypto n . left ((,) $ id2key $ nodeId ni)) msg
-    return (msg', ni)
+    let decipher n c = Composed $ decryptMessage crypto n . left ((,) $ id2key $ nodeId ni) $ c
+    msg' <- sequenceMessage $ transcode decipher msg
+    return $ fmap (, ni) $ sequenceMessage msg'
 
 decryptMessage :: Serialize x =>
                   TransportCrypto
                   -> Nonce24
                   -> Either (PublicKey, Encrypted8 x) (Asymm (Encrypted8 x))
-                  -> (Either String ∘ ((,) Nonce8)) x
-decryptMessage crypto n arg = plain8 $ ToxCrypto.decrypt secret e
- where
-    secret           = computeSharedSecret (transportSecret crypto) remotekey n
-    (remotekey,E8 e) = either id (senderKey &&& asymmData) arg
-    plain8           = Composed . fmap swap . (>>= decodePlain)
+                  -> IO ((Either String ∘ ((,) Nonce8)) x)
+decryptMessage crypto n arg = do
+    let (remotekey,E8 e) = either id (senderKey &&& asymmData) arg
+        plain8           = Composed . fmap swap . (>>= decodePlain)
+    secret <- lookupSharedSecret crypto (transportSecret crypto) remotekey n
+    return $ plain8 $ ToxCrypto.decrypt secret e
 
 sequenceMessage :: Applicative m => DHTMessage (m ∘ f) -> m (DHTMessage f)
 sequenceMessage (DHTPing asym)              = fmap DHTPing          $ sequenceA $ fmap uncomposed asym