From 0f494fd756b1af6176e50d436258dad0455d2a95 Mon Sep 17 00:00:00 2001 From: bensynapse <118375461+bensynapse@users.noreply.github.com> Date: Sun, 11 Oct 2026 08:13:34 +0300 Subject: [PATCH] close response bodies after loading compressed HTML --- query.go | 15 +++----- query_http_test.go | 94 ++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 100 insertions(+), 9 deletions(-) create mode 100644 query_http_test.go diff --git a/query.go b/query.go index ab22207..cb8b337 100644 --- a/query.go +++ b/query.go @@ -97,12 +97,12 @@ func LoadURL(url string) (*html.Node, error) { // LoadURLWithCustomClient loads the HTML document from the specified URL. Default enabling gzip on a HTTP request. // Custom HTTP Client. -func LoadURLWithClient(url string, client *http.Client)(*html.Node, error){ +func LoadURLWithClient(url string, client *http.Client) (*html.Node, error) { req, err := http.NewRequest("GET", url, nil) if err != nil { return nil, err } - if client == nil{ + if client == nil { client = http.DefaultClient } // Enable gzip compression. @@ -111,13 +111,8 @@ func LoadURLWithClient(url string, client *http.Client)(*html.Node, error){ if err != nil { return nil, err } - var reader io.ReadCloser - - defer func() { - if reader != nil { - reader.Close() - } - }() + defer resp.Body.Close() + var reader io.Reader encoding := resp.Header.Get("Content-Encoding") switch encoding { case "gzip": @@ -125,12 +120,14 @@ func LoadURLWithClient(url string, client *http.Client)(*html.Node, error){ if err != nil { return nil, err } + defer gzipReader.Close() reader = gzipReader case "deflate": zlibReader, err := zlib.NewReader(resp.Body) if err != nil { return nil, err } + defer zlibReader.Close() reader = zlibReader case "": reader = resp.Body diff --git a/query_http_test.go b/query_http_test.go new file mode 100644 index 0000000..0e588c0 --- /dev/null +++ b/query_http_test.go @@ -0,0 +1,94 @@ +package htmlquery + +import ( + "bytes" + "compress/gzip" + "compress/zlib" + "io" + "net/http" + "testing" +) + +type trackedResponseBody struct { + io.Reader + closes int +} + +func (b *trackedResponseBody) Close() error { + b.closes++ + return nil +} + +type responseTransport func(*http.Request) (*http.Response, error) + +func (f responseTransport) RoundTrip(req *http.Request) (*http.Response, error) { + return f(req) +} + +type failedResponseReader struct{} + +func (failedResponseReader) Read([]byte) (int, error) { + return 0, io.ErrClosedPipe +} + +func TestLoadURLWithClientClosesResponseBody(t *testing.T) { + const document = "examplepage" + var gz, def bytes.Buffer + gw := gzip.NewWriter(&gz) + if _, err := io.WriteString(gw, document); err != nil { + t.Fatal(err) + } + if err := gw.Close(); err != nil { + t.Fatal(err) + } + dw := zlib.NewWriter(&def) + if _, err := io.WriteString(dw, document); err != nil { + t.Fatal(err) + } + if err := dw.Close(); err != nil { + t.Fatal(err) + } + + for _, tc := range []struct { + name, encoding string + reader io.Reader + wantErr bool + }{ + {"plain", "", bytes.NewBufferString(document), false}, + {"gzip", "gzip", bytes.NewReader(gz.Bytes()), false}, + {"deflate", "deflate", bytes.NewReader(def.Bytes()), false}, + {"invalid gzip", "gzip", bytes.NewBufferString("invalid"), true}, + {"invalid deflate", "deflate", bytes.NewBufferString("invalid"), true}, + {"unsupported encoding", "br", bytes.NewBufferString(document), true}, + {"read failure", "", failedResponseReader{}, true}, + } { + t.Run(tc.name, func(t *testing.T) { + body := &trackedResponseBody{Reader: tc.reader} + client := &http.Client{Transport: responseTransport(func(req *http.Request) (*http.Response, error) { + return &http.Response{ + StatusCode: http.StatusOK, + Header: http.Header{ + "Content-Encoding": []string{tc.encoding}, + "Content-Type": []string{"text/html; charset=utf-8"}, + }, + Body: body, + Request: req, + }, nil + })} + + node, err := LoadURLWithClient("https://example.invalid/page", client) + if (err != nil) != tc.wantErr { + t.Fatalf("error = %v, want error %t", err, tc.wantErr) + } + if err == nil { + title := FindOne(node, "//title") + if title == nil || InnerText(title) != "example" { + t.Fatal("parsed document changed") + } + } + if body.closes != 1 { + t.Errorf("response body Close calls = %d, want 1", body.closes) + } + }) + } +}