1 {-# LANGUAGE OverloadedStrings #-}
2 module FuncTorrent.Peer
7 import Prelude hiding (lookup, concat, replicate, splitAt, take)
9 import System.IO (Handle, BufferMode(..), hSetBuffering)
10 import Data.ByteString (ByteString, pack, unpack, concat, hGet, hPut, singleton, take, empty)
11 import Data.ByteString.Lazy (fromStrict, fromChunks, toStrict)
12 import qualified Data.ByteString.Char8 as BC (replicate, pack, length)
13 import Network (connectTo, PortID(..))
14 import Data.Binary (Binary(..), decode, encode)
15 import Data.Binary.Put (putWord32be, putWord16be, putWord8)
16 import Data.Binary.Get (getWord32be, getWord16be, getWord8, runGet)
17 import Control.Monad (replicateM, liftM, forM)
18 import Control.Applicative ((<$>), liftA3)
20 import Data.Word (Word8)
21 import Data.Map (Map, fromList, toList, (!), mapWithKey, adjust)
22 import qualified Crypto.Hash.SHA1 as SHA1 (hash)
24 import FuncTorrent.Metainfo (Info(..), Metainfo(..))
25 import FuncTorrent.Utils (splitN, splitNum)
26 import FuncTorrent.Fileops (createDummyFile, writeFileAtOffset)
32 -- PeerState is a misnomer
33 data PeerState = PeerState { handle :: Handle
36 , meInterested :: Bool
38 , heInterested :: Bool}
40 data PieceDlState = Pending
45 -- todo - map with index to a new data structure (peers who have that piece amd state)
46 data PieceData = PieceData { peers :: [Peer] -- ^ list of peers who have this piece
47 , state :: PieceDlState -- ^ state of the piece from download perspective.
48 , hash :: ByteString -- ^ piece hash
49 , len :: Integer } -- ^ piece length
51 -- which piece is with which peers
52 type PieceMap = Map Integer PieceData
54 -- | Peer is a PeerID, IP address, port tuple
55 data Peer = Peer ID IP Port
58 data PeerMsg = KeepAliveMsg
64 | BitFieldMsg ByteString
65 | RequestMsg Integer Integer Integer
66 | PieceMsg Integer Integer ByteString
67 | CancelMsg Integer Integer Integer
71 -- Make the initial Piece map, with the assumption that no peer has the
72 -- piece and that every piece is pending download.
73 mkPieceMap :: Integer -> ByteString -> [Integer] -> PieceMap
74 mkPieceMap numPieces pieceHash pLengths = fromList kvs
75 where kvs = [(i, PieceData { peers = []
79 | (i, h, pLen) <- zip3 [0..numPieces] hashes pLengths]
80 hashes = splitN 20 pieceHash
82 havePiece :: PieceMap -> Integer -> Bool
84 state (pm ! index) == Have
86 genHandShakeMsg :: ByteString -> String -> ByteString
87 genHandShakeMsg infoHash peer_id = concat [pstrlen, pstr, reserved, infoHash, peerID]
88 where pstrlen = singleton 19
89 pstr = BC.pack "BitTorrent protocol"
90 reserved = BC.replicate 8 '\0'
91 peerID = BC.pack peer_id
93 connectToPeer :: Peer -> IO Handle
94 connectToPeer peer@(Peer _ ip port) = do
95 h <- connectTo ip (PortNumber (fromIntegral port))
96 hSetBuffering h LineBuffering
99 doHandShake :: Handle -> Peer -> ByteString -> String -> IO ()
100 doHandShake h peer@(Peer _ ip port) infoHash peerid = do
101 let hs = genHandShakeMsg infoHash peerid
103 putStrLn $ "--> handhake to peer: " ++ show peer
104 _ <- hGet h (length (unpack hs))
105 putStrLn $ "<-- handshake from peer: " ++ show peer
108 instance Binary PeerMsg where
109 put msg = case msg of
110 KeepAliveMsg -> putWord32be 0
111 ChokeMsg -> do putWord32be 1
113 UnChokeMsg -> do putWord32be 1
115 InterestedMsg -> do putWord32be 1
117 NotInterestedMsg -> do putWord32be 1
119 HaveMsg i -> do putWord32be 5
121 putWord32be (fromIntegral i)
122 BitFieldMsg bf -> do putWord32be $ fromIntegral (1 + bfListLen)
124 mapM_ putWord8 bfList
125 where bfList = unpack bf
126 bfListLen = length bfList
127 RequestMsg i o l -> do putWord32be 13
129 putWord32be (fromIntegral i)
130 putWord32be (fromIntegral o)
131 putWord32be (fromIntegral l)
132 PieceMsg i o b -> do putWord32be $ fromIntegral (9 + blocklen)
134 putWord32be (fromIntegral i)
135 putWord32be (fromIntegral o)
136 mapM_ putWord8 blockList
137 where blockList = unpack b
138 blocklen = length blockList
139 CancelMsg i o l -> do putWord32be 13
141 putWord32be (fromIntegral i)
142 putWord32be (fromIntegral o)
143 putWord32be (fromIntegral l)
144 PortMsg p -> do putWord32be 3
146 putWord16be (fromIntegral p)
152 1 -> return UnChokeMsg
153 2 -> return InterestedMsg
154 3 -> return NotInterestedMsg
155 4 -> liftM (HaveMsg . fromIntegral) getWord32be
156 5 -> liftM (BitFieldMsg . pack) (replicateM (fromIntegral l - 1) getWord8)
157 6 -> liftA3 RequestMsg getInteger getInteger getInteger
158 where getInteger = fromIntegral <$> getWord32be
159 7 -> liftA3 PieceMsg getInteger getInteger (pack <$> replicateM (fromIntegral l - 9) getWord8)
160 where getInteger = fromIntegral <$> getWord32be
161 8 -> liftA3 CancelMsg getInteger getInteger getInteger
162 where getInteger = fromIntegral <$> getWord32be
163 9 -> liftM (PortMsg . fromIntegral) getWord16be
164 _ -> error ("unknown message ID: " ++ show msgid)
166 getMsg :: Handle -> IO PeerMsg
171 then return KeepAliveMsg
174 return $ decode $ fromStrict $ concat [lBS, msg]
176 sendMsg :: Handle -> PeerMsg -> IO ()
177 sendMsg h msg = hPut h bsMsg
178 where bsMsg = toStrict $ encode msg
180 bsToInt :: ByteString -> Int
181 bsToInt x = fromIntegral (runGet getWord32be (fromChunks (return x)))
183 bitfieldToList :: [Word8] -> [Integer]
184 bitfieldToList bs = go bs 0
187 let setBits = [pos*8 + toInteger i | i <- [0..8], testBit b i]
189 setBits ++ go bs' (pos + 1)
191 -- recvMsg :: Peer -> Handle -> Msg
192 msgLoop :: PeerState -> PieceMap -> IO ()
193 msgLoop pState pieceStatus | not (meInterested pState) && heChoking pState = do
194 -- if me NOT Interested and she is Choking, tell her that
196 let h = handle pState
197 sendMsg h InterestedMsg
198 putStrLn $ "--> InterestedMsg to peer: " ++ show (peer pState)
199 msgLoop (pState { meInterested = True }) pieceStatus
200 | meInterested pState && not (heChoking pState) =
201 -- if me Interested and she not Choking, send her a request
203 case pickPiece pieceStatus of
204 Nothing -> putStrLn "Nothing to download"
206 let pLen = len (pieceStatus ! workPiece)
207 putStrLn $ "piece length = " ++ show pLen
208 pBS <- downloadPiece (handle pState) workPiece pLen
209 if not $ verifyHash pBS (hash (pieceStatus ! workPiece))
211 putStrLn $ "Hash mismatch: " ++ show (hash (pieceStatus ! workPiece)) ++ " vs " ++ show (take 20 (SHA1.hash pBS))
213 let fileOffset = if workPiece == 0 then 0 else workPiece * len (pieceStatus ! (workPiece - 1))
214 putStrLn $ "Write into file at offset: " ++ show fileOffset
215 writeFileAtOffset "/tmp/download.file" fileOffset pBS
216 msgLoop pState (adjust (\pieceData -> pieceData { state = Have }) workPiece pieceStatus)
218 msg <- getMsg (handle pState)
219 putStrLn $ "<-- " ++ show msg ++ "from peer: " ++ show (peer pState)
222 sendMsg (handle pState) KeepAliveMsg
223 putStrLn $ "--> " ++ "KeepAliveMsg to peer: " ++ show (peer pState)
224 msgLoop pState pieceStatus
225 BitFieldMsg bss -> do
226 let pieceList = bitfieldToList (unpack bss)
227 pieceStatus' = updatePieceAvailability pieceStatus (peer pState) pieceList
228 putStrLn $ show (length pieceList) ++ " Pieces"
229 -- for each pieceIndex in pieceList, make an entry in the pieceStatus
230 -- map with pieceIndex as the key and modify the value to add the peer.
231 -- download each of the piece in order
232 msgLoop pState pieceStatus'
234 msgLoop (pState { heChoking = False }) pieceStatus
236 msgLoop pState pieceStatus
238 -- simple algorithm to pick piece.
239 -- pick the first piece from 0 that is not downloaded yet.
240 pickPiece :: PieceMap -> Maybe Integer
242 let pieceList = toList m
243 allPending = filter (\(_, v) -> state v == Pending) pieceList
249 updatePieceAvailability :: PieceMap -> Peer -> [Integer] -> PieceMap
250 updatePieceAvailability pieceStatus p pieceList =
251 mapWithKey (\k pd -> if k `elem` pieceList
252 then (pd { peers = p : peers pd })
255 handlePeerMsgs :: Peer -> Metainfo -> String -> IO ()
256 handlePeerMsgs p m peerId = do
258 doHandShake h p (infoHash m) peerId
259 let state = PeerState { handle = h
261 , heInterested = False
263 , meInterested = False
265 pieceHash = pieces (info m)
266 numPieces = (toInteger . (`quot` 20) . BC.length) pieceHash
267 pLen = pieceLength (info m)
268 fileLen = lengthInBytes (info m)
269 pieceStatus = mkPieceMap numPieces pieceHash (splitNum fileLen pLen)
270 createDummyFile "/tmp/download.file" (fromIntegral fileLen)
271 msgLoop state pieceStatus
273 downloadPiece :: Handle -> Integer -> Integer -> IO ByteString
274 downloadPiece h index pieceLength = do
275 let chunks = splitNum pieceLength 16384
276 liftM concat $ forM (zip [0..] chunks) (\(i, pLen) -> do
277 sendMsg h (RequestMsg index (i*pLen) pLen)
278 putStrLn $ "--> " ++ "RequestMsg for Piece "
279 ++ show index ++ ", part: " ++ show i ++ " of length: "
283 PieceMsg index begin block -> do
284 putStrLn $ " <-- PieceMsg for Piece: "
290 putStrLn "ignoring irrelevant msg"
293 verifyHash :: ByteString -> ByteString -> Bool
294 verifyHash bs pieceHash =
295 take 20 (SHA1.hash bs) == pieceHash