1 files changed, 39 insertions, 9 deletions

diff --git a/src/Network/BitTorrent/DHT/Session.hs b/src/Network/BitTorrent/DHT/Session.hs
index 71400609..9243ef49 100644
--- a/src/Network/BitTorrent/DHT/Session.hs
+++ b/src/Network/BitTorrent/DHT/Session.hs
@@ -5,8 +5,7 @@
 {-# LANGUAGE MultiParamTypeClasses      #-}
 {-# LANGUAGE ScopedTypeVariables        #-}
 {-# LANGUAGE TypeFamilies               #-}
-
-{-# LANGUAGE RankNTypes #-} -- TODO remove
+{-# LANGUAGE TemplateHaskell            #-}
 module Network.BitTorrent.DHT.Session
        ( -- * Session
          DHT
@@ -35,19 +34,27 @@ module Network.BitTorrent.DHT.Session
 import Control.Applicative
 import Control.Concurrent.STM
 import Control.Exception hiding (Handler)
-import Control.Monad.Reader
 import Control.Monad.Base
+import Control.Monad.Logger
+import Control.Monad.Reader
 import Control.Monad.Trans.Control
 import Control.Monad.Trans.Resource
 import Data.Default
 import Data.Hashable
 import Data.List as L
+import Data.Monoid
+import Data.Text as T
+import Data.Text.Encoding as T
 import Data.Time
 import Data.Time.Clock.POSIX
+import System.Log.FastLogger
 import System.Random (randomIO)
+import Text.PrettyPrint as PP hiding ((<>))
+import Text.PrettyPrint.Class
 
 import Data.Torrent.InfoHash
 import Network.KRPC
+import Network.KRPC.Method
 import Network.BitTorrent.Core
 import Network.BitTorrent.Core.PeerAddr as P
 import Network.BitTorrent.DHT.Message
@@ -89,6 +96,7 @@ data Node ip = Node
   , routingTable  :: !(TVar    (Table     ip))
   , contactInfo   :: !(TVar    (PeerStore ip))
   , sessionTokens :: !(TVar     SessionTokens)
+  , loggerFun     :: !(Loc -> LogSource -> LogLevel -> LogStr -> IO ())
   }
 
 newtype DHT ip a = DHT { unDHT :: ReaderT (Node ip) (ResourceT IO) a }
@@ -96,6 +104,7 @@ newtype DHT ip a = DHT { unDHT :: ReaderT (Node ip) (ResourceT IO) a }
            , MonadIO, MonadBase IO
            , MonadReader (Node ip)
            )
+
 instance MonadBaseControl IO (DHT ip) where
   newtype StM (DHT ip) a = StM {
       unSt :: StM (ReaderT (Node ip) (ResourceT IO)) a
@@ -110,19 +119,26 @@ instance MonadBaseControl IO (DHT ip) where
 instance MonadKRPC (DHT ip) (DHT ip) where
   getManager = asks manager
 
+instance MonadLogger (DHT ip) where
+  monadLoggerLog loc src lvl msg = do
+    logger <- asks loggerFun
+    liftIO $ logger loc src lvl (toLogStr msg)
+
 runDHT :: forall ip a. Address ip
        => NodeAddr ip        -- ^ node address to bind;
        -> [Handler (DHT ip)] -- ^ handlers to run on accepted queries;
        -> DHT ip a           -- ^ DHT action to run;
        -> IO a               -- ^ result.
 runDHT naddr handlers action = runResourceT $ do
-  (_, m) <- allocate (newManager (toSockAddr naddr) handlers) closeManager
-  myId   <- liftIO genNodeId
-  node   <- liftIO $ Node m
+  runStderrLoggingT $ LoggingT $  \ logger -> do
+    (_, m) <- allocate (newManager (toSockAddr naddr) handlers) closeManager
+    myId   <- liftIO genNodeId
+    node   <- liftIO $ Node m
              <$> newTVarIO (nullTable myId)
              <*> newTVarIO def
              <*> (newTVarIO =<< nullSessionTokens)
-  runReaderT (unDHT (listen >> action)) node
+             <*> pure logger
+    runReaderT (unDHT (listen >> action)) node
 
 {-----------------------------------------------------------------------
 --  Routing
@@ -135,12 +151,14 @@ routing = runRouting ping refreshNodes getTimestamp
 -- TODO add timeout
 ping :: Address ip => NodeAddr ip -> DHT ip Bool
 ping addr = do
+  $(logDebugS) "routing.questionable_node" (T.pack (render (pretty addr)))
   Ping <- Ping <@> addr
   return True
 
 -- 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 (pretty nid)))
   nodes <- getClosest nid
   nss <- forM (nodeAddr <$> nodes) $ \ addr -> do
     NodeFound ns <- FindNode nid <@> addr
@@ -148,7 +166,10 @@ refreshNodes nid = do
   return $ L.concat nss
 
 getTimestamp :: DHT ip Timestamp
-getTimestamp = liftIO $ utcTimeToPOSIXSeconds <$> getCurrentTime
+getTimestamp = do
+  timestamp <- liftIO $ getCurrentTime
+  $(logDebugS) "routing.make_timestamp" (T.pack (render (pretty timestamp)))
+  return $ utcTimeToPOSIXSeconds timestamp
 
 {-----------------------------------------------------------------------
 -- Tokens
@@ -204,6 +225,9 @@ insertNode info = do
   t' <- routing (R.insert info t)
   putTable t'
 
+  let logMsg = "Routing table updated: " <> pretty t <> " -> " <> pretty t'
+  $(logDebugS) "insertNode" (T.pack (render logMsg))
+
 {-----------------------------------------------------------------------
 -- Peer storage
 -----------------------------------------------------------------------}
@@ -231,11 +255,17 @@ getPeerList ih = do
 -- Messaging
 -----------------------------------------------------------------------}
 
-(<@>) :: Address ip => KRPC (Query a) (Response b)
+(<@>) :: forall a b ip. Address ip => KRPC (Query a) (Response b)
       => a -> NodeAddr ip -> DHT ip b
 q <@> addr = do
   nid <- getNodeId
+
+  let Method name = method :: Method (Query a) (Response b)
+  let signature = T.decodeUtf8 name <> " @ " <> T.pack (render (pretty addr))
+  $(logDebugS) "queryNode" $ "Query sent | " <> signature
   Response remoteId r <- query (toSockAddr addr) (Query nid q)
+  $(logDebugS) "queryNode" $ "Query recv | " <> signature
+
   insertNode (NodeInfo remoteId addr)
   return r