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 -> FilePath -> IO ()
  87 msgLoop pState@(PeerState { meInterested = False, heChoking = True }) pieceStatus file =
  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 file
  95 msgLoop pState@(PeerState { meInterested = True, heChoking = False }) pieceStatus file =
  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 file fileOffset pBS
 111        msgLoop pState (adjust (\pieceData -> pieceData { state = Have }) workPiece pieceStatus) file
 112 msgLoop pState pieceStatus file = 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 file
 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' file
 128    UnChokeMsg ->
 129      msgLoop (pState { heChoking = False }) pieceStatus file
 130    _ ->
 131      msgLoop pState pieceStatus file
 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       fileName = name (info m)
 165       pieceStatus = mkPieceMap numPieces pieceHash (splitNum fileLen pLen)
 166   createDummyFile fileName (fromIntegral fileLen)
 167   msgLoop state pieceStatus fileName
 168
 169 downloadPiece :: Handle -> Integer -> Integer -> IO ByteString
 170 downloadPiece h index pieceLength = do
 171   let chunks = splitNum pieceLength 16384
 172   liftM concat $ forM (zip [0..] chunks) (\(i, pLen) -> do
 173                                              sendMsg h (RequestMsg index (i*pLen) pLen)
 174                                              putStrLn $ "--> " ++ "RequestMsg for Piece "
 175                                                ++ show index ++ ", part: " ++ show i ++ " of length: "
 176                                                ++ show pLen
 177                                              msg <- getMsg h
 178                                              case msg of
 179                                               PieceMsg index begin block -> do
 180                                                 putStrLn $ " <-- PieceMsg for Piece: "
 181                                                   ++ show index
 182                                                   ++ ", offset: "
 183                                                   ++ show begin
 184                                                 return block
 185                                               _ -> do
 186                                                 putStrLn "ignoring irrelevant msg"
 187                                                 return empty)
 188
 189 verifyHash :: ByteString -> ByteString -> Bool
 190 verifyHash bs pieceHash =
 191   take 20 (SHA1.hash bs) == pieceHash