6 files changed, 61 insertions, 23 deletions

diff --git a/examples/dhtd.hs b/examples/dhtd.hs
index bd2df9db..1e2c1467 100644
--- a/examples/dhtd.hs
+++ b/examples/dhtd.hs
@@ -45,7 +45,7 @@ import Control.Concurrent.Lifted
 import GHC.Conc                  (labelThread)
 #endif
 
-import Crypto.Tox (zeros32,SecretKey,PublicKey, generateSecretKey, toPublic, encodeSecret, decodeSecret)
+import Crypto.Tox (zeros32,SecretKey,PublicKey, generateSecretKey, toPublic, encodeSecret, decodeSecret, userKeys)
 import Network.UPNP as UPNP
 import Network.Address hiding (NodeId, NodeInfo(..))
 import Network.Kademlia.Search
@@ -669,8 +669,9 @@ main = do
 
     keysdb <- Tox.newKeysDatabase
 
-    (quitTox,toxdhts,toxips,taddrs) <- case porttox opts of
-        "" -> return (return (), Map.empty, return [],[])
+    (toxids,quitTox,toxdhts,toxips,taddrs) <- case porttox opts of
+        "" -> do keys <- atomically $ newTVar []
+                 return (keys,return (), Map.empty, return [],[])
         toxport -> do
             addrTox <- getBindAddress toxport (ip6tox opts)
             tox <- Tox.newTox keysdb addrTox
@@ -732,7 +733,7 @@ main = do
                 ips :: IO [SockAddr]
                 ips = readExternals Tox.nodeAddr [ Tox.routing4 $ Tox.toxRouting tox
                                                  , Tox.routing6 $ Tox.toxRouting tox ]
-            return (quitTox, dhts, ips, [addrTox])
+            return (userKeys (Tox.toxCryptoKeys tox), quitTox, dhts, ips, [addrTox])
 
     _ <- UPNP.requestPorts "dht-client" $ map (Datagram,) $ baddrs ++ taddrs
 
@@ -740,14 +741,13 @@ main = do
 
     waitForSignal <- do
         signalQuit <- newEmptyMVar
-        userkeys0 <- atomically (newTVar [])
         let session = clientSession $ Session
                 { netname           = concat $ take 1 $ Map.keys dhts -- initial default DHT
                 , dhts              = dhts  -- all DHTs
                 , signalQuit        = signalQuit
                 , swarms            = swarms
                 , toxkeys           = keysdb
-                , userkeys          = userkeys0
+                , userkeys          = toxids
                 , externalAddresses = liftM2 (++) btips toxips
                 }
         srv <- streamServer (withSession session) (SockAddrUnix "dht.sock")
diff --git a/src/Crypto/Tox.hs b/src/Crypto/Tox.hs
index b86a5395..645ca53e 100644
--- a/src/Crypto/Tox.hs
+++ b/src/Crypto/Tox.hs
@@ -366,6 +366,7 @@ data TransportCrypto = TransportCrypto
     , rendezvousPublic :: PublicKey
     , transportSymmetric :: STM SymmetricKey
     , transportNewNonce :: STM Nonce24
+    , userKeys :: TVar [(SecretKey,PublicKey)]
     }
 
 getPublicKey :: S.Get PublicKey
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)
 
diff --git a/todo.txt b/todo.txt
index 1a3bf8b4..287f276a 100644
--- a/todo.txt
+++ b/todo.txt
@@ -13,6 +13,8 @@ tox4: 6 buckets, tox6: 3 buckets
 
 handle exception: dhtd: Network.Socket.sendTo: does not exist (Network is unreachable
 
+tox: rename assym to asymm (short for asymmetric)
+
 tox: fallback to https://nodes.tox.chat/json
 
 tox: key search