1 {-# LANGUAGE OverloadedStrings #-}
2 module FuncTorrent.Peer
7 import Prelude hiding (lookup, concat, replicate, splitAt, take, filter)
9 import System.IO (Handle, BufferMode(..), hSetBuffering)
10 import Data.ByteString (ByteString, unpack, concat, hGet, hPut, take, empty)
11 import qualified Data.ByteString.Char8 as BC (length)
12 import Network (connectTo, PortID(..))
13 import Control.Monad.State
15 import Data.Word (Word8)
16 import Data.Map (Map, fromList, toList, (!), mapWithKey, adjust, filter)
17 import qualified Crypto.Hash.SHA1 as SHA1 (hash)
20 import FuncTorrent.Metainfo (Info(..), Metainfo(..))
21 import FuncTorrent.Utils (splitN, splitNum)
22 import FuncTorrent.Fileops (createDummyFile, writeFileAtOffset)
23 import FuncTorrent.PeerMsgs (Peer(..), PeerMsg(..), sendMsg, getMsg, genHandshakeMsg)
25 data PState = PState { handle :: Handle
28 , meInterested :: Bool
30 , heInterested :: Bool}
32 type PeerState = State PState
34 data PieceDlState = Pending
39 -- todo - map with index to a new data structure (peers who have that piece amd 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 mkPieceMap :: Integer -> ByteString -> [Integer] -> PieceMap
52 mkPieceMap numPieces pieceHash pLengths = fromList kvs
53 where kvs = [(i, PieceData { peers = []
57 | (i, h, pLen) <- zip3 [0..numPieces] hashes pLengths]
58 hashes = splitN 20 pieceHash
60 havePiece :: PieceMap -> Integer -> Bool
62 dlstate (pm ! index) == Have
64 connectToPeer :: Peer -> IO Handle
65 connectToPeer (Peer _ ip port) = do
66 h <- connectTo ip (PortNumber (fromIntegral port))
67 hSetBuffering h LineBuffering
70 doHandshake :: Handle -> Peer -> ByteString -> String -> IO ()
71 doHandshake h peer infoHash peerid = do
72 let hs = genHandshakeMsg infoHash peerid
74 putStrLn $ "--> handhake to peer: " ++ show peer
75 _ <- hGet h (length (unpack hs))
76 putStrLn $ "<-- handshake from peer: " ++ show peer
79 bitfieldToList :: [Word8] -> [Integer]
80 bitfieldToList bs = go bs 0
83 let setBits = [pos*8 + toInteger i | i <- [0..8], testBit b i]
85 setBits ++ go bs' (pos + 1)
87 -- helper functions to manipulate PeerState
90 -> Bool -- ^ meChoking
91 -> Bool -- ^ meInterested
92 -> Bool -- ^ heChoking
93 -> Bool -- ^ heInterested
95 toPeerState h p meCh meIn heCh heIn =
101 , meInterested = meIn }
103 -- simple algorithm to pick piece.
104 -- pick the first piece from 0 that is not downloaded yet.
105 pickPiece :: PieceMap -> Maybe Integer
107 (fst `liftM`) . headMay . toList . filter (\v -> dlstate v == Pending)
109 updatePieceAvailability :: PieceMap -> Peer -> [Integer] -> PieceMap
110 updatePieceAvailability pieceStatus p pieceList =
111 mapWithKey (\k pd -> if k `elem` pieceList
112 then (pd { peers = p : peers pd })
115 handlePeerMsgs :: Peer -> Metainfo -> String -> IO ()
116 handlePeerMsgs p m peerId = do
118 doHandshake h p (infoHash m) peerId
119 let pstate = toPeerState h p False False True True
120 pieceHash = pieces (info m)
121 numPieces = (toInteger . (`quot` 20) . BC.length) pieceHash
122 pLen = pieceLength (info m)
123 fileLen = lengthInBytes (info m)
124 fileName = name (info m)
125 pieceStatus = mkPieceMap numPieces pieceHash (splitNum fileLen pLen)
126 createDummyFile fileName (fromIntegral fileLen)
127 (r, _) <- runStateT (msgLoop pieceStatus fileName) pstate
130 msgLoop :: PieceMap -> FilePath -> StateT PState IO ()
131 msgLoop pieceStatus file = do
135 PState { meInterested = False, heChoking = True } -> do
136 liftIO $ sendMsg h InterestedMsg
137 gets peer >>= (\p -> liftIO $ putStrLn $ "--> InterestedMsg to peer: " ++ show p)
138 modify (\st -> st { meInterested = True })
139 msgLoop pieceStatus file
140 PState { meInterested = True, heChoking = False } -> do
141 case pickPiece pieceStatus of
142 Nothing -> liftIO $ putStrLn "Nothing to download"
144 let pLen = len (pieceStatus ! workPiece)
145 liftIO $ putStrLn $ "piece length = " ++ show pLen
146 pBS <- liftIO $ downloadPiece h workPiece pLen
147 if not $ verifyHash pBS (hash (pieceStatus ! workPiece))
149 liftIO $ putStrLn $ "Hash mismatch: " ++ show (hash (pieceStatus ! workPiece)) ++ " vs " ++ show (take 20 (SHA1.hash pBS))
151 let fileOffset = if workPiece == 0 then 0 else workPiece * len (pieceStatus ! (workPiece - 1))
152 liftIO $ putStrLn $ "Write into file at offset: " ++ show fileOffset
153 liftIO $ writeFileAtOffset file fileOffset pBS
154 msgLoop (adjust (\pieceData -> pieceData { dlstate = Have }) workPiece pieceStatus) file
156 msg <- liftIO $ getMsg h
157 gets peer >>= (\p -> liftIO $ putStrLn $ "<-- " ++ show msg ++ "from peer: " ++ show p)
160 liftIO $ sendMsg h KeepAliveMsg
161 gets peer >>= (\p -> liftIO $ putStrLn $ "--> " ++ "KeepAliveMsg to peer: " ++ show p)
162 msgLoop pieceStatus file
163 BitFieldMsg bss -> do
165 let pieceList = bitfieldToList (unpack bss)
166 pieceStatus' = updatePieceAvailability pieceStatus p pieceList
167 liftIO $ putStrLn $ show (length pieceList) ++ " Pieces"
168 -- for each pieceIndex in pieceList, make an entry in the pieceStatus
169 -- map with pieceIndex as the key and modify the value to add the peer.
170 -- download each of the piece in order
171 msgLoop pieceStatus' file
173 modify (\st -> st {heChoking = False })
174 msgLoop pieceStatus file
177 downloadPiece :: Handle -> Integer -> Integer -> IO ByteString
178 downloadPiece h index pieceLength = do
179 let chunks = splitNum pieceLength 16384
180 liftM concat $ forM (zip [0..] chunks) (\(i, pLen) -> do
181 sendMsg h (RequestMsg index (i*pLen) pLen)
182 putStrLn $ "--> " ++ "RequestMsg for Piece "
183 ++ show index ++ ", part: " ++ show i ++ " of length: "
187 PieceMsg index begin block -> do
188 putStrLn $ " <-- PieceMsg for Piece: "
194 putStrLn "ignoring irrelevant msg"
197 verifyHash :: ByteString -> ByteString -> Bool
198 verifyHash bs pieceHash =
199 take 20 (SHA1.hash bs) == pieceHash