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 me NOT Interested and she is Choking, tell her that
 203                               -- I am interested.
 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                               -- if me Interested and she not Choking, send her a request
 211                               -- for a piece.
 212                               case pickPiece pieceStatus of
 213                                Nothing -> putStrLn "Nothing to download"
 214                                Just workPiece -> do
 215                                  let pLen = len (pieceStatus ! workPiece)
 216                                  sendMsg (handle state) (RequestMsg workPiece 0 pLen)
 217                                  putStrLn $ "--> RequestMsg for Piece " ++ (show workPiece) ++ "to peer: " ++ show (peer state)
 218                                  msg <- getMsg (handle state)
 219                                  putStrLn $ "<-- " ++ show msg ++ "from peer: " ++ show (peer state)
 220                                  -- msgLoop state pieceStatus
 221                           | otherwise = do
 222                               msg <- getMsg (handle state)
 223                               putStrLn $ "<-- " ++ show msg ++ "from peer: " ++ show (peer state)
 224                               case msg of
 225                                KeepAliveMsg -> do
 226                                  sendMsg (handle state) KeepAliveMsg
 227                                  putStrLn $ "--> " ++ "KeepAliveMsg to peer: " ++ show (peer state)
 228                                  msgLoop state pieceStatus
 229                                BitFieldMsg bss -> do
 230                                  let pieceList = bitfieldToList (unpack bss)
 231                                      pieceStatus' = updatePieceAvailability pieceStatus (peer state) pieceList
 232                                  print pieceList
 233                                  -- for each pieceIndex in pieceList, make an entry in the pieceStatus
 234                                  -- map with pieceIndex as the key and modify the value to add the peer.
 235                                  -- download each of the piece in order
 236                                  msgLoop state pieceStatus'
 237                                UnChokeMsg -> do
 238                                  msgLoop (state { heChoking = False }) pieceStatus
 239                                _ -> do
 240                                  msgLoop state pieceStatus
 241
 242 -- simple algorithm to pick piece.
 243 -- pick the first piece from 0 that is not downloaded yet.
 244 pickPiece :: PieceMap -> Maybe Integer
 245 pickPiece m =
 246   let pieceList = toList m
 247       allPending = filter (\(k, v) -> state v == Pending) pieceList
 248   in
 249    case allPending of
 250     [] -> Nothing
 251     ((i, _):_) -> Just i
 252
 253 updatePieceAvailability :: PieceMap -> Peer -> [Integer] -> PieceMap
 254 updatePieceAvailability pieceStatus p pieceList =
 255   mapWithKey (\k pd -> if k `elem` pieceList
 256                        then (pd { peers = p : (peers pd) })
 257                        else pd) pieceStatus
 258
 259 handlePeerMsgs :: Peer -> Metainfo -> String -> (String -> IO ()) -> IO ()
 260 handlePeerMsgs p m peerId logFn = do
 261   h <- handShake p (infoHash m) peerId
 262   -- logFn "handShake"
 263   let state = PeerState { handle = h
 264                         , peer = p
 265                         , heInterested = False
 266                         , heChoking = True
 267                         , meInterested = False
 268                         , meChoking = True }
 269       pieceHash = pieces (info m)
 270       numPieces = (toInteger . (`quot` 20) . BC.length) pieceHash
 271       pLen = pieceLength (info m) :: Integer
 272       pieceStatus = mkPieceMap numPieces pieceHash pLen
 273   msgLoop state pieceStatus
 274