1 {-# LANGUAGE OverloadedStrings #-}
2 module FuncTorrent.Peer
7 import Prelude hiding (lookup, concat, replicate, splitAt, writeFile)
9 import System.IO (Handle, BufferMode(..), hSetBuffering)
10 import Data.ByteString (ByteString, pack, unpack, concat, hGet, hPut, singleton, writeFile)
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)
23 import FuncTorrent.Metainfo (Info(..), Metainfo(..))
24 import FuncTorrent.Utils (splitN, splitNum)
30 -- PeerState is a misnomer
31 data PeerState = PeerState { handle :: Handle
34 , meInterested :: Bool
36 , heInterested :: Bool}
38 data PieceDlState = Pending
43 -- todo - map with index to a new data structure (peers who have that piece amd state)
44 data PieceData = PieceData { peers :: [Peer] -- ^ list of peers who have this piece
45 , state :: PieceDlState -- ^ state of the piece from download perspective.
46 , hash :: ByteString -- ^ piece hash
47 , len :: Integer } -- ^ piece length
49 -- which piece is with which peers
50 type PieceMap = Map Integer PieceData
52 -- | Peer is a PeerID, IP address, port tuple
53 data Peer = Peer ID IP Port
56 data PeerMsg = KeepAliveMsg
62 | BitFieldMsg ByteString
63 | RequestMsg Integer Integer Integer
64 | PieceMsg Integer Integer ByteString
65 | CancelMsg Integer Integer Integer
69 -- Make the initial Piece map, with the assumption that no peer has the
70 -- piece and that every piece is pending download.
71 mkPieceMap :: Integer -> ByteString -> [Integer] -> PieceMap
72 mkPieceMap numPieces pieceHash pLengths = fromList kvs
73 where kvs = [(i, PieceData { peers = []
77 | (i, h, pLen) <- zip3 [0..numPieces] hashes pLengths]
78 hashes = splitN (fromIntegral numPieces) pieceHash
80 havePiece :: PieceMap -> Integer -> Bool
82 state (pm ! index) == Have
84 genHandShakeMsg :: ByteString -> String -> ByteString
85 genHandShakeMsg infoHash peer_id = concat [pstrlen, pstr, reserved, infoHash, peerID]
86 where pstrlen = singleton 19
87 pstr = BC.pack "BitTorrent protocol"
88 reserved = BC.replicate 8 '\0'
89 peerID = BC.pack peer_id
91 handShake :: Peer -> ByteString -> String -> IO Handle
92 handShake peer@(Peer _ ip port) infoHash peerid = do
93 let hs = genHandShakeMsg infoHash peerid
94 h <- connectTo ip (PortNumber (fromIntegral port))
95 hSetBuffering h LineBuffering
97 putStrLn $ "--> handhake to peer: " ++ show peer
98 _ <- hGet h (length (unpack hs))
99 putStrLn $ "<-- handshake from peer: " ++ show peer
102 instance Binary PeerMsg where
103 put msg = case msg of
104 KeepAliveMsg -> putWord32be 0
105 ChokeMsg -> do putWord32be 1
107 UnChokeMsg -> do putWord32be 1
109 InterestedMsg -> do putWord32be 1
111 NotInterestedMsg -> do putWord32be 1
113 HaveMsg i -> do putWord32be 5
115 putWord32be (fromIntegral i)
116 BitFieldMsg bf -> do putWord32be $ fromIntegral (1 + bfListLen)
118 mapM_ putWord8 bfList
119 where bfList = unpack bf
120 bfListLen = length bfList
121 RequestMsg i o l -> do putWord32be 13
123 putWord32be (fromIntegral i)
124 putWord32be (fromIntegral o)
125 putWord32be (fromIntegral l)
126 PieceMsg i o b -> do putWord32be $ fromIntegral (9 + blocklen)
128 putWord32be (fromIntegral i)
129 putWord32be (fromIntegral o)
130 mapM_ putWord8 blockList
131 where blockList = unpack b
132 blocklen = length blockList
133 CancelMsg i o l -> do putWord32be 13
135 putWord32be (fromIntegral i)
136 putWord32be (fromIntegral o)
137 putWord32be (fromIntegral l)
138 PortMsg p -> do putWord32be 3
140 putWord16be (fromIntegral p)
146 1 -> return UnChokeMsg
147 2 -> return InterestedMsg
148 3 -> return NotInterestedMsg
149 4 -> liftM (HaveMsg . fromIntegral) getWord32be
150 5 -> liftM (BitFieldMsg . pack) (replicateM (fromIntegral l - 1) getWord8)
151 6 -> liftA3 RequestMsg getInteger getInteger getInteger
152 where getInteger = fromIntegral <$> getWord32be
153 7 -> liftA3 PieceMsg getInteger getInteger (pack <$> replicateM (fromIntegral l - 9) getWord8)
154 where getInteger = fromIntegral <$> getWord32be
155 8 -> liftA3 CancelMsg getInteger getInteger getInteger
156 where getInteger = fromIntegral <$> getWord32be
157 9 -> liftM (PortMsg . fromIntegral) getWord16be
158 _ -> error ("unknown message ID: " ++ show msgid)
160 getMsg :: Handle -> IO PeerMsg
165 then return KeepAliveMsg
168 return $ decode $ fromStrict $ concat [lBS, msg]
170 sendMsg :: Handle -> PeerMsg -> IO ()
172 let bsMsg = toStrict $ encode msg
176 bsToInt :: ByteString -> Int
177 bsToInt x = fromIntegral (runGet getWord32be (fromChunks (return x)))
179 bitfieldToList :: [Word8] -> [Integer]
180 bitfieldToList bs = go bs 0
183 let setBits = [pos*8 + toInteger i | i <- [0..8], testBit b i]
185 setBits ++ go bs' (pos + 1)
187 createDummyFile :: FilePath -> Int -> IO ()
188 createDummyFile path size =
189 writeFile path (BC.replicate size '\0')
191 -- loop1 :: shake hands with all peers, find out the pieces they have, form PieceData.
192 -- recvMsg :: Peer -> Handle -> Msg
193 msgLoop :: PeerState -> PieceMap -> IO ()
194 msgLoop pState pieceStatus | meInterested pState == False &&
195 heChoking pState == True = do
196 -- if me NOT Interested and she is Choking, tell her that
198 let h = handle pState
199 sendMsg h InterestedMsg
200 putStrLn $ "--> InterestedMsg to peer: " ++ show (peer pState)
201 msgLoop (pState { meInterested = True }) pieceStatus
202 | meInterested pState == True &&
203 heChoking pState == False =
204 -- if me Interested and she not Choking, send her a request
206 case pickPiece pieceStatus of
207 Nothing -> putStrLn "Nothing to download"
209 let pLen = len (pieceStatus ! workPiece)
210 _ <- downloadPiece (handle pState) workPiece pLen
211 msgLoop pState (adjust (\pieceData -> pieceData { state = Have }) workPiece pieceStatus)
213 msg <- getMsg (handle pState)
214 putStrLn $ "<-- " ++ show msg ++ "from peer: " ++ show (peer pState)
217 sendMsg (handle pState) KeepAliveMsg
218 putStrLn $ "--> " ++ "KeepAliveMsg to peer: " ++ show (peer pState)
219 msgLoop pState pieceStatus
220 BitFieldMsg bss -> do
221 let pieceList = bitfieldToList (unpack bss)
222 pieceStatus' = updatePieceAvailability pieceStatus (peer pState) pieceList
224 -- for each pieceIndex in pieceList, make an entry in the pieceStatus
225 -- map with pieceIndex as the key and modify the value to add the peer.
226 -- download each of the piece in order
227 msgLoop pState pieceStatus'
229 msgLoop (pState { heChoking = False }) pieceStatus
231 msgLoop pState pieceStatus
233 -- simple algorithm to pick piece.
234 -- pick the first piece from 0 that is not downloaded yet.
235 pickPiece :: PieceMap -> Maybe Integer
237 let pieceList = toList m
238 allPending = filter (\(_, v) -> state v == Pending) pieceList
244 updatePieceAvailability :: PieceMap -> Peer -> [Integer] -> PieceMap
245 updatePieceAvailability pieceStatus p pieceList =
246 mapWithKey (\k pd -> if k `elem` pieceList
247 then (pd { peers = p : (peers pd) })
250 handlePeerMsgs :: Peer -> Metainfo -> String -> IO ()
251 handlePeerMsgs p m peerId = do
252 h <- handShake p (infoHash m) peerId
253 let state = PeerState { handle = h
255 , heInterested = False
257 , meInterested = False
259 pieceHash = pieces (info m)
260 numPieces = (toInteger . (`quot` 20) . BC.length) pieceHash
261 pLen = pieceLength (info m)
262 fileLen = lengthInBytes (info m)
263 pieceStatus = mkPieceMap numPieces pieceHash (splitNum fileLen pLen)
264 msgLoop state pieceStatus
266 downloadPiece :: Handle -> Integer -> Integer -> IO [ByteString]
267 downloadPiece h index pieceLength = do
268 let chunks = splitNum pieceLength 16384
269 forM (zip [0..] chunks) (\(i, pLen) -> do
270 sendMsg h (RequestMsg index (i*pLen) pLen)
271 putStrLn $ "--> " ++ "RequestMsg for Piece "
272 ++ (show index) ++ ", part: " ++ show i ++ " of length: "
274 PieceMsg index begin block <- getMsg h
275 putStrLn $ " <-- PieceMsg for Piece: "