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, writeFile, take)
   8
   9 import System.IO (Handle, BufferMode(..), IOMode(..), SeekMode(..), withFile, hSeek, hSetBuffering)
  10 import Data.ByteString (ByteString, pack, unpack, concat, hGet, hPut, singleton, writeFile, take, empty)
  11 import Data.ByteString.Lazy (fromStrict, fromChunks, toStrict)
  12 import qualified Data.ByteString.Char8 as BC (replicate, pack, length)
  13 import Network (connectTo, PortID(..))
  14 import Data.Binary (Binary(..), decode, encode)
  15 import Data.Binary.Put (putWord32be, putWord16be, putWord8)
  16 import Data.Binary.Get (getWord32be, getWord16be, getWord8, runGet)
  17 import Control.Monad (replicateM, liftM, forM)
  18 import Control.Applicative ((<$>), liftA3)
  19 import Data.Bits
  20 import Data.Word (Word8)
  21 import Data.Map (Map, fromList, toList, (!), mapWithKey, adjust)
  22 import qualified Crypto.Hash.SHA1 as SHA1 (hash)
  23
  24 import FuncTorrent.Metainfo (Info(..), Metainfo(..))
  25 import FuncTorrent.Utils (splitN, splitNum)
  26
  27 type ID = String
  28 type IP = String
  29 type Port = Integer
  30
  31 -- PeerState is a misnomer
  32 data PeerState = PeerState { handle :: Handle
  33                            , peer :: Peer
  34                            , meChoking :: Bool
  35                            , meInterested :: Bool
  36                            , heChoking :: Bool
  37                            , heInterested :: Bool}
  38
  39 data PieceDlState = Pending
  40                   | InProgress
  41                   | Have
  42                   deriving (Show, Eq)
  43
  44 -- todo - map with index to a new data structure (peers who have that piece amd state)
  45 data PieceData = PieceData { peers :: [Peer]        -- ^ list of peers who have this piece
  46                            , state :: PieceDlState  -- ^ state of the piece from download perspective.
  47                            , hash  :: ByteString    -- ^ piece hash
  48                            , len :: Integer }       -- ^ piece length
  49
  50 -- which piece is with which peers
  51 type PieceMap = Map Integer PieceData
  52
  53 -- | Peer is a PeerID, IP address, port tuple
  54 data Peer = Peer ID IP Port
  55           deriving (Show, Eq)
  56
  57 data PeerMsg = KeepAliveMsg
  58              | ChokeMsg
  59              | UnChokeMsg
  60              | InterestedMsg
  61              | NotInterestedMsg
  62              | HaveMsg Integer
  63              | BitFieldMsg ByteString
  64              | RequestMsg Integer Integer Integer
  65              | PieceMsg Integer Integer ByteString
  66              | CancelMsg Integer Integer Integer
  67              | PortMsg Port
  68              deriving (Show)
  69
  70 -- Make the initial Piece map, with the assumption that no peer has the
  71 -- piece and that every piece is pending download.
  72 mkPieceMap :: Integer -> ByteString -> [Integer] -> PieceMap
  73 mkPieceMap numPieces pieceHash pLengths = fromList kvs
  74   where kvs = [(i, PieceData { peers = []
  75                              , state = Pending
  76                              , hash = h
  77                              , len = pLen })
  78               | (i, h, pLen) <- zip3 [0..numPieces] hashes pLengths]
  79         hashes = splitN 20 pieceHash
  80
  81 havePiece :: PieceMap -> Integer -> Bool
  82 havePiece pm index =
  83   state (pm ! index) == Have
  84
  85 genHandShakeMsg :: ByteString -> String -> ByteString
  86 genHandShakeMsg infoHash peer_id = concat [pstrlen, pstr, reserved, infoHash, peerID]
  87   where pstrlen = singleton 19
  88         pstr = BC.pack "BitTorrent protocol"
  89         reserved = BC.replicate 8 '\0'
  90         peerID = BC.pack peer_id
  91
  92 handShake :: Peer -> ByteString -> String -> IO Handle
  93 handShake peer@(Peer _ ip port) infoHash peerid = do
  94   let hs = genHandShakeMsg infoHash peerid
  95   h <- connectTo ip (PortNumber (fromIntegral port))
  96   hSetBuffering h LineBuffering
  97   hPut h hs
  98   putStrLn $ "--> handhake to peer: " ++ show peer
  99   _ <- hGet h (length (unpack hs))
 100   putStrLn $ "<-- handshake from peer: " ++ show peer
 101   return h
 102
 103 instance Binary PeerMsg where
 104   put msg = case msg of
 105              KeepAliveMsg -> putWord32be 0
 106              ChokeMsg -> do putWord32be 1
 107                             putWord8 0
 108              UnChokeMsg -> do putWord32be 1
 109                               putWord8 1
 110              InterestedMsg -> do putWord32be 1
 111                                  putWord8 2
 112              NotInterestedMsg -> do putWord32be 1
 113                                     putWord8 3
 114              HaveMsg i -> do putWord32be 5
 115                              putWord8 4
 116                              putWord32be (fromIntegral i)
 117              BitFieldMsg bf -> do putWord32be $ fromIntegral (1 + bfListLen)
 118                                   putWord8 5
 119                                   mapM_ putWord8 bfList
 120                                     where bfList = unpack bf
 121                                           bfListLen = length bfList
 122              RequestMsg i o l -> do putWord32be 13
 123                                     putWord8 6
 124                                     putWord32be (fromIntegral i)
 125                                     putWord32be (fromIntegral o)
 126                                     putWord32be (fromIntegral l)
 127              PieceMsg i o b -> do putWord32be $ fromIntegral (9 + blocklen)
 128                                   putWord8 7
 129                                   putWord32be (fromIntegral i)
 130                                   putWord32be (fromIntegral o)
 131                                   mapM_ putWord8 blockList
 132                                     where blockList = unpack b
 133                                           blocklen = length blockList
 134              CancelMsg i o l -> do putWord32be 13
 135                                    putWord8 8
 136                                    putWord32be (fromIntegral i)
 137                                    putWord32be (fromIntegral o)
 138                                    putWord32be (fromIntegral l)
 139              PortMsg p -> do putWord32be 3
 140                              putWord8 9
 141                              putWord16be (fromIntegral p)
 142   get = do
 143     l <- getWord32be
 144     msgid <- getWord8
 145     case msgid of
 146      0 -> return ChokeMsg
 147      1 -> return UnChokeMsg
 148      2 -> return InterestedMsg
 149      3 -> return NotInterestedMsg
 150      4 -> liftM (HaveMsg . fromIntegral) getWord32be
 151      5 -> liftM (BitFieldMsg . pack) (replicateM (fromIntegral l - 1) getWord8)
 152      6 -> liftA3 RequestMsg getInteger getInteger getInteger
 153        where getInteger = fromIntegral <$> getWord32be
 154      7 -> liftA3 PieceMsg getInteger getInteger (pack  <$> replicateM (fromIntegral l - 9) getWord8)
 155        where getInteger = fromIntegral <$> getWord32be
 156      8 -> liftA3 CancelMsg getInteger getInteger getInteger
 157        where getInteger = fromIntegral <$> getWord32be
 158      9 -> liftM (PortMsg . fromIntegral) getWord16be
 159      _ -> error ("unknown message ID: " ++ show msgid)
 160
 161 getMsg :: Handle -> IO PeerMsg
 162 getMsg h = do
 163   lBS <- hGet h 4
 164   let l = bsToInt lBS
 165   if l == 0
 166     then return KeepAliveMsg
 167     else do
 168     msg <- hGet h l
 169     return $ decode $ fromStrict $ concat [lBS, msg]
 170
 171 sendMsg :: Handle -> PeerMsg -> IO ()
 172 sendMsg h msg =
 173   let bsMsg = toStrict $ encode msg
 174   in
 175    hPut h bsMsg
 176
 177 bsToInt :: ByteString -> Int
 178 bsToInt x = fromIntegral (runGet getWord32be (fromChunks (return x)))
 179
 180 bitfieldToList :: [Word8] -> [Integer]
 181 bitfieldToList bs = go bs 0
 182   where go [] _ = []
 183         go (b:bs') pos =
 184           let setBits = [pos*8 + toInteger i | i <- [0..8], testBit b i]
 185           in
 186            setBits ++ go bs' (pos + 1)
 187
 188 createDummyFile :: FilePath -> Int -> IO ()
 189 createDummyFile path size =
 190   writeFile path (BC.replicate size '\0')
 191
 192 -- write into a file at a specific offet
 193 writeFileAtOffset :: FilePath -> Integer -> ByteString -> IO ()
 194 writeFileAtOffset path offset block =
 195   withFile path ReadWriteMode $ (\h -> do
 196                                  _ <- hSeek h AbsoluteSeek offset
 197                                  hPut h block)
 198
 199 -- recvMsg :: Peer -> Handle -> Msg
 200 msgLoop :: PeerState -> PieceMap -> IO ()
 201 msgLoop pState pieceStatus | meInterested pState == False &&
 202                             heChoking pState == True = do
 203                               -- if me NOT Interested and she is Choking, tell her that
 204                               -- I am interested.
 205                               let h = handle pState
 206                               sendMsg h InterestedMsg
 207                               putStrLn $ "--> InterestedMsg to peer: " ++ show (peer pState)
 208                               msgLoop (pState { meInterested = True }) pieceStatus
 209                           | meInterested pState == True &&
 210                             heChoking pState == False =
 211                               -- if me Interested and she not Choking, send her a request
 212                               -- for a piece.
 213                               case pickPiece pieceStatus of
 214                                Nothing -> putStrLn "Nothing to download"
 215                                Just workPiece -> do
 216                                  let pLen = len (pieceStatus ! workPiece)
 217                                  pBS <- downloadPiece (handle pState) workPiece pLen
 218                                  -- if not $ verifyHash pBS (hash (pieceStatus ! workPiece))
 219                                  --  then
 220                                  --  putStrLn $ "Hash mismatch: " ++ show (hash (pieceStatus ! workPiece)) ++ " vs " ++ show (take 20 (SHA1.hash pBS))
 221                                  --  else do
 222                                  writeFileAtOffset "/tmp/download.file" (workPiece * pLen) pBS
 223                                  msgLoop pState (adjust (\pieceData -> pieceData { state = Have }) workPiece pieceStatus)
 224                           | otherwise = do
 225                               msg <- getMsg (handle pState)
 226                               putStrLn $ "<-- " ++ show msg ++ "from peer: " ++ show (peer pState)
 227                               case msg of
 228                                KeepAliveMsg -> do
 229                                  sendMsg (handle pState) KeepAliveMsg
 230                                  putStrLn $ "--> " ++ "KeepAliveMsg to peer: " ++ show (peer pState)
 231                                  msgLoop pState pieceStatus
 232                                BitFieldMsg bss -> do
 233                                  let pieceList = bitfieldToList (unpack bss)
 234                                      pieceStatus' = updatePieceAvailability pieceStatus (peer pState) pieceList
 235                                  putStrLn $ show (length pieceList) ++ " Pieces"
 236                                  -- for each pieceIndex in pieceList, make an entry in the pieceStatus
 237                                  -- map with pieceIndex as the key and modify the value to add the peer.
 238                                  -- download each of the piece in order
 239                                  msgLoop pState pieceStatus'
 240                                UnChokeMsg -> do
 241                                  msgLoop (pState { heChoking = False }) pieceStatus
 242                                _ -> do
 243                                  msgLoop pState pieceStatus
 244
 245 -- simple algorithm to pick piece.
 246 -- pick the first piece from 0 that is not downloaded yet.
 247 pickPiece :: PieceMap -> Maybe Integer
 248 pickPiece m =
 249   let pieceList = toList m
 250       allPending = filter (\(_, v) -> state v == Pending) pieceList
 251   in
 252    case allPending of
 253     [] -> Nothing
 254     ((i, _):_) -> Just i
 255
 256 updatePieceAvailability :: PieceMap -> Peer -> [Integer] -> PieceMap
 257 updatePieceAvailability pieceStatus p pieceList =
 258   mapWithKey (\k pd -> if k `elem` pieceList
 259                        then (pd { peers = p : (peers pd) })
 260                        else pd) pieceStatus
 261
 262 handlePeerMsgs :: Peer -> Metainfo -> String -> IO ()
 263 handlePeerMsgs p m peerId = do
 264   h <- handShake p (infoHash m) peerId
 265   let state = PeerState { handle = h
 266                         , peer = p
 267                         , heInterested = False
 268                         , heChoking = True
 269                         , meInterested = False
 270                         , meChoking = True }
 271       pieceHash = pieces (info m)
 272       numPieces = (toInteger . (`quot` 20) . BC.length) pieceHash
 273       pLen = pieceLength (info m)
 274       fileLen = lengthInBytes (info m)
 275       pieceStatus = mkPieceMap numPieces pieceHash (splitNum fileLen pLen)
 276   createDummyFile "/tmp/download.file" (fromIntegral fileLen)
 277   msgLoop state pieceStatus
 278
 279 downloadPiece :: Handle -> Integer -> Integer -> IO ByteString
 280 downloadPiece h index pieceLength = do
 281   let chunks = splitNum pieceLength 16384
 282   liftM concat $ forM (zip [0..] chunks) (\(i, pLen) -> do
 283                                              sendMsg h (RequestMsg index (i*pLen) pLen)
 284                                              putStrLn $ "--> " ++ "RequestMsg for Piece "
 285                                                ++ (show index) ++ ", part: " ++ show i ++ " of length: "
 286                                                ++ show pLen
 287                                              msg <- getMsg h
 288                                              case msg of
 289                                               PieceMsg index begin block -> do
 290                                                 putStrLn $ " <-- PieceMsg for Piece: "
 291                                                   ++ show index
 292                                                   ++ ", offset: "
 293                                                   ++ show begin
 294                                                 return block
 295                                               _ -> do
 296                                                 putStrLn "ignoring irrelevant msg"
 297                                                 return empty)
 298
 299 verifyHash :: ByteString -> ByteString -> Bool
 300 verifyHash bs pieceHash =
 301   take 20 (SHA1.hash bs) == pieceHash