1 files changed, 23 insertions, 10 deletions

diff --git a/src/Network/Tox/Onion/Transport.hs b/src/Network/Tox/Onion/Transport.hs
index 3e3596a6..34ba23f6 100644
--- a/src/Network/Tox/Onion/Transport.hs
+++ b/src/Network/Tox/Onion/Transport.hs
@@ -598,18 +598,31 @@ instance Sized OnionData where
             ConstSize n -> n
             VarSize f   -> f req
 
-encrypt :: TransportCrypto -> OnionMessage Identity -> OnionDestination r -> (OnionMessage Encrypted, OnionDestination r)
-encrypt crypto msg rpath = ( transcode ( (. (runIdentity . either id assymData))
-                                         . encryptMessage skey okey)
-                                       msg
-                           , rpath)
-    where
-        skey = fst $ aliasKey crypto rpath
+
+selectKey :: TransportCrypto -> OnionMessage Identity -> OnionDestination r -> IO (SecretKey, PublicKey)
+selectKey crypto
+          (OnionAnnounce a@Assym { assymData = Identity (AnnounceRequest _ pkey akey, _) })
+          rpath
+    | (akey /= zeroID) = atomically $ do
+                            ks <- filter (\(sk,pk) -> pk == id2key pkey)
+                                    <$> readTVar (userKeys crypto)
+                            maybe (return $ aliasKey crypto rpath)
+                                  return
+                                  (listToMaybe ks)
+selectKey crypto msg rpath = return $ aliasKey crypto rpath
+
+encrypt :: TransportCrypto -> OnionMessage Identity -> OnionDestination r -> IO (OnionMessage Encrypted, OnionDestination r)
+encrypt crypto msg rpath = do
+    (skey,pkey) <- selectKey crypto msg rpath
+    let skey = fst $ aliasKey crypto rpath
 
         -- The OnionToMe case shouldn't happen, but we'll use our own public
         -- key in this situation.
         okey = fromMaybe (transportPublic crypto) $ onionKey rpath
-
+    return ( transcode ( (. (runIdentity . either id assymData))
+                                         . encryptMessage skey okey)
+                                       msg
+           , rpath)
 
 encryptMessage :: Serialize a =>
                   SecretKey -> PublicKey -> Nonce24 -> a -> Encrypted a
@@ -618,8 +631,8 @@ encryptMessage skey destKey n a = ToxCrypto.encrypt secret plain
     secret = computeSharedSecret skey destKey n
     plain  = encodePlain a
 
-decrypt :: TransportCrypto -> OnionMessage Encrypted -> OnionDestination r -> Either String (OnionMessage Identity, OnionDestination r)
-decrypt crypto msg addr = do
+decrypt :: TransportCrypto -> OnionMessage Encrypted -> OnionDestination r -> IO (Either String (OnionMessage Identity, OnionDestination r))
+decrypt crypto msg addr = return $ do
     msg <- sequenceMessage $ transcode (\n -> decryptMessage (aliasKey crypto addr) n . left (senderkey addr)) msg
     Right (msg, addr)