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