1 files changed, 8 insertions, 6 deletions

diff --git a/dht/HandshakeCache.hs b/dht/HandshakeCache.hs
index d9ffacab..00836fc8 100644
--- a/dht/HandshakeCache.hs
+++ b/dht/HandshakeCache.hs
@@ -27,13 +27,15 @@ data HandshakeCache = HandshakeCache
     { -- Note that currently we are storing sent handshakes keyed by the
       -- locally issued cookie nonce.
       hscTable  :: TVar (MinMaxPSQ' Nonce24 POSIXTime (SecretKey,HandshakeData))
-    , hscSend   :: Multi.SessionAddress -> Handshake Encrypted -> IO ()
+    , hscSend   :: Multi.SessionAddress -> Handshake Encrypted -> IO Multi.SessionAddress
     , hscCrypto :: TransportCrypto
     , hscPendingCookies :: TVar (Map (PublicKey,PublicKey) ())
     }
 
 
-newHandshakeCache :: TransportCrypto -> (Multi.SessionAddress -> Handshake Encrypted -> IO ()) -> IO HandshakeCache
+newHandshakeCache :: TransportCrypto
+                        -> (Multi.SessionAddress -> Handshake Encrypted -> IO Multi.SessionAddress)
+                        -> IO HandshakeCache
 newHandshakeCache crypto send = atomically $ do
     tbl <- newTVar MM.empty
     pcs <- newTVar Map.empty
@@ -49,13 +51,13 @@ getSentHandshake :: HandshakeCache
                     -> Multi.SessionAddress
                     -> Cookie Identity -- locally issued
                     -> Cookie Encrypted -- remotely issued
-                    -> IO (Maybe (SecretKey, HandshakeData))
+                    -> IO (Maybe (Multi.SessionAddress, (SecretKey, HandshakeData)))
 getSentHandshake hscache me their_addr (Cookie n24 (Identity cd)) ecookie = do
     now <- getPOSIXTime
     io <- atomically $ do
         m <- checkExpiry now . MM.lookup' n24 <$> readTVar (hscTable hscache)
         case m of
-            Just s  -> return $ return $ Just s
+            Just s  -> return $ return $ Just (their_addr, s)
             Nothing -> do
                 let them = longTermKey cd
                 case Multi.nodeInfo (key2id $ dhtKey cd) their_addr of
@@ -64,8 +66,8 @@ getSentHandshake hscache me their_addr (Cookie n24 (Identity cd)) ecookie = do
                         (s,hs) <- cacheHandshakeSTM hscache me them their_node ecookie now
                         return $ do
                             dput XNetCrypto $ "getSentHandshake sending new handshake."
-                            hscSend hscache their_addr hs
-                            return $ Just s
+                            addr' <- hscSend hscache their_addr hs
+                            return $ Just (addr', s)
     r <- io
     dput XNetCrypto $ "getSentHandshake me="++show (key2id $ toPublic me)++" their_addr="++show their_addr++" --> " ++ show r
     return r