1 files changed, 54 insertions, 37 deletions

diff --git a/examples/toxrelay.hs b/examples/toxrelay.hs
index f03605f9..953b230b 100644
--- a/examples/toxrelay.hs
+++ b/examples/toxrelay.hs
@@ -15,6 +15,7 @@ import qualified Data.Map        as Map
          ;import Data.Map        (Map)
 import Data.Serialize
 import Data.Word
+import Network.Socket (SockAddr)
 import System.IO
 import System.IO.Error
 import System.Timeout
@@ -26,6 +27,7 @@ import Data.Tox.Relay
 import Network.Address (getBindAddress)
 import Network.StreamServer
 import Network.Tox (newCrypto)
+import Network.Tox.Onion.Transport hiding (encrypt,decrypt)
 
 
 
@@ -41,9 +43,9 @@ hGetSized h = runGet get <$> B.hGet h len -- We treat parse-fail the same as EOF
     ConstSize len = size :: Size x
 
 data RelaySession = RelaySession
-    { indexPool  :: IntSet                        -- ^ Ints that are either solicited or associated.
-    , solicited  :: Map PublicKey Int             -- ^ Reserved ids, not yet in associated.
-    , associated :: IntMap (RelayPacket -> IO ()) -- ^ Peers this session is connected to.
+    { indexPool  :: IntSet                                   -- ^ Ints that are either solicited or associated.
+    , solicited  :: Map PublicKey Int                        -- ^ Reserved ids, not yet in associated.
+    , associated :: IntMap ((ConId -> RelayPacket) -> IO ()) -- ^ Peers this session is connected to.
     }
 
 freshSession :: RelaySession
@@ -63,16 +65,17 @@ disconnect cons who = join $ atomically $ do
         Just (_,session) -> do
             modifyTVar' cons $ Map.delete who
             RelaySession { associated = cs } <- readTVar session
-            return $ let notifyPeer i send = (send (DisconnectNotification $ key2c i) >>)
+            return $ let notifyPeer i send = ((send DisconnectNotification) >>)
                      in IntMap.foldrWithKey notifyPeer (return ()) cs
 
 relaySession :: TransportCrypto
                 -> TVar (Map PublicKey (RelayPacket -> IO (),TVar RelaySession))
+                -> (SockAddr -> OnionRequest N1 -> IO ())
                 -> sock
                 -> Int
                 -> Handle
                 -> IO ()
-relaySession crypto cons _ conid h = do
+relaySession crypto cons sendOnion _ conid h = do
     -- atomically $ modifyTVar' cons $ IntMap.insert conid h
 
     -- mhello <- fmap (>>= \h -> decryptPayload (computeSharedSecret me (helloFrom h) (helloNonce h)) h) $ hGetSized h
@@ -98,16 +101,6 @@ relaySession crypto cons _ conid h = do
         noncef' <- let them' = sessionPublicKey (runIdentity $ helloData hello)
                    in lookupNonceFunction crypto me' them'
 
-        sendPacket <- do
-            v <- newMVar (sessionBaseNonce $ runIdentity $ welcomeData welcome)
-            return $ \p -> do
-                n24 <- takeMVar v
-                let bs = encode $ encrypt (noncef' n24) $ encodePlain (p :: RelayPacket)
-                do B.hPut h $ encode (fromIntegral (B.length bs) :: Word16)
-                   B.hPut h bs
-                 `catchIOError` \_ -> return ()
-                putMVar v (incrementNonce24 n24)
-
         let readPacket n24 = (>>= decrypt (noncef' n24) >=> decodePlain) <$> hGetPrefixed h
             base = sessionBaseNonce $ runIdentity $ helloData hello
 
@@ -116,33 +109,47 @@ relaySession crypto cons _ conid h = do
         forM_ mpkt0 $ \pkt0 -> do
 
             disconnect cons (helloFrom hello)
-            session <- atomically $ do
-                session <- newTVar freshSession
-                modifyTVar' cons $ Map.insert (helloFrom hello) (sendPacket,session)
-                return session
-
-            handlePacket cons (helloFrom hello) sendPacket session pkt0
+            (sendPacket,session) <- do
+                session <- atomically $ newTVar freshSession
+                sendPacket <- do
+                    v <- newMVar (sessionBaseNonce $ runIdentity $ welcomeData welcome)
+                    return $ \p -> do
+                        case p of
+                            DisconnectNotification con -> atomically $ do
+                                modifyTVar' session $ \s -> s
+                                    { indexPool = maybe id IntSet.delete (c2key con) (indexPool s)
+                                    , associated = maybe id IntMap.delete (c2key con) (associated s)
+                                    }
+                            _ -> return ()
+                        n24 <- takeMVar v
+                        let bs = encode $ encrypt (noncef' n24) $ encodePlain (p :: RelayPacket)
+                        do B.hPut h $ encode (fromIntegral (B.length bs) :: Word16)
+                           B.hPut h bs
+                         `catchIOError` \_ -> return ()
+                        putMVar v (incrementNonce24 n24)
+                atomically $ modifyTVar' cons $ Map.insert (helloFrom hello) (sendPacket,session)
+                return (sendPacket,session)
+
+            handlePacket cons conid (helloFrom hello) crypto sendOnion sendPacket session pkt0
 
             flip fix (incrementNonce24 base)  $ \loop n24 -> do
                 m <- readPacket n24
                 forM_ m $ \p -> do
-                    handlePacket cons (helloFrom hello) sendPacket session p
+                    handlePacket cons conid (helloFrom hello) crypto sendOnion sendPacket session p
                     loop (incrementNonce24 n24)
          `finally`
             disconnect cons (helloFrom hello)
 
-data R = R { routingRequest :: PublicKey -> IO ConId
-           , reply :: RelayPacket -> IO ()
-           , routeOOB :: PublicKey -> IO (Maybe (RelayPacket -> IO ()))
-           }
-
 handlePacket :: TVar (Map PublicKey (RelayPacket -> IO (), TVar RelaySession))
+                      -> Int
                       -> PublicKey
+                      -> TransportCrypto
+                      -> (SockAddr -> OnionRequest N1 -> IO ())
                       -> (RelayPacket -> IO ())
                       -> TVar RelaySession
                       -> RelayPacket
                       -> IO ()
-handlePacket cons me sendToMe session = \case
+handlePacket cons thistcp me crypto sendOnion sendToMe session = \case
     RoutingRequest them -> join $ atomically $ do
         mySession <- readTVar session
         mi <- case Map.lookup them (solicited mySession) of
@@ -161,16 +168,18 @@ handlePacket cons me sendToMe session = \case
             forM mp $ \(sendToThem,peer) -> do
                 theirSession <- readTVar peer
                 forM (Map.lookup me $ solicited theirSession) $ \reserved_id -> do
+                    let sendToThem' f = sendToThem $ f $ key2c reserved_id
+                        sendToMe'   f = sendToMe   $ f $ key2c i
                     writeTVar peer theirSession
                         { solicited  = Map.delete me (solicited theirSession)
-                        , associated = IntMap.insert reserved_id sendToMe (associated theirSession)
+                        , associated = IntMap.insert reserved_id sendToMe' (associated theirSession)
                         }
                     writeTVar session mySession
                         { solicited  = Map.delete them (solicited mySession)
-                        , associated = IntMap.insert i sendToThem (associated mySession)
+                        , associated = IntMap.insert i sendToThem' (associated mySession)
                         }
-                    return $ do sendToThem $ ConnectNotification (key2c reserved_id)
-                                sendToMe   $ ConnectNotification (key2c i)
+                    return $ do sendToThem' ConnectNotification
+                                sendToMe'   ConnectNotification
         return $ do sendToMe $ RoutingResponse (maybe badcon key2c mi) them
                     sequence_ notifyConnect
 
@@ -180,7 +189,7 @@ handlePacket cons me sendToMe session = \case
         m <- atomically $ Map.lookup them <$> readTVar cons
         forM_ m $ \(sendToThem,_) -> sendToThem $ OOBRecv me bs
 
-    RelayData con bs -> join $ atomically $ do
+    RelayData bs con -> join $ atomically $ do
         -- Data: Data packets can only be sent and received if the
         -- corresponding connection_id is connection (a Connect notification
         -- has been received from it) if the server receives a Data packet for
@@ -188,8 +197,14 @@ handlePacket cons me sendToMe session = \case
         mySession <- readTVar session
         return $ sequence_ $ do
             i <- c2key con
-            sendToThem <- IntMap.lookup i $ associated mySession
-            return $ sendToThem $ RelayData _todo bs
+            sendToThem' <- IntMap.lookup i $ associated mySession
+            return $ sendToThem' $ RelayData bs
+
+    OnionPacket p -> do
+        mp <- rewrap crypto (TCPIndex thistcp) p
+        case mp of
+            Right (p',addr) -> sendOnion addr p'
+            _               -> return ()
 
     _ -> return ()
 
@@ -199,9 +214,11 @@ main = do
     crypto <- newCrypto
     cons <- newTVarIO Map.empty
     a <- getBindAddress "33445" True
+    let sendOnion :: SockAddr -> OnionRequest N1 -> IO ()
+        sendOnion _ _ = return ()
     h <- streamServer ServerConfig
-                        { serverWarn = hPutStrLn stderr
-                        , serverSession = relaySession crypto cons
+                        { serverWarn    = hPutStrLn stderr
+                        , serverSession = relaySession crypto cons sendOnion
                         }
                       a