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