From 8c4e7129f33c4397147d1b05bf06bdf4f8d9906c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=88=98=E7=A5=A5=E8=B6=85?= <iwind.liu@gmail.com>
Date: Sat, 9 Apr 2022 20:37:05 +0800
Subject: [PATCH] =?UTF-8?q?=E5=A6=82=E6=9E=9CHeader=E4=B8=ADLocation?=
 =?UTF-8?q?=E5=AD=97=E6=AE=B5=E5=90=AB=E6=9C=89=E8=B7=9F=E6=BA=90=E7=AB=99?=
 =?UTF-8?q?=E4=B8=80=E6=A0=B7=E7=9A=84=E5=9C=B0=E5=9D=80=EF=BC=8C=E5=88=99?=
 =?UTF-8?q?=E8=87=AA=E5=8A=A8=E4=BF=AE=E6=94=B9=E4=B8=BA=E5=BD=93=E5=89=8D?=
 =?UTF-8?q?=E5=9F=9F=E5=90=8D?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 internal/nodes/http_client_pool.go           |  5 +++--
 internal/nodes/http_request_reverse_proxy.go | 16 ++++++++++++++++
 2 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/internal/nodes/http_client_pool.go b/internal/nodes/http_client_pool.go
index 6f5cb26..2503247 100644
--- a/internal/nodes/http_client_pool.go
+++ b/internal/nodes/http_client_pool.go
@@ -142,11 +142,12 @@ func (this *HTTPClientPool) Client(req *HTTPRequest,
 	rawClient = &http.Client{
 		Timeout:   readTimeout,
 		Transport: transport,
-		CheckRedirect: func(req *http.Request, via []*http.Request) error {
+		CheckRedirect: func(targetReq *http.Request, via []*http.Request) error {
+			// 是否跟随
 			if followRedirects {
 				var schemeIsSame = true
 				for _, r := range via {
-					if r.URL.Scheme != req.URL.Scheme {
+					if r.URL.Scheme != targetReq.URL.Scheme {
 						schemeIsSame = false
 						break
 					}
diff --git a/internal/nodes/http_request_reverse_proxy.go b/internal/nodes/http_request_reverse_proxy.go
index a79234a..37ea03d 100644
--- a/internal/nodes/http_request_reverse_proxy.go
+++ b/internal/nodes/http_request_reverse_proxy.go
@@ -279,6 +279,22 @@ func (this *HTTPRequest) doReverseProxy() {
 		}
 	}
 
+	// 替换Location中的源站地址
+	var locationHeader = resp.Header.Get("Location")
+	if len(locationHeader) > 0 {
+		locationURL, err := url.Parse(locationHeader)
+		if err == nil && (locationURL.Host == originAddr || strings.HasPrefix(originAddr, locationURL.Host+":")) {
+			locationURL.Host = this.ReqHost
+			if this.IsHTTP {
+				locationURL.Scheme = "http"
+			} else if this.IsHTTPS {
+				locationURL.Scheme = "https"
+			}
+
+			resp.Header.Set("Location", locationURL.String())
+		}
+	}
+
 	// 响应Header
 	this.writer.AddHeaders(resp.Header)
 	this.processResponseHeaders(resp.StatusCode)