1 {-# LANGUAGE OverloadedStrings #-}
2 module FuncTorrent.Peer
11 import Prelude hiding (lookup, concat, replicate, splitAt, take, filter)
13 import System.IO (Handle, BufferMode(..), hSetBuffering)
14 import Data.ByteString (ByteString, unpack, concat, hGet, hPut, take, 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 type PeerState = State PState
36 data PieceDlState = Pending
41 -- todo - map with index to a new data structure (peers who have that piece and state)
42 data PieceData = PieceData { peers :: [Peer] -- ^ list of peers who have this piece
43 , dlstate :: PieceDlState -- ^ state of the piece from download perspective.
44 , hash :: ByteString -- ^ piece hash
45 , len :: Integer } -- ^ piece length
47 -- which piece is with which peers
48 type PieceMap = Map Integer PieceData
51 -- Make the initial Piece map, with the assumption that no peer has the
52 -- piece and that every piece is pending download.
53 initPieceMap :: ByteString -> Integer -> Integer -> PieceMap
54 initPieceMap pieceHash fileLen pieceLen = fromList kvs
56 numPieces = (toInteger . (`quot` 20) . BC.length) pieceHash
57 kvs = [(i, PieceData { peers = []
61 | (i, h, pLen) <- zip3 [0..numPieces] hashes pLengths]
62 hashes = splitN 20 pieceHash
63 pLengths = (splitNum fileLen pieceLen)
65 pieceMapFromFile :: FilePath -> PieceMap -> IO PieceMap
66 pieceMapFromFile filePath pieceMap = do
67 traverseWithKey f pieceMap
70 let offset = if k == 0 then 0 else k * len (pieceMap ! (k - 1))
71 isHashValid <- (flip verifyHash) (hash v) <$> (readFileAtOffset filePath offset (len v))
73 then return $ v { dlstate = Have }
76 havePiece :: PieceMap -> Integer -> Bool
78 dlstate (pm ! index) == Have
80 connectToPeer :: Peer -> IO Handle
81 connectToPeer (Peer _ ip port) = do
82 h <- connectTo ip (PortNumber (fromIntegral port))
83 hSetBuffering h LineBuffering
86 doHandshake :: Handle -> Peer -> ByteString -> String -> IO ()
87 doHandshake h peer infoHash peerid = do
88 let hs = genHandshakeMsg infoHash peerid
90 putStrLn $ "--> handhake to peer: " ++ show peer
91 _ <- hGet h (length (unpack hs))
92 putStrLn $ "<-- handshake from peer: " ++ show peer
95 bitfieldToList :: [Word8] -> [Integer]
96 bitfieldToList bs = go bs 0
99 let setBits = [pos*8 + toInteger i | i <- [0..8], testBit b i]
101 setBits ++ go bs' (pos + 1)
103 -- helper functions to manipulate PeerState
104 toPeerState :: Handle
106 -> Bool -- ^ meChoking
107 -> Bool -- ^ meInterested
108 -> Bool -- ^ heChoking
109 -> Bool -- ^ heInterested
111 toPeerState h p meCh meIn heCh heIn =
115 , heInterested = heIn
117 , meInterested = meIn }
119 -- simple algorithm to pick piece.
120 -- pick the first piece from 0 that is not downloaded yet.
121 pickPiece :: PieceMap -> Maybe Integer
123 (fst `liftM`) . headMay . toList . filter (\v -> dlstate v == Pending)
125 bytesDownloaded :: PieceMap -> Integer
127 sum . map (len . snd) . toList . filter (\v -> dlstate v == Have)
129 updatePieceAvailability :: PieceMap -> Peer -> [Integer] -> PieceMap
130 updatePieceAvailability pieceStatus p pieceList =
131 mapWithKey (\k pd -> if k `elem` pieceList
132 then (pd { peers = p : peers pd })
135 handlePeerMsgs :: Peer -> String -> Metainfo -> PieceMap -> IO ()
136 handlePeerMsgs p peerId m pieceMap = do
138 doHandshake h p (infoHash m) peerId
139 let pstate = toPeerState h p False False True True
140 filePath = name (info m)
141 _ <- runStateT (msgLoop pieceMap filePath) pstate
144 msgLoop :: PieceMap -> FilePath -> StateT PState IO ()
145 msgLoop pieceStatus file = do
149 PState { meInterested = False, heChoking = True } -> do
150 liftIO $ sendMsg h InterestedMsg
151 gets peer >>= (\p -> liftIO $ putStrLn $ "--> InterestedMsg to peer: " ++ show p)
152 modify (\st -> st { meInterested = True })
153 msgLoop pieceStatus file
154 PState { meInterested = True, heChoking = False } ->
155 case pickPiece pieceStatus of
156 Nothing -> liftIO $ putStrLn "Nothing to download"
158 let pLen = len (pieceStatus ! workPiece)
159 liftIO $ putStrLn $ "piece length = " ++ show pLen
160 pBS <- liftIO $ downloadPiece h workPiece pLen
161 if not $ verifyHash pBS (hash (pieceStatus ! workPiece))
163 liftIO $ putStrLn $ "Hash mismatch"
165 let fileOffset = if workPiece == 0 then 0 else workPiece * len (pieceStatus ! (workPiece - 1))
166 liftIO $ putStrLn $ "Write into file at offset: " ++ show fileOffset
167 liftIO $ writeFileAtOffset file fileOffset pBS
168 msgLoop (adjust (\pieceData -> pieceData { dlstate = Have }) workPiece pieceStatus) file
170 msg <- liftIO $ getMsg h
171 gets peer >>= (\p -> liftIO $ putStrLn $ "<-- " ++ show msg ++ "from peer: " ++ show p)
174 liftIO $ sendMsg h KeepAliveMsg
175 gets peer >>= (\p -> liftIO $ putStrLn $ "--> " ++ "KeepAliveMsg to peer: " ++ show p)
176 msgLoop pieceStatus file
177 BitFieldMsg bss -> do
179 let pieceList = bitfieldToList (unpack bss)
180 pieceStatus' = updatePieceAvailability pieceStatus p pieceList
181 liftIO $ putStrLn $ show (length pieceList) ++ " Pieces"
182 -- for each pieceIndex in pieceList, make an entry in the pieceStatus
183 -- map with pieceIndex as the key and modify the value to add the peer.
184 -- download each of the piece in order
185 msgLoop pieceStatus' file
187 modify (\st -> st {heChoking = False })
188 msgLoop pieceStatus file
191 downloadPiece :: Handle -> Integer -> Integer -> IO ByteString
192 downloadPiece h index pieceLength = do
193 let chunks = splitNum pieceLength 16384
194 concat `liftM` forM (zip [0..] chunks) (\(i, pLen) -> do
195 sendMsg h (RequestMsg index (i*pLen) pLen)
196 putStrLn $ "--> " ++ "RequestMsg for Piece "
197 ++ show index ++ ", part: " ++ show i ++ " of length: "
201 PieceMsg index begin block -> do
202 putStrLn $ " <-- PieceMsg for Piece: "
208 putStrLn "ignoring irrelevant msg"