smallwebwaf/internal/proxy/limits_test.go

package proxy_test

import (
	"bytes"
	"errors"
	"io"
	"net/http"
	"strconv"
	"sync/atomic"
	"testing"

	"sneak.berlin/go/smallwebwaf/internal/requestlog"
)

// sizeLimit is the size limit the tests set, 1K as a setting.
const (
	sizeLimit        = 1 << 10
	sizeLimitSetting = "1K"
)

func TestRequestBodyLimit(t *testing.T) {
	t.Parallel()

	for _, tc := range []struct {
		name string
		size int
		// announced sends the size in Content-Length; otherwise the body
		// is sent in chunks with no length given.
		announced bool
		want      int
		action    string
		// refusedBeforeApp is a refusal before anything reaches the app.
		// A body over the limit with no length given has already partly
		// reached the app when it is refused.
		refusedBeforeApp bool
	}{
		{"announced, over the limit", 2 * sizeLimit, true,
			http.StatusRequestEntityTooLarge, requestlog.ActionTooLarge, true},
		{"announced, at the limit", sizeLimit, true,
			http.StatusOK, requestlog.ActionForward, false},
		{"not announced, over the limit", 4 * sizeLimit, false,
			http.StatusRequestEntityTooLarge, requestlog.ActionTooLarge, false},
		{"not announced, at the limit", sizeLimit, false,
			http.StatusOK, requestlog.ActionForward, false},
	} {
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()

			var calls atomic.Int32

			app := startApp(t, func(_ http.ResponseWriter, r *http.Request) {
				calls.Add(1)

				_, _ = io.Copy(io.Discard, r.Body)
			})
			addr, out := startProxy(t, app.URL, map[string]string{
				requestMaxBytes: sizeLimitSetting,
				metricsToken:    token,
			})

			var body io.Reader = bytes.NewReader(make([]byte, tc.size))
			if !tc.announced {
				body = io.MultiReader(body) // hides the length
			}

			wantStatus(t, do(t, newRequest(t, http.MethodPost, addr, "/upload", body)),
				tc.want)
			wantLine(t, out.requestLine(t), tc.want, tc.action)

			hits := 0
			if tc.action == requestlog.ActionTooLarge {
				hits = 1
			}

			wantLimitHits(t, addr, requestMaxBytes, hits)

			if tc.refusedBeforeApp && calls.Load() != 0 {
				t.Errorf("the app was called %d times, want never", calls.Load())
			}
		})
	}
}

func TestResponseBodyLimit(t *testing.T) {
	t.Parallel()

	for _, tc := range []struct {
		name string
		size int
		// announced sends the size in Content-Length; otherwise the body
		// is sent in chunks with no length given.
		announced bool
		want      int
		action    string
		// received is how much of a body the client gets, and cutOff
		// whether the connection is then cut.
		received int
		cutOff   bool
	}{
		{"announced, over the limit", 2 * sizeLimit, true, http.StatusBadGateway,
			requestlog.ActionTooLarge, len("Bad Gateway\n"), false},
		{"announced, at the limit", sizeLimit, true, http.StatusOK,
			requestlog.ActionForward, sizeLimit, false},
		{"not announced, over the limit", 4 * sizeLimit, false, http.StatusOK,
			requestlog.ActionTooLarge, sizeLimit, true},
		{"not announced, at the limit", sizeLimit, false, http.StatusOK,
			requestlog.ActionForward, sizeLimit, false},
	} {
		t.Run(tc.name, func(t *testing.T) {
			t.Parallel()

			app := startApp(t, func(w http.ResponseWriter, _ *http.Request) {
				answerWithSize(w, tc.size, tc.announced)
			})
			addr, out := startProxy(t, app.URL, map[string]string{
				responseMaxBytes: sizeLimitSetting,
				metricsToken:     token,
			})

			got := get(t, addr, "/download")
			wantStatus(t, got, tc.want)

			if len(got.body) != tc.received ||
				errors.Is(got.err, io.ErrUnexpectedEOF) != tc.cutOff {
				t.Errorf("client got %d bytes (%v), want %d",
					len(got.body), got.err, tc.received)
			}

			line := out.requestLine(t)
			wantLine(t, line, tc.want, tc.action)

			if line.UpstreamStatus != http.StatusOK {
				t.Errorf("log line has upstream_status %d", line.UpstreamStatus)
			}

			hits := 0
			if tc.action == requestlog.ActionTooLarge {
				hits = 1
			}

			wantLimitHits(t, addr, responseMaxBytes, hits)
		})
	}
}

// answerWithSize answers with a body of size bytes, announced in
// Content-Length or sent in chunks with no length given.
func answerWithSize(w http.ResponseWriter, size int, announced bool) {
	body := make([]byte, size)
	if announced {
		w.Header().Set("Content-Length", strconv.Itoa(size))
		_, _ = w.Write(body)

		return
	}

	// Sending part of it before the end keeps Go's server from working
	// out the length.
	_, _ = w.Write(body[:size/2])
	_ = http.NewResponseController(w).Flush()
	_, _ = w.Write(body[size/2:])
}