1 files changed, 15 insertions, 6 deletions

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)