From 7cc6f4c513f9e0d3686142e0a1a5aa2f76b3194a Mon Sep 17 00:00:00 2001
From: Erik Dubbelboer <erik@dubbelboer.com>
Date: Tue, 26 Apr 2022 00:17:15 +0200
Subject: [PATCH] Fix DoTimeout Streaming body bug

---
 http.go      | 11 +++++++++++
 server.go    |  1 +
 streaming.go |  1 +
 3 files changed, 13 insertions(+)

diff --git a/http.go b/http.go
index e8f9ec1c17..384a7b3ffd 100644
--- 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
index 44c973025b..396efdea59 100644
--- 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
index 11750a9d25..fc04916d5e 100644
--- 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)
 }