Repository: zeromicro/go-zero
Pull request #3498 (2023-08-16 06:22:56)

update: limit logBrief http body size
It seems unnecessary to dump full Http Body to `logBrief`. 
This PR move limitBodyBytes to top layer using. This ensures that the size of the http body passed to logBrief has been limited.

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/zeromicro/go-zero/pull/3498?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) Report
> Merging [#3498](https://app.codecov.io/gh/zeromicro/go-zero/pull/3498?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) (33c0805) into [master](https://app.codecov.io/gh/zeromicro/go-zero/commit/cb3ffc76a3e3014eadbca4a389f9801177766d8b?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) (cb3ffc7) will **decrease** coverage by `0.01%`.
> The diff coverage is `100.00%`.


| [Files Changed](https://app.codecov.io/gh/zeromicro/go-zero/pull/3498?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None) | Coverage Δ | |
|---|---|---|
| [core/iox/read.go](https://app.codecov.io/gh/zeromicro/go-zero/pull/3498?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None#diff-Y29yZS9pb3gvcmVhZC5nbw==) | `90.00% <100.00%> (+0.71%)` | :arrow_up: |
| [core/iox/tee.go](https://app.codecov.io/gh/zeromicro/go-zero/pull/3498?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None#diff-Y29yZS9pb3gvdGVlLmdv) | `100.00% <100.00%> (ø)` | |
| [rest/handler/loghandler.go](https://app.codecov.io/gh/zeromicro/go-zero/pull/3498?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None#diff-cmVzdC9oYW5kbGVyL2xvZ2hhbmRsZXIuZ28=) | `100.00% <100.00%> (+2.09%)` | :arrow_up: |

... and [3 files with indirect coverage changes](https://app.codecov.io/gh/zeromicro/go-zero/pull/3498/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=None)


Patch:
diff --git a/core/iox/read.go b/core/iox/read.go
index 530597fc..1f27f445 100644
--- a/core/iox/read.go
+++ b/core/iox/read.go
@@ -28,6 +28,16 @@ func DupReadCloser(reader io.ReadCloser) (io.ReadCloser, io.ReadCloser) {
 	return io.NopCloser(tee), io.NopCloser(&buf)
 }
 
+// LimitDupReadCloser returns two io.ReadCloser that read from the first will be written to the second.
+// But the second io.ReadCloser is limited to up to n bytes.
+// The first returned reader needs to be read first, because the content
+// read from it will be written to the underlying buffer of the second reader.
+func LimitDupReadCloser(reader io.ReadCloser, n int64) (io.ReadCloser, io.ReadCloser) {
+	var buf bytes.Buffer
+	tee := LimitTeeReader(reader, &buf, n)
+	return io.NopCloser(tee), io.NopCloser(&buf)
+}
+
 // KeepSpace customizes the reading functions to keep leading and tailing spaces.
 func KeepSpace() TextReadOption {
 	return func(o *textReadOptions) {
diff --git a/core/iox/tee.go b/core/iox/tee.go
new file mode 100644
index 00000000..07b00158
--- /dev/null
+++ b/core/iox/tee.go
@@ -0,0 +1,33 @@
+package iox
+
+import "io"
+
+// LimitTeeReader returns a Reader that writes up to n bytes to w what it reads from r.
+// First n bytes reads from r performed through it are matched with
+// corresponding writes to w. There is no internal buffering -
+// the write must complete before the first n bytes read completes.
+// Any error encountered while writing is reported as a read error.
+func LimitTeeReader(r io.Reader, w io.Writer, n int64) io.Reader {
+	return &limitTeeReader{r, w, n}
+}
+
+type limitTeeReader struct {
+	r io.Reader
+	w io.Writer
+	n int64 // limit bytes remaining
+}
+
+func (t *limitTeeReader) Read(p []byte) (n int, err error) {
+	n, err = t.r.Read(p)
+	if n > 0 && t.n > 0 {
+		limit := int64(n)
+		if limit > t.n {
+			limit = t.n
+		}
+		if n, err := t.w.Write(p[:limit]); err != nil {
+			return n, err
+		}
+		t.n -= limit
+	}
+	return
+}
diff --git a/rest/handler/loghandler.go b/rest/handler/loghandler.go
index 07cbdf19..eaf59a16 100644
--- a/rest/handler/loghandler.go
+++ b/rest/handler/loghandler.go
@@ -10,7 +10,6 @@ import (
 	"net/http"
 	"net/http/httputil"
 	"strconv"
-	"strings"
 	"time"
 
 	"github.com/zeromicro/go-zero/core/color"
@@ -39,7 +38,7 @@ func LogHandler(next http.Handler) http.Handler {
 		lrw := response.NewWithCodeResponseWriter(w)
 
 		var dup io.ReadCloser
-		r.Body, dup = iox.DupReadCloser(r.Body)
+		r.Body, dup = iox.LimitDupReadCloser(r.Body, limitBodyBytes)
 		next.ServeHTTP(lrw, r.WithContext(internal.WithLogCollector(r.Context(), logs)))
 		r.Body = dup
 		logBrief(r, lrw.Code, timer, logs)
@@ -136,14 +135,7 @@ func logBrief(r *http.Request, code int, timer *utils.ElapsedTimer, logs *intern
 
 	ok := isOkResponse(code)
 	if !ok {
-		fullReq := dumpRequest(r)
-		limitReader := io.LimitReader(strings.NewReader(fullReq), limitBodyBytes)
-		body, err := io.ReadAll(limitReader)
-		if err != nil {
-			buf.WriteString(fmt.Sprintf("\n%s", fullReq))
-		} else {
-			buf.WriteString(fmt.Sprintf("\n%s", string(body)))
-		}
+		buf.WriteString(fmt.Sprintf("\n%s", dumpRequest(r)))
 	}
 
 	body := logs.Flush()


Test patch:
diff --git a/core/iox/read_test.go b/core/iox/read_test.go
index 8ab2dd59..b31ce283 100644
--- a/core/iox/read_test.go
+++ b/core/iox/read_test.go
@@ -108,6 +108,29 @@ func TestDupReadCloser(t *testing.T) {
 	verify(r2)
 }
 
+func TestLimitDupReadCloser(t *testing.T) {
+	input := "hello world"
+	limitBytes := int64(4)
+	reader := io.NopCloser(bytes.NewBufferString(input))
+	r1, r2 := LimitDupReadCloser(reader, limitBytes)
+	verify := func(r io.Reader) {
+		output, err := io.ReadAll(r)
+		assert.Nil(t, err)
+		assert.Equal(t, input, string(output))
+	}
+	verifyLimit := func(r io.Reader, limit int64) {
+		output, err := io.ReadAll(r)
+		if limit < int64(len(input)) {
+			input = input[:limit]
+		}
+		assert.Nil(t, err)
+		assert.Equal(t, input, string(output))
+	}
+
+	verify(r1)
+	verifyLimit(r2, limitBytes)
+}
+
 func TestReadBytes(t *testing.T) {
 	reader := io.NopCloser(bytes.NewBufferString("helloworld"))
 	buf := make([]byte, 5)
diff --git a/core/iox/tee_test.go b/core/iox/tee_test.go
new file mode 100644
index 00000000..56e457c7
--- /dev/null
+++ b/core/iox/tee_test.go
@@ -0,0 +1,35 @@
+package iox
+
+import (
+	"bytes"
+	"io"
+	"testing"
+)
+
+func TestLimitTeeReader(t *testing.T) {
+	limit := int64(4)
+	src := []byte("hello, world")
+	dst := make([]byte, len(src))
+	rb := bytes.NewBuffer(src)
+	wb := new(bytes.Buffer)
+	r := LimitTeeReader(rb, wb, limit)
+	if n, err := io.ReadFull(r, dst); err != nil || n != len(src) {
+		t.Fatalf("ReadFull(r, dst) = %d, %v; want %d, nil", n, err, len(src))
+	}
+	if !bytes.Equal(dst, src) {
+		t.Errorf("bytes read = %q want %q", dst, src)
+	}
+	if !bytes.Equal(wb.Bytes(), src[:limit]) {
+		t.Errorf("bytes written = %q want %q", wb.Bytes(), src)
+	}
+	if n, err := r.Read(dst); n != 0 || err != io.EOF {
+		t.Errorf("r.Read at EOF = %d, %v want 0, EOF", n, err)
+	}
+	rb = bytes.NewBuffer(src)
+	pr, pw := io.Pipe()
+	pr.Close()
+	r = LimitTeeReader(rb, pw, limit)
+	if n, err := io.ReadFull(r, dst); n != 0 || err != io.ErrClosedPipe {
+		t.Errorf("closed tee: ReadFull(r, dst) = %d, %v; want 0, EPIPE", n, err)
+	}
+}