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.Time.Clock.POSIX
import Data.Tox.Relay
import qualified Data.Wrapper.PSQInt as Int
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 are assigned.
    , assigned   :: Map PublicKey Int                        -- ^ Assignments
    , associated :: IntMap ((ConId -> RelayPacket) -> IO ()) -- ^ Peers this session is connected to.
    -- TODO: Timestamp PSQ for reclaiming indices.  The Bool will indicate
    -- whether the index was ever associated.
    , timestamps :: Int.PSQ (Bool,POSIXTime)
    }

freshSession :: RelaySession
freshSession = RelaySession
    { indexPool  = IntSet.empty
    , assigned   = Map.empty
    , associated = IntMap.empty
    , timestamps = Int.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 _ thistcp h = do
    -- atomically $ modifyTVar' cons $ IntMap.insert conid h

    -- mhello <- fmap (>>= \h -> decryptPayload (computeSharedSecret me (helloFrom h) (helloNonce h)) h) $ hGetSized h

    dput XRelay $ "Relay client session tcp=" ++ show thistcp
    (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

        -- dput XRelay $ "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)

        B.hPut h $ encode $ encryptPayload (noncef $ welcomeNonce welcome) welcome
        dput XRelay $ unlines
            [ "Relay welcomes (tcp=" ++ show thistcp ++ ") " ++ showKey256 them
            -- , "   hello=" ++ show hello
            -- , " welcome=" ++ show 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 thistcp (helloFrom hello) crypto sendOnion sendPacket session pkt0

            atomically $ modifyTVar' clients $ IntMap.insert thistcp $
                \p -> do
                    dput XOnion $ unlines
                        [ "Sending onion reply to TCP client tcp="++show thistcp
                        , " pkt0=" ++ show pkt0
                        ]
                    sendPacket p

            flip fix (incrementNonce24 base)  $ \loop n24 -> do
                m <- readPacket n24
                forM_ m $ \p -> do
                    handlePacket cons thistcp (helloFrom hello) crypto sendOnion sendPacket session p
                    loop (incrementNonce24 n24)
         `finally` do
            atomically $ modifyTVar' clients $ IntMap.delete thistcp
            disconnect cons (helloFrom hello)
            dput XRelay $ "Relay client session closed. tcp=" ++ show thistcp

handlePacket :: TVar (Map PublicKey (RelayPacket -> IO (), TVar RelaySession)) -- ^ All sessions.
                      -> Int -- ^ TCP client number.
                      -> PublicKey -- ^ Public key of client.
                      -> TransportCrypto
                      -> (SockAddr -> OnionRequest N1 -> IO ()) -- ^ Forward onion packet.
                      -> (RelayPacket -> IO ()) -- ^ Send to this client.
                      -> TVar RelaySession -- ^ Session for this client.
                      -> RelayPacket
                      -> IO ()
handlePacket cons thistcp thisKey crypto sendOnion sendToClient session = \case
    RoutingRequest them -> join $ atomically $ do
        mySession <- readTVar session
        -- TODO: Do we already have an association?
        mi <- case Map.lookup them (assigned  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)
                            , assigned  = Map.insert them i (assigned  mySession)
                            }
                        return $ Just i
                    else
                        -- TODO: Attempt to reclaim an assigned but not yet associated connection.
                        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 thisKey $ assigned  theirSession) $ \reserved_id -> do
                    let sendToThem'   f = sendToThem $ f $ key2c reserved_id
                        sendToClient' f = sendToClient $ f $ key2c i
                    writeTVar peer theirSession
                        { -- assigned   = Map.insert thisKey reserved_id (assigned theirSession)
                          associated = IntMap.insert reserved_id sendToClient' (associated theirSession)
                        }
                    writeTVar session mySession
                        { -- assigned   = Map.insert them i (assigned mySession)
                          associated = IntMap.insert i sendToThem' (associated mySession)
                        }
                    return $ do
                        let showSession n k = "("++ show (key2c n) ++ ")" ++ showKey256 k
                        dput XRelay $
                            "Relay session " ++ showSession reserved_id thisKey
                                 ++ " <--> " ++ showSession i them
                        sendToThem' ConnectNotification
                        sendToClient' ConnectNotification
        return $ do sendToClient $ RoutingResponse (maybe badcon key2c mi) them
                    sequence_ notifyConnect

    RelayPing x -> sendToClient $ RelayPong x -- TODO x==0 is invalid.  Do we care?

    OOBSend them bs -> do
        dput XRelayVerbose $ "OOB send from " ++ showKey256 thisKey ++ " to " ++ showKey256 them
        m <- atomically $ Map.lookup them <$> readTVar cons
        forM_ m $ \(sendToThem,_) -> sendToThem $ OOBRecv thisKey bs

    RelayData bs con -> do
        mySession <- atomically $ readTVar session
        -- 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.
        let mbSendIt = do
                i <- c2key con
                sendToThem' <- IntMap.lookup i $ associated mySession
                return $ sendToThem' $ RelayData bs
        dput XRelayVerbose $ "RelayData from " ++ showKey256 thisKey ++ " to conid="
                      ++ show con ++ maybe " (no key)" (\io -> " (associated key)") mbSendIt
        sequence_ mbSendIt

    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 <- forkStreamServer ServerConfig
                        { serverWarn    = dput XMisc
                        , serverSession = \s n h -> relaySession crypto clients cons sendOnion s n h
                                            `catchIOError` \e -> do
                                                dput XRelay $ "relaySession died: " ++ show e
                        }
                      [b443,b80,b3389,udp_addr,b33445,bany]
    return (h,sendTCP_ clients)