From 651eabb4c6ba64b19b36d2dfc45cec564ae0eb58 Mon Sep 17 00:00:00 2001
From: Kevin Wan <wanjunfeng@gmail.com>
Date: Sat, 8 Feb 2025 21:18:07 +0800
Subject: [PATCH] chore: refactor gateway http context (#4636)

---
 gateway/server.go | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/gateway/server.go b/gateway/server.go
index e1af48cef..fe2703953 100644
--- a/gateway/server.go
+++ b/gateway/server.go
@@ -185,13 +185,13 @@ func (s *Server) buildHttpHandler(target *HttpClientConf) http.HandlerFunc {
 			return
 		}
 
+		// set the timeout if it's configured, take effect only if it's greater than 0
+		// and less than the deadline of the original request
 		if target.Timeout > 0 {
 			timeout := time.Duration(target.Timeout) * time.Millisecond
 			ctx, cancel := context.WithTimeout(r.Context(), timeout)
 			defer cancel()
 			req = req.WithContext(ctx)
-		} else {
-			req = req.WithContext(r.Context())
 		}
 
 		resp, err := httpc.DoRequest(req)
@@ -278,7 +278,7 @@ func buildRequestWithNewTarget(r *http.Request, target *HttpClientConf) (*http.R
 		}
 	}
 
-	return &http.Request{
+	newReq := &http.Request{
 		Method:        r.Method,
 		URL:           &u,
 		Header:        r.Header.Clone(),
@@ -287,7 +287,10 @@ func buildRequestWithNewTarget(r *http.Request, target *HttpClientConf) (*http.R
 		ProtoMinor:    r.ProtoMinor,
 		ContentLength: r.ContentLength,
 		Body:          io.NopCloser(r.Body),
-	}, nil
+	}
+
+	// make sure the context is passed to the new request
+	return newReq.WithContext(r.Context()), nil
 }
 
 func createDescriptorSource(cli zrpc.Client, up Upstream) (grpcurl.DescriptorSource, error) {