From cc21f5fae2356382e1564e43bb8192b0812b7d32 Mon Sep 17 00:00:00 2001 From: Awadabang <49111776+Awadabang@users.noreply.github.com> Date: Sat, 16 Sep 2023 19:58:21 +0800 Subject: [PATCH] update: limit logBrief http body size (#3498) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: 常公征 --- core/iox/read.go | 10 ++++++++++ core/iox/read_test.go | 23 +++++++++++++++++++++++ core/iox/tee.go | 33 +++++++++++++++++++++++++++++++++ core/iox/tee_test.go | 35 +++++++++++++++++++++++++++++++++++ rest/handler/loghandler.go | 12 ++---------- 5 files changed, 103 insertions(+), 10 deletions(-) create mode 100644 core/iox/tee.go create mode 100644 core/iox/tee_test.go 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/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.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/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) + } +} 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()