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, 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, toList, (!), mapWithKey)
  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                            , len :: Integer }       -- ^ piece length
  52
  53 -- which piece is with which peers
  54 type PieceMap = Map Integer PieceData
  55
  56 -- | Peer is a PeerID, IP address, port tuple
  57 data Peer = Peer ID IP Port
  58           deriving (Show, Eq)
  59
  60 data PeerMsg = KeepAliveMsg
  61              | ChokeMsg
  62              | UnChokeMsg
  63              | InterestedMsg
  64              | NotInterestedMsg
  65              | HaveMsg Integer
  66              | BitFieldMsg ByteString
  67              | RequestMsg Integer Integer Integer
  68              | PieceMsg Integer Integer ByteString
  69              | CancelMsg Integer Integer Integer
  70              | PortMsg Port
  71              deriving (Show)
  72
  73 -- Make the initial Piece map, with the assumption that no peer has the
  74 -- piece and that every piece is pending download.
  75 mkPieceMap :: Integer -> ByteString -> Integer -> PieceMap
  76 mkPieceMap numPieces pieceHash pLen = fromList kvs
  77   where kvs = [(i, PieceData { peers = []
  78                              , state = Pending
  79                              , hash = h
  80                              , len = pLen })
  81               | (i, h) <- zip [0..numPieces] hashes]
  82         hashes = splitN (fromIntegral numPieces) pieceHash
  83
  84 havePiece :: PieceMap -> Integer -> Bool
  85 havePiece pm index =
  86   state (pm ! index) == Have
  87
  88 genHandShakeMsg :: ByteString -> String -> ByteString
  89 genHandShakeMsg infoHash peer_id = concat [pstrlen, pstr, reserved, infoHash, peerID]
  90   where pstrlen = singleton 19
  91         pstr = BC.pack "BitTorrent protocol"
  92         reserved = BC.replicate 8 '\0'
  93         peerID = BC.pack peer_id
  94
  95 handShake :: Peer -> ByteString -> String -> IO Handle
  96 handShake peer@(Peer _ ip port) infoHash peerid = do
  97   let hs = genHandShakeMsg infoHash peerid
  98   h <- connectTo ip (PortNumber (fromIntegral port))
  99   hSetBuffering h LineBuffering
 100   hPut h hs
 101   putStrLn $ "--> handhake to peer: " ++ show peer
 102   _ <- hGet h (length (unpack hs))
 103   putStrLn $ "<-- handshake from peer: " ++ show peer
 104   return h
 105
 106 instance Binary PeerMsg where
 107   put msg = case msg of
 108              KeepAliveMsg -> putWord32be 0
 109              ChokeMsg -> do putWord32be 1
 110                             putWord8 0
 111              UnChokeMsg -> do putWord32be 1
 112                               putWord8 1
 113              InterestedMsg -> do putWord32be 1
 114                                  putWord8 2
 115              NotInterestedMsg -> do putWord32be 1
 116                                     putWord8 3
 117              HaveMsg i -> do putWord32be 5
 118                              putWord8 4
 119                              putWord32be (fromIntegral i)
 120              BitFieldMsg bf -> do putWord32be $ fromIntegral (1 + bfListLen)
 121                                   putWord8 5
 122                                   mapM_ putWord8 bfList
 123                                     where bfList = unpack bf
 124                                           bfListLen = length bfList
 125              RequestMsg i o l -> do putWord32be 13
 126                                     putWord8 6
 127                                     putWord32be (fromIntegral i)
 128                                     putWord32be (fromIntegral o)
 129                                     putWord32be (fromIntegral l)
 130              PieceMsg i o b -> do putWord32be $ fromIntegral (9 + blocklen)
 131                                   putWord8 7
 132                                   putWord32be (fromIntegral i)
 133                                   putWord32be (fromIntegral o)
 134                                   mapM_ putWord8 blockList
 135                                     where blockList = unpack b
 136                                           blocklen = length blockList
 137              CancelMsg i o l -> do putWord32be 13
 138                                    putWord8 8
 139                                    putWord32be (fromIntegral i)
 140                                    putWord32be (fromIntegral o)
 141                                    putWord32be (fromIntegral l)
 142              PortMsg p -> do putWord32be 3
 143                              putWord8 9
 144                              putWord16be (fromIntegral p)
 145   get = do
 146     l <- getWord32be
 147     msgid <- getWord8
 148     case msgid of
 149      0 -> return ChokeMsg
 150      1 -> return UnChokeMsg
 151      2 -> return InterestedMsg
 152      3 -> return NotInterestedMsg
 153      4 -> liftM (HaveMsg . fromIntegral) getWord32be
 154      5 -> liftM (BitFieldMsg . pack) (replicateM (fromIntegral l - 1) getWord8)
 155      6 -> liftA3 RequestMsg getInteger getInteger getInteger
 156        where getInteger = fromIntegral <$> getWord32be
 157      7 -> liftA3 PieceMsg getInteger getInteger (pack  <$> replicateM (fromIntegral l - 9) getWord8)
 158        where getInteger = fromIntegral <$> getWord32be
 159      8 -> liftA3 CancelMsg getInteger getInteger getInteger
 160        where getInteger = fromIntegral <$> getWord32be
 161      9 -> liftM (PortMsg . fromIntegral) getWord16be
 162      _ -> error ("unknown message ID: " ++ show msgid)
 163
 164 getMsg :: Handle -> IO PeerMsg
 165 getMsg h = do
 166   lBS <- hGet h 4
 167   let l = bsToInt lBS
 168   if l == 0
 169     then return KeepAliveMsg
 170     else do
 171     msg <- hGet h l
 172     return $ decode $ fromStrict $ concat [lBS, msg]
 173
 174 sendMsg :: Handle -> PeerMsg -> IO ()
 175 sendMsg h msg =
 176   let bsMsg = toStrict $ encode msg
 177   in
 178    hPut h bsMsg
 179
 180 bsToInt :: ByteString -> Int
 181 bsToInt x = fromIntegral (runGet getWord32be (fromChunks (return x)))
 182
 183 bitfieldToList :: [Word8] -> [Integer]
 184 bitfieldToList bs = go bs 0
 185   where go [] _ = []
 186         go (b:bs') pos =
 187           let setBits = [pos*8 + toInteger i | i <- [0..8], testBit b i]
 188           in
 189            setBits ++ go bs' (pos + 1)
 190
 191 -- downloadPiece :: Integer -> Handle -> IO ()
 192
 193 createDummyFile :: FilePath -> Int -> IO ()
 194 createDummyFile path size =
 195   writeFile path (BC.replicate size '\0')
 196
 197 -- loop1 :: shake hands with all peers, find out the pieces they have, form PieceData.
 198 -- recvMsg :: Peer -> Handle -> Msg
 199 msgLoop :: PeerState -> PieceMap -> IO ()
 200 msgLoop state pieceStatus | meInterested state == False &&
 201                             heChoking state == True = do
 202                               -- if meInterested and he NOT Choking, pick a piece to download
 203                               -- and send a requestmsg.
 204                               let h = handle state
 205                               sendMsg h InterestedMsg
 206                               putStrLn $ "--> InterestedMsg to peer: " ++ show (peer state)
 207                               msgLoop (state { meInterested = True }) pieceStatus
 208                           | meInterested state == True &&
 209                             heChoking state == False =
 210                               case pickPiece pieceStatus of
 211                                Nothing -> putStrLn "Nothing to download"
 212                                Just workPiece -> do
 213                                  let pLen = len (pieceStatus ! workPiece)
 214                                  sendMsg (handle state) (RequestMsg workPiece 0 pLen)
 215                                  putStrLn $ "--> RequestMsg for Piece " ++ (show workPiece) ++ "to peer: " ++ show (peer state)
 216                                  msg <- getMsg (handle state)
 217                                  putStrLn $ "<-- " ++ show msg ++ "from peer: " ++ show (peer state)
 218                                  -- msgLoop state pieceStatus
 219                           | otherwise = do
 220                               msg <- getMsg (handle state)
 221                               putStrLn $ "<-- " ++ show msg ++ "from peer: " ++ show (peer state)
 222                               case msg of
 223                                KeepAliveMsg -> do
 224                                  sendMsg (handle state) KeepAliveMsg
 225                                  putStrLn $ "--> " ++ "KeepAliveMsg to peer: " ++ show (peer state)
 226                                  msgLoop state pieceStatus
 227                                BitFieldMsg bss -> do
 228                                  let pieceList = bitfieldToList (unpack bss)
 229                                      pieceStatus' = updatePieceAvailability pieceStatus (peer state) pieceList
 230                                  print pieceList
 231                                  -- for each pieceIndex in pieceList, make an entry in the pieceStatus
 232                                  -- map with pieceIndex as the key and modify the value to add the peer.
 233                                  -- download each of the piece in order
 234                                  msgLoop state pieceStatus'
 235                                UnChokeMsg -> do
 236                                  msgLoop (state { heChoking = False }) pieceStatus
 237                                _ -> do
 238                                  msgLoop state pieceStatus
 239
 240 -- simple algorithm to pick piece.
 241 -- pick the first piece from 0 that is not downloaded yet.
 242 pickPiece :: PieceMap -> Maybe Integer
 243 pickPiece m =
 244   let pieceList = toList m
 245       allPending = filter (\(k, v) -> state v == Pending) pieceList
 246   in
 247    case allPending of
 248     [] -> Nothing
 249     ((i, _):_) -> Just i
 250
 251 updatePieceAvailability :: PieceMap -> Peer -> [Integer] -> PieceMap
 252 updatePieceAvailability pieceStatus p pieceList =
 253   mapWithKey (\k pd -> if k `elem` pieceList
 254                        then (pd { peers = p : (peers pd) })
 255                        else pd) pieceStatus
 256
 257 handlePeerMsgs :: Peer -> Metainfo -> String -> (String -> IO ()) -> IO ()
 258 handlePeerMsgs p m peerId logFn = do
 259   h <- handShake p (infoHash m) peerId
 260   -- logFn "handShake"
 261   let state = PeerState { handle = h
 262                         , peer = p
 263                         , heInterested = False
 264                         , heChoking = True
 265                         , meInterested = False
 266                         , meChoking = True }
 267       pieceHash = pieces (info m)
 268       numPieces = (toInteger . (`quot` 20) . BC.length) pieceHash
 269       pLen = pieceLength (info m) :: Integer
 270       pieceStatus = mkPieceMap numPieces pieceHash pLen
 271   msgLoop state pieceStatus
 272