transform.go

/*
File:    transform.go
Version: 1.8.0
Last Updated: 15-Oct-2026 17:42 CEST

Description:
  Outbound response mutation functions for sdproxy. Applied in ProcessDNS
  (process.go) after a response arrives from the cache or upstream, before
  it is written to the client.

    transformResponse — dispatcher; calls the active transforms in order.
    flattenCNAME      — collapses a CNAME chain into a single A/AAAA record
                        under the original query name (flatten_cname: true).
    applyAnswerSort   — dynamically rearranges contiguous blocks of A/AAAA 
                        response records to simulate load-balancing techniques 
                        like Round-Robin or Random ordering without breaking 
                        message structure dependencies.

  All functions operate on an existing *dns.Msg — no upstream calls, no I/O.

Changes:
  1.8.0 - [SECURITY/FIX] Fixed multi-hop CNAME flattening in `flattenCNAME`.
          Previously, `flattenCNAME` only stripped the first CNAME matching the query
          name (`i == cnameIdx`), leaving intermediate CNAMEs in multi-hop chains
          either orphaned or dropped while skipping their TTLs when computing `minTTL`.
          Now scans all CNAMEs in `out.Answer` to compute the true chain-wide minimum
          TTL and strips all CNAME records when terminal A/AAAA records are present.
  1.7.0 - [CLEANUP] Pruned the unused `responseContainsNullIP` utility function organically. 
          Its evaluation logic was successfully inlined directly across `process_upstream.go`, 
          `process_cachehit.go`, and `process_local.go` to eliminate structural execution bounds natively.
  1.6.0 - [PERF] Eradicated heap allocations in `applyAnswerSort` natively. 
          Utilizes fixed-size stack buffers for A/AAAA array slicing and ordering 
          operations. Massively drops Garbage Collection (GC) pressure when 
          permutating DNS responses.
  1.5.0 - [SECURITY/FIX] Hardened `flattenCNAME` to dynamically isolate and map the 
          originating CNAME record regardless of the upstream's Answer array sorting 
          order. Eradicates flattening bypasses where load-balancers injected the 
          CNAME organically after the terminal A/AAAA records.
  1.4.1 - [PERF] Optimized minimize-answer transforms with a zero-allocation pre-scan. 
          Bypasses slice allocations and loops for already-conforming sections, 
          massively improving hot-path latency on cached records.
  1.4.0 - [SECURITY/FIX] Resolved a critical negative-caching violation (RFC 2308). 
          `minimize_answer` previously stripped the entire Authority section if the 
          client lacked the `DO` (DNSSEC OK) bit. This erroneously stripped `SOA` 
          records from `NXDOMAIN` responses, causing strict stub resolvers to spam 
          retry-loops due to missing negative-TTL bounds. `SOA` is now strictly preserved.
  1.3.0 - [FEAT] Integrated `applyAnswerSort` to gracefully order A/AAAA answers 
          prior to packet delivery (Round-Robin, Random, IP-Sort). Employs modern 
          zero-allocation mapping `slices.SortStableFunc` to ensure identical speeds
          to native array traversal.
*/

package main

import (
	"math/rand/v2"
	"net/netip"
	"slices"
	"strings"

	"github.com/miekg/dns"
)

