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
 202
 203 {-
 204  -- Extension messages support (BEP-0010) --
 205
 206
 207    In the regular peer handshake, adventise support for extension protocol. Protocol
 208    extensions are done via the reserved bytes (8 of them) in the handshake message
 209    as detailed in BEP-0003. For this particular "Extension Protocol" extension, we use
 210    20th bit (counted from the right, from 0) is set to 1.
 211
 212    Once support for the extension protocol is established by the peer, the Peer is supposed
 213    to support one message with the ID 20. This is sent like a regular message with 4-byte
 214    length prefix and the msg id (20) in this case.
 215
 216    First byte of the payload of this message is either 0, which means it is a handshake
 217    msg.
 218
 219    The rest of the payload is a dictionary with various keys. All of them are optional. The
 220    one of interest at the moment for me is the one with key 'm' whose value is another
 221    dictionary of all supported extensions.
 222
 223    Here is where it gets interesting for us (to support magneturi. When the torrent client
 224    has only got a magneturi to look at, it has only got the list of trackers with it (we
 225    are not looking at the DHT case for the time being). So, it somehow needs to get the info
 226    dictionary. It gets this by talking to another peer in the network. To do that, the client
 227    needs to talk tracker protocol, get the list of peers and talk to peers using the above
 228    extension protocol to get the infodict as payload. Let us see how we can do that now.
 229
 230    If a peer already has the full infodict, then, the handshake message sent by that peer
 231    is something like this:
 232
 233      {'m': {'ut_metadata', 3}, 'metadata_size': 31235}
 234
 235    Note that the 'metadata_size' is not part of the value of the key 'm'.
 236    If we are a new client and are requesting the handshake to a peer, then we don't have
 237    the infodict yet, in which case, we only send the first part:
 238
 239      {'m': {'ut_metadata', 3}}
 240
 241    This is bencoded and sent across the wire. The value "3" (integer) against the key
 242    'ut_metadata" is an ordered integer within a client that identifies the extention.
 243    No two extension supported by the same client shares the same value. If the value is
 244    '0', then the extension is unsupported.
 245
 246    Here we use the BEP-0009, the metadata extension protocol. The metadata in this case
 247    is the infodict. The infodict itself is divided into 16KB sized pieces.
 248
 249    Here is a possible interaction between two peers:
 250
 251    1. Peer Pn comes up, gets the ip/ports of other peers, P0, P1.... Pn does not have the
 252    size of the infodict. Pn has advertised itself as supporting the extension protocol.
 253    It sends the handshake msg to other peers with this bit on in the reserved bytes.
 254    2. Let us say, P1 replied with a handshake. We check if it also supports the extension
 255    mechanism.
 256    3. Now we get into the extension message passing so that we have the info dict.
 257    To do that, we send the extension handshake (ut_metadata) m dict without the
 258    metadata_size. We get back the extension handshake with metadata_size. We take
 259    note of the size.
 260    4. We calculate the number of 16384 chunks in the total size of the metadata. That
 261    gives us the number of pieces the metadata has.
 262    5. We send a "request" extension msg:
 263      {'msg_type': 0, 'piece': 0}
 264    6. We recieve the "data" message.
 265      {'msg_type': 1, 'piece': 0, 'total_size': 3425} in bencoded format, followed by
 266      total_size bytes. total_size is 16KiB except perhaps for the last piece.
 267    7. If the peer does not have the requested piece, it sends the "reject" message.
 268      {'msg_type': 2, 'piece': 0}
 269    8. Repeat 5, 6/7 for every piece.
 270
 271    At this point, we have the infodict.
 272
 273 -)
 274