1 {-# LANGUAGE OverloadedStrings #-}
2 module FuncTorrent.Peer
8 import Prelude hiding (lookup, concat, replicate, splitAt, take, drop)
10 import Control.Monad.State
11 import Data.ByteString (ByteString, unpack, concat, hGet, hPut, take, drop, empty)
13 import Data.Word (Word8)
14 import Data.Map ((!), adjust)
15 import Network (connectTo, PortID(..))
16 import System.IO (Handle, BufferMode(..), hSetBuffering, hClose)
18 import FuncTorrent.Metainfo (Metainfo(..))
19 import FuncTorrent.PeerMsgs (Peer(..), PeerMsg(..), sendMsg, getMsg, genHandshakeMsg)
20 import FuncTorrent.Utils (splitNum, verifyHash)
21 import FuncTorrent.PieceManager (PieceDlState(..), PieceData(..), PieceMap, pickPiece, updatePieceAvailability)
22 import qualified FuncTorrent.FileSystem as FS (MsgChannel, writePieceToDisk, Piece(..))
24 data PState = PState { handle :: Handle
27 , meInterested :: Bool
29 , heInterested :: Bool}
31 havePiece :: PieceMap -> Integer -> Bool
33 dlstate (pm ! index) == Have
35 connectToPeer :: Peer -> IO Handle
36 connectToPeer (Peer _ ip port) = do
37 h <- connectTo ip (PortNumber (fromIntegral port))
38 hSetBuffering h LineBuffering
41 doHandshake :: Bool -> Handle -> Peer -> ByteString -> String -> IO ()
42 doHandshake True h p infohash peerid = do
43 let hs = genHandshakeMsg infohash peerid
45 putStrLn $ "--> handhake to peer: " ++ show p
46 _ <- hGet h (length (unpack hs))
47 putStrLn $ "<-- handshake from peer: " ++ show p
49 doHandshake False h p infohash peerid = do
50 let hs = genHandshakeMsg infohash peerid
51 putStrLn "waiting for a handshake"
52 hsMsg <- hGet h (length (unpack hs))
53 putStrLn $ "<-- handshake from peer: " ++ show p
54 let rxInfoHash = take 20 $ drop 28 hsMsg
55 if rxInfoHash /= infohash
57 putStrLn "infoHashes does not match"
62 putStrLn $ "--> handhake to peer: " ++ show p
65 bitfieldToList :: [Word8] -> [Integer]
66 bitfieldToList bs = go bs 0
69 let setBits = [pos*8 + toInteger i | i <- [0..8], testBit b i]
71 setBits ++ go bs' (pos + 1)
73 -- helper functions to manipulate PeerState
76 -> Bool -- ^ meChoking
77 -> Bool -- ^ meInterested
78 -> Bool -- ^ heChoking
79 -> Bool -- ^ heInterested
81 toPeerState h p meCh meIn heCh heIn =
87 , meInterested = meIn }
89 handlePeerMsgs :: Peer -> String -> Metainfo -> PieceMap -> Bool -> FS.MsgChannel -> IO ()
90 handlePeerMsgs p peerId m pieceMap isClient c = do
92 doHandshake isClient h p (infoHash m) peerId
93 let pstate = toPeerState h p False False True True
94 _ <- runStateT (msgLoop pieceMap c) pstate
97 msgLoop :: PieceMap -> FS.MsgChannel -> StateT PState IO ()
98 msgLoop pieceStatus msgchannel = do
102 PState { meInterested = False, heChoking = True } -> do
103 liftIO $ sendMsg h InterestedMsg
104 gets peer >>= (\p -> liftIO $ putStrLn $ "--> InterestedMsg to peer: " ++ show p)
105 modify (\st' -> st' { meInterested = True })
106 msgLoop pieceStatus msgchannel
107 PState { meInterested = True, heChoking = False } ->
108 case pickPiece pieceStatus of
109 Nothing -> liftIO $ putStrLn "Nothing to download"
111 let pLen = len (pieceStatus ! workPiece)
112 liftIO $ putStrLn $ "piece length = " ++ show pLen
113 pBS <- liftIO $ downloadPiece h workPiece pLen
114 if not $ verifyHash pBS (hash (pieceStatus ! workPiece))
116 liftIO $ putStrLn "Hash mismatch"
118 liftIO $ putStrLn $ "Write piece: " ++ show workPiece
119 liftIO $ FS.writePieceToDisk msgchannel workPiece pBS
120 msgLoop (adjust (\pieceData -> pieceData { dlstate = Have }) workPiece pieceStatus) msgchannel
122 msg <- liftIO $ getMsg h
123 gets peer >>= (\p -> liftIO $ putStrLn $ "<-- " ++ show msg ++ "from peer: " ++ show p)
126 liftIO $ sendMsg h KeepAliveMsg
127 gets peer >>= (\p -> liftIO $ putStrLn $ "--> " ++ "KeepAliveMsg to peer: " ++ show p)
128 msgLoop pieceStatus msgchannel
129 BitFieldMsg bss -> do
131 let pieceList = bitfieldToList (unpack bss)
132 pieceStatus' = updatePieceAvailability pieceStatus p pieceList
133 liftIO $ putStrLn $ show (length pieceList) ++ " Pieces"
134 -- for each pieceIndex in pieceList, make an entry in the pieceStatus
135 -- map with pieceIndex as the key and modify the value to add the peer.
136 -- download each of the piece in order
137 msgLoop pieceStatus' msgchannel
139 modify (\st' -> st' {heChoking = False })
140 msgLoop pieceStatus msgchannel
142 modify (\st' -> st' {heChoking = True })
143 msgLoop pieceStatus msgchannel
145 modify (\st' -> st' {heInterested = True})
146 msgLoop pieceStatus msgchannel
147 NotInterestedMsg -> do
148 modify (\st' -> st' {heInterested = False})
149 msgLoop pieceStatus msgchannel
150 CancelMsg _ _ _ -> -- check if valid index, begin, length
151 msgLoop pieceStatus msgchannel
153 msgLoop pieceStatus msgchannel
156 let pieceStatus' = updatePieceAvailability pieceStatus p [idx]
157 msgLoop pieceStatus' msgchannel
158 -- handle RequestMsg. No need to handle PieceMsg here.
162 downloadPiece :: Handle -> Integer -> Integer -> IO ByteString
163 downloadPiece h index pieceLength = do
164 let chunks = splitNum pieceLength 16384
165 concat `liftM` forM (zip [0..] chunks) (\(i, pLen) -> do
166 sendMsg h (RequestMsg index (i*pLen) pLen)
167 putStrLn $ "--> " ++ "RequestMsg for Piece "
168 ++ show index ++ ", part: " ++ show i ++ " of length: "
172 PieceMsg index begin block -> do
173 putStrLn $ " <-- PieceMsg for Piece: "
179 putStrLn $ "ignoring irrelevant msg: " ++ show msg