diff --git a/CHANGELOG.md b/CHANGELOG.md --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,14 @@ # Changelog +## 1.0.2.3 + +* Code cleanup and performance tweaks in encoder and decoder, + mostly around domain name compression. +* Fixed corner case in ad hoc "Mbox" presentation form with + single label mailbox names with interior dots, these are + now properly escaped. +* Fixed test suite type inference failure with GHC 9.10.3 + ## 1.0.2.2 * Changed default-language to GHC2021, GHC2024 forces newer cabal. diff --git a/dnsbase.cabal b/dnsbase.cabal --- a/dnsbase.cabal +++ b/dnsbase.cabal @@ -1,6 +1,6 @@ cabal-version: 3.0 name: dnsbase -version: 1.0.2.2 +version: 1.0.2.3 build-type: Simple synopsis: Stub DNS resolver with a typed RData model and value-based extension API @@ -53,6 +53,7 @@ category: Network tested-with: GHC == 9.10.3 , GHC == 9.12.3 + , GHC == 9.14.1 extra-source-files: CHANGELOG.md README.md @@ -83,6 +84,7 @@ , crypton >=0.30 && <1.2 , containers >=0.6 && <0.9 , deepseq ^>=1.5 + , hashable ^>=1.5 , hashtables >=1.2 && <1.6 , hourglass ^>=0.2.12 , iproute ^>=1.7.9 @@ -92,7 +94,7 @@ , primitive >=0.8 && <0.10 , template-haskell >=2.22 && <2.25 , text >=2.0 && <2.2 - , time >=1.11 && <1.16 + , time >=1.11 && <1.17 , transformers >=0.5 && <2.7 , unordered-containers ^>=0.2 diff --git a/internal/Net/DNSBase/Decode/Internal/Domain.hs b/internal/Net/DNSBase/Decode/Internal/Domain.hs --- a/internal/Net/DNSBase/Decode/Internal/Domain.hs +++ b/internal/Net/DNSBase/Decode/Internal/Domain.hs @@ -11,88 +11,155 @@ ) where import qualified Data.ByteString as B -import qualified Data.ByteString.Builder as B +import qualified Data.ByteString.Unsafe as BU +import qualified Data.Primitive.ByteArray as A +import Control.Monad.ST (ST, runST) +import Data.ByteString.Internal (ByteString(..)) +import Data.Primitive.ByteArray (MutableByteArray) import Net.DNSBase.Decode.Internal.State import Net.DNSBase.Internal.Domain import Net.DNSBase.Internal.Util --- | Wire form length limit, sans final empty root label. +-- | Wire-form length limit, including the terminal empty root label. maxWireLen :: Int maxWireLen = 255 --- | Parse a wire-form domain with \"No Compression\" (i.e. treat pointer labels --- as invalid) +-- | Initial allocation for the output buffer. Most names fit, so +-- typical decodes avoid resizing; the few that exceed it grow by +-- @cap + cap \`shiftR\` 1@ (1.5x) up to 'maxWireLen'. +initCap :: Int +initCap = 32 + +-- | Parse a wire-form domain with \"No Compression\" (i.e. treat pointer +-- labels as invalid). -- -- When defining the decoders for newly standardized RData types, it is --- generally required to use this function to decode transparent domain fields, --- as name compression is explicitly forbidden for domain fields of future RData --- types (see 'getDomain' for reference) +-- generally required to use this function to decode transparent domain +-- fields, as name compression is explicitly forbidden for domain fields +-- of future RData types (see 'getDomain' for reference) getDomainNC :: SGet Domain -getDomainNC = do - (_, bldr) <- getDomain' False =<< getPosition - case buildDomain (Just bldr) of - Just dom -> pure dom - Nothing -> failSGet "Internal error" +getDomainNC = decodeAt False --- | Parse a wire-form domain with name compression (pointer labels) allowed +-- | Parse a wire-form domain with name compression (pointer labels) +-- allowed. -- --- This function should only be used when decoding the owner name of resource --- records, as well as for fields of the initial set of RData types defined in --- [RFC 1035](https://tools.ietf.org/html/rfc1035) and several others listed in --- section 4 of [RFC 3597](https://tools.ietf.org/html/rfc3597#section-4), --- which also states that future RData types MUST NOT use name compression +-- This function should only be used when decoding the owner name of +-- resource records, as well as for fields of the initial set of RData +-- types defined in [RFC 1035](https://tools.ietf.org/html/rfc1035) and +-- several others listed in section 4 of +-- [RFC 3597](https://tools.ietf.org/html/rfc3597#section-4), which also +-- states that future RData types MUST NOT use name compression. getDomain :: SGet Domain -getDomain = do - -- No name (de)compression if the input is only a message fragment. - nc <- getNameComp - (_, bldr) <- getDomain' nc =<< getPosition - case buildDomain (Just bldr) of - Just dom -> pure dom - Nothing -> failSGet "Internal error" +getDomain = getNameComp >>= decodeAt --- | First octet of a label determines the interpretation of the rest of the label; --- 11XX_XXXX indicates a 14-bit compression pointer composed of the low 6 bits of --- that octet and the entirety of the next octet, while 00XX_XXXX is used for a --- standard label to encode its length (<=63). 01XX_XXXX was proposed for extended --- labels but remains experimental, and 10XX_XXXX is presently undefined. Latest --- status can be found in --- [IANA registry](https://www.iana.org/assignments/dns-parameters/dns-parameters.xhtml#dns-parameters-10) -getDomain' :: Bool -> Int -> SGet (Int, B.Builder) -getDomain' allowPtr start = do - vl <- get8 - if | vl == 0 -> do - end <- getPosition - -- Including the terminal empty label length byte - getSlice start (end+1) - | vl <= 63 -> do - let len = fromIntegral vl - skipNBytes len - getDomain' allowPtr start - | vl >= 0b1100_0000 -> do - unless allowPtr $ failSGet "domain name compression not allowed in current context" - end <- getPosition - (plen, prefix) <- getSlice start end - vl' <- get8 - let offset :: Word16 - offset = (fromIntegral (vl .&. 0x3f) `shiftL` 8) .|. (fromIntegral vl') - when (fromIntegral offset >= start) $ failSGet "invalid compression pointer" - (slen, suffix) <- getPtr offset - let len = plen + slen - when (len > maxWireLen) do - failSGet "domain name too long" - return $ (len, prefix <> suffix) - | otherwise -> failSGet "unsupported label type" +-- The shared body of 'getDomain' and 'getDomainNC'. Snapshots the +-- input packet and the SGet cursor, runs the pointer-following walk in +-- 'ST' over a growable 'MutableByteArray', and advances the SGet +-- cursor by however many bytes were consumed at the starting position +-- (the prefix labels plus either the terminal NUL byte or the 2-byte +-- pointer label that ended the inline portion). +decodeAt :: Bool -> SGet Domain +decodeAt allowPtr = do + pkt <- getPacket + start <- getPosition + case runST (decodeDomain allowPtr pkt start) of + Right (consumed, dom) -> dom <$ skipNBytes consumed + Left msg -> failSGet msg + +-- The wire-form first octet of a label determines the interpretation +-- of the rest of the label. @11XX_XXXX@ indicates a 14-bit compression +-- pointer composed of the low 6 bits of that octet and the entirety of +-- the next octet, while @00XX_XXXX@ encodes the length (<=63) of a +-- standard label. @01XX_XXXX@ was proposed for extended labels but +-- remains experimental, and @10XX_XXXX@ is presently undefined. +-- Current status is in the +-- [IANA registry](https://www.iana.org/assignments/dns-parameters/dns-parameters.xhtml#dns-parameters-10). +-- +-- The walk reads labels at @pos@ from the packet, copying their wire +-- bytes (including each leading length byte) into the growable output +-- buffer. A pointer label switches reading to its target offset +-- without writing anything itself; the @firstPtr@ position remembers +-- the absolute offset of the /first/ pointer encountered, so the SGet +-- cursor advance at the end can charge the caller for the prefix plus +-- the 2-byte pointer label, not for any bytes read via the pointer. +-- +-- Pointer cycles are prevented by the strict-decreasing @minPtr@ +-- invariant: each pointer target must be strictly less than the +-- previous one (initially the starting offset). Since pointer offsets +-- are non-negative and bounded, the chain terminates. +decodeDomain :: forall s. Bool -> ByteString -> Int + -> ST s (Either String (Int, Domain)) +decodeDomain allowPtr pkt@(BS fp _) !start = do + buf <- A.newByteArray initCap + runExceptT $ walk start start initCap buf 0 Nothing where - getPtr :: Word16 -> SGet (Int, B.Builder) - getPtr off = seekSGet off $ getPosition >>= getDomain' allowPtr + !pktLen = B.length pkt - -- get a bytestring slice from position i to position j-1 - getSlice :: Int -> Int -> SGet (Int, B.Builder) - getSlice off ((subtract (off + 1)) -> len) - | len < 0 = failSGet "negative-length domain name slice" - | len > maxWireLen = failSGet "domain name too long" - | otherwise = do - buf <- getPacket - let slice = B.take len $ B.drop off buf - return $ (len, B.byteString slice) + walk :: Int -- ^ current read offset in packet + -> Int -- ^ strict-decreasing pointer bound + -> Int -- ^ current buffer capacity + -> MutableByteArray s -- ^ output buffer + -> Int -- ^ current write position in buffer + -> Maybe Int -- ^ packet offset of first pointer, if any + -> ExceptT String (ST s) (Int, Domain) + walk !pos !minPtr !cap !buf !cursor !firstPtr = do + when (pos >= pktLen) $ + throwE "domain name overruns message" + let !vl = BU.unsafeIndex pkt pos + if | vl == 0 -> do + -- Terminal root label: write the NUL byte and freeze. + sbs <- lift do + (buf', _) <- ensureCap buf cap cursor 1 + A.writeByteArray @Word8 buf' cursor 0 + unsafeMBAToSBS buf' (cursor + 1) + let !consumed = case firstPtr of + Just p -> p - start + 2 + Nothing -> pos - start + 1 + pure (consumed, Domain_ sbs) + | vl <= 63 -> do + -- Inline label: copy its length byte plus content bytes. + -- The @>=@ in the wire-length check reserves space for + -- the terminal NUL the name will eventually need. + let !labLen = fromIntegral vl + !need = 1 + labLen + when (cursor + need >= maxWireLen) $ + throwE "domain name too long" + when (pos + need > pktLen) $ + throwE "label extends past end of message" + (buf', cap') <- lift $ ensureCap buf cap cursor need + lift $ copyFpToMBA (fp `plusForeignPtr` pos) buf' cursor need + walk (pos + need) minPtr cap' buf' (cursor + need) firstPtr + | vl >= 0b1100_0000 -> do + -- Compression pointer: jump to the target, with the + -- target offset constrained to be strictly less than + -- the previous bound to rule out cycles. + unless allowPtr $ + throwE "domain name compression not allowed in current context" + when (pos + 2 > pktLen) $ + throwE "compression pointer extends past end of message" + let !vl' = BU.unsafeIndex pkt (pos + 1) + !off = (fromIntegral (vl .&. 0x3f) `shiftL` 8) + .|. fromIntegral vl' + when (off >= minPtr) $ throwE "invalid compression pointer" + let !firstPtr' = case firstPtr of + Just _ -> firstPtr + Nothing -> Just pos + walk off off cap buf cursor firstPtr' + | otherwise -> throwE "unsupported label type" + +-- Grow @buf@ when @cursor + need@ would overflow @cap@. Growth is +-- @max needed (cap + cap \`shiftR\` 1)@ (1.5x), capped at 'maxWireLen' +-- so we never allocate beyond the wire-form limit. Returns the +-- (possibly resized) buffer and its new capacity. +ensureCap :: MutableByteArray s -> Int -> Int -> Int + -> ST s (MutableByteArray s, Int) +ensureCap buf cap cursor need + | needed <= cap = pure (buf, cap) + | otherwise = do + let !grown = cap + cap `shiftR` 1 + !newCap = min maxWireLen (max needed grown) + buf' <- A.resizeMutableByteArray buf newCap + pure (buf', newCap) + where + !needed = cursor + need diff --git a/internal/Net/DNSBase/Decode/Internal/Message.hs b/internal/Net/DNSBase/Decode/Internal/Message.hs --- a/internal/Net/DNSBase/Decode/Internal/Message.hs +++ b/internal/Net/DNSBase/Decode/Internal/Message.hs @@ -88,7 +88,7 @@ getQueries n = replicateM n getQuery where getQuery :: SGet DnsTriple - getQuery = DnsTriple <$> getDomain <*> getType <*> getClass + getQuery = DnsTriple <$> getDomainNC <*> getType <*> getClass where getType = RRTYPE <$> get16 getClass = RRCLASS <$> get16 diff --git a/internal/Net/DNSBase/Encode/Internal/State.hs b/internal/Net/DNSBase/Encode/Internal/State.hs --- a/internal/Net/DNSBase/Encode/Internal/State.hs +++ b/internal/Net/DNSBase/Encode/Internal/State.hs @@ -40,6 +40,9 @@ , passLen , failWith , setContext + , encoderOffset + , setQNameHint + , setLastOwnerHint ) where import qualified Control.Monad.Trans.RWS.CPS @@ -70,18 +73,24 @@ ---------------------------------------------------------------- --- | Encoder state, the NCTree (DNS name compression tree) is mutable in the ST --- monad. +-- | Encoder state. The QNAME and last-owner hints use 'RootDomain' +-- as the unset sentinel. data EncState s = EncState - { encOffset :: Int - , encDoNC :: Bool - , encNCTree :: NC.NCTree s + { encOffset :: Int + , encDoNC :: Bool + , encNCTree :: NC.NCMap s + , encQName :: Domain + , encQNameOff :: Int + , encLastOwner :: Domain + , encLastOwnerOff :: Int } -- | 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) +encInit donamecomp = do + nc <- stToSTE NC.empty + pure $ EncState 0 donamecomp nc RootDomain 0 RootDomain 0 ---------------------------------------------------------------- @@ -145,23 +154,57 @@ 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 + let !sb = shortBytes domain + !wlen = SB.length sb - 1 + if wlen > 0 && encDoNC + then case tryHintMatch domain encQName encQNameOff + encLastOwner encLastOwnerOff of + Just ptr -> put16 ptr + Nothing -> do + m <- liftSPut . stToSTE $ + NC.lookupAndRegister sb encOffset encNCTree + case m of + Nothing -> putWireForm domain + Just (hb, p) -> do + putShortByteStringPrefix hb sb + put16 p + else putWireForm domain + +-- QNAME is checked before the last-owner hint: its offset always +-- points at verbatim bytes, whereas the last-owner offset may point +-- at a head-bytes-plus-pointer emission, so preferring QNAME keeps +-- pointer chains one hop shorter when both hints match. +tryHintMatch + :: Domain -> Domain -> Int -> Domain -> Int -> Maybe Word16 +tryHintMatch dom qn qOff lo loOff + | qn == dom = Just (mkPtr qOff) + | lo == dom = Just (mkPtr loOff) + | otherwise = Nothing 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 + mkPtr off = fromIntegral (0xC000 .|. off) +{-# INLINE tryHintMatch #-} +-- | Record the QNAME compression hint. +setQNameHint :: ErrorContext r => Domain -> Int -> SPut s r +setQNameHint !dom !off = do + s <- get + put $! s { encQName = dom, encQNameOff = off } + +-- | Record the last-RR-owner compression hint. The offset is +-- updated only when the owner changes, so consecutive RRs sharing +-- the same owner all point at the first emission rather than +-- chaining through each other. +setLastOwnerHint :: ErrorContext r => Domain -> Int -> SPut s r +setLastOwnerHint !dom !off = do + s <- get + when (encLastOwner s /= dom) $ + put $! s { encLastOwner = dom, encLastOwnerOff = off } + +-- | The encoder offset at which the next emitted byte will land. +encoderOffset :: SPutM s r Int +encoderOffset = gets encOffset +{-# INLINE encoderOffset #-} + -- | Encode a domain name verbatim, without name compression. putWireForm :: ErrorContext r => Domain -> SPut s r putWireForm = encVar (SB.length . shortBytes) (B.shortByteString . shortBytes) @@ -172,9 +215,6 @@ pattern MaxPos :: Int pattern MaxPos = 0xffff -pattern MaxPtr :: Int -pattern MaxPtr = 0x3fff - {-# INLINE addPos #-} addPos :: ErrorContext r => Int -> SPut s r addPos n = do @@ -239,6 +279,13 @@ putShortByteString :: ErrorContext r => ShortByteString -> SPut s r putShortByteString b = unless (SB.null b) $ encVar SB.length B.shortByteString b + +-- | Write the first @n@ bytes of a 'ShortByteString' verbatim. +putShortByteStringPrefix :: ErrorContext r => Int -> ShortByteString -> SPut s r +putShortByteStringPrefix !n !sbs + | n > 0 = addPos n >> tell (shortByteStringSliceB 0 n sbs) + | n == 0 = pure () + | otherwise = failWith CantEncode -- should never happen -- | Write a DNS /character-string/: an 8-bit length prefix followed -- by the bytes. Fails with 'EncodeTooLong' if the input exceeds 255 bytes. diff --git a/internal/Net/DNSBase/Internal/Domain.hs b/internal/Net/DNSBase/Internal/Domain.hs --- a/internal/Net/DNSBase/Internal/Domain.hs +++ b/internal/Net/DNSBase/Internal/Domain.hs @@ -44,7 +44,6 @@ -- ** Binary serialization functions , wireBytes , mbWireForm - , buildDomain -- ** Predicates , isLDHLabel , isLDHName @@ -56,10 +55,8 @@ ) where 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.Short as SB -import qualified Data.ByteString.Lazy as LB import qualified Data.ByteString.Unsafe as B import qualified Data.List as L import qualified Data.Primitive.ByteArray as A @@ -282,19 +279,6 @@ instance Show Mbox where showsPrec p m = showsPrec p $ presentString m mempty ----------------------------------------- Wire-form assembly - --- | Execute a builder to produce a 'Domain'. Used by the wire-form --- decoder to turn the @\\@ Builder produced by --- pointer-following into a 'Domain'; the presentation-form parsers --- live in "Net.DNSBase.Domain" and write directly into a fresh --- 'A.MutableByteArray' rather than going via Builder. -buildDomain :: Maybe B.Builder -> Maybe Domain -buildDomain mb = mb >>= \b -> do - let buf = LB.toStrict $ B.toLazyByteStringWith domainStrat mempty b - guard $ B.length buf < 256 - pure $! Domain_ $ SB.toShort buf - ---------------------------------------- Conversions @@ -968,66 +952,71 @@ -- | Build the standard (dot-terminated) /presentation form/ of 'Domain'. presentDomain :: Domain -> Builder -> Builder -presentDomain = fromWire dotB W_dot +presentDomain RootDomain = presentByte W_dot +presentDomain (shortBytes -> sb) = go 0 + where + go :: Int -> Builder -> Builder + go !pos@(fromIntegral . SB.index sb -> !llen) + | llen == 0 = id + | otherwise = presentDomainLabelSlice W_dot sb (pos + 1) llen + . presentByte W_dot + . go (pos + llen + 1) --- | Build the /presentation form/ of a 'Domain' without a trailing dot. --- The root domain is nevertheless presented as a single @.@ byte. +-- | Build the canonical /presentation form/ of a 'Host' as a +-- lower-case name without a trailing dot. The root domain is +-- nevertheless presented as a single @.@ byte. +-- presentHost :: Domain -> Builder -> Builder -presentHost = toCanonical W_dot +presentHost RootDomain = presentByte W_dot +presentHost (shortBytes -> sb) = go 0 + where + !end = SB.length sb - 1 + go :: Int -> Builder -> Builder + go !pos@(fromIntegral . SB.index sb -> !llen) + | !next <- pos + llen + 1 + , next /= end = presentHostLabelSlice W_dot sb (pos + 1) llen + . presentByte W_dot . go next + | otherwise = presentHostLabelSlice W_dot sb (pos + 1) llen --- | Build an ad hoc /mailbox form/ of a 'Domain', without a trailing dot, --- and with @\'\@\'@ as the first label separator. +-- | Build the /mailbox form/ of a 'Domain', without a trailing +-- dot, and with @\'\@\'@ as the first label separator. +-- +-- With single-label names we need to escape both 'W_dot' and +-- 'W_at', so the separator passed to 'canon' must be 'W_dot' +-- in that case, even if it would otherwise be 'W_at'. +-- presentMbox :: Domain -> Builder -> Builder -presentMbox = toCanonical W_at - --- | Build a presentation form. -fromWire :: (Builder -> Builder) -> Word8 -> Domain -> Builder -> Builder -fromWire dterm sep0 (B.uncons . wireBytes -> ht) k - | Just (len, bs) <- ht = go sep0 len bs - | otherwise = impossible +presentMbox RootDomain = presentByte W_dot +presentMbox (shortBytes -> sb) = go W_at 0 where - go :: Word8 -> Word8 -> ByteString -> Builder - go _ 0 _ = dotB k - go sep len bytes = - let (label, suffix) = B.splitAt (fromEnum len) bytes - in case B.uncons suffix of - Just (slen, sbytes) - | slen > 0 -> presentDomainLabel sep label - . presentByte sep - $ go W_dot slen sbytes - _ -> presentDomainLabel W_dot label $ dterm k + !end = SB.length sb - 1 + go :: Word8 -> Int -> Builder -> Builder + go !sep !pos@(fromIntegral . SB.index sb -> !llen) + | !next <- pos + llen + 1 + , next /= end = presentDomainLabelSlice sep sb (pos + 1) llen + . presentByte sep . go W_dot next + | otherwise = presentDomainLabelSlice W_dot sb (pos + 1) llen --- | Build a canonical presentation form (folded to lower case) -toCanonical :: Word8 -> Domain -> Builder -> Builder -toCanonical sep0 (B.uncons . wireBytes -> ht) k - | Just (len, bs) <- ht = go sep0 len bs - | otherwise = impossible - where - go :: Word8 -> Word8 -> ByteString -> Builder - go _ 0 _ = dotB k - go sep len bytes = - let (label, suffix) = B.splitAt (fromEnum len) bytes - in canon sep label - $ case B.uncons suffix of - Just (slen, sbytes) - | slen > 0 -> presentByte sep (go W_dot slen sbytes) - _ -> k - where - canon W_dot = presentHostLabel W_dot - canon w = presentDomainLabel w +-- | Walk a slice of a 'ShortByteString' applying 'domainLabelBP'. +presentDomainLabelSlice + :: Word8 -> ShortByteString -> Int -> Int -> Builder -> Builder +{-# INLINE presentDomainLabelSlice #-} +presentDomainLabelSlice sep sb off len k = + primMapShortByteStringSliceBounded (domainLabelBP sep) off len sb <> k +-- | Walk a slice of a 'ShortByteString' applying 'hostLabelBP'. +presentHostLabelSlice + :: Word8 -> ShortByteString -> Int -> Int -> Builder -> Builder +{-# INLINE presentHostLabelSlice #-} +presentHostLabelSlice sep sb off len k = + primMapShortByteStringSliceBounded (hostLabelBP sep) off len sb <> k + ---------------------------------------- Util -- | Most domain names are short, use small buffers, but no need to make them -- too tight since we ultimately copy again into a short bytestring. domainStrat :: B.AllocationStrategy domainStrat = B.untrimmedStrategy 32 128 - -pattern W_dot :: Word8; pattern W_dot = 0x2e -pattern W_at :: Word8; pattern W_at = 0x40 - -dotB :: Builder -> Builder -dotB = presentByte W_dot w2i :: Word8 -> Int w2i = fromIntegral diff --git a/internal/Net/DNSBase/Internal/Message.hs b/internal/Net/DNSBase/Internal/Message.hs --- a/internal/Net/DNSBase/Internal/Message.hs +++ b/internal/Net/DNSBase/Internal/Message.hs @@ -111,7 +111,9 @@ -- putQuestion :: DnsTriple -> SPut s RData putQuestion DnsTriple{..} = do + qOff <- encoderOffset putDomain dnsTripleName + setQNameHint dnsTripleName qOff put32 $ fromIntegral @Word16 (coerce dnsTripleType) `unsafeShiftL` 16 .|. fromIntegral @Word16 (coerce dnsTripleClass) diff --git a/internal/Net/DNSBase/Internal/NameComp.hs b/internal/Net/DNSBase/Internal/NameComp.hs --- a/internal/Net/DNSBase/Internal/NameComp.hs +++ b/internal/Net/DNSBase/Internal/NameComp.hs @@ -1,67 +1,256 @@ -- | -- Module : Net.DNSBase.Internal.NameComp --- Description : TBD +-- Description : DNS name-compression state for the wire-format encoder -- Copyright : (c) Viktor Dukhovni, 2026 -- License : BSD-3-Clause -- Maintainer : ietf-dane@dukhovni.org -- Stability : unstable +-- +-- A mutable trie whose keys ('NCNode') pair a parent suffix's +-- message offset with a chunk slice of the owning name's wire form, +-- mapping each entry to the offset at which the corresponding suffix +-- first appears in the output message. Used by +-- 'Net.DNSBase.Encode.Internal.State.putDomain' to emit DNS name +-- pointers per RFC 1035 ยง4.1.4. +-- +-- A /chunk/ is a contiguous run of consecutive DNS labels (with their +-- length bytes) covering at least 'minChunkBytes' bytes of wire form. +-- Chunks are anchored at the tail of the name: starting from the +-- rightmost label, each chunk extends leftward until it reaches the +-- minimum-size threshold, at which point a new chunk begins. The +-- leftmost chunk may be smaller than the threshold (it inherits +-- whatever labels are left over). A label too short to make a +-- worthwhile pointer target on its own is therefore grouped with its +-- right neighbours. +-- +-- The trie itself is a single flat hashtable; each name with @C@ +-- chunks contributes @C@ entries, and each traversal step is one +-- hashtable lookup matching whole chunks. +-- +-- The chunk component of each key is a slice (offset + length) into +-- a 'ShortByteString' belonging to the 'Domain' that first introduced +-- the chunk. The slice retains a reference to that buffer for the +-- duration of the encode, so chunk bytes are not copied into freshly +-- allocated keys. +-- {-# LANGUAGE BangPatterns #-} +{-# LANGUAGE MagicHash #-} +-- {-# LANGUAGE PatternSynonyms #-} module Net.DNSBase.Internal.NameComp - ( NCTree + ( NCMap , empty - , insert - , lookup + , lookupAndRegister ) where -import Prelude hiding (lookup) -import Control.Monad.ST as ST -import qualified Data.ByteString as B -import qualified Data.HashTable.ST.Basic as LH +import Net.DNSBase.Internal.Util + +import qualified Data.ByteString.Short as SB import qualified Data.HashTable.Class as H +import qualified Data.HashTable.ST.Basic as LH +import Control.Monad.ST (ST) +import Data.Foldable (foldlM) +import Data.Hashable (Hashable(..), hashByteArrayWithSalt) +import Data.Primitive.ByteArray + ( ByteArray# + , MutableByteArray + , compareByteArrays + , newByteArray + , readByteArray + , writeByteArray + ) -type Path = [B.ByteString] -type Map s = LH.HashTable s B.ByteString -data NCTree s = NCTree (Map s (NCTree s)) Int +-------------------------------------------------------------------------------- +-- Compression-trie node +-------------------------------------------------------------------------------- --- | Create a root node with given value -empty :: Int -> ST.ST s (NCTree s) -empty n = flip NCTree n <$> H.new +-- A compression-trie node identified by its parent suffix's +-- message-offset and a slice into the owning name's 'ShortByteString' +-- holding the chunk's wire bytes. The byte array reference keeps +-- the source buffer alive as long as the node is reachable from the +-- map; 'Eq' and 'Hashable' read the byte range in place. +-- +-- Field widths reflect the bounds: parent offsets are clamped at the +-- 14-bit pointer range (with @0xFFFF@ as the @ROOT@ sentinel, +-- distinct from any real offset because real offsets never exceed +-- @0x3FFF@), and a DNS name is at most 255 wire bytes so the chunk +-- offset and length both fit in a byte. +data NCNode = NCNode + !ByteArray# -- owning byte array + {-# UNPACK #-} !Word16 -- parent offset + {-# UNPACK #-} !Word8 -- chunk byte offset + {-# UNPACK #-} !Word8 -- chunk byte length --- | Insert a domain with the given labels with last label ending at given --- offset (if stored uncompressed). The caller MUST not insert any paths whose --- tail lies beyond the first 16k of the DNS message. That is the end offset --- must not exceed @0x3fff@. -insert :: Path -> Int -> NCTree s -> ST.ST s () -insert = go +instance Eq NCNode where + NCNode aA pA offA lenA == NCNode aB pB offB lenB + | pA /= pB || lenA /= lenB = False + | otherwise = + compareByteArrays (ByteArray aA) (fromIntegral offA) + (ByteArray aB) (fromIntegral offB) + (fromIntegral lenA) == EQ + +instance Hashable NCNode where + hashWithSalt salt (NCNode ba parent off len) = + hashByteArrayWithSalt ba (fromIntegral off) (fromIntegral len) + (salt `hashWithSalt` (fromIntegral parent :: Int)) + +-------------------------------------------------------------------------------- +-- Public API +-------------------------------------------------------------------------------- + +-- | The compression state: a single mutable hashtable keyed on +-- 'NCNode' (parent offset plus chunk slice), mapping each entry to +-- the message offset at which the suffix @chunk.parent-suffix@ +-- first appears. +newtype NCMap s = NCMap (LH.HashTable s NCNode Int) + +-- | Create a fresh empty compression state. +empty :: ST s (NCMap s) +empty = NCMap <$> H.new + +-- The 14-bit pointer cap. Children with absolute offsets exceeding +-- this can't be the target of an RFC 1035 compression pointer, so +-- registering them is pointless. +maxPtr :: Int +maxPtr = 0x3fff + +-- Sentinel parent-offset for the root level of the trie. Distinct +-- from any real wire-message offset, which is non-negative. +parentRoot :: Int +parentRoot = -1 + +-- The pointer high-bits mask: a name pointer is @0xC000 .|. target@. +pointerMark :: Int +pointerMark = 0xC000 + +-- Minimum byte count for a chunk. A trie entry pointing at @n@ wire +-- bytes pays back @n - 2@ bytes on a hit, so chunks shorter than +-- about 4 bytes barely earn their keep; 8 puts every registered +-- chunk solidly into "worth pointing at" territory while still +-- exploiting the long zone-suffix sharing that dominates real DNS +-- traffic. +minChunkBytes :: Word8 +minChunkBytes = 8 + +-- Max chunks per name. Each chunk has at least 'minChunkBytes' wire +-- bytes except possibly the leftmost leftover, so at most +-- @floor(255 \/ minChunkBytes) + 1@. +maxChunks :: Int +maxChunks = fromIntegral $ 1 + 255 `quot` minChunkBytes + +-- | Look up and register a domain in one walk. +-- +-- Given the wire form @dname@ of a domain (including its terminal NUL +-- byte) about to be emitted at message offset @baseOff@, walk its +-- chunks root-first. In the prefix that already matches the map, +-- thread parent offsets through hits; at the first miss (or running +-- off the end), register the new entries the input contributes. +-- Registrations whose absolute offsets exceed the 14-bit pointer +-- range ('maxPtr') are silently dropped: the input can still benefit +-- from compression against existing entries, but won't itself be the +-- target of future pointers. +-- +-- Returns 'Nothing' when the name has no compressible suffix in the +-- map (so the caller emits the full wire form), or +-- @'Just' (headBytes, ptr)@ to indicate the caller should emit the +-- first @headBytes@ bytes of the wire form verbatim, then the 2-byte +-- pointer word @ptr@. +lookupAndRegister :: forall s. ShortByteString + -> Int + -> NCMap s + -> ST s (Maybe (Int, Word16)) +lookupAndRegister !dname@(SBS nameArr) !baseOff (NCMap m) = do + if | wlen == 1 -> pure Nothing -- root domain + | otherwise -> do + chunkArr <- newByteArray maxChunks + nChunks <- buildChunks chunkArr + phase1 chunkArr nChunks 0 (wlen - 1) parentRoot Nothing where - go :: Path -> Int -> NCTree s -> ST.ST s () - go (!l:ls) !end !(NCTree m _) = - H.mutateST m l $ alter ls $! end - B.length l - 1 - go _ _ _ = pure () + !wlen = SB.length dname - -- | Alter (create or update) the node with given index and start position - alter :: Path -> Int -> Maybe (NCTree s) -> ST.ST s (Maybe (NCTree s), ()) - alter !ls !start !old = case old of - -- At existing intermediate nodes recurse to store the rest of the path - Just n | null ls -> pure $ node n - | otherwise -> node n <$ go ls start n - -- In new intermediate nodes store the tip offset + distance from tip - Nothing | null ls -> node <$> empty start - | otherwise -> do - e <- empty start - node e <$ go ls start e + -- Compute offsets of one-or-more chunks of consecutive labels, + -- each at least 'minChunkBytes' wire bytes (the leftmost + -- leftover may be smaller). Chunks are emitted in root-first + -- iteration order: @chunkArr[0]@ is the outermost (rightmost) + -- chunk, @chunkArr[nChunks - 1]@ is the innermost (leftmost). + -- + buildChunks :: MutableByteArray s -> ST s Int + buildChunks carr = walk 0 [] where - node n = (Just n, ()) + walk :: Word8 -> [Word8] -> ST s Int + walk !pos@((+ 1) . SB.index dname . fromIntegral -> !llen) !lens + | llen /= 1 = walk (pos + llen) (llen : lens) + | otherwise = foldlM step (0, 1, fromIntegral $ wlen - 1) lens >>= \ case + (!ix, 0, _) -> pure ix + (!ix, _, _) -> (ix + 1) <$ writeByteArray @Word8 carr ix 0 --- | Return the length of the path prefix (domain suffix) and corresponding --- offset for the input path (reversed list of wire-form labels), counting both --- the length (1) and payload size of each label, not including the terminal --- NUL label. -lookup :: Path -> (NCTree s) -> ST.ST s (Int, Int) -lookup labels root = go labels root 0 - where - go (!l:ls) !(NCTree m off) !slen = do - mn <- H.lookup m l - case mn of - Just n -> go ls n $! slen + B.length l + 1 - _ -> pure (slen, off) - go _ (NCTree _ off) slen = pure (slen, off) + step :: (Int, Word8, Word8) -> Word8 -> ST s (Int, Word8, Word8) + step (!ix, !acc, !end) !llen + | acc' < minChunkBytes = pure (ix, acc', start) + | otherwise = (ix + 1, 0, start) <$ writeByteArray carr ix start + where + !acc' = acc + llen + !start = end - llen + + -- Lookup phase. Walks chunks root-first; each step either + -- advances the parent offset on a hit, or transitions to phase 2 + -- on the first miss. @chunkEnd@ is the exclusive end position + -- of the current chunk in the wire form (one past its last + -- byte); each step's @chunkEnd@ becomes the next step's via the + -- current chunk's start (chunks are contiguous in the wire form, + -- with the next-outer chunk lying immediately to the right). + -- The 'acc' carries @(target, matchedChunkStart)@ for the + -- deepest match seen so far, or 'Nothing' if no chunk has + -- matched yet. + phase1 + :: MutableByteArray s + -> Int -- chunk count + -> Int -- current chunk index + -> Int -- chunk end (exclusive) + -> Int -- current parent_off + -> Maybe (Int, Int) -- (target, matchedChunkStart) + -> ST s (Maybe (Int, Word16)) + phase1 carr nChunks i chunkEnd parent acc + | i >= nChunks = pure (finalize acc) + | otherwise = do + cs <- readByteArray carr i :: ST s Word8 + let !chunkStart = fromIntegral cs :: Int + !chunkLen = chunkEnd - chunkStart + !node = NCNode nameArr (fromIntegral parent) + cs (fromIntegral chunkLen) + mhit <- H.lookup m node + case mhit of + Just hitOff -> + phase1 carr nChunks (i + 1) chunkStart hitOff + (Just (hitOff, chunkStart)) + Nothing -> phase2 carr nChunks i chunkEnd parent acc + + -- Insert phase. Walks the remainder of chunks from index @i@ + -- onward, registering each at its absolute offset. Drops + -- out-of-range registrations silently. Returns whatever + -- compression decision was frozen at the transition. + phase2 + :: MutableByteArray s + -> Int + -> Int + -> Int + -> Int + -> Maybe (Int, Int) + -> ST s (Maybe (Int, Word16)) + phase2 carr nChunks i chunkEnd parent acc + | i >= nChunks = pure (finalize acc) + | otherwise = do + cs <- readByteArray carr i :: ST s Word8 + let !chunkStart = fromIntegral cs :: Int + !chunkLen = chunkEnd - chunkStart + !node = NCNode nameArr (fromIntegral parent) + cs (fromIntegral chunkLen) + !newOff = baseOff + chunkStart + when (newOff <= maxPtr) $ + H.insert m node newOff + phase2 carr nChunks (i + 1) chunkStart newOff acc + + finalize :: Maybe (Int, Int) -> Maybe (Int, Word16) + finalize Nothing = Nothing + finalize (Just (tgt, cs)) = Just (cs, mkPtr tgt) + + mkPtr :: Int -> Word16 + mkPtr off = fromIntegral (pointerMark .|. off) diff --git a/internal/Net/DNSBase/Internal/Present.hs b/internal/Net/DNSBase/Internal/Present.hs --- a/internal/Net/DNSBase/Internal/Present.hs +++ b/internal/Net/DNSBase/Internal/Present.hs @@ -27,6 +27,24 @@ , hPutBuilder -- *** 'hPutBuilder' specialised to @stdout@ , putBuilder + -- ** Zone-file escape primitives + , charBP + , bsEscBP + , decEscBP + -- ** ASCII byte-name patterns + , pattern W_space + , pattern W_dquote + , pattern W_dollar + , pattern W_open + , pattern W_close + , pattern W_comma + , pattern W_dot + , pattern W_0 + , pattern W_semi + , pattern W_at + , pattern W_A + , pattern W_bslash + , pattern W_delete ) where import qualified Data.ByteString.Builder as B @@ -280,3 +298,43 @@ do (> 9) do P.word8Dec do ((), ) >$< (const 0x30 >$< P.liftFixedToBounded P.word8) >*< P.word8Dec + +------------- Zone-file escape primitives + +charBP :: P.BoundedPrim Word8 +{-# INLINE charBP #-} +charBP = P.liftFixedToBounded P.word8 + +bsEscBP :: P.BoundedPrim Word8 +{-# INLINE bsEscBP #-} +bsEscBP = (W_bslash,) >$< charBP >*< charBP + +decEscBP :: P.BoundedPrim Word8 +{-# INLINE decEscBP #-} +decEscBP = P.condB (> 99) dec3 $ P.condB (> 9) dec2 dec1 + where + dec3 = (W_bslash,) + >$< charBP >*< P.word8Dec + dec2 = (W_bslash,) . (W_0,) + >$< charBP >*< charBP >*< P.word8Dec + dec1 = ((W_bslash, W_0),) . (W_0,) + >$< (charBP >*< charBP) >*< (charBP >*< P.word8Dec) + {-# INLINE dec3 #-} + {-# INLINE dec2 #-} + {-# INLINE dec1 #-} + +------------- ASCII byte-name patterns + +pattern W_space :: Word8; pattern W_space = 0x20 +pattern W_dquote :: Word8; pattern W_dquote = 0x22 +pattern W_dollar :: Word8; pattern W_dollar = 0x24 +pattern W_open :: Word8; pattern W_open = 0x28 +pattern W_close :: Word8; pattern W_close = 0x29 +pattern W_comma :: Word8; pattern W_comma = 0x2c +pattern W_dot :: Word8; pattern W_dot = 0x2e +pattern W_0 :: Word8; pattern W_0 = 0x30 +pattern W_semi :: Word8; pattern W_semi = 0x3b +pattern W_at :: Word8; pattern W_at = 0x40 +pattern W_A :: Word8; pattern W_A = 0x41 +pattern W_bslash :: Word8; pattern W_bslash = 0x5c +pattern W_delete :: Word8; pattern W_delete = 0x7f diff --git a/internal/Net/DNSBase/Internal/RR.hs b/internal/Net/DNSBase/Internal/RR.hs --- a/internal/Net/DNSBase/Internal/RR.hs +++ b/internal/Net/DNSBase/Internal/RR.hs @@ -70,7 +70,9 @@ putRR :: RR -> SPut s RData putRR RR{..} = do + ownerOff <- encoderOffset putDomain rrOwner + setLastOwnerHint rrOwner ownerOff putSizedBuilder $! mbWord16 (coerce $ rdataType rrData) <> mbWord16 (coerce rrClass) diff --git a/internal/Net/DNSBase/Internal/Text.hs b/internal/Net/DNSBase/Internal/Text.hs --- a/internal/Net/DNSBase/Internal/Text.hs +++ b/internal/Net/DNSBase/Internal/Text.hs @@ -13,6 +13,8 @@ , presentDomainLabel , presentHostLabel , presentCSVList + , domainLabelBP + , hostLabelBP ) where import qualified Data.ByteString.Builder as B @@ -69,21 +71,30 @@ -> B.Builder {-# INLINE presentDomainLabel #-} presentDomainLabel sep bytes k = - P.primMapByteStringBounded bp bytes <> k + P.primMapByteStringBounded (domainLabelBP sep) bytes <> k + +-- | 'BoundedPrim' driving 'presentDomainLabel'-style per-byte +-- presentation-form escaping with the given label separator +-- (typically 0x2e, @\'.\'@). Exposed so that callers walking an +-- unpinned byte source (e.g.\ a 'ShortByteString' slice via +-- 'primMapShortByteStringSliceBounded') can apply the same escape +-- rules without going through a pinned 'ByteString'. +domainLabelBP :: Word8 -> P.BoundedPrim Word8 +{-# INLINE domainLabelBP #-} +domainLabelBP sep = bp where isgraph = \w -> w > W_space && w < W_delete bp = P.condB isgraph sp decEscBP sp = P.condB special bsEscBP charBP - where - special = \ case - W_dquote -> True - W_bslash -> True - W_dollar -> True - W_open -> True - W_close -> True - W_semi -> True - W_at -> True - w -> w == sep + special = \ case + W_dquote -> True + W_bslash -> True + W_dollar -> True + W_open -> True + W_close -> True + W_semi -> True + W_at -> True + w -> w == sep -- | Present a 'Net.DNSBase.Domain.Host' label folded to lower case, with the given continuation. -- @@ -93,7 +104,14 @@ -> B.Builder {-# INLINE presentHostLabel #-} presentHostLabel sep bytes k = - P.primMapByteStringBounded bp bytes <> k + P.primMapByteStringBounded (hostLabelBP sep) bytes <> k + +-- | 'BoundedPrim' driving 'presentHostLabel'-style per-byte +-- presentation-form escaping with lowercase folding of ASCII +-- uppercase letters. +hostLabelBP :: Word8 -> P.BoundedPrim Word8 +{-# INLINE hostLabelBP #-} +hostLabelBP sep = bp where isgraph = \w -> w > W_space && w < W_delete bp = P.condB isgraph sp decEscBP @@ -145,43 +163,6 @@ >$< charBP >*< charBP >*< charBP bsEsc''' = (W_bslash,) . (W_bslash,) . (W_bslash,) >$< charBP >*< charBP >*< charBP >*< charBP - -------------- Text encoding BoundedPrim helpers - -charBP :: P.BoundedPrim Word8 -{-# INLINE charBP #-} -charBP = P.liftFixedToBounded P.word8 - -bsEscBP :: P.BoundedPrim Word8 -{-# INLINE bsEscBP #-} -bsEscBP = (W_bslash,) >$< charBP >*< charBP - -decEscBP :: P.BoundedPrim Word8 -{-# INLINE decEscBP #-} -decEscBP = P.condB (> 99) dec3 $ P.condB (> 9) dec2 dec1 - where - dec3 = (W_bslash,) - >$< charBP >*< P.word8Dec - dec2 = (W_bslash,) . (W_0,) - >$< charBP >*< charBP >*< P.word8Dec - dec1 = ((W_bslash, W_0),) . (W_0,) - >$< (charBP >*< charBP) >*< (charBP >*< P.word8Dec) - {-# INLINE dec3 #-} - {-# INLINE dec2 #-} - {-# INLINE dec1 #-} - -pattern W_space :: Word8; pattern W_space = 0x20 -pattern W_dquote :: Word8; pattern W_dquote = 0x22 -pattern W_dollar :: Word8; pattern W_dollar = 0x24 -pattern W_open :: Word8; pattern W_open = 0x28 -pattern W_close :: Word8; pattern W_close = 0x29 -pattern W_comma :: Word8; pattern W_comma = 0x2c -pattern W_0 :: Word8; pattern W_0 = 0x30 -pattern W_semi :: Word8; pattern W_semi = 0x3b -pattern W_at :: Word8; pattern W_at = 0x40 -pattern W_A :: Word8; pattern W_A = 0x41 -pattern W_bslash :: Word8; pattern W_bslash = 0x5c -pattern W_delete :: Word8; pattern W_delete = 0x7f ----- diff --git a/internal/Net/DNSBase/Internal/Util.hs b/internal/Net/DNSBase/Internal/Util.hs --- a/internal/Net/DNSBase/Internal/Util.hs +++ b/internal/Net/DNSBase/Internal/Util.hs @@ -10,7 +10,7 @@ , bool, cond , compose4 , ByteArray(..), baToShortByteString, modifyArray - , sbsToByteArray, sbsToMutableByteArray + , sbsToByteArray, sbsToMutableByteArray, copyFpToMBA, unsafeMBAToSBS , Down(..), comparing , (.|.), (.&.), clearBit, countLeadingZeros, complement, setBit , shiftL, shiftR, testBit, unsafeShiftL, unsafeShiftR @@ -29,13 +29,19 @@ , allocaBytesAligned, castPtr, copyBytes, byteSwap32 , fillBytes, minusPtr, peek, peekElemOff, plusForeignPtr , unsafePerformFPIO + , shortByteStringSliceB + , primMapShortByteStringSliceBounded ) where import qualified Data.Primitive.ByteArray as A import qualified Data.ByteString.Short as SB +import qualified Data.ByteString.Builder.Internal as BI +import qualified Data.ByteString.Builder.Prim as BP +import qualified Data.ByteString.Builder.Prim.Internal as BPI import Control.Applicative ((<|>)) import Control.Monad ( (>=>), forM, forM_, guard, join, mzero, replicateM ) import Control.Monad ( unless, void, when ) +import Control.Monad.Primitive (unsafeIOToPrim, unsafePrimToIO) import Control.Monad.ST (ST) import Control.Monad.Trans.Class (lift) import Control.Monad.Trans.Except (ExceptT(ExceptT), throwE, catchE, runExceptT, withExceptT) @@ -63,7 +69,7 @@ import Data.Typeable (Typeable, cast) import Data.Word (Word8, Word16, Word32, Word64, byteSwap16, byteSwap32, byteSwap64) import Foreign (ForeignPtr, Ptr, allocaBytesAligned, castPtr, copyBytes) -import Foreign (fillBytes, minusPtr, peek, peekElemOff, plusForeignPtr) +import Foreign (fillBytes, minusPtr, peek, peekElemOff, plusForeignPtr, plusPtr) import GHC.ByteOrder (ByteOrder(..), targetByteOrder) import GHC.ForeignPtr (unsafeWithForeignPtr) import Type.Reflection (TypeRep, pattern TypeRep) @@ -170,3 +176,70 @@ sbsToMutableByteArray :: ShortByteString -> ST s (MutableByteArray s) sbsToMutableByteArray sb@(SBS ba) = A.thawByteArray (ByteArray ba) 0 (SB.length sb) + +-- Copy @n@ bytes from the input pointer into the mutable output +-- buffer at offset @dstOff@. Crosses the IO\/ST bridge via +-- 'unsafeIOToPrim' \/ 'unsafePrimToIO': the 'ForeignPtr' unwrap +-- requires IO, while the 'MutableByteArray' lives in 'ST s'. +copyFpToMBA :: forall s. ForeignPtr Word8 + -> MutableByteArray s + -> Int + -> Int + -> ST s () +copyFpToMBA fp !dst !dstOff !n = + unsafeIOToPrim @(ST s) $ unsafeWithForeignPtr @Word8 fp + $ \ p -> unsafePrimToIO $ cp dst dstOff p n + where + cp :: MutableByteArray s -> Int -> Ptr Word8 -> Int -> ST s () + cp = A.copyPtrToMutableByteArray + +-- Shrink a mutable byte array to the payload size and free in +-- place, invalidating the original, which must not be used again. +-- +unsafeMBAToSBS :: MutableByteArray s -> Int -> ST s ShortByteString +unsafeMBAToSBS !m !n = do + A.shrinkMutableByteArray m n + baToShortByteString <$> A.unsafeFreezeByteArray m + +----- 'ShortByteString' slice builders + +-- | Emit @len@ bytes starting at @off@ within @sbs@ as a 'Builder', +-- copying directly from the underlying byte array into the builder's +-- output buffer. The caller is responsible for ensuring +-- @0 <= off@ and @off + len <= SB.length sbs@. +shortByteStringSliceB :: Int -> Int -> ShortByteString -> Builder +shortByteStringSliceB !off !len sbs + | len <= 0 = mempty + | otherwise = BI.ensureFree len <> BI.builder step + where + step :: forall r. BI.BuildStep r -> BI.BuildStep r + step k (BI.BufferRange op opEnd) = do + A.copyByteArrayToPtr op (sbsToByteArray sbs) off len + k (BI.BufferRange (op `plusPtr` len) opEnd) + +-- | Walk @len@ bytes starting at @off@ within @sbs@, applying a +-- 'BP.BoundedPrim Word8' to each byte and emitting the result into a +-- 'Builder'. Equivalent to bytestring's 'BP.primMapByteStringBounded' +-- but sourced from an unpinned 'ShortByteString' slice rather than a +-- pinned 'ByteString'. Handles output-buffer-full mid-slice by +-- resuming from the next unprocessed input byte. The caller is +-- responsible for ensuring @0 <= off@ and @off + len <= SB.length sbs@. +primMapShortByteStringSliceBounded + :: BP.BoundedPrim Word8 -> Int -> Int -> ShortByteString -> Builder +primMapShortByteStringSliceBounded bp !off !len sbs + | len <= 0 = mempty + | otherwise = BI.builder (step off) + where + !bound = (BPI.sizeBound bp) + !sEnd = off + len + + step :: Int -> BI.BuildStep r -> BI.BuildStep r + step !i k (BI.BufferRange op0 ope) = go i op0 + where + go !j !op + | j >= sEnd = k (BI.BufferRange op ope) + | op `plusPtr` bound > ope = + return $ BI.bufferFull bound op (step j k) + | otherwise = do + op' <- BPI.runB bp (SB.index sbs j) op + go (j + 1) op' diff --git a/tests/test.hs b/tests/test.hs --- a/tests/test.hs +++ b/tests/test.hs @@ -203,8 +203,8 @@ , ( mkRR zone $ T_SOA $$(dnLit8 "dns.example.org") $$(mbLit8 "postmaster@dns.example.org") 2023111301 1800 900 604800 86400 , "example.org. 300 IN SOA dns.example.org. postmaster.dns.example.org. 2023111301 1800 900 604800 86400" , "076578616d706c65036f726700" - <> "0006" <> "0001" <> "0000012c" <> "0027" - <> "03646e73c000" <> "0a706f73746d6173746572c017" + <> "0006" <> "0001" <> "0000012c" <> "002b" + <> "03646e73c000" <> "0a706f73746d617374657203646e73c000" <> "78963a85" <> "00000708" <> "00000384" <> "00093a80" <> "00015180" ) , ( mkRR zone $ T_MB $$(dnLit8 "madname.example.org") @@ -1213,11 +1213,13 @@ pure $ EdnsOption $ O_EDE code name text where getName :: Word16 -> Maybe OptionCodec -> ShortByteString - getName (fromIntegral -> code) + getName code (Just (OptionCodec (_ :: Proxy b) (v :: OptionExtensionVal b))) = case eqT @O_ede @b of - Just Refl -> IM.findWithDefault mempty code v - Nothing -> mempty + Just Refl + | i <- fromIntegral code + , Just s <- IM.lookup i v -> s + _ -> mempty getName _ _ = mempty genOpaqueOpt :: OptionMap -> Gen EdnsOption