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
26 import Prelude hiding (lookup, concat, replicate, splitAt, take, drop)
28 import Control.Monad.State
29 import Data.ByteString (ByteString, unpack, concat, hGet, hPut, take, drop, empty)
31 import Data.Word (Word8)
32 import Data.Map ((!), adjust)
33 import Network (connectTo, PortID(..))
34 import System.IO (Handle, BufferMode(..), hSetBuffering, hClose)
36 import FuncTorrent.Metainfo (Metainfo(..))
37 import FuncTorrent.PeerMsgs (Peer(..), PeerMsg(..), sendMsg, getMsg, genHandshakeMsg)
38 import FuncTorrent.Utils (splitNum, verifyHash)
39 import FuncTorrent.PieceManager (PieceDlState(..), PieceData(..), PieceMap, pickPiece, updatePieceAvailability)
40 import qualified FuncTorrent.FileSystem as FS (MsgChannel, writePieceToDisk)
42 data PState = PState { handle :: Handle
45 , meInterested :: Bool
47 , heInterested :: Bool}
49 havePiece :: PieceMap -> Integer -> Bool
51 dlstate (pm ! index) == Have
53 connectToPeer :: Peer -> IO Handle
54 connectToPeer (Peer ip port) = do
55 h <- connectTo ip (PortNumber (fromIntegral port))
56 hSetBuffering h LineBuffering
59 doHandshake :: Bool -> Handle -> Peer -> ByteString -> String -> IO ()
60 doHandshake True h p infohash peerid = do
61 let hs = genHandshakeMsg infohash peerid
63 putStrLn $ "--> handhake to peer: " ++ show p
64 _ <- hGet h (length (unpack hs))
65 putStrLn $ "<-- handshake from peer: " ++ show p
67 doHandshake False h p infohash peerid = do
68 let hs = genHandshakeMsg infohash peerid
69 putStrLn "waiting for a handshake"
70 hsMsg <- hGet h (length (unpack hs))
71 putStrLn $ "<-- handshake from peer: " ++ show p
72 let rxInfoHash = take 20 $ drop 28 hsMsg
73 if rxInfoHash /= infohash
75 putStrLn "infoHashes does not match"
80 putStrLn $ "--> handhake to peer: " ++ show p
83 bitfieldToList :: [Word8] -> [Integer]
84 bitfieldToList bs = go bs 0
87 let setBits = [pos*8 + toInteger i | i <- [0..8], testBit b i]
89 setBits ++ go bs' (pos + 1)
91 -- helper functions to manipulate PeerState
94 -> Bool -- ^ meChoking
95 -> Bool -- ^ meInterested
96 -> Bool -- ^ heChoking
97 -> Bool -- ^ heInterested
99 toPeerState h p meCh meIn heCh heIn =
103 , heInterested = heIn
105 , meInterested = meIn }
107 handlePeerMsgs :: Peer -> String -> Metainfo -> PieceMap -> Bool -> FS.MsgChannel -> IO ()
108 handlePeerMsgs p peerId m pieceMap isClient c = do
110 doHandshake isClient h p (infoHash m) peerId
111 let pstate = toPeerState h p False False True True
112 _ <- runStateT (msgLoop pieceMap c) pstate
115 msgLoop :: PieceMap -> FS.MsgChannel -> StateT PState IO ()
116 msgLoop pieceStatus msgchannel = do
120 PState { meInterested = False, heChoking = True } -> do
121 liftIO $ sendMsg h InterestedMsg
122 gets peer >>= (\p -> liftIO $ putStrLn $ "--> InterestedMsg to peer: " ++ show p)
123 modify (\st' -> st' { meInterested = True })
124 msgLoop pieceStatus msgchannel
125 PState { meInterested = True, heChoking = False } ->
126 case pickPiece pieceStatus of
127 Nothing -> liftIO $ putStrLn "Nothing to download"
129 let pLen = len (pieceStatus ! workPiece)
130 liftIO $ putStrLn $ "piece length = " ++ show pLen
131 pBS <- liftIO $ downloadPiece h workPiece pLen
132 if not $ verifyHash pBS (hash (pieceStatus ! workPiece))
134 liftIO $ putStrLn "Hash mismatch"
136 liftIO $ putStrLn $ "Write piece: " ++ show workPiece
137 liftIO $ FS.writePieceToDisk msgchannel workPiece pBS
138 msgLoop (adjust (\pieceData -> pieceData { dlstate = Have }) workPiece pieceStatus) msgchannel
140 msg <- liftIO $ getMsg h
141 gets peer >>= (\p -> liftIO $ putStrLn $ "<-- " ++ show msg ++ " from peer: " ++ show p)
144 liftIO $ sendMsg h KeepAliveMsg
145 gets peer >>= (\p -> liftIO $ putStrLn $ "--> " ++ "KeepAliveMsg to peer: " ++ show p)
146 msgLoop pieceStatus msgchannel
147 BitFieldMsg bss -> do
149 let pieceList = bitfieldToList (unpack bss)
150 pieceStatus' = updatePieceAvailability pieceStatus p pieceList
151 liftIO $ putStrLn $ show (length pieceList) ++ " Pieces"
152 -- for each pieceIndex in pieceList, make an entry in the pieceStatus
153 -- map with pieceIndex as the key and modify the value to add the peer.
154 -- download each of the piece in order
155 msgLoop pieceStatus' msgchannel
157 modify (\st' -> st' {heChoking = False })
158 msgLoop pieceStatus msgchannel
160 modify (\st' -> st' {heChoking = True })
161 msgLoop pieceStatus msgchannel
163 modify (\st' -> st' {heInterested = True})
164 msgLoop pieceStatus msgchannel
165 NotInterestedMsg -> do
166 modify (\st' -> st' {heInterested = False})
167 msgLoop pieceStatus msgchannel
168 CancelMsg {} -> -- check if valid index, begin, length
169 msgLoop pieceStatus msgchannel
171 msgLoop pieceStatus msgchannel
174 let pieceStatus' = updatePieceAvailability pieceStatus p [idx]
175 msgLoop pieceStatus' msgchannel
177 liftIO $ putStrLn ".. not doing anything with the msg"
178 msgLoop pieceStatus msgchannel
179 -- No need to handle PieceMsg and RequestMsg here.
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: " ++ show msg