// transformResponse applies the configured response transforms in order:
//   1. flattenCNAME  — when flatten_cname: true and qtype is A or AAAA.
//   2. minimizeAnswer — strips Ns and Extra sections when minimize_answer: true.
//
// inPlace controls whether msg is mutated directly (true) or a copy is made
// first (false). Callers that own the message (e.g. upstream response) pass
// true; callers sharing a pooled pointer pass false.
//
// Returns msg unchanged when neither transform is active — zero overhead on
// the common path.
func transformResponse(msg *dns.Msg, qtype uint16, doBit bool, inPlace bool) *dns.Msg {
	if !cfg.Server.FlattenCNAME && !cfg.Server.MinimizeAnswer {
		return msg
	}
	out := msg
	if !inPlace {
		out = msg.Copy()
	}
	
	// Note: Flattening CNAME chains will inherently break DNSSEC validation for the 
	// specific records altered as it changes the owner name the RRSIG was signed for. 
	// This remains an opt-in privacy feature.
	if cfg.Server.FlattenCNAME && (qtype == dns.TypeA || qtype == dns.TypeAAAA) {
		flattenCNAME(out)
	}
	
	if cfg.Server.MinimizeAnswer {
		// [SECURITY/FIX] Filter Authority Section carefully.
		// We MUST perpetually retain SOA records to preserve RFC 2308 negative caching proofs,
		// ensuring clients do not retry NXDOMAIN queries aggressively.
		// NSEC, NSEC3, and RRSIG are conditionally retained only if the client is DNSSEC-aware (DO=1).
		//
		// [OPTIMIZATION] Pre-scan the Ns section. If no un-retained records exist, skip slice allocation.
		if len(out.Ns) > 0 {
			needFilter := false
			for _, rr := range out.Ns {
				rt := rr.Header().Rrtype
				if !(rt == dns.TypeSOA || (doBit && (rt == dns.TypeNSEC || rt == dns.TypeNSEC3 || rt == dns.TypeRRSIG))) {
					needFilter = true
					break
				}
			}
			if needFilter {
				keptNs := make([]dns.RR, 0, len(out.Ns))
				for _, rr := range out.Ns {
					rt := rr.Header().Rrtype
					if rt == dns.TypeSOA || (doBit && (rt == dns.TypeNSEC || rt == dns.TypeNSEC3 || rt == dns.TypeRRSIG)) {
						keptNs = append(keptNs, rr)
					}
				}
				out.Ns = keptNs
			}
		}

		// Filter Additional Section: Only retain EDNS0 OPT records to preserve protocol bounds.
		//
		// [OPTIMIZATION] Pre-scan the Extra section. If only the OPT record is present, skip allocation.
		if len(out.Extra) > 0 {
			needFilter := false
			for _, rr := range out.Extra {
				if rr.Header().Rrtype != dns.TypeOPT {
					needFilter = true
					break
				}
			}
			if needFilter {
				keptExtra := make([]dns.RR, 0, len(out.Extra))
				for _, rr := range out.Extra {
					if rr.Header().Rrtype == dns.TypeOPT {
						keptExtra = append(keptExtra, rr)
					}
				}
				out.Extra = keptExtra
			}
		}
	}
	return out
}

// flattenCNAME collapses a CNAME chain in out.Answer into a single A or AAAA
// record under the original query name. TTL is set to the chain minimum so
// the result expires no later than the shortest-lived link in the chain.
//
// Dynamically scans for the originating CNAME record natively, ensuring 
// out-of-order upstream responses are successfully flattened without evasion.
// No-ops when no terminal addresses are found or no matching CNAME is present.
func flattenCNAME(out *dns.Msg) {
	if len(out.Answer) < 2 {
		return
	}
	
	var cname *dns.CNAME
	
	queryName := ""
	if len(out.Question) > 0 {
		queryName = out.Question[0].Name
	}

	// Scan for the CNAME that matches the original query organically
	for _, rr := range out.Answer {
		if c, ok := rr.(*dns.CNAME); ok {
			if queryName == "" || strings.EqualFold(c.Hdr.Name, queryName) {
				cname = c
				if queryName == "" {
					queryName = c.Hdr.Name
				}
				break
			}
		}
	}
	
	if cname == nil {
		return // No matching CNAME found to flatten
	}

	minTTL := cname.Hdr.Ttl
	finals := make([]dns.RR, 0, len(out.Answer)-1)

	for _, rr := range out.Answer {
		switch r := rr.(type) {
		case *dns.CNAME:
			if r.Hdr.Ttl < minTTL {
				minTTL = r.Hdr.Ttl
			}
		case *dns.A:
			if r.Hdr.Ttl < minTTL {
				minTTL = r.Hdr.Ttl
			}
			cp := *r
			cp.Hdr.Name = queryName
			finals = append(finals, &cp)
		case *dns.AAAA:
			if r.Hdr.Ttl < minTTL {
				minTTL = r.Hdr.Ttl
			}
			cp := *r
			cp.Hdr.Name = queryName
			finals = append(finals, &cp)
		}
	}
	
	if len(finals) == 0 {
		return // Chain did not resolve to terminal IPs, leave intact
	}
	
	// Apply the chain-minimum TTL to every final record securely
	for _, rr := range finals {
		rr.Header().Ttl = minTTL
	}
	out.Answer = finals
}

