From 3db20563d25b171195e33f035c61512a6af4cdd9 Mon Sep 17 00:00:00 2001
From: ginuerzh <ginuerzh@gmail.com>
Date: Thu, 7 Nov 2024 11:04:42 +0800
Subject: [PATCH] add custom http response header for reverse proxy

---
 config/config.go                   |  2 ++
 config/parsing/node/parse.go       |  5 +++--
 go.mod                             |  2 +-
 go.sum                             |  4 ++--
 internal/util/forwarder/sniffer.go | 11 +++++++++++
 5 files changed, 19 insertions(+), 5 deletions(-)

diff --git a/config/config.go b/config/config.go
index 957ab8e1..edc8546f 100644
--- a/config/config.go
+++ b/config/config.go
@@ -398,6 +398,8 @@ type HTTPNodeConfig struct {
 	Host string `yaml:",omitempty" json:"host,omitempty"`
 	// additional request header
 	Header map[string]string `yaml:",omitempty" json:"header,omitempty"`
+	// additional response header
+	ResponseHeader map[string]string `yaml:",omitempty" json:"responseHeader,omitempty"`
 	// rewrite URL
 	Rewrite []HTTPURLRewriteConfig `yaml:",omitempty" json:"rewrite,omitempty"`
 	// rewrite response body
diff --git a/config/parsing/node/parse.go b/config/parsing/node/parse.go
index 69a9b1ea..413107f6 100644
--- a/config/parsing/node/parse.go
+++ b/config/parsing/node/parse.go
@@ -177,8 +177,9 @@ func ParseNode(hop string, cfg *config.NodeConfig, log logger.Logger) (*chain.No
 
 	if cfg.HTTP != nil {
 		settings := &chain.HTTPNodeSettings{
-			Host:   cfg.HTTP.Host,
-			Header: cfg.HTTP.Header,
+			Host:           cfg.HTTP.Host,
+			Header:         cfg.HTTP.Header,
+			ResponseHeader: cfg.HTTP.ResponseHeader,
 		}
 
 		if auth := cfg.HTTP.Auth; auth != nil && auth.Username != "" {
diff --git a/go.mod b/go.mod
index 17a0e7f5..dd8b5965 100644
--- a/go.mod
+++ b/go.mod
@@ -9,7 +9,7 @@ require (
 	github.com/asaskevich/govalidator v0.0.0-20210307081110-f21760c49a8d
 	github.com/gin-contrib/cors v1.7.2
 	github.com/gin-gonic/gin v1.10.0
-	github.com/go-gost/core v0.1.8
+	github.com/go-gost/core v0.1.9
 	github.com/go-gost/gosocks4 v0.0.1
 	github.com/go-gost/gosocks5 v0.4.2
 	github.com/go-gost/plugin v0.1.1
diff --git a/go.sum b/go.sum
index 5636cb24..b050ff07 100644
--- a/go.sum
+++ b/go.sum
@@ -50,8 +50,8 @@ github.com/gin-contrib/sse v0.1.0 h1:Y/yl/+YNO8GZSjAhjMsSuLt29uWRFHdHYUb5lYOV9qE
 github.com/gin-contrib/sse v0.1.0/go.mod h1:RHrZQHXnP2xjPF+u1gW/2HnVO7nvIa9PG3Gm+fLHvGI=
 github.com/gin-gonic/gin v1.10.0 h1:nTuyha1TYqgedzytsKYqna+DfLos46nTv2ygFy86HFU=
 github.com/gin-gonic/gin v1.10.0/go.mod h1:4PMNQiOhvDRa013RKVbsiNwoyezlm2rm0uX/T7kzp5Y=
-github.com/go-gost/core v0.1.8 h1:OtE30bodK9v4lCCTCM89Wa3rYJD31pRgQQut1dZyFmk=
-github.com/go-gost/core v0.1.8/go.mod h1:WGI43jOka7FAsSAwi/fSMaqxdR+E339ycb4NBGlFr6A=
+github.com/go-gost/core v0.1.9 h1:qrKzhE9X2+zuNc8d4hRv4I8ZklwHAfja5HngrLucM2c=
+github.com/go-gost/core v0.1.9/go.mod h1:WGI43jOka7FAsSAwi/fSMaqxdR+E339ycb4NBGlFr6A=
 github.com/go-gost/gosocks4 v0.0.1 h1:+k1sec8HlELuQV7rWftIkmy8UijzUt2I6t+iMPlGB2s=
 github.com/go-gost/gosocks4 v0.0.1/go.mod h1:3B6L47HbU/qugDg4JnoFPHgJXE43Inz8Bah1QaN9qCc=
 github.com/go-gost/gosocks5 v0.4.2 h1:IianxHTkACPqCwiOAT3MHoMdSUl+SEPSRu1ikawC1Pc=
diff --git a/internal/util/forwarder/sniffer.go b/internal/util/forwarder/sniffer.go
index 50f98098..b1bc5fe3 100644
--- a/internal/util/forwarder/sniffer.go
+++ b/internal/util/forwarder/sniffer.go
@@ -404,6 +404,7 @@ func (h *Sniffer) httpRoundTrip(ctx context.Context, rw, cc io.ReadWriter, node
 		req.Header.Set("Connection", "close")
 	}
 
+	var responseHeader map[string]string
 	var bodyRewrites []chain.HTTPBodyRewriteSettings
 	if httpSettings := node.Options().HTTP; httpSettings != nil {
 		if auther := httpSettings.Auther; auther != nil {
@@ -437,6 +438,7 @@ func (h *Sniffer) httpRoundTrip(ctx context.Context, rw, cc io.ReadWriter, node
 			}
 		}
 
+		responseHeader = httpSettings.ResponseHeader
 		bodyRewrites = httpSettings.RewriteBody
 	}
 
@@ -477,6 +479,15 @@ func (h *Sniffer) httpRoundTrip(ctx context.Context, rw, cc io.ReadWriter, node
 	defer resp.Body.Close()
 	xio.SetReadDeadline(cc, time.Time{})
 
+	if len(responseHeader) > 0 {
+		if resp.Header == nil {
+			resp.Header = http.Header{}
+		}
+		for k, v := range responseHeader {
+			resp.Header.Set(k, v)
+		}
+	}
+
 	ro.HTTP.StatusCode = resp.StatusCode
 	ro.HTTP.Response.Header = resp.Header
 	ro.HTTP.Response.ContentLength = resp.ContentLength