State.hs

-- |
-- Module      : Net.DNSBase.Encode.Internal.State
-- Description : Encoder state monad and wire-format primitives
-- Copyright   : (c) IIJ Innovation Institute Inc., 2009
--               (c) Viktor Dukhovni, 2020-2026
-- License     : BSD-3-Clause
-- Maintainer  : ietf-dane@dukhovni.org
-- Stability   : unstable
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE GeneralisedNewtypeDeriving #-}

module Net.DNSBase.Encode.Internal.State
    ( EncodeErr(..)
    , ErrorContext
    , SPut, SPutM, localSPut
    , buildCompressed
    , encodeCompressed
    , buildVerbatim
    , encodeVerbatim
    , putDomain
    , putWireForm
    , put8
    , put16
    , put32
    , put64
    , putIPv4
    , putIPv6
    , putByteString
    , putByteStringLen8
    , putByteStringLen16
    , putShortByteString
    , putShortByteStringLen8
    , putShortByteStringLen16
    , putUtf8Text
    , putUtf8TextLen8
    , putUtf8TextLen16
    , putSizedBuilder
    , putReplicate
    -- Encoder state mutation
    , passLen
    , failWith
    , setContext
    ) where

import qualified Control.Monad.Trans.RWS.CPS
        as R ( RWST, evalRWST, ask, get, gets, local
             , pass, put, tell )
import qualified Control.Monad.STE as STE
import qualified Control.Monad.STE.Internal as STE
import qualified Data.ByteString as B
import qualified Data.ByteString.Builder as B
import qualified Data.ByteString.Builder.Extra as B
import qualified Data.ByteString.Lazy as LB
import qualified Data.ByteString.Short as SB
import qualified Data.Text as T
import qualified Data.Text.Encoding as T
import qualified Data.Text.Unsafe as T
import GHC.ST as G (ST(..))

import qualified Net.DNSBase.Internal.NameComp as NC
import Net.DNSBase.Internal.Domain
import Net.DNSBase.Encode.Internal.Metric
import Net.DNSBase.Internal.Error
import Net.DNSBase.Internal.Util

----------------------------------------------------------------

stToSTE :: G.ST s a -> STE.STE e s a
stToSTE = coerce

----------------------------------------------------------------

-- | Encoder state, the NCTree (DNS name compression tree) is mutable in the ST
-- monad.
data EncState s = EncState
    { encOffset :: Int
    , encDoNC   :: Bool
    , encNCTree :: NC.NCTree s
    }

-- | Initial encoder state.
encInit :: Bool -- ^ If "True", DNS name compression is enabled
        -> STE.STE e s (EncState s)
encInit donamecomp = EncState 0 donamecomp <$> stToSTE (NC.empty 0)

----------------------------------------------------------------

type ErrorContext r = (Typeable r, Show r, Eq r)

buildSPut :: ErrorContext r
        => (forall s. SPut s r)
        -> Bool
        -> Either (EncodeErr (Maybe r)) (Int, Builder)
buildSPut m donc = STE.handleSTE id do
    st <- encInit donc
    evalSPutM (m >> gets encOffset) Nothing st

-- | Execute the composed 'Builder' endomorphisms to encode a packet of the
-- cumulative length.
runSPut :: ErrorContext r
        => (forall s. SPut s r)
        -> Bool
        -> Either (EncodeErr (Maybe r)) ByteString
runSPut m donc = do
    (len, builder) <- buildSPut m donc
    pure $ LB.toStrict
         $ B.toLazyByteStringWith (strat len) mempty builder
  where
    strat len = B.untrimmedStrategy len len

-- | Perform a stateful encoding with DNS name compression.  The initial error
-- context is "Nothing".  Specific values can be provided during the
-- computation by using 'localSPut'.
buildCompressed :: ErrorContext r
                => (forall s. SPut s r)
                -> Either (EncodeErr (Maybe r)) Builder
buildCompressed m = snd <$> buildSPut m True

-- | Perform a stateful encoding with DNS name compression.  The initial error
-- context is "Nothing".  Specific values can be provided during the
-- computation by using 'localSPut'.
encodeCompressed :: ErrorContext r
                 => (forall s. SPut s r)
                 -> Either (EncodeErr (Maybe r)) ByteString
encodeCompressed m = runSPut m True

-- | Perform a stateful encoding without DNS name compression.  The initial
-- error context is "Nothing".  Specific values can be provided during the
-- computation by using 'localSPut'.
buildVerbatim :: ErrorContext r
              => (forall s. SPut s r)
              -> Either (EncodeErr (Maybe r)) Builder
buildVerbatim m = snd <$> buildSPut m False

