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, liftM)
  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              deriving (Show)
  67
  68 instance Binary PeerMsg where
  69   put msg = case msg of
  70              KeepAliveMsg -> putWord32be 0
  71              ChokeMsg -> do putWord32be 1
  72                             putWord8 0
  73              UnChokeMsg -> do putWord32be 1
  74                               putWord8 1
  75              InterestedMsg -> do putWord32be 1
  76                                  putWord8 2
  77              NotInterestedMsg -> do putWord32be 1
  78                                     putWord8 3
  79              HaveMsg i -> do putWord32be 5
  80                              putWord8 4
  81                              putWord32be (fromIntegral i)
  82              BitFieldMsg bf -> do putWord32be $ fromIntegral (1 + bfListLen)
  83                                   putWord8 5
  84                                   mapM_ putWord8 bfList
  85                                     where bfList = unpack bf
  86                                           bfListLen = length bfList
  87              RequestMsg i o l -> do putWord32be 13
  88                                     putWord8 6
  89                                     putWord32be (fromIntegral i)
  90                                     putWord32be (fromIntegral o)
  91                                     putWord32be (fromIntegral l)
  92              PieceMsg i o b -> do putWord32be $ fromIntegral (9 + blocklen)
  93                                   putWord8 7
  94                                   putWord32be (fromIntegral i)
  95                                   putWord32be (fromIntegral o)
  96                                   mapM_ putWord8 blockList
  97                                     where blockList = unpack b
  98                                           blocklen = length blockList
  99              CancelMsg i o l -> do putWord32be 13
 100                                    putWord8 8
 101                                    putWord32be (fromIntegral i)
 102                                    putWord32be (fromIntegral o)
 103                                    putWord32be (fromIntegral l)
 104              PortMsg p -> do putWord32be 3
 105                              putWord8 9
 106                              putWord16be (fromIntegral p)
 107   get = do
 108     l <- getWord32be
 109     msgid <- getWord8
 110     case msgid of
 111      0 -> return ChokeMsg
 112      1 -> return UnChokeMsg
 113      2 -> return InterestedMsg
 114      3 -> return NotInterestedMsg
 115      4 -> liftM (HaveMsg . fromIntegral) getWord32be
 116      5 -> liftM (BitFieldMsg . pack) (replicateM (fromIntegral l - 1) getWord8)
 117      6 -> liftA3 RequestMsg getInteger getInteger getInteger
 118        where getInteger = fromIntegral <$> getWord32be
 119      7 -> liftA3 PieceMsg getInteger getInteger (pack  <$> replicateM (fromIntegral l - 9) getWord8)
 120        where getInteger = fromIntegral <$> getWord32be
 121      8 -> liftA3 CancelMsg getInteger getInteger getInteger
 122        where getInteger = fromIntegral <$> getWord32be
 123      9 -> liftM (PortMsg . fromIntegral) getWord16be
 124      _ -> error ("unknown message ID: " ++ show msgid)
 125
 126 getMsg :: Handle -> IO PeerMsg
 127 getMsg h = do
 128   lBS <- hGet h 4
 129   let l = bsToInt lBS
 130   if l == 0
 131     then return KeepAliveMsg
 132     else do
 133     msg <- hGet h l
 134     return $ decode $ fromStrict $ concat [lBS, msg]
 135
 136 sendMsg :: Handle -> PeerMsg -> IO ()
 137 sendMsg h msg = hPut h bsMsg
 138   where bsMsg = toStrict $ encode msg
 139
 140 genHandshakeMsg :: ByteString -> String -> ByteString
 141 genHandshakeMsg infoHash peer_id = concat [pstrlen, pstr, reserved, infoHash, peerID]
 142   where pstrlen = singleton 19
 143         pstr = BC.pack "BitTorrent protocol"
 144         reserved = BC.replicate 8 '\0'
 145         peerID = BC.pack peer_id
 146
 147 bsToInt :: ByteString -> Int
 148 bsToInt x = fromIntegral (runGet getWord32be (fromChunks (return x)))
 149
 150 makePeer :: ByteString -> Peer
 151 makePeer peer = Peer (toIP ip') (toPort port')
 152   where (ip', port') = splitAt 4 peer