From 23bdedc7426377564f1ae8b6b26b910f1958be31 Mon Sep 17 00:00:00 2001 From: samwaf Date: Tue, 6 Jan 2026 12:47:17 +0800 Subject: [PATCH] feat:custom header #592 --- model/hosts.go | 32 +++++++++++++ model/request/waf_host_req.go | 2 + service/waf_service/waf_host.go | 2 + wafdb/migrations_core.go | 34 ++++++++++++++ wafdb/patch_sql.go | 3 +- wafenginecore/proxy.go | 81 +++++++++++++++++++++++++++++++++ 6 files changed, 153 insertions(+), 1 deletion(-) diff --git a/model/hosts.go b/model/hosts.go index a75f51b..9aaf1f8 100644 --- a/model/hosts.go +++ b/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 +} diff --git a/model/request/waf_host_req.go b/model/request/waf_host_req.go index 09df199..efec6fa 100644 --- a/model/request/waf_host_req.go +++ b/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"` diff --git a/service/waf_service/waf_host.go b/service/waf_service/waf_host.go index 821a356..72c9e5b 100644 --- a/service/waf_service/waf_host.go +++ b/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 diff --git a/wafdb/migrations_core.go b/wafdb/migrations_core.go index 18853f4..7792e23 100644 --- a/wafdb/migrations_core.go +++ b/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 + }, + }, }) // 执行迁移 diff --git a/wafdb/patch_sql.go b/wafdb/patch_sql.go index 4946fc1..9fc54b8 100644 --- a/wafdb/patch_sql.go +++ b/wafdb/patch_sql.go @@ -7,8 +7,9 @@ import ( "SamWaf/global" "SamWaf/model" "SamWaf/model/baseorm" - "gorm.io/gorm" "time" + + "gorm.io/gorm" ) /* diff --git a/wafenginecore/proxy.go b/wafenginecore/proxy.go index 70d571d..da59420 100644 --- a/wafenginecore/proxy.go +++ b/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 +}