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
|
{-# LANGUAGE BangPatterns #-}
module Network.BitTorrent.DHT.ContactInfo
( PeerStore
, Network.BitTorrent.DHT.ContactInfo.lookup
, Network.BitTorrent.DHT.ContactInfo.freshPeers
, Network.BitTorrent.DHT.ContactInfo.insertPeer
, knownSwarms
) where
import Control.Applicative
import Data.Default
import Data.List as L
import Data.Maybe
import Data.HashMap.Strict as HM
import Data.Serialize
import Data.Wrapper.PSQ as PSQ
import Data.Time.Clock.POSIX
import Data.ByteString (ByteString)
import Data.Word
import Network.Socket (SockAddr(..))
import Data.Torrent
import Network.BitTorrent.Address
-- {-
-- import Data.HashMap.Strict as HM
--
-- import Data.Torrent.InfoHash
-- import Network.BitTorrent.Address
--
-- -- increase prefix when table is too large
-- -- decrease prefix when table is too small
-- -- filter outdated peers
--
-- {-----------------------------------------------------------------------
-- -- PeerSet
-- -----------------------------------------------------------------------}
--
-- type PeerSet a = [(PeerAddr a, NodeInfo a, Timestamp)]
--
-- -- compare PSQueue vs Ordered list
--
-- takeNewest :: PeerSet a -> [PeerAddr a]
-- takeNewest = undefined
--
-- dropOld :: Timestamp -> PeerSet a -> PeerSet a
-- dropOld = undefined
--
-- insert :: PeerAddr a -> Timestamp -> PeerSet a -> PeerSet a
-- insert = undefined
--
-- type Mask = Int
-- type Size = Int
-- type Timestamp = Int
--
-- {-----------------------------------------------------------------------
-- -- InfoHashMap
-- -----------------------------------------------------------------------}
--
-- -- compare handwritten prefix tree versus IntMap
--
-- data Tree a
-- = Nil
-- | Tip !InfoHash !(PeerSet a)
-- | Bin !InfoHash !Mask !Size !Timestamp (Tree a) (Tree a)
--
-- insertTree :: InfoHash -> a -> Tree a -> Tree a
-- insertTree = undefined
--
-- type Prio = Int
--
-- --shrink :: ContactInfo ip -> Int
-- shrink Nil = Nil
-- shrink (Tip _ _) = undefined
-- shrink (Bin _ _) = undefined
--
-- {-----------------------------------------------------------------------
-- -- InfoHashMap
-- -----------------------------------------------------------------------}
--
-- -- compare new design versus HashMap
--
-- data IntMap k p a
-- type ContactInfo = Map InfoHash Timestamp (Set (PeerAddr IP) Timestamp)
--
-- data ContactInfo ip = PeerStore
-- { maxSize :: Int
-- , prefixSize :: Int
-- , thisNodeId :: NodeId
--
-- , count :: Int -- ^ Cached size of the 'peerSet'
-- , peerSet :: HashMap InfoHash [PeerAddr ip]
-- }
--
-- size :: ContactInfo ip -> Int
-- size = undefined
--
-- prefixSize :: ContactInfo ip -> Int
-- prefixSize = undefined
--
-- lookup :: InfoHash -> ContactInfo ip -> [PeerAddr ip]
-- lookup = undefined
--
-- insert :: InfoHash -> PeerAddr ip -> ContactInfo ip -> ContactInfo ip
-- insert = undefined
--
-- -- | Limit in size.
-- prune :: NodeId -> Int -> ContactInfo ip -> ContactInfo ip
-- prune pref targetSize Nil = Nil
-- prune pref targetSize (Tip _ _) = undefined
--
-- -- | Remove expired entries.
-- splitGT :: Timestamp -> ContactInfo ip -> ContactInfo ip
-- splitGT = undefined
-- -}
-- | Storage used to keep track a set of known peers in client,
-- tracker or DHT sessions.
newtype PeerStore ip = PeerStore (HashMap InfoHash (SwarmData ip))
type Timestamp = POSIXTime
data SwarmData ip = SwarmData
{ peers :: !(PSQ (PeerAddr ip) Timestamp)
, name :: !(Maybe ByteString)
}
newtype SerializeAddress a = SerializeAddress { unserializeAddress :: a }
instance Address a => Serialize (SerializeAddress a) where
get = SerializeAddress <$> do
c <- get
case (c::Word8) of
0x34 -> do ip4 <- get
return $ fromJust $ fromAddr (ip4::IPv4)
0x36 -> do ip6 <- get
return $ fromJust $ fromAddr (ip6::IPv6)
_ -> return $ error "cannot deserialize non-IP SerializeAddress"
put (SerializeAddress a)
| Just ip4 <- fromAddr a
= put (0x34::Word8) >> put (ip4::IPv4)
| Just ip6 <- fromAddr a
= put (0x36::Word8) >> put (ip6::IPv6)
| otherwise = return $ error "cannot serialize non-IP SerializeAddress"
instance (Ord ip, Address ip) => Serialize (SwarmData ip) where
get = flip SwarmData <$> get
<*> ( PSQ.fromList . L.map parseAddr <$> get )
where
parseAddr addr = (unserializeAddress <$> addr)
:-> 0
put SwarmData{..} = do
put name
put $ L.map (\(addr :-> _) -> (SerializeAddress <$> addr))
-- XXX: should we serialize the timestamp?
$ PSQ.toList peers
knownSwarms :: PeerStore ip -> [ (InfoHash, Int, Maybe ByteString) ]
knownSwarms (PeerStore m) = L.map (\(ih,SwarmData q n) -> (ih,PSQ.size q,n)) $ HM.toList m
swarmSingleton :: Ord ip => PeerAddr ip -> SwarmData ip
swarmSingleton a = SwarmData
{ peers = PSQ.singleton a 0
, name = Nothing }
swarmInsert :: Ord ip => SwarmData ip -> SwarmData ip -> SwarmData ip
swarmInsert old new = SwarmData
{ peers = L.foldl' (\q (a :-> t) -> PSQ.insertWith (\p _ -> p) a t q) (peers old) (PSQ.toList $ peers new)
, name = name new <|> name old -- TODO: decodeUtf8' check
}
isSwarmOccupied SwarmData{..} = not $ PSQ.null peers
-- | Empty store.
instance Default (PeerStore a) where
def = PeerStore HM.empty
{-# INLINE def #-}
-- | Monoid under union operation.
instance Ord a => Monoid (PeerStore a) where
mempty = def
{-# INLINE mempty #-}
mappend (PeerStore a) (PeerStore b) =
PeerStore (HM.unionWith swarmInsert a b)
{-# INLINE mappend #-}
-- | Can be used to store peers between invocations of the client
-- software.
instance (Ord a, Address a) => Serialize (PeerStore a) where
get = PeerStore . HM.fromList <$> get
put (PeerStore m) = put (L.filter (isSwarmOccupied . snd) $ HM.toList m)
-- | Returns all peers associated with a given info hash.
lookup :: Ord a => InfoHash -> PeerStore a -> [PeerAddr a]
lookup ih (PeerStore m) = maybe [] (PSQ.keys . peers) $ HM.lookup ih m
batchSize = 64
-- | Used in 'get_peers' DHT queries.
freshPeers :: Ord a => InfoHash -> Timestamp -> PeerStore a -> ([PeerAddr a], PeerStore a)
freshPeers ih tm (PeerStore m) = fromMaybe ([],PeerStore m) $ do
swarm <- HM.lookup ih m
let ps0 = take batchSize $ unfoldr (incomp minView) (peers swarm)
peers' = case reverse ps0 of
(_,psq):_ -> psq
_ -> peers swarm
ps = L.map (key . fst) ps0
m' = HM.insert ih swarm { peers = L.foldl' (\q p -> PSQ.insert p tm q) peers' ps } m
return $! (ps,PeerStore m')
incomp :: (x -> Maybe (r,x)) -> x -> Maybe ((r,x),x)
incomp !f !x = do
(result,x') <- f x
pure $! ( (result,x'), x' )
-- | Used in 'announce_peer' DHT queries.
insertPeer :: Ord a => InfoHash -> Maybe ByteString -> PeerAddr a -> PeerStore a -> PeerStore a
insertPeer !ih !name !a !(PeerStore m) = seq a' $ PeerStore (HM.insertWith swarmInsert ih a' m)
where
a' = SwarmData { peers = PSQ.singleton a 0
, name = name }
|
