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