go-omni/app/bodyPolicy_test.go

package app

import (
	"bytes"
	"context"
	"errors"
	"io"
	"net"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync/atomic"
	"testing"
	"time"

	"git.yuntuo.hk/k2safe/go-omni/config"
	"git.yuntuo.hk/k2safe/go-omni/httpx"
	"git.yuntuo.hk/k2safe/go-omni/internal/testutil"
	"git.yuntuo.hk/k2safe/go-omni/response"
)

func serveBodyPolicy(t *testing.T, a *App, transport string) (string, *http.Client) {
	t.Helper()
	cert, key, clientTLS := testutil.CertificateFiles(t)
	a.config.Transport = transport
	a.config.TLS = config.TLSConfig{CertFile: cert, KeyFile: key}
	l, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		t.Fatal(err)
	}
	ctx, cancel := context.WithCancel(context.Background())
	done := make(chan error, 1)
	go func() { done <- a.ServeContext(ctx, l) }()
	client := &http.Client{Timeout: 5 * time.Second, Transport: &http.Transport{TLSClientConfig: clientTLS, ForceAttemptHTTP2: transport == config.TransportHTTP2}}
	t.Cleanup(func() {
		client.CloseIdleConnections()
		cancel()
		select {
		case err := <-done:
			if err != nil {
				t.Error(err)
			}
		case <-time.After(5 * time.Second):
			t.Error("server did not stop")
		}
	})
	return "https://" + l.Addr().String(), client
}

func consumeBody(c *httpx.Context) error {
	_, err := io.Copy(io.Discard, c.Request().Body)
	if err != nil {
		return err
	}
	return c.Text(200, "accepted")
}

func TestUploadPolicyAcrossTransports(t *testing.T) {
	for _, transport := range []string{config.TransportNetHTTP, config.TransportNative, config.TransportHTTP2} {
		t.Run(transport, func(t *testing.T) {
			a := New()
			a.POST("/default", consumeBody)
			a.HandleWithBodyLimit("POST", "/large", 16<<20, nil, consumeBody)
			a.Group("/v1").Group("files").HandleWithBodyLimit("POST", "/large", 16<<20, nil, consumeBody)
			base, client := serveBodyPolicy(t, a, transport)
			payload := bytes.Repeat([]byte{'x'}, 9<<20)
			for _, chunked := range []bool{false, true} {
				for _, path := range []string{"/default", "/large", "/v1/files/large"} {
					var body io.Reader = bytes.NewReader(payload)
					if chunked {
						body = struct{ io.Reader }{body}
					}
					req, _ := http.NewRequest("POST", base+path, body)
					res, err := client.Do(req)
					if err != nil {
						t.Fatal(err)
					}
					data, err := io.ReadAll(res.Body)
					res.Body.Close()
					if err != nil {
						t.Fatal(err)
					}
					want := 200
					if path == "/default" {
						want = 413
					}
					if res.StatusCode != want || (want == 413 && !strings.Contains(string(data), "BODY_TOO_LARGE")) {
						t.Fatalf("%s unknownLength=%v: %d %s", path, chunked, res.StatusCode, data)
					}
					if transport == config.TransportHTTP2 && res.ProtoMajor != 2 {
						t.Fatal("HTTP/2 not negotiated")
					}
				}
			}
		})
	}
}

