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