fix: critical memory leak (#183)

Signed-off-by: Neko Ayaka <[email protected]>
nekomeowww · Dec 10, 2023 · 46d728e · 46d728e
1 parent e48122a
commit 46d728e
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 6 deletions.
diff --git a/internal/models/smr/smr.go b/internal/models/smr/smr.go
@@ -173,20 +173,34 @@ func (m *Model) extractContentFromURL(ctx context.Context, urlString string) (*r
 	resp, err := m.req.
 		R().
 		EnableDumpTo(dumpBuffer).
+		DisableAutoReadResponse().
 		SetContext(ctx).
 		Get(parsedURL.String())
 	if err != nil {
 		return nil, fmt.Errorf("failed to get url %s, %w: %v", parsedURL.String(), ErrNetworkError, err)
 	}
 	if !resp.IsSuccessState() {
+		errorBuf := new(bytes.Buffer)
+		defer errorBuf.Reset()
+
+		_, err = io.Copy(errorBuf, resp.Body)
+		if err != nil {
+			fmt.Fprintf(errorBuf, "failed to read response body: %v", err)
+		}
+
+		dumpBuffer.WriteString("\n")
+		dumpBuffer.Write(errorBuf.Bytes())
+
 		return nil, fmt.Errorf("failed to get url %s, %w, status code: %d, dump: %s", parsedURL.String(), ErrRequestFailed, resp.StatusCode, dumpBuffer.String())
 	}
 	if !strings.Contains(resp.Header.Get("Content-Type"), "text/html") {
 		return nil, fmt.Errorf("url fetched, but content-type not supported yet, %w, content-type: %s", ErrContentNotSupported, resp.Header.Get("Content-Type"))
 	}
 
 	defer resp.Body.Close()
+
 	buffer := new(bytes.Buffer)
+	defer buffer.Reset()
 
 	_, err = io.Copy(buffer, resp.Body)
 	if err != nil {

diff --git a/pkg/linkprev/linkprev.go b/pkg/linkprev/linkprev.go
@@ -26,10 +26,7 @@ type Client struct {
 
 func NewClient() *Client {
 	return &Client{
-		reqClient: req.
-			C().
-			SetUserAgent("Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/111.0.0.0 Safari/537.36 Edg/111.0.1661.54").
-			EnableDumpEachRequest(),
+		reqClient: req.C().SetUserAgent("Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/111.0.0.0 Safari/537.36 Edg/111.0.1661.54"),
 	}
 }
 
@@ -41,6 +38,8 @@ func (c *Client) Preview(ctx context.Context, urlStr string) (Meta, error) {
 		return Meta{}, err
 	}
 
+	defer body.Reset()
+
 	doc, err := goquery.NewDocumentFromReader(body)
 	if err != nil {
 		return Meta{}, fmt.Errorf("failed to parse response body with goquery: %v", err)
@@ -85,18 +84,30 @@ func (c *Client) alterRequestForTwitter(request *req.Request, urlStr string) *re
 	return request.SetHeader("User-Agent", "Mozilla/5.0 (compatible; Googlebot/2.1; +http://www.google.com/bot.html)")
 }
 
-func (c *Client) request(r *req.Request, urlStr string) (io.Reader, error) {
+func (c *Client) request(r *req.Request, urlStr string) (*bytes.Buffer, error) {
 	dumpBuffer := new(bytes.Buffer)
 	defer dumpBuffer.Reset()
 
 	resp, err := r.
 		EnableDumpTo(dumpBuffer).
+		DisableAutoReadResponse().
 		Get(urlStr)
 	if err != nil {
 		return nil, fmt.Errorf("failed to get a preview of url %s, %w: %v", urlStr, ErrNetworkError, err)
 	}
 	if !resp.IsSuccessState() {
-		return nil, fmt.Errorf("failed to get url %s, %w, status code: %d, dump: %s", urlStr, ErrRequestFailed, resp.StatusCode, dumpBuffer.String())
+		errorBuf := new(bytes.Buffer)
+		defer errorBuf.Reset()
+
+		_, err = io.Copy(errorBuf, resp.Body)
+		if err != nil {
+			fmt.Fprintf(errorBuf, "failed to read response body: %v", err)
+		}
+
+		dumpBuffer.WriteString("\n")
+		dumpBuffer.Write(errorBuf.Bytes())
+
+		return nil, fmt.Errorf("failed to get url %s, %w, status code: %d, dump:\n%s", urlStr, ErrRequestFailed, resp.StatusCode, dumpBuffer.String())
 	}
 
 	defer resp.Body.Close()