1 files changed, 564 insertions, 1 deletions

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