1 {-# LANGUAGE OverloadedStrings #-}
2 module FuncTorrent.Peer
11 import Prelude hiding (lookup, concat, replicate, splitAt, take, drop, filter)
13 import System.IO (Handle, BufferMode(..), hSetBuffering, hClose)
14 import Data.ByteString (ByteString, unpack, concat, hGet, hPut, take, drop, empty)
15 import qualified Data.ByteString.Char8 as BC (length)
16 import Network (connectTo, PortID(..))
17 import Control.Monad.State
19 import Data.Word (Word8)
20 import Data.Map (Map, fromList, toList, (!), mapWithKey, traverseWithKey, adjust, filter)
23 import FuncTorrent.Metainfo (Info(..), Metainfo(..))
24 import FuncTorrent.Utils (splitN, splitNum, writeFileAtOffset, readFileAtOffset, verifyHash)
25 import FuncTorrent.PeerMsgs (Peer(..), PeerMsg(..), sendMsg, getMsg, genHandshakeMsg)
27 data PState = PState { handle :: Handle
30 , meInterested :: Bool
32 , heInterested :: Bool}
34 data PieceDlState = Pending
39 -- todo - map with index to a new data structure (peers who have that piece and state)
40 data PieceData = PieceData { peers :: [Peer] -- ^ list of peers who have this piece
41 , dlstate :: PieceDlState -- ^ state of the piece from download perspective.
42 , hash :: ByteString -- ^ piece hash
43 , len :: Integer } -- ^ piece length
45 -- which piece is with which peers
46 type PieceMap = Map Integer PieceData
49 -- Make the initial Piece map, with the assumption that no peer has the
50 -- piece and that every piece is pending download.
51 initPieceMap :: ByteString -> Integer -> Integer -> PieceMap
52 initPieceMap pieceHash fileLen pieceLen = fromList kvs
54 numPieces = (toInteger . (`quot` 20) . BC.length) pieceHash
55 kvs = [(i, PieceData { peers = []
59 | (i, h, pLen) <- zip3 [0..numPieces] hashes pLengths]
60 hashes = splitN 20 pieceHash
61 pLengths = splitNum fileLen pieceLen
63 pieceMapFromFile :: FilePath -> PieceMap -> IO PieceMap
64 pieceMapFromFile filePath pieceMap =
65 traverseWithKey f pieceMap
68 let offset = if k == 0 then 0 else k * len (pieceMap ! (k - 1))
69 isHashValid <- flip verifyHash (hash v) <$> readFileAtOffset filePath offset (len v)
71 then return $ v { dlstate = Have }
74 havePiece :: PieceMap -> Integer -> Bool
76 dlstate (pm ! index) == Have
78 connectToPeer :: Peer -> IO Handle
79 connectToPeer (Peer _ ip port) = do
80 h <- connectTo ip (PortNumber (fromIntegral port))
81 hSetBuffering h LineBuffering
84 doHandshake :: Bool -> Handle -> Peer -> ByteString -> String -> IO ()
85 doHandshake True h peer infoHash peerid = do
86 let hs = genHandshakeMsg infoHash peerid
88 putStrLn $ "--> handhake to peer: " ++ show peer
89 _ <- hGet h (length (unpack hs))
90 putStrLn $ "<-- handshake from peer: " ++ show peer
92 doHandshake False h peer infoHash peerid = do
93 let hs = genHandshakeMsg infoHash peerid
94 putStrLn "waiting for a handshake"
95 hsMsg <- hGet h (length (unpack hs))
96 putStrLn $ "<-- handshake from peer: " ++ show peer
97 let rxInfoHash = take 20 $ drop 28 hsMsg
98 if rxInfoHash /= infoHash
100 putStrLn "infoHashes does not match"
105 putStrLn $ "--> handhake to peer: " ++ show peer
108 bitfieldToList :: [Word8] -> [Integer]
109 bitfieldToList bs = go bs 0
112 let setBits = [pos*8 + toInteger i | i <- [0..8], testBit b i]
114 setBits ++ go bs' (pos + 1)
116 -- helper functions to manipulate PeerState
117 toPeerState :: Handle
119 -> Bool -- ^ meChoking
120 -> Bool -- ^ meInterested
121 -> Bool -- ^ heChoking
122 -> Bool -- ^ heInterested
124 toPeerState h p meCh meIn heCh heIn =
128 , heInterested = heIn
130 , meInterested = meIn }
132 -- simple algorithm to pick piece.
133 -- pick the first piece from 0 that is not downloaded yet.
134 pickPiece :: PieceMap -> Maybe Integer
136 (fst `liftM`) . headMay . toList . filter (\v -> dlstate v == Pending)
138 bytesDownloaded :: PieceMap -> Integer
140 sum . map (len . snd) . toList . filter (\v -> dlstate v == Have)
142 updatePieceAvailability :: PieceMap -> Peer -> [Integer] -> PieceMap
143 updatePieceAvailability pieceStatus p pieceList =
144 mapWithKey (\k pd -> if k `elem` pieceList
145 then (pd { peers = p : peers pd })
148 handlePeerMsgs :: Peer -> String -> Metainfo -> PieceMap -> Bool -> IO ()
149 handlePeerMsgs p peerId m pieceMap isClient = do
151 doHandshake isClient h p (infoHash m) peerId
152 let pstate = toPeerState h p False False True True
153 filePath = name (info m)
154 _ <- runStateT (msgLoop pieceMap filePath) pstate
157 msgLoop :: PieceMap -> FilePath -> StateT PState IO ()
158 msgLoop pieceStatus file = do
162 PState { meInterested = False, heChoking = True } -> do
163 liftIO $ sendMsg h InterestedMsg
164 gets peer >>= (\p -> liftIO $ putStrLn $ "--> InterestedMsg to peer: " ++ show p)
165 modify (\st -> st { meInterested = True })
166 msgLoop pieceStatus file
167 PState { meInterested = True, heChoking = False } ->
168 case pickPiece pieceStatus of
169 Nothing -> liftIO $ putStrLn "Nothing to download"
171 let pLen = len (pieceStatus ! workPiece)
172 liftIO $ putStrLn $ "piece length = " ++ show pLen
173 pBS <- liftIO $ downloadPiece h workPiece pLen
174 if not $ verifyHash pBS (hash (pieceStatus ! workPiece))
176 liftIO $ putStrLn "Hash mismatch"
178 let fileOffset = if workPiece == 0 then 0 else workPiece * len (pieceStatus ! (workPiece - 1))
179 liftIO $ putStrLn $ "Write into file at offset: " ++ show fileOffset
180 liftIO $ writeFileAtOffset file fileOffset pBS
181 msgLoop (adjust (\pieceData -> pieceData { dlstate = Have }) workPiece pieceStatus) file
183 msg <- liftIO $ getMsg h
184 gets peer >>= (\p -> liftIO $ putStrLn $ "<-- " ++ show msg ++ "from peer: " ++ show p)
187 liftIO $ sendMsg h KeepAliveMsg
188 gets peer >>= (\p -> liftIO $ putStrLn $ "--> " ++ "KeepAliveMsg to peer: " ++ show p)
189 msgLoop pieceStatus file
190 BitFieldMsg bss -> do
192 let pieceList = bitfieldToList (unpack bss)
193 pieceStatus' = updatePieceAvailability pieceStatus p pieceList
194 liftIO $ putStrLn $ show (length pieceList) ++ " Pieces"
195 -- for each pieceIndex in pieceList, make an entry in the pieceStatus
196 -- map with pieceIndex as the key and modify the value to add the peer.
197 -- download each of the piece in order
198 msgLoop pieceStatus' file
200 modify (\st -> st {heChoking = False })
201 msgLoop pieceStatus file
203 modify (\st -> st {heChoking = True })
204 msgLoop pieceStatus file
206 modify (\st -> st {heInterested = True})
207 msgLoop pieceStatus file
208 NotInterestedMsg -> do
209 modify (\st -> st {heInterested = False})
210 msgLoop pieceStatus file
211 CancelMsg _ _ _ -> -- check if valid index, begin, length
212 msgLoop pieceStatus file
214 msgLoop pieceStatus file
215 -- handle RequestMsg, HaveMsg. No need to handle PieceMsg here.
219 downloadPiece :: Handle -> Integer -> Integer -> IO ByteString
220 downloadPiece h index pieceLength = do
221 let chunks = splitNum pieceLength 16384
222 concat `liftM` forM (zip [0..] chunks) (\(i, pLen) -> do
223 sendMsg h (RequestMsg index (i*pLen) pLen)
224 putStrLn $ "--> " ++ "RequestMsg for Piece "
225 ++ show index ++ ", part: " ++ show i ++ " of length: "
229 PieceMsg index begin block -> do
230 putStrLn $ " <-- PieceMsg for Piece: "
236 putStrLn "ignoring irrelevant msg"