func TestConfiguredBodyPolicyBeforeMiddleware(t *testing.T) {
	for _, transport := range []string{config.TransportNetHTTP, config.TransportNative, config.TransportHTTP2} {
		t.Run(transport, func(t *testing.T) {
			a := New(config.WithHTTPConfig(config.HTTPConfig{DefaultMaxBodyBytes: 32, MaxBodyBytes: 128}))
			var calls atomic.Int32
			a.Use(func(next httpx.Handler) httpx.Handler {
				return func(c *httpx.Context) error {
					data, err := io.ReadAll(c.Request().Body)
					if err != nil {
						return err
					}
					c.Request().Body = io.NopCloser(bytes.NewReader(data))
					return next(c)
				}
			})
			h := func(c *httpx.Context) error { calls.Add(1); return consumeBody(c) }
			a.POST("/default", h)
			a.HandleWithBodyLimit("POST", "/override", 64, nil, h)
			a.HandleWithBodyLimit("POST", "/hard", 128, nil, h)
			base, client := serveBodyPolicy(t, a, transport)
			for _, tc := range []struct {
				path         string
				size, status int
			}{
				{"/default", 32, 200}, {"/default", 33, 413}, {"/override", 64, 200}, {"/override", 65, 413}, {"/hard", 128, 200}, {"/hard", 129, 413},
			} {
				for _, unknown := range []bool{false, true} {
					before := calls.Load()
					var body io.Reader = strings.NewReader(strings.Repeat("a", tc.size))
					if unknown {
						body = struct{ io.Reader }{body}
					}
					res, err := client.Post(base+tc.path, "application/octet-stream", body)
					if err != nil {
						t.Fatal(err)
					}
					data, err := io.ReadAll(res.Body)
					res.Body.Close()
					if err != nil || res.StatusCode != tc.status {
						t.Fatalf("%+v unknown=%v: %d %s %v", tc, unknown, res.StatusCode, data, err)
					}
					if tc.status == 413 && (calls.Load() != before || !strings.Contains(string(data), "BODY_TOO_LARGE")) {
						t.Fatal("oversize reached business or lost error", string(data))
					}
				}
			}
		})
	}
}

func TestBodyHardLimitConflictBeforeServing(t *testing.T) {
	for _, entry := range []string{"RunContext", "ServeContext", "RunWith", "NewServer", "Engine"} {
		t.Run(entry, func(t *testing.T) {
			a := New(config.WithHTTPConfig(config.HTTPConfig{MaxBodyBytes: 8 << 20}))
			a.Group("/v1").Group("files").HandleWithBodyLimit("POST", "/upload/", 16<<20, nil, consumeBody)
			closed := 0
			a.OnShutdown(func() error { closed++; return nil })
			if err := a.ValidateHTTP(); err == nil {
				t.Fatal("external preflight accepted conflicting route limit")
			}
			var err error
			switch entry {
			case "RunContext":
				err = a.RunContext(context.Background(), "bad address")
			case "ServeContext":
				l, e := net.Listen("tcp", "127.0.0.1:0")
				if e != nil {
					t.Fatal(e)
				}
				err = a.ServeContext(context.Background(), l)
				if _, e = l.Accept(); !errors.Is(e, net.ErrClosed) {
					t.Fatal("listener leaked", e)
				}
			case "RunWith":
				err = a.RunWith(context.Background(), runnerFunc(func(context.Context, http.Handler, string) error { t.Error("runner started"); return nil }))
			case "NewServer":
				func() {
					defer func() {
						if p := recover(); p != nil {
							err, _ = p.(error)
						}
					}()
					a.NewServer("")
				}()
				_ = a.Close()
			case "Engine":
				w := httptest.NewRecorder()
				a.Engine().ServeHTTP(w, httptest.NewRequest("POST", "/v1/files/upload/", nil))
				if w.Code != 503 {
					t.Fatal(w.Code)
				}
				err = a.validateErr
				_ = a.Close()
			}
			for _, part := range []string{"POST /v1/files/upload/", "16777216", "8388608", "http.maxBodyBytes"} {
				if err == nil || !strings.Contains(err.Error(), part) {
					t.Fatalf("missing %q: %v", part, err)
				}
			}
			if closed != 1 {
				t.Fatal("cleanup count", closed)
			}
		})
	}
}

func TestDefaultBodyLimitRegistrationPaths(t *testing.T) {
	a := New(config.WithHTTPConfig(config.HTTPConfig{DefaultMaxBodyBytes: 4, MaxBodyBytes: 16}))
	h := func(*httpx.Context) error { t.Error("oversize reached business"); return nil }
	encoder, _ := response.New(response.Config{})
	a.Handle("POST", "/handle", h)
	a.HandleWithResponse("POST", "/response", encoder, h)
	a.Any("/any", h)
	g := a.Group("/group")
	g.POST("/post", h)
	g.Any("/any", h)
	g.HandleWithResponse("POST", "/response", encoder, h)
	for _, path := range []string{"/handle", "/response", "/any", "/group/post", "/group/any", "/group/response"} {
		w := httptest.NewRecorder()
		r := httptest.NewRequest("POST", path, strings.NewReader("12345"))
		a.ServeHTTP(w, r)
		if w.Code != 413 {
			t.Fatal(path, w.Code, w.Body.String())
		}
	}
}