1 {-# LANGUAGE OverloadedStrings #-}
2 module FuncTorrent.Peer
7 import Prelude hiding (lookup, concat, replicate, splitAt, writeFile, take)
9 import System.IO (Handle, BufferMode(..), IOMode(..), SeekMode(..), withFile, hSeek, hSetBuffering)
10 import Data.ByteString (ByteString, pack, unpack, concat, hGet, hPut, singleton, writeFile, 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)
31 -- PeerState is a misnomer
32 data PeerState = PeerState { handle :: Handle
35 , meInterested :: Bool
37 , heInterested :: Bool}
39 data PieceDlState = Pending
44 -- todo - map with index to a new data structure (peers who have that piece amd state)
45 data PieceData = PieceData { peers :: [Peer] -- ^ list of peers who have this piece
46 , state :: PieceDlState -- ^ state of the piece from download perspective.
47 , hash :: ByteString -- ^ piece hash
48 , len :: Integer } -- ^ piece length
50 -- which piece is with which peers
51 type PieceMap = Map Integer PieceData
53 -- | Peer is a PeerID, IP address, port tuple
54 data Peer = Peer ID IP Port
57 data PeerMsg = KeepAliveMsg
63 | BitFieldMsg ByteString
64 | RequestMsg Integer Integer Integer
65 | PieceMsg Integer Integer ByteString
66 | CancelMsg Integer Integer Integer
70 -- Make the initial Piece map, with the assumption that no peer has the
71 -- piece and that every piece is pending download.
72 mkPieceMap :: Integer -> ByteString -> [Integer] -> PieceMap
73 mkPieceMap numPieces pieceHash pLengths = fromList kvs
74 where kvs = [(i, PieceData { peers = []
78 | (i, h, pLen) <- zip3 [0..numPieces] hashes pLengths]
79 hashes = splitN 20 pieceHash
81 havePiece :: PieceMap -> Integer -> Bool
83 state (pm ! index) == Have
85 genHandShakeMsg :: ByteString -> String -> ByteString
86 genHandShakeMsg infoHash peer_id = concat [pstrlen, pstr, reserved, infoHash, peerID]
87 where pstrlen = singleton 19
88 pstr = BC.pack "BitTorrent protocol"
89 reserved = BC.replicate 8 '\0'
90 peerID = BC.pack peer_id
92 handShake :: Peer -> ByteString -> String -> IO Handle
93 handShake peer@(Peer _ ip port) infoHash peerid = do
94 let hs = genHandShakeMsg infoHash peerid
95 h <- connectTo ip (PortNumber (fromIntegral port))
96 hSetBuffering h LineBuffering
98 putStrLn $ "--> handhake to peer: " ++ show peer
99 _ <- hGet h (length (unpack hs))
100 putStrLn $ "<-- handshake from peer: " ++ show peer
103 instance Binary PeerMsg where
104 put msg = case msg of
105 KeepAliveMsg -> putWord32be 0
106 ChokeMsg -> do putWord32be 1
108 UnChokeMsg -> do putWord32be 1
110 InterestedMsg -> do putWord32be 1
112 NotInterestedMsg -> do putWord32be 1
114 HaveMsg i -> do putWord32be 5
116 putWord32be (fromIntegral i)
117 BitFieldMsg bf -> do putWord32be $ fromIntegral (1 + bfListLen)
119 mapM_ putWord8 bfList
120 where bfList = unpack bf
121 bfListLen = length bfList
122 RequestMsg i o l -> do putWord32be 13
124 putWord32be (fromIntegral i)
125 putWord32be (fromIntegral o)
126 putWord32be (fromIntegral l)
127 PieceMsg i o b -> do putWord32be $ fromIntegral (9 + blocklen)
129 putWord32be (fromIntegral i)
130 putWord32be (fromIntegral o)
131 mapM_ putWord8 blockList
132 where blockList = unpack b
133 blocklen = length blockList
134 CancelMsg i o l -> do putWord32be 13
136 putWord32be (fromIntegral i)
137 putWord32be (fromIntegral o)
138 putWord32be (fromIntegral l)
139 PortMsg p -> do putWord32be 3
141 putWord16be (fromIntegral p)
147 1 -> return UnChokeMsg
148 2 -> return InterestedMsg
149 3 -> return NotInterestedMsg
150 4 -> liftM (HaveMsg . fromIntegral) getWord32be
151 5 -> liftM (BitFieldMsg . pack) (replicateM (fromIntegral l - 1) getWord8)
152 6 -> liftA3 RequestMsg getInteger getInteger getInteger
153 where getInteger = fromIntegral <$> getWord32be
154 7 -> liftA3 PieceMsg getInteger getInteger (pack <$> replicateM (fromIntegral l - 9) getWord8)
155 where getInteger = fromIntegral <$> getWord32be
156 8 -> liftA3 CancelMsg getInteger getInteger getInteger
157 where getInteger = fromIntegral <$> getWord32be
158 9 -> liftM (PortMsg . fromIntegral) getWord16be
159 _ -> error ("unknown message ID: " ++ show msgid)
161 getMsg :: Handle -> IO PeerMsg
166 then return KeepAliveMsg
169 return $ decode $ fromStrict $ concat [lBS, msg]
171 sendMsg :: Handle -> PeerMsg -> IO ()
173 let bsMsg = toStrict $ encode msg
177 bsToInt :: ByteString -> Int
178 bsToInt x = fromIntegral (runGet getWord32be (fromChunks (return x)))
180 bitfieldToList :: [Word8] -> [Integer]
181 bitfieldToList bs = go bs 0
184 let setBits = [pos*8 + toInteger i | i <- [0..8], testBit b i]
186 setBits ++ go bs' (pos + 1)
188 createDummyFile :: FilePath -> Int -> IO ()
189 createDummyFile path size =
190 writeFile path (BC.replicate size '\0')
192 -- write into a file at a specific offet
193 writeFileAtOffset :: FilePath -> Integer -> ByteString -> IO ()
194 writeFileAtOffset path offset block =
195 withFile path ReadWriteMode $ (\h -> do
196 _ <- hSeek h AbsoluteSeek offset
199 -- recvMsg :: Peer -> Handle -> Msg
200 msgLoop :: PeerState -> PieceMap -> IO ()
201 msgLoop pState pieceStatus | meInterested pState == False &&
202 heChoking pState == True = do
203 -- if me NOT Interested and she is Choking, tell her that
205 let h = handle pState
206 sendMsg h InterestedMsg
207 putStrLn $ "--> InterestedMsg to peer: " ++ show (peer pState)
208 msgLoop (pState { meInterested = True }) pieceStatus
209 | meInterested pState == True &&
210 heChoking pState == False =
211 -- if me Interested and she not Choking, send her a request
213 case pickPiece pieceStatus of
214 Nothing -> putStrLn "Nothing to download"
216 let pLen = len (pieceStatus ! workPiece)
217 putStrLn $ "piece length = " ++ show pLen
218 pBS <- downloadPiece (handle pState) workPiece pLen
219 if not $ verifyHash pBS (hash (pieceStatus ! workPiece))
221 putStrLn $ "Hash mismatch: " ++ show (hash (pieceStatus ! workPiece)) ++ " vs " ++ show (take 20 (SHA1.hash pBS))
223 let fileOffset = if workPiece == 0 then 0 else (len (pieceStatus ! (workPiece - 1)))
224 writeFileAtOffset "/tmp/download.file" fileOffset pBS
225 msgLoop pState (adjust (\pieceData -> pieceData { state = Have }) workPiece pieceStatus)
227 msg <- getMsg (handle pState)
228 putStrLn $ "<-- " ++ show msg ++ "from peer: " ++ show (peer pState)
231 sendMsg (handle pState) KeepAliveMsg
232 putStrLn $ "--> " ++ "KeepAliveMsg to peer: " ++ show (peer pState)
233 msgLoop pState pieceStatus
234 BitFieldMsg bss -> do
235 let pieceList = bitfieldToList (unpack bss)
236 pieceStatus' = updatePieceAvailability pieceStatus (peer pState) pieceList
237 putStrLn $ show (length pieceList) ++ " Pieces"
238 -- for each pieceIndex in pieceList, make an entry in the pieceStatus
239 -- map with pieceIndex as the key and modify the value to add the peer.
240 -- download each of the piece in order
241 msgLoop pState pieceStatus'
243 msgLoop (pState { heChoking = False }) pieceStatus
245 msgLoop pState pieceStatus
247 -- simple algorithm to pick piece.
248 -- pick the first piece from 0 that is not downloaded yet.
249 pickPiece :: PieceMap -> Maybe Integer
251 let pieceList = toList m
252 allPending = filter (\(_, v) -> state v == Pending) pieceList
258 updatePieceAvailability :: PieceMap -> Peer -> [Integer] -> PieceMap
259 updatePieceAvailability pieceStatus p pieceList =
260 mapWithKey (\k pd -> if k `elem` pieceList
261 then (pd { peers = p : (peers pd) })
264 handlePeerMsgs :: Peer -> Metainfo -> String -> IO ()
265 handlePeerMsgs p m peerId = do
266 h <- handShake p (infoHash m) peerId
267 let state = PeerState { handle = h
269 , heInterested = False
271 , meInterested = False
273 pieceHash = pieces (info m)
274 numPieces = (toInteger . (`quot` 20) . BC.length) pieceHash
275 pLen = pieceLength (info m)
276 fileLen = lengthInBytes (info m)
277 pieceStatus = mkPieceMap numPieces pieceHash (splitNum fileLen pLen)
278 createDummyFile "/tmp/download.file" (fromIntegral fileLen)
279 msgLoop state pieceStatus
281 downloadPiece :: Handle -> Integer -> Integer -> IO ByteString
282 downloadPiece h index pieceLength = do
283 let chunks = splitNum pieceLength 16384
284 liftM concat $ forM (zip [0..] chunks) (\(i, pLen) -> do
285 sendMsg h (RequestMsg index (i*pLen) pLen)
286 putStrLn $ "--> " ++ "RequestMsg for Piece "
287 ++ (show index) ++ ", part: " ++ show i ++ " of length: "
291 PieceMsg index begin block -> do
292 putStrLn $ " <-- PieceMsg for Piece: "
298 putStrLn "ignoring irrelevant msg"
301 verifyHash :: ByteString -> ByteString -> Bool
302 verifyHash bs pieceHash =
303 take 20 (SHA1.hash bs) == pieceHash