path: root/dht/src/Network/Tox/Relay.hs

{-# 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.Onion.Transport hiding (encrypt,decrypt)

import DPut
import DebugTag


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

-- This type manages ConId assignments.
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 (IntMap (RelayPacket -> IO ()))
                -> TVar (Map PublicKey (RelayPacket -> IO (),TVar RelaySession))
                -> (SockAddr -> OnionRequest N1 -> IO ())
                -> sock
                -> Int
                -> Handle
                -> IO ()
relaySession crypto clients 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

    dput XOnion $ "Relay client session conid=" ++ show conid
    (hGetSized h >>=) $ mapM_ $ \helloE -> do

    let me = transportSecret crypto
        them = helloFrom helloE

    noncef <- lookupNonceFunction crypto me them
    let mhello = decryptPayload (noncef $ helloNonce helloE) helloE
    dput XOnion $ "Relay client (conid=" ++ show conid ++ ") decrypted hello = " ++ show mhello
    forM_ mhello $ \hello -> do
        let _ = hello :: Hello Identity

        dput XOnion $ "Relay client sent hello. conid=" ++ show conid
        (me',welcome) <- atomically $ do
            skey <- transportNewKey crypto
            dta <- HelloData (toPublic skey) <$> transportNewNonce crypto
            w24 <- transportNewNonce crypto
            return (skey, Welcome w24 $ pure dta)

        dput XOnion $ unlines [ "Relay client to receive welcome. conid=" ++ show conid
                              , show welcome
                              ]
        B.hPut h $ encode $ encryptPayload (noncef $ welcomeNonce welcome) welcome

        dput XOnion $ "Relay client welcome sent. conid=" ++ show conid

        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

            atomically $ modifyTVar' clients $ IntMap.insert conid $
                \p -> do
                    dput XOnion $ "Sending onion reply to TCP client conid="++show conid
                    sendPacket p

            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` do
            atomically $ modifyTVar' clients $ IntMap.delete conid
            disconnect cons (helloFrom hello)
            dput XOnion $ "Relay client session closed. conid=" ++ show conid

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 n24 (Addressed addr req) -> do
        dput XOnion $ "Received onion request via TCP client conid="++show thistcp
        rpath <- atomically $ do
                    sym <- transportSymmetric crypto
                    n <- transportNewNonce crypto
                    return $ wrapSymmetric sym n (TCPIndex thistcp) NoReturnPath
        sendOnion addr $ OnionRequest n24 req rpath

    _ -> return ()


sendTCP_ :: TVar (IntMap (RelayPacket -> IO ())) -> Int -> OnionMessage Encrypted -> IO ()
sendTCP_ st addr x = join $ atomically
    $ IntMap.lookup addr <$> readTVar st >>= \case
        Nothing   -> return $ return ()
        Just send -> return $ send $ OnionPacketResponse x

tcpRelay :: TransportCrypto
            -> SockAddr -- ^ UDP bind address (this port may be tried for TCP if hardcoded defaults dont work).
            -> (SockAddr -> OnionRequest N1 -> IO ()) -- ^ This callback will be used to forward onion messages over udp.
            -> IO ( ServerHandle -- Handle to the Tox Tcp-Relay server.
                  , Int -> OnionMessage Encrypted -> IO () -- forward onion response to tcp client.
                  )
tcpRelay crypto udp_addr sendOnion = do
    cons <- newTVarIO Map.empty
    clients <- newTVarIO IntMap.empty
    b443   <- getBindAddress "443"   True
    b80    <- getBindAddress "80"    True
    b3389  <- getBindAddress "3389"  True
    b33445 <- getBindAddress "33445" True
    bany   <- getBindAddress ""      True
    h <- streamServer ServerConfig
                        { serverWarn    = hPutStrLn stderr
                        , serverSession = relaySession crypto clients cons sendOnion
                        }
                      [b443,b80,b3389,udp_addr,b33445,bany]
    return (h,sendTCP_ clients)