summaryrefslogtreecommitdiff
path: root/src/Network/BitTorrent/DHT/Routing.hs
blob: 5f00a92462b706ac787313b9ab004d89e74bdf9a (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
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