fix for body field writer exhaustion

This commit is contained in:
Pavel 2024-09-26 15:29:39 +03:00
parent 78fadaa520
commit 9c73cabd21
2 changed files with 66 additions and 20 deletions

View File

@ -96,6 +96,11 @@ func Body(val any) zap.Field {
if err != nil { if err != nil {
return zap.String(BodyAttr, fmt.Sprintf("%#v", val)) 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{} var m interface{}
if err := json.Unmarshal(data, &m); err == nil { if err := json.Unmarshal(data, &m); err == nil {
return zap.Any(BodyAttr, m) return zap.Any(BodyAttr, m)

View File

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