|
| 1 | +package httpx |
| 2 | + |
| 3 | +import ( |
| 4 | + "bytes" |
| 5 | + "net/http" |
| 6 | + "net/http/httptest" |
| 7 | + "strings" |
| 8 | + "testing" |
| 9 | + "time" |
| 10 | + |
| 11 | + "github.com/projectdiscovery/retryablehttp-go" |
| 12 | + "github.com/stretchr/testify/require" |
| 13 | + "golang.org/x/text/encoding/simplifiedchinese" |
| 14 | + "golang.org/x/text/transform" |
| 15 | +) |
| 16 | + |
| 17 | +// newLocalHTTPX builds an HTTPX instance suitable for hitting a local test |
| 18 | +// server only (no external network, CDN checks disabled). |
| 19 | +func newLocalHTTPX(t *testing.T) *HTTPX { |
| 20 | + t.Helper() |
| 21 | + options := DefaultOptions |
| 22 | + options.CdnCheck = "false" |
| 23 | + options.Timeout = 5 * time.Second |
| 24 | + options.RetryMax = 0 |
| 25 | + // NB: relies on DefaultOptions.MaxResponseBodySizeToRead being non-zero |
| 26 | + // (see TestDefaultOptionsHasNonZeroReadSize) so the body is actually read. |
| 27 | + |
| 28 | + ht, err := New(&options) |
| 29 | + require.NoError(t, err) |
| 30 | + return ht |
| 31 | +} |
| 32 | + |
| 33 | +// doLocal issues a GET against a local httptest server and returns the parsed |
| 34 | +// httpx Response. |
| 35 | +func doLocal(t *testing.T, ht *HTTPX, url string) *Response { |
| 36 | + t.Helper() |
| 37 | + req, err := retryablehttp.NewRequest(http.MethodGet, url, nil) |
| 38 | + require.NoError(t, err) |
| 39 | + resp, err := ht.Do(req, UnsafeOptions{}) |
| 40 | + require.NoError(t, err) |
| 41 | + return resp |
| 42 | +} |
| 43 | + |
| 44 | +// legacyWordsLines reproduces the exact word/line computation that existed |
| 45 | +// before the refactor, so we can assert the new byte-based path is equivalent. |
| 46 | +func legacyWordsLines(body []byte) (words, lines int) { |
| 47 | + s := string(body) |
| 48 | + if s != "" { |
| 49 | + words = len(strings.Split(s, " ")) |
| 50 | + lines = len(strings.Split(strings.TrimSpace(s), "\n")) |
| 51 | + } |
| 52 | + return |
| 53 | +} |
| 54 | + |
| 55 | +// TestDefaultOptionsHasNonZeroReadSize guards against the package var-init |
| 56 | +// ordering regression where DefaultOptions was initialized before |
| 57 | +// DefaultMaxResponseBodySize, leaving MaxResponseBodySizeToRead at 0 (which made |
| 58 | +// LimitReader read zero bytes and produced empty bodies for library users). |
| 59 | +func TestDefaultOptionsHasNonZeroReadSize(t *testing.T) { |
| 60 | + require.NotZero(t, DefaultMaxResponseBodySize) |
| 61 | + require.Equal(t, DefaultMaxResponseBodySize, DefaultOptions.MaxResponseBodySizeToRead) |
| 62 | +} |
| 63 | + |
| 64 | +func TestDoBodyNoDecodePreservesRawAndData(t *testing.T) { |
| 65 | + body := []byte("hello world\nsecond line\n") |
| 66 | + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { |
| 67 | + w.Header().Set("Content-Type", "text/html") |
| 68 | + _, _ = w.Write(body) |
| 69 | + })) |
| 70 | + defer ts.Close() |
| 71 | + |
| 72 | + resp := doLocal(t, newLocalHTTPX(t), ts.URL) |
| 73 | + |
| 74 | + require.Equal(t, body, resp.Data, "decoded data must equal body") |
| 75 | + require.Equal(t, body, resp.RawData, "raw data must equal undecoded body") |
| 76 | + |
| 77 | + wantWords, wantLines := legacyWordsLines(body) |
| 78 | + require.Equal(t, wantWords, resp.Words) |
| 79 | + require.Equal(t, wantLines, resp.Lines) |
| 80 | +} |
| 81 | + |
| 82 | +func TestDoBodyEmpty(t *testing.T) { |
| 83 | + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { |
| 84 | + w.Header().Set("Content-Type", "text/html") |
| 85 | + })) |
| 86 | + defer ts.Close() |
| 87 | + |
| 88 | + resp := doLocal(t, newLocalHTTPX(t), ts.URL) |
| 89 | + require.Empty(t, resp.Data) |
| 90 | + require.Empty(t, resp.RawData) |
| 91 | + require.Equal(t, 0, resp.Words) |
| 92 | + require.Equal(t, 0, resp.Lines) |
| 93 | +} |
| 94 | + |
| 95 | +// TestDoBodyGBKDecodeKeepsRawUndecoded ensures that when DecodeData actually |
| 96 | +// transcodes the body, RawData still holds the original (undecoded) bytes while |
| 97 | +// Data holds the decoded UTF-8 bytes. |
| 98 | +func TestDoBodyGBKDecodeKeepsRawUndecoded(t *testing.T) { |
| 99 | + utf8Body := "<html><head></head><body>你好世界 测试</body></html>" |
| 100 | + gbkBody, _, err := transform.Bytes(simplifiedchinese.GBK.NewEncoder(), []byte(utf8Body)) |
| 101 | + require.NoError(t, err) |
| 102 | + require.NotEqual(t, []byte(utf8Body), gbkBody, "precondition: gbk bytes differ from utf8") |
| 103 | + |
| 104 | + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { |
| 105 | + w.Header().Set("Content-Type", "text/html; charset=gbk") |
| 106 | + _, _ = w.Write(gbkBody) |
| 107 | + })) |
| 108 | + defer ts.Close() |
| 109 | + |
| 110 | + resp := doLocal(t, newLocalHTTPX(t), ts.URL) |
| 111 | + |
| 112 | + require.Equal(t, gbkBody, resp.RawData, "RawData must hold the original undecoded bytes") |
| 113 | + require.Equal(t, []byte(utf8Body), resp.Data, "Data must hold the decoded UTF-8 bytes") |
| 114 | +} |
| 115 | + |
| 116 | +// TestDoBodyNoDecodeSharesBacking documents the memory optimization: on the |
| 117 | +// no-decode hot path RawData and Data share the same backing array (no extra |
| 118 | +// full-body copy is made). |
| 119 | +func TestDoBodyNoDecodeSharesBacking(t *testing.T) { |
| 120 | + body := []byte("shared backing array body") |
| 121 | + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { |
| 122 | + w.Header().Set("Content-Type", "text/plain") |
| 123 | + _, _ = w.Write(body) |
| 124 | + })) |
| 125 | + defer ts.Close() |
| 126 | + |
| 127 | + resp := doLocal(t, newLocalHTTPX(t), ts.URL) |
| 128 | + |
| 129 | + require.NotEmpty(t, resp.Data) |
| 130 | + require.NotEmpty(t, resp.RawData) |
| 131 | + require.Equal(t, resp.Data, resp.RawData) |
| 132 | + require.Same(t, &resp.Data[0], &resp.RawData[0], |
| 133 | + "RawData and Data should share the backing array on the no-decode path") |
| 134 | +} |
| 135 | + |
| 136 | +// TestWordsLinesEquivalence is the core guard for the refactor: the byte-based |
| 137 | +// counting used on the hot path must be identical to the previous |
| 138 | +// strings.Split-based counting for a wide range of inputs and edge cases. |
| 139 | +func TestWordsLinesEquivalence(t *testing.T) { |
| 140 | + cases := []string{ |
| 141 | + "", |
| 142 | + "a", |
| 143 | + "a b c", |
| 144 | + " ", // only spaces |
| 145 | + "a b", // consecutive spaces |
| 146 | + "line1\nline2\nline3", // multiple lines |
| 147 | + "\n\n\n", // only newlines |
| 148 | + " leading and trailing ", // surrounding whitespace |
| 149 | + "\n mixed \t whitespace \n", // tabs/newlines around |
| 150 | + "trailing newline\n", |
| 151 | + "word", |
| 152 | + "tab\tseparated values", |
| 153 | + "unicode \u00a0 nbsp space", |
| 154 | + "emoji 😀 and spaces ", |
| 155 | + } |
| 156 | + |
| 157 | + for _, c := range cases { |
| 158 | + body := []byte(c) |
| 159 | + wantWords, wantLines := legacyWordsLines(body) |
| 160 | + |
| 161 | + var gotWords, gotLines int |
| 162 | + if len(body) > 0 { |
| 163 | + gotWords = bytes.Count(body, []byte{' '}) + 1 |
| 164 | + gotLines = bytes.Count(bytes.TrimSpace(body), []byte{'\n'}) + 1 |
| 165 | + } |
| 166 | + |
| 167 | + require.Equalf(t, wantWords, gotWords, "words mismatch for %q", c) |
| 168 | + require.Equalf(t, wantLines, gotLines, "lines mismatch for %q", c) |
| 169 | + } |
| 170 | +} |
| 171 | + |
| 172 | +// TestBodyMetricsCountingDoesNotAllocate locks in the optimization: the |
| 173 | +// byte-based word/line counting used on the hot path must not allocate (the |
| 174 | +// previous string(respbody) + strings.Split approach allocated O(len(body))). |
| 175 | +// If someone reintroduces a full-body string copy or Split-based counting, this |
| 176 | +// test fails. |
| 177 | +func TestBodyMetricsCountingDoesNotAllocate(t *testing.T) { |
| 178 | + body := bytes.Repeat([]byte("lorem ipsum dolor sit amet\n"), 40000) // ~1MB |
| 179 | + var words, lines int |
| 180 | + |
| 181 | + allocs := testing.AllocsPerRun(50, func() { |
| 182 | + // identical expressions to the hot path in Do() |
| 183 | + words = bytes.Count(body, []byte{' '}) + 1 |
| 184 | + lines = bytes.Count(bytes.TrimSpace(body), []byte{'\n'}) + 1 |
| 185 | + }) |
| 186 | + |
| 187 | + require.NotZero(t, words) |
| 188 | + require.NotZero(t, lines) |
| 189 | + require.Zerof(t, allocs, "word/line counting must not allocate, got %v allocs/op", allocs) |
| 190 | +} |
0 commit comments