3 files changed, 172 insertions, 69 deletions

diff --git a/src/Network/BitTorrent/DHT/Query.hs b/src/Network/BitTorrent/DHT/Query.hs
index 9162abdc..68c67900 100644
--- a/src/Network/BitTorrent/DHT/Query.hs
+++ b/src/Network/BitTorrent/DHT/Query.hs
@@ -86,6 +86,7 @@ import Network.KRPC.Method    as KRPC
 import Network.DatagramServer.Mainline (ReflectedIP(..))
 import Network.DatagramServer (QueryFailure(..))
 import Data.Torrent
+import qualified Network.DHT as DHT
 import Network.DHT.Mainline
 import Network.DHT.Routing as R
 import Network.BitTorrent.DHT.Session
@@ -100,6 +101,7 @@ import Network.DatagramServer.Tox
 #endif
 import Network.Address hiding (NodeId)
 import Network.DatagramServer.Types as RPC hiding (Query,Response)
+import Control.Monad.Trans.Control
 
 {-----------------------------------------------------------------------
 --  Handlers
@@ -312,79 +314,42 @@ refreshNodes nid = do
     return ()
   return () -- \$ L.concat nss
 
+logc :: Char -> String -> DHT ip ()
+logc 'D' = $(logDebugS) "insertNode" . T.pack
+logc 'W' = $(logWarnS)  "insertNode" . T.pack
+logc 'I' = $(logInfoS)  "insertNode" . T.pack
+logc c   = $(logInfoS)  "insertNode" . T.pack . (c :) . (':' :)
+
 -- | This operation do not block but acquire exclusive access to
 --   routing table.
 insertNode :: forall ip. Address ip => NodeInfo KMessageOf ip () -> Maybe ReflectedIP -> DHT ip ()
 insertNode info witnessed_ip0 = do
-  var <- asks routingInfo
-  tm <- getTimestamp
-  let showTable = do
-          t <- getTable
-          let logMsg = "Routing table: " <> pPrint t
-          $(logDebugS) "insertNode" (T.pack (render logMsg))
-  let arrival0 = TryInsert info
-      arrival4 = TryInsert (mapAddress fromAddr info) :: Event _ (Maybe IPv4) _
-  $(logDebugS) "insertNode" $ T.pack (show arrival4)
-  maxbuckets <- asks (optBucketCount . options)
-  fallbackid <- asks tentativeNodeId
-  let atomicInsert arrival witnessed_ip = do
-          minfo <- readTVar var
-          let change ip0 = fromMaybe fallbackid $ do
-                            ip <- fromSockAddr ip0 :: Maybe ip
-                            listToMaybe
-                                $ rank id (nodeId $ foreignNode arrival)
-                                $ bep42s ip fallbackid
-          case minfo of
-            Just inf -> do
-              (ps,t') <- R.insert tm arrival $ myBuckets inf
-              writeTVar var $ Just $ inf { myBuckets = t' }
-              return $ do
-                case witnessed_ip of
-                    Just (ReflectedIP ip)
-                      | ip /= myAddress inf
-                        -> $(logInfo) ( T.pack $ L.unwords
-                              $ [ "Possible NAT?"
-                                , show (toSockAddr $ nodeAddr $ foreignNode arrival)
-                                , "reports my address:"
-                                , show ip ] )
-                            -- TODO: Let routing table vote on my IP/NodeId.
-                    _   -> return ()
-                return ps
-            Nothing ->
-                let dropped = return $ do
-                      -- Ignore non-witnessing nodes until somebody tells
-                      -- us our ip address.
-                      $(logWarnS) "insertNode" ("Dropped "
-                          <> T.pack (show (toSockAddr $ nodeAddr $ foreignNode arrival)))
-                      return []
-                in fromMaybe dropped $ do
-                      ReflectedIP ip <- witnessed_ip
-                      let nil = nullTable (change ip) maxbuckets
-                      return $ do
-                        (ps,t') <- R.insert tm arrival nil
-                        let new_info = R.Info t' (change ip) ip
-                        writeTVar var $ Just new_info
-                        return $ do
-                           $(logInfo) ( T.pack $ L.unwords
-                                [ "External IP address:"
-                                , show ip
-                                , "(reported by"
-                                , show (toSockAddr $ nodeAddr $ foreignNode arrival)
-                                    <> ")"
-                                ] )
-                           return ps
-  ps <- join $ liftIO $ atomically $ atomicInsert arrival0 witnessed_ip0
-  showTable
-  _ <- fork $ do
-    myThreadId >>= liftIO . flip labelThread "DHT.insertNode.pingResults"
-    forM_ ps $ \(CheckPing ns)-> do
-      forM_ ns $ \n -> do
-        (b,mip) <- probeNode (nodeAddr n)
-        let alive = PingResult n b
-        $(logDebugS) "insertNode" $ T.pack ("PingResult "++show (nodeId n,b))
-        _ <- join $ liftIO $ atomically $ atomicInsert alive mip
-        showTable
-  return ()
+  bc <- optBucketCount <$> asks options
+  nid <- asks tentativeNodeId
+  logm0 <- embed_ (uncurry logc)
+  let logm c = logm0 . (c,)
+  dht_node_state <- ask -- XXX: This prevents ping probe from modifying the Node state.
+  probe0 <- embed probeNode -- probeNode :: Address ip => NodeAddr ip -> DHT ip (Bool, Maybe ReflectedIP)
+  let probe n = probe0 n >>= runDHT dht_node_state . restoreM
+      changeip ip0 arrival = fromMaybe (DHT.fallbackID params) $ do -- warning: recursive
+                    ip <- fromSockAddr ip0 :: Maybe ip
+                    listToMaybe
+                        $ rank id (nodeId $ foreignNode arrival)
+                        $ bep42s ip (DHT.fallbackID params) -- warning: recursive
+      params = DHT.TableParameters
+        { maxBuckets = bc :: Int
+        , fallbackID = nid :: NodeId KMessageOf
+        , adjustID = changeip :: SockAddr -> Event KMessageOf ip () -> NodeId KMessageOf
+        , logMessage = logm :: Char -> String -> IO ()
+        , pingProbe = probe :: NodeAddr ip -> IO (Bool, Maybe ReflectedIP)
+        }
+  tbl <- asks routingInfo
+  let state = DHT.TableKeeper
+        { routingInfo = tbl
+        , grokNode    = DHT.insertNode params state
+        , grokAddress = \_ _ -> return () -- :: Maybe SockAddr -> ReflectedIP -> IO ()
+        }
+  liftIO $ DHT.insertNode params state info witnessed_ip0
 
 -- | Throws exception if node is not responding.
 queryNode :: forall a b ip. Address ip => KRPC (Query a) (Response b)
diff --git a/src/Network/DHT.hs b/src/Network/DHT.hs
new file mode 100644
index 00000000..0dab29cd
--- /dev/null
+++ b/src/Network/DHT.hs
@@ -0,0 +1,125 @@
+{-# LANGUAGE CPP, ScopedTypeVariables, PartialTypeSignatures, FlexibleContexts #-}
+module Network.DHT
+    ( -- makeTableKeeper
+    -- , TableKeeper(..)
+      module Network.DHT -- for now
+    , module Network.DHT.Types
+    ) where
+
+import Data.Bits
+import Data.Maybe
+import Data.Monoid
+import Network.Address
+import Network.DHT.Types
+import Network.DatagramServer.Types
+import Network.DHT.Routing
+import Control.Concurrent.STM
+#ifdef THREAD_DEBUG
+import Control.Concurrent.Lifted.Instrument
+#else
+import GHC.Conc (labelThread)
+import Control.Concurrent.Lifted
+#endif
+import Text.PrettyPrint as PP hiding ((<>), ($$))
+import Text.PrettyPrint.HughesPJClass hiding ((<>),($$))
+
+import Control.Monad
+import Data.Time.Clock       (getCurrentTime)
+import Data.Time.Clock.POSIX (utcTimeToPOSIXSeconds)
+
+data TableKeeper msg ip u = TableKeeper
+    { routingInfo  :: TVar    (Maybe (Info msg ip u))
+    , grokNode :: NodeInfo msg ip u -> Maybe (ReflectedIP) -> IO ()
+    , grokAddress :: Maybe SockAddr -> ReflectedIP -> IO ()
+    }
+
+makeTableKeeper :: forall msg ip u.
+                ( Address ip
+                , Show u
+                , Show (NodeId msg)
+                , Ord (NodeId msg)
+                , FiniteBits (NodeId msg)
+                ) => TableParameters msg ip u -> IO (TableKeeper msg ip u)
+makeTableKeeper param@TableParameters{..} = do
+    error "TODO makeTableKeeper" -- kick off table-updating thread
+    ri <- atomically (newTVar Nothing)
+    let tk = TableKeeper{ routingInfo  = ri
+                        , grokNode    = insertNode param tk
+                        , grokAddress = error "todo"
+                        }
+    return tk
+
+atomicInsert :: ( Eq ip, Address ip, Ord (NodeId msg), FiniteBits (NodeId msg)
+                ) => TableParameters msg ip u -> TableKeeper msg ip u -> Timestamp -> Event msg ip u -> Maybe ReflectedIP -> STM (IO [CheckPing msg ip u])
+atomicInsert param@TableParameters{..} state tm arrival witnessed_ip = do
+    minfo <- readTVar (routingInfo state)
+    case minfo of
+        Just inf -> do
+          (ps,t') <- insert tm arrival $ myBuckets inf
+          writeTVar (routingInfo state) $ Just $ inf { myBuckets = t' }
+          return $ do
+            case witnessed_ip of
+                Just (ReflectedIP ip)
+                  | ip /= myAddress inf
+                    -> logMessage 'I' $ unwords
+                          $ [ "Possible NAT?"
+                            , show (toSockAddr $ nodeAddr $ foreignNode arrival)
+                            , "reports my address:"
+                            , show ip ]
+                        -- TODO: Let routing table vote on my IP/NodeId.
+                _   -> return ()
+            return ps
+        Nothing ->
+            let dropped = return $ do
+                  -- Ignore non-witnessing nodes until somebody tells
+                  -- us our ip address.
+                  logMessage 'W' ("Dropped " ++ show (toSockAddr $ nodeAddr $ foreignNode arrival))
+                  return []
+            in fromMaybe dropped $ do
+                  ReflectedIP ip <- witnessed_ip
+                  let nil = nullTable (adjustID ip arrival) maxBuckets
+                  return $ do
+                    (ps,t') <- insert tm arrival nil
+                    let new_info = Info t' (adjustID ip arrival) ip
+                    writeTVar (routingInfo state) $ Just new_info
+                    return $ do
+                       logMessage 'I' $ unwords
+                            [ "External IP address:"
+                            , show ip
+                            , "(reported by"
+                            , show (toSockAddr $ nodeAddr $ foreignNode arrival)
+                                ++ ")"
+                            ]
+                       return ps
+
+-- | This operation do not block but acquire exclusive access to
+--   routing table.
+insertNode :: forall msg ip u.
+                ( Address ip
+                , Show u
+                , Show (NodeId msg)
+                , Ord (NodeId msg)
+                , FiniteBits (NodeId msg)
+                ) => TableParameters msg ip u -> TableKeeper msg ip u -> NodeInfo msg ip u -> Maybe ReflectedIP -> IO ()
+insertNode param@TableParameters{..} state info witnessed_ip0 = do
+  tm <- utcTimeToPOSIXSeconds <$> getCurrentTime -- Network.DHT.Routing.TimeStamp = POSIXTime
+  let showTable = do
+        t <- atomically $ fmap myBuckets <$> readTVar (routingInfo state)
+        let logMsg = "Routing table: " <> pPrint t
+        logMessage 'D' (render logMsg)
+  let arrival = TryInsert info
+  logMessage 'D' $ show ( TryInsert (mapAddress fromAddr info) :: Event _ (Maybe IPv4) _ )
+  ps <- join $ atomically $ atomicInsert param state tm arrival witnessed_ip0
+  showTable
+  _ <- fork $ do
+    myThreadId >>= flip labelThread "DHT.insertNode.pingResults"
+    forM_ ps $ \(CheckPing ns)-> do
+      forM_ ns $ \n -> do
+        (b,mip) <- pingProbe (nodeAddr n)
+        let alive = PingResult n b
+        logMessage 'D' $ "PingResult "++show (nodeId n,b)
+        _ <- join $ atomically $ atomicInsert param state tm alive mip
+        showTable
+        return ()
+  return ()
+
diff --git a/src/Network/DHT/Types.hs b/src/Network/DHT/Types.hs
new file mode 100644
index 00000000..ed2dc175
--- /dev/null
+++ b/src/Network/DHT/Types.hs
@@ -0,0 +1,13 @@
+module Network.DHT.Types where
+
+import Network.Socket (SockAddr)
+import Network.DatagramServer.Types
+import Network.DHT.Routing
+
+data TableParameters msg ip u = TableParameters
+    { maxBuckets :: Int
+    , fallbackID :: NodeId msg
+    , pingProbe :: NodeAddr ip -> IO (Bool, Maybe ReflectedIP)
+    , logMessage :: Char -> String -> IO ()
+    , adjustID :: SockAddr -> Event msg ip u -> NodeId msg
+    }