Messaging.hs

{-# LANGUAGE OverloadedStrings #-}

module Test.Network.Transport.QUIC.Internal.Messaging (tests) where

import Control.Monad.IO.Class (MonadIO (liftIO))
import Data.ByteString (ByteString)
import Data.ByteString qualified as BS
import Data.IORef (atomicModifyIORef, newIORef)
import Hedgehog (forAll, property, (===))
import Hedgehog.Gen qualified as Gen
import Hedgehog.Range qualified as Range
import Network.Transport.QUIC.Internal (MessageReceived (..), decodeMessage, encodeMessage)
import Test.Tasty (TestTree, testGroup)
import Test.Tasty.Hedgehog (testProperty)

tests :: TestTree
tests =
  testGroup
    "Messaging"
    [testMessageEncodingAndDecoding]

testMessageEncodingAndDecoding :: TestTree
testMessageEncodingAndDecoding = testProperty "Encoded messages can be decoded" $ property $ do
  -- The message length is encoded and decoded as a Word32. Generate data above
  -- a Word8 (255) to exercise the Word32 parsing of the number of bytes in each
  -- message.
  messages <- forAll (Gen.list (Range.linear 0 3) (Gen.bytes (Range.linear 1 4096)))
  let encoded = mconcat $ encodeMessage messages

  getBytes <- liftIO $ messageDecoder encoded

  decoded <- liftIO $ decodeMessage getBytes
  Right (Message messages) === decoded

messageDecoder :: ByteString -> IO (Int -> IO ByteString)
messageDecoder allBytes = do
  ref <- newIORef allBytes
  pure
    ( \nbytes -> do
        atomicModifyIORef
          ref
          ( \remainingBytes ->
              ( BS.drop nbytes remainingBytes,
                BS.take nbytes remainingBytes
              )
          )
    )