From 7ff777ad0d97312b850adf0fdedcb12d689ad22d Mon Sep 17 00:00:00 2001
From: zhyee <1054948153@qq.com>
Date: Thu, 12 Oct 2023 15:10:29 +0800
Subject: [PATCH] Feat:     - Let "ReadBody" return more readable error     -
 Optimize the performance of "CatURL" func

---
 network/http/http.go      | 55 +++++++++++++++++++++-----
 network/http/http_test.go | 83 +++++++++++++++++++++++++++++++++++++++
 point/category.go         | 17 ++++++--
 3 files changed, 141 insertions(+), 14 deletions(-)

diff --git a/network/http/http.go b/network/http/http.go
index 0a1056cc..ff32bf3d 100644
--- a/network/http/http.go
+++ b/network/http/http.go
@@ -7,22 +7,57 @@
 package http
 
 import (
-	"io/ioutil"
+	"bufio"
+	"bytes"
+	"compress/gzip"
+	"fmt"
+	"io"
 	"net/http"
 )
 
-// ReadBody will automatically unzip body.
+var (
+	// ZIPMagic see https://en.wikipedia.org/wiki/ZIP_(file_format)#Local_file_header
+	ZIPMagic = []byte{0x50, 0x4b, 0x3, 0x4} //
+
+	// LZ4Magic see https://android.googlesource.com/platform/external/lz4/+/HEAD/doc/lz4_Frame_format.md#general-structure-of-lz4-frame-format
+	LZ4Magic = []byte{0x4, 0x22, 0x4d, 0x18}
+
+	// GzipMagic see https://en.wikipedia.org/wiki/Gzip#File_format
+	GzipMagic = []byte{0x1f, 0x8b}
+)
+
+// ReadBody will automatically unzip the body.
 func ReadBody(req *http.Request) ([]byte, error) {
-	buf, err := ioutil.ReadAll(req.Body)
-	if err != nil {
-		return nil, err
-	}
+	var (
+		rc  io.ReadCloser = req.Body
+		err error
+	)
 
-	// as HTTP server, we do not need to close body
+	// as an HTTP server, we do not need to close the Body
 	switch req.Header.Get("Content-Encoding") {
 	case "gzip":
-		return Unzip(buf)
-	default:
-		return buf, err
+		bufReader := bufio.NewReader(rc)
+		magic, err := bufReader.Peek(len(GzipMagic))
+		if err != nil {
+			return nil, fmt.Errorf("unable to peek 2 bytes from Body: %w", err)
+		}
+
+		if bytes.Compare(GzipMagic, magic) == 0 {
+			rc, err = gzip.NewReader(bufReader)
+			if err != nil {
+				return nil, fmt.Errorf("unable to init gzip reader: %w", err)
+			}
+			defer rc.Close()
+		} else {
+			l.Warnf(`illegal gzip format while Content-Encoding = gzip, magic %v expected, got %v`, GzipMagic, magic)
+			rc = io.NopCloser(bufReader)
+		}
 	}
+
+	buf, err := io.ReadAll(rc)
+	if err != nil {
+		return nil, fmt.Errorf("unable to successfully read: %w, bytes has read: %d, http Content-Length: %d", err, len(buf), req.ContentLength)
+	}
+
+	return buf, nil
 }
diff --git a/network/http/http_test.go b/network/http/http_test.go
index 9146255d..fac8e4e9 100644
--- a/network/http/http_test.go
+++ b/network/http/http_test.go
@@ -6,14 +6,21 @@
 package http
 
 import (
+	"archive/zip"
+	"bytes"
+	"compress/gzip"
+	"io"
 
 	// nolint:gosec
 	"crypto/md5"
 	"fmt"
 	"log"
 	"net/http"
+	"strings"
 	"testing"
 	"time"
+
+	"github.com/GuanceCloud/cliutils/testutil"
 )
 
 func TestSign(t *testing.T) {
@@ -58,3 +65,79 @@ func TestSign(t *testing.T) {
 	}
 	log.Printf("[debug] %+#v", o)
 }
+
+func TestMagic(t *testing.T) {
+	buf := &bytes.Buffer{}
+	w := gzip.NewWriter(buf)
+	_, err := w.Write(LZ4Magic)
+	testutil.Ok(t, err)
+	err = w.Close()
+	testutil.Ok(t, err)
+
+	if bytes.Compare(GzipMagic, buf.Bytes()[:len(GzipMagic)]) != 0 {
+		t.Fatalf("gzip magic: %v expected, got: %v", GzipMagic, buf.Bytes()[:len(GzipMagic)])
+	}
+
+	zipBuf := &bytes.Buffer{}
+	zw := zip.NewWriter(zipBuf)
+	f, err := zw.Create("tmp")
+	testutil.Ok(t, err)
+	_, err = f.Write(GzipMagic)
+	testutil.Ok(t, err)
+	err = zw.Close()
+	testutil.Ok(t, err)
+
+	if bytes.Compare(ZIPMagic, zipBuf.Bytes()[:len(ZIPMagic)]) != 0 {
+		t.Fatalf("zip magic: %v expected, got: %v", ZIPMagic, buf.Bytes()[:len(ZIPMagic)])
+	}
+}
+
+func TestReadBody(t *testing.T) {
+	text := `少小离家老大回，
+乡音无改鬓毛衰。
+`
+	type testCase struct {
+		name            string
+		reqBody         io.Reader
+		contentEncoding string
+	}
+
+	buf := &bytes.Buffer{}
+	gw := gzip.NewWriter(buf)
+	_, err := gw.Write([]byte(text))
+	testutil.Ok(t, err)
+	err = gw.Close()
+	testutil.Ok(t, err)
+
+	cases := []testCase{
+		{
+			name:            "plain",
+			reqBody:         strings.NewReader(text),
+			contentEncoding: "",
+		},
+		{
+			name:            "gzip",
+			reqBody:         buf,
+			contentEncoding: "gzip",
+		},
+		{
+			name:            "invalid gzip",
+			reqBody:         strings.NewReader(text),
+			contentEncoding: "gzip",
+		},
+	}
+
+	for _, c := range cases {
+		t.Run(c.name, func(t *testing.T) {
+			req, err := http.NewRequest(http.MethodPost, "/", c.reqBody)
+			testutil.Ok(t, err)
+			if c.contentEncoding != "" {
+				req.Header.Set("Content-Encoding", c.contentEncoding)
+			}
+			body, err := ReadBody(req)
+			testutil.Ok(t, err)
+			testutil.Equals(t, text, string(body))
+		})
+	}
+
+}
diff --git a/point/category.go b/point/category.go
index e9a64753..b1527b0b 100644
--- a/point/category.go
+++ b/point/category.go
@@ -50,10 +50,8 @@ func CatString(c string) Category {
 }
 
 func CatURL(c string) Category {
-	for k, v := range categoryURL {
-		if c == v {
-			return k
-		}
+	if category, ok := URLCategoryMap[c]; ok {
+		return category
 	}
 	return UnknownCategory
 }
@@ -152,6 +150,8 @@ var (
 		UnknownCategory: URLUnknownCategory,
 	}
 
+	URLCategoryMap = flipMap(categoryURL)
+
 	categoryAias = map[Category]string{
 		Metric:            CM,
 		Network:           CN,
@@ -183,3 +183,12 @@ var (
 		DynamicDWCategory: SDynamicDWCategory,
 	}
 )
+
+// Exchanges all keys with their associated values in a map
+func flipMap(m map[Category]string) map[string]Category {
+	fm := make(map[string]Category, len(m))
+	for c, url := range m {
+		fm[url] = c
+	}
+	return fm
+}