Parser.hs

-----------------------------------------------------------------------------
-- |
-- Module      :  Conjure.Protocol.PWP.Parser
-- Copyright   :  (c) Lemmih 2005-2006
-- License     :  BSD-like
--
-- Maintainer  :  lemmih@gmail.com
-- Stability   :  experimental
-- Portability :  non-portable (requires Foreign.*)
--
-- Input handler for the Peer Wire Protocol.
-----------------------------------------------------------------------------
module Conjure.Protocol.PWP.Parser
    ( hGetHandshake
    , hGetMessage
    -- For tests
    , parseMessage
    ) where

import Control.Monad ( liftM )
import System.IO
import System.IO.Unsafe
import Foreign

import qualified Data.ByteString.Char8 as Char8
import qualified Data.ByteString as BS
import qualified Data.ByteString.Internal as BS
import Data.ByteString (ByteString)

import Conjure.Protocol.PWP.Types
import Conjure.Utils

-- | Read a 4-byte big-endian @Int@ off the wire
hGetInt :: Handle -> IO Int
hGetInt h = do [a,b,c,d] <- mapM (liftM fromEnum) (replicate 4 (hGetChar h))
               return (a `shiftL` 24 + b `shiftL` 16 + c `shiftL` 8 + d)

-- | Parse a 4-byte big-endian @Int@
parseInt :: ByteString -> (Int, ByteString)
parseInt str
    = let [a,b,c,d] = map (fromIntegral . BS.index str) [0,1,2,3]
      in (a `shiftL` 24 .|. b `shiftL` 16 .|. c `shiftL` 8 .|. d, BS.drop 4 str)


hGetHandshake :: Handle -> IO Handshake
hGetHandshake h = do c <- hGetChar h
                     let len = fromEnum c
                     if len /= 19 then error "Unknown protocol in handshake" else do
                     msg <- BS.hGet h len
                     if msg /= Char8.pack "BitTorrent protocol" then error "Unknown protocol in handshake" else do
                     BS.hGet h 8 -- Those 8 reserved bytes ...
                     infohash <- BS.hGet h 20
                     client <- unsafeInterleaveIO $ BS.hGet h 20
                     -- 'unsafeInterleaveIO' is required because
                     -- handshake could be without client id. This
                     -- happens when tracker wants to check whether
                     -- we are alive and not behind the NAT. This
                     -- way we could postpone reading client id for
                     -- later, after we have sent our own handshake
                     return $ Handshake infohash (PeerId client)

-- This a bit ugly.
-- | Read a @Message@ from a @Handle@ and report when chunks are downloaded.
hGetMessage :: (Integer -- ^ Current time in milliseconds.
                -> Int  -- ^ Bytes per second.
                -> IO ())
            -> Handle -> IO Message
hGetMessage report h
    = do len <- hGetInt h
         if len == 0 then return KeepAlive else do
         let loop l ptr
                 = do hWaitForInput h (-1)
                      realLen <- hGetBufNonBlocking h ptr l
                      now <- getCurrentTime
                      report now realLen
                      if realLen == l
                         then return ()
                         else loop (l-realLen) (advancePtr ptr realLen)
         msgID <- fmap BS.head $ BS.hGet h 1
         msg <- if msgID == 7
                then do fp <- mallocForeignPtrBytes (len-1)
                        withForeignPtr fp (loop (len-1))
                        return $ BS.fromForeignPtr fp 0 (len-1)
                else BS.hGet h (len - 1)
         return $! parseMessage msgID msg

-- FIXME: error handling.
parseMessage :: Word8 -> ByteString -> Message
parseMessage msgID payload
    = case msgID of
        0 -> Choke
        1 -> Unchoke
        2 -> Interested
        3 -> NotInterested
        4 -> Have (fst $ parseInt payload)
        5 -> BitField payload
        6 -> let [idx, offset, len] = take 3 (combine parseInt)
             in Request idx offset len
        7 -> let (idx,str') = parseInt payload
                 (offset,str'') = parseInt str'
             in RequestedPiece idx offset str''
        8 -> let [idx, offset, len] = take 3 (combine parseInt)
             in Cancel idx offset len
        _ -> Unknown msgID
    where combine f = let loop str = let (val,str') = f str in val : loop str'
                      in loop payload