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()) } } }