TCP Relay Server library module.

1 files changed, 235 insertions, 0 deletions

diff --git a/src/Network/Tox/Relay.hs b/src/Network/Tox/Relay.hs
new file mode 100644
index 00000000..17bbc379
--- /dev/null
+++ b/src/Network/Tox/Relay.hs
@@ -0,0 +1,235 @@
+{-# LANGUAGE LambdaCase          #-}
+{-# LANGUAGE RecordWildCards     #-}
+{-# LANGUAGE ScopedTypeVariables #-}
+module Network.Tox.Relay (tcpRelay) where
+
+import Control.Concurrent.MVar
+import Control.Concurrent.STM
+import Control.Exception
+import Control.Monad
+import qualified Data.ByteString as B
+import Data.Function
+import Data.Functor.Identity
+import qualified Data.IntMap     as IntMap
+         ;import Data.IntMap     (IntMap)
+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
+
+import Crypto.Tox
+import qualified Data.IntervalSet as IntSet
+         ;import Data.IntervalSet (IntSet)
+import Data.Tox.Relay
+import Network.Address (getBindAddress)
+import Network.SocketLike
+import Network.StreamServer
+import Network.Tox (newCrypto)
+import Network.Tox.Onion.Transport hiding (encrypt,decrypt)
+
+
+
+hGetPrefixed :: Serialize a => Handle -> IO (Either String a)
+hGetPrefixed h = do
+    mlen <- runGet getWord16be <$> B.hGet h 2
+    -- We treat parse-fail the same as EOF.
+    fmap join $ forM mlen $ \len -> runGet get <$> B.hGet h (fromIntegral len)
+
+hGetSized :: forall x. (Sized x, Serialize x) => Handle -> IO (Either String x)
+hGetSized h = runGet get <$> B.hGet h len -- We treat parse-fail the same as EOF.
+ where
+    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 ((ConId -> RelayPacket) -> IO ()) -- ^ Peers this session is connected to.
+    }
+
+freshSession :: RelaySession
+freshSession = RelaySession
+    { indexPool  = IntSet.empty
+    , solicited  = Map.empty
+    , associated = IntMap.empty
+    }
+
+disconnect :: TVar (Map PublicKey (RelayPacket -> IO (),TVar RelaySession))
+                -> PublicKey
+                -> IO ()
+disconnect cons who = join $ atomically $ do
+    Map.lookup who <$> readTVar cons
+        >>= \case
+        Nothing -> return $ return ()
+        Just (_,session) -> do
+            modifyTVar' cons $ Map.delete who
+            RelaySession { associated = cs } <- readTVar session
+            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 sendOnion _ conid h = do
+    -- atomically $ modifyTVar' cons $ IntMap.insert conid h
+
+    -- mhello <- fmap (>>= \h -> decryptPayload (computeSharedSecret me (helloFrom h) (helloNonce h)) h) $ hGetSized h
+
+    (hGetSized h >>=) $ mapM_ $ \helloE -> do
+
+    let me = transportSecret crypto
+        them = helloFrom helloE
+
+    noncef <- lookupNonceFunction crypto me them
+    let mhello = decryptPayload (noncef $ helloNonce helloE) helloE
+    forM_ mhello $ \hello -> do
+        let _ = hello :: Hello Identity
+
+        (me',welcome) <- atomically $ do
+            skey <- transportNewKey crypto
+            dta <- HelloData (toPublic skey) <$> transportNewNonce crypto
+            w24 <- transportNewNonce crypto
+            return (skey, Welcome w24 $ pure dta)
+
+        B.hPut h $ encode $ encryptPayload (noncef $ welcomeNonce welcome) welcome
+
+        noncef' <- let them' = sessionPublicKey (runIdentity $ helloData hello)
+                   in lookupNonceFunction crypto me' them'
+
+        let readPacket n24 = (>>= decrypt (noncef' n24) >=> decodePlain) <$> hGetPrefixed h
+            base = sessionBaseNonce $ runIdentity $ helloData hello
+
+        -- You get 3 seconds to send a session packet.
+        mpkt0 <- join <$> timeout 3000000 (either (const Nothing) Just <$> readPacket base)
+        forM_ mpkt0 $ \pkt0 -> do
+
+            disconnect cons (helloFrom hello)
+            (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 conid (helloFrom hello) crypto sendOnion sendPacket session p
+                    loop (incrementNonce24 n24)
+         `finally`
+            disconnect cons (helloFrom hello)
+
+handlePacket :: TVar (Map PublicKey (RelayPacket -> IO (), TVar RelaySession))
+                      -> Int
+                      -> PublicKey
+                      -> TransportCrypto
+                      -> (SockAddr -> OnionRequest N1 -> IO ())
+                      -> (RelayPacket -> IO ())
+                      -> TVar RelaySession
+                      -> RelayPacket
+                      -> IO ()
+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
+            Nothing -> fmap join $ forM (IntSet.nearestOutsider 0 (indexPool mySession)) $ \i -> do
+                if -120 <= i && i <= 119
+                    then do
+                        writeTVar session mySession
+                            { indexPool = IntSet.insert i (indexPool mySession)
+                            , solicited = Map.insert them i (solicited mySession)
+                            }
+                        return $ Just i
+                    else return Nothing -- No more slots available.
+            Just i  -> return $ Just i
+        notifyConnect <- fmap (join . join) $ forM mi $ \i -> do
+            mp <- Map.lookup them <$> readTVar cons
+            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)
+                        }
+                    writeTVar session mySession
+                        { solicited  = Map.delete them (solicited mySession)
+                        , associated = IntMap.insert i sendToThem' (associated mySession)
+                        }
+                    return $ do sendToThem' ConnectNotification
+                                sendToMe'   ConnectNotification
+        return $ do sendToMe $ RoutingResponse (maybe badcon key2c mi) them
+                    sequence_ notifyConnect
+
+    RelayPing x -> sendToMe $ RelayPong x -- TODO x==0 is invalid.  Do we care?
+
+    OOBSend them bs -> do
+        m <- atomically $ Map.lookup them <$> readTVar cons
+        forM_ m $ \(sendToThem,_) -> sendToThem $ OOBRecv me bs
+
+    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
+        -- a non connected or existent connection it will discard it.
+        mySession <- readTVar session
+        return $ sequence_ $ do
+            i <- c2key con
+            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 ()
+
+
+sendTCP_ :: TVar (IntMap (RelayPacket -> IO ())) -> Int -> OnionResponse N1 -> IO ()
+sendTCP_ st addr x = join $ atomically
+    $ IntMap.lookup addr <$> readTVar st >>= \case
+        Nothing   -> return $ return ()
+        Just send -> return $ send $ OnionPacketResponse x
+
+tcpRelay :: SockAddr -> (SockAddr -> OnionRequest N1 -> IO ()) -> IO (ServerHandle, Int -> OnionResponse N1 -> IO ())
+tcpRelay udp_addr sendOnion = do
+    crypto <- newCrypto
+    cons <- newTVarIO Map.empty
+    clients <- newTVarIO IntMap.empty
+    b443   <- getBindAddress "443"   True
+    b80    <- getBindAddress "80"    True
+    b33445 <- getBindAddress "33445" True
+    bany   <- getBindAddress ""      True
+    h <- streamServer ServerConfig
+                        { serverWarn    = hPutStrLn stderr
+                        , serverSession = relaySession crypto cons sendOnion
+                        }
+                      [b443,b80,udp_addr,b33445,bany]
+    return (h,sendTCP_ clients)
+