-
Notifications
You must be signed in to change notification settings - Fork 31
/
Torrent.hs
154 lines (126 loc) · 4.38 KB
/
Torrent.hs
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
-- | The following module is responsible for general types used
-- throughout the system.
module Torrent (
-- * Types
InfoHash
, PeerId
, AnnounceURL
, TorrentState(..)
, TorrentInfo(..)
, PieceNum
, PieceSize
, PieceMap
, PiecesDoneMap
, PieceInfo(..)
, BlockSize
, Block(..)
, Capabilities(..)
-- * Interface
, bytesLeft
, defaultBlockSize
, defaultOptimisticSlots
, defaultPort
, mkPeerId
, mkTorrentInfo
)
where
import Control.Applicative
import Control.DeepSeq
import Data.Array
import Data.List
import Data.Maybe (fromMaybe)
import qualified Data.ByteString as B
import qualified Data.Map as M
import Data.Word
import Numeric
import System.Random
import System.Random.Shuffle
import Test.QuickCheck
import Protocol.BCode
import Digest
import Version
-- | The type of Infohashes as used in torrents. These are identifiers
-- of torrents
type InfoHash = Digest
-- | The peerId is the ID of a client. It is used to identify clients
-- from each other
type PeerId = String
-- | The internal type of Announce URLs
type AnnounceURL = B.ByteString
-- | Internal type for a torrent. It identifies a torrent in various places of the system.
data TorrentInfo = TorrentInfo {
infoHash :: InfoHash,
pieceCount :: Int, -- Number of pieces in torrent
announceURLs :: [[AnnounceURL]]
} deriving Show
data TorrentState = Seeding | Leeching
deriving Show
instance NFData TorrentState
----------------------------------------------------------------------
-- Capabilities
data Capabilities = Fast | Extended
deriving (Show, Eq)
-- PIECES
----------------------------------------------------------------------
type PieceNum = Int
type PieceSize = Int
data PieceInfo = PieceInfo {
offset :: !Integer, -- ^ Offset of the piece, might be greater than Int
len :: !Integer, -- ^ Length of piece; usually a small value
digest :: !B.ByteString -- ^ Digest of piece; taken from the .torret file
} deriving (Eq, Show)
type PieceMap = Array PieceNum PieceInfo
-- | The PiecesDoneMap is a map which is true if we have the piece and false otherwise
type PiecesDoneMap = M.Map PieceNum Bool
-- | Return the amount of bytes left on a torrent given what pieces are done and the
-- map of the shape of the torrent in question.
bytesLeft :: PiecesDoneMap -> PieceMap -> Integer
bytesLeft done pm =
foldl' (\accu (k,v) ->
case M.lookup k done of
Just False -> (len v) + accu
_ -> accu) 0 $ Data.Array.assocs pm
-- BLOCKS
----------------------------------------------------------------------
type BlockSize = Int
data Block = Block { blockOffset :: !Int -- ^ offset of this block within the piece
, blockSize :: !BlockSize -- ^ size of this block within the piece
} deriving (Eq, Ord, Show)
instance NFData Block where
rnf (Block bo sz) = rnf bo `seq` rnf sz `seq` ()
instance Arbitrary Block where
arbitrary = Block <$> pos <*> pos
where pos = choose (0, 4294967296 - 1)
defaultBlockSize :: BlockSize
defaultBlockSize = 16384 -- Bytes
-- | Default number of optimistic slots
defaultOptimisticSlots :: Int
defaultOptimisticSlots = 2
-- | Default port to communicate on
defaultPort :: Word16
defaultPort = 1579
-- | Convert a BCode block into its corresponding TorrentInfo block, perhaps
-- failing in the process.
mkTorrentInfo :: BCode -> IO TorrentInfo
mkTorrentInfo bc = do
(ann, np) <- case queryInfo bc of Nothing -> fail "Could not create torrent info"
Just x -> return x
ih <- hashInfoDict bc
let alist = fromMaybe [[ann]] $ announceList bc
-- BEP012 says that lists of URL inside each tier must be shuffled
gen <- newStdGen
let alist' = map (\xs -> shuffle' xs (length xs) gen) alist
return TorrentInfo { infoHash = ih, pieceCount = np, announceURLs = alist'}
where
queryInfo b =
do ann <- announce b
np <- numberPieces b
return (ann, np)
-- | Create a new PeerId for this client
mkPeerId :: StdGen -> PeerId
mkPeerId gen = header ++ take (20 - length header) ranString
where randomList :: Int -> StdGen -> [Int]
randomList n = take n . unfoldr (Just . random)
rs = randomList 10 gen
ranString = concatMap (\i -> showHex (abs i) "") rs
header = "-CT" ++ protoVersion ++ "-"