1 files changed, 14 insertions, 7 deletions

diff --git a/src/Network/Tox/Crypto/Handlers.hs b/src/Network/Tox/Crypto/Handlers.hs
index b8e99d2d..ac3d1ef0 100644
--- a/src/Network/Tox/Crypto/Handlers.hs
+++ b/src/Network/Tox/Crypto/Handlers.hs
@@ -94,9 +94,9 @@ freshCryptoSession  sessions
     ncState0 <- atomically $ newTVar Accepted
     ncTheirBaseNonce0 <- atomically $ newTVar theirBaseNonce
     n24 <- atomically $ transportNewNonce crypto
+    state <- lookupSharedSecret crypto key remoteDhtPublicKey n24
     let myhandshakeData = newHandShakeData crypto hp
         plain = encodePlain myhandshakeData
-        state = computeSharedSecret key remoteDhtPublicKey n24
         encrypted = encrypt state plain
         myhandshake = Handshake { handshakeCookie = otherCookie
                                 , handshakeNonce = n24
@@ -150,14 +150,19 @@ cryptoNetHandler sessions addr (NetHandshake (Handshake (Cookie n24 ecookie) non
     -- Handle Handshake Message
     let crypto = transportCrypto sessions
         allsessions = netCryptoSessions sessions
-        anyRight xs f = foldr1 (<|>) $ map f xs
+        anyRight [] f     = return $ Left "missing key"
+        anyRight (x:xs) f = f x >>= either (const $ anyRight xs f) (return . Right)
     seckeys <- map fst <$> atomically (readTVar (userKeys crypto))
     symkey <- atomically $ transportSymmetric crypto
     now <- getPOSIXTime
-    let lr = do -- Either Monad
+    lr <- fmap join . sequence $ do -- Either Monad
             (CookieData cookieTime remotePubkey remoteDhtkey) <- (decodePlain =<< decryptSymmetric symkey n24 ecookie)
-            (key,HandshakeData { baseNonce, sessionKey, cookieHash, otherCookie })
-                <- anyRight seckeys $ \key -> (key,) <$> (decodePlain =<< decrypt (computeSharedSecret key remotePubkey nonce24) encrypted)
+            Right $ do -- IO Monad
+            decrypted <- anyRight seckeys $ \key -> do
+                            secret <- lookupSharedSecret crypto key remotePubkey nonce24
+                            return $ (key,) <$> (decodePlain =<< decrypt secret encrypted)
+            return $ do -- Either Monad
+            (key,HandshakeData { baseNonce, sessionKey, cookieHash, otherCookie }) <- decrypted
             -- check cookie time < 15 seconds ago
             guard (now - fromIntegral cookieTime < 15)
             -- cookie hash is valid? sha512 of ecookie
@@ -208,9 +213,11 @@ cryptoNetHandler sessions addr (NetCrypto (CryptoPacket nonce16 encrypted)) = do
             let diff :: Word16
                 diff = nonce16 - fromIntegral (last2Bytes theirBaseNonce) -- truncating to Word16
             tempNonce <- addtoNonce24 theirBaseNonce (fromIntegral diff) -- expanding to Word
-            let lr = do -- Either Monad --
+            lr <- fmap join $ sequence $ do -- Either Monad --
                         pubkey <- maybeToEither ncTheirSessionPublic
-                        decodePlain =<< decrypt (computeSharedSecret ncSessionSecret pubkey tempNonce) encrypted
+                        Right $ do -- IO Monad
+                        secret <- lookupSharedSecret crypto ncSessionSecret pubkey tempNonce
+                        return $ decodePlain =<< decrypt secret encrypted
             case lr of
                 Left _ -> return Nothing -- decryption failed, ignore packet
                 Right cd@(CryptoData {bufferStart, bufferEnd, bufferData=cm}) -> do -- decryption succeeded,