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