path: root/src/Network/KRPC/Manager.hs

-- |
--   Copyright   :  (c) Sam Truzjan 2013, 2014
--   License     :  BSD3
--   Maintainer  :  pxqr.sta@gmail.com
--   Stability   :  experimental
--   Portability :  portable
--
--   Normally, you don't need to import this module.
--
{-# LANGUAGE CPP                    #-}
{-# LANGUAGE OverloadedStrings      #-}
{-# LANGUAGE FlexibleInstances      #-}
{-# LANGUAGE FlexibleContexts       #-}
{-# LANGUAGE ScopedTypeVariables    #-}
{-# LANGUAGE DefaultSignatures      #-}
{-# LANGUAGE MultiParamTypeClasses  #-}
{-# LANGUAGE FunctionalDependencies #-}
{-# LANGUAGE DeriveDataTypeable     #-}
{-# LANGUAGE TemplateHaskell        #-}
module Network.KRPC.Manager
       ( -- * Manager
         MonadKRPC (..)
       , Options (..)
       , Manager
       , newManager
       , closeManager
       , withManager
       , isActive
       , listen

         -- * Queries
       , QueryFailure (..)
       , query
       , query'
       , queryRaw
       , getQueryCount

         -- * Handlers
       , HandlerFailure (..)
       , Handler
       , handler
       ) where

import Control.Applicative
#ifdef THREAD_DEBUG
import Control.Concurrent.Lifted.Instrument
#else
import GHC.Conc (labelThread)
import Control.Concurrent.Lifted
#endif
import Control.Exception hiding (Handler)
import qualified Control.Exception.Lifted as E (Handler (..))
import Control.Exception.Lifted as Lifted (catches, finally)
import Control.Monad
import Control.Monad.Logger
import Control.Monad.Reader
import Control.Monad.Trans.Control
#ifdef VERSION_bencoding
import Data.BEncode as BE
import Data.BEncode.Internal as BE
import Data.BEncode.Pretty (showBEncode)
#else
import qualified Data.Tox as Tox
#endif
import qualified Data.ByteString.Base16 as Base16
import Data.ByteString as BS
import Data.ByteString.Char8 as BC
import Data.ByteString.Lazy as BL
import Data.Default.Class
import Data.IORef
import Data.List as L
import Data.Map as M
import Data.Monoid
import Data.Serialize as S
import Data.Text as T
import Data.Text.Encoding as T
import Data.Tuple
import Data.Typeable
import Network.RPC
import Network.KRPC.Message
import Network.KRPC.Method hiding (Envelope)
import qualified Network.KRPC.Method as KRPC (Envelope)
import Network.Socket hiding (listen)
import Network.Socket.ByteString as BS
import System.IO.Error
import System.Timeout
#ifdef VERSION_bencoding
import Network.DHT.Mainline
#endif


{-----------------------------------------------------------------------
--  Options
-----------------------------------------------------------------------}

-- | RPC manager options.
data Options = Options
  { -- | Initial 'TransactionId' incremented with each 'query';
    optSeedTransaction :: {-# UNPACK #-} !Int

    -- | Time to wait for response from remote node, in seconds.
  , optQueryTimeout    :: {-# UNPACK #-} !Int

    -- | Maximum number of bytes to receive.
  , optMaxMsgSize      :: {-# UNPACK #-} !Int
  } deriving (Show, Eq)

defaultSeedTransaction :: Int
defaultSeedTransaction = 0

defaultQueryTimeout :: Int
defaultQueryTimeout = 120

defaultMaxMsgSize :: Int
defaultMaxMsgSize = 64 * 1024

-- | Permissive defaults.
instance Default Options where
  def = Options
    { optSeedTransaction = defaultSeedTransaction
    , optQueryTimeout    = defaultQueryTimeout
    , optMaxMsgSize      = defaultMaxMsgSize
    }

validateOptions :: Options -> IO ()
validateOptions Options {..}
  | optQueryTimeout < 1
  = throwIO (userError "krpc: non-positive query timeout")
  | optMaxMsgSize   < 1
  = throwIO (userError "krpc: non-positive buffer size")
  |       otherwise     = return ()

{-----------------------------------------------------------------------
--  Options
-----------------------------------------------------------------------}

type KResult = Either KError KMessage -- Response

type TransactionCounter = IORef Int
type CallId             = (TransactionId, SockAddr)
type CallRes            = MVar (KQueryArgs, KResult) -- (raw response, decoded response)
type PendingCalls       = IORef (Map CallId CallRes)

type HandlerBody h msg v = SockAddr -> msg v -> h (Either String (msg v))

-- | Handler is a function which will be invoked then some /remote/
-- node querying /this/ node.
type Handler     h msg v = (MethodName, HandlerBody h msg v)

-- | Keep track pending queries made by /this/ node and handle queries
-- made by /remote/ nodes.
data Manager h = Manager
  { sock               :: !Socket
  , options            :: !Options
  , listenerThread     :: !(MVar ThreadId)
  , transactionCounter :: {-# UNPACK #-} !TransactionCounter
  , pendingCalls       :: {-# UNPACK #-} !PendingCalls
#ifdef VERSION_bencoding
  , handlers           :: [Handler h KMessageOf BValue]
#else
  , handlers           :: [Handler h KMessageOf BC.ByteString]
#endif
  }

-- | A monad which can perform or handle queries.
class (MonadBaseControl IO m, MonadLogger m, MonadIO m)
    => MonadKRPC h m | m -> h where

  -- | Ask for manager.
  getManager :: m (Manager h)

  default getManager :: MonadReader (Manager h) m => m (Manager h)
  getManager = ask

  -- | Can be used to add logging for instance.
  liftHandler :: h a -> m a

  default liftHandler :: m a -> m a
  liftHandler = id

instance (MonadBaseControl IO h, MonadLogger h, MonadIO h)
      => MonadKRPC h (ReaderT (Manager h) h) where

  liftHandler = lift

sockAddrFamily :: SockAddr -> Family
sockAddrFamily (SockAddrInet  _ _    ) = AF_INET
sockAddrFamily (SockAddrInet6 _ _ _ _) = AF_INET6
sockAddrFamily (SockAddrUnix  _      ) = AF_UNIX
sockAddrFamily (SockAddrCan _        ) = AF_CAN

-- | Bind socket to the specified address. To enable query handling
-- run 'listen'.
newManager :: Options                       -- ^ various protocol options;
           -> SockAddr                      -- ^ address to listen on;
#ifdef VERSION_bencoding
           -> [Handler h KMessageOf BValue] -- ^ handlers to run on incoming queries.
#else
           -> [Handler h KMessageOf BC.ByteString] -- ^ handlers to run on incoming queries.
#endif
           -> IO (Manager h)                -- ^ new rpc manager.
newManager opts @ Options {..} servAddr handlers = do
    validateOptions opts
    sock  <- bindServ
    tref  <- newEmptyMVar
    tran  <- newIORef optSeedTransaction
    calls <- newIORef M.empty
    return $ Manager sock opts tref tran calls handlers
  where
    bindServ = do
      let family = sockAddrFamily servAddr
      sock <- socket family Datagram defaultProtocol
      when (family == AF_INET6) $ do
        setSocketOption sock IPv6Only 0
      bindSocket sock servAddr
      return sock

-- | Unblock all pending calls and close socket.
closeManager :: Manager m -> IO ()
closeManager Manager {..} = do
  maybe (return ()) killThread =<< tryTakeMVar listenerThread
  -- TODO unblock calls
  close sock

-- | Check if the manager is still active. Manager becomes active
-- until 'closeManager' called.
isActive :: Manager m -> IO Bool
isActive Manager {..} = liftIO $ isBound sock
{-# INLINE isActive #-}

-- | Normally you should use Control.Monad.Trans.Resource.allocate
-- function.
#ifdef VERSION_bencoding
withManager :: Options -> SockAddr -> [Handler h KMessageOf BValue]
#else
withManager :: Options -> SockAddr -> [Handler h KMessageOf BC.ByteString]
#endif
            -> (Manager h -> IO a) -> IO a
withManager opts addr hs = bracket (newManager opts addr hs) closeManager

{-----------------------------------------------------------------------
--  Logging
-----------------------------------------------------------------------}

-- TODO prettify log messages
querySignature :: MethodName -> TransactionId -> SockAddr -> Text
querySignature name transaction addr = T.concat
#ifdef VERSION_bencoding
  [  "&", T.decodeUtf8 name
  , " #", T.decodeUtf8 (Base16.encode transaction) -- T.decodeUtf8 transaction
#else
  [  "&", T.pack (show name)
  , " #", T.decodeUtf8 (Base16.encode $ S.encode transaction)
#endif
  , " @", T.pack (show addr)
  ]

{-----------------------------------------------------------------------
--  Client
-----------------------------------------------------------------------}
-- we don't need to know about TransactionId while performing query,
-- so we introduce QueryFailure exceptions

-- | Used to signal 'query' errors.
data QueryFailure
  = SendFailed                  -- ^ unable to send query;
  | QueryFailed  ErrorCode Text -- ^ remote node return error;
  | TimeoutExpired              -- ^ remote node not responding.
  deriving (Show, Eq, Typeable)

instance Exception QueryFailure

sendMessage :: MonadIO m => Socket -> SockAddr -> BC.ByteString -> m ()
sendMessage sock addr a = do
  liftIO $ sendManyTo sock [a] addr

genTransactionId :: TransactionCounter -> IO TransactionId
genTransactionId ref = do
  cur <- atomicModifyIORef' ref $ \ cur -> (succ cur, cur)
#ifdef VERSION_bencoding
  return $ BC.pack (show cur)
#else
  return $ either (error "failed to create TransactionId") id $ S.decode $ BC.pack (L.take 24 $ show cur ++ L.repeat ' ')
#endif

-- | How many times 'query' call have been performed.
getQueryCount :: MonadKRPC h m => m Int
getQueryCount = do
  Manager {..} <- getManager
  curTrans <- liftIO $ readIORef transactionCounter
  return $ curTrans - optSeedTransaction options

registerQuery :: CallId -> PendingCalls -> IO CallRes
registerQuery cid ref = do
  ares <- newEmptyMVar
  atomicModifyIORef' ref $ \ m -> (M.insert cid ares m, ())
  return ares

-- simultaneous M.lookup and M.delete guarantees that we never get two
-- or more responses to the same query
unregisterQuery :: CallId -> PendingCalls -> IO (Maybe CallRes)
unregisterQuery cid ref = do
  atomicModifyIORef' ref $ swap .
    M.updateLookupWithKey (const (const Nothing)) cid


-- (sendmsg EINVAL)
sendQuery :: Socket -> SockAddr -> BC.ByteString -> IO ()
sendQuery sock addr q = handle sockError $ sendMessage sock addr q
  where
    sockError :: IOError -> IO ()
    sockError _ = throwIO SendFailed

-- | Enqueue query to the given node.
--
--  This function should throw 'QueryFailure' exception if quered node
--  respond with @error@ message or the query timeout expires.
--
query :: forall h m a b. (MonadKRPC h m, KRPC a b) => SockAddr -> a -> m b
query addr params = queryK addr params (\_ x _ -> x)

-- | Like 'query' but possibly returns your externally routable IP address.
query' :: forall h m a b. (MonadKRPC h m, KRPC a b) => SockAddr -> a -> m (b, Maybe ReflectedIP)
query' addr params = queryK addr params (const (,))

-- | Enqueue a query, but give us the complete BEncoded content sent by the
-- remote Node.  This is useful for handling extensions that this library does
-- not otherwise support.
queryRaw :: forall h m a b. (MonadKRPC h m, KRPC a b) => SockAddr -> a -> m (b, KQueryArgs)
queryRaw addr params = queryK addr params (\raw x _ -> (x,raw))

queryK :: forall h m a b x. (MonadKRPC h m, KRPC a b) =>
    SockAddr -> a -> (KQueryArgs -> b -> Maybe ReflectedIP -> x) -> m x
queryK addr params kont = do
  Manager {..} <- getManager
  tid <- liftIO $ genTransactionId transactionCounter
  let queryMethod = method :: Method a b
  let signature = querySignature (methodName queryMethod) tid addr
  $(logDebugS) "query.sending" signature

  mres <- liftIO $ do
    ares <- registerQuery (tid, addr) pendingCalls

#ifdef VERSION_bencoding
    let q = Q (KQuery (toBEncode params) (methodName queryMethod) tid)
        qb = encodePayload q :: KMessage
        qbs = encodeHeaders () qb :: BC.ByteString
#else
    let q = Tox.Message (methodName queryMethod) cli tid params
        cli = error "TODO TOX client node id"
        ctx = error "TODO TOX ToxCipherContext"
        qb = encodePayload q :: Tox.Message BC.ByteString
        qbs = encodeHeaders ctx qb :: BC.ByteString
#endif
    sendQuery sock addr qbs
      `onException` unregisterQuery (tid, addr) pendingCalls

    timeout (optQueryTimeout options * 10 ^ (6 :: Int)) $ do
      (raw,res) <- readMVar ares -- MVar (KQueryArgs, KResult)
      case res of
#ifdef VERSION_bencoding
        Left  (KError   c m _) -> throwIO $ QueryFailed c (T.decodeUtf8 m)
        Right (R (KResponse {..})) ->
          case fromBEncode respVals of
            Right r -> pure $ kont raw r respIP
#else
        Left  _                  -> throwIO $ QueryFailed GenericError "TODO: TOX ERROR"
        Right (Tox.Message {..}) ->
          case S.decode msgPayload of
            Right r -> pure $ kont raw r Nothing
#endif
            Left  e -> throwIO $ QueryFailed ProtocolError (T.pack e)

  case mres of
    Just res -> do
      $(logDebugS) "query.responded" $ signature
      return res

    Nothing -> do
      _ <- liftIO $ unregisterQuery (tid, addr) pendingCalls
      $(logWarnS) "query.not_responding" $ signature <> " for " <>
             T.pack (show (optQueryTimeout options)) <> " seconds"
      throw $ TimeoutExpired

{-----------------------------------------------------------------------
--  Handlers
-----------------------------------------------------------------------}
-- we already throw:
--
--   * ErrorCode(MethodUnknown) in the 'dispatchHandler';
--
--   * ErrorCode(ServerError) in the 'runHandler';
--
--   * ErrorCode(GenericError) in the 'runHandler' (those can be
--   async exception too)
--
--  so HandlerFailure should cover *only* 'ProtocolError's.

-- | Used to signal protocol errors.
data HandlerFailure
  = BadAddress            -- ^ for e.g.: node calls herself;
  | InvalidParameter Text -- ^ for e.g.: bad session token.
    deriving (Show, Eq, Typeable)

instance Exception HandlerFailure

prettyHF :: HandlerFailure -> BS.ByteString
prettyHF  BadAddress               = T.encodeUtf8 "bad address"
prettyHF (InvalidParameter reason) = T.encodeUtf8 $
  "invalid parameter: " <> reason

prettyQF :: QueryFailure -> BS.ByteString
prettyQF e = T.encodeUtf8 $ "handler fail while performing query: "
  <> T.pack (show e)

-- | Make handler from handler function. Any thrown exception will be
-- supressed and send over the wire back to the querying node.
--
--   If the handler make some 'query' normally it /should/ handle
--   corresponding 'QueryFailure's.
--
handler :: forall h a b msg raw. (KRPC a b, Applicative h, Functor msg, WireFormat raw msg, SerializableTo raw a, SerializableTo raw b)
        => (SockAddr -> a -> h b) -> Handler h msg raw
handler body = (name, wrapper)
  where
    Method name = method :: Method a b
    wrapper :: SockAddr -> msg raw -> h (Either String (msg raw))
    wrapper addr args =
      case decodePayload args of
        Left  e -> pure $ Left e
        Right a -> Right . encodePayload . buildReply (error "self node-id") addr args <$> body addr (envelopePayload a)

runHandler :: MonadKRPC h m
#ifdef VERSION_bencoding
           => HandlerBody h KMessageOf BValue -> SockAddr -> KQuery -> m KResult
#else
           => HandlerBody h KMessageOf BC.ByteString -> SockAddr -> KQuery -> m KResult
#endif
runHandler h addr m = Lifted.catches wrapper failbacks
  where
    signature = querySignature (queryMethod m) (queryId m) addr

    wrapper = do
      $(logDebugS) "handler.quered" signature
#ifdef VERSION_bencoding
      result <- liftHandler (h addr (Q m))
#else
      result <- liftHandler (h addr m)
#endif

      case result of
        Left msg -> do
          $(logDebugS) "handler.bad_query" $ signature <> " !" <> T.pack msg
#ifdef VERSION_bencoding
          return $ Left $ KError ProtocolError (BC.pack msg) (queryId m)
#else
          return $ Left $ decodeError "TODO TOX ProtocolError" (queryId m)
#endif

        Right a -> do -- KQueryArgs
          $(logDebugS) "handler.success" signature
          return $ Right a

    failbacks =
      [ E.Handler $ \ (e :: HandlerFailure) -> do
          $(logDebugS) "handler.failed" signature
#ifdef VERSION_bencoding
          return $ Left $ KError ProtocolError (prettyHF e) (queryId m)
#else
          return $ Left $ decodeError "TODO TOX ProtocolError 2" (queryId m)
#endif


        -- may happen if handler makes query and fail
      , E.Handler $ \ (e :: QueryFailure) -> do
#ifdef VERSION_bencoding
          return $ Left $ KError ServerError (prettyQF e) (queryId m)
#else
          return $ Left $ decodeError "TODO TOX ServerError" (queryId m)
#endif

        -- since handler thread exit after sendMessage we can safely
        -- suppress async exception here
      , E.Handler $ \ (e :: SomeException) -> do
#ifdef VERSION_bencoding
          return $ Left $ KError GenericError (BC.pack (show e)) (queryId m)
#else
          return $ Left $ decodeError "TODO TOX GenericError" (queryId m)
#endif
      ]

dispatchHandler :: MonadKRPC h m => KQuery -> SockAddr -> m KResult
dispatchHandler q addr = do
  Manager {..} <- getManager
  case L.lookup (queryMethod q) handlers of
#ifdef VERSION_bencoding
    Nothing -> return $ Left $ KError MethodUnknown (queryMethod q) (queryId q)
#else
    Nothing -> return $ Left $ decodeError "TODO TOX Error MethodUnknown" (queryId q)
#endif
    Just h  -> runHandler h addr q

{-----------------------------------------------------------------------
--  Listener
-----------------------------------------------------------------------}

-- TODO bound amount of parallel handler *threads*:
--
-- peer A flooding with find_node
-- peer B trying to ping peer C
-- peer B fork too many threads
-- ... space leak
--
handleQuery :: MonadKRPC h m => KQueryArgs -> KQuery -> SockAddr -> m ()
handleQuery raw q addr = void $ fork $ do
  myThreadId >>= liftIO . flip labelThread "KRPC.handleQuery"
  Manager {..} <- getManager
  res <- dispatchHandler q addr
#ifdef VERSION_bencoding
  let  res' = either E id res
       resbe = either toBEncode toBEncode res
  $(logOther "q") $ T.unlines
    [ either (const "<unicode-fail>") id $ T.decodeUtf8' (BL.toStrict $ showBEncode raw)
    , "==>"
    , either (const "<unicode-fail>") id $ T.decodeUtf8' (BL.toStrict $ showBEncode resbe)
    ]
  sendMessage sock addr $ encodeHeaders () res'
#else
  -- Errors not sent for Tox.
  let ctx = error "TODO TOX ToxCipherContext 2"
  either (const $ return ()) (sendMessage sock addr . encodeHeaders ctx) res
#endif

handleResponse :: MonadKRPC h m => KQueryArgs -> KResult -> SockAddr -> m ()
handleResponse raw result addr = do
  Manager {..} <- getManager
  liftIO $ do
#ifdef VERSION_bencoding
    let resultId = either errorId envelopeTransaction result
#else
    let resultId = either Tox.msgNonce Tox.msgNonce result
#endif
    mcall <- unregisterQuery (resultId, addr) pendingCalls
    case mcall of
      Nothing   -> return ()
      Just ares -> putMVar ares (raw,result)

#ifdef VERSION_bencoding
handleMessage :: MonadKRPC h m => KQueryArgs -> KMessage -> SockAddr -> m ()
handleMessage raw (Q q) = handleQuery    raw q
handleMessage raw (R r) = handleResponse raw (Right (R r))
handleMessage raw (E e) = handleResponse raw (Left  e)
#else
handleMessage :: MonadKRPC h m => KQueryArgs -> Tox.Message BC.ByteString -> SockAddr -> m ()
handleMessage raw q | Tox.isQuery q    = handleQuery    raw q
handleMessage raw r | Tox.isResponse r = handleResponse raw (Right r)
handleMessage raw e | Tox.isError e    = handleResponse raw (Left  e)
#endif

listener :: MonadKRPC h m => m ()
listener = do
  Manager {..} <- getManager
  fix $ \again -> do
    let ctx = error "TODO TOX ToxCipherContext 3"
    (bs, addr) <- liftIO $ do
      handle exceptions $ BS.recvFrom sock (optMaxMsgSize options)
#ifdef VERSION_bencoding
    case BE.parse bs >>= \r -> (,) r <$> BE.decode bs of
#else
    case return bs >>= \r -> (,) r <$> decodeHeaders ctx bs of
#endif
      -- TODO ignore unknown messages at all?
#ifdef VERSION_bencoding
      Left  e       -> liftIO $ sendMessage sock addr $ encodeHeaders () (E (unknownMessage e) :: KMessage)
#else
      Left  _       -> return () -- TODO TOX send unknownMessage error
#endif
      Right (raw,m) -> handleMessage raw m addr
    again
  where
    exceptions :: IOError -> IO (BS.ByteString, SockAddr)
    exceptions e
        -- packets with empty payload may trigger eof exception
      | isEOFError e = return  ("", SockAddrInet 0 0)
      |   otherwise  = throwIO e

-- | Should be run before any 'query', otherwise they will never
-- succeed.
listen :: MonadKRPC h m => m ()
listen = do
  Manager {..} <- getManager
  tid <- fork $ do
    myThreadId >>= liftIO . flip labelThread "KRPC.listen"
    listener `Lifted.finally`
      liftIO (takeMVar listenerThread)
  liftIO $ putMVar listenerThread tid