From 4a6bed77284590fa3ba2236854511790d0ad4e87 Mon Sep 17 00:00:00 2001
From: yusing <yusing@6uo.me>
Date: Mon, 26 May 2025 19:57:39 +0800
Subject: [PATCH] refactor: headers handling

---
 .../net/gphttp/middleware/modify_request.go   | 64 ++++++++++---------
 .../net/gphttp/middleware/modify_response.go  |  6 +-
 .../net/gphttp/reverseproxy/reverse_proxy.go  |  4 +-
 3 files changed, 41 insertions(+), 33 deletions(-)

diff --git a/internal/net/gphttp/middleware/modify_request.go b/internal/net/gphttp/middleware/modify_request.go
index ee665be..81c3b4f 100644
--- a/internal/net/gphttp/middleware/modify_request.go
+++ b/internal/net/gphttp/middleware/modify_request.go
@@ -33,8 +33,14 @@ func (mr *ModifyRequestOpts) finalize() {
 func (mr *modifyRequest) before(w http.ResponseWriter, r *http.Request) (proceed bool) {
 	mr.AddTraceRequest("before modify request", r)
 
-	mr.addPrefix(r, nil, r.URL.Path)
-	mr.modifyHeaders(r, nil, r.Header)
+	if len(mr.AddPrefix) != 0 {
+		mr.addPrefix(r, r.URL.Path)
+	}
+	if !mr.needVarSubstitution {
+		mr.modifyHeaders(r, r.Header)
+	} else {
+		mr.modifyHeadersWithVarSubstitution(r, nil, r.Header)
+	}
 	mr.AddTraceRequest("after modify request", r)
 	return true
 }
@@ -50,42 +56,40 @@ func (mr *ModifyRequestOpts) checkVarSubstitution() {
 	}
 }
 
-func (mr *ModifyRequestOpts) modifyHeaders(req *http.Request, resp *http.Response, headers http.Header) {
-	if !mr.needVarSubstitution {
-		for k, v := range mr.SetHeaders {
-			if req != nil && strings.EqualFold(k, "host") {
-				defer func() {
-					req.Host = v
-				}()
-			}
-			headers[k] = []string{v}
-		}
-		for k, v := range mr.AddHeaders {
-			headers[k] = append(headers[k], v)
-		}
-	} else {
-		for k, v := range mr.SetHeaders {
-			if req != nil && strings.EqualFold(k, "host") {
-				defer func() {
-					req.Host = varReplace(req, resp, v)
-				}()
-			}
-			headers[k] = []string{varReplace(req, resp, v)}
-		}
-		for k, v := range mr.AddHeaders {
-			headers[k] = append(headers[k], varReplace(req, resp, v))
+func (mr *ModifyRequestOpts) modifyHeaders(req *http.Request, headers http.Header) {
+	for k, v := range mr.SetHeaders {
+		if req != nil && strings.EqualFold(k, "host") {
+			defer func() {
+				req.Host = v
+			}()
 		}
+		headers[k] = []string{v}
+	}
+	for k, v := range mr.AddHeaders {
+		headers[k] = append(headers[k], v)
 	}
-
 	for _, k := range mr.HideHeaders {
 		delete(headers, k)
 	}
 }
 
-func (mr *modifyRequest) addPrefix(r *http.Request, _ *http.Response, path string) {
-	if len(mr.AddPrefix) == 0 {
-		return
+func (mr *ModifyRequestOpts) modifyHeadersWithVarSubstitution(req *http.Request, resp *http.Response, headers http.Header) {
+	for k, v := range mr.SetHeaders {
+		if req != nil && strings.EqualFold(k, "host") {
+			defer func() {
+				req.Host = varReplace(req, resp, v)
+			}()
+		}
+		headers[k] = []string{varReplace(req, resp, v)}
 	}
+	for k, v := range mr.AddHeaders {
+		headers[k] = append(headers[k], varReplace(req, resp, v))
+	}
+	for _, k := range mr.HideHeaders {
+		delete(headers, k)
+	}
+}
 
+func (mr *modifyRequest) addPrefix(r *http.Request, path string) {
 	r.URL.Path = filepath.Join(mr.AddPrefix, path)
 }
diff --git a/internal/net/gphttp/middleware/modify_response.go b/internal/net/gphttp/middleware/modify_response.go
index 3285d9b..e110cc5 100644
--- a/internal/net/gphttp/middleware/modify_response.go
+++ b/internal/net/gphttp/middleware/modify_response.go
@@ -14,7 +14,11 @@ var ModifyResponse = NewMiddleware[modifyResponse]()
 // modifyResponse implements ResponseModifier.
 func (mr *modifyResponse) modifyResponse(resp *http.Response) error {
 	mr.AddTraceResponse("before modify response", resp)
-	mr.modifyHeaders(resp.Request, resp, resp.Header)
+	if !mr.needVarSubstitution {
+		mr.modifyHeaders(resp.Request, resp.Header)
+	} else {
+		mr.modifyHeadersWithVarSubstitution(resp.Request, resp, resp.Header)
+	}
 	mr.AddTraceResponse("after modify response", resp)
 	return nil
 }
diff --git a/internal/net/gphttp/reverseproxy/reverse_proxy.go b/internal/net/gphttp/reverseproxy/reverse_proxy.go
index 42ddcbc..3ede2c3 100644
--- a/internal/net/gphttp/reverseproxy/reverse_proxy.go
+++ b/internal/net/gphttp/reverseproxy/reverse_proxy.go
@@ -271,7 +271,7 @@ func (p *ReverseProxy) handler(rw http.ResponseWriter, req *http.Request) {
 		return
 	}
 
-	req.Header.Del("Forwarded")
+	outreq.Header.Del("Forwarded")
 	httpheaders.RemoveHopByHopHeaders(outreq.Header)
 
 	// Issue 21096: tell backend applications that care about trailer support
@@ -312,7 +312,7 @@ func (p *ReverseProxy) handler(rw http.ResponseWriter, req *http.Request) {
 	}
 
 	var reqScheme string
-	if req.TLS != nil {
+	if req.TLS != nil || req.Header.Get("X-Forwarded-Proto") == "https" {
 		reqScheme = "https"
 	} else {
 		reqScheme = "http"