feat: stream native request bodies
This commit is contained in:
@@ -410,12 +410,49 @@ func (s *Server) serveConnection(baseCtx context.Context, state *connection) {
|
||||
}
|
||||
}
|
||||
handler := s.handler
|
||||
disconnectWatch := newClientDisconnectWatch(reader, state.conn, cancel)
|
||||
var disconnectWatch *clientDisconnectWatch
|
||||
if body, ok := request.Body.(*requestBody); ok {
|
||||
// A body reader and the disconnect watcher must never read the
|
||||
// bufio.Reader concurrently. The body starts the watcher once it
|
||||
// reaches EOF; body read errors cancel the request directly.
|
||||
body.cancel = cancel
|
||||
body.onComplete = func() {
|
||||
disconnectWatch = newClientDisconnectWatch(reader, state.conn, cancel)
|
||||
}
|
||||
} else {
|
||||
disconnectWatch = newClientDisconnectWatch(reader, state.conn, cancel)
|
||||
}
|
||||
handler.ServeHTTP(writer, request)
|
||||
disconnectWatch.stop()
|
||||
bodyErr := request.Body.Close()
|
||||
if disconnectWatch != nil {
|
||||
disconnectWatch.stop()
|
||||
}
|
||||
cancel()
|
||||
|
||||
closing := s.isClosing() || state.closeWhenIdle.Load()
|
||||
closing := s.isClosing() || state.closeWhenIdle.Load() || bodyErr != nil
|
||||
if bodyErr != nil {
|
||||
var protocolErr *protocolError
|
||||
if !errors.As(bodyErr, &protocolErr) {
|
||||
if requestSlot {
|
||||
<-s.requestSlots
|
||||
}
|
||||
writer.reset("")
|
||||
s.writerPool.Put(writer)
|
||||
state.active.Store(false)
|
||||
return
|
||||
}
|
||||
// Request bodies are streamed into the handler, so a malformed
|
||||
// chunk or an over-limit body can be discovered while draining
|
||||
// after the handler returns. Replace the buffered response with
|
||||
// the same protocol error that eager parsing used to return.
|
||||
if !writer.streaming {
|
||||
writer.reset(request.Method)
|
||||
writer.configure(request, state.conn, true, s.WriteTimeout, cancel)
|
||||
writer.Header().Set("Content-Type", "text/plain; charset=utf-8")
|
||||
writer.WriteHeader(protocolErr.status)
|
||||
writer.body.WriteString(protocolErr.message + "\n")
|
||||
}
|
||||
}
|
||||
if err := writer.writeTo(state.conn, closeAfter || closing, s.WriteTimeout); err != nil {
|
||||
if requestSlot {
|
||||
<-s.requestSlots
|
||||
@@ -577,27 +614,23 @@ func readRequest(reader *bufio.Reader, conn net.Conn, maxHeaderBytes int, maxBod
|
||||
if !chunked && contentLength > maxBodyBytes {
|
||||
return nil, true, &protocolError{status: http.StatusRequestEntityTooLarge, message: "request body too large"}
|
||||
}
|
||||
var body []byte
|
||||
var body io.ReadCloser = http.NoBody
|
||||
var trailer http.Header
|
||||
if chunked {
|
||||
body, trailer, err = readChunkedBody(reader, maxBodyBytes, setBodyDeadline, readLine)
|
||||
if err != nil {
|
||||
return nil, true, err
|
||||
trailer = make(http.Header)
|
||||
body = &requestBody{
|
||||
reader: reader, maxBytes: maxBodyBytes, chunked: true,
|
||||
trailer: trailer, setDeadline: setBodyDeadline, readLine: readLine,
|
||||
}
|
||||
// net/http uses -1 to indicate that the decoded length was not known
|
||||
// from the request headers. The body is already buffered for the
|
||||
// handler, but retaining this value keeps the request contract
|
||||
// compatible with the standard server.
|
||||
// from the request headers.
|
||||
contentLength = -1
|
||||
} else {
|
||||
body = make([]byte, contentLength)
|
||||
if contentLength > 0 {
|
||||
if setBodyDeadline != nil {
|
||||
setBodyDeadline()
|
||||
}
|
||||
if _, err := io.ReadFull(reader, body); err != nil {
|
||||
return nil, true, err
|
||||
}
|
||||
body = &requestBody{reader: reader, remaining: contentLength, setDeadline: setBodyDeadline}
|
||||
}
|
||||
}
|
||||
request := &http.Request{
|
||||
@@ -607,7 +640,7 @@ func readRequest(reader *bufio.Reader, conn net.Conn, maxHeaderBytes int, maxBod
|
||||
ProtoMajor: protoMajor,
|
||||
ProtoMinor: protoMinor,
|
||||
Header: header,
|
||||
Body: io.NopCloser(bytes.NewReader(body)),
|
||||
Body: body,
|
||||
ContentLength: contentLength,
|
||||
Host: header.Get("Host"),
|
||||
RequestURI: target,
|
||||
@@ -627,78 +660,199 @@ func readRequest(reader *bufio.Reader, conn net.Conn, maxHeaderBytes int, maxBod
|
||||
return request, closeAfter, nil
|
||||
}
|
||||
|
||||
func readChunkedBody(reader *bufio.Reader, maxBodyBytes int64, setBodyDeadline func(), readLine func() (string, error)) ([]byte, http.Header, error) {
|
||||
var body bytes.Buffer
|
||||
initialCapacity := int64(64 << 10)
|
||||
if maxBodyBytes < initialCapacity {
|
||||
initialCapacity = maxBodyBytes
|
||||
}
|
||||
if initialCapacity > 0 && initialCapacity <= int64(int(^uint(0)>>1)) {
|
||||
body.Grow(int(initialCapacity))
|
||||
}
|
||||
type requestBody struct {
|
||||
reader *bufio.Reader
|
||||
remaining int64
|
||||
maxBytes int64
|
||||
readBytes int64
|
||||
chunked bool
|
||||
chunkRemain int64
|
||||
needSuffix bool
|
||||
trailer http.Header
|
||||
readLine func() (string, error)
|
||||
setDeadline func()
|
||||
started bool
|
||||
cancel context.CancelFunc
|
||||
onComplete func()
|
||||
done bool
|
||||
closed bool
|
||||
err error
|
||||
}
|
||||
|
||||
func (b *requestBody) Read(p []byte) (int, error) {
|
||||
if len(p) == 0 {
|
||||
return 0, nil
|
||||
}
|
||||
if b.err != nil {
|
||||
return 0, b.err
|
||||
}
|
||||
if b.done {
|
||||
return 0, io.EOF
|
||||
}
|
||||
if !b.chunked {
|
||||
return b.readContentLength(p)
|
||||
}
|
||||
return b.readChunked(p)
|
||||
}
|
||||
|
||||
func (b *requestBody) readContentLength(p []byte) (int, error) {
|
||||
if b.remaining == 0 {
|
||||
b.complete()
|
||||
return 0, io.EOF
|
||||
}
|
||||
if !b.started {
|
||||
b.started = true
|
||||
if b.setDeadline != nil {
|
||||
b.setDeadline()
|
||||
}
|
||||
}
|
||||
if int64(len(p)) > b.remaining {
|
||||
p = p[:b.remaining]
|
||||
}
|
||||
n, err := b.reader.Read(p)
|
||||
b.remaining -= int64(n)
|
||||
b.readBytes += int64(n)
|
||||
if b.remaining == 0 {
|
||||
b.complete()
|
||||
if n > 0 {
|
||||
return n, nil
|
||||
}
|
||||
return 0, io.EOF
|
||||
}
|
||||
if err != nil {
|
||||
b.fail(err)
|
||||
}
|
||||
return n, err
|
||||
}
|
||||
|
||||
func (b *requestBody) readChunked(p []byte) (int, error) {
|
||||
for {
|
||||
line, err := readLine()
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
chunkSizeText := line
|
||||
if extension := strings.IndexByte(chunkSizeText, ';'); extension >= 0 {
|
||||
chunkSizeText = chunkSizeText[:extension]
|
||||
}
|
||||
chunkSizeText = strings.TrimSpace(chunkSizeText)
|
||||
if chunkSizeText == "" || !validHexNumber(chunkSizeText) {
|
||||
return nil, nil, &protocolError{status: http.StatusBadRequest, message: "invalid chunk size"}
|
||||
}
|
||||
chunkSize, parseErr := strconv.ParseUint(chunkSizeText, 16, 63)
|
||||
if parseErr != nil {
|
||||
return nil, nil, &protocolError{status: http.StatusBadRequest, message: "invalid chunk size"}
|
||||
}
|
||||
if chunkSize == 0 {
|
||||
trailer := make(http.Header)
|
||||
for {
|
||||
line, err = readLine()
|
||||
if err != nil {
|
||||
return nil, nil, err
|
||||
if b.chunkRemain == 0 {
|
||||
if b.needSuffix {
|
||||
var suffix [2]byte
|
||||
if _, err := io.ReadFull(b.reader, suffix[:]); err != nil {
|
||||
b.fail(err)
|
||||
return 0, err
|
||||
}
|
||||
if line == "" {
|
||||
return body.Bytes(), trailer, nil
|
||||
if suffix[0] != '\r' || suffix[1] != '\n' {
|
||||
err := &protocolError{status: http.StatusBadRequest, message: "malformed chunk terminator"}
|
||||
b.fail(err)
|
||||
return 0, err
|
||||
}
|
||||
if strings.HasPrefix(line, " ") || strings.HasPrefix(line, "\t") {
|
||||
return nil, nil, &protocolError{status: http.StatusBadRequest, message: "folded trailers are not supported"}
|
||||
}
|
||||
index := strings.IndexByte(line, ':')
|
||||
if index <= 0 {
|
||||
return nil, nil, &protocolError{status: http.StatusBadRequest, message: "malformed trailer"}
|
||||
}
|
||||
name := strings.TrimSpace(line[:index])
|
||||
value := strings.TrimSpace(line[index+1:])
|
||||
if !validToken(name) || !validHeaderValue(value) {
|
||||
return nil, nil, &protocolError{status: http.StatusBadRequest, message: "malformed trailer"}
|
||||
}
|
||||
trailer.Add(name, value)
|
||||
b.needSuffix = false
|
||||
}
|
||||
if err := b.beginChunk(); err != nil {
|
||||
b.fail(err)
|
||||
return 0, err
|
||||
}
|
||||
if b.done {
|
||||
return 0, io.EOF
|
||||
}
|
||||
}
|
||||
if int64(body.Len()) > maxBodyBytes-int64(chunkSize) {
|
||||
return nil, nil, &protocolError{status: http.StatusRequestEntityTooLarge, message: "request body too large"}
|
||||
if int64(len(p)) > b.chunkRemain {
|
||||
p = p[:b.chunkRemain]
|
||||
}
|
||||
if setBodyDeadline != nil {
|
||||
setBodyDeadline()
|
||||
n, err := b.reader.Read(p)
|
||||
b.chunkRemain -= int64(n)
|
||||
b.readBytes += int64(n)
|
||||
if b.chunkRemain == 0 {
|
||||
b.needSuffix = true
|
||||
}
|
||||
if _, err := io.CopyN(&body, reader, int64(chunkSize)); err != nil {
|
||||
return nil, nil, err
|
||||
if n > 0 {
|
||||
if err != nil && b.chunkRemain > 0 {
|
||||
b.fail(err)
|
||||
return n, err
|
||||
}
|
||||
return n, nil
|
||||
}
|
||||
if setBodyDeadline != nil {
|
||||
setBodyDeadline()
|
||||
if err != nil {
|
||||
b.fail(err)
|
||||
return 0, err
|
||||
}
|
||||
var suffix [2]byte
|
||||
if _, err := io.ReadFull(reader, suffix[:]); err != nil {
|
||||
return nil, nil, err
|
||||
}
|
||||
if suffix[0] != '\r' || suffix[1] != '\n' {
|
||||
return nil, nil, &protocolError{status: http.StatusBadRequest, message: "malformed chunk terminator"}
|
||||
}
|
||||
}
|
||||
|
||||
func (b *requestBody) beginChunk() error {
|
||||
if b.setDeadline != nil {
|
||||
b.setDeadline()
|
||||
}
|
||||
line, err := b.readLine()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
chunkSizeText := line
|
||||
if extension := strings.IndexByte(chunkSizeText, ';'); extension >= 0 {
|
||||
chunkSizeText = chunkSizeText[:extension]
|
||||
}
|
||||
chunkSizeText = strings.TrimSpace(chunkSizeText)
|
||||
if chunkSizeText == "" || !validHexNumber(chunkSizeText) {
|
||||
return &protocolError{status: http.StatusBadRequest, message: "invalid chunk size"}
|
||||
}
|
||||
chunkSize, parseErr := strconv.ParseUint(chunkSizeText, 16, 63)
|
||||
if parseErr != nil {
|
||||
return &protocolError{status: http.StatusBadRequest, message: "invalid chunk size"}
|
||||
}
|
||||
if chunkSize == 0 {
|
||||
for {
|
||||
line, err = b.readLine()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if line == "" {
|
||||
b.complete()
|
||||
return nil
|
||||
}
|
||||
if strings.HasPrefix(line, " ") || strings.HasPrefix(line, "\t") {
|
||||
return &protocolError{status: http.StatusBadRequest, message: "folded trailers are not supported"}
|
||||
}
|
||||
index := strings.IndexByte(line, ':')
|
||||
if index <= 0 {
|
||||
return &protocolError{status: http.StatusBadRequest, message: "malformed trailer"}
|
||||
}
|
||||
name := strings.TrimSpace(line[:index])
|
||||
value := strings.TrimSpace(line[index+1:])
|
||||
if !validToken(name) || !validHeaderValue(value) {
|
||||
return &protocolError{status: http.StatusBadRequest, message: "malformed trailer"}
|
||||
}
|
||||
b.trailer.Add(name, value)
|
||||
}
|
||||
}
|
||||
if int64(chunkSize) > b.maxBytes-b.readBytes {
|
||||
return &protocolError{status: http.StatusRequestEntityTooLarge, message: "request body too large"}
|
||||
}
|
||||
b.chunkRemain = int64(chunkSize)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (b *requestBody) Close() error {
|
||||
if b.closed {
|
||||
return b.err
|
||||
}
|
||||
b.closed = true
|
||||
if !b.done && b.err == nil {
|
||||
_, _ = io.Copy(io.Discard, b)
|
||||
}
|
||||
return b.err
|
||||
}
|
||||
|
||||
func (b *requestBody) complete() {
|
||||
if b.done {
|
||||
return
|
||||
}
|
||||
b.done = true
|
||||
if b.onComplete != nil {
|
||||
b.onComplete()
|
||||
}
|
||||
}
|
||||
|
||||
func (b *requestBody) fail(err error) {
|
||||
if b.err == nil {
|
||||
b.err = err
|
||||
}
|
||||
b.done = true
|
||||
if b.cancel != nil {
|
||||
b.cancel()
|
||||
}
|
||||
}
|
||||
|
||||
func validHexNumber(value string) bool {
|
||||
|
||||
Reference in New Issue
Block a user