ProtocolSpec.hs

{-# LANGUAGE OverloadedStrings #-}

module Hegel.ProtocolSpec (spec) where

import Codec.CBOR.Term (Term (..))
import qualified Data.ByteString as BS
import Data.Bits (shiftL, shiftR, (.|.))
import Data.Word (Word32)
import System.IO (Handle, hFlush, hSetBinaryMode)
import System.Process (createPipe)
import Test.Hspec

import Hegel.Protocol

-- | Helper: create a pipe pair of handles for testing.
createPipePair :: IO (Handle, Handle)
createPipePair = createPipe

-- | Pack a Word32 as 4 big-endian bytes (reimplemented for testing since
-- the library does not export this).
testPackWord32BE :: Word32 -> BS.ByteString
testPackWord32BE w = BS.pack
  [ fromIntegral (w `shiftR` 24)
  , fromIntegral (w `shiftR` 16)
  , fromIntegral (w `shiftR` 8)
  , fromIntegral w
  ]

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

spec :: Spec
spec = do
  describe "CRC32" $ do
    it "computes the known check value for '123456789'" $ do
      let input = "123456789"
      computeCRC32 input `shouldBe` 0xCBF43926

  describe "pack/unpack Word32 BE" $ do
    it "round-trips uint32 values" $ do
      let values = [0, 1, 255, 256, 65535, 0x12345678, 0xFFFFFFFF] :: [Word32]
      mapM_
        ( \w -> do
            let packed = testPackWord32BE w
            BS.length packed `shouldBe` 4
            testUnpackWord32BE packed `shouldBe` w
        )
        values

  describe "Packet round-trip" $ do
    it "writes and reads a packet through a pipe pair" $ do
      (readH, writeH) <- createPipePair
      hSetBinaryMode readH True
      hSetBinaryMode writeH True
      let pkt =
            Packet
              { packetChannelId = 42,
                packetMessageId = 7,
                packetIsReply = False,
                packetPayload = encodeTerm (TString "hello")
              }
      writePacket writeH pkt
      result <- readPacket readH
      packetChannelId result `shouldBe` 42
      packetMessageId result `shouldBe` 7
      packetIsReply result `shouldBe` False
      decodeTerm (packetPayload result) `shouldBe` TString "hello"

    it "round-trips a reply packet" $ do
      (readH, writeH) <- createPipePair
      hSetBinaryMode readH True
      hSetBinaryMode writeH True
      let pkt =
            Packet
              { packetChannelId = 10,
                packetMessageId = 3,
                packetIsReply = True,
                packetPayload = encodeTerm (TBool True)
              }
      writePacket writeH pkt
      result <- readPacket readH
      packetChannelId result `shouldBe` 10
      packetMessageId result `shouldBe` 3
      packetIsReply result `shouldBe` True

  describe "Invalid magic" $ do
    it "raises on invalid magic number" $ do
      (readH, writeH) <- createPipePair
      hSetBinaryMode readH True
      hSetBinaryMode writeH True
      -- Write a header with bad magic (0xDEADBEEF), followed by terminator
      let badHeader =
            BS.concat
              [ testPackWord32BE 0xDEADBEEF,
                testPackWord32BE 0,
                testPackWord32BE 1,
                testPackWord32BE 1,
                testPackWord32BE 0
              ]
      BS.hPut writeH badHeader
      BS.hPut writeH (BS.singleton 0x0A)
      hFlush writeH
      readPacket readH `shouldThrow` anyIOException

  describe "Bad CRC" $ do
    it "raises on checksum mismatch" $ do
      (readH, writeH) <- createPipePair
      hSetBinaryMode readH True
      hSetBinaryMode writeH True
      let payload = encodeTerm (TString "test")
      let payloadLen = fromIntegral (BS.length payload) :: Word32
      -- Compute correct CRC then corrupt it
      let headerForCRC =
            BS.concat
              [ testPackWord32BE magic,
                testPackWord32BE 0,
                testPackWord32BE 1,
                testPackWord32BE 1,
                testPackWord32BE payloadLen
              ]
      let correctCRC = computeCRC32Parts headerForCRC payload
      let wrongCRC = correctCRC + 1
      let badHeader =
            BS.concat
              [ testPackWord32BE magic,
                testPackWord32BE wrongCRC,
                testPackWord32BE 1,
                testPackWord32BE 1,
                testPackWord32BE payloadLen
              ]
      BS.hPut writeH badHeader
      BS.hPut writeH payload
      BS.hPut writeH (BS.singleton 0x0A)
      hFlush writeH
      readPacket readH `shouldThrow` anyIOException