package server

import (
	"bytes"
	"crypto/tls"
	"encoding/base64"
	"errors"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"github.com/miekg/dns"
)

// countingReader counts the bytes readCapped actually pulled out of a body, so
// the "refused on the header" branch can be told apart from the "refused while
// reading" branch. Both return an error; only one of them let the client send
// the body.
type countingReader struct {
	r io.Reader
	n int
}

func (c *countingReader) Read(p []byte) (int, error) {
	n, err := c.r.Read(p)
	c.n += n
	return n, err
}

// errReader emits a prefix and then fails, standing in for a client that
// disconnects mid-body. readCapped has to surface that as an error: a truncated
// DNS message handed to dns.Msg.Unpack would be unpacked as garbage.
type errReader struct {
	prefix string
	err    error
	done   bool
}

func (f *errReader) Read(p []byte) (int, error) {
	if f.done {
		return 0, f.err
	}
	f.done = true
	return copy(p, f.prefix), nil
}

// TestReadCapped covers the DoH POST body reader. It is a security control, not
// a convenience: the handler is unauthenticated, and the bound is what stops a
// client from making the server allocate an arbitrary amount of memory per
// request. Two branches have to both work, because the whole cap is only as
// strong as the weaker one:
//
//   - a declared Content-Length over the cap is refused WITHOUT reading the
//     body at all (the header is a fast path, and the counting reader proves
//     nothing was read);
//   - a missing or understated Content-Length is caught by the reader itself,
//     which takes max+1 bytes and only then says "too large".
func TestReadCapped(t *testing.T) {
	const max = 16
	full := strings.Repeat("a", max)
	cases := []struct {
		name string
		// body is what the client sends.
		body string
		// failAfter, when non-zero, makes the body stop with an error after
		// that many bytes.
		failAfter int
		// length is the Content-Length the request advertises. -1 is
		// net/http's value for a chunked request: no declared length at all.
		length int64
		want   string
		// wantReads is the number of body bytes readCapped should have
		// consumed. -1 means the case is not about how much was read.
		wantReads int
		wantErr   bool
	}{
		{
			name: "an empty body is legal", body: "", length: 0,
			want: "", wantReads: 0,
		},
		{
			name: "a body exactly at the cap is accepted whole", body: full, length: max,
			want: full, wantReads: max,
		},
		{
			name: "one byte over the cap with an honest Content-Length is refused before the body is read",
			body: full + "a", length: max + 1,
			want: "", wantReads: 0, wantErr: true,
		},
		{
			// A client may claim any length it likes, so an over-declared
			// length must not be the only thing that decides. It is still
			// the fast path, so it must not cost a read.
			name: "a wildly over-declared Content-Length is refused before the body is read",
			body: full + "a", length: 64 << 20,
			want: "", wantReads: 0, wantErr: true,
		},
		{
			// ContentLength == -1 is what net/http reports for a chunked
			// request, so this is the branch a length-less client hits.
			name: "one byte over the cap with no declared length is caught while reading",
			body: full + "a", length: -1,
			want: "", wantReads: max + 1, wantErr: true,
		},
		{
			name: "a body under the cap with no declared length is accepted whole",
			body: "chunked-body", length: -1,
			want: "chunked-body", wantReads: len("chunked-body"),
		},
		{
			// The cap is not the only thing that can end a read early, and a
			// half-read body must not be handed on as if it were complete.
			name: "a body that fails mid-read surfaces the read error",
			body: "abcde", failAfter: 5, length: -1,
			want: "", wantReads: 5, wantErr: true,
		},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			var src io.Reader = strings.NewReader(tc.body)
			if tc.failAfter > 0 {
				src = &errReader{prefix: tc.body, err: errors.New("connection reset by peer")}
			}
			cr := &countingReader{r: src}
			r := httptest.NewRequest(http.MethodPost, "/dns-query", nil)
			r.Body = io.NopCloser(cr)
			r.ContentLength = tc.length

			got, err := readCapped(r, max)
			if tc.wantErr {
				if err == nil {
					t.Fatalf("readCapped accepted %q, want an error", got)
				}
				if got != nil {
					t.Errorf("readCapped returned %q alongside its error, want nil", got)
				}
			} else {
				if err != nil {
					t.Fatalf("readCapped(%q) = %v, want it accepted", tc.want, err)
				}
				if string(got) != tc.want {
					t.Errorf("readCapped = %q, want %q", got, tc.want)
				}
			}
			if tc.wantReads >= 0 && cr.n != tc.wantReads {
				t.Errorf("readCapped read %d body bytes, want %d", cr.n, tc.wantReads)
			}
		})
	}
}

