Skip to content

Commit

Permalink
Feat:
Browse files Browse the repository at this point in the history
    - Let "ReadBody" return more readable error
    - Optimize the performance of "CatURL" func
  • Loading branch information
zhyee committed Oct 13, 2023
1 parent abf20c8 commit 7770763
Show file tree
Hide file tree
Showing 5 changed files with 381 additions and 18 deletions.
106 changes: 106 additions & 0 deletions network/http/gin.go
Original file line number Diff line number Diff line change
Expand Up @@ -6,10 +6,14 @@
package http

import (
"bufio"
"bytes"
"compress/gzip"
"crypto/md5" //nolint:gosec
"encoding/hex"
"fmt"
"hash"
"io"
"io/ioutil"
"net"
"net/http"
Expand Down Expand Up @@ -268,6 +272,7 @@ func RequestLoggerMiddleware(c *gin.Context) {
body[:len(body)%MaxRequestBodyLen]+"...")
}

// Deprecated, use GzipReadWithMD5 instead
func GinReadWithMD5(c *gin.Context) (buf []byte, md5str string, err error) {
buf, err = readBody(c)
if err != nil {
Expand All @@ -283,6 +288,7 @@ func GinReadWithMD5(c *gin.Context) (buf []byte, md5str string, err error) {
return
}

// Deprecated, use GzipRead instead.
func GinRead(c *gin.Context) (buf []byte, err error) {
buf, err = readBody(c)
if err != nil {
Expand All @@ -307,6 +313,106 @@ func GinGetArg(c *gin.Context, hdr, param string) (v string, err error) {
return
}

type ReaderWithHash struct {
r io.Reader
h hash.Hash
}

func NewReaderWithHash(r io.Reader, h hash.Hash) *ReaderWithHash {
return &ReaderWithHash{
r: r,
h: h,
}
}

func (h *ReaderWithHash) Read(p []byte) (int, error) {
n, err := h.r.Read(p)
if n > 0 {
_, we := h.h.Write(p[:n])
if we != nil {
return n, fmt.Errorf("unable to write data to hasher: %w", we)
}
}
return n, err
}

func (h *ReaderWithHash) Sum() []byte {
return h.h.Sum(nil)
}

func (h *ReaderWithHash) SumHex() string {
return hex.EncodeToString(h.Sum())
}

func (h *ReaderWithHash) Close() error {
return nil
}

func gzipReadMD5AndClose(req *http.Request, md5Sum bool, closeBody bool) ([]byte, string, error) {

var (
rc io.ReadCloser
rh *ReaderWithHash
)

rc = req.Body
if closeBody {
defer req.Body.Close()
}

if md5Sum {
rh = NewReaderWithHash(rc, md5.New())
defer rh.Close()
rc = rh
}

// as an HTTP server, we do not need to close the Body
switch req.Header.Get("Content-Encoding") {
case "gzip":
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)
}
}

body, 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(body), req.ContentLength)
}

if md5Sum && rh != nil {
return body, rh.SumHex(), nil
}

return body, "", nil

}

// GzipReadWithMD5 will automatically unzip the Request.Body and calculate its MD5 sum,
// it WILL close the Request.Body on return.
func GzipReadWithMD5(req *http.Request) ([]byte, string, error) {
return gzipReadMD5AndClose(req, true, true)
}

// GzipRead will automatically unzip the Request.Body, it WILL close the Request.Body on return.
func GzipRead(req *http.Request) ([]byte, error) {
body, _, err := gzipReadMD5AndClose(req, false, true)
return body, err
}