-- | Perform a stateful encoding without DNS name compression.  The initial
-- error context is "Nothing".  Specific values can be provided during the
-- computation by using 'localSPut'.
encodeVerbatim :: ErrorContext r
               => (forall s. SPut s r)
               -> Either (EncodeErr (Maybe r)) ByteString
encodeVerbatim m = runSPut m False

-- | Encode a domain with possible name compression if the entire name fits in
-- the first 16K of the output.
putDomain :: ErrorContext r => Domain -> SPut s r
putDomain domain = do
    EncState{..} <- get
    let !wlen = B.length (wireBytes domain) - 1
    if | wlen > 0 && encDoNC
       , !end <- encOffset + wlen
       , !ls <- revLabels domain
        -> do (!slen, !off) <- liftSPut . stToSTE $ NC.lookup ls encNCTree
              when (end <= MaxPtr) $
                  liftSPut . stToSTE $ NC.insert ls end encNCTree
              putCompressed domain wlen slen off
       | otherwise -> putWireForm domain
  where
    putCompressed !dom !dlen !slen !off
        | slen == 0 = putWireForm dom
        | otherwise = do
              when (slen < dlen) do
                  putByteString $ B.take (dlen - slen) $ wireBytes domain
              put16 $ toEnum $ (MaxPos - MaxPtr) + off

