upstream_net.go

/*
File:    upstream_net.go
Version: 1.4.0 (Split)
Updated: 21-Apr-2026 11:55 CEST

Description:
  TCP, DoT, HTTP, and QUIC stream network dialers and protocol implementations.
  Extracted from upstream.go.

Changes:
  1.4.0  - [SECURITY] Hardened DoH payload stream extraction. Introduced a strict
           length boundary check inside the `io.LimitReader` loop to sever connections 
           attempting to stream payloads exceeding the maximum DNS buffer size,
           preventing silent memory truncations.
  1.3.0  - [TRACEABILITY] Rigorously applied `%w` error wrapping across all 
           upstream network implementations (`exchangeHTTP`, `exchangeDoQ`, `exchangeStream`).
           This provides highly descriptive context to downstream handlers, allowing 
           the exact failure points (e.g., "pack buffer", "unpack response") to 
           be traced in the system logs natively.
  1.2.0  - [PERF] Tuned HTTP IdleConnTimeout to optimally balance TLS session resumption 
           pools and eliminate connection stagnation.
  1.1.0  - [LOGGING] Integrated httptrace.ClientTrace into exchangeHTTP to
           accurately extract the connected remote IP address for DoH and DoH3 logging.
*/

package main

import (
	"bytes"
	"context"
	"encoding/base64"
	"encoding/binary"
	"errors"
	"fmt"
	"io"
	"net/http"
	"net/http/httptrace"
	"time"

	"github.com/miekg/dns"
	"github.com/quic-go/quic-go"
)

// ---------------------------------------------------------------------------
// TCP / DoT connection reuse
// ---------------------------------------------------------------------------

func (u *Upstream) takeStreamConn(addr string) *dns.Conn {
	u.streamMu.Lock()
	entry := u.streamConns[addr]
	delete(u.streamConns, addr)
	u.streamMu.Unlock()

	if entry == nil {
		return nil
	}
	if time.Since(entry.idleAt) > streamIdleMax {
		entry.conn.Close()
		return nil
	}
	return entry.conn
}

func (u *Upstream) putStreamConn(addr string, conn *dns.Conn) {
	u.streamMu.Lock()
	if old := u.streamConns[addr]; old != nil {
		old.conn.Close()
	}
	u.streamConns[addr] = &streamConnEntry{conn: conn, idleAt: time.Now()}
	u.streamMu.Unlock()
}

func (u *Upstream) dialStream(addr, tlsHost string) (*dns.Conn, error) {
	if u.Proto == "dot" {
		tlsConf := u.baseTLSConf.Clone()
		if tlsHost != "" {
			tlsConf.ServerName = tlsHost
		}
		conn, err := dns.DialTimeoutWithTLS("tcp-tls", addr, tlsConf, streamTimeout)
		if err != nil {
			return nil, fmt.Errorf("dot dial: %w", err)
		}
		return conn, nil
	}
	
	conn, err := dns.DialTimeout("tcp", addr, streamTimeout)
	if err != nil {
		return nil, fmt.Errorf("tcp dial: %w", err)
	}
	return conn, nil
}

func (u *Upstream) exchangeStream(ctx context.Context, req *dns.Msg, dialAddr, tlsHost string) (*dns.Msg, error) {
	deadline, ok := ctx.Deadline()
	if !ok {
		deadline = time.Now().Add(streamTimeout)
	}

	if conn := u.takeStreamConn(dialAddr); conn != nil {
		conn.SetDeadline(deadline)
		if err := conn.WriteMsg(req); err == nil {
			if resp, err := conn.ReadMsg(); err == nil && resp != nil {
				u.putStreamConn(dialAddr, conn)
				return resp, nil
			}
		}
		conn.Close()
	}

	conn, err := u.dialStream(dialAddr, tlsHost)
	if err != nil {
		return nil, err // Already wrapped natively in dialStream
	}
	
	conn.SetDeadline(deadline)
	
	if err := conn.WriteMsg(req); err != nil {
		conn.Close()
		return nil, fmt.Errorf("write msg: %w", err)
	}
	
	resp, err := conn.ReadMsg()
	if err != nil {
		conn.Close()
		return nil, fmt.Errorf("read msg: %w", err)
	}
	
	if resp == nil {
		conn.Close()
		return nil, errors.New("empty response")
	}
	
	u.putStreamConn(dialAddr, conn)
	return resp, nil
}

// ---------------------------------------------------------------------------
// HTTP (DoH / DoH3)
// ---------------------------------------------------------------------------

