src/FuncTorrent/PeerMsgs.hs

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