From 33c988a952c5ce053f1bfb7f397028f21375f616 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E7=A8=8B=E5=B9=BF?= <chengg@dacentsoft.com>
Date: Thu, 14 Dec 2023 11:25:31 +0800
Subject: [PATCH] proxy add header process

---
 hander/proxy.go | 19 +++++++++++++++++--
 model/model.go  | 22 +++++++++++++++++-----
 2 files changed, 34 insertions(+), 7 deletions(-)

diff --git a/hander/proxy.go b/hander/proxy.go
index 7cfe163..fbb6db0 100644
--- a/hander/proxy.go
+++ b/hander/proxy.go
@@ -31,12 +31,13 @@ func (p *ProxyHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 }
 
 // init httputil.ReverseProxy instance and add path rewrite and add session-sticky cookie to response
-func makeProxy(upstream string, pw model.PathRewrite, index int) *httputil.ReverseProxy {
+func makeProxy(upstream string, path *model.HttpPath, index int) *httputil.ReverseProxy {
 	p := &httputil.ReverseProxy{}
 	p.Director = func(req *http.Request) {
 		turl, _ := url.Parse(upstream)
 		req.URL.Host = turl.Host
 		req.URL.Scheme = turl.Scheme
+		pw := path.Rewrite
 		if pw.Replace != "" {
 			req.URL.Path = strings.TrimPrefix(req.URL.Path, pw.Replace)
 			if req.URL.RawPath != "" {
@@ -49,6 +50,20 @@ func makeProxy(upstream string, pw model.PathRewrite, index int) *httputil.Rever
 				req.URL.RawPath = strings.TrimSuffix(pw.With, "/") + "/" + req.URL.RawPath
 			}
 		}
+
+		if len(path.Headers) > 0 {
+			for _, header := range path.Headers {
+				//req.Header.Add(header.Name, header.Value)
+				value := ""
+				switch header.Value {
+				case string(model.ProxyHost):
+					value = turl.Host
+				default:
+					value = header.Value
+				}
+				req.Header.Set(header.Name, value)
+			}
+		}
 	}
 
 	p.ModifyResponse = func(resp *http.Response) error {
@@ -80,7 +95,7 @@ func NewProxyHandler(p *model.HttpPath) *ProxyHandler {
 	ph.proxy = make([]*httputil.ReverseProxy, upstreamCount)
 
 	for index, upstream := range p.Upstreams {
-		ph.proxy[index] = makeProxy(upstream, p.Rewrite, index)
+		ph.proxy[index] = makeProxy(upstream, p, index)
 	}
 	return ph
 }
diff --git a/model/model.go b/model/model.go
index 396d306..eed9738 100644
--- a/model/model.go
+++ b/model/model.go
@@ -9,11 +9,12 @@ import (
 )
 
 type HttpPath struct {
-	Path      string      `json:"path"`
-	Root      string      `json:"root"`
-	Default   string      `json:"default"`
-	Upstreams []string    `json:"upstreams"`
-	Rewrite   PathRewrite `json:"pathrewrite"`
+	Path      string         `json:"path"`
+	Root      string         `json:"root"`
+	Default   string         `json:"default"`
+	Upstreams []string       `json:"upstreams"`
+	Rewrite   PathRewrite    `json:"pathrewrite"`
+	Headers   []HeaderDefine `json:"headers"`
 }
 
 type PathRewrite struct {
@@ -21,6 +22,17 @@ type PathRewrite struct {
 	With    string `json:"with"`
 }
 
+type HeaderDefine struct {
+	Name  string `json:"name"`
+	Value string `json:"value"`
+}
+
+type HeaderValueConst string
+
+const (
+	ProxyHost HeaderValueConst = "$ProxyHost"
+)
+
 type HttpServerConfig struct {
 	Name       string `json:"name"`
 	ServerName string `json:"server"`