node_probe_body_cap_test.go 1.4 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950
  1. package service
  2. import (
  3. "context"
  4. "net/http"
  5. "net/http/httptest"
  6. "net/url"
  7. "strconv"
  8. "strings"
  9. "testing"
  10. "github.com/mhsanaei/3x-ui/v3/internal/database/model"
  11. )
  12. // An oversized status body must be rejected, not buffered whole by encoding/json.
  13. func TestProbeRejectsOversizedStatusBody(t *testing.T) {
  14. srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
  15. w.Header().Set("Content-Type", "application/json")
  16. _, _ = w.Write([]byte(`{"success":true,"obj":{"cpuPct":1,"panelVersion":"`))
  17. pad := strings.Repeat("x", 1<<20)
  18. for i := 0; i < 3; i++ {
  19. _, _ = w.Write([]byte(pad))
  20. }
  21. _, _ = w.Write([]byte(`"}}`))
  22. }))
  23. defer srv.Close()
  24. u, err := url.Parse(srv.URL)
  25. if err != nil {
  26. t.Fatalf("parse url: %v", err)
  27. }
  28. port, err := strconv.Atoi(u.Port())
  29. if err != nil {
  30. t.Fatalf("parse port: %v", err)
  31. }
  32. n := &model.Node{
  33. Id: 1, Name: "big", Scheme: "http", Address: u.Hostname(), Port: port,
  34. BasePath: "/", Enable: true, AllowPrivateAddress: true, TlsVerifyMode: "skip",
  35. }
  36. svc := &NodeService{}
  37. patch, err := svc.Probe(context.Background(), n)
  38. if err == nil {
  39. t.Fatal("Probe accepted a 3 MiB status body, want an error")
  40. }
  41. // Pin the rejection to the capped decode, not a transport or envelope failure.
  42. if !strings.HasPrefix(patch.LastError, "decode response: ") {
  43. t.Fatalf("LastError = %q, want a \"decode response: \" rejection", patch.LastError)
  44. }
  45. }