// TestDecodeBase64URL covers the RFC 8484 §4.1 GET parameter decoder. The two
// encodings it accepts are base64url with and without padding: the RFC asks for
// the unpadded form, every real client sends the padded one, and refusing the
// padded form would break all of them. What it must NOT accept is the standard
// alphabet, because '+' and '/' have to be percent-encoded in a query string
// and a decoder that took them would be a decoder for a different encoding.
func TestDecodeBase64URL(t *testing.T) {
	msg := []byte{0xde, 0xad, 0xbe, 0xef, 0x00, 0x01, 0x02, 0x03}
	// Standard base64 of these bytes contains both '+' and '/'.
	std := base64.StdEncoding.EncodeToString(msg)
	if !strings.ContainsAny(std, "+/") {
		t.Fatalf("fixture %q contains no + or /; the standard-alphabet case would prove nothing", std)
	}
	cases := []struct {
		name    string
		in      string
		want    []byte
		wantErr bool
	}{
		{name: "unpadded base64url", in: base64.RawURLEncoding.EncodeToString(msg), want: msg},
		{name: "padded base64url", in: base64.URLEncoding.EncodeToString(msg), want: msg},
		{name: "standard base64 with + and / is rejected", in: std, wantErr: true},
		{name: "a space inside the parameter is rejected", in: "QQ BB", wantErr: true},
		{name: "a tab inside the parameter is rejected", in: "QQ\tBB", wantErr: true},
		{
			// encoding/base64 skips CR and LF in every encoding, padded or
			// not. A newline in ?dns= therefore decodes rather than 400s.
			// It is harmless — the bytes are the same message — and it is
			// pinned here so nobody "fixes" it by assuming the decoder
			// rejects all whitespace. The size cap is applied to the raw
			// parameter before decoding, so this cannot be used to smuggle
			// an oversized message past the cap.
			name: "a newline inside the parameter is ignored by the decoder",
			in:   "QQ\n==", want: []byte{0x41},
		},
		{name: "an empty parameter decodes to nothing", in: "", want: []byte{}},
		{name: "a truncated quantum is rejected", in: "QQQQQ", wantErr: true},
		{name: "a single character is rejected", in: "A", wantErr: true},
		{
			// Two characters is a legal UNPADDED encoding of one byte, so it
			// decodes on the RawURLEncoding path before the padded one is
			// tried. It is not an error case; the case above ("QQQQQ") is,
			// because five characters is not a whole number of quanta.
			name: "two characters are a valid unpadded encoding of one byte",
			in:   "QQ", want: []byte{0x41},
		},
		{
			// decodeBase64URL is a pure decoder: it imposes no size limit,
			// because the limit that matters (the longest possible encoding
			// of a 65535-byte DNS message) is a property of the caller's
			// parameter, and the DoH handler applies it before calling here.
			// The bound itself is pinned in TestDoHGETQueryStringOverTheCapIsRefused.
			name: "a long but legal parameter decodes; the cap is the handler's job",
			in:   base64.RawURLEncoding.EncodeToString(bytes.Repeat([]byte{0x41}, 1<<20)),
			want: bytes.Repeat([]byte{0x41}, 1<<20),
		},
	}
	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			got, err := decodeBase64URL(tc.in)
			if tc.wantErr {
				if err == nil {
					t.Fatalf("decodeBase64URL(%q) = %x, want an error", tc.in, got)
				}
				return
			}
			if err != nil {
				t.Fatalf("decodeBase64URL(%q) = %v, want it accepted", tc.in, err)
			}
			if !bytes.Equal(got, tc.want) {
				t.Errorf("decodeBase64URL(%q) = %x, want %x", tc.in, got, tc.want)
			}
		})
	}
	// The cases above are the decoder's edge cases; this is the case that
	// actually happens. A real query must survive the round trip byte for
	// byte, in both encodings, or a client using the padded form (which is
	// what curl and most browsers send) gets a 400 for a well-formed query.
	wire, err := question("roundtrip.test.", dns.TypeA).Pack()
	if err != nil {
		t.Fatalf("pack: %v", err)
	}
	for _, enc := range []struct {
		name string
		in   string
	}{
		{"unpadded", base64.RawURLEncoding.EncodeToString(wire)},
		{"padded", base64.URLEncoding.EncodeToString(wire)},
	} {
		t.Run("a real query round trips, "+enc.name, func(t *testing.T) {
			got, err := decodeBase64URL(enc.in)
			if err != nil {
				t.Fatalf("decodeBase64URL: %v", err)
			}
			if !bytes.Equal(got, wire) {
				t.Fatalf("round trip changed the message: %d bytes in, %d out", len(wire), len(got))
			}
			if m := new(dns.Msg); m.Unpack(got) != nil {
				t.Fatalf("the decoded bytes are not a DNS message")
			}
		})
	}
}

