path: root/src/Network/BitTorrent/DHT/Query.hs

-- |
--   Copyright   :  (c) Sam Truzjan 2014
--   License     :  BSD3
--   Maintainer  :  pxqr.sta@gmail.com
--   Stability   :  experimental
--   Portability :  portable
--
--   This module provides functions to interact with other nodes.
--   Normally, you don't need to import this module, use
--   "Network.BitTorrent.DHT" instead.
--
{-# LANGUAGE FlexibleContexts    #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TemplateHaskell     #-}
module Network.BitTorrent.DHT.Query
       ( -- * Handler
         -- | To bind specific set of handlers you need to pass
         -- handler list to the 'startNode' function.
         pingH
       , findNodeH
       , getPeersH
       , announceH
       , defaultHandlers

         -- * Query
         -- ** Basic
         -- | A basic query perform a single request expecting a
         -- single response.
       , Iteration
       , pingQ
       , findNodeQ
       , getPeersQ
       , announceQ

         -- ** Iterative
         -- | An iterative query perform multiple basic queries,
         -- concatenate its responses, optionally yielding result and
         -- continue to the next iteration.
       , Search
       , search
       , publish

       -- ** Routing table
       , restoreTable
       , insertNode
       , refreshNodes

       -- ** Messaging
       , queryNode
       , (<@>)
       ) where

import Control.Applicative
import Control.Concurrent.Lifted hiding (yield)
import Control.Exception.Lifted hiding (Handler)
import Control.Monad.Reader
import Control.Monad.Logger
import Data.Conduit
import Data.Conduit.List as C hiding (mapMaybe, mapM_)
import Data.Either
import Data.List as L
import Data.Monoid
import Data.Text as T
import Network
import Text.PrettyPrint as PP hiding ((<>), ($$))
import Text.PrettyPrint.HughesPJClass hiding ((<>),($$))
import Data.Time
import Data.Time.Clock.POSIX

import Network.KRPC hiding (Options, def)
import Data.Torrent
import Network.BitTorrent.Address
import Network.BitTorrent.DHT.Message
import Network.BitTorrent.DHT.Routing as R
import Network.BitTorrent.DHT.Session
import Control.Concurrent.STM

{-----------------------------------------------------------------------
--  Handlers
-----------------------------------------------------------------------}

nodeHandler :: Address ip => KRPC (Query a) (Response b)
           => (NodeAddr ip -> a -> DHT ip b) -> NodeHandler ip
nodeHandler action = handler $ \ sockAddr (Query remoteId read_only q) -> do
  case fromSockAddr sockAddr of
    Nothing    -> throwIO BadAddress
    Just naddr -> do
      let ni = NodeInfo remoteId naddr
      -- Do not route read-only nodes. (bep 43)
      if read_only
        then $(logWarnS) "nodeHandler" $ "READ-ONLY " <> T.pack (show $ pPrint ni)
        else insertNode ni >> return () -- TODO need to block. why?
      Response <$> asks thisNodeId <*> action naddr q

-- | Default 'Ping' handler.
pingH :: Address ip => NodeHandler ip
pingH = nodeHandler $ \ _ Ping -> do
  return Ping

-- | Default 'FindNode' handler.
findNodeH :: Address ip => NodeHandler ip
findNodeH = nodeHandler $ \ _ (FindNode nid) -> do
  NodeFound <$> getClosest nid

-- | Default 'GetPeers' handler.
getPeersH :: Address ip => NodeHandler ip
getPeersH = nodeHandler $ \ naddr (GetPeers ih) -> do
  ps <- getPeerList ih
  tok <- grantToken naddr
  $(logDebugS) "getPeersH" $ "INFO-HASH " <> T.pack (show (ih,fmap fromAddr naddr :: NodeAddr (Maybe IP)))
  return $ GotPeers ps tok

-- | Default 'Announce' handler.
announceH :: Address ip => NodeHandler ip
announceH = nodeHandler $ \ naddr @ NodeAddr {..} (Announce {..}) -> do
  valid <- checkToken naddr sessionToken
  unless valid $ do
    throwIO $ InvalidParameter "token"

  let annPort  = if impliedPort then nodePort else port
  let peerAddr = PeerAddr Nothing nodeHost annPort
  insertPeer topic peerAddr
  return Announced

-- | Includes all default query handlers.
defaultHandlers :: Address ip => [NodeHandler ip]
defaultHandlers = [pingH, findNodeH, getPeersH, announceH]

{-----------------------------------------------------------------------
--  Basic queries
-----------------------------------------------------------------------}

type Iteration ip o = NodeInfo ip -> DHT ip (Either [NodeInfo ip] [o ip])

-- | The most basic query. May be used to check if the given node is
-- alive or get its 'NodeId'.
pingQ :: Address ip => NodeAddr ip -> DHT ip (NodeInfo ip)
pingQ addr = do
  (nid, Ping) <- queryNode addr Ping
  return (NodeInfo nid addr)

-- TODO [robustness] match range of returned node ids with the
-- expected range and either filter bad nodes or discard response at
-- all throwing an exception
findNodeQ :: Address ip => TableKey key => key -> Iteration ip NodeInfo
findNodeQ key NodeInfo {..} = do
  NodeFound closest <- FindNode (toNodeId key) <@> nodeAddr
  $(logInfoS) "findNodeQ" $ "NodeFound " <> T.pack (show $ L.map pPrint closest)
  return $ Right closest

getPeersQ :: Address ip => InfoHash -> Iteration ip PeerAddr
getPeersQ topic NodeInfo {..} = do
  GotPeers {..} <- GetPeers topic <@> nodeAddr
  let dist = distance (toNodeId topic) nodeId
  $(logInfoS) "getPeersQ" $ T.pack
         $ "distance: " <> render (pPrint dist) <> " , result: "
        <> case peers of { Left _ -> "NODES"; Right _ -> "PEERS" }
  return peers

announceQ :: Address ip => InfoHash -> PortNumber -> Iteration ip NodeAddr
announceQ ih p NodeInfo {..} = do
  GotPeers {..} <- GetPeers ih <@> nodeAddr
  case peers of
    Left  ns
      | False     -> undefined -- TODO check if we can announce
      | otherwise -> return (Left ns)
    Right ps -> do -- TODO *probably* add to peer cache
      Announced <- Announce False ih p grantedToken <@> nodeAddr
      return (Right [nodeAddr])

{-----------------------------------------------------------------------
--  Iterative queries
-----------------------------------------------------------------------}

type Search    ip o = Conduit [NodeInfo ip] (DHT ip) [o ip]

-- TODO: use reorder and filter (Traversal option) leftovers
search :: TableKey k => Address ip => k -> Iteration ip o -> Search ip o
search _ action = do
  awaitForever $ \ batch -> unless (L.null batch) $ do
    $(logWarnS) "search" "start query"
    responses <- lift $ queryParallel (action <$> batch)
    let (nodes, results) = partitionEithers responses
    $(logWarnS) "search" ("done query more:" <> T.pack (show (L.length nodes, L.length results)))
    leftover $ L.concat nodes
    mapM_ yield results

publish :: Address ip => InfoHash -> PortNumber -> DHT ip ()
publish ih p = do
  nodes <- getClosest ih
  r     <- asks (optReplication . options)
  _ <- sourceList [nodes] $= search ih (announceQ ih p) $$ C.take r
  return ()

republish :: DHT ip ThreadId
republish = fork $ do
  i <- asks (optReannounce . options)
  error "DHT.republish: not implemented"

routing :: Address ip => Routing ip a -> DHT ip (Maybe a)
routing = runRouting probeNode refreshNodes getTimestamp

getTimestamp :: DHT ip Timestamp
getTimestamp = do
  utcTime <- liftIO $ getCurrentTime
  $(logDebugS) "routing.make_timestamp" (T.pack (render (pPrint utcTime)))
  return $ utcTimeToPOSIXSeconds utcTime


probeNode :: Address ip => NodeAddr ip -> DHT ip Bool
probeNode addr = do
  $(logDebugS) "routing.questionable_node" (T.pack (render (pPrint addr)))
  result <- try $ Ping <@> addr
  let _ = result :: Either SomeException Ping
  return $ either (const False) (const True) result


-- FIXME do not use getClosest sinse we should /refresh/ them
refreshNodes :: Address ip => NodeId -> DHT ip () -- [NodeInfo ip]
refreshNodes nid = do
  $(logDebugS) "routing.refresh_bucket" (T.pack (render (pPrint nid)))
  nodes <- getClosest nid
  do
    -- forM (L.take 1 nodes) $ \ addr -> do
    -- NodeFound ns <- FindNode nid <@> addr
    -- Expected type: ConduitM [NodeAddr ip] [NodeInfo ip] (DHT ip) ()
    --   Actual type: ConduitM [NodeInfo ip] [NodeInfo ip] (DHT ip) ()
    -- nss <- sourceList [[addr]] $= search nid (findNodeQ nid) $$ C.consume
    nss <- sourceList [nodes] $= search nid (findNodeQ nid) $$ C.consume
    $(logWarnS) "refreshNodes" $ "received " <> T.pack (show (L.length (L.concat nss))) <> " nodes."
    queryParallel $ flip L.map (L.concat nss) $ \n -> do
        $(logWarnS) "refreshNodes" $ "received node: " <> T.pack (show (pPrint n))
        pingQ (nodeAddr n)
        insertNode n
    return ()
  return () -- $ L.concat nss

-- | This operation do not block but acquire exclusive access to
--   routing table.
insertNode :: Address ip => NodeInfo ip -> DHT ip ThreadId
insertNode info = fork $ do
  var <- asks routingTable
  tm <- getTimestamp
  let showTable = do
          t <- liftIO $ atomically $ readTVar var
          let logMsg = "Routing table: " <> pPrint t
          $(logDebugS) "insertNode" (T.pack (render logMsg))
  t <- liftIO $ atomically $ readTVar var
  let arrival  = TryInsert info
      arrival4 = TryInsert (fmap fromAddr info) :: Event (Maybe IPv4)
  $(logDebugS) "insertNode" $ T.pack (show arrival4)
  ps <- liftIO $ atomically $ do
      t <- readTVar var
      (ps,t') <- R.insert tm arrival t
      writeTVar var t'
      return ps
  showTable
  fork $ forM_ ps $ \(CheckPing ns)-> do
    forM_ ns $ \n -> do
        alive <- PingResult n <$> probeNode (nodeAddr n)
        let PingResult _ b = alive
        $(logDebugS) "insertNode" $ T.pack ("PingResult "++show (nodeId n,b))
        tm <- getTimestamp
        liftIO $ atomically $ do
            t <- readTVar var
            (_,t') <- R.insert tm alive t
            writeTVar var t'
        showTable
  return ()

-- | Throws exception if node is not responding.
queryNode :: forall a b ip. Address ip => KRPC (Query a) (Response b)
          => NodeAddr ip -> a -> DHT ip (NodeId, b)
queryNode addr q = do
  nid <- asks thisNodeId
  let read_only = False -- TODO: check for NAT issues. (BEP 43)
  Response remoteId r <- query (toSockAddr addr) (Query nid read_only q)
  insertNode (NodeInfo remoteId addr)
  return (remoteId, r)

-- | Infix version of 'queryNode' function.
(<@>) :: Address ip => KRPC (Query a) (Response b)
      => a -> NodeAddr ip -> DHT ip b
q <@> addr = snd <$> queryNode addr q
{-# INLINE (<@>) #-}

restoreTable :: Address ip => Table ip -> DHT ip ()
restoreTable tbl = mapM_ (insertNode . fst) $ L.concat $ R.toList tbl