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     (Peer(..),
  24      PieceMap,
  25      handlePeerMsgs
  26     ) where
  27
  28 import Prelude hiding (lookup, concat, replicate, splitAt, take, drop)
  29
  30 import Control.Monad.State
  31 import Data.ByteString (ByteString, unpack, concat, hGet, hPut, take, drop, empty)
  32 import Data.Bits
  33 import Data.Word (Word8)
  34 import Data.Map ((!), adjust)
  35 import Network (connectTo, PortID(..))
  36 import System.IO (Handle, BufferMode(..), hSetBuffering, hClose)
  37
  38 import FuncTorrent.Metainfo (Metainfo(..))
  39 import FuncTorrent.PeerMsgs (Peer(..), PeerMsg(..), sendMsg, getMsg, genHandshakeMsg)
  40 import FuncTorrent.Utils (splitNum, verifyHash)
  41 import FuncTorrent.PieceManager (PieceDlState(..), PieceData(..), PieceMap, pickPiece, updatePieceAvailability)
  42 import qualified FuncTorrent.FileSystem as FS (MsgChannel, writePieceToDisk)
  43
  44 data PState = PState { handle :: Handle
  45                      , peer :: Peer
  46                      , meChoking :: Bool
  47                      , meInterested :: Bool
  48                      , heChoking :: Bool
  49                      , heInterested :: Bool}
  50
  51 havePiece :: PieceMap -> Integer -> Bool
  52 havePiece pm index =
  53   dlstate (pm ! index) == Have
  54
  55 connectToPeer :: Peer -> IO Handle
  56 connectToPeer (Peer _ ip port) = do
  57   h <- connectTo ip (PortNumber (fromIntegral port))
  58   hSetBuffering h LineBuffering
  59   return h
  60
  61 doHandshake :: Bool -> Handle -> Peer -> ByteString -> String -> IO ()
  62 doHandshake True h p infohash peerid = do
  63   let hs = genHandshakeMsg infohash peerid
  64   hPut h hs
  65   putStrLn $ "--> handhake to peer: " ++ show p
  66   _ <- hGet h (length (unpack hs))
  67   putStrLn $ "<-- handshake from peer: " ++ show p
  68   return ()
  69 doHandshake False h p infohash peerid = do
  70   let hs = genHandshakeMsg infohash peerid
  71   putStrLn "waiting for a handshake"
  72   hsMsg <- hGet h (length (unpack hs))
  73   putStrLn $ "<-- handshake from peer: " ++ show p
  74   let rxInfoHash = take 20 $ drop 28 hsMsg
  75   if rxInfoHash /= infohash
  76     then do
  77     putStrLn "infoHashes does not match"
  78     hClose h
  79     return ()
  80     else do
  81     _ <- hPut h hs
  82     putStrLn $ "--> handhake to peer: " ++ show p
  83     return ()
  84
  85 bitfieldToList :: [Word8] -> [Integer]
  86 bitfieldToList bs = go bs 0
  87   where go [] _ = []
  88         go (b:bs') pos =
  89           let setBits = [pos*8 + toInteger i | i <- [0..8], testBit b i]
  90           in
  91            setBits ++ go bs' (pos + 1)
  92
  93 -- helper functions to manipulate PeerState
  94 toPeerState :: Handle
  95             -> Peer
  96             -> Bool  -- ^ meChoking
  97             -> Bool  -- ^ meInterested
  98             -> Bool  -- ^ heChoking
  99             -> Bool  -- ^ heInterested
 100             -> PState
 101 toPeerState h p meCh meIn heCh heIn =
 102   PState { handle = h
 103          , peer = p
 104          , heChoking = heCh
 105          , heInterested = heIn
 106          , meChoking = meCh
 107          , meInterested = meIn }
 108
 109 handlePeerMsgs :: Peer -> String -> Metainfo -> PieceMap -> Bool -> FS.MsgChannel -> IO ()
 110 handlePeerMsgs p peerId m pieceMap isClient c = do
 111   h <- connectToPeer p
 112   doHandshake isClient h p (infoHash m) peerId
 113   let pstate = toPeerState h p False False True True
 114   _ <- runStateT (msgLoop pieceMap c) pstate
 115   return ()
 116
 117 msgLoop :: PieceMap -> FS.MsgChannel -> StateT PState IO ()
 118 msgLoop pieceStatus msgchannel = do
 119   h <- gets handle
 120   st <- get
 121   case st of
 122     PState { meInterested = False, heChoking = True } -> do
 123       liftIO $ sendMsg h InterestedMsg
 124       gets peer >>= (\p -> liftIO $ putStrLn $ "--> InterestedMsg to peer: " ++ show p)
 125       modify (\st' -> st' { meInterested = True })
 126       msgLoop pieceStatus msgchannel
 127     PState { meInterested = True, heChoking = False } ->
 128       case pickPiece pieceStatus of
 129         Nothing -> liftIO $ putStrLn "Nothing to download"
 130         Just workPiece -> do
 131           let pLen = len (pieceStatus ! workPiece)
 132           liftIO $ putStrLn $ "piece length = " ++ show pLen
 133           pBS <- liftIO $ downloadPiece h workPiece pLen
 134           if not $ verifyHash pBS (hash (pieceStatus ! workPiece))
 135             then
 136             liftIO $ putStrLn "Hash mismatch"
 137             else do
 138             liftIO $ putStrLn $ "Write piece: " ++ show workPiece
 139             liftIO $ FS.writePieceToDisk msgchannel workPiece pBS
 140             msgLoop (adjust (\pieceData -> pieceData { dlstate = Have }) workPiece pieceStatus) msgchannel
 141     _ -> do
 142       msg <- liftIO $ getMsg h
 143       gets peer >>= (\p -> liftIO $ putStrLn $ "<-- " ++ show msg ++ " from peer: " ++ show p)
 144       case msg of
 145         KeepAliveMsg -> do
 146           liftIO $ sendMsg h KeepAliveMsg
 147           gets peer >>= (\p -> liftIO $ putStrLn $ "--> " ++ "KeepAliveMsg to peer: " ++ show p)
 148           msgLoop pieceStatus msgchannel
 149         BitFieldMsg bss -> do
 150           p <- gets peer
 151           let pieceList = bitfieldToList (unpack bss)
 152               pieceStatus' = updatePieceAvailability pieceStatus p pieceList
 153           liftIO $ putStrLn $ show (length pieceList) ++ " Pieces"
 154           -- for each pieceIndex in pieceList, make an entry in the pieceStatus
 155           -- map with pieceIndex as the key and modify the value to add the peer.
 156           -- download each of the piece in order
 157           msgLoop pieceStatus' msgchannel
 158         UnChokeMsg -> do
 159           modify (\st' -> st' {heChoking = False })
 160           msgLoop pieceStatus msgchannel
 161         ChokeMsg -> do
 162           modify (\st' -> st' {heChoking = True })
 163           msgLoop pieceStatus msgchannel
 164         InterestedMsg -> do
 165           modify (\st' -> st' {heInterested = True})
 166           msgLoop pieceStatus msgchannel
 167         NotInterestedMsg -> do
 168           modify (\st' -> st' {heInterested = False})
 169           msgLoop pieceStatus msgchannel
 170         CancelMsg _ _ _ -> -- check if valid index, begin, length
 171           msgLoop pieceStatus msgchannel
 172         PortMsg _ ->
 173           msgLoop pieceStatus msgchannel
 174         HaveMsg idx -> do
 175           p <- gets peer
 176           let pieceStatus' = updatePieceAvailability pieceStatus p [idx]
 177           msgLoop pieceStatus' msgchannel
 178         _ -> do
 179           liftIO $ putStrLn $ ".. not doing anything with the msg"
 180           msgLoop pieceStatus msgchannel
 181         -- No need to handle PieceMsg and RequestMsg here.
 182
 183
 184 downloadPiece :: Handle -> Integer -> Integer -> IO ByteString
 185 downloadPiece h index pieceLength = do
 186   let chunks = splitNum pieceLength 16384
 187   concat `liftM` forM (zip [0..] chunks) (\(i, pLen) -> do
 188                                               sendMsg h (RequestMsg index (i*pLen) pLen)
 189                                               putStrLn $ "--> " ++ "RequestMsg for Piece "
 190                                                 ++ show index ++ ", part: " ++ show i ++ " of length: "
 191                                                 ++ show pLen
 192                                               msg <- getMsg h
 193                                               case msg of
 194                                                 PieceMsg index begin block -> do
 195                                                   putStrLn $ " <-- PieceMsg for Piece: "
 196                                                     ++ show index
 197                                                     ++ ", offset: "
 198                                                     ++ show begin
 199                                                   return block
 200                                                 _ -> do
 201                                                   putStrLn $ "ignoring irrelevant msg: " ++ show msg
 202                                                   return empty)
 203