Protocol.hs

{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE StrictData #-}

-- | Wire protocol for Hegel.
--
-- This module defines the binary packet format used for communication between
-- the Hegel client and the Hegel server. Messages are serialized as packets
-- with a 20-byte header, a CBOR-encoded payload, and a terminator byte.
--
-- The header consists of 5 big-endian unsigned 32-bit integers:
--
--   * Magic number (0x4845474C, \"HEGL\")
--   * CRC32 checksum
--   * Channel ID
--   * Message ID (high bit = reply flag)
--   * Payload length
module Hegel.Protocol
  ( -- * Packet type
    Packet (..)

    -- * Constants
  , magic
  , headerSize
  , terminator
  , replyBit
  , closeChannelMessageId
  , closeChannelPayload

    -- * Packet I\/O
  , readPacket
  , writePacket

    -- * CRC32
  , computeCRC32
  , computeCRC32Parts

    -- * CBOR helpers
  , extractInt
  , extractFloat
  , extractBool
  , extractText
  , extractBytes
  , extractList
  , extractMap
  , encodeTerm
  , decodeTerm
  ) where

import Codec.CBOR.Read (deserialiseFromBytes)
import Codec.CBOR.Term (Term (..))
import qualified Codec.CBOR.Term as CBOR
import qualified Codec.CBOR.Write as CBORWrite
import Control.Monad (when)
import Data.Bits ((.&.), (.|.), shiftL, shiftR, testBit, xor)
import Data.ByteString (ByteString)
import qualified Data.ByteString as BS
import qualified Data.ByteString.Internal as BSInternal
import qualified Data.ByteString.Lazy as LBS
import Data.Digest.CRC32 (crc32, crc32Update)
import Data.Text (Text)
import Data.Word (Word32, Word8)
import Foreign.Ptr (Ptr)
import Foreign.Storable (pokeByteOff)
import System.IO (Handle, hFlush)

-- ---------------------------------------------------------------------------
-- Constants
-- ---------------------------------------------------------------------------

-- | Magic number identifying Hegel packets (\"HEGL\" in hex).
magic :: Word32
magic = 0x4845474C

-- | Size of the packet header in bytes (5 big-endian uint32 fields).
headerSize :: Int
headerSize = 20

-- | Terminator byte appended after the payload.
terminator :: Word8
terminator = 0x0A

-- | Bit flag indicating a reply message. When set in the message ID field,
-- the message is a response to a previous request.
replyBit :: Word32
replyBit = 1 `shiftL` 31

-- | Special message ID used when closing a channel.
closeChannelMessageId :: Word32
closeChannelMessageId = (1 `shiftL` 31) - 1

-- | Special payload byte sent when closing a channel. Chosen to be invalid
-- CBOR per RFC 8949 (reserved tag byte 0xFE).
closeChannelPayload :: ByteString
closeChannelPayload = BS.singleton 0xFE

-- ---------------------------------------------------------------------------
-- Packet type
-- ---------------------------------------------------------------------------

-- | A single message in the wire protocol.
data Packet = Packet
  { packetChannelId :: Word32
  -- ^ The logical channel this packet belongs to.
  , packetMessageId :: Word32
  -- ^ Identifier for request\/response correlation (without reply bit).
  , packetIsReply :: Bool
  -- ^ Whether this is a reply to a previous request.
  , packetPayload :: ByteString
  -- ^ The raw payload bytes (typically CBOR-encoded).
  }
  deriving (Show, Eq)

-- ---------------------------------------------------------------------------
-- Byte packing helpers
-- ---------------------------------------------------------------------------

-- | Write a 'Word32' in big-endian order at the given byte offset.
pokeWord32BE :: Ptr Word8 -> Int -> Word32 -> IO ()
pokeWord32BE ptr off w = do
  pokeByteOff ptr off (fromIntegral (w `shiftR` 24) :: Word8)
  pokeByteOff ptr (off + 1) (fromIntegral (w `shiftR` 16) :: Word8)
  pokeByteOff ptr (off + 2) (fromIntegral (w `shiftR` 8) :: Word8)
  pokeByteOff ptr (off + 3) (fromIntegral w :: Word8)

-- | Unpack a 'Word32' from 4 big-endian bytes starting at offset.
unpackWord32BE :: ByteString -> Int -> Word32
unpackWord32BE bs off =
  let b0 = fromIntegral (BS.index bs off) :: Word32
      b1 = fromIntegral (BS.index bs (off + 1)) :: Word32
      b2 = fromIntegral (BS.index bs (off + 2)) :: Word32
      b3 = fromIntegral (BS.index bs (off + 3)) :: Word32
  in (b0 `shiftL` 24) .|. (b1 `shiftL` 16) .|. (b2 `shiftL` 8) .|. b3

-- | Build a packet header from its wire-format fields.
buildHeader :: Word32 -> Word32 -> Word32 -> Word32 -> ByteString
buildHeader checksum channelId messageId payloadLen =
  BSInternal.unsafeCreate headerSize $ \ptr -> do
    pokeWord32BE ptr 0 magic
    pokeWord32BE ptr 4 checksum
    pokeWord32BE ptr 8 channelId
    pokeWord32BE ptr 12 messageId
    pokeWord32BE ptr 16 payloadLen

-- ---------------------------------------------------------------------------
-- CRC32
-- ---------------------------------------------------------------------------

-- | Compute the CRC32 checksum of a 'ByteString', returned as 'Word32'.
computeCRC32 :: ByteString -> Word32
computeCRC32 = crc32

-- | Compute the CRC32 checksum over the concatenation of two byte strings.
computeCRC32Parts :: ByteString -> ByteString -> Word32
computeCRC32Parts a = crc32Update (crc32 a)

-- ---------------------------------------------------------------------------
-- Handle-based exact read
-- ---------------------------------------------------------------------------

-- | Read exactly @n@ bytes from a 'Handle'. Raises an 'IOError' on short
-- read (i.e., if the handle reaches EOF before @n@ bytes are available).
hGetExact :: Handle -> Int -> IO ByteString
hGetExact _ 0 = return BS.empty
hGetExact h n = go BS.empty
  where
    go acc
      | BS.length acc >= n = return acc
      | otherwise = do
          chunk <- BS.hGetSome h (n - BS.length acc)
          if BS.null chunk
            then
              if BS.null acc
                then ioError $ userError "Connection closed partway through reading packet."
                else ioError $ userError "Connection closed while reading data"
            else go (BS.append acc chunk)

-- ---------------------------------------------------------------------------
-- Packet I\/O
-- ---------------------------------------------------------------------------

-- | Read and parse a single packet from the given 'Handle'.
--
-- Raises 'IOError' if the magic number is invalid, the terminator byte is
-- invalid, or the CRC32 checksum does not match.
readPacket :: Handle -> IO Packet
readPacket h = do
  headerBuf <- hGetExact h headerSize
  let pktMagic = unpackWord32BE headerBuf 0
      checksum = unpackWord32BE headerBuf 4
      channelId = unpackWord32BE headerBuf 8
      rawMessageId = unpackWord32BE headerBuf 12
      payloadLen = unpackWord32BE headerBuf 16
      isReply = testBit rawMessageId 31
      messageId = if isReply then rawMessageId `xor` replyBit else rawMessageId

  when (pktMagic /= magic) $
    ioError $ userError $
      "Invalid magic number: expected 0x" ++ showHex32 magic
        ++ ", got 0x" ++ showHex32 pktMagic

  payloadBuf <- hGetExact h (fromIntegral payloadLen)
  termBuf <- hGetExact h 1
  let termByte = BS.index termBuf 0

  when (termByte /= terminator) $
    ioError $ userError $
      "Invalid terminator: expected 0x" ++ showHex8 terminator
        ++ ", got 0x" ++ showHex8 termByte

  -- Verify checksum: CRC32 over header with checksum field zeroed + payload
  let headerForCheck = buildHeader 0 channelId rawMessageId payloadLen
      computedCRC = computeCRC32Parts headerForCheck payloadBuf

  when (computedCRC /= checksum) $
    ioError $ userError $
      "Checksum mismatch: expected 0x" ++ showHex32 checksum
        ++ ", got 0x" ++ showHex32 computedCRC

  return Packet
    { packetChannelId = channelId
    , packetMessageId = messageId
    , packetIsReply = isReply
    , packetPayload = payloadBuf
    }

-- | Serialize and write a packet to the given 'Handle'.
writePacket :: Handle -> Packet -> IO ()
writePacket h pkt = do
  let wireMessageId =
        if packetIsReply pkt
          then packetMessageId pkt .|. replyBit
          else packetMessageId pkt
      payloadLen = fromIntegral (BS.length (packetPayload pkt)) :: Word32

  -- Build header with zeroed checksum for CRC computation
  let headerForCRC = buildHeader 0 (packetChannelId pkt) wireMessageId payloadLen
      checksum = computeCRC32Parts headerForCRC (packetPayload pkt)

  -- Build final header with real checksum
  let headerFinal = buildHeader checksum (packetChannelId pkt) wireMessageId payloadLen
  BS.hPut h headerFinal
  BS.hPut h (packetPayload pkt)
  BS.hPut h (BS.singleton terminator)
  hFlush h

-- ---------------------------------------------------------------------------
-- Hex formatting helpers
-- ---------------------------------------------------------------------------

-- | Format a 'Word32' as an 8-digit uppercase hex string.
showHex32 :: Word32 -> String
showHex32 w = map toHexDigit (reverse [nibble i | i <- [0..7]])
  where
    nibble i = fromIntegral ((w `shiftR` (i * 4)) .&. 0xF) :: Int
    toHexDigit n
      | n < 10    = toEnum (fromEnum '0' + n)
      | otherwise = toEnum (fromEnum 'A' + n - 10)

-- | Format a 'Word8' as a 2-digit uppercase hex string.
showHex8 :: Word8 -> String
showHex8 w = [toHexDigit hi, toHexDigit lo]
  where
    hi = fromIntegral (w `shiftR` 4) :: Int
    lo = fromIntegral (w .&. 0xF) :: Int
    toHexDigit n
      | n < 10    = toEnum (fromEnum '0' + n)
      | otherwise = toEnum (fromEnum 'A' + n - 10)

-- ---------------------------------------------------------------------------
-- CBOR helpers
-- ---------------------------------------------------------------------------

-- | Extract an integer from a CBOR 'Term'.
--
-- Accepts 'TInt' and 'TInteger'. Raises an error for other types.
extractInt :: Term -> Int
extractInt (TInt i) = i
extractInt (TInteger i) = fromIntegral i
extractInt other = error $ "Expected CBOR int, got " ++ termTypeName other

-- | Extract a double-precision float from a CBOR 'Term'.
--
-- Accepts 'TFloat', 'TDouble', and 'THalf'. Raises an error for other types.
extractFloat :: Term -> Double
extractFloat (TFloat f) = realToFrac f
extractFloat (TDouble d) = d
extractFloat (THalf f) = realToFrac f
extractFloat other = error $ "Expected CBOR float, got " ++ termTypeName other

-- | Extract a boolean from a CBOR 'Term'.
--
-- Raises an error if the term is not 'TBool'.
extractBool :: Term -> Bool
extractBool (TBool b) = b
extractBool other = error $ "Expected CBOR bool, got " ++ termTypeName other

-- | Extract a text string from a CBOR 'Term'.
--
-- Raises an error if the term is not 'TString'.
extractText :: Term -> Text
extractText (TString t) = t
extractText other = error $ "Expected CBOR text, got " ++ termTypeName other

-- | Extract a byte string from a CBOR 'Term'.
--
-- Raises an error if the term is not 'TBytes'.
extractBytes :: Term -> ByteString
extractBytes (TBytes b) = b
extractBytes other = error $ "Expected CBOR bytes, got " ++ termTypeName other

-- | Extract a list from a CBOR 'Term'.
--
-- Raises an error if the term is not 'TList'.
extractList :: Term -> [Term]
extractList (TList xs) = xs
extractList other = error $ "Expected CBOR array, got " ++ termTypeName other

-- | Extract a key-value map from a CBOR 'Term'.
--
-- Raises an error if the term is not 'TMap'.
extractMap :: Term -> [(Term, Term)]
extractMap (TMap kvs) = kvs
extractMap other = error $ "Expected CBOR map, got " ++ termTypeName other

-- | Encode a CBOR 'Term' to a strict 'ByteString'.
encodeTerm :: Term -> ByteString
encodeTerm = CBORWrite.toStrictByteString . CBOR.encodeTerm

-- | Decode a strict 'ByteString' to a CBOR 'Term'.
--
-- Raises an error if the bytes are not valid CBOR.
decodeTerm :: ByteString -> Term
decodeTerm bs =
  case deserialiseFromBytes CBOR.decodeTerm (LBS.fromStrict bs) of
    Left err -> error $ "CBOR decode error: " ++ show err
    Right (_, term) -> term

-- | Return a human-readable name for the type of a CBOR 'Term'.
termTypeName :: Term -> String
termTypeName (TInt _) = "int"
termTypeName (TInteger _) = "integer"
termTypeName (TFloat _) = "float"
termTypeName (TDouble _) = "double"
termTypeName (THalf _) = "half"
termTypeName (TBool _) = "bool"
termTypeName TNull = "null"
termTypeName (TBytes _) = "bytes"
termTypeName (TBytesI _) = "bytes (indefinite)"
termTypeName (TString _) = "text"
termTypeName (TStringI _) = "text (indefinite)"
termTypeName (TList _) = "array"
termTypeName (TListI _) = "array (indefinite)"
termTypeName (TMap _) = "map"
termTypeName (TMapI _) = "map (indefinite)"
termTypeName (TTagged _ _) = "tag"
termTypeName (TSimple _) = "simple"