// exchangeHTTP sends a DNS query over HTTP/2 (DoH) or HTTP/3 (DoH3).
//
// Returns the resolved DNS message and the underlying IP address (remoteAddr)
// connected to via httptrace.ClientTrace if successfully triggered.
//
// HTTP method selection (u.useGET, set once at ParseUpstream time):
//
//	POST (default, useGET=false):
//	  Body is the raw packed DNS message. bytes.NewReader aliases the pool buffer,
//	  so the buffer must NOT be returned to the pool until after client.Do returns
//	  (defer handles this).
//
//	GET (useGET=true, RFC 8484 §4.1):
//	  Packed message is base64url-encoded (no padding) and appended as ?dns=...
//	  EncodeToString returns an independent string copy, so the pool buffer is
//	  released immediately after encoding — before the HTTP round-trip starts.
//	  Responses may be cached by HTTP intermediaries.
func (u *Upstream) exchangeHTTP(ctx context.Context, req *dns.Msg, targetURL string, client *http.Client) (*dns.Msg, string, error) {
	bufPtr := smallBufPool.Get().(*[]byte)
	packed, err := req.PackBuffer((*bufPtr)[:0])
	if err != nil {
		smallBufPool.Put(bufPtr)
		return nil, "", fmt.Errorf("pack buffer: %w", err)
	}

	// Capture the actual remote address successfully dialed for enhanced telemetry.
	var remoteAddr string
	trace := &httptrace.ClientTrace{
		GotConn: func(connInfo httptrace.GotConnInfo) {
			if connInfo.Conn != nil {
				remoteAddr = connInfo.Conn.RemoteAddr().String()
			}
		},
	}
	ctx = httptrace.WithClientTrace(ctx, trace)

	var hReq *http.Request
	if u.useGET {
		encoded := base64.RawURLEncoding.EncodeToString(packed)
		smallBufPool.Put(bufPtr)
		hReq, err = http.NewRequestWithContext(ctx, http.MethodGet, targetURL+"?dns="+encoded, nil)
		if err != nil {
			return nil, "", fmt.Errorf("new http get request: %w", err)
		}
		hReq.Header.Set("Accept", "application/dns-message")
	} else {
		defer smallBufPool.Put(bufPtr)
		hReq, err = http.NewRequestWithContext(ctx, http.MethodPost, targetURL, bytes.NewReader(packed))
		if err != nil {
			return nil, "", fmt.Errorf("new http post request: %w", err)
		}
		hReq.Header.Set("Content-Type", "application/dns-message")
		hReq.Header.Set("Accept", "application/dns-message")
	}

	resp, err := client.Do(hReq)
	if err != nil {
		return nil, remoteAddr, fmt.Errorf("http request: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return nil, remoteAddr, fmt.Errorf("bad status code: %d", resp.StatusCode)
	}

	rBufPtr := largeBufPool.Get().(*[]byte)
	defer largeBufPool.Put(rBufPtr)
	rBuf := *rBufPtr

	n := 0
	lr := io.LimitReader(resp.Body, int64(len(rBuf)))
	for {
		c, readErr := lr.Read(rBuf[n:])
		n += c
		if readErr == io.EOF {
			break
		}
		if readErr != nil {
			return nil, remoteAddr, fmt.Errorf("read body: %w", readErr)
		}
		// [SECURITY] Protect against maliciously sized or infinitely streaming payloads.
		// If the buffer reaches maximum capacity without encountering an EOF, sever the 
		// connection to prevent silent unpacking of truncated, corrupted DNS payloads.
		if n == len(rBuf) {
			return nil, remoteAddr, errors.New("security constraint: response payload exceeded maximum buffer capacity")
		}
	}

	dnsResp := new(dns.Msg)
	if err := dnsResp.Unpack(rBuf[:n]); err != nil {
		return nil, remoteAddr, fmt.Errorf("unpack response: %w", err)
	}
	return dnsResp, remoteAddr, nil
}

// ---------------------------------------------------------------------------
// QUIC (DoQ)
// ---------------------------------------------------------------------------

// exchangeDoQ tries a cached QUIC connection first, dials fresh on miss/stale,
// and stores the connection in the pool ONLY after a successful exchange.
// It leverages QUIC stream multiplexing by retaining the active connection in the pool.
//
// host is the fully expanded hostname ({client-name} already substituted).
// It is used as both the TLS SNI and the connection-pool key.
func (u *Upstream) exchangeDoQ(ctx context.Context, req *dns.Msg, host string) (*dns.Msg, string, error) {
	key := host

	// Check for a healthy, pooled multiplexed connection
	if entry := u.getDoQConn(key); entry != nil {
		resp, err := u.doqStreamExchange(ctx, entry.conn, req)
		if err == nil {
			u.updateDoQIdle(key, entry)
			return resp, entry.dialAddr, nil
		}
		// Conn failed, evict it. Next parallel query will trigger dial
		u.removeDoQConn(key, entry, "stale")
	}

	// Use singleflight to prevent Thundering Herd. Parallel cache-misses will wait
	// for the single winner to dial the DoQ connection and return it.
	resI, err, _ := u.doqDialGroup.Do(key, func() (any, error) {
		// Verify pool again (a parallel goroutine might have just populated it)
		if entry := u.getDoQConn(key); entry != nil {
			return entry, nil
		}
		newConn, dialAddr, err := u.dialDoQ(ctx, host, u.dialAddrs)
		if err != nil {
			return nil, err // Wrapped natively inside dialDoQ
		}

		entry := &doqConnEntry{conn: newConn, dialAddr: dialAddr, idleAt: time.Now()}
		u.doqMu.Lock()
		if old := u.doqConns[key]; old != nil && old.conn != newConn {
			old.conn.CloseWithError(0, "replaced")
		}
		u.doqConns[key] = entry
		u.doqMu.Unlock()
		return entry, nil
	})

	if err != nil {
		return nil, u.RawURL, err
	}

	entry := resI.(*doqConnEntry)
	resp, err := u.doqStreamExchange(ctx, entry.conn, req)
	if err != nil {
		u.removeDoQConn(key, entry, "exchange failed")

		// RFC 9250 §10.5: latch doqNo0RTT on DOQ_PROTOCOL_ERROR, retry once
		// without 0-RTT. This handles strict servers that reject early data.
		if !u.doqNo0RTT.Swap(true) {
			// Retry through singleflight with a unique retry key
			resI, retryErr, _ := u.doqDialGroup.Do(key+"_retry", func() (any, error) {
				retryConn, retryAddr, err := u.dialDoQ(ctx, host, u.dialAddrs)
				if err != nil {
					return nil, err // Wrapped natively inside dialDoQ
				}
				e := &doqConnEntry{conn: retryConn, dialAddr: retryAddr, idleAt: time.Now()}
				u.doqMu.Lock()
				if old := u.doqConns[key]; old != nil && old.conn != retryConn {
					old.conn.CloseWithError(0, "replaced")
				}
				u.doqConns[key] = e
				u.doqMu.Unlock()
				return e, nil
			})

			if retryErr == nil {
				retryEntry := resI.(*doqConnEntry)
				resp, err = u.doqStreamExchange(ctx, retryEntry.conn, req)
				if err == nil {
					u.updateDoQIdle(key, retryEntry)
					return resp, retryEntry.dialAddr, nil
				}
				u.removeDoQConn(key, retryEntry, "retry failed")
			}
		}
		return nil, entry.dialAddr, err // Return the wrapped streamExchange error directly
	}

	u.updateDoQIdle(key, entry)
	return resp, entry.dialAddr, nil
}

// getDoQConn fetches the active QUIC connection without deleting it from the pool.
// This allows true stream multiplexing across concurrent DNS queries.
func (u *Upstream) getDoQConn(key string) *doqConnEntry {
	u.doqMu.Lock()
	entry := u.doqConns[key]
	u.doqMu.Unlock()

	if entry == nil {
		return nil
	}
	if time.Since(entry.idleAt) > doqIdleMax {
		u.removeDoQConn(key, entry, "idle timeout")
		return nil
	}
	return entry
}

func (u *Upstream) updateDoQIdle(key string, entry *doqConnEntry) {
	u.doqMu.Lock()
	// Only update if it's still the reigning connection for this key
	if u.doqConns[key] == entry {
		entry.idleAt = time.Now()
	}
	u.doqMu.Unlock()
}

func (u *Upstream) removeDoQConn(key string, entry *doqConnEntry, reason string) {
	u.doqMu.Lock()
	if u.doqConns[key] == entry {
		delete(u.doqConns, key)
	}
	u.doqMu.Unlock()
	entry.conn.CloseWithError(0, reason)
}

// dialDoQ establishes a new QUIC connection for DoQ (RFC 9250).
// Tries each address in addrs in order; returns on first success.
// 0-RTT is disabled permanently after a DOQ_PROTOCOL_ERROR (u.doqNo0RTT).
func (u *Upstream) dialDoQ(ctx context.Context, host string, addrs []string) (*quic.Conn, string, error) {
	tlsConf := getHardenedTLSConfig()
	tlsConf.NextProtos = []string{"doq"}
	if host != "" {
		tlsConf.ServerName = host
	}

	qConf := &quic.Config{
		MaxIdleTimeout:     doqIdleMax,
		MaxIncomingStreams: 0, // client-only: we never accept streams from the server
	}

	var lastErr error
	for _, addr := range addrs {
		var conn *quic.Conn
		if !u.doqNo0RTT.Load() {
			qConf.Allow0RTT = true
			conn, lastErr = quic.DialAddrEarly(ctx, addr, tlsConf, qConf)
		} else {
			qConf.Allow0RTT = false
			conn, lastErr = quic.DialAddr(ctx, addr, tlsConf, qConf)
		}
		if lastErr == nil {
			return conn, addr, nil
		}
	}
	
	if lastErr != nil {
		return nil, "", fmt.Errorf("doq: all dial addresses failed for %s: %w", host, lastErr)
	}
	return nil, "", fmt.Errorf("doq: all dial addresses failed for %s", host)
}

// doqStreamExchange opens a single QUIC stream, writes the DNS query with a
// 2-byte length prefix (RFC 9250 §4.2), closes the write side first (FIN),
// then reads the response. stream.Close() before ReadFull is required —
// strict servers wait for the write-side FIN before sending the response.
//
// Both the pack buffer and the response read buffer use smallBufPool for
// typical DNS sizes (≤4096 bytes), eliminating two heap allocations per DoQ
// upstream query. Oversized messages fall back to a direct allocation.
// dns.Msg.Unpack copies all fields from its input so the pool buffer is
// safely released immediately after Unpack returns.
func (u *Upstream) doqStreamExchange(ctx context.Context, conn *quic.Conn, req *dns.Msg) (*dns.Msg, error) {
	stream, err := conn.OpenStreamSync(ctx)
	if err != nil {
		return nil, fmt.Errorf("open stream: %w", err)
	}
	defer stream.CancelRead(0) // clean up read side on any error path

	// Prevent stream operations from leaking goroutines when context cancels
	deadline, ok := ctx.Deadline()
	if !ok {
		deadline = time.Now().Add(streamTimeout)
	}
	stream.SetDeadline(deadline)

	pBufPtr := smallBufPool.Get().(*[]byte)
	packed, err := req.PackBuffer((*pBufPtr)[:0])
	if err != nil {
		smallBufPool.Put(pBufPtr)
		return nil, fmt.Errorf("pack buffer: %w", err)
	}

	// Write the 2-byte length prefix and packed payload separately.
	var lenBuf [2]byte
	binary.BigEndian.PutUint16(lenBuf[:], uint16(len(packed)))
	if _, err := stream.Write(lenBuf[:]); err != nil {
		smallBufPool.Put(pBufPtr)
		return nil, fmt.Errorf("write length: %w", err)
	}
	if _, err := stream.Write(packed); err != nil {
		smallBufPool.Put(pBufPtr)
		return nil, fmt.Errorf("write payload: %w", err)
	}
	// Send write-side FIN BEFORE reading — strict DoQ servers stall otherwise.
	if err := stream.Close(); err != nil {
		smallBufPool.Put(pBufPtr)
		return nil, fmt.Errorf("close write stream: %w", err)
	}
	smallBufPool.Put(pBufPtr) // packed fully written; pool buffer safely released

	// Read the 2-byte response length prefix.
	if _, err := io.ReadFull(stream, lenBuf[:]); err != nil {
		return nil, fmt.Errorf("read length: %w", err)
	}
	respLen := int(binary.BigEndian.Uint16(lenBuf[:]))
	if respLen == 0 {
		return nil, errors.New("doq: zero-length response")
	}

	// Use smallBufPool for responses that fit (≤4096 bytes). Fall back to a
	// direct allocation for the rare oversized case (e.g. DNSKEY responses).
	var (
		rBufPtr *[]byte
		respBuf []byte
	)
	if respLen <= 4096 {
		rBufPtr = smallBufPool.Get().(*[]byte)
		respBuf = (*rBufPtr)[:respLen]
	} else {
		respBuf = make([]byte, respLen)
	}

	_, readErr := io.ReadFull(stream, respBuf)
	resp := new(dns.Msg)
	var unpackErr error
	if readErr == nil {
		unpackErr = resp.Unpack(respBuf)
	}
	
	if rBufPtr != nil {
		smallBufPool.Put(rBufPtr) // Unpack copied all fields; safe to release now
	}
	
	if readErr != nil {
		return nil, fmt.Errorf("read payload: %w", readErr)
	}
	if unpackErr != nil {
		return nil, fmt.Errorf("unpack response: %w", unpackErr)
	}
	
	return resp, nil
}