src/FuncTorrent/Peer.hs

   1 {-# LANGUAGE OverloadedStrings #-}
   2 module FuncTorrent.Peer
   3     (Peer(..),
   4      handlePeerMsgs,
   5      bytesDownloaded
   6     ) where
   7
   8 import Prelude hiding (lookup, concat, replicate, splitAt, take, filter)
   9
  10 import System.IO (Handle, BufferMode(..), hSetBuffering)
  11 import Data.ByteString (ByteString, unpack, concat, hGet, hPut, take, empty)
  12 import qualified Data.ByteString.Char8 as BC (length)
  13 import Network (connectTo, PortID(..))
  14 import Control.Monad.State
  15 import Data.Bits
  16 import Data.Word (Word8)
  17 import Data.Map (Map, fromList, toList, (!), mapWithKey, adjust, filter)
  18 import qualified Crypto.Hash.SHA1 as SHA1 (hash)
  19 import Safe (headMay)
  20
  21 import FuncTorrent.Metainfo (Info(..), Metainfo(..))
  22 import FuncTorrent.Utils (splitN, splitNum)
  23 import FuncTorrent.Fileops (createDummyFile, writeFileAtOffset)
  24 import FuncTorrent.PeerMsgs (Peer(..), PeerMsg(..), sendMsg, getMsg, genHandshakeMsg)
  25
  26 data PState = PState { handle :: Handle
  27                      , peer :: Peer
  28                      , meChoking :: Bool
  29                      , meInterested :: Bool
  30                      , heChoking :: Bool
  31                      , heInterested :: Bool}
  32
  33 type PeerState = State PState
  34
  35 data PieceDlState = Pending
  36                   | Downloading
  37                   | Have
  38                   deriving (Show, Eq)
  39
  40 -- todo - map with index to a new data structure (peers who have that piece and state)
  41 data PieceData = PieceData { peers :: [Peer]        -- ^ list of peers who have this piece
  42                            , dlstate :: PieceDlState  -- ^ state of the piece from download perspective.
  43                            , hash  :: ByteString    -- ^ piece hash
  44                            , len :: Integer }       -- ^ piece length
  45
  46 -- which piece is with which peers
  47 type PieceMap = Map Integer PieceData
  48
  49
  50 -- Make the initial Piece map, with the assumption that no peer has the
  51 -- piece and that every piece is pending download.
  52 mkPieceMap :: Integer -> ByteString -> [Integer] -> PieceMap
  53 mkPieceMap numPieces pieceHash pLengths = fromList kvs
  54   where kvs = [(i, PieceData { peers = []
  55                              , dlstate = Pending
  56                              , hash = h
  57                              , len = pLen })
  58               | (i, h, pLen) <- zip3 [0..numPieces] hashes pLengths]
  59         hashes = splitN 20 pieceHash
  60
  61 havePiece :: PieceMap -> Integer -> Bool
  62 havePiece pm index =
  63   dlstate (pm ! index) == Have
  64
  65 connectToPeer :: Peer -> IO Handle
  66 connectToPeer (Peer _ ip port) = do
  67   h <- connectTo ip (PortNumber (fromIntegral port))
  68   hSetBuffering h LineBuffering
  69   return h
  70
  71 doHandshake :: Handle -> Peer -> ByteString -> String -> IO ()
  72 doHandshake h peer infoHash peerid = do
  73   let hs = genHandshakeMsg infoHash peerid
  74   hPut h hs
  75   putStrLn $ "--> handhake to peer: " ++ show peer
  76   _ <- hGet h (length (unpack hs))
  77   putStrLn $ "<-- handshake from peer: " ++ show peer
  78   return ()
  79
  80 bitfieldToList :: [Word8] -> [Integer]
  81 bitfieldToList bs = go bs 0
  82   where go [] _ = []
  83         go (b:bs') pos =
  84           let setBits = [pos*8 + toInteger i | i <- [0..8], testBit b i]
  85           in
  86            setBits ++ go bs' (pos + 1)
  87
  88 -- helper functions to manipulate PeerState
  89 toPeerState :: Handle
  90             -> Peer
  91             -> Bool  -- ^ meChoking
  92             -> Bool  -- ^ meInterested
  93             -> Bool  -- ^ heChoking
  94             -> Bool  -- ^ heInterested
  95             -> PState
  96 toPeerState h p meCh meIn heCh heIn =
  97   PState { handle = h
  98          , peer = p
  99          , heChoking = heCh
 100          , heInterested = heIn
 101          , meChoking = meCh
 102          , meInterested = meIn }
 103
 104 -- simple algorithm to pick piece.
 105 -- pick the first piece from 0 that is not downloaded yet.
 106 pickPiece :: PieceMap -> Maybe Integer
 107 pickPiece =
 108   (fst `liftM`) . headMay . toList . filter (\v -> dlstate v == Pending)
 109
 110 bytesDownloaded :: PieceMap -> Integer
 111 bytesDownloaded =
 112   sum . (map (len . snd)) . toList . filter (\v -> dlstate v == Have)
 113
 114 updatePieceAvailability :: PieceMap -> Peer -> [Integer] -> PieceMap
 115 updatePieceAvailability pieceStatus p pieceList =
 116   mapWithKey (\k pd -> if k `elem` pieceList
 117                        then (pd { peers = p : peers pd })
 118                        else pd) pieceStatus
 119
 120 handlePeerMsgs :: Peer -> Metainfo -> String -> IO ()
 121 handlePeerMsgs p m peerId = do
 122   h <- connectToPeer p
 123   doHandshake h p (infoHash m) peerId
 124   let pstate = toPeerState h p False False True True
 125       pieceHash = pieces (info m)
 126       numPieces = (toInteger . (`quot` 20) . BC.length) pieceHash
 127       pLen = pieceLength (info m)
 128       fileLen = lengthInBytes (info m)
 129       fileName = name (info m)
 130       pieceStatus = mkPieceMap numPieces pieceHash (splitNum fileLen pLen)
 131   createDummyFile fileName (fromIntegral fileLen)
 132   _ <- runStateT (msgLoop pieceStatus fileName) pstate
 133   return ()
 134
 135 msgLoop :: PieceMap -> FilePath -> StateT PState IO ()
 136 msgLoop pieceStatus file = do
 137   h <- gets handle
 138   st <- get
 139   case st of
 140     PState { meInterested = False, heChoking = True } -> do
 141       liftIO $ sendMsg h InterestedMsg
 142       gets peer >>= (\p -> liftIO $ putStrLn $ "--> InterestedMsg to peer: " ++ show p)
 143       modify (\st -> st { meInterested = True })
 144       msgLoop pieceStatus file
 145     PState { meInterested = True, heChoking = False } ->
 146       case pickPiece pieceStatus of
 147         Nothing -> liftIO $ putStrLn "Nothing to download"
 148         Just workPiece -> do
 149           let pLen = len (pieceStatus ! workPiece)
 150           liftIO $ putStrLn $ "piece length = " ++ show pLen
 151           pBS <- liftIO $ downloadPiece h workPiece pLen
 152           if not $ verifyHash pBS (hash (pieceStatus ! workPiece))
 153             then
 154             liftIO $ putStrLn $ "Hash mismatch: " ++ show (hash (pieceStatus ! workPiece)) ++ " vs " ++ show (take 20 (SHA1.hash pBS))
 155             else do
 156             let fileOffset = if workPiece == 0 then 0 else workPiece * len (pieceStatus ! (workPiece - 1))
 157             liftIO $ putStrLn $ "Write into file at offset: " ++ show fileOffset
 158             liftIO $ writeFileAtOffset file fileOffset pBS
 159             msgLoop (adjust (\pieceData -> pieceData { dlstate = Have }) workPiece pieceStatus) file
 160     _ -> do
 161       msg <- liftIO $ getMsg h
 162       gets peer >>= (\p -> liftIO $ putStrLn $ "<-- " ++ show msg ++ "from peer: " ++ show p)
 163       case msg of
 164         KeepAliveMsg -> do
 165           liftIO $ sendMsg h KeepAliveMsg
 166           gets peer >>= (\p -> liftIO $ putStrLn $ "--> " ++ "KeepAliveMsg to peer: " ++ show p)
 167           msgLoop pieceStatus file
 168         BitFieldMsg bss -> do
 169           p <- gets peer
 170           let pieceList = bitfieldToList (unpack bss)
 171               pieceStatus' = updatePieceAvailability pieceStatus p pieceList
 172           liftIO $ putStrLn $ show (length pieceList) ++ " Pieces"
 173           -- for each pieceIndex in pieceList, make an entry in the pieceStatus
 174           -- map with pieceIndex as the key and modify the value to add the peer.
 175           -- download each of the piece in order
 176           msgLoop pieceStatus' file
 177         UnChokeMsg -> do
 178           modify (\st -> st {heChoking = False })
 179           msgLoop pieceStatus file
 180
 181
 182 downloadPiece :: Handle -> Integer -> Integer -> IO ByteString
 183 downloadPiece h index pieceLength = do
 184   let chunks = splitNum pieceLength 16384
 185   liftM concat $ forM (zip [0..] chunks) (\(i, pLen) -> do
 186                                              sendMsg h (RequestMsg index (i*pLen) pLen)
 187                                              putStrLn $ "--> " ++ "RequestMsg for Piece "
 188                                                ++ show index ++ ", part: " ++ show i ++ " of length: "
 189                                                ++ show pLen
 190                                              msg <- getMsg h
 191                                              case msg of
 192                                               PieceMsg index begin block -> do
 193                                                 putStrLn $ " <-- PieceMsg for Piece: "
 194                                                   ++ show index
 195                                                   ++ ", offset: "
 196                                                   ++ show begin
 197                                                 return block
 198                                               _ -> do
 199                                                 putStrLn "ignoring irrelevant msg"
 200                                                 return empty)
 201
 202 verifyHash :: ByteString -> ByteString -> Bool
 203 verifyHash bs pieceHash =
 204   take 20 (SHA1.hash bs) == pieceHash