Merge pull request #80 from Neur0toxine/fix-body-field

fix for body field writer exhaustion
This commit is contained in:
Pavel 2024-09-26 15:34:16 +03:00 committed by GitHub
commit ac8765ef1e
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 66 additions and 20 deletions

View File

@ -96,6 +96,11 @@ func Body(val any) zap.Field {
if err != nil {
return zap.String(BodyAttr, fmt.Sprintf("%#v", val))
}
if seeker, ok := item.(io.Seeker); ok {
_, _ = seeker.Seek(0, 0)
} else if writer, ok := item.(io.Writer); ok {
_, _ = writer.Write(data)
}
var m interface{}
if err := json.Unmarshal(data, &m); err == nil {
return zap.Any(BodyAttr, m)

View File

@ -4,6 +4,8 @@ import (
"bytes"
"errors"
"fmt"
"github.com/stretchr/testify/require"
"go.uber.org/zap"
"io"
"net/http"
"testing"
@ -62,83 +64,117 @@ func TestHTTPStatusName(t *testing.T) {
func TestStreamID(t *testing.T) {
var cases = []struct {
name string
input interface{}
result interface{}
}{
{
name: "empty",
input: "",
result: "",
},
{
name: "string",
input: "test body",
result: "test body",
},
}
for _, c := range cases {
t.Run(c.name, func(t *testing.T) {
val := StreamID(c.input)
assert.Equal(t, StreamIDAttr, val.Key)
assert.Equal(t, c.result, val.String)
})
}
}
func TestBody(t *testing.T) {
var cases = []struct {
name string
input interface{}
result interface{}
asserts func(t *testing.T, field zap.Field, input, result interface{})
}{
{
name: "empty string input",
input: "",
result: nil,
},
{
name: "nil input",
input: nil,
result: nil,
},
{
name: "string input",
input: "test body",
result: "test body",
},
{
name: "json input",
input: `{"success":true}`,
result: map[string]interface{}{"success": true},
},
{
name: "empty byte slice input",
input: []byte{},
result: nil,
},
{
input: nil,
result: nil,
},
{
name: "byte slice input",
input: []byte("test body"),
result: "test body",
},
{
name: "json byte slice input",
input: []byte(`{"success":true}`),
result: map[string]interface{}{"success": true},
},
{
name: "eof reader input",
input: newReaderMock(func(p []byte) (n int, err error) {
return 0, io.EOF
}),
result: nil,
},
{
name: "empty reader input",
input: newReaderMockData([]byte{}),
result: nil,
},
{
name: "data reader input",
input: newReaderMockData([]byte("ooga booga")),
result: "ooga booga",
},
{
name: "json data reader input",
input: newReaderMockData([]byte(`{"success":true}`)),
result: map[string]interface{}{"success": true},
},
{
name: "check that seeker is rewound",
input: bytes.NewReader([]byte(`{"success":true}`)),
result: map[string]interface{}{"success": true},
asserts: func(t *testing.T, val zap.Field, input, result interface{}) {
data, err := io.ReadAll(input.(io.Reader))
require.NoError(t, err)
assert.Equal(t, []byte(`{"success":true}`), data)
},
},
{
name: "check that writer is rebuilt",
input: bytes.NewBuffer([]byte(`{"success":true}`)),
result: map[string]interface{}{"success": true},
asserts: func(t *testing.T, val zap.Field, input, result interface{}) {
data, err := io.ReadAll(input.(io.Reader))
require.NoError(t, err)
assert.Equal(t, []byte(`{"success":true}`), data)
},
},
}
for _, c := range cases {
t.Run(c.name, func(t *testing.T) {
val := Body(c.input)
assert.Equal(t, BodyAttr, val.Key)
@ -150,6 +186,11 @@ func TestBody(t *testing.T) {
default:
assert.Equal(t, c.result, val.Interface)
}
if c.asserts != nil {
c.asserts(t, val, c.input, c.result)
}
})
}
}