// TestParseTLSVersion pins the fallback, which is a security setting: an
// operator who typos min_version, or leaves it empty, must get TLS 1.2 — never
// TLS 1.0 and never "whatever the Go default happens to be this release".
func TestParseTLSVersion(t *testing.T) {
	cases := []struct {
		in   string
		want uint16
		why  string
	}{
		{in: "1.2", want: tls.VersionTLS12, why: "the documented value"},
		{in: "1.3", want: tls.VersionTLS13, why: "the documented value"},
		{in: "1.0", want: tls.VersionTLS10, why: "still reachable, though config documents only 1.2/1.3"},
		{in: "1.1", want: tls.VersionTLS12, why: "unknown values fail UP to 1.2, never down to 1.1"},
		{in: "", want: tls.VersionTLS12, why: "an unset min_version means the secure default"},
		{in: "junk", want: tls.VersionTLS12, why: "a typo must not weaken the floor"},
		{in: "13", want: tls.VersionTLS12, why: "only the exact string 1.3 is honoured"},
		{in: "1.3 ", want: tls.VersionTLS12, why: "no trimming: an odd value falls back, it does not panic"},
	}
	for _, tc := range cases {
		t.Run("min_version="+tc.in, func(t *testing.T) {
			if got := parseTLSVersion(tc.in); got != tc.want {
				t.Errorf("parseTLSVersion(%q) = %d, want %d (%s)", tc.in, got, tc.want, tc.why)
			}
		})
	}
}

// TestNormaliseBase covers the DoH base path normaliser. What it does is
// deliberately narrow — add the leading slash, trim the trailing ones — because
// it runs on operator config at Start and must never fail a startup. What it
// must NOT grow is a second parser: the path-traversal cases are here to pin
// that the string is passed through verbatim, because the thing that actually
// keeps a traversal out of the handler is net/http's ServeMux, which cleans
// the path before it matches. That is pinned by
// TestDoHBasePathTraversalNeverReachesTheHandler below; a sanitiser here would
// be a second, weaker implementation of the same rule.
func TestNormaliseBase(t *testing.T) {
	cases := []struct{ in, want, why string }{
		{in: "", want: "/dns-query", why: "the default, so an unset base path still serves DoH"},
		{in: "query", want: "/query", why: "a path written without its leading slash is what operators type"},
		{in: "/query", want: "/query", why: "already normalised"},
		{in: "/query/", want: "/query", why: "a trailing slash would register a different mux pattern"},
		{in: "/query///", want: "/query", why: "all of them, and never the empty string"},
		{in: "/", want: "/", why: "serve at the root is a legal (if blunt) choice"},
		{in: "///", want: "/", why: "trimming stops at a single slash"},
		{in: "///query///", want: "///query", why: "only a MISSING leading slash is added; runs of them are the mux's business"},
		{in: "/a//b", want: "/a//b", why: "only the ends are touched; an inner // is left for the mux"},
		{in: "/query?x=1", want: "/query?x=1", why: "no query splitting: this is a path, not a URL"},
		{in: "/../etc", want: "/../etc", why: "passed through; ServeMux cleans the path (see the traversal test)"},
		{in: "/dns-query/..", want: "/dns-query/..", why: "passed through, same reason"},
	}
	for _, tc := range cases {
		t.Run(tc.in, func(t *testing.T) {
			if got := normaliseBase(tc.in); got != tc.want {
				t.Errorf("normaliseBase(%q) = %q, want %q (%s)", tc.in, got, tc.want, tc.why)
			}
		})
	}
}