// applyAnswerSort sorts the A and AAAA records in the Answer section natively.
// Supports "round-robin", "random", and "ip-sort". Returns true if the order was modified.
//
// [PERF 1.6.0] Deploys fixed-size 16-element stack buffers to organically bypass 
// slice GC allocations. Safe fallback applies standard `make` sizing if the limits are reached.
func applyAnswerSort(msg *dns.Msg, method string) bool {
	if method == "" || method == "none" || len(msg.Answer) < 2 {
		return false
	}
	
	type rrItem struct {
		idx int
		rr  dns.RR
	}

	// [PERF] Zero-allocation arrays effectively absorb 99% of normal responses
	var aRecordsBuf [16]rrItem
	var aaaaRecordsBuf [16]rrItem
	aRecords := aRecordsBuf[:0]
	aaaaRecords := aaaaRecordsBuf[:0]
	
	for i, rr := range msg.Answer {
		if rr.Header().Rrtype == dns.TypeA {
			aRecords = append(aRecords, rrItem{i, rr})
		} else if rr.Header().Rrtype == dns.TypeAAAA {
			aaaaRecords = append(aaaaRecords, rrItem{i, rr})
		}
	}
	
	changed := false
	
	processGroup := func(group []rrItem) {
		if len(group) < 2 {
			return
		}
		
		var origBuf [16]dns.RR
		var original []dns.RR
		if len(group) <= 16 {
			original = origBuf[:len(group)]
		} else {
			original = make([]dns.RR, len(group))
		}
		
		for i, item := range group {
			original[i] = item.rr
		}
		
		var sortBuf [16]dns.RR
		var sorted []dns.RR
		if len(group) <= 16 {
			sorted = sortBuf[:len(group)]
		} else {
			sorted = make([]dns.RR, len(group))
		}
		copy(sorted, original)
		
		switch method {
		case "round-robin":
			first := sorted[0]
			copy(sorted, sorted[1:])
			sorted[len(sorted)-1] = first
		case "random":
			rand.Shuffle(len(sorted), func(i, j int) {
				sorted[i], sorted[j] = sorted[j], sorted[i]
			})
		case "ip-sort":
			slices.SortStableFunc(sorted, func(a, b dns.RR) int {
				var ipI, ipJ netip.Addr
				if rr, ok := a.(*dns.A); ok {
					ipI, _ = netip.AddrFromSlice(rr.A)
				} else if rr, ok := a.(*dns.AAAA); ok {
					ipI, _ = netip.AddrFromSlice(rr.AAAA)
				}
				if rr, ok := b.(*dns.A); ok {
					ipJ, _ = netip.AddrFromSlice(rr.A)
				} else if rr, ok := b.(*dns.AAAA); ok {
					ipJ, _ = netip.AddrFromSlice(rr.AAAA)
				}
				return ipI.Compare(ipJ)
			})
		}
		
		groupChanged := false
		for i := range sorted {
			if sorted[i] != original[i] {
				groupChanged = true
				break
			}
		}
		
		if groupChanged {
			changed = true
			for i, item := range group {
				msg.Answer[item.idx] = sorted[i]
			}
		}
	}
	
	processGroup(aRecords)
	processGroup(aaaaRecords)
	
	return changed
}