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
  88             -> Peer
  89             -> Bool  -- ^ meChoking
  90             -> Bool  -- ^ meInterested
  91             -> Bool  -- ^ heChoking
  92             -> Bool  -- ^ heInterested
  93             -> PState
  94 toPeerState h p meCh meIn heCh heIn =
  95   PState { handle = h
  96          , peer = p
  97          , heChoking = heCh
  98          , heInterested = heIn
  99          , meChoking = meCh
 100          , meInterested = meIn }
 101
 102 -- -- recvMsg :: Peer -> Handle -> Msg
 103 -- msgLoop :: PeerState -> PieceMap -> FilePath -> IO ()
 104 -- msgLoop pState@(PeerState { meInterested = False, heChoking = True }) pieceStatus file = do
 105 --   -- if me NOT Interested and she is Choking, tell her that
 106 --   -- I am interested.
 107 --   let h = handle pState
 108 --   sendMsg h InterestedMsg
 109 --   putStrLn $ "--> InterestedMsg to peer: " ++ show (peer pState)
 110 --   msgLoop (pState { meInterested = True }) pieceStatus file
 111 -- msgLoop pState@(PeerState { meInterested = True, heChoking = False }) pieceStatus file =
 112 --   -- if me Interested and she not Choking, send her a request
 113 --   -- for a piece.
 114 --   case pickPiece pieceStatus of
 115 --    Nothing -> putStrLn "Nothing to download"
 116 --    Just workPiece -> do
 117 --      let pLen = len (pieceStatus ! workPiece)
 118 --      putStrLn $ "piece length = " ++ show pLen
 119 --      pBS <- downloadPiece (handle pState) workPiece pLen
 120 --      if not $ verifyHash pBS (hash (pieceStatus ! workPiece))
 121 --        then
 122 --        putStrLn $ "Hash mismatch: " ++ show (hash (pieceStatus ! workPiece)) ++ " vs " ++ show (take 20 (SHA1.hash pBS))
 123 --        else do
 124 --        let fileOffset = if workPiece == 0 then 0 else workPiece * len (pieceStatus ! (workPiece - 1))
 125 --        putStrLn $ "Write into file at offset: " ++ show fileOffset
 126 --        writeFileAtOffset file fileOffset pBS
 127 --        msgLoop pState (adjust (\pieceData -> pieceData { state = Have }) workPiece pieceStatus) file
 128 -- msgLoop pState pieceStatus file = do
 129 --   msg <- getMsg (handle pState)
 130 --   putStrLn $ "<-- " ++ show msg ++ "from peer: " ++ show (peer pState)
 131 --   case msg of
 132 --    KeepAliveMsg -> do
 133 --      sendMsg (handle pState) KeepAliveMsg
 134 --      putStrLn $ "--> " ++ "KeepAliveMsg to peer: " ++ show (peer pState)
 135 --      msgLoop pState pieceStatus file
 136 --    BitFieldMsg bss -> do
 137 --      let pieceList = bitfieldToList (unpack bss)
 138 --          pieceStatus' = updatePieceAvailability pieceStatus (peer pState) pieceList
 139 --      putStrLn $ show (length pieceList) ++ " Pieces"
 140 --      -- for each pieceIndex in pieceList, make an entry in the pieceStatus
 141 --      -- map with pieceIndex as the key and modify the value to add the peer.
 142 --      -- download each of the piece in order
 143 --      msgLoop pState pieceStatus' file
 144 --    UnChokeMsg ->
 145 --      msgLoop (pState { heChoking = False }) pieceStatus file
 146 --    _ ->
 147 --      msgLoop pState pieceStatus file
 148
 149 -- simple algorithm to pick piece.
 150 -- pick the first piece from 0 that is not downloaded yet.
 151 pickPiece :: PieceMap -> Maybe Integer
 152 pickPiece m =
 153   let pieceList = toList m
 154       allPending = filter (\(_, v) -> dlstate v == Pending) pieceList
 155   in
 156    case allPending of
 157     [] -> Nothing
 158     ((i, _):_) -> Just i
 159
 160 updatePieceAvailability :: PieceMap -> Peer -> [Integer] -> PieceMap
 161 updatePieceAvailability pieceStatus p pieceList =
 162   mapWithKey (\k pd -> if k `elem` pieceList
 163                        then (pd { peers = p : peers pd })
 164                        else pd) pieceStatus
 165
 166 handlePeerMsgs :: Peer -> Metainfo -> String -> IO ()
 167 handlePeerMsgs p m peerId = do
 168   h <- connectToPeer p
 169   doHandshake h p (infoHash m) peerId
 170   let pstate = toPeerState h p False True False True
 171       pieceHash = pieces (info m)
 172       numPieces = (toInteger . (`quot` 20) . BC.length) pieceHash
 173       pLen = pieceLength (info m)
 174       fileLen = lengthInBytes (info m)
 175       fileName = name (info m)
 176       pieceStatus = mkPieceMap numPieces pieceHash (splitNum fileLen pLen)
 177   createDummyFile fileName (fromIntegral fileLen)
 178   (r, _) <- runStateT (msgLoop pieceStatus fileName) pstate
 179   return ()
 180
 181 msgLoop :: PieceMap -> FilePath -> StateT PState IO ()
 182 msgLoop pieceStatus file = do
 183   h <- gets handle
 184   msg <- liftIO $ getMsg h
 185   gets peer >>= (\p -> liftIO $ putStrLn $ "<-- " ++ show msg ++ "from peer: " ++ show p)
 186   case msg of
 187     KeepAliveMsg -> do
 188       liftIO $ sendMsg h KeepAliveMsg
 189       gets peer >>= (\p -> liftIO $ putStrLn $ "--> " ++ "KeepAliveMsg to peer: " ++ show p)
 190       msgLoop pieceStatus file
 191
 192 downloadPiece :: Handle -> Integer -> Integer -> IO ByteString
 193 downloadPiece h index pieceLength = do
 194   let chunks = splitNum pieceLength 16384
 195   liftM concat $ forM (zip [0..] chunks) (\(i, pLen) -> do
 196                                              sendMsg h (RequestMsg index (i*pLen) pLen)
 197                                              putStrLn $ "--> " ++ "RequestMsg for Piece "
 198                                                ++ show index ++ ", part: " ++ show i ++ " of length: "
 199                                                ++ show pLen
 200                                              msg <- getMsg h
 201                                              case msg of
 202                                               PieceMsg index begin block -> do
 203                                                 putStrLn $ " <-- PieceMsg for Piece: "
 204                                                   ++ show index
 205                                                   ++ ", offset: "
 206                                                   ++ show begin
 207                                                 return block
 208                                               _ -> do
 209                                                 putStrLn "ignoring irrelevant msg"
 210                                                 return empty)
 211
 212 verifyHash :: ByteString -> ByteString -> Bool
 213 verifyHash bs pieceHash =
 214   take 20 (SHA1.hash bs) == pieceHash