Printer.hs

-----------------------------------------------------------------------------
-- |
-- Module      :  Conjure.Protocol.PWP.Printer
-- Copyright   :  (c) Lemmih 2005-2006
-- License     :  BSD-like
--
-- Maintainer  :  lemmih@gmail.com
-- Stability   :  experimental
-- Portability :  portable
--
-- Output handler for the Peer Wire Protocol.
-----------------------------------------------------------------------------
module Conjure.Protocol.PWP.Printer
    ( hPutHandshake
    , hPutMessage
    , serializeMessage
    ) where

import Data.Word
import System.IO
import Data.Bits

import qualified Data.ByteString as BS
import qualified Data.ByteString.Lazy as LBS

import Conjure.Protocol.PWP.Types

-- | Write a 4-byte big-endian @Int@ to the wire
hPutInt :: Handle -> Int -> IO ()
hPutInt h x = BS.hPut h (BS.pack (serializeInt x))

-- | Convert an @Int@ to 4-byte big-endian format
serializeInt :: Int -> [Word8]
serializeInt x
    = [ fromIntegral $ x `shiftR` shft .&. 255 | shft <- [24,16,8,0]]



hPutHandshake :: Handle -> Handshake -> IO ()
hPutHandshake h (Handshake i (PeerId c))
    = do hPutStr h $ "\19BitTorrent protocol" ++ replicate 8 '\0'
         BS.hPut h i
         BS.hPut h c
         hFlush h


-- | Write a @Message@ to a @Handle@
hPutMessage :: Handle -> Message -> IO ()
hPutMessage h msg = do let str = serializeMessage msg
                       hPutInt h (myLength str)
                       LBS.hPut h str
                       hFlush h
  where myLength = sum . map BS.length . LBS.toChunks

-- | Return a message\'s body in wire form
serializeMessage :: Message -> LBS.ByteString
serializeMessage KeepAlive = LBS.empty
serializeMessage msg       = LBS.singleton tag `LBS.append` rest 
    where
    (tag, rest) = case msg of 
                  Choke         -> (0, LBS.empty)
                  Unchoke       -> (1, LBS.empty)
                  Interested    -> (2, LBS.empty)
                  NotInterested -> (3, LBS.empty)
                  Have n        -> (4, LBS.pack (serializeInt n))
                  BitField pieces  -> (5, LBS.fromChunks [pieces])
                  Request p off size -> (6, LBS.pack (concatMap serializeInt [p,off,size]))
                  RequestedPiece p off str 
                                     -> (7, LBS.fromChunks $ BS.pack (concatMap serializeInt [p,off]):[str])
                  Cancel  p off size -> (8, LBS.pack (concatMap serializeInt [p,off,size]))
                  _ -> error "Conjure.Protocol.PWP.Printer.serializeMessage: can't happen"