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 System.Directory (doesFileExist)
12 import Data.ByteString (ByteString, unpack, concat, hGet, hPut, take, empty)
13 import qualified Data.ByteString.Char8 as BC (length)
14 import Network (connectTo, PortID(..))
15 import Control.Monad.State
17 import Data.Word (Word8)
18 import Data.Map (Map, fromList, toList, (!), mapWithKey, traverseWithKey, adjust, filter)
19 import qualified Crypto.Hash.SHA1 as SHA1 (hash)
22 import FuncTorrent.Metainfo (Info(..), Metainfo(..))
23 import FuncTorrent.Utils (splitN, splitNum, createDummyFile, writeFileAtOffset, readFileAtOffset)
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 initPieceMap :: ByteString -> Integer -> Integer -> PieceMap
53 initPieceMap pieceHash fileLen pieceLen = fromList kvs
55 numPieces = (toInteger . (`quot` 20) . BC.length) pieceHash
56 kvs = [(i, PieceData { peers = []
60 | (i, h, pLen) <- zip3 [0..numPieces] hashes pLengths]
61 hashes = splitN 20 pieceHash
62 pLengths = (splitNum fileLen pieceLen)
64 updatePieceMap :: FilePath -> PieceMap -> IO PieceMap
65 updatePieceMap filePath pieceMap = do
66 dfe <- doesFileExist filePath
67 -- TODO: this is not enough, file should have the same size as well
69 then pieceMapFromFile filePath pieceMap
72 pieceMapFromFile :: FilePath -> PieceMap -> IO PieceMap
73 pieceMapFromFile filePath pieceMap = do
74 traverseWithKey f pieceMap
77 let offset = if k == 0 then 0 else k * len (pieceMap ! (k - 1))
78 isHashValid <- (flip verifyHash) (hash v) <$> (readFileAtOffset filePath offset (len v))
80 then return $ v { dlstate = Have }
83 havePiece :: PieceMap -> Integer -> Bool
85 dlstate (pm ! index) == Have
87 connectToPeer :: Peer -> IO Handle
88 connectToPeer (Peer _ ip port) = do
89 h <- connectTo ip (PortNumber (fromIntegral port))
90 hSetBuffering h LineBuffering
93 doHandshake :: Handle -> Peer -> ByteString -> String -> IO ()
94 doHandshake h peer infoHash peerid = do
95 let hs = genHandshakeMsg infoHash peerid
97 putStrLn $ "--> handhake to peer: " ++ show peer
98 _ <- hGet h (length (unpack hs))
99 putStrLn $ "<-- handshake from peer: " ++ show peer
102 bitfieldToList :: [Word8] -> [Integer]
103 bitfieldToList bs = go bs 0
106 let setBits = [pos*8 + toInteger i | i <- [0..8], testBit b i]
108 setBits ++ go bs' (pos + 1)
110 -- helper functions to manipulate PeerState
111 toPeerState :: Handle
113 -> Bool -- ^ meChoking
114 -> Bool -- ^ meInterested
115 -> Bool -- ^ heChoking
116 -> Bool -- ^ heInterested
118 toPeerState h p meCh meIn heCh heIn =
122 , heInterested = heIn
124 , meInterested = meIn }
126 -- simple algorithm to pick piece.
127 -- pick the first piece from 0 that is not downloaded yet.
128 pickPiece :: PieceMap -> Maybe Integer
130 (fst `liftM`) . headMay . toList . filter (\v -> dlstate v == Pending)
132 bytesDownloaded :: PieceMap -> Integer
134 sum . map (len . snd) . toList . filter (\v -> dlstate v == Have)
136 updatePieceAvailability :: PieceMap -> Peer -> [Integer] -> PieceMap
137 updatePieceAvailability pieceStatus p pieceList =
138 mapWithKey (\k pd -> if k `elem` pieceList
139 then (pd { peers = p : peers pd })
142 handlePeerMsgs :: Peer -> String -> Metainfo -> IO ()
143 handlePeerMsgs p peerId m = do
145 doHandshake h p (infoHash m) peerId
146 let pstate = toPeerState h p False False True True
147 pieceHash = pieces (info m)
148 pLen = pieceLength (info m)
149 fileLen = lengthInBytes (info m)
150 fileName = name (info m)
151 pieceStatus = initPieceMap pieceHash fileLen pLen
152 pieceStatus' <- updatePieceMap fileName pieceStatus
153 createDummyFile fileName (fromIntegral fileLen)
154 _ <- runStateT (msgLoop pieceStatus' fileName) 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: " ++ show (hash (pieceStatus ! workPiece)) ++ " vs " ++ show (take 20 (SHA1.hash pBS))
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
204 downloadPiece :: Handle -> Integer -> Integer -> IO ByteString
205 downloadPiece h index pieceLength = do
206 let chunks = splitNum pieceLength 16384
207 concat `liftM` forM (zip [0..] chunks) (\(i, pLen) -> do
208 sendMsg h (RequestMsg index (i*pLen) pLen)
209 putStrLn $ "--> " ++ "RequestMsg for Piece "
210 ++ show index ++ ", part: " ++ show i ++ " of length: "
214 PieceMsg index begin block -> do
215 putStrLn $ " <-- PieceMsg for Piece: "
221 putStrLn "ignoring irrelevant msg"
224 verifyHash :: ByteString -> ByteString -> Bool
225 verifyHash bs pieceHash =
226 take 20 (SHA1.hash bs) == pieceHash