fix: assert gzip content by http.Header

gogf · Mar 2, 2024 · a07fca8 · a07fca8
1 parent 4b1c473
commit a07fca8
Show file tree

Hide file tree

Showing 6 changed files with 78 additions and 101 deletions.
diff --git a/net/gclient/gclient_tracing.go b/net/gclient/gclient_tracing.go
@@ -91,7 +91,7 @@ func internalMiddlewareTracing(c *Client, r *http.Request) (response *Response,
 	reqBodyContentBytes, _ := io.ReadAll(response.Body)
 	response.Body = utils.NewReadCloser(reqBodyContentBytes, false)
 
-	resBodyContent, err := gtrace.SafeContent(reqBodyContentBytes)
+	resBodyContent, err := gtrace.SafeContent(reqBodyContentBytes, response.Header)
 	if err != nil {
 		span.SetStatus(codes.Error, fmt.Sprintf(`safe res content:%+v`, err))
 	}

diff --git a/net/gclient/gclient_tracing_tracer.go b/net/gclient/gclient_tracing_tracer.go
@@ -152,7 +152,7 @@ func (ct *clientTracer) wroteRequest(info httptrace.WroteRequestInfo) {
 		ct.span.SetStatus(codes.Error, fmt.Sprintf(`%+v`, info.Err))
 	}
 
-	reqBodyContent, err := gtrace.SafeContent(ct.requestBody)
+	reqBodyContent, err := gtrace.SafeContent(ct.requestBody, ct.request.Header)
 	if err != nil {
 		ct.span.SetStatus(codes.Error, fmt.Sprintf(`safe req content:%+v`, err))
 	}

diff --git a/net/ghttp/ghttp_middleware_tracing.go b/net/ghttp/ghttp_middleware_tracing.go
@@ -88,7 +88,7 @@ func internalMiddlewareServerTracing(r *Request) {
 		return
 	}
 	r.Body = utils.NewReadCloser(reqBodyContentBytes, false)
-	reqBodyContent, err := gtrace.SafeContent(reqBodyContentBytes)
+	reqBodyContent, err := gtrace.SafeContent(reqBodyContentBytes, r.Header)
 	if err != nil {
 		span.SetStatus(codes.Error, fmt.Sprintf(`safe req content:%+v`, err))
 	}
@@ -109,7 +109,7 @@ func internalMiddlewareServerTracing(r *Request) {
 	}
 
 	// Response content logging.
-	resBodyContent, err := gtrace.SafeContent(r.Response.Buffer())
+	resBodyContent, err := gtrace.SafeContent(r.Response.Buffer(), r.Response.Header())
 	if err != nil {
 		span.SetStatus(codes.Error, fmt.Sprintf(`safe res content:%+v`, err))
 	}

diff --git a/net/gtrace/gtrace_content.go b/net/gtrace/gtrace_content.go
@@ -7,19 +7,19 @@
 package gtrace
 
 import (
-	"bytes"
-	"compress/gzip"
-	"fmt"
-	"io"
+	"net/http"
+	"strings"
+
+	"github.com/gogf/gf/v2/encoding/gcompress"
 
 	"github.com/gogf/gf/v2/text/gstr"
 )
 
 // SafeContent safe trace content
-func SafeContent(data []byte) (string, error) {
+func SafeContent(data []byte, header http.Header) (string, error) {
 	var err error
-	if IsGzipped(data) {
-		if data, err = unCompressGzip(data); err != nil {
+	if GzipAccepted(header) {
+		if data, err = gcompress.UnGzip(data); err != nil {
 			return string(data), err
 		}
 	}
@@ -32,27 +32,16 @@ func SafeContent(data []byte) (string, error) {
 	return content, nil
 }
 
-// unCompressGzip decode gzip content
-func unCompressGzip(data []byte) ([]byte, error) {
-	reader, err := gzip.NewReader(bytes.NewReader(data))
-	if err != nil {
-		return nil, fmt.Errorf(`read gzip response err:%+v`, err)
-	}
-	defer reader.Close()
-
-	uncompressed, err := io.ReadAll(reader)
-	if err != nil {
-		return data, fmt.Errorf(`get uncompress value err:%+v`, err)
-	}
-
-	return uncompressed, nil
-}
-
-// IsGzipped check content was gzip
-func IsGzipped(input []byte) bool {
-	if len(input) < 2 {
-		return false
+// GzipAccepted returns whether the client will accept gzip-encoded content.
+func GzipAccepted(header http.Header) bool {
+	a := header.Get("Content-Encoding")
+	parts := strings.Split(a, ",")
+	for _, part := range parts {
+		part = strings.TrimSpace(part)
+		if part == "gzip" || strings.HasPrefix(part, "gzip;") {
+			return true
+		}
 	}
 
-	return (int64(input[0])&0xff | (int64(int64(input[1])<<8) & 0xff00)) == 0x8b1f
+	return false
 }
diff --git a/net/gtrace/gtrace_z_unit_content_test.go b/net/gtrace/gtrace_z_unit_content_test.go
diff --git a/net/gtrace/gtrace_z_unit_test.go b/net/gtrace/gtrace_z_unit_test.go
@@ -8,8 +8,12 @@ package gtrace_test
 
 import (
 	"context"
+	"net/http"
+	"strings"
 	"testing"
 
+	"github.com/gogf/gf/v2/encoding/gcompress"
+
 	"github.com/gogf/gf/v2/net/gtrace"
 	"github.com/gogf/gf/v2/test/gtest"
 	"github.com/gogf/gf/v2/text/gstr"
@@ -53,3 +57,56 @@ func TestWithUUID(t *testing.T) {
 		t.Assert(gtrace.GetTraceID(newCtx), gstr.Replace(uuid, "-", ""))
 	})
 }
+
+func TestSafeContent(t *testing.T) {
+	var (
+		defText    = "中"
+		shortData  = strings.Repeat(defText, gtrace.MaxContentLogSize()-1)
+		standData  = strings.Repeat(defText, gtrace.MaxContentLogSize())
+		longData   = strings.Repeat(defText, gtrace.MaxContentLogSize()+1)
+		header     = http.Header{}
+		gzipHeader = http.Header{
+			"Content-Encoding": []string{"gzip"},
+		}
+	)
+	gtest.C(t, func(t *gtest.T) {
+		t.AssertEQ(gtrace.GzipAccepted(gzipHeader), true)
+	})
+
+	// safe content
+	gtest.C(t, func(t *gtest.T) {
+
+		t1, err := gtrace.SafeContent([]byte(shortData), header)
+		t.AssertNil(err)
+		t.Assert(t1, shortData)
+
+		t2, err := gtrace.SafeContent([]byte(standData), header)
+		t.AssertNil(err)
+		t.Assert(t2, standData)
+
+		t3, err := gtrace.SafeContent([]byte(longData), header)
+		t.AssertNil(err)
+		t.Assert(t3, standData+"...")
+	})
+
+	// compress content
+	var (
+		compressShortData, _ = gcompress.Gzip([]byte(shortData))
+		compressStandData, _ = gcompress.Gzip([]byte(standData))
+		compressLongData, _  = gcompress.Gzip([]byte(longData))
+	)
+	gtest.C(t, func(t *gtest.T) {
+
+		t1, err := gtrace.SafeContent(compressShortData, gzipHeader)
+		t.AssertNil(err)
+		t.Assert(t1, shortData)
+
+		t2, err := gtrace.SafeContent(compressStandData, gzipHeader)
+		t.AssertNil(err)
+		t.Assert(t2, standData)
+
+		t3, err := gtrace.SafeContent(compressLongData, gzipHeader)
+		t.AssertNil(err)
+		t.Assert(t3, standData+"...")
+	})
+}