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, forM, forever)
  18 import Control.Applicative ((<$>), liftA3)
  19 import Data.Bits
  20 import Data.Word (Word8)
  21 import Data.Map (Map(..), fromList, toList, (!), mapWithKey, adjust)
  22
  23 import FuncTorrent.Metainfo (Info(..), Metainfo(..))
  24 import FuncTorrent.Utils (splitN, splitNum)
  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 pLengths = fromList kvs
  77   where kvs = [(i, PieceData { peers = []
  78                              , state = Pending
  79                              , hash = h
  80                              , len = pLen })
  81               | (i, h, pLen) <- zip3 [0..numPieces] hashes pLengths]
  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 createDummyFile :: FilePath -> Int -> IO ()
 192 createDummyFile path size =
 193   writeFile path (BC.replicate size '\0')
 194
 195 -- loop1 :: shake hands with all peers, find out the pieces they have, form PieceData.
 196 -- recvMsg :: Peer -> Handle -> Msg
 197 msgLoop :: PeerState -> PieceMap -> IO ()
 198 msgLoop state pieceStatus | meInterested state == False &&
 199                             heChoking state == True = do
 200                               -- if me NOT Interested and she is Choking, tell her that
 201                               -- I am interested.
 202                               let h = handle state
 203                               sendMsg h InterestedMsg
 204                               putStrLn $ "--> InterestedMsg to peer: " ++ show (peer state)
 205                               msgLoop (state { meInterested = True }) pieceStatus
 206                           | meInterested state == True &&
 207                             heChoking state == False =
 208                               -- if me Interested and she not Choking, send her a request
 209                               -- for a piece.
 210                               case pickPiece pieceStatus of
 211                                Nothing -> putStrLn "Nothing to download"
 212                                Just workPiece -> do
 213                                  let pLen = len (pieceStatus ! workPiece)
 214                                  pBS <- downloadPiece (handle state) workPiece pLen
 215                                  -- sendMsg (handle state) (RequestMsg workPiece 0 pLen)
 216                                  -- putStrLn $ "--> RequestMsg for Piece " ++ (show workPiece) ++ "to peer: " ++ show (peer state) ++ " of length: " ++ show pLen
 217                                  -- msg <- getMsg (handle state)
 218                                  -- putStrLn $ "<-- " ++ show msg ++ "from peer: " ++ show (peer state)
 219                                  msgLoop state (adjust (\pieceData -> pieceData { state = Have }) workPiece pieceStatus)
 220                           | otherwise = do
 221                               msg <- getMsg (handle state)
 222                               putStrLn $ "<-- " ++ show msg ++ "from peer: " ++ show (peer state)
 223                               case msg of
 224                                KeepAliveMsg -> do
 225                                  sendMsg (handle state) KeepAliveMsg
 226                                  putStrLn $ "--> " ++ "KeepAliveMsg to peer: " ++ show (peer state)
 227                                  msgLoop state pieceStatus
 228                                BitFieldMsg bss -> do
 229                                  let pieceList = bitfieldToList (unpack bss)
 230                                      pieceStatus' = updatePieceAvailability pieceStatus (peer state) pieceList
 231                                  print pieceList
 232                                  -- for each pieceIndex in pieceList, make an entry in the pieceStatus
 233                                  -- map with pieceIndex as the key and modify the value to add the peer.
 234                                  -- download each of the piece in order
 235                                  msgLoop state pieceStatus'
 236                                UnChokeMsg -> do
 237                                  msgLoop (state { heChoking = False }) pieceStatus
 238                                _ -> do
 239                                  msgLoop state pieceStatus
 240
 241 -- simple algorithm to pick piece.
 242 -- pick the first piece from 0 that is not downloaded yet.
 243 pickPiece :: PieceMap -> Maybe Integer
 244 pickPiece m =
 245   let pieceList = toList m
 246       allPending = filter (\(k, v) -> state v == Pending) pieceList
 247   in
 248    case allPending of
 249     [] -> Nothing
 250     ((i, _):_) -> Just i
 251
 252 updatePieceAvailability :: PieceMap -> Peer -> [Integer] -> PieceMap
 253 updatePieceAvailability pieceStatus p pieceList =
 254   mapWithKey (\k pd -> if k `elem` pieceList
 255                        then (pd { peers = p : (peers pd) })
 256                        else pd) pieceStatus
 257
 258 handlePeerMsgs :: Peer -> Metainfo -> String -> (String -> IO ()) -> IO ()
 259 handlePeerMsgs p m peerId logFn = do
 260   h <- handShake p (infoHash m) peerId
 261   -- logFn "handShake"
 262   let state = PeerState { handle = h
 263                         , peer = p
 264                         , heInterested = False
 265                         , heChoking = True
 266                         , meInterested = False
 267                         , meChoking = True }
 268       pieceHash = pieces (info m)
 269       numPieces = (toInteger . (`quot` 20) . BC.length) pieceHash
 270       pLen = pieceLength (info m)
 271       fileLen = lengthInBytes (info m)
 272       pieceStatus = mkPieceMap numPieces pieceHash (splitNum fileLen pLen)
 273   msgLoop state pieceStatus
 274
 275 downloadPiece :: Handle -> Integer -> Integer -> IO [PeerMsg]
 276 downloadPiece h index pieceLength = do
 277   let chunks = splitNum pieceLength 16384
 278   forM (zip [0..] chunks) (\(i, pLen) -> do
 279                               sendMsg h (RequestMsg index (i*pLen) pLen)
 280                               putStrLn $ "--> " ++ "RequestMsg for Piece " ++ (show index) ++ ", part: " ++ show i ++ " of length: " ++ show pLen
 281                               -- putStrLn $ "--> RequestMsg for Piece " ++ (show workPiece) ++ "to peer: " ++ show (peer state) ++ " of length: " ++ show pLen
 282                               getMsg h)