// TestDoHBasePathTraversalNeverReachesTheHandler is the security half of the
// normaliseBase comment above. The DoH mux is built from a configured path, so
// the question is what happens when a request arrives with dot segments in it:
// it must never be answered by the DNS handler, and it must never reach the
// engine. net/http's ServeMux cleans the path and issues a 307 to the cleaned
// form, which is why there is no sanitiser in normaliseBase. Redirects are not
// followed here: following one would answer this test with the redirect
// TARGET's behaviour rather than the traversal's.
func TestDoHBasePathTraversalNeverReachesTheHandler(t *testing.T) {
	h := newHarness(t, nil)
	cl := noRedirectClient(t, h)
	cases := []struct {
		path string
		want int
	}{
		// Cleaned to /healthz, which is a registered path: a 307 to it, and
		// the DNS handler never sees the request. ServeMux is what keeps the
		// traversal out; normaliseBase deliberately does not try.
		{path: "/dns-query/../healthz", want: http.StatusTemporaryRedirect},
		{path: "/dns-query/./x", want: http.StatusTemporaryRedirect},
		// Encoded dot segments are not cleaned by the mux, so they match no
		// pattern at all.
		{path: "/dns-query/%2e%2e/healthz", want: http.StatusNotFound},
		{path: "/dns-query/..%2fhealthz", want: http.StatusNotFound},
		{path: "/..%2f..%2fetc/passwd", want: http.StatusNotFound},
	}
	for _, tc := range cases {
		t.Run(tc.path, func(t *testing.T) {
			res, err := cl.Get(h.dohURL(tc.path))
			if err != nil {
				t.Fatalf("GET %s: %v", tc.path, err)
			}
			_ = res.Body.Close()
			if res.StatusCode != tc.want {
				t.Errorf("GET %s = %d, want %d", tc.path, res.StatusCode, tc.want)
			}
		})
	}
	// Nothing above reached the engine, and the plain path still works, so
	// the 404s are not the DNS handler being broken. Exactly one request
	// reaches the engine: the control query.
	h.fake.reset()
	if a := answerAddr(h.askDoHGET(question("control.test.", dns.TypeA))); a != "203.0.113.9" {
		t.Errorf("control query answered %q, want 203.0.113.9", a)
	}
	if n := len(h.fake.seen()); n != 1 {
		t.Errorf("%d requests reached the engine, want exactly the control query: %+v", n, h.fake.seen())
	}
}

// noRedirectClient is h.dohClient without redirect following, so a test can
// assert on the status the server produced rather than on whatever the
// Location header points at.
func noRedirectClient(t *testing.T, h *harness) *http.Client {
	t.Helper()
	tr := &http.Transport{TLSClientConfig: h.clientTLS()}
	t.Cleanup(tr.CloseIdleConnections)
	return &http.Client{
		Transport: tr,
		Timeout:   5 * time.Second,
		CheckRedirect: func(*http.Request, []*http.Request) error {
			return http.ErrUseLastResponse
		},
	}
}