2 - Copyright (C) 2015-2016 Ramakrishnan Muthukrishnan <ram@rkrishnan.org>
4 - This file is part of FuncTorrent.
6 - FuncTorrent is free software; you can redistribute it and/or modify
7 - it under the terms of the GNU General Public License as published by
8 - the Free Software Foundation; either version 3 of the License, or
9 - (at your option) any later version.
11 - FuncTorrent is distributed in the hope that it will be useful,
12 - but WITHOUT ANY WARRANTY; without even the implied warranty of
13 - MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
14 - GNU General Public License for more details.
16 - You should have received a copy of the GNU General Public License
17 - along with FuncTorrent; if not, see <http://www.gnu.org/licenses/>
20 {-# LANGUAGE OverloadedStrings #-}
22 module FuncTorrent.Peer
27 import Prelude hiding (lookup, concat, replicate, splitAt, take, drop)
29 import Control.Monad.State
30 import Data.ByteString (ByteString, unpack, concat, hGet, hPut, take, drop, empty)
32 import Data.Word (Word8)
33 import Data.Map ((!), adjust)
34 import Network (connectTo, PortID(..))
35 import System.IO (Handle, BufferMode(..), hSetBuffering, hClose)
37 import FuncTorrent.Metainfo (Metainfo(..))
38 import FuncTorrent.PeerMsgs (Peer(..), PeerMsg(..), sendMsg, getMsg, genHandshakeMsg)
39 import FuncTorrent.Utils (splitNum, verifyHash)
40 import FuncTorrent.PieceManager (PieceDlState(..), PieceData(..), PieceMap, pickPiece, updatePieceAvailability)
41 import qualified FuncTorrent.FileSystem as FS (MsgChannel, writePieceToDisk)
43 data PState = PState { handle :: Handle
46 , meInterested :: Bool
48 , heInterested :: Bool}
50 havePiece :: PieceMap -> Integer -> Bool
52 dlstate (pm ! index) == Have
54 connectToPeer :: Peer -> IO Handle
55 connectToPeer (Peer ip port) = do
56 h <- connectTo ip (PortNumber (fromIntegral port))
57 hSetBuffering h LineBuffering
60 doHandshake :: Bool -> Handle -> Peer -> ByteString -> String -> IO ()
61 doHandshake True h p infohash peerid = do
62 let hs = genHandshakeMsg infohash peerid
64 putStrLn $ "--> handhake to peer: " ++ show p
65 _ <- hGet h (length (unpack hs))
66 putStrLn $ "<-- handshake from peer: " ++ show p
68 doHandshake False h p infohash peerid = do
69 let hs = genHandshakeMsg infohash peerid
70 putStrLn "waiting for a handshake"
71 hsMsg <- hGet h (length (unpack hs))
72 putStrLn $ "<-- handshake from peer: " ++ show p
73 let rxInfoHash = take 20 $ drop 28 hsMsg
74 if rxInfoHash /= infohash
76 putStrLn "infoHashes does not match"
81 putStrLn $ "--> handhake to peer: " ++ show p
84 bitfieldToList :: [Word8] -> [Integer]
85 bitfieldToList bs = go bs 0
88 let setBits = [pos*8 + toInteger i | i <- [0..8], testBit b i]
90 setBits ++ go bs' (pos + 1)
92 -- helper functions to manipulate PeerState
95 -> Bool -- ^ meChoking
96 -> Bool -- ^ meInterested
97 -> Bool -- ^ heChoking
98 -> Bool -- ^ heInterested
100 toPeerState h p meCh meIn heCh heIn =
104 , heInterested = heIn
106 , meInterested = meIn }
108 handlePeerMsgs :: Peer -> String -> Metainfo -> PieceMap -> Bool -> FS.MsgChannel -> IO ()
109 handlePeerMsgs p peerId m pieceMap isClient c = do
111 doHandshake isClient h p (infoHash m) peerId
112 let pstate = toPeerState h p False False True True
113 _ <- runStateT (msgLoop pieceMap c) pstate
116 msgLoop :: PieceMap -> FS.MsgChannel -> StateT PState IO ()
117 msgLoop pieceStatus msgchannel = do
121 PState { meInterested = False, heChoking = True } -> do
122 liftIO $ sendMsg h InterestedMsg
123 gets peer >>= (\p -> liftIO $ putStrLn $ "--> InterestedMsg to peer: " ++ show p)
124 modify (\st' -> st' { meInterested = True })
125 msgLoop pieceStatus msgchannel
126 PState { meInterested = True, heChoking = False } ->
127 case pickPiece pieceStatus of
128 Nothing -> liftIO $ putStrLn "Nothing to download"
130 let pLen = len (pieceStatus ! workPiece)
131 liftIO $ putStrLn $ "piece length = " ++ show pLen
132 pBS <- liftIO $ downloadPiece h workPiece pLen
133 if not $ verifyHash pBS (hash (pieceStatus ! workPiece))
135 liftIO $ putStrLn "Hash mismatch"
137 liftIO $ putStrLn $ "Write piece: " ++ show workPiece
138 liftIO $ FS.writePieceToDisk msgchannel workPiece pBS
139 msgLoop (adjust (\pieceData -> pieceData { dlstate = Have }) workPiece pieceStatus) msgchannel
141 msg <- liftIO $ getMsg h
142 gets peer >>= (\p -> liftIO $ putStrLn $ "<-- " ++ show msg ++ " from peer: " ++ show p)
145 liftIO $ sendMsg h KeepAliveMsg
146 gets peer >>= (\p -> liftIO $ putStrLn $ "--> " ++ "KeepAliveMsg to peer: " ++ show p)
147 msgLoop pieceStatus msgchannel
148 BitFieldMsg bss -> do
150 let pieceList = bitfieldToList (unpack bss)
151 pieceStatus' = updatePieceAvailability pieceStatus p pieceList
152 liftIO $ putStrLn $ show (length pieceList) ++ " Pieces"
153 -- for each pieceIndex in pieceList, make an entry in the pieceStatus
154 -- map with pieceIndex as the key and modify the value to add the peer.
155 -- download each of the piece in order
156 msgLoop pieceStatus' msgchannel
158 modify (\st' -> st' {heChoking = False })
159 msgLoop pieceStatus msgchannel
161 modify (\st' -> st' {heChoking = True })
162 msgLoop pieceStatus msgchannel
164 modify (\st' -> st' {heInterested = True})
165 msgLoop pieceStatus msgchannel
166 NotInterestedMsg -> do
167 modify (\st' -> st' {heInterested = False})
168 msgLoop pieceStatus msgchannel
169 CancelMsg _ _ _ -> -- check if valid index, begin, length
170 msgLoop pieceStatus msgchannel
172 msgLoop pieceStatus msgchannel
175 let pieceStatus' = updatePieceAvailability pieceStatus p [idx]
176 msgLoop pieceStatus' msgchannel
178 liftIO $ putStrLn $ ".. not doing anything with the msg"
179 msgLoop pieceStatus msgchannel
180 -- No need to handle PieceMsg and RequestMsg here.
183 downloadPiece :: Handle -> Integer -> Integer -> IO ByteString
184 downloadPiece h index pieceLength = do
185 let chunks = splitNum pieceLength 16384
186 concat `liftM` forM (zip [0..] chunks) (\(i, pLen) -> do
187 sendMsg h (RequestMsg index (i*pLen) pLen)
188 putStrLn $ "--> " ++ "RequestMsg for Piece "
189 ++ show index ++ ", part: " ++ show i ++ " of length: "
193 PieceMsg index begin block -> do
194 putStrLn $ " <-- PieceMsg for Piece: "
200 putStrLn $ "ignoring irrelevant msg: " ++ show msg