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