src/FuncTorrent/Peer.hs

   1 {-# LANGUAGE OverloadedStrings #-}
   2 module FuncTorrent.Peer
   3     (Peer(..),
   4      handlePeerMsgs
   5     ) where
   6
   7 import Prelude hiding (lookup, concat, replicate, splitAt, empty, writeFile)
   8
   9 import System.IO (Handle, BufferMode(..), hSetBuffering)
  10 import Data.ByteString (ByteString, pack, unpack, concat, hGet, hPut, singleton, writeFile)
  11 import Data.ByteString.Lazy (fromStrict, fromChunks)
  12 import qualified Data.ByteString.Char8 as BC (replicate, pack, length)
  13 import Network (connectTo, PortID(..))
  14 import Data.Binary (Binary(..), decode)
  15 import Data.Binary.Put (putWord32be, putWord16be, putWord8)
  16 import Data.Binary.Get (getWord32be, getWord16be, getWord8, runGet)
  17 import Control.Monad (replicateM, liftM, forever)
  18 import Control.Applicative ((<$>), liftA3)
  19 import Data.Bits
  20 import Data.Word (Word8)
  21 import Data.Map (Map(..), fromList)
  22
  23 import FuncTorrent.Metainfo (Info(..), Metainfo(..))
  24 import FuncTorrent.Utils (splitN)
  25
  26 type ID = String
  27 type IP = String
  28 type Port = Integer
  29
  30 data PeerState = PeerState { handle :: Handle
  31                            , am_choking :: Bool
  32                            , am_interested :: Bool
  33                            , peer_choking :: Bool
  34                            , peer_interested :: Bool}
  35
  36 -- Maintain info on every piece and the current state of it.
  37 -- should probably be a TVar.
  38 type Pieces = [PieceData]
  39
  40 data PieceDlState = Pending
  41                   | InProgress
  42                   | Have
  43                   deriving (Show)
  44
  45 -- todo - map with index to a new data structure (peers who have that piece amd state)
  46 data PieceData = PieceData { peers :: [Peer]        -- ^ list of peers who have this piece
  47                            , state :: PieceDlState  -- ^ state of the piece from download perspective.
  48                            , hash  :: ByteString }      -- ^ piece hash
  49
  50 -- which piece is with which peers
  51 type PieceMap = Map Integer PieceData
  52
  53 -- | Peer is a PeerID, IP address, port tuple
  54 data Peer = Peer ID IP Port
  55           deriving (Show, Eq)
  56
  57 data PeerMsg = KeepAliveMsg
  58              | ChokeMsg
  59              | UnChokeMsg
  60              | InterestedMsg
  61              | NotInterestedMsg
  62              | HaveMsg Integer
  63              | BitFieldMsg ByteString
  64              | RequestMsg Integer Integer Integer
  65              | PieceMsg Integer Integer ByteString
  66              | CancelMsg Integer Integer Integer
  67              | PortMsg Port
  68              deriving (Show)
  69
  70 -- Make the initial Piece map, with the assumption that no peer has the
  71 -- piece and that every piece is pending download.
  72 mkPieceMap :: Integer -> ByteString -> Map Integer PieceData
  73 mkPieceMap numPieces pieceHash = fromList kvs
  74   where kvs = [(i, PieceData { peers = []
  75                              , state = Pending
  76                              , hash = h }) | (i, h) <- zip [0..numPieces] hashes]
  77         hashes = splitN (fromIntegral numPieces) pieceHash
  78
  79 genHandShakeMsg :: ByteString -> String -> ByteString
  80 genHandShakeMsg infoHash peer_id = concat [pstrlen, pstr, reserved, infoHash, peerID]
  81   where pstrlen = singleton 19
  82         pstr = BC.pack "BitTorrent protocol"
  83         reserved = BC.replicate 8 '\0'
  84         peerID = BC.pack peer_id
  85
  86 handShake :: Peer -> ByteString -> String -> IO Handle
  87 handShake (Peer _ ip port) infoHash peerid = do
  88   let hs = genHandShakeMsg infoHash peerid
  89   h <- connectTo ip (PortNumber (fromIntegral port))
  90   hSetBuffering h LineBuffering
  91   hPut h hs
  92   rlenBS <- hGet h (length (unpack hs))
  93   putStrLn $ "got handshake from peer: " ++ show rlenBS
  94   return h
  95
  96 instance Binary PeerMsg where
  97   put msg = case msg of
  98              KeepAliveMsg -> putWord32be 0
  99              ChokeMsg -> do putWord32be 1
 100                             putWord8 0
 101              UnChokeMsg -> do putWord32be 1
 102                               putWord8 1
 103              InterestedMsg -> do putWord32be 1
 104                                  putWord8 2
 105              NotInterestedMsg -> do putWord32be 1
 106                                     putWord8 3
 107              HaveMsg i -> do putWord32be 5
 108                              putWord8 4
 109                              putWord32be (fromIntegral i)
 110              BitFieldMsg bf -> do putWord32be $ fromIntegral (1 + bfListLen)
 111                                   putWord8 5
 112                                   mapM_ putWord8 bfList
 113                                     where bfList = unpack bf
 114                                           bfListLen = length bfList
 115              RequestMsg i o l -> do putWord32be 13
 116                                     putWord8 6
 117                                     putWord32be (fromIntegral i)
 118                                     putWord32be (fromIntegral o)
 119                                     putWord32be (fromIntegral l)
 120              PieceMsg i o b -> do putWord32be $ fromIntegral (9 + blocklen)
 121                                   putWord8 7
 122                                   putWord32be (fromIntegral i)
 123                                   putWord32be (fromIntegral o)
 124                                   mapM_ putWord8 blockList
 125                                     where blockList = unpack b
 126                                           blocklen = length blockList
 127              CancelMsg i o l -> do putWord32be 13
 128                                    putWord8 8
 129                                    putWord32be (fromIntegral i)
 130                                    putWord32be (fromIntegral o)
 131                                    putWord32be (fromIntegral l)
 132              PortMsg p -> do putWord32be 3
 133                              putWord8 9
 134                              putWord16be (fromIntegral p)
 135   get = do
 136     l <- getWord32be
 137     msgid <- getWord8
 138     case msgid of
 139      0 -> return ChokeMsg
 140      1 -> return UnChokeMsg
 141      2 -> return InterestedMsg
 142      3 -> return NotInterestedMsg
 143      4 -> liftM (HaveMsg . fromIntegral) getWord32be
 144      5 -> liftM (BitFieldMsg . pack) (replicateM (fromIntegral l - 1) getWord8)
 145      6 -> liftA3 RequestMsg getInteger getInteger getInteger
 146        where getInteger = fromIntegral <$> getWord32be
 147      7 -> liftA3 PieceMsg getInteger getInteger (pack  <$> replicateM (fromIntegral l - 9) getWord8)
 148        where getInteger = fromIntegral <$> getWord32be
 149      8 -> liftA3 CancelMsg getInteger getInteger getInteger
 150        where getInteger = fromIntegral <$> getWord32be
 151      9 -> liftM (PortMsg . fromIntegral) getWord16be
 152      _ -> error ("unknown message ID: " ++ show msgid)
 153
 154 getMsg :: Handle -> IO PeerMsg
 155 getMsg h = do
 156   lBS <- hGet h 4
 157   let l = bsToInt lBS
 158   if l == 0
 159     then return KeepAliveMsg
 160     else do
 161     msg <- hGet h l
 162     return $ decode $ fromStrict $ concat [lBS, msg]
 163
 164
 165 bsToInt :: ByteString -> Int
 166 bsToInt x = fromIntegral (runGet getWord32be (fromChunks (return x)))
 167
 168 bitfieldToList :: [Word8] -> [Integer]
 169 bitfieldToList bs = go bs 0
 170   where go [] _ = []
 171         go (b:bs') pos =
 172           let setBits = [pos*8 + (toInteger i) | i <- [0..8], testBit b i]
 173           in
 174            setBits ++ (go bs' (pos + 1))
 175
 176 -- downloadPiece :: Integer -> Handle -> IO ()
 177
 178 createDummyFile :: FilePath -> Int -> IO ()
 179 createDummyFile path size = do
 180   writeFile path (BC.replicate size '\0')
 181
 182 -- loop1 :: shake hands with all peers, find out the pieces they have, form PieceData.
 183 -- recvMsg :: Peer -> Handle -> Msg
 184 msgLoop :: Handle -> ByteString -> IO ()
 185 msgLoop h pieceHash =
 186   let numPieces = (toInteger . (`quot` 20) . BC.length) pieceHash
 187       pieceStatus = mkPieceMap numPieces pieceHash
 188   in
 189    forever $ do
 190      msg <- getMsg h
 191      putStrLn $ "got a " ++ show msg
 192      case msg of
 193       BitFieldMsg bss -> do
 194         let pieceList = bitfieldToList (unpack bss)
 195         putStrLn (show pieceList)
 196         -- download each of the piece in order
 197       _ -> putStrLn (show msg)
 198
 199 handlePeerMsgs :: Peer -> Metainfo -> String -> (String -> IO ()) -> IO ()
 200 handlePeerMsgs p m peerId logFn = do
 201   h <- handShake p (infoHash m) peerId
 202   logFn $ "handShake"
 203   msgLoop h (pieces (info m))
 204