1 files changed, 24 insertions, 17 deletions

diff --git a/src/Network/Tox/DHT/Transport.hs b/src/Network/Tox/DHT/Transport.hs
index 79643fad..dd2838f2 100644
--- a/src/Network/Tox/DHT/Transport.hs
+++ b/src/Network/Tox/DHT/Transport.hs
@@ -38,12 +38,14 @@ import qualified Crypto.Tox as ToxCrypto
 import Network.QueryResponse
 
 import Control.Arrow
+import Control.Concurrent.STM
 import Control.Monad
 import Data.Bool
 import qualified Data.ByteString.Char8  as B8
 import qualified Data.ByteString  as B
          ;import Data.ByteString  (ByteString)
 import Data.Functor.Contravariant
+import Data.Maybe
 import Data.Monoid
 import Data.Serialize             as S
 import Data.Tuple
@@ -89,23 +91,28 @@ mapMessage f (DHTCookie nonce fcookie) = f nonce fcookie
 instance Sized Ping where size = ConstSize 1
 instance Sized Pong where size = ConstSize 1
 
-parseDHTAddr :: (ByteString, SockAddr) -> Either (DHTMessage Encrypted8,NodeInfo) (ByteString,SockAddr)
-parseDHTAddr (msg,saddr)
+parseDHTAddr :: TransportCrypto -> (ByteString, SockAddr) -> IO (Either (DHTMessage Encrypted8,NodeInfo) (ByteString,SockAddr))
+parseDHTAddr crypto (msg,saddr)
     | Just (typ,bs) <- B.uncons msg
-    , let right = Right (msg,saddr)
-          left  = either (const right) Left
+    , let right = return $ Right (msg,saddr)
+          left  = either (const right) (return . Left)
     = case typ of
     0x00 -> left $ direct bs saddr DHTPing
     0x01 -> left $ direct bs saddr DHTPong
     0x02 -> left $ direct bs saddr DHTGetNodes
     0x04 -> left $ direct bs saddr DHTSendNodes
     0x18 -> left $ direct bs saddr DHTCookieRequest
-    0x19 -> left $ fanGet bs getCookie (uncurry DHTCookie) (const $ noReplyAddr saddr)
+    0x19 -> do
+        cs <- atomically $ readTVar (pendingCookies crypto)
+        let ni = fromMaybe (noReplyAddr saddr) $ do
+                    (cnt,key) <- lookup saddr cs
+                    either (const Nothing) Just $ nodeInfo (key2id key) saddr
+        left $ fanGet bs getCookie (uncurry DHTCookie) (const $ ni)
     0x20 -> left $ fanGet bs getDHTReqest (uncurry DHTDHTRequest) (asymNodeInfo saddr . snd)
     _ -> right
 
-encodeDHTAddr :: (DHTMessage Encrypted8,NodeInfo) -> (ByteString, SockAddr)
-encodeDHTAddr (msg,ni) = (runPut $ putMessage msg, nodeAddr ni)
+encodeDHTAddr :: (DHTMessage Encrypted8,NodeInfo) -> IO (ByteString, SockAddr)
+encodeDHTAddr (msg,ni) = return (runPut $ putMessage msg, nodeAddr ni)
 
 dhtMessageType :: DHTMessage Encrypted8 -> ( Word8, Put )
 dhtMessageType (DHTPing           a) = (0x00, putAsymm a)
@@ -416,26 +423,26 @@ encryptMessage :: Serialize a =>
                   TransportCrypto ->
                   PublicKey ->
                   Nonce24 -> Either (Nonce8,a) (Asymm (Nonce8,a)) -> Encrypted8 a
-encryptMessage crypto destKey n (Right asymm) = E8 $ ToxCrypto.encrypt secret plain
+encryptMessage crypto destKey n arg = E8 $ ToxCrypto.encrypt secret plain
  where
     secret = computeSharedSecret (transportSecret crypto) destKey n
-    plain  = encodePlain $ swap $ asymmData asymm
-encryptMessage crypto destKey n (Left plain) = _todo -- need cached public key.
+    plain  = encodePlain $ swap $ either id asymmData arg
 
 decrypt :: TransportCrypto -> DHTMessage Encrypted8 -> NodeInfo -> Either String (DHTMessage ((,) Nonce8), NodeInfo)
-decrypt crypto msg ni = (, ni) <$> (sequenceMessage $ transcode (decryptMessage crypto) msg)
+decrypt crypto msg ni = do
+    msg' <- sequenceMessage $ transcode (\n -> decryptMessage crypto n . left ((,) $ id2key $ nodeId ni)) msg
+    return (msg', ni)
 
 decryptMessage :: Serialize x =>
                   TransportCrypto
                   -> Nonce24
-                  -> Either (Encrypted8 x) (Asymm (Encrypted8 x))
+                  -> Either (PublicKey, Encrypted8 x) (Asymm (Encrypted8 x))
                   -> (Either String ∘ ((,) Nonce8)) x
-decryptMessage crypto n (Right asymmE) = plain8 $ ToxCrypto.decrypt secret e
+decryptMessage crypto n arg = plain8 $ ToxCrypto.decrypt secret e
  where
-    secret = computeSharedSecret (transportSecret crypto) (senderKey asymmE) n
-    E8 e   = asymmData asymmE
-    plain8 = Composed . fmap swap . (>>= decodePlain)
-decryptMessage crypto n (Left (E8 e)) = _todo -- need cached public key
+    secret           = computeSharedSecret (transportSecret crypto) remotekey n
+    (remotekey,E8 e) = either id (senderKey &&& asymmData) arg
+    plain8           = Composed . fmap swap . (>>= decodePlain)
 
 sequenceMessage :: Applicative m => DHTMessage (m ∘ f) -> m (DHTMessage f)
 sequenceMessage (DHTPing asym)              = fmap DHTPing          $ sequenceA $ fmap uncomposed asym