Merge pull request #625 from samwafgo/feat_custom_header

feat:custom header

model/hosts.go

@@ -45,6 +45,7 @@ type Hosts struct {
 	TransportJSON        string `json:"transport_json"`           //传输配置 json
 	DefaultEncoding      string `json:"default_encoding"`         //默认编码 utf-8 或者 gbk  auto字符串自动选择
 	LogOnlyMode          int    `json:"log_only_mode"`            //仅记录模式 1开启 0关闭
+	CustomHeadersJSON    string `json:"custom_headers_json"`      //自定义头信息配置 json
 }
 
 type HostsDefense struct {
@@ -170,6 +171,18 @@ type TransportConfig struct {
 	ExpectContinueTimeout int `json:"expect_continue_timeout"` // Expect Continue超时时间(秒)
 }
 
+// CustomHeaderItem 自定义头信息项
+type CustomHeaderItem struct {
+	HeaderName  string `json:"header_name"`  // 头信息名称
+	HeaderValue string `json:"header_value"` // 头信息值，支持内置变量
+}
+
+// CustomHeadersConfig 自定义头信息配置
+type CustomHeadersConfig struct {
+	IsEnableCustomHeaders int                `json:"is_enable_custom_headers"` // 是否开启自定义头信息 1开启 0关闭
+	Headers               []CustomHeaderItem `json:"headers"`                  // 自定义头信息列表
+}
+
 // ParseTransportConfig 解析传输配置
 func ParseTransportConfig(transportJSON string) TransportConfig {
 	var config TransportConfig
@@ -217,3 +230,22 @@ func ParseHostsDefense(defenseJSON string) HostsDefense {
 	}
 	return defense
 }
+
+// ParseCustomHeadersConfig 解析自定义头信息配置
+func ParseCustomHeadersConfig(customHeadersJSON string) CustomHeadersConfig {
+	var config CustomHeadersConfig
+
+	// 设置默认值
+	config.IsEnableCustomHeaders = 0
+	config.Headers = []CustomHeaderItem{}
+
+	// 如果JSON不为空，则解析覆盖默认值
+	if customHeadersJSON != "" {
+		err := json.Unmarshal([]byte(customHeadersJSON), &config)
+		if err != nil {
+			// 解析失败时使用默认值，可以记录日志
+			return config
+		}
+	}
+	return config
+}

model/request/waf_host_req.go

@@ -39,6 +39,7 @@ type WafHostAddReq struct {
 	DefaultEncoding      string `json:"default_encoding"`         //默认编码 utf-8 或者 gbk auto字符串自动选择
 	LogOnlyMode          int    `json:"log_only_mode"`            //是否只记录日志 1 是 0 不是
 	TransportJSON        string `json:"transport_json"`           //Transport配置 json
+	CustomHeadersJSON    string `json:"custom_headers_json"`      //自定义头信息配置 json
 }
 
 type WafHostDelReq struct {
@@ -86,6 +87,7 @@ type WafHostEditReq struct {
 	DefaultEncoding      string `json:"default_encoding"`         //默认编码 utf-8 或者 gbk auto字符串自动选择
 	LogOnlyMode          int    `json:"log_only_mode"`            //是否只记录日志 1 是 0 不是
 	TransportJSON        string `json:"transport_json"`           //Transport配置 json
+	CustomHeadersJSON    string `json:"custom_headers_json"`      //自定义头信息配置 json
 }
 type WafHostGuardStatusReq struct {
 	CODE         string `json:"code"`

service/waf_service/waf_host.go

@@ -96,6 +96,7 @@ func (receiver *WafHostService) AddApi(wafHostAddReq request.WafHostAddReq) (str
 		DefaultEncoding:      wafHostAddReq.DefaultEncoding,
 		LogOnlyMode:          wafHostAddReq.LogOnlyMode,
 		TransportJSON:        wafHostAddReq.TransportJSON,
+		CustomHeadersJSON:    wafHostAddReq.CustomHeadersJSON,
 	}
 	global.GWAF_LOCAL_DB.Create(wafHost)
 	return wafHost.Code, nil
@@ -156,6 +157,7 @@ func (receiver *WafHostService) ModifyApi(wafHostEditReq request.WafHostEditReq)
 		"DefaultEncoding":      wafHostEditReq.DefaultEncoding,
 		"LogOnlyMode":          wafHostEditReq.LogOnlyMode,
 		"TransportJSON":        wafHostEditReq.TransportJSON,
+		"CustomHeadersJSON":    wafHostEditReq.CustomHeadersJSON,
 	}
 	err := global.GWAF_LOCAL_DB.Debug().Model(model.Hosts{}).Where("CODE=?", wafHostEditReq.CODE).Updates(hostMap).Error
 

wafdb/migrations_core.go

@@ -307,6 +307,40 @@ func RunCoreDBMigrations(db *gorm.DB) error {
 				return nil
 			},
 		},
+		// 迁移8: 为 hosts 表添加 custom_headers_json 字段（自定义头信息功能）
+		{
+			ID: "202601060001_add_hosts_custom_headers_json",
+			Migrate: func(tx *gorm.DB) error {
+				zlog.Info("迁移 202601060001: 为 hosts 表添加 custom_headers_json 字段")
+
+				// 检查字段是否已存在
+				if tx.Migrator().HasColumn(&model.Hosts{}, "custom_headers_json") {
+					zlog.Info("custom_headers_json 字段已存在，跳过添加")
+					return nil
+				}
+
+				// 添加字段
+				if err := tx.Migrator().AddColumn(&model.Hosts{}, "custom_headers_json"); err != nil {
+					return fmt.Errorf("添加 custom_headers_json 字段失败: %w", err)
+				}
+
+				// 设置默认值
+				defaultCustomHeadersConfig := `{"is_enable_custom_headers":0,"headers":[]}`
+				if err := tx.Exec("UPDATE hosts SET custom_headers_json = ? WHERE custom_headers_json IS NULL OR custom_headers_json = ''", defaultCustomHeadersConfig).Error; err != nil {
+					zlog.Warn("设置 custom_headers_json 默认值失败", "error", err.Error())
+				}
+
+				zlog.Info("custom_headers_json 字段添加成功")
+				return nil
+			},
+			Rollback: func(tx *gorm.DB) error {
+				zlog.Info("回滚 202601060001: 删除 hosts 表的 custom_headers_json 字段")
+				if tx.Migrator().HasColumn(&model.Hosts{}, "custom_headers_json") {
+					return tx.Migrator().DropColumn(&model.Hosts{}, "custom_headers_json")
+				}
+				return nil
+			},
+		},
 	})
 
 	// 执行迁移

