src/FuncTorrent/Peer.hs

   1 {-
   2  - Copyright (C) 2015-2016 Ramakrishnan Muthukrishnan <ram@rkrishnan.org>
   3  -
   4  - This file is part of FuncTorrent.
   5  -
   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.
  10  -
  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.
  15  -
  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/>
  18  -}
  19
  20 {-# LANGUAGE OverloadedStrings #-}
  21
  22 module FuncTorrent.Peer
  23     (PieceMap,
  24      handlePeerMsgs
  25     ) where
  26
  27 import Prelude hiding (lookup, concat, replicate, splitAt, take, drop)
  28
  29 import Control.Monad.State
  30 import Data.ByteString (ByteString, unpack, concat, hGet, hPut, take, drop, empty)
  31 import Data.Bits
  32 import Data.Word (Word8)
  33 import Data.Map ((!), adjust)
  34 import Network (connectTo, PortID(..))
  35 import System.IO (Handle, BufferMode(..), hSetBuffering, hClose)
  36
  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)
  42
  43 data PState = PState { handle :: Handle
  44                      , peer :: Peer
  45                      , meChoking :: Bool
  46                      , meInterested :: Bool
  47                      , heChoking :: Bool
  48                      , heInterested :: Bool}
  49
  50 havePiece :: PieceMap -> Integer -> Bool
  51 havePiece pm index =
  52   dlstate (pm ! index) == Have
  53
  54 connectToPeer :: Peer -> IO Handle
  55 connectToPeer (Peer ip port) = do
  56   h <- connectTo ip (PortNumber (fromIntegral port))
  57   hSetBuffering h LineBuffering
  58   return h
  59
  60 doHandshake :: Bool -> Handle -> Peer -> ByteString -> String -> IO ()
  61 doHandshake True h p infohash peerid = do
  62   let hs = genHandshakeMsg infohash peerid
  63   hPut h hs
  64   putStrLn $ "--> handhake to peer: " ++ show p
  65   _ <- hGet h (length (unpack hs))
  66   putStrLn $ "<-- handshake from peer: " ++ show p
  67   return ()
  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
  75     then do
  76     putStrLn "infoHashes does not match"
  77     hClose h
  78     return ()
  79     else do
  80     _ <- hPut h hs
  81     putStrLn $ "--> handhake to peer: " ++ show p
  82     return ()
  83
  84 bitfieldToList :: [Word8] -> [Integer]
  85 bitfieldToList bs = go bs 0
  86   where go [] _ = []
  87         go (b:bs') pos =
  88           let setBits = [pos*8 + toInteger i | i <- [0..8], testBit b i]
  89           in
  90            setBits ++ go bs' (pos + 1)
  91
  92 -- helper functions to manipulate PeerState
  93 toPeerState :: Handle
  94             -> Peer
  95             -> Bool  -- ^ meChoking
  96             -> Bool  -- ^ meInterested
  97             -> Bool  -- ^ heChoking
  98             -> Bool  -- ^ heInterested
  99             -> PState
 100 toPeerState h p meCh meIn heCh heIn =
 101   PState { handle = h
 102          , peer = p
 103          , heChoking = heCh
 104          , heInterested = heIn
 105          , meChoking = meCh
 106          , meInterested = meIn }
 107
 108 handlePeerMsgs :: Peer -> String -> Metainfo -> PieceMap -> Bool -> FS.MsgChannel -> IO ()
 109 handlePeerMsgs p peerId m pieceMap isClient c = do
 110   h <- connectToPeer p
 111   doHandshake isClient h p (infoHash m) peerId
 112   let pstate = toPeerState h p False False True True
 113   _ <- runStateT (msgLoop pieceMap c) pstate
 114   return ()
 115
 116 msgLoop :: PieceMap -> FS.MsgChannel -> StateT PState IO ()
 117 msgLoop pieceStatus msgchannel = do
 118   h <- gets handle
 119   st <- get
 120   case st of
 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"
 129         Just workPiece -> do
 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))
 134             then
 135             liftIO $ putStrLn "Hash mismatch"
 136             else do
 137             liftIO $ putStrLn $ "Write piece: " ++ show workPiece
 138             liftIO $ FS.writePieceToDisk msgchannel workPiece pBS
 139             msgLoop (adjust (\pieceData -> pieceData { dlstate = Have }) workPiece pieceStatus) msgchannel
 140     _ -> do
 141       msg <- liftIO $ getMsg h
 142       gets peer >>= (\p -> liftIO $ putStrLn $ "<-- " ++ show msg ++ " from peer: " ++ show p)
 143       case msg of
 144         KeepAliveMsg -> do
 145           liftIO $ sendMsg h KeepAliveMsg
 146           gets peer >>= (\p -> liftIO $ putStrLn $ "--> " ++ "KeepAliveMsg to peer: " ++ show p)
 147           msgLoop pieceStatus msgchannel
 148         BitFieldMsg bss -> do
 149           p <- gets peer
 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
 157         UnChokeMsg -> do
 158           modify (\st' -> st' {heChoking = False })
 159           msgLoop pieceStatus msgchannel
 160         ChokeMsg -> do
 161           modify (\st' -> st' {heChoking = True })
 162           msgLoop pieceStatus msgchannel
 163         InterestedMsg -> do
 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
 171         PortMsg _ ->
 172           msgLoop pieceStatus msgchannel
 173         HaveMsg idx -> do
 174           p <- gets peer
 175           let pieceStatus' = updatePieceAvailability pieceStatus p [idx]
 176           msgLoop pieceStatus' msgchannel
 177         _ -> do
 178           liftIO $ putStrLn $ ".. not doing anything with the msg"
 179           msgLoop pieceStatus msgchannel
 180         -- No need to handle PieceMsg and RequestMsg here.
 181
 182
 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: "
 190                                                 ++ show pLen
 191                                               msg <- getMsg h
 192                                               case msg of
 193                                                 PieceMsg index begin block -> do
 194                                                   putStrLn $ " <-- PieceMsg for Piece: "
 195                                                     ++ show index
 196                                                     ++ ", offset: "
 197                                                     ++ show begin
 198                                                   return block
 199                                                 _ -> do
 200                                                   putStrLn $ "ignoring irrelevant msg: " ++ show msg
 201                                                   return empty)
 202