Merged Network.KRPC.Manager into Network.KRPC.

3 files changed, 565 insertions, 598 deletions

diff --git a/src/Network/BitTorrent/DHT/Query.hs b/src/Network/BitTorrent/DHT/Query.hs
index 56ea262a..a872fc72 100644
--- a/src/Network/BitTorrent/DHT/Query.hs
+++ b/src/Network/BitTorrent/DHT/Query.hs
@@ -83,7 +83,7 @@ import Data.Time.Clock.POSIX
 
 import Network.KRPC as KRPC hiding (Options, def)
 import Network.KRPC.Message (ReflectedIP(..))
-import Network.KRPC.Manager (QueryFailure(..))
+import Network.KRPC (QueryFailure(..))
 import Data.Torrent
 import Network.BitTorrent.DHT.Message
 import Network.BitTorrent.DHT.Routing as R
diff --git a/src/Network/KRPC.hs b/src/Network/KRPC.hs
index d185fb4c..22ab81f2 100644
--- a/src/Network/KRPC.hs
+++ b/src/Network/KRPC.hs
@@ -50,6 +50,16 @@
 --
 --   For protocol details see "Network.KRPC.Message" module.
 --
+{-# LANGUAGE CPP                    #-}
+{-# LANGUAGE OverloadedStrings      #-}
+{-# LANGUAGE FlexibleInstances      #-}
+{-# LANGUAGE FlexibleContexts       #-}
+{-# LANGUAGE ScopedTypeVariables    #-}
+{-# LANGUAGE DefaultSignatures      #-}
+{-# LANGUAGE MultiParamTypeClasses  #-}
+{-# LANGUAGE FunctionalDependencies #-}
+{-# LANGUAGE DeriveDataTypeable     #-}
+{-# LANGUAGE TemplateHaskell        #-}
 module Network.KRPC
        ( -- * Methods
          Method
@@ -87,5 +97,558 @@ module Network.KRPC
 import Data.Default.Class
 import Network.KRPC.Message
 import Network.KRPC.Method
-import Network.KRPC.Manager
 import Network.Socket (SockAddr (..))
+
+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
diff --git a/src/Network/KRPC/Manager.hs b/src/Network/KRPC/Manager.hs
deleted file mode 100644
index efd59f32..00000000
--- a/src/Network/KRPC/Manager.hs
+++ /dev/null
@@ -1,596 +0,0 @@
--- |
---   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