mirror of
https://github.com/seaweedfs/seaweedfs.git
synced 2026-09-17 12:00:44 +02:00
Chris Lu
93d4a6aefd
* s3: drain request body before error response * s3: keep oversized request bodies drainable
119 lines
2.8 KiB
Go
119 lines
2.8 KiB
Go
package s3err
|
|
|
|
import (
|
|
"bytes"
|
|
"io"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"regexp"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/gorilla/mux"
|
|
"github.com/seaweedfs/seaweedfs/weed/util/request_id"
|
|
"github.com/stretchr/testify/assert"
|
|
)
|
|
|
|
func TestWriteErrorResponseReusesRequestID(t *testing.T) {
|
|
req := httptest.NewRequest(http.MethodGet, "/bucket/object", nil)
|
|
req = mux.SetURLVars(req, map[string]string{
|
|
"bucket": "bucket",
|
|
"object": "object",
|
|
})
|
|
req = req.WithContext(request_id.Set(req.Context(), "req-123"))
|
|
|
|
rr := httptest.NewRecorder()
|
|
WriteErrorResponse(rr, req, ErrNoSuchKey)
|
|
|
|
assert.Equal(t, "req-123", rr.Header().Get(request_id.AmzRequestIDHeader))
|
|
assert.Equal(t, "req-123", extractRequestIDFromBody(rr.Body.String()))
|
|
}
|
|
|
|
func TestWriteErrorResponseDrainsRequestBodyBeforeWriting(t *testing.T) {
|
|
body := &trackingReadCloser{data: bytes.Repeat([]byte("a"), 1024)}
|
|
req := httptest.NewRequest(http.MethodPut, "/bucket/object", nil)
|
|
req.Body = body
|
|
req.ContentLength = int64(body.remaining())
|
|
req = mux.SetURLVars(req, map[string]string{
|
|
"bucket": "bucket",
|
|
"object": "object",
|
|
})
|
|
|
|
rr := &drainCheckingResponseWriter{
|
|
header: make(http.Header),
|
|
body: body,
|
|
}
|
|
|
|
WriteErrorResponse(rr, req, ErrInternalError)
|
|
|
|
assert.Empty(t, rr.writeHeaderErr)
|
|
assert.Equal(t, 0, body.remaining())
|
|
assert.Equal(t, http.StatusInternalServerError, rr.status)
|
|
assert.Len(t, rr.readDeadlines, 2)
|
|
assert.False(t, rr.readDeadlines[0].IsZero())
|
|
assert.True(t, rr.readDeadlines[1].IsZero())
|
|
}
|
|
|
|
func extractRequestIDFromBody(body string) string {
|
|
re := regexp.MustCompile(`<RequestId>([^<]+)</RequestId>`)
|
|
matches := re.FindStringSubmatch(body)
|
|
if len(matches) < 2 {
|
|
return ""
|
|
}
|
|
return matches[1]
|
|
}
|
|
|
|
type trackingReadCloser struct {
|
|
data []byte
|
|
}
|
|
|
|
func (t *trackingReadCloser) Read(p []byte) (int, error) {
|
|
if len(t.data) == 0 {
|
|
return 0, io.EOF
|
|
}
|
|
n := copy(p, t.data)
|
|
t.data = t.data[n:]
|
|
return n, nil
|
|
}
|
|
|
|
func (t *trackingReadCloser) Close() error {
|
|
return nil
|
|
}
|
|
|
|
func (t *trackingReadCloser) remaining() int {
|
|
return len(t.data)
|
|
}
|
|
|
|
type drainCheckingResponseWriter struct {
|
|
header http.Header
|
|
body *trackingReadCloser
|
|
status int
|
|
writeHeaderErr string
|
|
readDeadlines []time.Time
|
|
}
|
|
|
|
func (d *drainCheckingResponseWriter) Header() http.Header {
|
|
return d.header
|
|
}
|
|
|
|
func (d *drainCheckingResponseWriter) Write(p []byte) (int, error) {
|
|
if d.status == 0 {
|
|
d.WriteHeader(http.StatusOK)
|
|
}
|
|
return len(p), nil
|
|
}
|
|
|
|
func (d *drainCheckingResponseWriter) WriteHeader(status int) {
|
|
d.status = status
|
|
if d.body.remaining() != 0 {
|
|
d.writeHeaderErr = "request body was not drained before WriteHeader"
|
|
}
|
|
}
|
|
|
|
func (d *drainCheckingResponseWriter) Flush() {}
|
|
|
|
func (d *drainCheckingResponseWriter) SetReadDeadline(t time.Time) error {
|
|
d.readDeadlines = append(d.readDeadlines, t)
|
|
return nil
|
|
}
|