src/FuncTorrent/PeerMsgs.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.PeerMsgs
  23        (genHandshakeMsg,
  24         sendMsg,
  25         getMsg,
  26         Peer(..),
  27         makePeer,
  28         PeerMsg(..)
  29        ) where
  30
  31 import Prelude hiding (lookup, concat, replicate, splitAt, take)
  32
  33 import System.IO (Handle)
  34 import Data.ByteString (ByteString, pack, unpack, concat, hGet, hPut, singleton)
  35 import Data.ByteString.Lazy (fromStrict, fromChunks, toStrict)
  36 import Data.ByteString.Char8 as BC (splitAt)
  37 import qualified Data.ByteString.Char8 as BC (replicate, pack)
  38 import Control.Monad (replicateM)
  39 import Control.Applicative (liftA3)
  40
  41 import Data.Binary (Binary(..), decode, encode)
  42 import Data.Binary.Put (putWord32be, putWord16be, putWord8)
  43 import Data.Binary.Get (getWord32be, getWord16be, getWord8, runGet)
  44
  45 import FuncTorrent.Utils (toIP, toPort)
  46
  47 -- | Peer is a IP address, port tuple
  48 data Peer = Peer IP Port
  49           deriving (Show, Eq)
  50
  51 type ID = String
  52 type IP = String
  53 type Port = Integer
  54
  55 data PeerMsg = KeepAliveMsg
  56              | ChokeMsg
  57              | UnChokeMsg
  58              | InterestedMsg
  59              | NotInterestedMsg
  60              | HaveMsg Integer
  61              | BitFieldMsg ByteString
  62              | RequestMsg Integer Integer Integer
  63              | PieceMsg Integer Integer ByteString
  64              | CancelMsg Integer Integer Integer
  65              | PortMsg Port
  66              | ExtendedMsg Integer ByteString
  67              deriving (Show)
  68
  69 instance Binary PeerMsg where
  70   put msg = case msg of
  71              KeepAliveMsg -> putWord32be 0
  72              ChokeMsg -> do putWord32be 1
  73                             putWord8 0
  74              UnChokeMsg -> do putWord32be 1
  75                               putWord8 1
  76              InterestedMsg -> do putWord32be 1
  77                                  putWord8 2
  78              NotInterestedMsg -> do putWord32be 1
  79                                     putWord8 3
  80              HaveMsg i -> do putWord32be 5
  81                              putWord8 4
  82                              putWord32be (fromIntegral i)
  83              BitFieldMsg bf -> do putWord32be $ fromIntegral (1 + bfListLen)
  84                                   putWord8 5
  85                                   mapM_ putWord8 bfList
  86                                     where bfList = unpack bf
  87                                           bfListLen = length bfList
  88              RequestMsg i o l -> do putWord32be 13
  89                                     putWord8 6
  90                                     putIndexOffsetLength i o l
  91              PieceMsg i o b -> do putWord32be $ fromIntegral (9 + blocklen)
  92                                   putWord8 7
  93                                   putWord32be (fromIntegral i)
  94                                   putWord32be (fromIntegral o)
  95                                   mapM_ putWord8 blockList
  96                                     where blockList = unpack b
  97                                           blocklen = length blockList
  98              CancelMsg i o l -> do putWord32be 13
  99                                    putWord8 8
 100                                    putIndexOffsetLength i o l
 101              PortMsg p -> do putWord32be 3
 102                              putWord8 9
 103                              putWord16be (fromIntegral p)
 104              ExtendedHandshakeMsg t b-> do putWord32be msgLen
 105                                            putWord8 20
 106                                            putWord8 t -- 0 => handshake msg
 107                                            -- actual extension msg follows
 108                                            mapM_ putWord8 blockList
 109                                              where blockList = unpack b
 110                                                    blockLen  = length blockList
 111
 112
 113     where putIndexOffsetLength i o l = do
 114             putWord32be (fromIntegral i)
 115             putWord32be (fromIntegral o)
 116             putWord32be (fromIntegral l)
 117
 118   get = do
 119     l <- getWord32be
 120     msgid <- getWord8
 121     case msgid of
 122      0 -> return ChokeMsg
 123      1 -> return UnChokeMsg
 124      2 -> return InterestedMsg
 125      3 -> return NotInterestedMsg
 126      4 -> fmap (HaveMsg . fromIntegral) getWord32be
 127      5 -> fmap (BitFieldMsg . pack) (replicateM (fromIntegral l - 1) getWord8)
 128      6 -> liftA3 RequestMsg getInteger getInteger getInteger
 129        where getInteger = fromIntegral <$> getWord32be
 130      7 -> liftA3 PieceMsg getInteger getInteger (pack  <$> replicateM (fromIntegral l - 9) getWord8)
 131        where getInteger = fromIntegral <$> getWord32be
 132      8 -> liftA3 CancelMsg getInteger getInteger getInteger
 133        where getInteger = fromIntegral <$> getWord32be
 134      9 -> fmap (PortMsg . fromIntegral) getWord16be
 135      _ -> error ("unknown message ID: " ++ show msgid)
 136
 137 getMsg :: Handle -> IO PeerMsg
 138 getMsg h = do
 139   lBS <- hGet h 4
 140   let l = bsToInt lBS
 141   if l == 0
 142     then return KeepAliveMsg
 143     else do
 144     msg <- hGet h l
 145     return $ decode $ fromStrict $ concat [lBS, msg]
 146
 147 sendMsg :: Handle -> PeerMsg -> IO ()
 148 sendMsg h msg = hPut h bsMsg
 149   where bsMsg = toStrict $ encode msg
 150
 151 genHandshakeMsg :: ByteString -> String -> ByteString
 152 genHandshakeMsg infoHash peer_id = concat [pstrlen, pstr, reserved1, reserved2, reserved3, infoHash, peerID]
 153   where pstrlen = singleton 19
 154         pstr = BC.pack "BitTorrent protocol"
 155         reserved1 = BC.replicate 4 '\0'
 156         reserved2 = singleton 0x10 -- support extension protocol
 157         reserved3 = BC.replicate 3 '\0'
 158         peerID = BC.pack peer_id
 159
 160 bsToInt :: ByteString -> Int
 161 bsToInt x = fromIntegral (runGet getWord32be (fromChunks (return x)))
 162
 163 makePeer :: ByteString -> Peer
 164 makePeer peer = Peer (toIP ip') (toPort port')
 165   where (ip', port') = splitAt 4 peer