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
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
|
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE DeriveDataTypeable #-}
module Network.BitTorrent.Exchange.Session
( -- * Session
Session
, Event (..)
, LogFun
, sessionLogger
-- * Construction
, newSession
, closeSession
, withSession
-- * Connection Set
, connect
, establish
-- * Query
, waitMetadata
, takeMetadata
) where
import Control.Applicative
import Control.Concurrent
import Control.Concurrent.Chan.Split as CS
import Control.Concurrent.STM
import Control.Exception hiding (Handler)
import Control.Lens
import Control.Monad.Logger
import Control.Monad.Reader
import Data.ByteString as BS
import Data.ByteString.Lazy as BL
import Data.Conduit
import Data.Conduit.List as CL (iterM)
import Data.Map as M
import Data.Monoid
import Data.Set as S
import Data.Text as T
import Data.Typeable
import Text.PrettyPrint hiding ((<>))
import Text.PrettyPrint.Class
import System.Log.FastLogger (LogStr, ToLogStr (..))
import Data.BEncode as BE
import Data.Torrent as Torrent
import Network.BitTorrent.Internal.Types
import Network.BitTorrent.Address
import Network.BitTorrent.Exchange.Bitfield as BF
import Network.BitTorrent.Exchange.Block as Block
import Network.BitTorrent.Exchange.Connection
import Network.BitTorrent.Exchange.Download as D
import Network.BitTorrent.Exchange.Message as Message
import System.Torrent.Storage
{-----------------------------------------------------------------------
-- Exceptions
-----------------------------------------------------------------------}
data ExchangeError
= InvalidRequest BlockIx StorageFailure
| CorruptedPiece PieceIx
deriving (Show, Typeable)
instance Exception ExchangeError
packException :: Exception e => (e -> ExchangeError) -> IO a -> IO a
packException f m = try m >>= either (throwIO . f) return
{-----------------------------------------------------------------------
-- Session state
-----------------------------------------------------------------------}
-- TODO unmap storage on zero connections
data Cached a = Cached
{ cachedValue :: !a
, cachedData :: BL.ByteString -- keep lazy
}
cache :: BEncode a => a -> Cached a
cache s = Cached s (BE.encode s)
-- | Logger function.
type LogFun = Loc -> LogSource -> LogLevel -> LogStr -> IO ()
--data SessionStatus = Seeder | Leecher
data SessionState
= WaitingMetadata
{ metadataDownload :: MVar MetadataDownload
, metadataCompleted :: MVar InfoDict -- ^ used to unblock waiters
, contentRootPath :: FilePath
}
| HavingMetadata
{ metadataCache :: Cached InfoDict
, contentDownload :: MVar ContentDownload
, contentStorage :: Storage
}
newSessionState :: FilePath -> Either InfoHash InfoDict -> IO SessionState
newSessionState rootPath (Left ih ) = do
WaitingMetadata <$> newMVar def <*> newEmptyMVar <*> pure rootPath
newSessionState rootPath (Right dict) = do
storage <- openInfoDict ReadWriteEx rootPath dict
download <- newMVar $ D.contentDownload (BF.haveNone (totalPieces storage))
(piPieceLength (idPieceInfo dict))
storage
return $ HavingMetadata (cache dict) download storage
closeSessionState :: SessionState -> IO ()
closeSessionState WaitingMetadata {..} = return ()
closeSessionState HavingMetadata {..} = close contentStorage
haveMetadata :: InfoDict -> SessionState -> IO SessionState
haveMetadata dict WaitingMetadata {..} = do
storage <- openInfoDict ReadWriteEx contentRootPath dict
download <- newMVar $ D.contentDownload (BF.haveNone (totalPieces storage))
(piPieceLength (idPieceInfo dict))
storage
return HavingMetadata
{ metadataCache = cache dict
, contentDownload = download
, contentStorage = storage
}
haveMetadata _ s = return s
{-----------------------------------------------------------------------
-- Session
-----------------------------------------------------------------------}
data Session = Session
{ sessionPeerId :: !(PeerId)
, sessionTopic :: !(InfoHash)
, sessionLogger :: !(LogFun)
, sessionEvents :: !(SendPort (Event Session))
, sessionState :: !(MVar SessionState)
------------------------------------------------------------------------
, connectionsPrefs :: !ConnectionPrefs
-- | Connections either waiting for TCP/uTP 'connect' or waiting
-- for BT handshake.
, connectionsPending :: !(TVar (Set (PeerAddr IP)))
-- | Connections successfully handshaked and data transfer can
-- take place.
, connectionsEstablished :: !(TVar (Map (PeerAddr IP) (Connection Session)))
-- | TODO implement choking mechanism
, connectionsUnchoked :: [PeerAddr IP]
-- | Messages written to this channel will be sent to the all
-- connections, including pending connections (but right after
-- handshake).
, connectionsBroadcast :: !(Chan Message)
}
instance EventSource Session where
data Event Session
= ConnectingTo (PeerAddr IP)
| ConnectionEstablished (PeerAddr IP)
| ConnectionAborted
| ConnectionClosed (PeerAddr IP)
| SessionClosed
deriving Show
listen Session {..} = CS.listen sessionEvents
newSession :: LogFun
-> PeerAddr (Maybe IP) -- ^ /external/ address of this peer;
-> FilePath -- ^ root directory for content files;
-> Either InfoHash InfoDict -- ^ torrent info dictionary;
-> IO Session
newSession logFun addr rootPath source = do
let ih = either id idInfoHash source
pid <- maybe genPeerId return (peerId addr)
eventStream <- newSendPort
sState <- newSessionState rootPath source
sStateVar <- newMVar sState
pSetVar <- newTVarIO S.empty
eSetVar <- newTVarIO M.empty
chan <- newChan
return Session
{ sessionPeerId = pid
, sessionTopic = ih
, sessionLogger = logFun
, sessionEvents = eventStream
, sessionState = sStateVar
, connectionsPrefs = def
, connectionsPending = pSetVar
, connectionsEstablished = eSetVar
, connectionsUnchoked = []
, connectionsBroadcast = chan
}
closeSession :: Session -> IO ()
closeSession Session {..} = do
s <- readMVar sessionState
closeSessionState s
{-
hSet <- atomically $ do
pSet <- swapTVar connectionsPending S.empty
eSet <- swapTVar connectionsEstablished S.empty
return pSet
mapM_ kill hSet
-}
withSession :: ()
withSession = error "withSession"
{-----------------------------------------------------------------------
-- Logging
-----------------------------------------------------------------------}
instance MonadLogger (Connected Session) where
monadLoggerLog loc src lvl msg = do
conn <- ask
ses <- asks connSession
addr <- asks connRemoteAddr
let addrSrc = src <> " @ " <> T.pack (render (pretty addr))
liftIO $ sessionLogger ses loc addrSrc lvl (toLogStr msg)
logMessage :: MonadLogger m => Message -> m ()
logMessage msg = logDebugN $ T.pack (render (pretty msg))
logEvent :: MonadLogger m => Text -> m ()
logEvent = logInfoN
{-----------------------------------------------------------------------
-- Connection set
-----------------------------------------------------------------------}
--- Connection status transition:
---
--- pending -> established -> finished -> closed
--- | \|/ /|\
--- \-------------------------------------|
---
--- Purpose of slots:
--- 1) to avoid duplicates
--- 2) connect concurrently
---
-- | Add connection to the pending set.
pendingConnection :: PeerAddr IP -> Session -> STM Bool
pendingConnection addr Session {..} = do
pSet <- readTVar connectionsPending
eSet <- readTVar connectionsEstablished
if (addr `S.member` pSet) || (addr `M.member` eSet)
then return False
else do
modifyTVar' connectionsPending (S.insert addr)
return True
-- | Pending connection successfully established, add it to the
-- established set.
establishedConnection :: Connected Session ()
establishedConnection = do
conn <- ask
addr <- asks connRemoteAddr
Session {..} <- asks connSession
liftIO $ atomically $ do
modifyTVar connectionsPending (S.delete addr)
modifyTVar connectionsEstablished (M.insert addr conn)
-- | Either this or remote peer decided to finish conversation
-- (conversation is alread /established/ connection), remote it from
-- the established set.
finishedConnection :: Connected Session ()
finishedConnection = do
Session {..} <- asks connSession
addr <- asks connRemoteAddr
liftIO $ atomically $ do
modifyTVar connectionsEstablished $ M.delete addr
-- | There are no state for this connection, remove it from the all
-- sets.
closedConnection :: PeerAddr IP -> Session -> STM ()
closedConnection addr Session {..} = do
modifyTVar connectionsPending $ S.delete addr
modifyTVar connectionsEstablished $ M.delete addr
getConnectionConfig :: Session -> IO (ConnectionConfig Session)
getConnectionConfig s @ Session {..} = do
chan <- dupChan connectionsBroadcast
let sessionLink = SessionLink {
linkTopic = sessionTopic
, linkPeerId = sessionPeerId
, linkMetadataSize = Nothing
, linkOutputChan = Just chan
, linkSession = s
}
return ConnectionConfig
{ cfgPrefs = connectionsPrefs
, cfgSession = sessionLink
, cfgWire = mainWire
}
type Finalizer = IO ()
type Runner = (ConnectionConfig Session -> IO ())
runConnection :: Runner -> Finalizer -> PeerAddr IP -> Session -> IO ()
runConnection runner finalize addr set @ Session {..} = do
_ <- forkIO (action `finally` cleanup)
return ()
where
action = do
notExist <- atomically $ pendingConnection addr set
when notExist $ do
cfg <- getConnectionConfig set
runner cfg
cleanup = do
finalize
-- runStatusUpdates status (SS.resetPending addr)
-- TODO Metata.resetPending addr
atomically $ closedConnection addr set
-- | Establish connection from scratch. If this endpoint is already
-- connected, no new connections is created. This function do not block.
connect :: PeerAddr IP -> Session -> IO ()
connect addr = runConnection (connectWire addr) (return ()) addr
-- | Establish connection with already pre-connected endpoint. If this
-- endpoint is already connected, no new connections is created. This
-- function do not block.
--
-- 'PendingConnection' will be closed automatically, you do not need
-- to call 'closePending'.
establish :: PendingConnection -> Session -> IO ()
establish conn = runConnection (acceptWire conn) (closePending conn)
(pendingPeer conn)
-- | Why do we need this message?
type BroadcastMessage = ExtendedCaps -> Message
broadcast :: BroadcastMessage -> Session -> IO ()
broadcast = error "broadcast"
{-----------------------------------------------------------------------
-- Helpers
-----------------------------------------------------------------------}
waitMVar :: MVar a -> IO ()
waitMVar m = withMVar m (const (return ()))
-- This function appear in new GHC "out of box". (moreover it is atomic)
tryReadMVar :: MVar a -> IO (Maybe a)
tryReadMVar m = do
ma <- tryTakeMVar m
maybe (return ()) (putMVar m) ma
return ma
readBlock :: BlockIx -> Storage -> IO (Block BL.ByteString)
readBlock bix @ BlockIx {..} s = do
p <- packException (InvalidRequest bix) $ do readPiece ixPiece s
let chunk = BL.take (fromIntegral ixLength) $
BL.drop (fromIntegral ixOffset) (pieceData p)
if BL.length chunk == fromIntegral ixLength
then return $ Block ixPiece ixOffset chunk
else throwIO $ InvalidRequest bix (InvalidSize ixLength)
-- |
tryReadMetadataBlock :: PieceIx
-> Connected Session (Maybe (Torrent.Piece BS.ByteString, Int))
tryReadMetadataBlock pix = do
Session {..} <- asks connSession
s <- liftIO (readMVar sessionState)
case s of
WaitingMetadata {..} -> error "tryReadMetadataBlock"
HavingMetadata {..} -> error "tryReadMetadataBlock"
sendBroadcast :: PeerMessage msg => msg -> Wire Session ()
sendBroadcast msg = do
Session {..} <- asks connSession
error "sendBroadcast"
-- liftIO $ msg `broadcast` sessionConnections
waitMetadata :: Session -> IO InfoDict
waitMetadata Session {..} = do
s <- readMVar sessionState
case s of
WaitingMetadata {..} -> readMVar metadataCompleted
HavingMetadata {..} -> return (cachedValue metadataCache)
takeMetadata :: Session -> IO (Maybe InfoDict)
takeMetadata Session {..} = do
s <- readMVar sessionState
case s of
WaitingMetadata {..} -> return Nothing
HavingMetadata {..} -> return (Just (cachedValue metadataCache))
{-----------------------------------------------------------------------
-- Triggers
-----------------------------------------------------------------------}
-- | Trigger is the reaction of a handler at some event.
type Trigger = Wire Session ()
interesting :: Trigger
interesting = do
addr <- asks connRemoteAddr
sendMessage (Interested True)
sendMessage (Choking False)
tryFillRequestQueue
fillRequestQueue :: Trigger
fillRequestQueue = do
maxN <- lift getMaxQueueLength
rbf <- use connBitfield
addr <- asks connRemoteAddr
-- blks <- withStatusUpdates $ do
-- n <- getRequestQueueLength addr
-- scheduleBlocks addr rbf (maxN - n)
-- mapM_ (sendMessage . Request) blks
return ()
tryFillRequestQueue :: Trigger
tryFillRequestQueue = do
allowed <- canDownload <$> use connStatus
when allowed $ do
fillRequestQueue
{-----------------------------------------------------------------------
-- Incoming message handling
-----------------------------------------------------------------------}
type Handler msg = msg -> Wire Session ()
handleStatus :: Handler StatusUpdate
handleStatus s = do
connStatus %= over remoteStatus (updateStatus s)
case s of
Interested _ -> return ()
Choking True -> do
addr <- asks connRemoteAddr
-- withStatusUpdates (SS.resetPending addr)
return ()
Choking False -> tryFillRequestQueue
handleAvailable :: Handler Available
handleAvailable msg = do
connBitfield %= case msg of
Have ix -> BF.insert ix
Bitfield bf -> const bf
--thisBf <- getThisBitfield
thisBf <- undefined
case msg of
Have ix
| ix `BF.member` thisBf -> return ()
| otherwise -> interesting
Bitfield bf
| bf `BF.isSubsetOf` thisBf -> return ()
| otherwise -> interesting
handleTransfer :: Handler Transfer
handleTransfer (Request bix) = do
Session {..} <- asks connSession
s <- liftIO $ readMVar sessionState
case s of
WaitingMetadata {..} -> return ()
HavingMetadata {..} -> do
bitfield <- undefined -- getThisBitfield
upload <- canUpload <$> use connStatus
when (upload && ixPiece bix `BF.member` bitfield) $ do
blk <- liftIO $ readBlock bix contentStorage
sendMessage (Message.Piece blk)
handleTransfer (Message.Piece blk) = do
Session {..} <- asks connSession
s <- liftIO $ readMVar sessionState
case s of
WaitingMetadata {..} -> return () -- TODO (?) break connection
HavingMetadata {..} -> do
isSuccess <- undefined -- withStatusUpdates (SS.pushBlock blk storage)
case isSuccess of
Nothing -> liftIO $ throwIO $ userError "block is not requested"
Just isCompleted -> do
when isCompleted $ do
sendBroadcast (Have (blkPiece blk))
-- maybe send not interested
tryFillRequestQueue
handleTransfer (Cancel bix) = filterQueue (not . (transferResponse bix))
where
transferResponse bix (Transfer (Message.Piece blk)) = blockIx blk == bix
transferResponse _ _ = False
{-----------------------------------------------------------------------
-- Metadata exchange
-----------------------------------------------------------------------}
-- TODO introduce new metadata exchange specific exceptions
waitForMetadata :: Trigger
waitForMetadata = do
Session {..} <- asks connSession
needFetch <- undefined --liftIO (isEmptyMVar infodict)
when needFetch $ do
canFetch <- allowed ExtMetadata <$> use connExtCaps
if canFetch
then tryRequestMetadataBlock
else undefined -- liftIO (waitMVar infodict)
tryRequestMetadataBlock :: Trigger
tryRequestMetadataBlock = do
mpix <- lift $ undefined --withMetadataUpdates Metadata.scheduleBlock
case mpix of
Nothing -> error "tryRequestMetadataBlock"
Just pix -> sendMessage (MetadataRequest pix)
handleMetadata :: Handler ExtendedMetadata
handleMetadata (MetadataRequest pix) =
lift (tryReadMetadataBlock pix) >>= sendMessage . mkResponse
where
mkResponse Nothing = MetadataReject pix
mkResponse (Just (piece, total)) = MetadataData piece total
handleMetadata (MetadataData {..}) = do
ih <- asks connTopic
mdict <- lift $ undefined --withMetadataUpdates (Metadata.pushBlock piece ih)
case mdict of
Nothing -> tryRequestMetadataBlock -- not completed, need all blocks
Just dict -> do -- complete, wake up payload fetch
Session {..} <- asks connSession
liftIO $ modifyMVar_ sessionState (haveMetadata dict)
handleMetadata (MetadataReject pix) = do
lift $ undefined -- withMetadataUpdates (Metadata.cancelPending pix)
handleMetadata (MetadataUnknown _ ) = do
logInfoN "Unknown metadata message"
{-----------------------------------------------------------------------
-- Main entry point
-----------------------------------------------------------------------}
acceptRehandshake :: ExtendedHandshake -> Trigger
acceptRehandshake ehs = error "acceptRehandshake"
handleExtended :: Handler ExtendedMessage
handleExtended (EHandshake ehs) = acceptRehandshake ehs
handleExtended (EMetadata _ msg) = handleMetadata msg
handleExtended (EUnknown _ _ ) = logWarnN "Unknown extension message"
handleMessage :: Handler Message
handleMessage KeepAlive = return ()
handleMessage (Status s) = handleStatus s
handleMessage (Available msg) = handleAvailable msg
handleMessage (Transfer msg) = handleTransfer msg
handleMessage (Port n) = error "handleMessage"
handleMessage (Fast _) = error "handleMessage"
handleMessage (Extended msg) = handleExtended msg
exchange :: Wire Session ()
exchange = do
waitForMetadata
bf <- undefined --getThisBitfield
sendMessage (Bitfield bf)
awaitForever handleMessage
mainWire :: Wire Session ()
mainWire = do
lift establishedConnection
Session {..} <- asks connSession
-- lift $ resizeBitfield (totalPieces storage)
logEvent "Connection established"
iterM logMessage =$= exchange =$= iterM logMessage
lift finishedConnection
|