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 12, 2023
1 parent abf20c8 commit 7ff777a
Show file tree
Hide file tree
Showing 3 changed files with 141 additions and 14 deletions.
55 changes: 45 additions & 10 deletions network/http/http.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
83 changes: 83 additions & 0 deletions network/http/http_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand Down Expand Up @@ -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))
})
}

}
17 changes: 13 additions & 4 deletions point/category.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
Expand Down Expand Up @@ -152,6 +150,8 @@ var (
UnknownCategory: URLUnknownCategory,
}

URLCategoryMap = flipMap(categoryURL)

categoryAias = map[Category]string{
Metric: CM,
Network: CN,
Expand Down Expand Up @@ -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
}

0 comments on commit 7ff777a

Please sign in to comment.