Fix DoTimeout Streaming body bug

valyala · Apr 25, 2022 · 7cc6f4c · 7cc6f4c
1 parent 9a0b4d0
commit 7cc6f4c
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 0 deletions.
diff --git a/http.go b/http.go
@@ -783,6 +783,14 @@ func swapRequestBody(a, b *Request) {
 	a.body, b.body = b.body, a.body
 	a.bodyRaw, b.bodyRaw = b.bodyRaw, a.bodyRaw
 	a.bodyStream, b.bodyStream = b.bodyStream, a.bodyStream
+
+	// This code assumes that if a requestStream was swapped the headers are also swapped or copied.
+	if rs, ok := a.bodyStream.(*requestStream); ok {
+		rs.header = &a.Header
+	}
+	if rs, ok := b.bodyStream.(*requestStream); ok {
+		rs.header = &b.Header
+	}
 }
 
 func swapResponseBody(a, b *Response) {
@@ -1881,6 +1889,9 @@ func (req *Request) closeBodyStream() error {
 	if bsc, ok := req.bodyStream.(io.Closer); ok {
 		err = bsc.Close()
 	}
+	if rs, ok := req.bodyStream.(*requestStream); ok {
+		releaseRequestStream(rs)
+	}
 	req.bodyStream = nil
 	return err
 }

diff --git a/server.go b/server.go
@@ -2440,6 +2440,7 @@ func (s *Server) serveConn(c net.Conn) (err error) {
 			if rs, ok := ctx.Request.bodyStream.(*requestStream); ok {
 				releaseRequestStream(rs)
 			}
+			ctx.Request.bodyStream = nil
 		}
 
 		s.setState(c, StateIdle)

diff --git a/streaming.go b/streaming.go
@@ -98,6 +98,7 @@ func releaseRequestStream(rs *requestStream) {
 	rs.totalBytesRead = 0
 	rs.chunkLeft = 0
 	rs.reader = nil
+	rs.header = nil
 	requestStreamPool.Put(rs)
 }