EvoBGP/internal/pipeline/doh_limits_test.go

package pipeline

import (
	"bytes"
	"compress/gzip"
	"context"
	"errors"
	"io"
	"net"
	"net/http"
	"net/http/httptest"
	"net/netip"
	"strings"
	"testing"
	"time"

	"evobgp/internal/httpclient"
	"github.com/miekg/dns"
)

type dohLimitRoundTrip func(*http.Request) (*http.Response, error)

func (f dohLimitRoundTrip) RoundTrip(r *http.Request) (*http.Response, error) { return f(r) }

type dohLimitBody struct {
	io.Reader
	read   int
	closed bool
}

func (b *dohLimitBody) Read(p []byte) (int, error) {
	n, e := b.Reader.Read(p)
	b.read += n
	return n, e
}
func (b *dohLimitBody) Close() error { b.closed = true; return nil }

func dohLimitCases(t *testing.T) []struct {
	name   string
	limit  int
	prefix []byte
	call   func(context.Context, *http.Client, string) ([]netip.Addr, error)
	pad    byte
} {
	t.Helper()
	msg := new(dns.Msg)
	msg.Response = true
	msg.Answer = []dns.RR{&dns.A{Hdr: dns.RR_Header{Name: "example.test.", Rrtype: dns.TypeA, Class: dns.ClassINET, Ttl: 60}, A: net.IPv4(192, 0, 2, 1)}}
	wire, err := msg.Pack()
	if err != nil {
		t.Fatal(err)
	}
	return []struct {
		name   string
		limit  int
		prefix []byte
		call   func(context.Context, *http.Client, string) ([]netip.Addr, error)
		pad    byte
	}{
		{"wire", 2 << 20, wire, func(ctx context.Context, hc *http.Client, url string) ([]netip.Addr, error) {
			return resolveDomainWithDOHMessage(ctx, hc, url, "example.test", dns.TypeA)
		}, 0},
		{"json", 1 << 20, []byte(`{"Answer":[{"type":1,"data":"192.0.2.1"}]}`), func(ctx context.Context, hc *http.Client, url string) ([]netip.Addr, error) {
			return resolveDomainWithDOHJSON(ctx, hc, url, "example.test", "A")
		}, ' '},
	}
}

func TestDOHCompleteResponseLimits(t *testing.T) {
	for _, tc := range dohLimitCases(t) {
		t.Run(tc.name, func(t *testing.T) {
			for _, known := range []bool{true, false} {
				for _, extra := range []int{0, 1} {
					data := append(append([]byte{}, tc.prefix...), bytes.Repeat([]byte{tc.pad}, tc.limit-len(tc.prefix)+extra)...)
					body := &dohLimitBody{Reader: bytes.NewReader(data)}
					hc := &http.Client{Transport: dohLimitRoundTrip(func(r *http.Request) (*http.Response, error) {
						if _, ok := r.Context().Deadline(); !ok {
							t.Error("missing deadline")
						}
						length := int64(-1)
						if known {
							length = int64(len(data))
						}
						return &http.Response{StatusCode: 200, ContentLength: length, Body: body, Header: make(http.Header)}, nil
					})}
					ips, err := tc.call(context.Background(), hc, "https://resolver.example/dns-query")
					if extra == 0 && (err != nil || len(ips) != 1) {
						t.Fatalf("exact limit rejected known=%v ips=%v: %v", known, ips, err)
					}
					if extra == 1 && !errors.Is(err, httpclient.ErrResponseTooLarge) {
						t.Fatalf("valid prefix bypassed size cap known=%v: %v", known, err)
					}
					if body.read > tc.limit+1 || !body.closed {
						t.Fatalf("read=%d closed=%v", body.read, body.closed)
					}
				}
			}
		})
	}
}

func TestDOHGzipDecodedLimit(t *testing.T) {
	for _, tc := range dohLimitCases(t) {
		t.Run(tc.name, func(t *testing.T) {
			server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				w.Header().Set("Content-Encoding", "gzip")
				gz := gzip.NewWriter(w)
				_, _ = gz.Write(tc.prefix)
				_, _ = gz.Write(bytes.Repeat([]byte{tc.pad}, tc.limit))
				_ = gz.Close()
			}))
			defer server.Close()
			if _, err := tc.call(context.Background(), server.Client(), server.URL); !errors.Is(err, httpclient.ErrResponseTooLarge) {
				t.Fatalf("gzip bypass: %v", err)
			}
		})
	}
}

func TestDOHNon200AndCancellationBounds(t *testing.T) {
	for _, tc := range dohLimitCases(t) {
		t.Run(tc.name, func(t *testing.T) {
			body := &dohLimitBody{Reader: strings.NewReader(strings.Repeat("x", 1<<20))}
			hc := &http.Client{Transport: dohLimitRoundTrip(func(*http.Request) (*http.Response, error) {
				return &http.Response{StatusCode: 400, Status: "400 Bad Request", Body: body, Header: make(http.Header)}, nil
			})}
			if _, err := tc.call(context.Background(), hc, "https://resolver.example/dns-query"); err == nil || body.read > 1024 || !body.closed {
				t.Fatalf("non200 read=%d closed=%v error=%v", body.read, body.closed, err)
			}
			ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)
			defer cancel()
			hc.Transport = dohLimitRoundTrip(func(r *http.Request) (*http.Response, error) { <-r.Context().Done(); return nil, r.Context().Err() })
			if _, err := tc.call(ctx, hc, "https://resolver.example/dns-query"); !errors.Is(err, context.DeadlineExceeded) {
				t.Fatalf("caller cancellation lost: %v", err)
			}
		})
	}
}

func TestDOHJSONRejectsTrailingObject(t *testing.T) {
	hc := &http.Client{Transport: dohLimitRoundTrip(func(*http.Request) (*http.Response, error) {
		return &http.Response{StatusCode: 200, ContentLength: -1, Body: io.NopCloser(strings.NewReader(`{"Answer":[]} {"ignored":true}`)), Header: make(http.Header)}, nil
	})}
	if _, err := resolveDomainWithDOHJSON(context.Background(), hc, "https://resolver.example/dns-query", "example.test", "A"); err == nil {
		t.Fatal("trailing JSON was ignored")
	}
}