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, toStrict)
  12 import qualified Data.ByteString.Char8 as BC (replicate, pack, length)
  13 import Network (connectTo, PortID(..))
  14 import Data.Binary (Binary(..), decode, encode)
  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 sendMsg :: Handle -> PeerMsg -> IO ()
 170 sendMsg h msg =
 171   let bsMsg = toStrict $ encode msg
 172   in
 173    hPut h bsMsg
 174
 175 bsToInt :: ByteString -> Int
 176 bsToInt x = fromIntegral (runGet getWord32be (fromChunks (return x)))
 177
 178 bitfieldToList :: [Word8] -> [Integer]
 179 bitfieldToList bs = go bs 0
 180   where go [] _ = []
 181         go (b:bs') pos =
 182           let setBits = [pos*8 + toInteger i | i <- [0..8], testBit b i]
 183           in
 184            setBits ++ go bs' (pos + 1)
 185
 186 -- downloadPiece :: Integer -> Handle -> IO ()
 187
 188 createDummyFile :: FilePath -> Int -> IO ()
 189 createDummyFile path size =
 190   writeFile path (BC.replicate size '\0')
 191
 192 -- loop1 :: shake hands with all peers, find out the pieces they have, form PieceData.
 193 -- recvMsg :: Peer -> Handle -> Msg
 194 msgLoop :: Handle -> ByteString -> PeerState -> IO ()
 195 msgLoop h pieceHash state =
 196   let numPieces = (toInteger . (`quot` 20) . BC.length) pieceHash
 197       pieceStatus = mkPieceMap numPieces pieceHash
 198   in
 199    forever $ do
 200      msg <- getMsg h
 201      putStrLn $ "got a " ++ show msg
 202      case msg of
 203       BitFieldMsg bss -> do
 204         let pieceList = bitfieldToList (unpack bss)
 205         print pieceList
 206         -- for each pieceIndex in pieceList, make an entry in the pieceStatus
 207         -- map with pieceIndex as the key and modify the value to add the peer.
 208
 209         -- download each of the piece in order
 210       UnChokeMsg -> do
 211         print msg
 212         msgLoop h pieceHash (state {heChoking = False})
 213       _ -> print msg
 214
 215 handlePeerMsgs :: Peer -> Metainfo -> String -> (String -> IO ()) -> IO ()
 216 handlePeerMsgs p m peerId logFn = do
 217   h <- handShake p (infoHash m) peerId
 218   logFn "handShake"
 219   let state = PeerState { peer = p
 220                         , heInterested = False
 221                         , heChoking = True
 222                         , meInterested = True
 223                         , meChoking = False }
 224   msgLoop h (pieces (info m)) state
 225