3 files changed, 330 insertions, 5 deletions

diff --git a/src/Network/BitTorrent/Core.hs b/src/Network/BitTorrent/Core.hs
index 7b8ff07d..9cfb3dd7 100644
--- a/src/Network/BitTorrent/Core.hs
+++ b/src/Network/BitTorrent/Core.hs
@@ -19,5 +19,6 @@ module Network.BitTorrent.Core
 import Data.IP
 
 import Network.BitTorrent.Core.Fingerprint as Core
+import Network.BitTorrent.Core.Node        as Core
 import Network.BitTorrent.Core.PeerId      as Core
 import Network.BitTorrent.Core.PeerAddr    as Core
diff --git a/src/Network/BitTorrent/Core/Node.hs b/src/Network/BitTorrent/Core/Node.hs
index e93c3586..5098d260 100644
--- a/src/Network/BitTorrent/Core/Node.hs
+++ b/src/Network/BitTorrent/Core/Node.hs
@@ -4,6 +4,7 @@
 module Network.BitTorrent.Core.Node
        (  -- * Node ID
          NodeId
+       , testIdBit
        , genNodeId
 
          -- * Node address
@@ -16,9 +17,11 @@ module Network.BitTorrent.Core.Node
 import Control.Applicative
 import Data.Aeson (ToJSON, FromJSON)
 import Data.Aeson.TH
+import Data.Bits
 import Data.ByteString as BS
 import Data.BEncode as BE
 import Data.Serialize as S
+import Data.Word
 import Network
 import System.Entropy
 
@@ -28,11 +31,12 @@ import Network.BitTorrent.Core.PeerAddr ()
 {-----------------------------------------------------------------------
 --  Node id
 -----------------------------------------------------------------------}
+-- TODO more compact representation ('ShortByteString's?)
 
 -- | Normally, /this/ node id should we saved between invocations of
 -- the client software.
 newtype NodeId = NodeId ByteString
-  deriving (Show, Eq, FromJSON, ToJSON)
+  deriving (Show, Eq, Ord, FromJSON, ToJSON)
 
 nodeIdSize :: Int
 nodeIdSize = 20
@@ -43,6 +47,15 @@ instance Serialize NodeId where
   put (NodeId bs) = putByteString bs
   {-# INLINE put #-}
 
+-- | Test if the nth bit is set.
+testIdBit :: NodeId -> Word -> Bool
+testIdBit (NodeId bs) i
+  | fromIntegral i < nodeIdSize * 8
+  , (q, r) <- quotRem (fromIntegral i) 8
+  = testBit (BS.index bs q) r
+  |     otherwise      = False
+{-# INLINE testIdBit #-}
+
 -- TODO WARN is the 'system' random suitable for this?
 -- | Generate random NodeID used for the entire session.
 --   Distribution of ID's should be as uniform as possible.
@@ -50,8 +63,6 @@ instance Serialize NodeId where
 genNodeId :: IO NodeId
 genNodeId = NodeId <$> getEntropy nodeIdSize
 
-type Distance = NodeId
-
 {-----------------------------------------------------------------------
 --  Node address
 -----------------------------------------------------------------------}
@@ -59,7 +70,7 @@ type Distance = NodeId
 data NodeAddr a = NodeAddr
   { nodeHost ::                !a
   , nodePort :: {-# UNPACK #-} !PortNumber
-  } deriving (Show, Eq)
+  } deriving (Show, Eq, Ord)
 
 $(deriveJSON omitRecordPrefix ''NodeAddr)
 
@@ -82,7 +93,7 @@ instance BEncode a => BEncode (NodeAddr a) where
 data NodeInfo a = NodeInfo
   { nodeId   :: !NodeId
   , nodeAddr :: !(NodeAddr a)
-  } deriving (Show, Eq)
+  } deriving (Show, Eq, Ord)
 
 $(deriveJSON omitRecordPrefix ''NodeInfo)
 
diff --git a/src/Network/BitTorrent/DHT/Routing.hs b/src/Network/BitTorrent/DHT/Routing.hs
new file mode 100644
index 00000000..5f00a924
--- /dev/null
+++ b/src/Network/BitTorrent/DHT/Routing.hs
@@ -0,0 +1,313 @@
+-- |
+--   Copyright   :  (c) Sam Truzjan 2013
+--   License     :  BSD3
+--   Maintainer  :  pxqr.sta@gmail.com
+--   Stability   :  experimental
+--   Portability :  portable
+--
+{-# LANGUAGE RecordWildCards #-}
+{-# LANGUAGE TypeOperators   #-}
+{-# LANGUAGE DeriveGeneric   #-}
+module Network.BitTorrent.DHT.Routing
+       ( -- * Routing table
+         Table
+       , BucketCount
+
+         -- * Routing
+       , Routing
+       , runRouting
+
+         -- * Query
+       , thisId
+       , Network.BitTorrent.DHT.Routing.size
+       , Network.BitTorrent.DHT.Routing.depth
+       , K
+       , Network.BitTorrent.DHT.Routing.kclosest
+       , Network.BitTorrent.DHT.Routing.kclosestHash
+
+         -- * Construction
+       , Network.BitTorrent.DHT.Routing.nullTable
+       , Network.BitTorrent.DHT.Routing.insert
+       ) where
+
+import Control.Applicative hiding (empty)
+import Control.Arrow
+import Control.Monad
+import Data.List as L hiding (insert)
+import Data.Maybe
+import Data.Function
+import Data.PSQueue as PSQ
+import Data.Serialize as S hiding (Result, Done)
+import Data.Time
+import Data.Time.Clock.POSIX
+import Data.Word
+import GHC.Generics
+
+import Data.Torrent.InfoHash
+import Network.BitTorrent.Core
+
+{-
+-- | Routing tree should contain key -> value pairs in this way:
+--
+--     * More keys that near to our node key, and less keys that far
+--     from our node key.
+--
+--     * Tree might be saturated. If this happen we can only update
+--     buckets, but we can't add new buckets.
+--
+--   Instead of using ordinary binary tree and keep track is it
+--   following restrictions above (that's somewhat non-trivial) we
+--   store distance -> value keys. This lead to simple data structure
+--   that actually isomorphic to non-empty list. So we first map our
+--   keys to distances using our node ID and store them in tree. When
+--   we need to extract a pair we map distances to keys back, again
+--   using our node ID. This normalization happen in routing table.
+--
+data Tree k v
+  = Tip (Bucket k v)
+  | Bin (Tree k v)   (Bucket k v)
+
+empty :: Int -> Tree k v
+empty = Tip . Bucket.empty
+
+insert :: Applicative f => Bits k
+       => (v -> f Bool) -> (k, v) -> Tree k v -> f (Tree k v)
+insert ping (k, v) = go 0
+  where
+    go n (Tip bucket)
+      | isFull bucket, (near, far) <- split n bucket
+                          = pure (Tip near `Bin` far)
+      |     otherwise     = Tip <$> Bucket.insert ping (k, v) bucket
+
+    go n (Bin near far)
+      | k `testBit` n = Bin <$> pure near <*> Bucket.insert ping (k, v) far
+      | otherwise     = Bin <$> go (succ n) near <*> pure far
+-}
+
+{-----------------------------------------------------------------------
+--  Insertion
+-----------------------------------------------------------------------}
+
+type Timestamp = POSIXTime
+type PingInterval = POSIXTime
+
+data Routing ip result
+  = Full      result
+  | Done     (Timestamp -> result)
+  | Refresh  (NodeAddr ip) (([NodeInfo ip], Timestamp) -> Routing ip result)
+  | NeedPing (NodeAddr ip) (Maybe Timestamp            -> Routing ip result)
+
+instance Functor (Routing ip) where
+  fmap f (Full          r) = Full          (     f   r)
+  fmap f (Done          r) = Done          (     f . r)
+  fmap f (Refresh  addr g) = Refresh  addr (fmap f . g)
+  fmap f (NeedPing addr g) = NeedPing addr (fmap f . g)
+
+runRouting :: (Monad m, Eq ip)
+             => (NodeAddr ip -> m Bool)          -- ^ ping_node
+             -> (NodeId      -> m [NodeInfo ip]) -- ^ find_nodes
+             -> m Timestamp                      -- ^ timestamper
+             -> Routing ip f
+             -> m f                              -- ^ result
+runRouting ping_node find_nodes timestamp = go
+  where
+    go (Full          r) = return r
+    go (Done          f) = liftM f timestamp
+    go (NeedPing addr f) = do
+      pong <- ping_node addr
+      if pong
+        then do
+             time <- timestamp
+             go (f (Just time))
+        else go (f Nothing)
+
+    go (Refresh nodes f) = do
+      let nid = undefined
+      go (f undefined)
+
+{-----------------------------------------------------------------------
+    Bucket
+-----------------------------------------------------------------------}
+
+-- | Timestamp - last time this node is pinged.
+type NodeEntry ip = Binding (NodeInfo ip) Timestamp
+
+instance (Serialize k, Serialize v) => Serialize (Binding k v) where
+  get = (:->) <$> get <*> get
+  put (k :-> v) = put k >> put v
+
+-- TODO instance Pretty where
+
+-- | Most clients use this value for maximum bucket size.
+defaultBucketSize :: Int
+defaultBucketSize = 20
+
+-- | Bucket is also limited in its length — thus it's called k-bucket.
+--   When bucket becomes full we should split it in two lists by
+--   current span bit. Span bit is defined by depth in the routing
+--   table tree. Size of the bucket should be choosen such that it's
+--   very unlikely that all nodes in bucket fail within an hour of
+--   each other.
+--
+type Bucket ip = PSQ (NodeInfo ip) Timestamp
+
+instance (Serialize k, Serialize v, Ord k, Ord v)
+       => Serialize (PSQ k v) where
+  get = PSQ.fromList <$> get
+  put = put . PSQ.toList
+
+-- | Get the most recently changed node entry, if any.
+lastChanged :: Eq ip => Bucket ip -> Maybe (NodeEntry ip)
+lastChanged bucket
+  | L.null timestamps = Nothing
+  |      otherwise    = Just (L.maximumBy (compare `on` prio) timestamps)
+  where
+    timestamps = PSQ.toList bucket
+
+leastRecently :: Eq ip => Bucket ip -> Maybe (NodeEntry ip, Bucket ip)
+leastRecently = minView
+
+-- | Update interval, in seconds.
+delta :: NominalDiffTime
+delta = 15
+
+-- | Max bucket size, in nodes.
+type Alpha = Int
+
+defaultAlpha :: Int
+defaultAlpha = 8
+
+insertNode :: Eq ip => NodeInfo ip -> Bucket ip -> ip `Routing` Bucket ip
+insertNode info bucket
+  -- just update timestamp if a node is already in bucket
+  | Just _ <- PSQ.lookup info bucket
+  = Done $ \ t -> PSQ.insertWith max info t bucket
+
+  -- update the all bucket if it is too outdated
+  | Just (NodeInfo {..} :-> lastSeen) <- lastChanged bucket
+  , lastSeen > delta
+  = Refresh nodeAddr $ \ (infos, t) ->
+      insertNode info $
+      L.foldr (\ x -> PSQ.insertWith max x t) bucket infos
+
+  -- update questionable nodes, if any; then try to insert our new node
+  -- this case can remove bad nodes from bucket, so we can insert a new one
+  | Just ((old @ NodeInfo {..} :-> leastSeen), rest) <- leastRecently bucket
+  , leastSeen > delta
+  = NeedPing nodeAddr $ insertNode info . maybe rest
+        (\ pong_time -> PSQ.insert old pong_time bucket)
+
+  -- bucket is good, but not full => we can insert a new node
+  | PSQ.size bucket < defaultAlpha  = Done (\ t -> PSQ.insert info t bucket)
+
+  -- bucket is full of good nodes => ignore new node
+  | otherwise = Full bucket
+
+type BitIx = Word
+
+split :: Eq ip => BitIx -> Bucket ip -> (Bucket ip, Bucket ip)
+split i = (PSQ.fromList *** PSQ.fromList) . partition spanBit . PSQ.toList
+  where
+    spanBit entry = testIdBit (nodeId (key entry)) i
+
+{-----------------------------------------------------------------------
+--  Table
+-----------------------------------------------------------------------}
+
+type BucketCount = Word
+
+defaultBucketCount :: BucketCount
+defaultBucketCount = 20
+
+data Table ip
+  = Tip  NodeId BucketCount (Bucket ip)
+  | Zero (Table  ip) (Bucket ip)
+  | One  (Bucket ip) (Table  ip)
+    deriving Generic
+
+instance Serialize NominalDiffTime where
+  put = putWord32be . fromIntegral   . fromEnum
+  get = (toEnum     . fromIntegral) <$> getWord32be
+
+-- | Normally, routing table should we saved between invocations of
+-- the client software. Note that you don't need store /this/ 'NodeId'
+-- since it is included in routing table.
+instance (Eq ip, Serialize ip) => Serialize (Table ip)
+
+nullTable :: Eq ip => NodeId -> Table ip
+nullTable nid = Tip nid defaultBucketCount PSQ.empty
+
+thisId :: Table ip -> NodeId
+thisId (Tip  nid _ _) = nid
+thisId (Zero table _) = thisId table
+thisId (One _  table) = thisId table
+
+-- | Get number of nodes in the table.
+size :: Table ip -> Int
+size = go
+  where
+    go (Tip _ _ bucket) = PSQ.size bucket
+    go (Zero t  bucket) = go t + PSQ.size bucket
+    go (One bucket t  ) = PSQ.size bucket + go t
+
+-- | Get number of buckets in the table.
+depth :: Table ip -> BucketCount
+depth = go
+  where
+    go (Tip _ _ _) = 1
+    go (Zero t  _) = succ (go t)
+    go (One  _  t) = succ (go t)
+
+lookupBucket :: NodeId -> Table ip -> Maybe (Bucket ip)
+lookupBucket nid = go 0
+  where
+    go i (Zero table bucket)
+      |  testIdBit nid i  = pure bucket
+      |     otherwise     = go (succ i) table
+    go i (One  bucket table)
+      |  testIdBit nid i  = go (succ i) table
+      |     otherwise     = pure bucket
+    go _ (Tip _ _ bucket) = pure bucket
+
+type K = Int
+
+-- | Used in 'find_node' queries.
+kclosest :: Eq ip => K -> NodeId -> Table ip -> [NodeInfo ip]
+kclosest k nid = L.map key . PSQ.toList . fromMaybe PSQ.empty
+               . lookupBucket nid
+
+coerseId :: (Serialize a, Serialize b) => a -> b
+coerseId = either (error msg) id . S.decode . S.encode
+  where
+    msg = "coerseId: impossible"
+
+-- | Used in 'get_peers' queries.
+kclosestHash :: Eq a => Alpha -> InfoHash -> Table a -> [NodeInfo a]
+kclosestHash k nid t = kclosest k (coerseId nid) t
+
+{-----------------------------------------------------------------------
+--  Routing
+-----------------------------------------------------------------------}
+
+splitTip :: Eq ip => NodeId -> BucketCount -> BitIx -> Bucket ip -> Table ip
+splitTip nid n i bucket
+  | testIdBit nid i = (One  zeros (Tip nid (pred n) ones))
+  |    otherwise    = (Zero (Tip nid (pred n) zeros) ones)
+  where
+    (zeros, ones) = split i bucket
+
+-- | Used in each query.
+insert :: Eq ip => NodeInfo ip -> Table ip -> ip `Routing` Table ip
+insert info @ NodeInfo {..} = go (0 :: BitIx)
+  where
+    go i (Zero table  bucket)
+      | testIdBit nodeId i  =   Zero  table   <$> insertNode info bucket
+      |     otherwise       = (`Zero` bucket) <$> go (succ i) table
+    go i (One  bucket table )
+      | testIdBit nodeId i  =   One  bucket   <$> go (succ i) table
+      |     otherwise       = (`One` table)   <$> insertNode info bucket
+    go i (Tip nid n bucket) = case insertNode info bucket of
+      Full kbucket
+        |   n == 0  -> Tip nid n <$> Full kbucket
+        | otherwise -> go (succ i) (splitTip nid n i kbucket)
+      result -> Tip nid n <$> result