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@(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   putStrLn $ "--> handhake to peer: " ++ show peer
  99   rlenBS <- hGet h (length (unpack hs))
 100   putStrLn $ "<-- handshake from peer: " ++ show peer
 101   return h
 102
 103 instance Binary PeerMsg where
 104   put msg = case msg of
 105              KeepAliveMsg -> putWord32be 0
 106              ChokeMsg -> do putWord32be 1
 107                             putWord8 0
 108              UnChokeMsg -> do putWord32be 1
 109                               putWord8 1
 110              InterestedMsg -> do putWord32be 1
 111                                  putWord8 2
 112              NotInterestedMsg -> do putWord32be 1
 113                                     putWord8 3
 114              HaveMsg i -> do putWord32be 5
 115                              putWord8 4
 116                              putWord32be (fromIntegral i)
 117              BitFieldMsg bf -> do putWord32be $ fromIntegral (1 + bfListLen)
 118                                   putWord8 5
 119                                   mapM_ putWord8 bfList
 120                                     where bfList = unpack bf
 121                                           bfListLen = length bfList
 122              RequestMsg i o l -> do putWord32be 13
 123                                     putWord8 6
 124                                     putWord32be (fromIntegral i)
 125                                     putWord32be (fromIntegral o)
 126                                     putWord32be (fromIntegral l)
 127              PieceMsg i o b -> do putWord32be $ fromIntegral (9 + blocklen)
 128                                   putWord8 7
 129                                   putWord32be (fromIntegral i)
 130                                   putWord32be (fromIntegral o)
 131                                   mapM_ putWord8 blockList
 132                                     where blockList = unpack b
 133                                           blocklen = length blockList
 134              CancelMsg i o l -> do putWord32be 13
 135                                    putWord8 8
 136                                    putWord32be (fromIntegral i)
 137                                    putWord32be (fromIntegral o)
 138                                    putWord32be (fromIntegral l)
 139              PortMsg p -> do putWord32be 3
 140                              putWord8 9
 141                              putWord16be (fromIntegral p)
 142   get = do
 143     l <- getWord32be
 144     msgid <- getWord8
 145     case msgid of
 146      0 -> return ChokeMsg
 147      1 -> return UnChokeMsg
 148      2 -> return InterestedMsg
 149      3 -> return NotInterestedMsg
 150      4 -> liftM (HaveMsg . fromIntegral) getWord32be
 151      5 -> liftM (BitFieldMsg . pack) (replicateM (fromIntegral l - 1) getWord8)
 152      6 -> liftA3 RequestMsg getInteger getInteger getInteger
 153        where getInteger = fromIntegral <$> getWord32be
 154      7 -> liftA3 PieceMsg getInteger getInteger (pack  <$> replicateM (fromIntegral l - 9) getWord8)
 155        where getInteger = fromIntegral <$> getWord32be
 156      8 -> liftA3 CancelMsg getInteger getInteger getInteger
 157        where getInteger = fromIntegral <$> getWord32be
 158      9 -> liftM (PortMsg . fromIntegral) getWord16be
 159      _ -> error ("unknown message ID: " ++ show msgid)
 160
 161 getMsg :: Handle -> IO PeerMsg
 162 getMsg h = do
 163   lBS <- hGet h 4
 164   let l = bsToInt lBS
 165   if l == 0
 166     then return KeepAliveMsg
 167     else do
 168     msg <- hGet h l
 169     return $ decode $ fromStrict $ concat [lBS, msg]
 170
 171 sendMsg :: Handle -> PeerMsg -> IO ()
 172 sendMsg h msg =
 173   let bsMsg = toStrict $ encode msg
 174   in
 175    hPut h bsMsg
 176
 177 bsToInt :: ByteString -> Int
 178 bsToInt x = fromIntegral (runGet getWord32be (fromChunks (return x)))
 179
 180 bitfieldToList :: [Word8] -> [Integer]
 181 bitfieldToList bs = go bs 0
 182   where go [] _ = []
 183         go (b:bs') pos =
 184           let setBits = [pos*8 + toInteger i | i <- [0..8], testBit b i]
 185           in
 186            setBits ++ go bs' (pos + 1)
 187
 188 -- downloadPiece :: Integer -> Handle -> IO ()
 189
 190 createDummyFile :: FilePath -> Int -> IO ()
 191 createDummyFile path size =
 192   writeFile path (BC.replicate size '\0')
 193
 194 -- loop1 :: shake hands with all peers, find out the pieces they have, form PieceData.
 195 -- recvMsg :: Peer -> Handle -> Msg
 196 msgLoop :: PeerState -> PieceMap -> IO ()
 197 msgLoop state pieceStatus = do
 198   -- if meInterested and he NOT Choking, pick a piece to download
 199   -- and send a requestmsg.
 200   let isMeInterested = meInterested state
 201       isHeChoking = heChoking state
 202   if (isMeInterested && isHeChoking)
 203     then
 204     do
 205       let h = handle state
 206       sendMsg h InterestedMsg
 207       putStrLn $ "--> InterestedMsg to peer: " ++ show (peer state)
 208       msgLoop state pieceStatus
 209     else
 210     do
 211       msg <- getMsg (handle state)
 212       putStrLn $ "<-- " ++ show msg ++ "from peer: " ++ show (peer state)
 213       case msg of
 214        KeepAliveMsg -> do
 215          sendMsg (handle state) KeepAliveMsg
 216          msgLoop state pieceStatus
 217        BitFieldMsg bss -> do
 218          let pieceList = bitfieldToList (unpack bss)
 219          print pieceList
 220          -- for each pieceIndex in pieceList, make an entry in the pieceStatus
 221          -- map with pieceIndex as the key and modify the value to add the peer.
 222          -- download each of the piece in order
 223
 224          msgLoop state pieceStatus
 225        UnChokeMsg -> do
 226          msgLoop (state {heChoking = False}) pieceStatus
 227        _ -> do
 228          msgLoop state pieceStatus
 229
 230 handlePeerMsgs :: Peer -> Metainfo -> String -> (String -> IO ()) -> IO ()
 231 handlePeerMsgs p m peerId logFn = do
 232   h <- handShake p (infoHash m) peerId
 233   -- logFn "handShake"
 234   let state = PeerState { handle = h
 235                         , peer = p
 236                         , heInterested = False
 237                         , heChoking = False
 238                         , meInterested = True
 239                         , meChoking = False }
 240       pieceHash = (pieces (info m))
 241       numPieces = (toInteger . (`quot` 20) . BC.length) pieceHash
 242       pieceStatus = mkPieceMap numPieces pieceHash
 243   msgLoop state pieceStatus
 244