From 2c23e5ba4f17a45acb66af14dfcf926267e9f01b Mon Sep 17 00:00:00 2001
From: zhangdongdong92 <zhangdongdong92@jd.com>
Date: Mon, 23 Jun 2025 11:49:32 +0800
Subject: [PATCH] feat: support Enabled-Redirect

---
 main.go          | 41 +++++++++++++++++++++++++++++++++++++++++
 proxy_handler.go |  2 ++
 2 files changed, 43 insertions(+)

diff --git a/main.go b/main.go
index 84f594dd..9063f991 100644
--- a/main.go
+++ b/main.go
@@ -1,10 +1,12 @@
 package main
 
 import (
+	"bytes"
 	"context"
 	"crypto/tls"
 	"flag"
 	"fmt"
+	"io"
 	"net"
 	"net/http"
 	"os"
@@ -271,6 +273,12 @@ func serveHTTP(rw http.ResponseWriter, r *http.Request) {
 			respondWith(rw, err, http.StatusForbidden)
 			return
 		}
+
+		if r.URL.Query().Get(xRedirectQuery) == "1" {
+			redirect(proxy, rw, r)
+			return
+		}
+
 		proxy.ServeHTTP(rw, r)
 	default:
 		badRequest.Inc()
@@ -280,6 +288,39 @@ func serveHTTP(rw http.ResponseWriter, r *http.Request) {
 	}
 }
 
+func redirect(proxy *reverseProxy, rw http.ResponseWriter, r *http.Request) {
+	name, password := getAuth(r)
+	found, _, c, _ := proxy.getUser(name, password)
+	if !found {
+		rw.Header().Set("Connection", "close")
+		err := fmt.Errorf("invalid username or password for user %q", name)
+		respondWith(rw, err, http.StatusUnauthorized)
+		return
+	}
+
+	targetURL := *r.URL
+	targetURL.Host = c.getReplica().getHost().Host()
+	targetURL.User = *&r.URL.User
+
+	query := targetURL.Query()
+	query.Del(xRedirectQuery)
+	targetURL.RawQuery = query.Encode()
+
+	if r.Body != nil && r.Body != http.NoBody {
+		bodyBytes, err := io.ReadAll(r.Body)
+		if err != nil {
+			respondWith(rw, fmt.Errorf("failed to read request body: %v", err), http.StatusInternalServerError)
+			return
+		}
+		r.Body.Close()
+		r.Body = io.NopCloser(bytes.NewReader(bodyBytes))
+		r.ContentLength = int64(len(bodyBytes))
+	}
+
+	rw.Header().Set("Location", targetURL.String())
+	rw.WriteHeader(http.StatusPermanentRedirect)
+}
+
 func loadConfig() (*config.Config, error) {
 	if *configFile == "" {
 		log.Fatalf("Missing -config flag")
diff --git a/proxy_handler.go b/proxy_handler.go
index bdf69542..81e79d7d 100644
--- a/proxy_handler.go
+++ b/proxy_handler.go
@@ -12,6 +12,8 @@ const (
 	xForwardedForHeader = "X-Forwarded-For"
 	xRealIPHeader       = "X-Real-Ip"
 	forwardedHeader     = "Forwarded"
+
+	xRedirectQuery = "X-Redirect" //Query
 )
 
 type ProxyHandler struct {