{-# LANGUAGE MagicHash #-}
{-# OPTIONS_GHC -O2 #-}
-- -O2: its loops over every byte sent run 15 to 20% faster than at -O1
-- (Bench.Encode "large params"); library-wide -O2 made row parsing slower.
-- | TDS packet framing.
--
-- [MS-TDS 2.2.3 Packets](https://learn.microsoft.com/en-us/openspecs/windows_protocols/ms-tds/f4f8a0d9-7c4a-4a36-a6f3-cd9f6f0d8a5b)
module Hmssql.Tds.Packet
  ( PacketType (..)
  , packetTypeWord
  , packetTypeFromWord
  , PacketStatus (..)
  , statusEndOfMessage
  , statusIgnore
  , statusResetConnection
  , Header (..)
  , headerLength
  , readHeader
  , getHeader
  , putHeader
  , pokePacketHeader
  , packetPayloadSize
  , splitIntoPackets
  , frameSegments
  , pokeFramedPacket
  , packetCount
  , maxSqlBatchPackets
  ) where

import Data.Binary.Get (Get)
import qualified Data.Binary.Get as Get
import Data.Binary.Put (Put)
import qualified Data.Binary.Put as Put
import Data.Bits (Bits, complement, shiftR, (.&.), (.|.))
import qualified Data.ByteString as B
import qualified Data.ByteString.Internal as BI
import qualified Data.ByteString.Lazy as LB
import Data.Int (Int64)
import Data.Word (Word16, Word32, Word8)
import Foreign.ForeignPtr (plusForeignPtr)
import Foreign.Marshal.Utils (copyBytes)
import Foreign.Ptr (Ptr, plusPtr)
import Foreign.Storable (pokeByteOff)

import Hmssql.Internal.Derived
import Hmssql.Tds.Bytes (word16BeAt, word8At)
import Hmssql.Tds.Context

-- | The type byte of a packet header.
data PacketType
  = PTSqlBatch          -- ^ 0x01
  | PTRpc               -- ^ 0x03
  | PTTabularResult     -- ^ 0x04
  | PTAttention         -- ^ 0x06
  | PTBulkLoad          -- ^ 0x07
  | PTTransactionManager -- ^ 0x0E
  | PTLogin7            -- ^ 0x10
  | PTSspi              -- ^ 0x11
  | PTPrelogin          -- ^ 0x12, also carries TLS handshake records for non-strict encryption
  | PTOther !Word8
  deriving stock (PacketType -> PacketType -> Bool
(PacketType -> PacketType -> Bool)
-> (PacketType -> PacketType -> Bool) -> Eq PacketType
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: PacketType -> PacketType -> Bool
== :: PacketType -> PacketType -> Bool
$c/= :: PacketType -> PacketType -> Bool
/= :: PacketType -> PacketType -> Bool
Eq)

instance Show PacketType where
  showsPrec :: Int -> PacketType -> ShowS
showsPrec Int
p = \case
    PacketType
PTSqlBatch -> String -> ShowS
showString String
"PTSqlBatch"
    PacketType
PTRpc -> String -> ShowS
showString String
"PTRpc"
    PacketType
PTTabularResult -> String -> ShowS
showString String
"PTTabularResult"
    PacketType
PTAttention -> String -> ShowS
showString String
"PTAttention"
    PacketType
PTBulkLoad -> String -> ShowS
showString String
"PTBulkLoad"
    PacketType
PTTransactionManager -> String -> ShowS
showString String
"PTTransactionManager"
    PacketType
PTLogin7 -> String -> ShowS
showString String
"PTLogin7"
    PacketType
PTSspi -> String -> ShowS
showString String
"PTSspi"
    PacketType
PTPrelogin -> String -> ShowS
showString String
"PTPrelogin"
    PTOther Word8
a -> Int -> Addr# -> Word8 -> ShowS
forall a. Show a => Int -> Addr# -> a -> ShowS
showCon1 Int
p Addr#
"PTOther"# Word8
a
  {-# OPAQUE showsPrec #-}

packetTypeWord :: PacketType -> Word8
packetTypeWord :: PacketType -> Word8
packetTypeWord = \case
  PacketType
PTSqlBatch -> Word8
0x01
  PacketType
PTRpc -> Word8
0x03
  PacketType
PTTabularResult -> Word8
0x04
  PacketType
PTAttention -> Word8
0x06
  PacketType
PTBulkLoad -> Word8
0x07
  PacketType
PTTransactionManager -> Word8
0x0E
  PacketType
PTLogin7 -> Word8
0x10
  PacketType
PTSspi -> Word8
0x11
  PacketType
PTPrelogin -> Word8
0x12
  PTOther Word8
w -> Word8
w

packetTypeFromWord :: Word8 -> PacketType
packetTypeFromWord :: Word8 -> PacketType
packetTypeFromWord = \case
  Word8
0x01 -> PacketType
PTSqlBatch
  Word8
0x03 -> PacketType
PTRpc
  Word8
0x04 -> PacketType
PTTabularResult
  Word8
0x06 -> PacketType
PTAttention
  Word8
0x07 -> PacketType
PTBulkLoad
  Word8
0x0E -> PacketType
PTTransactionManager
  Word8
0x10 -> PacketType
PTLogin7
  Word8
0x11 -> PacketType
PTSspi
  Word8
0x12 -> PacketType
PTPrelogin
  Word8
w -> Word8 -> PacketType
PTOther Word8
w

-- | The status byte of a packet header; 'statusEndOfMessage', 'statusIgnore'
-- and 'statusResetConnection' are its bits.
newtype PacketStatus = PacketStatus Word8
  deriving newtype (PacketStatus -> PacketStatus -> Bool
(PacketStatus -> PacketStatus -> Bool)
-> (PacketStatus -> PacketStatus -> Bool) -> Eq PacketStatus
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: PacketStatus -> PacketStatus -> Bool
== :: PacketStatus -> PacketStatus -> Bool
$c/= :: PacketStatus -> PacketStatus -> Bool
/= :: PacketStatus -> PacketStatus -> Bool
Eq, Eq PacketStatus
PacketStatus
Eq PacketStatus =>
(PacketStatus -> PacketStatus -> PacketStatus)
-> (PacketStatus -> PacketStatus -> PacketStatus)
-> (PacketStatus -> PacketStatus -> PacketStatus)
-> (PacketStatus -> PacketStatus)
-> (PacketStatus -> Int -> PacketStatus)
-> (PacketStatus -> Int -> PacketStatus)
-> PacketStatus
-> (Int -> PacketStatus)
-> (PacketStatus -> Int -> PacketStatus)
-> (PacketStatus -> Int -> PacketStatus)
-> (PacketStatus -> Int -> PacketStatus)
-> (PacketStatus -> Int -> Bool)
-> (PacketStatus -> Maybe Int)
-> (PacketStatus -> Int)
-> (PacketStatus -> Bool)
-> (PacketStatus -> Int -> PacketStatus)
-> (PacketStatus -> Int -> PacketStatus)
-> (PacketStatus -> Int -> PacketStatus)
-> (PacketStatus -> Int -> PacketStatus)
-> (PacketStatus -> Int -> PacketStatus)
-> (PacketStatus -> Int -> PacketStatus)
-> (PacketStatus -> Int)
-> Bits PacketStatus
Int -> PacketStatus
PacketStatus -> Bool
PacketStatus -> Int
PacketStatus -> Maybe Int
PacketStatus -> PacketStatus
PacketStatus -> Int -> Bool
PacketStatus -> Int -> PacketStatus
PacketStatus -> PacketStatus -> PacketStatus
forall a.
Eq a =>
(a -> a -> a)
-> (a -> a -> a)
-> (a -> a -> a)
-> (a -> a)
-> (a -> Int -> a)
-> (a -> Int -> a)
-> a
-> (Int -> a)
-> (a -> Int -> a)
-> (a -> Int -> a)
-> (a -> Int -> a)
-> (a -> Int -> Bool)
-> (a -> Maybe Int)
-> (a -> Int)
-> (a -> Bool)
-> (a -> Int -> a)
-> (a -> Int -> a)
-> (a -> Int -> a)
-> (a -> Int -> a)
-> (a -> Int -> a)
-> (a -> Int -> a)
-> (a -> Int)
-> Bits a
$c.&. :: PacketStatus -> PacketStatus -> PacketStatus
.&. :: PacketStatus -> PacketStatus -> PacketStatus
$c.|. :: PacketStatus -> PacketStatus -> PacketStatus
.|. :: PacketStatus -> PacketStatus -> PacketStatus
$cxor :: PacketStatus -> PacketStatus -> PacketStatus
xor :: PacketStatus -> PacketStatus -> PacketStatus
$ccomplement :: PacketStatus -> PacketStatus
complement :: PacketStatus -> PacketStatus
$cshift :: PacketStatus -> Int -> PacketStatus
shift :: PacketStatus -> Int -> PacketStatus
$crotate :: PacketStatus -> Int -> PacketStatus
rotate :: PacketStatus -> Int -> PacketStatus
$czeroBits :: PacketStatus
zeroBits :: PacketStatus
$cbit :: Int -> PacketStatus
bit :: Int -> PacketStatus
$csetBit :: PacketStatus -> Int -> PacketStatus
setBit :: PacketStatus -> Int -> PacketStatus
$cclearBit :: PacketStatus -> Int -> PacketStatus
clearBit :: PacketStatus -> Int -> PacketStatus
$ccomplementBit :: PacketStatus -> Int -> PacketStatus
complementBit :: PacketStatus -> Int -> PacketStatus
$ctestBit :: PacketStatus -> Int -> Bool
testBit :: PacketStatus -> Int -> Bool
$cbitSizeMaybe :: PacketStatus -> Maybe Int
bitSizeMaybe :: PacketStatus -> Maybe Int
$cbitSize :: PacketStatus -> Int
bitSize :: PacketStatus -> Int
$cisSigned :: PacketStatus -> Bool
isSigned :: PacketStatus -> Bool
$cshiftL :: PacketStatus -> Int -> PacketStatus
shiftL :: PacketStatus -> Int -> PacketStatus
$cunsafeShiftL :: PacketStatus -> Int -> PacketStatus
unsafeShiftL :: PacketStatus -> Int -> PacketStatus
$cshiftR :: PacketStatus -> Int -> PacketStatus
shiftR :: PacketStatus -> Int -> PacketStatus
$cunsafeShiftR :: PacketStatus -> Int -> PacketStatus
unsafeShiftR :: PacketStatus -> Int -> PacketStatus
$crotateL :: PacketStatus -> Int -> PacketStatus
rotateL :: PacketStatus -> Int -> PacketStatus
$crotateR :: PacketStatus -> Int -> PacketStatus
rotateR :: PacketStatus -> Int -> PacketStatus
$cpopCount :: PacketStatus -> Int
popCount :: PacketStatus -> Int
Bits)

instance Show PacketStatus where
  showsPrec :: Int -> PacketStatus -> ShowS
showsPrec Int
p (PacketStatus Word8
a) = Int -> Addr# -> Word8 -> ShowS
forall a. Show a => Int -> Addr# -> a -> ShowS
showCon1 Int
p Addr#
"PacketStatus"# Word8
a
  {-# OPAQUE showsPrec #-}

statusEndOfMessage, statusIgnore, statusResetConnection :: PacketStatus
statusEndOfMessage :: PacketStatus
statusEndOfMessage = Word8 -> PacketStatus
PacketStatus Word8
0x01
statusIgnore :: PacketStatus
statusIgnore = Word8 -> PacketStatus
PacketStatus Word8
0x02
statusResetConnection :: PacketStatus
statusResetConnection = Word8 -> PacketStatus
PacketStatus Word8
0x08

-- | A packet header.
data Header = Header
  { Header -> PacketType
headerType :: !PacketType
  , Header -> PacketStatus
headerStatus :: !PacketStatus
  , Header -> Word16
headerPacketLength :: !Word16
    -- ^ Total packet length including the header and payload, in bytes.
  , Header -> Word16
headerSpid :: !Word16
  , Header -> Word8
headerPacketId :: !Word8
  , Header -> Word8
headerWindow :: !Word8
  }

instance Show Header where
  showsPrec :: Int -> Header -> ShowS
showsPrec Int
p (Header PacketType
a PacketStatus
b Word16
c Word16
e Word8
f Word8
g) = Int -> Addr# -> [Field] -> ShowS
showRecord Int
p Addr#
"Header"#
    [ Addr#
"headerType"# Addr# -> PacketType -> Field
forall a. Show a => Addr# -> a -> Field
:= PacketType
a, Addr#
"headerStatus"# Addr# -> PacketStatus -> Field
forall a. Show a => Addr# -> a -> Field
:= PacketStatus
b, Addr#
"headerPacketLength"# Addr# -> Word16 -> Field
forall a. Show a => Addr# -> a -> Field
:= Word16
c, Addr#
"headerSpid"# Addr# -> Word16 -> Field
forall a. Show a => Addr# -> a -> Field
:= Word16
e
    , Addr#
"headerPacketId"# Addr# -> Word8 -> Field
forall a. Show a => Addr# -> a -> Field
:= Word8
f, Addr#
"headerWindow"# Addr# -> Word8 -> Field
forall a. Show a => Addr# -> a -> Field
:= Word8
g
    ]
  {-# OPAQUE showsPrec #-}

instance Eq Header where
  Header PacketType
a PacketStatus
b Word16
c Word16
e Word8
f Word8
g == :: Header -> Header -> Bool
== Header PacketType
a' PacketStatus
b' Word16
c' Word16
e' Word8
f' Word8
g' =
    PacketType -> PacketType -> Bool
forall a. Eq a => a -> a -> Bool
eq PacketType
a PacketType
a' Bool -> Bool -> Bool
&& PacketStatus -> PacketStatus -> Bool
forall a. Eq a => a -> a -> Bool
eq PacketStatus
b PacketStatus
b' Bool -> Bool -> Bool
&& Word16
c Word16 -> Word16 -> Bool
forall a. Eq a => a -> a -> Bool
== Word16
c' Bool -> Bool -> Bool
&& Word16
e Word16 -> Word16 -> Bool
forall a. Eq a => a -> a -> Bool
== Word16
e' Bool -> Bool -> Bool
&& Word8
f Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Word8
f' Bool -> Bool -> Bool
&& Word8
g Word8 -> Word8 -> Bool
forall a. Eq a => a -> a -> Bool
== Word8
g'
  {-# OPAQUE (==) #-}

-- | The fixed size of a TDS packet header, in bytes. The total packet
-- length is the separate 'headerPacketLength' field.
headerLength :: Int
headerLength :: Int
headerLength = Int
8

-- | The header at the start of a buffer of at least 'headerLength' bytes,
-- read in place: one per packet received ('error' on a shorter buffer).
readHeader :: B.ByteString -> Header
readHeader :: ByteString -> Header
readHeader ByteString
bs = Header
  { headerType :: PacketType
headerType = Word8 -> PacketType
packetTypeFromWord (ByteString -> Int -> Word8
word8At ByteString
bs Int
0)
  , headerStatus :: PacketStatus
headerStatus = Word8 -> PacketStatus
PacketStatus (ByteString -> Int -> Word8
word8At ByteString
bs Int
1)
  , headerPacketLength :: Word16
headerPacketLength = ByteString -> Int -> Word16
word16BeAt ByteString
bs Int
2
  , headerSpid :: Word16
headerSpid = ByteString -> Int -> Word16
word16BeAt ByteString
bs Int
4
  , headerPacketId :: Word8
headerPacketId = ByteString -> Int -> Word8
word8At ByteString
bs Int
6
  , headerWindow :: Word8
headerWindow = ByteString -> Int -> Word8
word8At ByteString
bs Int
7
  }

-- | 'readHeader' as a 'Get' (the tests' mock server).
getHeader :: Get Header
getHeader :: Get Header
getHeader = ByteString -> Header
readHeader (ByteString -> Header) -> Get ByteString -> Get Header
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> Get ByteString
Get.getByteString Int
headerLength

putHeader :: Header -> Put
putHeader :: Header -> Put
putHeader Header{Word8
Word16
PacketStatus
PacketType
headerType :: Header -> PacketType
headerStatus :: Header -> PacketStatus
headerPacketLength :: Header -> Word16
headerSpid :: Header -> Word16
headerPacketId :: Header -> Word8
headerWindow :: Header -> Word8
headerType :: PacketType
headerStatus :: PacketStatus
headerPacketLength :: Word16
headerSpid :: Word16
headerPacketId :: Word8
headerWindow :: Word8
..} = do
  Word8 -> Put
Put.putWord8 (Word8 -> Put) -> Word8 -> Put
forall a b. (a -> b) -> a -> b
$ PacketType -> Word8
packetTypeWord PacketType
headerType
  let PacketStatus Word8
st = PacketStatus
headerStatus
  Word8 -> Put
Put.putWord8 Word8
st
  Word16 -> Put
Put.putWord16be Word16
headerPacketLength
  Word16 -> Put
Put.putWord16be Word16
headerSpid
  Word8 -> Put
Put.putWord8 Word8
headerPacketId
  Word8 -> Put
Put.putWord8 Word8
headerWindow

-- | The body bytes each packet of 'splitIntoPackets' carries: the
-- negotiated packet size less the header (clamped as 'frameSegments'
-- clamps it).
packetPayloadSize :: TdsContext -> Int
packetPayloadSize :: TdsContext -> Int
packetPayloadSize TdsContext
ctx = Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
1 (Word32 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Word32 -> Word32 -> Word32
forall a. Ord a => a -> a -> a
min Word32
maxPacketSize (TdsContext -> Word32
tdsPacketSize TdsContext
ctx)) Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
headerLength)

-- | Write a packet header at the pointer: the type, the status (as given;
-- the caller sets 'statusEndOfMessage' on the last packet), the total
-- length (header included, at most 'maxPacketSize') and the packet ID.
-- For a message framed packet by packet as it is produced
-- ("Hmssql.Tds.PacketWriter"); 'frameSegments' frames a complete body.
pokePacketHeader :: Ptr Word8 -> PacketType -> PacketStatus -> Int -> Word8 -> IO ()
pokePacketHeader :: Ptr Word8 -> PacketType -> PacketStatus -> Int -> Word8 -> IO ()
pokePacketHeader Ptr Word8
ptr PacketType
ptype (PacketStatus Word8
st) Int
len Word8
pid = do
  Ptr Word8 -> Int -> Word8 -> IO ()
forall b. Ptr b -> Int -> Word8 -> IO ()
forall a b. Storable a => Ptr b -> Int -> a -> IO ()
pokeByteOff Ptr Word8
ptr Int
0 (PacketType -> Word8
packetTypeWord PacketType
ptype)
  Ptr Word8 -> Int -> Word8 -> IO ()
forall b. Ptr b -> Int -> Word8 -> IO ()
forall a b. Storable a => Ptr b -> Int -> a -> IO ()
pokeByteOff Ptr Word8
ptr Int
1 Word8
st
  Ptr Word8 -> Int -> Word8 -> IO ()
forall b. Ptr b -> Int -> Word8 -> IO ()
forall a b. Storable a => Ptr b -> Int -> a -> IO ()
pokeByteOff Ptr Word8
ptr Int
2 (Int -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
len Int -> Int -> Int
forall a. Bits a => a -> Int -> a
`shiftR` Int
8) :: Word8)
  Ptr Word8 -> Int -> Word8 -> IO ()
forall b. Ptr b -> Int -> Word8 -> IO ()
forall a b. Storable a => Ptr b -> Int -> a -> IO ()
pokeByteOff Ptr Word8
ptr Int
3 (Int -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
len :: Word8)
  Ptr Word8 -> Int -> Word32 -> IO ()
forall b. Ptr b -> Int -> Word32 -> IO ()
forall a b. Storable a => Ptr b -> Int -> a -> IO ()
pokeByteOff Ptr Word8
ptr Int
4 (Word32
0 :: Word32)  -- SPID (2 bytes), packet ID, window
  Ptr Word8 -> Int -> Word8 -> IO ()
forall b. Ptr b -> Int -> Word8 -> IO ()
forall a b. Storable a => Ptr b -> Int -> a -> IO ()
pokeByteOff Ptr Word8
ptr Int
6 Word8
pid

-- | How many packets 'splitIntoPackets' frames a message body into.
packetCount :: TdsContext -> LB.ByteString -> Int
packetCount :: TdsContext -> ByteString -> Int
packetCount TdsContext
ctx ByteString
body = Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
1 (Int64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral ((ByteString -> Int64
LB.length ByteString
body Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
+ Int64
n Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
- Int64
1) Int64 -> Int64 -> Int64
forall a. Integral a => a -> a -> a
`div` Int64
n))
  where n :: Int64
n = Int -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (TdsContext -> Int
packetPayloadSize TdsContext
ctx)

-- | The most packets SQL Server accepts in one SQLBatch message: it closes
-- the connection, without an error message, on a longer one.  (Hence the
-- documented batch size limit of "65,536 * network packet size"; SQL
-- Server 2025 accepts exactly 65535 full packets, at every packet size and
-- TDS version.  RPC messages have no such limit.)
maxSqlBatchPackets :: Int
maxSqlBatchPackets :: Int
maxSqlBatchPackets = Int
65535

-- | Frame a complete message body into packets of the negotiated size, one
-- strict chunk per packet.  ('frameSegments' with one packet per segment.)
--
-- The extra status bits go on every packet, except 'statusResetConnection',
-- which goes on the first packet only (MS-TDS 2.2.3.1.2), and EOM, which is
-- set on the last packet here.
--
-- Body lengths and the framed allocation sizes must fit in 'Int'.
-- A packet size outside what the header can describe is clamped: to at
-- least one byte of payload, so the split always makes progress, and to at
-- most 'maxPacketSize', so the 16-bit length field never truncates.  (The
-- connection never negotiates such a size; the clamp prevents a zero-size
-- split or a truncated packet header.)
splitIntoPackets :: TdsContext -> PacketType -> PacketStatus -> LB.ByteString -> [B.ByteString]
splitIntoPackets :: TdsContext
-> PacketType -> PacketStatus -> ByteString -> [ByteString]
splitIntoPackets = Int
-> TdsContext
-> PacketType
-> PacketStatus
-> ByteString
-> [ByteString]
frameSegments Int
0

-- | Frame a complete message body as 'splitIntoPackets' does, into strict
-- chunks of as many whole packets as fit in the given number of bytes (at
-- least one packet each). The message is the chunks concatenated. Body
-- lengths and framed allocation sizes must fit in 'Int'. No uninitialized
-- prefix is reserved by this entry point.
--
-- Each byte of the body is copied once, into a chunk of exactly its size;
-- the chunks are produced lazily, so a message sent as it is framed needs
-- the body and one chunk at a time.  The connection sends a chunk at a
-- time, sized to fill a TLS record, or a larger socket write.
frameSegments :: Int -> TdsContext -> PacketType -> PacketStatus -> LB.ByteString -> [B.ByteString]
frameSegments :: Int
-> TdsContext
-> PacketType
-> PacketStatus
-> ByteString
-> [ByteString]
frameSegments Int
target TdsContext
ctx PacketType
ptype PacketStatus
extraStatus ByteString
body0 = Bool -> Word8 -> Int64 -> ByteString -> [ByteString]
go Bool
True Word8
0 Int64
total0 ByteString
body0
  where
    total0 :: Int64
total0 = ByteString -> Int64
LB.length ByteString
body0
    payloadSize :: Int
payloadSize = TdsContext -> Int
packetPayloadSize TdsContext
ctx
    perSegment :: Int
perSegment = Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
1 (Int
target Int -> Int -> Int
forall a. Integral a => a -> a -> a
`div` (Int
payloadSize Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
headerLength))

    -- Requests are framed by "Hmssql.Tds.PacketWriter" instead, whose first
    -- packet ID is 1; SQL Server accepts either.  This remains for what is framed
    -- from memory, and as the tests' independent reference.
    go :: Bool -> Word8 -> Int64 -> LB.ByteString -> [B.ByteString]
    go :: Bool -> Word8 -> Int64 -> ByteString -> [ByteString]
go Bool
isFirst !Word8
pid !Int64
remaining ByteString
body =
      let full :: Int
full = Int64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int64
remaining Int64 -> Int64 -> Int64
forall a. Integral a => a -> a -> a
`div` Int -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
payloadSize)
          partial :: Int
partial = if Int64
remaining Int64 -> Int64 -> Int64
forall a. Integral a => a -> a -> a
`mod` Int -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
payloadSize Int64 -> Int64 -> Bool
forall a. Eq a => a -> a -> Bool
== Int64
0 then Int
0 else Int
1
          -- An empty message is one packet with no payload.
          left :: Int
left = Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
1 (Int
full Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
partial) :: Int
          n :: Int
n = Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
perSegment Int
left
          segPayload :: Int64
segPayload = Int64 -> Int64 -> Int64
forall a. Ord a => a -> a -> a
min Int64
remaining (Int -> Int64
forall a b. (Integral a, Num b) => a -> b
fromIntegral (Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
payloadSize))
          (ByteString
here, ByteString
rest) = Int64 -> ByteString -> (ByteString, ByteString)
LB.splitAt Int64
segPayload ByteString
body
          isLastSegment :: Bool
isLastSegment = Int
n Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
left
          !seg :: ByteString
seg = Int -> (Ptr Word8 -> IO ()) -> ByteString
BI.unsafeCreate (Int64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int64
segPayload Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
* Int
headerLength) ((Ptr Word8 -> IO ()) -> ByteString)
-> (Ptr Word8 -> IO ()) -> ByteString
forall a b. (a -> b) -> a -> b
$ \Ptr Word8
ptr ->
            Ptr Word8
-> Bool -> Word8 -> Int -> Bool -> Int -> [ByteString] -> IO ()
fill Ptr Word8
ptr Bool
isFirst Word8
pid Int
n Bool
isLastSegment (Int64 -> Int
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int64
segPayload) (ByteString -> [ByteString]
LB.toChunks ByteString
here)
       in if Bool
isLastSegment then [ByteString
seg] else ByteString
seg ByteString -> [ByteString] -> [ByteString]
forall a. a -> [a] -> [a]
: Bool -> Word8 -> Int64 -> ByteString -> [ByteString]
go Bool
False (Word8
pid Word8 -> Word8 -> Word8
forall a. Num a => a -> a -> a
+ Int -> Word8
forall a b. (Integral a, Num b) => a -> b
fromIntegral Int
n) (Int64
remaining Int64 -> Int64 -> Int64
forall a. Num a => a -> a -> a
- Int64
segPayload) ByteString
rest

    -- Write @n@ packets holding @len@ bytes of payload, taken from the chunks.
    fill :: Ptr Word8 -> Bool -> Word8 -> Int -> Bool -> Int -> [B.ByteString] -> IO ()
    fill :: Ptr Word8
-> Bool -> Word8 -> Int -> Bool -> Int -> [ByteString] -> IO ()
fill Ptr Word8
ptr0 Bool
isFirst0 Word8
pid0 Int
n Bool
isLastSegment Int
len0 [ByteString]
chunks0 = Ptr Word8 -> Bool -> Word8 -> Int -> Int -> [ByteString] -> IO ()
packets Ptr Word8
ptr0 Bool
isFirst0 Word8
pid0 Int
0 Int
len0 [ByteString]
chunks0
      where
        packets :: Ptr Word8 -> Bool -> Word8 -> Int -> Int -> [ByteString] -> IO ()
packets !Ptr Word8
ptr !Bool
isFirst !Word8
pid !Int
i !Int
len [ByteString]
chunks
          | Int
i Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n = () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
          | Bool
otherwise = do
              let size :: Int
size = Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
payloadSize Int
len
                  isLast :: Bool
isLast = Bool
isLastSegment Bool -> Bool -> Bool
&& Int
i Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
n Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
1
              chunks' <- Ptr Word8
-> PacketType
-> PacketStatus
-> Bool
-> Word8
-> Bool
-> Int
-> [ByteString]
-> IO [ByteString]
pokeFramedPacket Ptr Word8
ptr PacketType
ptype PacketStatus
extraStatus Bool
isFirst Word8
pid Bool
isLast Int
size [ByteString]
chunks
              packets (ptr `plusPtr` (headerLength + size)) False (pid + 1) (i + 1) (len - size) chunks'

-- | Write a whole packet of a message body at the pointer: its header
-- (with EOM if it is the last, and the extra status bits as
-- 'splitIntoPackets' describes them) and @size@ bytes of payload from the
-- front of the chunks; the chunks left.
pokeFramedPacket :: Ptr Word8 -> PacketType -> PacketStatus -> Bool -> Word8 -> Bool -> Int -> [B.ByteString] -> IO [B.ByteString]
pokeFramedPacket :: Ptr Word8
-> PacketType
-> PacketStatus
-> Bool
-> Word8
-> Bool
-> Int
-> [ByteString]
-> IO [ByteString]
pokeFramedPacket Ptr Word8
ptr PacketType
ptype PacketStatus
extraStatus Bool
isFirst Word8
pid Bool
isLast Int
size [ByteString]
chunks = do
  let eom :: PacketStatus
eom = if Bool
isLast then PacketStatus
statusEndOfMessage else Word8 -> PacketStatus
PacketStatus Word8
0
      PacketStatus Word8
reset = PacketStatus
statusResetConnection
      mask :: Word8
mask = if Bool
isFirst then Word8
0xFE else Word8
0xFE Word8 -> Word8 -> Word8
forall a. Bits a => a -> a -> a
.&. Word8 -> Word8
forall a. Bits a => a -> a
complement Word8
reset
  Ptr Word8 -> PacketType -> PacketStatus -> Int -> Word8 -> IO ()
pokePacketHeader Ptr Word8
ptr PacketType
ptype (PacketStatus
eom PacketStatus -> PacketStatus -> PacketStatus
forall a. Bits a => a -> a -> a
.|. (PacketStatus
extraStatus PacketStatus -> PacketStatus -> PacketStatus
forall a. Bits a => a -> a -> a
.&. Word8 -> PacketStatus
PacketStatus Word8
mask)) (Int
size Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
headerLength) Word8
pid
  Ptr Word8 -> Int -> [ByteString] -> IO [ByteString]
copyFrom (Ptr Word8
ptr Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
headerLength) Int
size [ByteString]
chunks
{-# INLINE pokeFramedPacket #-}

-- | Copy @k@ bytes from the front of the chunks; the chunks left.
copyFrom :: Ptr Word8 -> Int -> [B.ByteString] -> IO [B.ByteString]
copyFrom :: Ptr Word8 -> Int -> [ByteString] -> IO [ByteString]
copyFrom Ptr Word8
_ Int
0 [ByteString]
chunks = [ByteString] -> IO [ByteString]
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [ByteString]
chunks
copyFrom Ptr Word8
dst Int
k (BI.BS ForeignPtr Word8
fp Int
clen : [ByteString]
cs) = do
  let m :: Int
m = Int -> Int -> Int
forall a. Ord a => a -> a -> a
min Int
k Int
clen
  ForeignPtr Word8 -> (Ptr Word8 -> IO ()) -> IO ()
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
BI.unsafeWithForeignPtr ForeignPtr Word8
fp ((Ptr Word8 -> IO ()) -> IO ()) -> (Ptr Word8 -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Ptr Word8
src -> Ptr Word8 -> Ptr Word8 -> Int -> IO ()
forall a. Ptr a -> Ptr a -> Int -> IO ()
copyBytes Ptr Word8
dst Ptr Word8
src Int
m  -- a bounded memcpy
  if Int
m Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
clen
    then Ptr Word8 -> Int -> [ByteString] -> IO [ByteString]
copyFrom (Ptr Word8
dst Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
m) (Int
k Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
m) [ByteString]
cs
    else [ByteString] -> IO [ByteString]
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ForeignPtr Word8 -> Int -> ByteString
BI.BS (ForeignPtr Word8
fp ForeignPtr Word8 -> Int -> ForeignPtr Word8
forall a b. ForeignPtr a -> Int -> ForeignPtr b
`plusForeignPtr` Int
m) (Int
clen Int -> Int -> Int
forall a. Num a => a -> a -> a
- Int
m) ByteString -> [ByteString] -> [ByteString]
forall a. a -> [a] -> [a]
: [ByteString]
cs)
copyFrom Ptr Word8
_ Int
_ [] = [ByteString] -> IO [ByteString]
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure []  -- can't happen: the chunks hold the payload