path: root/src/Network/Tox/DHT/Handlers.hs

{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE PatternSynonyms            #-}
{-# LANGUAGE TupleSections              #-}
{-# LANGUAGE CPP                        #-}
module Network.Tox.DHT.Handlers where

import Network.Tox.DHT.Transport as DHTTransport
import Network.QueryResponse as QR hiding (Client)
import qualified Network.QueryResponse as QR (Client)
import Crypto.Tox
import Network.Kademlia.Search
import qualified Data.Wrapper.PSQInt as Int
import Network.Kademlia
import Network.Address               (WantIP (..), ipFamily, testIdBit,fromSockAddr, sockAddrPort)
import qualified Network.Kademlia.Routing as R
import Control.TriadCommittee
import System.Global6
import OnionRouter

import qualified Data.ByteArray as BA
import qualified Data.ByteString.Char8 as C8
import qualified Data.ByteString.Base16 as Base16
import Control.Arrow
import Control.Monad
import Control.Concurrent.Lifted.Instrument
import Control.Concurrent.STM
import Data.Time.Clock.POSIX  (POSIXTime, getPOSIXTime)
import Network.Socket
import Data.Hashable
import Data.IP
import Data.Ord
import Data.Maybe
import Data.Bits
import Data.Serialize (Serialize)
import Data.Word
import Data.List
import System.IO

data TransactionId = TransactionId
 { transactionKey :: Nonce8  -- ^ Used to lookup pending query.
 , cryptoNonce    :: Nonce24 -- ^ Used during the encryption layer.
 }
 deriving (Eq,Ord,Show)

newtype PacketKind = PacketKind Word8
 deriving (Eq, Ord, Serialize)

pattern OnionRequest0Type = PacketKind 128 -- 0x80     Onion Request 0
pattern OnionRequest1Type = PacketKind 129 -- 0x81     Onion Request 1
pattern OnionRequest2Type = PacketKind 130 -- 0x82     Onion Request 2
pattern AnnounceType  = PacketKind 131 -- 0x83     Announce Request
pattern AnnounceResponseType = PacketKind 132 -- 0x84     Announce Response

pattern DataRequestType = PacketKind 133 -- 0x85     Onion Data Request (data to route request packet)
pattern DataResponseType = PacketKind 134 -- 0x86     Onion Data Response (data to route response packet)
-- 0x8c     Onion Response 3
-- 0x8d     Onion Response 2
pattern OnionResponse3Type = PacketKind 140 -- 0x8c     Onion Response 3
pattern OnionResponse2Type = PacketKind 141 -- 0x8d     Onion Response 2
pattern OnionResponse1Type = PacketKind 142 -- 0x8e     Onion Response 1
-- 0xf0     Bootstrap Info

pattern DHTRequestType = PacketKind 32 -- 0x20     DHT Request

pattern CookieRequestType = PacketKind 0x18
pattern CookieResponseType = PacketKind 0x19

pattern PingType      = PacketKind 0 -- 0x00     Ping Request
pattern PongType      = PacketKind 1 -- 0x01     Ping Response
pattern GetNodesType  = PacketKind 2 -- 0x02     Nodes Request
pattern SendNodesType = PacketKind 4 -- 0x04     Nodes Response


instance Show PacketKind where
    showsPrec d PingType = mappend "PingType"
    showsPrec d PongType = mappend "PongType"
    showsPrec d GetNodesType = mappend "GetNodesType"
    showsPrec d SendNodesType = mappend "SendNodesType"
    showsPrec d DHTRequestType = mappend "DHTRequestType"
    showsPrec d OnionRequest0Type = mappend "OnionRequest0Type"
    showsPrec d OnionResponse1Type = mappend "OnionResponse1Type"
    showsPrec d OnionResponse3Type = mappend "OnionResponse3Type"
    showsPrec d AnnounceType = mappend "AnnounceType"
    showsPrec d AnnounceResponseType = mappend "AnnounceResponseType"
    showsPrec d DataRequestType = mappend "DataRequestType"
    showsPrec d DataResponseType = mappend "DataResponseType"
    showsPrec d CookieRequestType = mappend "CookieRequestType"
    showsPrec d CookieResponseType = mappend "CookieResponseType"
    showsPrec d (PacketKind x) = mappend "PacketKind " . showsPrec (d+1) x

msgType :: ( Serialize (f DHTRequest)
           , Serialize (f Cookie), Serialize (f CookieRequest)
           , Serialize (f SendNodes), Serialize (f GetNodes)
           , Serialize (f Pong), Serialize (f Ping)
           ) => DHTMessage f -> PacketKind
msgType msg = PacketKind $ fst $ dhtMessageType msg

classify :: Client -> Message -> MessageClass String PacketKind TransactionId NodeInfo Message
classify client (DHTLanDiscovery {}) = IsUnsolicited (lanDiscoveryH client)
classify client msg = fromMaybe (IsUnknown "unknown")
                      $ mapMessage (\nonce24 (nonce8,_) -> go msg (TransactionId nonce8 nonce24)) msg
 where
    go (DHTPing          {}) = IsQuery PingType
    go (DHTGetNodes      {}) = IsQuery GetNodesType
    go (DHTPong          {}) = IsResponse
    go (DHTSendNodes     {}) = IsResponse
    go (DHTCookieRequest {}) = IsQuery CookieRequestType
    go (DHTCookie        {}) = IsResponse
    go (DHTDHTRequest    {}) = IsQuery DHTRequestType

data Routing = Routing
    { tentativeId :: NodeInfo
    , sched4      :: !( TVar (Int.PSQ POSIXTime) )
    , routing4    :: !( TVar (R.BucketList NodeInfo) )
    , committee4  :: TriadCommittee NodeId SockAddr
    , sched6      :: !( TVar (Int.PSQ POSIXTime) )
    , routing6    :: !( TVar (R.BucketList NodeInfo) )
    , committee6  :: TriadCommittee NodeId SockAddr
    }

newRouting :: SockAddr -> TransportCrypto
                       -> (TVar (R.BucketList NodeInfo) -> SockAddr -> STM ()) -- ^ invoked on IPv4 change
                       -> (TVar (R.BucketList NodeInfo) -> SockAddr -> STM ()) -- ^ invoked on IPv6 change
                       -> IO Routing
newRouting addr crypto update4 update6 = do
    let tentative_ip4 = fromMaybe (IPv4 $ toEnum 0) (IPv4 <$> fromSockAddr addr)
        tentative_ip6 = fromMaybe (IPv6 $ toEnum 0) (IPv6 <$> fromSockAddr addr)
        tentative_info = NodeInfo
                { nodeId   = key2id $ transportPublic crypto
                , nodeIP   = fromMaybe (toEnum 0) (fromSockAddr addr)
                , nodePort = fromMaybe 0 $ sockAddrPort addr
                }
        tentative_info4 = tentative_info { nodeIP = tentative_ip4 }
    tentative_info6 <-
        maybe (tentative_info { nodeIP = tentative_ip6 })
              (\ip6 -> tentative_info { nodeIP = IPv6 ip6 })
            <$> case addr of
                    SockAddrInet {} -> return Nothing
                    _               -> global6
    atomically $ do
        let nobkts = R.defaultBucketCount :: Int
        tbl4 <- newTVar $ R.nullTable (comparing nodeId) (\s -> hashWithSalt s . nodeId) tentative_info4 nobkts
        tbl6 <- newTVar $ R.nullTable (comparing nodeId) (\s -> hashWithSalt s . nodeId) tentative_info6 nobkts
        committee4 <- newTriadCommittee (update4 tbl4) -- updateIPVote tbl4 addr4
        committee6 <- newTriadCommittee (update6 tbl6) -- updateIPVote tbl6 addr6
        sched4 <- newTVar Int.empty
        sched6 <- newTVar Int.empty
        return $ Routing tentative_info sched4 tbl4 committee4 sched6 tbl6 committee6


-- TODO: This should cover more cases
isLocal :: IP -> Bool
isLocal (IPv6 ip6) = (ip6 == toEnum 0)
isLocal (IPv4 ip4) = (ip4 == toEnum 0)

isGlobal :: IP -> Bool
isGlobal = not . isLocal

prefer4or6 :: NodeInfo -> Maybe WantIP -> WantIP
prefer4or6 addr iptyp = fromMaybe (ipFamily $ nodeIP addr) iptyp

toxSpace :: R.KademliaSpace NodeId NodeInfo
toxSpace = R.KademliaSpace
    { R.kademliaLocation = nodeId
    , R.kademliaTestBit  = testNodeIdBit
    , R.kademliaXor      = xorNodeId
    , R.kademliaSample   = sampleNodeId
    }


pingH :: NodeInfo -> Ping -> IO Pong
pingH _ Ping = return Pong

getNodesH :: Routing -> NodeInfo -> GetNodes -> IO SendNodes
getNodesH routing addr (GetNodes nid) = do
    let preferred = prefer4or6 addr Nothing

    (append4,append6) <- atomically $ do
        ni4 <- R.thisNode <$> readTVar (routing4 routing)
        ni6 <- R.thisNode <$> readTVar (routing6 routing)
        return $ case ipFamily (nodeIP addr) of
            Want_IP4 | isGlobal (nodeIP ni6) -> (id, (++ [ni6]))
            Want_IP6 | isGlobal (nodeIP ni4) -> ((++ [ni4]), id)
            _                                -> (id, id)
    ks  <- go append4 $ routing4 routing
    ks6 <- go append6 $ routing6 routing
    let (ns1,ns2) = case preferred of Want_IP6 -> (ks6,ks)
                                      Want_IP4 -> (ks,ks6)
                                      Want_Both -> error $ "BUG:unreachable at " ++ __FILE__ ++ ":" ++ show __LINE__
    return $ SendNodes
           $ if null ns2 then ns1
                         else take 4 (take 3 ns1 ++ ns2)
 where
    go f var = f . R.kclosest toxSpace k nid <$> atomically (readTVar var)

    k = 4

cookieRequestH :: TransportCrypto -> NodeInfo -> CookieRequest -> IO Cookie
cookieRequestH crypto ni (CookieRequest remoteUserKey) = do
    (n24,sym,us) <- atomically $ do
        n24 <- transportNewNonce crypto
        sym <- transportSymmetric crypto
        us <- readTVar $ userKeys crypto
        return (n24,sym,us)
    timestamp <- round . (* 1000000) <$> getPOSIXTime
    let dta = encodePlain $ CookieData
            { cookieTime  = timestamp
            , longTermKey = remoteUserKey
            , dhtKey      = transportPublic crypto
            }
        edta = encryptSymmetric sym n24 dta
    return $ Cookie n24 edta

lanDiscoveryH :: Client -> NodeInfo -> NodeInfo -> IO (Maybe (Message -> Message))
lanDiscoveryH client _ ni = do
    forkIO $ do
        myThreadId >>= flip labelThread "lan-discover-ping"
        ping client ni
        return ()
    return Nothing

type Message = DHTMessage ((,) Nonce8)

type Client = QR.Client String PacketKind TransactionId NodeInfo Message


wrapAsymm :: TransactionId -> NodeInfo -> NodeInfo -> (Nonce8 -> dta) -> Asymm dta
wrapAsymm (TransactionId n8 n24) src dst dta = Asymm
    { senderKey = id2key $ nodeId src
    , asymmNonce = n24
    , asymmData = dta n8
    }

serializer :: PacketKind
                -> (Asymm (Nonce8,ping) -> Message)
                -> (Message -> Maybe (Asymm (Nonce8,pong)))
                -> MethodSerializer TransactionId NodeInfo Message PacketKind ping (Maybe pong)
serializer pktkind mkping mkpong = MethodSerializer
        { methodTimeout  = \tid addr -> return (addr, 5000000)
        , method         = pktkind
        -- wrapQuery :: tid -> addr -> addr -> qry -> x
        , wrapQuery      = \tid src dst ping -> mkping $ wrapAsymm tid src dst (, ping)
        -- unwrapResponse :: x -> b
        , unwrapResponse = fmap (snd . asymmData) . mkpong
        }


unpong :: Message -> Maybe (Asymm (Nonce8,Pong))
unpong (DHTPong asymm) = Just asymm
unpong _               = Nothing

showHex :: BA.ByteArrayAccess ba => ba -> String
showHex bs = C8.unpack $ Base16.encode $ BA.convert bs

ping :: Client -> NodeInfo -> IO Bool
ping client addr = do
    hPutStrLn stderr $ show addr ++ " <-- ping"
    reply <- QR.sendQuery client (serializer PingType DHTPing unpong) Ping addr
    hPutStrLn stderr $ show addr ++ " -pong-> " ++ show reply
    maybe (return False) (\Pong -> return True) $ join reply


saveCookieKey :: TVar [(SockAddr, (Int, PublicKey))] -> SockAddr -> PublicKey -> STM ()
saveCookieKey var saddr pk = do
    cookiekeys <- readTVar var
    case break (\(stored,_) -> stored == saddr) cookiekeys of
        (xs,[])                               -> writeTVar var $ (saddr, (1  ,pk)) : xs
        (xs,(_,(c,stored)):ys) | stored == pk -> writeTVar var $ (saddr, (c+1,pk)) : xs ++ ys
        _                                     -> retry -- Wait for requests to this address
                                                       -- under a different key to time out
                                                       -- before we try this key.

loseCookieKey :: TVar [(SockAddr, (Int, PublicKey))] -> SockAddr -> PublicKey -> STM ()
loseCookieKey var saddr pk = do
    cookiekeys <- readTVar var
    case break (\(stored,_) -> stored == saddr) cookiekeys of
        (xs,(_,(1,stored)):ys) | stored == pk -> writeTVar var $                     xs ++ ys
        (xs,(_,(c,stored)):ys) | stored == pk -> writeTVar var $ (saddr, (c-1,pk)) : xs ++ ys
        _                                     -> return () -- unreachable?


cookieRequest :: TransportCrypto -> Client -> PublicKey -> NodeInfo -> IO (Maybe Cookie)
cookieRequest crypto client localUserKey addr = do
    let sockAddr = nodeAddr addr
        nid = id2key $ nodeId addr
        cookieSerializer
            = MethodSerializer
                   { methodTimeout  = \tid addr -> return (addr, 5000000)
                   , method    = CookieRequestType
                   , wrapQuery = \tid src dst cr -> DHTCookieRequest $ wrapAsymm tid src dst (, cr)
                   , unwrapResponse = fmap snd . unCookie
                   }
        cookieRequest = CookieRequest localUserKey
    atomically $ saveCookieKey (pendingCookies crypto) sockAddr nid
    hPutStrLn stderr $ show addr ++ " <-- cookieRequest"
    reply <- QR.sendQuery client cookieSerializer cookieRequest addr
    atomically $ loseCookieKey (pendingCookies crypto) sockAddr nid
    hPutStrLn stderr $ show addr ++ " -cookieResponse-> " ++ show reply
    return $ join reply

unCookie :: DHTMessage t -> Maybe (t Cookie)
unCookie (DHTCookie n24 fcookie) = Just fcookie
unCookie _ = Nothing

unsendNodes :: Message -> Maybe (Asymm (Nonce8,SendNodes))
unsendNodes (DHTSendNodes asymm) = Just asymm
unsendNodes _                    = Nothing

unwrapNodes :: SendNodes -> ( [NodeInfo], [NodeInfo], Maybe () )
unwrapNodes (SendNodes ns) = (ns,ns,Just ())

getNodes :: Client -> NodeId -> NodeInfo -> IO (Maybe ([NodeInfo],[NodeInfo],Maybe ()))
getNodes client nid addr = do
    -- hPutStrLn stderr $ show addr ++ " <-- getnodes " ++ show nid
    reply <- QR.sendQuery client (serializer GetNodesType DHTGetNodes unsendNodes) (GetNodes nid) addr
    -- hPutStrLn stderr $ show addr ++ " -sendnodes-> " ++ show reply
    return $ fmap unwrapNodes $ join reply

updateRouting :: Client -> Routing -> OnionRouter -> NodeInfo -> Message -> IO ()
updateRouting client routing orouter naddr msg
    | PacketKind 0x21 <- msgType msg = return () -- ignore lan discovery
    | otherwise                      = do
    case prefer4or6 naddr Nothing of
        Want_IP4 -> updateTable client naddr orouter (routing4 routing) (committee4 routing) (sched4 routing)
        Want_IP6 -> updateTable client naddr orouter (routing6 routing) (committee6 routing) (sched6 routing)
        Want_Both -> error $ "BUG:unreachable at " ++ __FILE__ ++ ":" ++ show __LINE__

updateTable :: Client -> NodeInfo -> OnionRouter -> TVar (R.BucketList NodeInfo) -> TriadCommittee NodeId SockAddr -> TVar (Int.PSQ POSIXTime) -> IO ()
updateTable client naddr orouter tbl committee sched = do
        self <- atomically $ R.thisNode <$> readTVar tbl
        when (nodeIP self /= nodeIP naddr) $ do
            -- TODO: IP address vote?
            insertNode (toxKademlia client committee orouter tbl sched) naddr

toxKademlia :: Client -> TriadCommittee NodeId SockAddr -> OnionRouter -> TVar (R.BucketList NodeInfo) -> TVar (Int.PSQ POSIXTime) -> Kademlia NodeId NodeInfo
toxKademlia client committee orouter var sched
    = Kademlia quietInsertions
               toxSpace
               (vanillaIO var $ ping client)
                   { tblTransition = \tr -> do
                        io1 <- transitionCommittee committee tr
                        io2 <- touchBucket toxSpace (15*60) var sched tr
                        hookBucketList toxSpace var orouter tr
                        return $ do
                            io1 >> io2
                            {-
                            hPutStrLn stderr $ unwords
                                [ show (transitionedTo tr)
                                , show (transitioningNode tr)
                                ]
                            -}
                            return ()
                   }

transitionCommittee :: TriadCommittee NodeId SockAddr -> RoutingTransition NodeInfo -> STM (IO ())
transitionCommittee committee (RoutingTransition ni Stranger) = do
    delVote committee (nodeId ni)
    return $ do
        -- hPutStrLn stderr $ "delVote "++show (nodeId ni)
        return ()
transitionCommittee committee _ = return $ return ()

type Handler = MethodHandler String TransactionId NodeInfo Message

isPing :: (f Ping -> Ping) ->  DHTMessage f -> Either String Ping
isPing unpack (DHTPing a) = Right $ unpack $ asymmData a
isPing _      _           = Left "Bad ping"

mkPong :: TransactionId -> NodeInfo -> NodeInfo -> Pong -> DHTMessage ((,) Nonce8)
mkPong tid src dst pong = DHTPong $ wrapAsymm tid src dst (, pong)

isGetNodes :: (f GetNodes -> GetNodes) ->  DHTMessage f -> Either String GetNodes
isGetNodes unpack (DHTGetNodes a) = Right $ unpack $ asymmData a
isGetNodes _      _               = Left "Bad GetNodes"

mkSendNodes :: TransactionId -> NodeInfo -> NodeInfo -> SendNodes -> DHTMessage ((,) Nonce8)
mkSendNodes tid src dst sendnodes = DHTSendNodes $ wrapAsymm tid src dst (, sendnodes)

isCookieRequest :: (f CookieRequest -> CookieRequest) -> DHTMessage f -> Either String CookieRequest
isCookieRequest unpack (DHTCookieRequest a) = Right $ unpack $ asymmData a
isCookieRequest _      _                    = Left "Bad cookie request"

mkCookie :: TransactionId -> NodeInfo -> NodeInfo -> Cookie -> DHTMessage ((,) Nonce8)
mkCookie (TransactionId n8 n24) src dst cookie = DHTCookie n24 (n8,cookie)

isDHTRequest :: (f DHTRequest -> DHTRequest) -> DHTMessage f -> Either String DHTRequest
isDHTRequest unpack (DHTDHTRequest pubkey a) = Right $ unpack $ asymmData a
isDHTRequest _      _                        = Left "Bad dht relay request"

dhtRequestH :: NodeInfo -> DHTRequest -> IO ()
dhtRequestH ni req = do
    hPutStrLn stderr $ "Unhandled DHT Request: " ++ show req

handlers :: TransportCrypto -> Routing -> PacketKind -> Maybe Handler
handlers _      routing PingType          = Just $ MethodHandler (isPing snd) mkPong            $ pingH
handlers _      routing GetNodesType      = Just $ MethodHandler (isGetNodes snd) mkSendNodes   $ getNodesH routing
handlers crypto _       CookieRequestType = Just $ MethodHandler (isCookieRequest snd) mkCookie $ cookieRequestH crypto
handlers _      _       DHTRequestType    = Just $ NoReply (isDHTRequest snd) $ dhtRequestH
handlers _      _       typ               = error $ "TODO DHT handlers " ++ show typ

nodeSearch :: Client -> Search NodeId (IP,PortNumber) () NodeInfo NodeInfo
nodeSearch client  = Search
    { searchSpace       = toxSpace
    , searchNodeAddress = nodeIP &&& nodePort
    , searchQuery       = getNodes client
    }