src/FuncTorrent/PeerMsgs.hs

   1 {-# LANGUAGE OverloadedStrings #-}
   2 {-
   3 Copyright (C) 2015-2016 Ramakrishnan Muthukrishnan <ram@rkrishnan.org>
   4
   5 This file is part of FuncTorrent.
   6
   7 FuncTorrent is free software; you can redistribute it and/or modify
   8 it under the terms of the GNU General Public License as published by
   9 the Free Software Foundation; either version 3 of the License, or
  10 (at your option) any later version.
  11
  12 FuncTorrent is distributed in the hope that it will be useful,
  13 but WITHOUT ANY WARRANTY; without even the implied warranty of
  14 MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  15 GNU General Public License for more details.
  16
  17 You should have received a copy of the GNU General Public License
  18 along with FuncTorrent; if not,  see <http://www.gnu.org/licenses/>
  19 -}
  20
  21 module FuncTorrent.PeerMsgs
  22        (genHandshakeMsg,
  23         sendMsg,
  24         getMsg,
  25         Peer(..),
  26         PeerMsg(..)
  27        ) where
  28
  29 import Prelude hiding (lookup, concat, replicate, splitAt, take)
  30
  31 import System.IO (Handle)
  32 import Data.ByteString (ByteString, pack, unpack, concat, hGet, hPut, singleton)
  33 import Data.ByteString.Lazy (fromStrict, fromChunks, toStrict)
  34 import qualified Data.ByteString.Char8 as BC (replicate, pack)
  35 import Control.Monad (replicateM, liftM)
  36 import Control.Applicative (liftA3)
  37
  38 import Data.Binary (Binary(..), decode, encode)
  39 import Data.Binary.Put (putWord32be, putWord16be, putWord8)
  40 import Data.Binary.Get (getWord32be, getWord16be, getWord8, runGet)
  41
  42 -- | Peer is a PeerID, IP address, port tuple
  43 data Peer = Peer ID IP Port
  44           deriving (Show, Eq)
  45
  46 type ID = String
  47 type IP = String
  48 type Port = Integer
  49
  50 data PeerMsg = KeepAliveMsg
  51              | ChokeMsg
  52              | UnChokeMsg
  53              | InterestedMsg
  54              | NotInterestedMsg
  55              | HaveMsg Integer
  56              | BitFieldMsg ByteString
  57              | RequestMsg Integer Integer Integer
  58              | PieceMsg Integer Integer ByteString
  59              | CancelMsg Integer Integer Integer
  60              | PortMsg Port
  61              deriving (Show)
  62
  63 instance Binary PeerMsg where
  64   put msg = case msg of
  65              KeepAliveMsg -> putWord32be 0
  66              ChokeMsg -> do putWord32be 1
  67                             putWord8 0
  68              UnChokeMsg -> do putWord32be 1
  69                               putWord8 1
  70              InterestedMsg -> do putWord32be 1
  71                                  putWord8 2
  72              NotInterestedMsg -> do putWord32be 1
  73                                     putWord8 3
  74              HaveMsg i -> do putWord32be 5
  75                              putWord8 4
  76                              putWord32be (fromIntegral i)
  77              BitFieldMsg bf -> do putWord32be $ fromIntegral (1 + bfListLen)
  78                                   putWord8 5
  79                                   mapM_ putWord8 bfList
  80                                     where bfList = unpack bf
  81                                           bfListLen = length bfList
  82              RequestMsg i o l -> do putWord32be 13
  83                                     putWord8 6
  84                                     putWord32be (fromIntegral i)
  85                                     putWord32be (fromIntegral o)
  86                                     putWord32be (fromIntegral l)
  87              PieceMsg i o b -> do putWord32be $ fromIntegral (9 + blocklen)
  88                                   putWord8 7
  89                                   putWord32be (fromIntegral i)
  90                                   putWord32be (fromIntegral o)
  91                                   mapM_ putWord8 blockList
  92                                     where blockList = unpack b
  93                                           blocklen = length blockList
  94              CancelMsg i o l -> do putWord32be 13
  95                                    putWord8 8
  96                                    putWord32be (fromIntegral i)
  97                                    putWord32be (fromIntegral o)
  98                                    putWord32be (fromIntegral l)
  99              PortMsg p -> do putWord32be 3
 100                              putWord8 9
 101                              putWord16be (fromIntegral p)
 102   get = do
 103     l <- getWord32be
 104     msgid <- getWord8
 105     case msgid of
 106      0 -> return ChokeMsg
 107      1 -> return UnChokeMsg
 108      2 -> return InterestedMsg
 109      3 -> return NotInterestedMsg
 110      4 -> liftM (HaveMsg . fromIntegral) getWord32be
 111      5 -> liftM (BitFieldMsg . pack) (replicateM (fromIntegral l - 1) getWord8)
 112      6 -> liftA3 RequestMsg getInteger getInteger getInteger
 113        where getInteger = fromIntegral <$> getWord32be
 114      7 -> liftA3 PieceMsg getInteger getInteger (pack  <$> replicateM (fromIntegral l - 9) getWord8)
 115        where getInteger = fromIntegral <$> getWord32be
 116      8 -> liftA3 CancelMsg getInteger getInteger getInteger
 117        where getInteger = fromIntegral <$> getWord32be
 118      9 -> liftM (PortMsg . fromIntegral) getWord16be
 119      _ -> error ("unknown message ID: " ++ show msgid)
 120
 121 getMsg :: Handle -> IO PeerMsg
 122 getMsg h = do
 123   lBS <- hGet h 4
 124   let l = bsToInt lBS
 125   if l == 0
 126     then return KeepAliveMsg
 127     else do
 128     msg <- hGet h l
 129     return $ decode $ fromStrict $ concat [lBS, msg]
 130
 131 sendMsg :: Handle -> PeerMsg -> IO ()
 132 sendMsg h msg = hPut h bsMsg
 133   where bsMsg = toStrict $ encode msg
 134
 135 genHandshakeMsg :: ByteString -> String -> ByteString
 136 genHandshakeMsg infoHash peer_id = concat [pstrlen, pstr, reserved, infoHash, peerID]
 137   where pstrlen = singleton 19
 138         pstr = BC.pack "BitTorrent protocol"
 139         reserved = BC.replicate 8 '\0'
 140         peerID = BC.pack peer_id
 141
 142 bsToInt :: ByteString -> Int
 143 bsToInt x = fromIntegral (runGet getWord32be (fromChunks (return x)))