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