3 files changed, 52 insertions, 17 deletions

diff --git a/src/Network/QueryResponse.hs b/src/Network/QueryResponse.hs
index 27c89674..1901e164 100644
--- a/src/Network/QueryResponse.hs
+++ b/src/Network/QueryResponse.hs
@@ -53,24 +53,44 @@ data Transport err addr x = Transport
 -- packet representations.  It could be used to change UDP 'ByteString's into
 -- bencoded syntax trees or to add an encryption layer in which addresses have
 -- associated public keys.
-layerTransport ::
-        (x -> addr -> Either err (x', addr'))
+layerTransportM ::
+        (x -> addr -> IO (Either err (x', addr')))
         -- ^ Function that attempts to transform a low-level address/packet
         -- pair into a higher level representation.
-        -> (x' -> addr' -> (x, addr))
+        -> (x' -> addr' -> IO (x, addr))
         -- ^ Function to encode a high-level address/packet into a lower level
         -- representation.
         -> Transport err addr x
         -- ^ The low-level transport to be transformed.
         -> Transport err addr' x'
-layerTransport parse encode tr =
+layerTransportM parse encode tr =
     tr { awaitMessage = \kont ->
-                awaitMessage tr $ \m -> kont $ fmap (>>= uncurry parse) m
+                awaitMessage tr $ \m -> mapM (mapM $ uncurry parse) m >>= kont . fmap join
        , sendMessage = \addr' msg' -> do
-                let (msg,addr) = encode msg' addr'
+                (msg,addr) <- encode msg' addr'
                 sendMessage tr addr msg
        }
 
+
+-- | This function modifies a 'Transport' to use higher-level addresses and
+-- packet representations.  It could be used to change UDP 'ByteString's into
+-- bencoded syntax trees or to add an encryption layer in which addresses have
+-- associated public keys.
+layerTransport ::
+        (x -> addr -> Either err (x', addr'))
+        -- ^ Function that attempts to transform a low-level address/packet
+        -- pair into a higher level representation.
+        -> (x' -> addr' -> (x, addr))
+        -- ^ Function to encode a high-level address/packet into a lower level
+        -- representation.
+        -> Transport err addr x
+        -- ^ The low-level transport to be transformed.
+        -> Transport err addr' x'
+layerTransport parse encode tr =
+    layerTransportM (\x addr -> return $ parse x addr)
+                    (\x' addr' -> return $ encode x' addr')
+                    tr
+
 -- | Paritions a 'Transport' into two higher-level transports.  Note: An 'MVar'
 -- is used to share the same underlying socket, so be sure to fork a thread for
 -- both returned 'Transport's to avoid hanging.
diff --git a/src/Network/Tox.hs b/src/Network/Tox.hs
index 98e9691b..f41b0f25 100644
--- a/src/Network/Tox.hs
+++ b/src/Network/Tox.hs
@@ -105,6 +105,7 @@ newCrypto = do
     let pubkey    = toPublic secret
         aliaspub  = toPublic alias
         raliaspub = toPublic ralias
+    ukeys <- atomically $ newTVar []
     (symkey, drg) <- do
         drg0 <- getSystemDRG
         return $ randomBytesGenerate 32 drg0 :: IO (ByteString, SystemDRG)
@@ -125,6 +126,7 @@ newCrypto = do
             let (nonce, drg2) = withDRG drg1 (Nonce24 <$> getRandomBytes 24)
             writeTVar noncevar drg2
             return nonce
+        , userKeys = ukeys
         }
 
 updateIP :: TVar (R.BucketList NodeInfo) -> SockAddr -> STM ()
@@ -266,7 +268,7 @@ newTox keydb addr = do
         nil <- nullSessionTokens
         atomically $ newTVar nil { maxInterval = 20 } -- 20 second timeout on announce ping-ids.
     oniondrg <- drgNew
-    let onionnet = layerTransport (Onion.decrypt crypto) (Onion.encrypt crypto) onioncrypt
+    let onionnet = layerTransportM (Onion.decrypt crypto) (Onion.encrypt crypto) onioncrypt
     onionclient <- newClient oniondrg onionnet Onion.classify
                     (getOnionAlias crypto $ R.thisNode <$> readTVar (DHT.routing4 routing))
                     (Onion.handlers onionnet routing toks keydb)
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)