1 {-# LANGUAGE OverloadedStrings #-}
2 module FuncTorrent.Peer
8 import Prelude hiding (lookup, concat, replicate, splitAt, take, filter)
10 import System.IO (Handle, BufferMode(..), hSetBuffering)
11 import Data.ByteString (ByteString, unpack, concat, hGet, hPut, take, empty)
12 import qualified Data.ByteString.Char8 as BC (length)
13 import Network (connectTo, PortID(..))
14 import Control.Monad.State
16 import Data.Word (Word8)
17 import Data.Map (Map, fromList, toList, (!), mapWithKey, adjust, filter)
18 import qualified Crypto.Hash.SHA1 as SHA1 (hash)
21 import FuncTorrent.Metainfo (Info(..), Metainfo(..))
22 import FuncTorrent.Utils (splitN, splitNum)
23 import FuncTorrent.Fileops (createDummyFile, writeFileAtOffset)
24 import FuncTorrent.PeerMsgs (Peer(..), PeerMsg(..), sendMsg, getMsg, genHandshakeMsg)
26 data PState = PState { handle :: Handle
29 , meInterested :: Bool
31 , heInterested :: Bool}
33 type PeerState = State PState
35 data PieceDlState = Pending
40 -- todo - map with index to a new data structure (peers who have that piece and state)
41 data PieceData = PieceData { peers :: [Peer] -- ^ list of peers who have this piece
42 , dlstate :: PieceDlState -- ^ state of the piece from download perspective.
43 , hash :: ByteString -- ^ piece hash
44 , len :: Integer } -- ^ piece length
46 -- which piece is with which peers
47 type PieceMap = Map Integer PieceData
50 -- Make the initial Piece map, with the assumption that no peer has the
51 -- piece and that every piece is pending download.
52 mkPieceMap :: Integer -> ByteString -> [Integer] -> PieceMap
53 mkPieceMap numPieces pieceHash pLengths = fromList kvs
54 where kvs = [(i, PieceData { peers = []
58 | (i, h, pLen) <- zip3 [0..numPieces] hashes pLengths]
59 hashes = splitN 20 pieceHash
61 havePiece :: PieceMap -> Integer -> Bool
63 dlstate (pm ! index) == Have
65 connectToPeer :: Peer -> IO Handle
66 connectToPeer (Peer _ ip port) = do
67 h <- connectTo ip (PortNumber (fromIntegral port))
68 hSetBuffering h LineBuffering
71 doHandshake :: Handle -> Peer -> ByteString -> String -> IO ()
72 doHandshake h peer infoHash peerid = do
73 let hs = genHandshakeMsg infoHash peerid
75 putStrLn $ "--> handhake to peer: " ++ show peer
76 _ <- hGet h (length (unpack hs))
77 putStrLn $ "<-- handshake from peer: " ++ show peer
80 bitfieldToList :: [Word8] -> [Integer]
81 bitfieldToList bs = go bs 0
84 let setBits = [pos*8 + toInteger i | i <- [0..8], testBit b i]
86 setBits ++ go bs' (pos + 1)
88 -- helper functions to manipulate PeerState
91 -> Bool -- ^ meChoking
92 -> Bool -- ^ meInterested
93 -> Bool -- ^ heChoking
94 -> Bool -- ^ heInterested
96 toPeerState h p meCh meIn heCh heIn =
100 , heInterested = heIn
102 , meInterested = meIn }
104 -- simple algorithm to pick piece.
105 -- pick the first piece from 0 that is not downloaded yet.
106 pickPiece :: PieceMap -> Maybe Integer
108 (fst `liftM`) . headMay . toList . filter (\v -> dlstate v == Pending)
110 bytesDownloaded :: PieceMap -> Integer
112 sum . map (len . snd) . toList . filter (\v -> dlstate v == Have)
114 updatePieceAvailability :: PieceMap -> Peer -> [Integer] -> PieceMap
115 updatePieceAvailability pieceStatus p pieceList =
116 mapWithKey (\k pd -> if k `elem` pieceList
117 then (pd { peers = p : peers pd })
120 handlePeerMsgs :: Peer -> String -> Metainfo -> IO ()
121 handlePeerMsgs p peerId m = do
123 doHandshake h p (infoHash m) peerId
124 let pstate = toPeerState h p False False True True
125 pieceHash = pieces (info m)
126 numPieces = (toInteger . (`quot` 20) . BC.length) pieceHash
127 pLen = pieceLength (info m)
128 fileLen = lengthInBytes (info m)
129 fileName = name (info m)
130 pieceStatus = mkPieceMap numPieces pieceHash (splitNum fileLen pLen)
131 createDummyFile fileName (fromIntegral fileLen)
132 _ <- runStateT (msgLoop pieceStatus fileName) pstate
135 msgLoop :: PieceMap -> FilePath -> StateT PState IO ()
136 msgLoop pieceStatus file = do
140 PState { meInterested = False, heChoking = True } -> do
141 liftIO $ sendMsg h InterestedMsg
142 gets peer >>= (\p -> liftIO $ putStrLn $ "--> InterestedMsg to peer: " ++ show p)
143 modify (\st -> st { meInterested = True })
144 msgLoop pieceStatus file
145 PState { meInterested = True, heChoking = False } ->
146 case pickPiece pieceStatus of
147 Nothing -> liftIO $ putStrLn "Nothing to download"
149 let pLen = len (pieceStatus ! workPiece)
150 liftIO $ putStrLn $ "piece length = " ++ show pLen
151 pBS <- liftIO $ downloadPiece h workPiece pLen
152 if not $ verifyHash pBS (hash (pieceStatus ! workPiece))
154 liftIO $ putStrLn $ "Hash mismatch: " ++ show (hash (pieceStatus ! workPiece)) ++ " vs " ++ show (take 20 (SHA1.hash pBS))
156 let fileOffset = if workPiece == 0 then 0 else workPiece * len (pieceStatus ! (workPiece - 1))
157 liftIO $ putStrLn $ "Write into file at offset: " ++ show fileOffset
158 liftIO $ writeFileAtOffset file fileOffset pBS
159 msgLoop (adjust (\pieceData -> pieceData { dlstate = Have }) workPiece pieceStatus) file
161 msg <- liftIO $ getMsg h
162 gets peer >>= (\p -> liftIO $ putStrLn $ "<-- " ++ show msg ++ "from peer: " ++ show p)
165 liftIO $ sendMsg h KeepAliveMsg
166 gets peer >>= (\p -> liftIO $ putStrLn $ "--> " ++ "KeepAliveMsg to peer: " ++ show p)
167 msgLoop pieceStatus file
168 BitFieldMsg bss -> do
170 let pieceList = bitfieldToList (unpack bss)
171 pieceStatus' = updatePieceAvailability pieceStatus p pieceList
172 liftIO $ putStrLn $ show (length pieceList) ++ " Pieces"
173 -- for each pieceIndex in pieceList, make an entry in the pieceStatus
174 -- map with pieceIndex as the key and modify the value to add the peer.
175 -- download each of the piece in order
176 msgLoop pieceStatus' file
178 modify (\st -> st {heChoking = False })
179 msgLoop pieceStatus file
182 downloadPiece :: Handle -> Integer -> Integer -> IO ByteString
183 downloadPiece h index pieceLength = do
184 let chunks = splitNum pieceLength 16384
185 concat `liftM` forM (zip [0..] chunks) (\(i, pLen) -> do
186 sendMsg h (RequestMsg index (i*pLen) pLen)
187 putStrLn $ "--> " ++ "RequestMsg for Piece "
188 ++ show index ++ ", part: " ++ show i ++ " of length: "
192 PieceMsg index begin block -> do
193 putStrLn $ " <-- PieceMsg for Piece: "
199 putStrLn "ignoring irrelevant msg"
202 verifyHash :: ByteString -> ByteString -> Bool
203 verifyHash bs pieceHash =
204 take 20 (SHA1.hash bs) == pieceHash