wafdb/patch_sql.go

@@ -7,8 +7,9 @@ import (
 	"SamWaf/global"
 	"SamWaf/model"
 	"SamWaf/model/baseorm"
-	"gorm.io/gorm"
 	"time"
+
+	"gorm.io/gorm"
 )
 
 /*

wafenginecore/proxy.go

@@ -13,7 +13,9 @@ import (
 	"net"
 	"net/http"
 	"net/url"
+	"regexp"
 	"strconv"
+	"strings"
 	"time"
 )
 
@@ -246,5 +248,84 @@ func (waf *WafEngine) getCustomHeaders(r *http.Request, host string, isEnableLoa
 	if r.TLS != nil {
 		customHeaders["X-FORWARDED-PROTO"] = "https"
 	}
+
+	// 解析自定义头信息配置
+	customHeadersConfig := model.ParseCustomHeadersConfig(hostTarget.Host.CustomHeadersJSON)
+
+	// 如果启用了自定义头信息
+	if customHeadersConfig.IsEnableCustomHeaders == 1 {
+		// 获取客户端IP
+		clientIP := r.RemoteAddr
+		if idx := strings.LastIndex(clientIP, ":"); idx != -1 {
+			clientIP = clientIP[:idx]
+		}
+
+		// 遍历自定义头信息列表
+		for _, header := range customHeadersConfig.Headers {
+			if header.HeaderName == "" {
+				continue
+			}
+
+			// 解析并替换变量
+			headerValue := waf.parseHeaderValue(header.HeaderValue, r, clientIP)
+
+			// 添加到自定义头信息中
+			customHeaders[header.HeaderName] = headerValue
+		}
+	}
+
 	return customHeaders
 }
+
+// parseHeaderValue 解析头信息值中的变量
+// 支持的变量：
+// ${header:HeaderName} - 获取原始请求头
+// ${client_ip} - 获取客户端IP
+// ${remote_addr} - 获取远程地址
+// ${host} - 获取主机名
+// ${scheme} - 获取协议（http/https）
+// ${request_uri} - 获取请求URI
+// ${request_method} - 获取请求方法
+func (waf *WafEngine) parseHeaderValue(value string, r *http.Request, clientIP string) string {
+	// 正则表达式匹配 ${变量名} 或 ${header:头名称}
+	varPattern := regexp.MustCompile(`\$\{([^}]+)\}`)
+
+	result := varPattern.ReplaceAllStringFunc(value, func(match string) string {
+		// 提取变量名（去掉 ${ 和 }）
+		varName := match[2 : len(match)-1]
+
+		// 检查是否是header变量
+		if strings.HasPrefix(varName, "header:") {
+			headerName := strings.TrimPrefix(varName, "header:")
+			headerValue := r.Header.Get(headerName)
+			if headerValue == "" {
+				return "" // 如果header不存在，返回空字符串
+			}
+			return headerValue
+		}
+
+		// 处理其他内置变量
+		switch varName {
+		case "client_ip":
+			return clientIP
+		case "remote_addr":
+			return r.RemoteAddr
+		case "host":
+			return r.Host
+		case "scheme":
+			if r.TLS != nil {
+				return "https"
+			}
+			return "http"
+		case "request_uri":
+			return r.RequestURI
+		case "request_method":
+			return r.Method
+		default:
+			// 未知变量，保持原样
+			return match
+		}
+	})
+
+	return result
+}