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 | not (meInterested pState) && heChoking pState = do
202 -- if me NOT Interested and she is Choking, tell her that
204 let h = handle pState
205 sendMsg h InterestedMsg
206 putStrLn $ "--> InterestedMsg to peer: " ++ show (peer pState)
207 msgLoop (pState { meInterested = True }) pieceStatus
208 | meInterested pState && not (heChoking pState) =
209 -- if me Interested and she not Choking, send her a request
211 case pickPiece pieceStatus of
212 Nothing -> putStrLn "Nothing to download"
214 let pLen = len (pieceStatus ! workPiece)
215 putStrLn $ "piece length = " ++ show pLen
216 pBS <- downloadPiece (handle pState) workPiece pLen
217 if not $ verifyHash pBS (hash (pieceStatus ! workPiece))
219 putStrLn $ "Hash mismatch: " ++ show (hash (pieceStatus ! workPiece)) ++ " vs " ++ show (take 20 (SHA1.hash pBS))
221 let fileOffset = if workPiece == 0 then 0 else workPiece * len (pieceStatus ! (workPiece - 1))
222 putStrLn $ "Write into file at offset: " ++ show fileOffset
223 writeFileAtOffset "/tmp/download.file" fileOffset pBS
224 msgLoop pState (adjust (\pieceData -> pieceData { state = Have }) workPiece pieceStatus)
226 msg <- getMsg (handle pState)
227 putStrLn $ "<-- " ++ show msg ++ "from peer: " ++ show (peer pState)
230 sendMsg (handle pState) KeepAliveMsg
231 putStrLn $ "--> " ++ "KeepAliveMsg to peer: " ++ show (peer pState)
232 msgLoop pState pieceStatus
233 BitFieldMsg bss -> do
234 let pieceList = bitfieldToList (unpack bss)
235 pieceStatus' = updatePieceAvailability pieceStatus (peer pState) pieceList
236 putStrLn $ show (length pieceList) ++ " Pieces"
237 -- for each pieceIndex in pieceList, make an entry in the pieceStatus
238 -- map with pieceIndex as the key and modify the value to add the peer.
239 -- download each of the piece in order
240 msgLoop pState pieceStatus'
242 msgLoop (pState { heChoking = False }) pieceStatus
244 msgLoop pState pieceStatus
246 -- simple algorithm to pick piece.
247 -- pick the first piece from 0 that is not downloaded yet.
248 pickPiece :: PieceMap -> Maybe Integer
250 let pieceList = toList m
251 allPending = filter (\(_, v) -> state v == Pending) pieceList
257 updatePieceAvailability :: PieceMap -> Peer -> [Integer] -> PieceMap
258 updatePieceAvailability pieceStatus p pieceList =
259 mapWithKey (\k pd -> if k `elem` pieceList
260 then (pd { peers = p : peers pd })
263 handlePeerMsgs :: Peer -> Metainfo -> String -> IO ()
264 handlePeerMsgs p m peerId = do
265 h <- handShake p (infoHash m) peerId
266 let state = PeerState { handle = h
268 , heInterested = False
270 , meInterested = False
272 pieceHash = pieces (info m)
273 numPieces = (toInteger . (`quot` 20) . BC.length) pieceHash
274 pLen = pieceLength (info m)
275 fileLen = lengthInBytes (info m)
276 pieceStatus = mkPieceMap numPieces pieceHash (splitNum fileLen pLen)
277 createDummyFile "/tmp/download.file" (fromIntegral fileLen)
278 msgLoop state pieceStatus
280 downloadPiece :: Handle -> Integer -> Integer -> IO ByteString
281 downloadPiece h index pieceLength = do
282 let chunks = splitNum pieceLength 16384
283 liftM concat $ forM (zip [0..] chunks) (\(i, pLen) -> do
284 sendMsg h (RequestMsg index (i*pLen) pLen)
285 putStrLn $ "--> " ++ "RequestMsg for Piece "
286 ++ show index ++ ", part: " ++ show i ++ " of length: "
290 PieceMsg index begin block -> do
291 putStrLn $ " <-- PieceMsg for Piece: "
297 putStrLn "ignoring irrelevant msg"
300 verifyHash :: ByteString -> ByteString -> Bool
301 verifyHash bs pieceHash =
302 take 20 (SHA1.hash bs) == pieceHash