// Deprecated, you should first consider using GzipRead or GzipReadWithMD5
func Unzip(in []byte) (out []byte, err error) {
gzr, err := gzip.NewReader(bytes.NewBuffer(in))
if err != nil {
Expand Down
163 changes: 163 additions & 0 deletions network/http/gin_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -6,15 +6,44 @@
package http

import (
"bytes"
"compress/gzip"
"crypto/md5"
"encoding/hex"
"fmt"
"github.com/GuanceCloud/cliutils/testutil"
"io"
"net/http"
"strings"
"testing"
"time"

"github.com/gin-gonic/gin"
)

const testText = `观测云提供的系统全链路可观测解决方案,
可实现从底层基础设施到通用技术组件,
再到业务应用系统的全链路可观测,
将不可预知性变为确定已知性。
观测云提供快速实现系统可观测的解决方案,满足云、云原生、应用和业务上的监测需求。
通过自定义监测方案,实现实时可交互仪表板、高效观测基础设施、全链路应用性能可观测等功能,保障系统稳定性
观测云、全链路可观测、实时监测、自定义监测、云原生
观测云提供的系统全链路可观测解决方案,
可实现从底层基础设施到通用技术组件,
再到业务应用系统的全链路可观测,
将不可预知性变为确定已知性。
观测云提供快速实现系统可观测的解决方案,满足云、云原生、应用和业务上的监测需求。
通过自定义监测方案,实现实时可交互仪表板、高效观测基础设施、全链路应用性能可观测等功能,保障系统稳定性
观测云、全链路可观测、实时监测、自定义监测、云原生
观测云提供的系统全链路可观测解决方案,
可实现从底层基础设施到通用技术组件,
再到业务应用系统的全链路可观测,
将不可预知性变为确定已知性。
观测云提供快速实现系统可观测的解决方案,满足云、云原生、应用和业务上的监测需求。
通过自定义监测方案,实现实时可交互仪表板、高效观测基础设施、全链路应用性能可观测等功能,保障系统稳定性
观测云、全链路可观测、实时监测、自定义监测、云原生
`

func BenchmarkAllMiddlewares(b *testing.B) {
cases := []struct {
name string
Expand Down Expand Up @@ -200,3 +229,137 @@ func TestMiddlewares(t *testing.T) {
resp.Body.Close()
}
}

func TestNewHashReader(t *testing.T) {
src := []byte(testText)

r := NewReaderWithHash(bytes.NewReader(src), md5.New())

all, err := io.ReadAll(r)
testutil.Ok(t, err)
testutil.Equals(t, src, all)

md5Sum := md5.Sum(src)
s := r.Sum()

testutil.Equals(t, md5Sum[:], s)

fmt.Println(r.SumHex())
fmt.Println(hex.EncodeToString(md5Sum[:]))
testutil.Equals(t, hex.EncodeToString(md5Sum[:]), r.SumHex())
}

type testCase struct {
name string
body []byte
contentEncoding string
}

func TestGzipReadWithMD5(t *testing.T) {

gzipOut := &bytes.Buffer{}
gw := gzip.NewWriter(gzipOut)
_, err := gw.Write([]byte(testText))
testutil.Ok(t, err)
err = gw.Close()
testutil.Ok(t, err)

testCases := []testCase{
{
name: "plain body",
body: []byte(testText),
contentEncoding: "",
},
{
name: "gzip body",
body: gzipOut.Bytes(),
contentEncoding: "gzip",
},
}

for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {

req1, err := http.NewRequest(http.MethodPost, "/", bytes.NewReader(tc.body))
testutil.Ok(t, err)
if tc.contentEncoding != "" {
req1.Header.Set("Content-Encoding", tc.contentEncoding)
}

req2, err := http.NewRequest(http.MethodPost, "/", bytes.NewReader(tc.body))
testutil.Ok(t, err)
if tc.contentEncoding != "" {
req2.Header.Set("Content-Encoding", tc.contentEncoding)
}

ginCtx := &gin.Context{Request: req1}
body1, sum1, err := GinReadWithMD5(ginCtx)
testutil.Ok(t, err)

body2, sum2, err := GzipReadWithMD5(req2)
testutil.Ok(t, err)

testutil.Equals(t, body1, body2)
testutil.Equals(t, sum1, sum2)

})
}
}

func BenchmarkGzipReadWithMD5(b *testing.B) {

text := strings.Repeat(testText, 100000)

gzipOut := &bytes.Buffer{}
gw := gzip.NewWriter(gzipOut)
_, err := gw.Write([]byte(text))
testutil.Ok(b, err)
err = gw.Close()
testutil.Ok(b, err)

sr := strings.NewReader(text)

b.Run("GinRead", func(t *testing.B) {
sr.Reset(text)

req, err := http.NewRequest(http.MethodPost, "/", sr)
testutil.Ok(t, err)

_, err = GinRead(&gin.Context{Request: req})
testutil.Ok(t, err)
})

b.Run("GzipRead", func(t *testing.B) {

sr.Reset(text)

req, err := http.NewRequest(http.MethodPost, "/", sr)
testutil.Ok(t, err)

_, err = GzipRead(req)
testutil.Ok(t, err)
})

b.Run("GinReadWithMD5", func(t *testing.B) {

req, err := http.NewRequest(http.MethodPost, "/", bytes.NewReader(gzipOut.Bytes()))
testutil.Ok(t, err)
req.Header.Set("Content-Encoding", "gzip")

body, _, err := GinReadWithMD5(&gin.Context{Request: req})
testutil.Ok(t, err)
testutil.Equals(t, string(body), text)
})

b.Run("GzipReadWithMD5", func(t *testing.B) {

req, err := http.NewRequest(http.MethodPost, "/", bytes.NewReader(gzipOut.Bytes()))
testutil.Ok(t, err)
req.Header.Set("Content-Encoding", "gzip")

body, _, err := GzipReadWithMD5(req)
testutil.Ok(t, err)
testutil.Equals(t, string(body), text)
})

}
28 changes: 14 additions & 14 deletions network/http/http.go
Original file line number Diff line number Diff line change
Expand Up @@ -7,22 +7,22 @@
package http

import (
"io/ioutil"
"net/http"
)

// ReadBody will automatically unzip body.
func ReadBody(req *http.Request) ([]byte, error) {
buf, err := ioutil.ReadAll(req.Body)
if err != nil {
return nil, err
}
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}

// as HTTP server, we do not need to close body
switch req.Header.Get("Content-Encoding") {
case "gzip":
return Unzip(buf)
default:
return buf, err
}
// GzipMagic see https://en.wikipedia.org/wiki/Gzip#File_format
GzipMagic = []byte{0x1f, 0x8b}
)

// ReadBody will automatically unzip the body, it doesn't close the Request.Body.
func ReadBody(req *http.Request) ([]byte, error) {
body, _, err := gzipReadMD5AndClose(req, false, false)
return body, err
}
Loading

0 comments on commit 7770763

Please sign in to comment.