package server

import (
	"bytes"
	"encoding/base64"
	"io"
	"net/http"
	"strings"
	"testing"

	"github.com/miekg/dns"
)

// The DoH handler is the one unauthenticated, TCP-terminated entry point in
// the process, so its size limits are a security control and not a nicety:
// AGENTS.md lists "amplification" under the areas that have been wrong before
// and must not lose their cap. Two paths, one cap each:
//
//   - POST: readCapped refuses a body over dns.MaxMsgSize, both from an honest
//     Content-Length and from a streamed body whose length is not declared.
//   - GET: the ?dns= parameter must be refused when it cannot possibly be a
//     DNS message — the encoding of a legal message is at most
//     base64.RawURLEncoding.EncodedLen(dns.MaxMsgSize) characters, and a DNS
//     message cannot exceed 65535 bytes on the wire (RFC 1035 §4.2.1, the
//     16-bit length prefix). Without that bound the only limit is net/http's
//     MaxHeaderBytes, a megabyte of attacker-chosen base64 decoded per
//     request, against a cap of 64 KiB on the POST path beside it.

// oversizedBody returns a body just over the cap. Its content is irrelevant:
// nothing past the cap may be read, let alone unpacked.
func oversizedBody() []byte { return bytes.Repeat([]byte{0x00}, dns.MaxMsgSize+1) }

// TestDoHPostBodyOverTheCapIsRefused covers the declared-length branch: an
// honest Content-Length over the cap is refused without reading the body.
func TestDoHPostBodyOverTheCapIsRefused(t *testing.T) {
	h := newHarness(t, nil)
	body := oversizedBody()
	res, err := h.dohClient(false).Post(h.dohURL(h.base()), contentTypeDNS, bytes.NewReader(body))
	if err != nil {
		t.Fatalf("POST: %v", err)
	}
	defer res.Body.Close()
	got, _ := io.ReadAll(res.Body)
	if res.StatusCode != http.StatusBadRequest {
		t.Errorf("status = %d, want 400: %s", res.StatusCode, got)
	}
	if n := len(h.fake.seen()); n != 0 {
		t.Errorf("an over-cap body reached the engine %d times", n)
	}
}

// TestDoHPostBodyOverTheCapIsRefusedWithoutAContentLength is the branch that
// matters: a client may lie about, or omit, Content-Length. readCapped has to
// detect "too large" by reading max+1 bytes itself, not by trusting the
// header, or the cap is a suggestion.
func TestDoHPostBodyOverTheCapIsRefusedWithoutAContentLength(t *testing.T) {
	h := newHarness(t, nil)
	// An unknown-length body makes net/http send it chunked, so the server
	// sees ContentLength == -1 and has to use its own reader.
	body := io.NopCloser(bytes.NewReader(oversizedBody()))
	req, err := http.NewRequest(http.MethodPost, h.dohURL(h.base()), body)
	if err != nil {
		t.Fatalf("request: %v", err)
	}
	req.ContentLength = -1
	req.Header.Set("Content-Type", contentTypeDNS)
	res, err := h.dohClient(false).Do(req)
	if err != nil {
		t.Fatalf("POST: %v", err)
	}
	defer res.Body.Close()
	got, _ := io.ReadAll(res.Body)
	if res.StatusCode != http.StatusBadRequest {
		t.Errorf("status = %d, want 400: %s", res.StatusCode, got)
	}
	if n := len(h.fake.seen()); n != 0 {
		t.Errorf("an over-cap chunked body reached the engine %d times", n)
	}
}

// TestDoHGETQueryStringOverTheCapIsRefused: a ?dns= parameter longer than any
// legal encoding is refused before it is decoded.
func TestDoHGETQueryStringOverTheCapIsRefused(t *testing.T) {
	h := newHarness(t, nil)
	cl := h.dohClient(false)
	// One character past the longest possible encoding of a DNS message.
	param := base64.RawURLEncoding.EncodeToString(make([]byte, dns.MaxMsgSize+1))
	for _, size := range []int{len(param), 512 * 1024} {
		long := param
		if size > len(param) {
			long = param + strings.Repeat("A", size-len(param))
		}
		res, err := cl.Get("https://" + h.addr("doh") + h.base() + "?dns=" + long)
		if err != nil {
			t.Fatalf("GET (%d chars): %v", len(long), err)
		}
		body, _ := io.ReadAll(res.Body)
		res.Body.Close()
		// 413 rather than 400: the parameter is a payload that is too big,
		// and it keeps the two paths symmetric — readCapped on POST answers
		// "cannot read body" for the same reason. A conforming client can
		// never reach this status: the cap is the longest possible encoding
		// of a legal DNS message.
		if res.StatusCode != http.StatusRequestEntityTooLarge {
			t.Errorf("GET with a %d-character dns parameter = %d, want 413: %s", len(long), res.StatusCode, body)
		}
	}
	if n := len(h.fake.seen()); n != 0 {
		t.Errorf("an over-cap query string reached the engine %d times", n)
	}
}

// TestDoHServesALargeButLegalQuery guards the cap from being set too tight. A
// resolver may legitimately be handed a message near the 65535-byte wire
// limit — here an EDNS0 padding option, the same trick the fake uses for a big
// answer — and a cap that rejected it would break real clients, not attackers.
func TestDoHServesALargeButLegalQuery(t *testing.T) {
	h := newHarness(t, nil)
	q := question("large.test.", dns.TypeA)
	q.SetEdns0(4096, false)
	// Size the padding so the packed message lands just under the 65535-byte
	// wire limit: a header (12) + question (18) + OPT (11) + the padding
	// option's own 4 bytes of code/length, plus the padding itself.
	q.IsEdns0().Option = append(q.IsEdns0().Option, &dns.EDNS0_NSID{Nsid: "00"})
	q.IsEdns0().Option = append(q.IsEdns0().Option, &dns.EDNS0_PADDING{Padding: make([]byte, 65400)})
	wire, err := q.Pack()
	if err != nil {
		t.Fatalf("pack: %v", err)
	}
	if len(wire) <= 65000 || len(wire) > dns.MaxMsgSize {
		t.Fatalf("fixture is %d bytes; it must be near the %d-byte limit and legal", len(wire), dns.MaxMsgSize)
	}
	if got := len(base64.RawURLEncoding.EncodeToString(wire)); got > base64.RawURLEncoding.EncodedLen(dns.MaxMsgSize) {
		t.Fatalf("fixture's encoding is %d characters, over the cap by construction", got)
	}

	// GET.
	res, body := h.doHGET(q)
	if res.StatusCode != http.StatusOK {
		t.Errorf("GET with a %d-byte message = %d, want 200: %s", len(wire), res.StatusCode, body)
	}
	// POST.
	res, body = h.doHPOST(q)
	if res.StatusCode != http.StatusOK {
		t.Errorf("POST with a %d-byte message = %d, want 200: %s", len(wire), res.StatusCode, body)
	}
	if n := len(h.fake.seen()); n != 2 {
		t.Errorf("engine saw %d requests, want 2", n)
	}
}