2 files changed, 23 insertions, 11 deletions

diff --git a/DHTTransport.hs b/DHTTransport.hs
index 013fa322..ac263cdc 100644
--- a/DHTTransport.hs
+++ b/DHTTransport.hs
@@ -307,15 +307,18 @@ forwardDHTRequests crypto closeLookup dht = dht { awaitMessage = await' }
         m -> pass m
 
 encrypt :: TransportCrypto -> DHTMessage ((,) Nonce8) -> NodeInfo -> (DHTMessage Encrypted8, NodeInfo)
-encrypt crypto msg ni = (transcode (encryptMessage crypto) msg, ni)
+encrypt crypto msg ni = ( transcode (encryptMessage crypto (id2key $ nodeId ni)) msg
+                        , ni )
 
 encryptMessage :: Serialize a =>
-                  TransportCrypto -> Nonce24 -> Either (Nonce8,a) (Assym (Nonce8,a)) -> Encrypted8 a
-encryptMessage crypto n (Right assym) = E8 $ ToxCrypto.encrypt secret plain
+                  TransportCrypto ->
+                  PublicKey ->
+                  Nonce24 -> Either (Nonce8,a) (Assym (Nonce8,a)) -> Encrypted8 a
+encryptMessage crypto destKey n (Right assym) = E8 $ ToxCrypto.encrypt secret plain
  where
-    secret = computeSharedSecret (transportSecret crypto) (senderKey assym) n
+    secret = computeSharedSecret (transportSecret crypto) destKey n
     plain  = encodePlain $ swap $ assymData assym
-encryptMessage crypto n (Left plain) = _todo -- need cached public key.
+encryptMessage crypto destKey n (Left plain) = _todo -- need cached public key.
 
 decrypt :: TransportCrypto -> DHTMessage Encrypted8 -> NodeInfo -> Either String (DHTMessage ((,) Nonce8), NodeInfo)
 decrypt crypto msg ni = (, ni) <$> (sequenceMessage $ transcode (decryptMessage crypto) msg)
diff --git a/OnionTransport.hs b/OnionTransport.hs
index 6901038d..ce1063d2 100644
--- a/OnionTransport.hs
+++ b/OnionTransport.hs
@@ -81,6 +81,10 @@ data OnionToOwner = OnionToOwner NodeInfo (ReturnPath N3)
                   | OnionToMe SockAddr -- SockAddr is immediate peer in route
     deriving Show
 
+onionKey :: OnionToOwner -> Maybe PublicKey
+onionKey (OnionToOwner ni _) = Just $ id2key (nodeId ni)
+onionKey _                   = Nothing
+
 instance Sized (OnionMessage Encrypted) where
     size = VarSize $ \case
         OnionAnnounce a -> case size of ConstSize n -> n + 1
@@ -443,17 +447,22 @@ instance Sized OnionData where
             VarSize f   -> f dhtpk
 
 encrypt :: TransportCrypto -> OnionMessage Identity -> OnionToOwner -> (OnionMessage Encrypted, OnionToOwner)
-encrypt crypto msg rpath = (transcode (encryptMessage crypto) msg, rpath)
+encrypt crypto msg rpath = ( transcode (encryptMessage crypto okey) msg
+                           , rpath)
+    where
+        -- The OnionToMe case shouldn't happen, but we'll use our own public
+        -- key in this situation.
+        okey = fromMaybe (transportPublic crypto) $ onionKey rpath
 
 encryptMessage :: Serialize a =>
-                  TransportCrypto -> Nonce24 -> Either (Identity a) (Assym (Identity a)) -> Encrypted a
-encryptMessage crypto n (Right a) = ToxCrypto.encrypt secret plain
+                  TransportCrypto -> PublicKey -> Nonce24 -> Either (Identity a) (Assym (Identity a)) -> Encrypted a
+encryptMessage crypto destKey n (Right a) = ToxCrypto.encrypt secret plain
  where
-    secret = computeSharedSecret (transportSecret crypto) (senderKey a) n
+    secret = computeSharedSecret (transportSecret crypto) destKey n
     plain  = encodePlain $ runIdentity $ assymData a
-encryptMessage crypto n (Left x) = ToxCrypto.encrypt secret plain
+encryptMessage crypto destKey n (Left x) = ToxCrypto.encrypt secret plain
  where
-    secret = computeSharedSecret (transportSecret crypto) _todo n -- OnionAnnounceResponse has no sender key
+    secret = computeSharedSecret (transportSecret crypto) destKey n
     plain  = encodePlain $ runIdentity $ x
 
 decrypt :: TransportCrypto -> OnionMessage Encrypted -> OnionToOwner -> Either String (OnionMessage Identity, OnionToOwner)