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