-- | Encode a domain name verbatim, without name compression.
putWireForm :: ErrorContext r => Domain -> SPut s r
putWireForm = encVar (SB.length . shortBytes) (B.shortByteString . shortBytes)
{-# INLINE putWireForm #-}

----------------------------------------------------------------

pattern MaxPos :: Int
pattern MaxPos = 0xffff

pattern MaxPtr :: Int
pattern MaxPtr = 0x3fff

{-# INLINE addPos #-}
addPos :: ErrorContext r => Int -> SPut s r
addPos n = do
    !s@EncState{ encOffset = pos } <- get
    let !pos' = pos + n
    when (n > MaxPos || pos' > MaxPos) do
        ask >>= liftSPut . STE.throwSTE . EncodeTooLong
    put $! s { encOffset = pos' }

{-# INLINE encFix #-}
encFix :: ErrorContext r => Int -> (a -> Builder) -> a -> SPut s r
encFix size enc a = addPos size >> tell (enc a)

{-# INLINE encVar #-}
encVar :: ErrorContext r => (a -> Int) -> (a -> Builder) -> a -> SPut s r
encVar getSize enc a = encFix (getSize a) enc a

----------------------------------------------------------------

-- | Write a single octet.
put8 :: ErrorContext r => Word8 -> SPut s r
put8 = encFix 1 B.word8
{-# INLINE put8 #-}

-- | Write a big-endian 16-bit word.
put16 :: ErrorContext r => Word16 -> SPut s r
put16 = encFix 2 B.word16BE
{-# INLINE put16 #-}

-- | Write a big-endian 32-bit word.
put32 :: ErrorContext r => Word32 -> SPut s r
put32 = encFix 4 B.word32BE
{-# INLINE put32 #-}

-- | Write a big-endian 64-bit word.
put64 :: ErrorContext r => Word64 -> SPut s r
put64 = encFix 8 B.word64BE
{-# INLINE put64 #-}

-- | Write the four octets of an IPv4 address in network order.
putIPv4 :: ErrorContext r => IPv4 -> SPut s r
putIPv4 = put32 . fromIPv4w
{-# INLINE putIPv4 #-}

-- | Write the sixteen octets of an IPv6 address in network order.
putIPv6 :: ErrorContext r => IPv6 -> SPut s r
putIPv6 ip6 =
    putSizedBuilder $! mbWord32 w0
                    <> mbWord32 w1
                    <> mbWord32 w2
                    <> mbWord32 w3
  where
    (w0, w1, w2, w3) = fromIPv6w ip6
{-# INLINE putIPv6 #-}

-- | Write the bytes of a 'ByteString' verbatim, with no length prefix.
putByteString :: ErrorContext r => ByteString -> SPut s r
putByteString b =
    unless (B.null b) $ encVar B.length B.byteString b

-- | Write the bytes of a 'ShortByteString' verbatim, with no length prefix.
putShortByteString :: ErrorContext r => ShortByteString -> SPut s r
putShortByteString b =
    unless (SB.null b) $ encVar SB.length B.shortByteString b

-- | Write a DNS /character-string/: an 8-bit length prefix followed
-- by the bytes.  Fails with 'EncodeTooLong' if the input exceeds 255 bytes.
putByteStringLen8 :: ErrorContext r => ByteString -> SPut s r
putByteStringLen8 bs@(B.length -> len) | len <= 0xff = do
    addPos (len + 1)
    tell $ B.word8 (iw8 len) <> B.byteString bs
putByteStringLen8 _ =
    failWith EncodeTooLong

-- | 'ShortByteString' counterpart of 'putByteStringLen8'.
putShortByteStringLen8 :: ErrorContext r => ShortByteString -> SPut s r
putShortByteStringLen8 bs@(SB.length -> len) | len <= 0xff = do
    addPos $ len + 1
    tell $ B.word8 (iw8 len) <> B.shortByteString bs
putShortByteStringLen8 _ = failWith EncodeTooLong

-- | Write a 16-bit-length-prefixed 'ByteString'.  Fails with
-- 'EncodeTooLong' if the input exceeds 65535 bytes.
putByteStringLen16 :: ErrorContext r => ByteString -> SPut s r
putByteStringLen16 bs@(B.length -> len) | len <= 0xffff = do
    addPos $ len + 2
    tell $ B.word16BE (iw16 len) <> B.byteString bs
putByteStringLen16 _ = failWith EncodeTooLong

-- | 'ShortByteString' counterpart of 'putByteStringLen16'.
putShortByteStringLen16 :: ErrorContext r => ShortByteString -> SPut s r
putShortByteStringLen16 bs@(SB.length -> len) | len <= 0xffff = do
    addPos $ len + 2
    tell $ B.word16BE (iw16 len) <> B.shortByteString bs
putShortByteStringLen16 _ = failWith EncodeTooLong

-- | Write the UTF-8 encoding of a 'Text' verbatim, with no
-- length prefix.
putUtf8Text :: ErrorContext r => Text -> SPut s r
putUtf8Text t =
    unless (T.null t) $ encVar T.lengthWord8 T.encodeUtf8Builder t

-- | Write a UTF-8-encoded 'Text' with an 8-bit length prefix
-- (length in /bytes/, not codepoints).  Fails with 'EncodeTooLong'
-- if the UTF-8 encoding exceeds 255 bytes.
putUtf8TextLen8 :: ErrorContext r => Text -> SPut s r
putUtf8TextLen8 t@(T.lengthWord8-> len) | len <= 0xff = do
    addPos $ len + 1
    tell $ B.word8 (iw8 len) <> T.encodeUtf8Builder t
putUtf8TextLen8 _ = failWith EncodeTooLong

-- | Write a UTF-8-encoded 'Text' with a 16-bit length prefix
-- (length in /bytes/).  Fails with 'EncodeTooLong' if the UTF-8
-- encoding exceeds 65535 bytes.
putUtf8TextLen16 :: ErrorContext r => Text -> SPut s r
putUtf8TextLen16 t@(T.lengthWord8-> len) | len <= 0xffff = do
    addPos $ len + 2
    tell $ B.word16BE (iw16 len) <> T.encodeUtf8Builder t
putUtf8TextLen16 _ = failWith EncodeTooLong

iw8 :: Int -> Word8
iw8 = fromIntegral
{-# INLINE iw8 #-}

iw16 :: Int -> Word16
iw16 = fromIntegral
{-# INLINE iw16 #-}

-- | Write @n@ copies of the byte @w@ (used for fixed-width
-- zero-padded fields).
putReplicate :: ErrorContext r => Word8 -> Word8 -> SPut s r
putReplicate n w =
    encFix (fromEnum n) B.lazyByteString $
        LB.replicate (fromIntegral n) w

-- | Write a length-tracked 'SizedBuilder' verbatim, advancing the
-- encoder offset by the builder's recorded length.  Fails with
-- 'EncodeTooLong' when the builder itself carries the overflow
-- sentinel.
putSizedBuilder :: ErrorContext r => SizedBuilder -> SPut s r
putSizedBuilder (SizedBuilder len b) = addPos len >> tell b
putSizedBuilder _                    = failWith EncodeTooLong
{-# INLINE putSizedBuilder #-}

------------------------------------------

-- | Wrap a writer with an outer 16-bit length prefix that is
-- computed automatically from the writer's output.  Used for
-- @RDLENGTH@ and the various length-prefixed sub-fields of RR
-- data (SVCB option values, EDNS option values, and so on).
passLen :: ErrorContext r => SPutM s r a -> SPutM s r a
passLen m = pass $ do
        pos <- addPos 2 >> gets encOffset
        x <- m
        len <- subtract pos <$> gets encOffset
        return (x, prependLen len)

prependLen :: Int -> B.Builder -> B.Builder
prependLen = mappend . B.word16BE . fromIntegral

-- | Abort the encoder with the given error, attaching the current
-- 'ErrorContext' from the reader environment.  Used by individual
-- writers when an out-of-range or otherwise invalid value cannot
-- be serialised.
failWith :: ErrorContext r => (forall a. ErrorContext a => a -> EncodeErr a) -> SPut s r
failWith f = ask >>= liftSPut . STE.throwSTE . f

-- | Run a sub-encoder with the given context value installed, so
-- that any 'failWith' inside reports its error against that
-- context.  Used to label sections of the encode (RR header, RR
-- data, EDNS option, ...) so that error messages identify which
-- field went wrong.
setContext :: ErrorContext r => r -> SPutM s r a -> SPutM s r a
setContext r = localSPut (const $ Just r)
{-# INLINE setContext #-}

----------------------------------------------------------------

type BaseM s r = STE.STE (EncodeErr (Maybe r)) s
type RWST s r = R.RWST (Maybe r) Builder (EncState s) (BaseM s r)

-- | Encode an output packet in the ST monad, with @r@ as an
-- optional error context (typically the RData being encoded, when
-- applicable).
type SPut s r = SPutM s r ()

-- | The underlying SPut monad
newtype SPutM s r a = SPutM { _unSPutM :: RWST s r a }

instance Functor (SPutM s r) where
    fmap :: forall a b. (a -> b) -> SPutM s r a -> SPutM s r b
    fmap = coerce (fmap :: (a -> b) -> RWST s r a -> RWST s r b)
    {-# INLINE fmap #-}

    (<$) :: forall a b. a -> SPutM s r b -> SPutM s r a
    (<$) = coerce ((<$) :: a -> RWST s r b -> RWST s r a)
    {-# INLINE (<$) #-}

instance Applicative (SPutM s r) where
    pure :: forall a. a -> SPutM s r a
    pure = coerce (pure :: a -> RWST s r a)
    {-# INLINE pure #-}

    liftA2 :: forall a b c. (a -> b -> c) -> SPutM s r a -> SPutM s r b -> SPutM s r c
    liftA2 = coerce (liftA2 :: (a -> b -> c) -> RWST s r a -> RWST s r b -> RWST s r c)
    {-# INLINE liftA2 #-}

    (<*>) :: forall a b. SPutM s r (a -> b) -> SPutM s r a -> SPutM s r b
    (<*>) = coerce ((<*>) :: RWST s r (a -> b) -> RWST s r a -> RWST s r b)
    {-# INLINE (<*>) #-}

    (*>) :: forall a b. SPutM s r a -> SPutM s r b -> SPutM s r b
    (*>) = coerce ((*>) :: RWST s r a -> RWST s r b -> RWST s r b)
    {-# INLINE (*>) #-}

    (<*) :: forall a b. SPutM s r a -> SPutM s r b -> SPutM s r a
    (<*) = coerce ((<*) :: RWST s r a -> RWST s r b -> RWST s r a)
    {-# INLINE (<*) #-}

instance Monad (SPutM s r) where
    (>>=) :: forall a b. SPutM s r a -> (a -> SPutM s r b) -> SPutM s r b
    (>>=) = coerce ((>>=) :: RWST s r a -> (a -> RWST s r b) -> RWST s r b)
    {-# INLINE (>>=) #-}

evalSPutM :: forall s r a. (forall t. SPutM t r a) -> (Maybe r) -> (EncState s) -> BaseM s r (a, Builder)
evalSPutM m = coerce (R.evalRWST (coerce m :: (forall t. RWST t r a)) :: Maybe r -> EncState s -> BaseM s r (a, Builder))
{-# INLINE evalSPutM #-}

liftSPut :: forall s r a. BaseM s r a -> SPutM s r a
liftSPut = coerce (lift :: BaseM s r a -> RWST s r a)

ask :: forall s r. SPutM s r (Maybe r)
ask = coerce (R.ask :: RWST s r (Maybe r))
{-# INLINE ask #-}

get :: forall s r. SPutM s r (EncState s)
get = coerce (R.get :: RWST s r (EncState s))
{-# INLINE get #-}

gets :: forall s r a. (EncState s -> a) -> SPutM s r a
gets = coerce (R.gets :: (EncState s -> a) -> RWST s r a)
{-# INLINE gets #-}

-- | Run the encoder with a modified context
localSPut :: forall s r a. (Maybe r -> Maybe r) -> SPutM s r a -> SPutM s r a
localSPut = coerce (R.local :: (Maybe r -> Maybe r) -> RWST s r a -> RWST s r a)
{-# INLINE localSPut #-}

pass :: forall s r a. SPutM s r (a, Builder -> Builder) -> SPutM s r a
pass = coerce (R.pass :: RWST s r (a, Builder -> Builder) -> RWST s r a)
{-# INLINE pass #-}

put :: forall s r. EncState s -> SPutM s r ()
put = coerce (R.put :: EncState s -> RWST s r ())
{-# INLINE put #-}

tell :: forall s r. Builder -> SPutM s r ()
tell = coerce (R.tell :: Builder -> RWST s r ())
{-# INLINE tell #-}