From 86cbdac36f727bcf074a700a78c8de0e86a2d794 Mon Sep 17 00:00:00 2001
From: kim <89579420+NyaaaWhatsUpDoc@users.noreply.github.com>
Date: Fri, 2 Jun 2023 09:34:52 +0100
Subject: [PATCH] [performance] wrap httpclient response body to ensure drained
 before close (#1854)

Signed-off-by: kim <grufwub@gmail.com>
---
 go.mod                                       |  2 +-
 go.sum                                       |  4 +--
 internal/httpclient/client.go                | 27 +++++++++++++++-----
 vendor/codeberg.org/gruf/go-iotools/close.go |  8 ++++++
 vendor/modules.txt                           |  2 +-
 5 files changed, 33 insertions(+), 10 deletions(-)

diff --git a/go.mod b/go.mod
index 949421d95..7fc28c576 100644
--- a/go.mod
+++ b/go.mod
@@ -9,7 +9,7 @@ require (
 	codeberg.org/gruf/go-debug v1.3.0
 	codeberg.org/gruf/go-errors/v2 v2.2.0
 	codeberg.org/gruf/go-fastcopy v1.1.2
-	codeberg.org/gruf/go-iotools v0.0.0-20221224124424-3386841cb225
+	codeberg.org/gruf/go-iotools v0.0.0-20230601182242-d933b07dcbef
 	codeberg.org/gruf/go-kv v1.6.1
 	codeberg.org/gruf/go-logger/v2 v2.2.1
 	codeberg.org/gruf/go-mutexes v1.1.5
diff --git a/go.sum b/go.sum
index 66909040f..f515cd51a 100644
--- a/go.sum
+++ b/go.sum
@@ -64,8 +64,8 @@ codeberg.org/gruf/go-fastpath/v2 v2.0.0 h1:iAS9GZahFhyWEH0KLhFEJR+txx1ZhMXxYzu2q
 codeberg.org/gruf/go-fastpath/v2 v2.0.0/go.mod h1:3pPqu5nZjpbRrOqvLyAK7puS1OfEtQvjd6342Cwz56Q=
 codeberg.org/gruf/go-hashenc v1.0.2 h1:U3jH6zMXZiL96czD/qaJd8OR2h7LlBzGv/2WxnMHI/g=
 codeberg.org/gruf/go-hashenc v1.0.2/go.mod h1:eK+A8clLcEN/m1nftNsRId0kfYDQnETnuIfBGZ8Gvsg=
-codeberg.org/gruf/go-iotools v0.0.0-20221224124424-3386841cb225 h1:tP9YvEBfADGG3mXkfrALLadlcbrZsFsWKZvFtUZtrt8=
-codeberg.org/gruf/go-iotools v0.0.0-20221224124424-3386841cb225/go.mod h1:B8uq4yHtIcKXhBZT9C/SYisz25lldLHMVpwZPz4ADLQ=
+codeberg.org/gruf/go-iotools v0.0.0-20230601182242-d933b07dcbef h1:3Ydviw47TFEk27FRCOXkRxU3MfgyNzoicLzq8J3NbtI=
+codeberg.org/gruf/go-iotools v0.0.0-20230601182242-d933b07dcbef/go.mod h1:B8uq4yHtIcKXhBZT9C/SYisz25lldLHMVpwZPz4ADLQ=
 codeberg.org/gruf/go-kv v1.6.1 h1:HsCZEy0zfGq1oFGOEOO2qnpooiGefCZBbcUpa3KdXn8=
 codeberg.org/gruf/go-kv v1.6.1/go.mod h1:O/YkSvKiS9XsRolM3rqCd9YJmND7dAXu9z+PrlYO4bc=
 codeberg.org/gruf/go-logger/v2 v2.2.1 h1:RP2u059EQKTBFV3cN8X6xDxNk2RkzqdgXGKflKqB7Oc=
diff --git a/internal/httpclient/client.go b/internal/httpclient/client.go
index efbf4cd18..2a2485561 100644
--- a/internal/httpclient/client.go
+++ b/internal/httpclient/client.go
@@ -34,6 +34,7 @@ import (
 	"codeberg.org/gruf/go-byteutil"
 	"codeberg.org/gruf/go-cache/v3"
 	errorsv2 "codeberg.org/gruf/go-errors/v2"
+	"codeberg.org/gruf/go-iotools"
 	"codeberg.org/gruf/go-kv"
 	"github.com/superseriousbusiness/gotosocial/internal/gtscontext"
 	"github.com/superseriousbusiness/gotosocial/internal/gtserror"
@@ -265,7 +266,8 @@ func (c *Client) DoSigned(r *http.Request, sign SignFunc) (rsp *http.Response, e
 				}
 			}
 
-			// Ensure unset.
+			// Close + unset rsp.
+			_ = rsp.Body.Close()
 			rsp = nil
 
 		} else if errorsv2.Comparable(err,
@@ -326,11 +328,6 @@ func (c *Client) do(req *http.Request) (*http.Response, error) {
 		return nil, err
 	}
 
-	// Check response body not too large.
-	if rsp.ContentLength > c.bodyMax {
-		return nil, ErrBodyTooLarge
-	}
-
 	// Seperate the body implementers.
 	rbody := (io.Reader)(rsp.Body)
 	cbody := (io.Closer)(rsp.Body)
@@ -345,11 +342,29 @@ func (c *Client) do(req *http.Request) (*http.Response, error) {
 	// Don't trust them, limit body reads.
 	rbody = io.LimitReader(rbody, limit)
 
+	// Wrap closer to ensure entire body drained BEFORE close.
+	cbody = iotools.CloserAfterCallback(cbody, func() {
+		_, _ = discard.ReadFrom(rbody)
+	})
+
 	// Wrap body with limit.
 	rsp.Body = &struct {
 		io.Reader
 		io.Closer
 	}{rbody, cbody}
 
+	// Check response body not too large.
+	if rsp.ContentLength > c.bodyMax {
+		_ = rsp.Body.Close()
+		return nil, ErrBodyTooLarge
+	}
+
 	return rsp, nil
 }
+
+// cast discard writer to full interface it supports.
+var discard = io.Discard.(interface { //nolint
+	io.Writer
+	io.StringWriter
+	io.ReaderFrom
+})
diff --git a/vendor/codeberg.org/gruf/go-iotools/close.go b/vendor/codeberg.org/gruf/go-iotools/close.go
index fbed7f33c..3f0ee7780 100644
--- a/vendor/codeberg.org/gruf/go-iotools/close.go
+++ b/vendor/codeberg.org/gruf/go-iotools/close.go
@@ -17,6 +17,14 @@ func CloserCallback(c io.Closer, cb func()) io.Closer {
 	})
 }
 
+func CloserAfterCallback(c io.Closer, cb func()) io.Closer {
+	return CloserFunc(func() (err error) {
+		defer func() { err = c.Close() }()
+		cb()
+		return
+	})
+}
+
 // CloseOnce wraps an io.Closer to ensure it only performs the close logic once.
 func CloseOnce(c io.Closer) io.Closer {
 	return CloserFunc(func() error {
diff --git a/vendor/modules.txt b/vendor/modules.txt
index 38af992c4..9d4b237e4 100644
--- a/vendor/modules.txt
+++ b/vendor/modules.txt
@@ -36,7 +36,7 @@ codeberg.org/gruf/go-fastpath/v2
 # codeberg.org/gruf/go-hashenc v1.0.2
 ## explicit; go 1.16
 codeberg.org/gruf/go-hashenc
-# codeberg.org/gruf/go-iotools v0.0.0-20221224124424-3386841cb225
+# codeberg.org/gruf/go-iotools v0.0.0-20230601182242-d933b07dcbef
 ## explicit; go 1.19
 codeberg.org/gruf/go-iotools
 # codeberg.org/gruf/go-kv v1.6.1