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