xfy
c95f474539
新增 internal/variable 包，提供 Nginx 风格的变量展开功能： - 支持 $remote_addr、$host、$uri、$args 等 30+ 内置变量 - 使用 sync.Pool 优化 VariableContext 分配 - 支持 set_response_info 存储响应状态信息 添加 github.com/google/uuid 依赖用于请求 ID 生成。 Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
386 lines
8.4 KiB
Go
386 lines
8.4 KiB
Go
// Package variable 提供高性能的变量系统，支持 nginx 风格的变量展开。
|
||
//
|
||
// 该包实现了统一的变量存储和展开机制，用于：
|
||
// - 访问日志格式模板
|
||
// - 代理请求头设置
|
||
// - URL 重写规则
|
||
//
|
||
// 支持的变量格式：
|
||
// - $var: 简单变量
|
||
// - ${var}: 带花括号的变量（用于变量后有字符的场景）
|
||
//
|
||
// 性能特性：
|
||
// - 使用快速字符串扫描（非正则表达式）
|
||
// - sync.Pool 复用 VariableContext
|
||
// - 内置变量惰性求值并缓存
|
||
//
|
||
// 作者：xfy
|
||
package variable
|
||
|
||
import (
|
||
"strconv"
|
||
"strings"
|
||
"sync"
|
||
|
||
"github.com/valyala/fasthttp"
|
||
)
|
||
|
||
// VariableStore 变量存储接口
|
||
type VariableStore interface {
|
||
// Get 获取变量值
|
||
Get(name string) (string, bool)
|
||
// Set 设置变量值（用于自定义变量）
|
||
Set(name string, value string)
|
||
}
|
||
|
||
// BuiltinVariable 内置变量定义
|
||
type BuiltinVariable struct {
|
||
Name string
|
||
Description string
|
||
Getter func(ctx *fasthttp.RequestCtx) string
|
||
}
|
||
|
||
// VariableContext 变量上下文，绑定到请求
|
||
type VariableContext struct {
|
||
ctx *fasthttp.RequestCtx
|
||
store map[string]string // 自定义变量存储
|
||
cache map[string]string // 内置变量缓存
|
||
status int // HTTP 状态码（由外部设置）
|
||
bodySize int64 // 响应体大小（由外部设置）
|
||
duration int64 // 请求处理时间纳秒（由外部设置）
|
||
serverName string // 服务器名称
|
||
}
|
||
|
||
// pool 用于复用 VariableContext
|
||
var pool = sync.Pool{
|
||
New: func() interface{} {
|
||
return &VariableContext{
|
||
store: make(map[string]string),
|
||
cache: make(map[string]string),
|
||
}
|
||
},
|
||
}
|
||
|
||
// builtinVars 内置变量注册表
|
||
var builtinVars = make(map[string]*BuiltinVariable)
|
||
|
||
// RegisterBuiltin 注册内置变量
|
||
func RegisterBuiltin(v *BuiltinVariable) {
|
||
builtinVars[v.Name] = v
|
||
}
|
||
|
||
// GetBuiltin 获取内置变量定义
|
||
func GetBuiltin(name string) *BuiltinVariable {
|
||
return builtinVars[name]
|
||
}
|
||
|
||
// NewVariableContext 从池中获取 VariableContext
|
||
func NewVariableContext(ctx *fasthttp.RequestCtx) *VariableContext {
|
||
vc := pool.Get().(*VariableContext)
|
||
vc.ctx = ctx
|
||
vc.status = 0
|
||
vc.bodySize = 0
|
||
vc.duration = 0
|
||
vc.serverName = ""
|
||
// 清空缓存
|
||
for k := range vc.cache {
|
||
delete(vc.cache, k)
|
||
}
|
||
// 清空自定义变量
|
||
for k := range vc.store {
|
||
delete(vc.store, k)
|
||
}
|
||
return vc
|
||
}
|
||
|
||
// ReleaseVariableContext 释放 VariableContext 回池中
|
||
func ReleaseVariableContext(vc *VariableContext) {
|
||
if vc == nil {
|
||
return
|
||
}
|
||
vc.ctx = nil
|
||
vc.status = 0
|
||
vc.bodySize = 0
|
||
vc.duration = 0
|
||
vc.serverName = ""
|
||
pool.Put(vc)
|
||
}
|
||
|
||
// SetResponseInfo 设置响应信息（用于需要 status、body_bytes_sent、request_time 的场景）
|
||
func (vc *VariableContext) SetResponseInfo(status int, bodySize int64, durationNs int64) {
|
||
vc.status = status
|
||
vc.bodySize = bodySize
|
||
vc.duration = durationNs
|
||
}
|
||
|
||
// SetServerName 设置服务器名称
|
||
func (vc *VariableContext) SetServerName(name string) {
|
||
vc.serverName = name
|
||
}
|
||
|
||
// Get 获取变量值（优先自定义变量，再查内置变量）
|
||
func (vc *VariableContext) Get(name string) (string, bool) {
|
||
// 1. 先查自定义变量
|
||
if v, ok := vc.store[name]; ok {
|
||
return v, true
|
||
}
|
||
|
||
// 2. 检查从 SetResponseInfo/SetServerName 设置的值
|
||
// 优先检查 struct 字段，再检查 ctx.UserValue（兼容 SetResponseInfoInContext）
|
||
switch name {
|
||
case VarStatus:
|
||
if vc.status > 0 {
|
||
return strconv.Itoa(vc.status), true
|
||
}
|
||
if v := vc.ctx.UserValue(VarStatus); v != nil {
|
||
if i, ok := v.(int); ok {
|
||
return strconv.Itoa(i), true
|
||
}
|
||
}
|
||
case VarBodyBytesSent:
|
||
if vc.bodySize > 0 {
|
||
return strconv.FormatInt(vc.bodySize, 10), true
|
||
}
|
||
if v := vc.ctx.UserValue(VarBodyBytesSent); v != nil {
|
||
if i, ok := v.(int64); ok {
|
||
return strconv.FormatInt(i, 10), true
|
||
}
|
||
}
|
||
return "0", true
|
||
case VarRequestTime:
|
||
if vc.duration > 0 {
|
||
// 转换为秒，保留 3 位小数
|
||
seconds := float64(vc.duration) / 1e9
|
||
return strconv.FormatFloat(seconds, 'f', 3, 64), true
|
||
}
|
||
if v := vc.ctx.UserValue(VarRequestTime); v != nil {
|
||
if i, ok := v.(int64); ok {
|
||
seconds := float64(i) / 1e9
|
||
return strconv.FormatFloat(seconds, 'f', 3, 64), true
|
||
}
|
||
}
|
||
return "0.000", true
|
||
case VarServerName:
|
||
if vc.serverName != "" {
|
||
return vc.serverName, true
|
||
}
|
||
}
|
||
|
||
// 3. 查内置变量缓存
|
||
if v, ok := vc.cache[name]; ok {
|
||
return v, true
|
||
}
|
||
|
||
// 4. 求值内置变量并缓存
|
||
if v, ok := vc.evalBuiltin(name); ok {
|
||
vc.cache[name] = v
|
||
return v, true
|
||
}
|
||
|
||
return "", false
|
||
}
|
||
|
||
// Set 设置自定义变量
|
||
func (vc *VariableContext) Set(name string, value string) {
|
||
vc.store[name] = value
|
||
}
|
||
|
||
// evalBuiltin 求值内置变量
|
||
func (vc *VariableContext) evalBuiltin(name string) (string, bool) {
|
||
builtin := builtinVars[name]
|
||
if builtin == nil || builtin.Getter == nil {
|
||
return "", false
|
||
}
|
||
return builtin.Getter(vc.ctx), true
|
||
}
|
||
|
||
// Expand 展开模板字符串中的变量
|
||
// 支持 $var 和 ${var} 两种格式
|
||
// 对于未定义的变量，保持原样不变
|
||
func (vc *VariableContext) Expand(template string) string {
|
||
if template == "" {
|
||
return ""
|
||
}
|
||
|
||
// 快速路径：没有变量
|
||
hasVar := false
|
||
for i := 0; i < len(template); i++ {
|
||
if template[i] == '$' {
|
||
hasVar = true
|
||
break
|
||
}
|
||
}
|
||
if !hasVar {
|
||
return template
|
||
}
|
||
|
||
var result strings.Builder
|
||
result.Grow(len(template) * 2)
|
||
|
||
i := 0
|
||
for i < len(template) {
|
||
if template[i] != '$' {
|
||
result.WriteByte(template[i])
|
||
i++
|
||
continue
|
||
}
|
||
|
||
// 到达末尾，保留 $
|
||
if i+1 >= len(template) {
|
||
result.WriteByte('$')
|
||
i++
|
||
continue
|
||
}
|
||
|
||
// ${var} 格式
|
||
if template[i+1] == '{' {
|
||
// 查找匹配的 }
|
||
end := strings.IndexByte(template[i+2:], '}')
|
||
if end == -1 {
|
||
result.WriteByte('$')
|
||
i++
|
||
continue
|
||
}
|
||
// end 是相对 i+2 的偏移量
|
||
varName := template[i+2 : i+2+end]
|
||
if varName == "" {
|
||
// 空变量名，保持 ${}
|
||
result.WriteString("${}")
|
||
i += 2 + end + 1
|
||
continue
|
||
}
|
||
// 获取变量值
|
||
if v, ok := vc.Get(varName); ok {
|
||
result.WriteString(v)
|
||
} else {
|
||
// 未定义变量，保持原样
|
||
result.WriteString("${")
|
||
result.WriteString(varName)
|
||
result.WriteByte('}')
|
||
}
|
||
// i+2 是变量名开始，+end 是 } 的位置，+1 跳过 }
|
||
i += 2 + end + 1
|
||
continue
|
||
}
|
||
|
||
// $var 格式（变量名由字母、数字、下划线组成）
|
||
j := i + 1
|
||
for j < len(template) {
|
||
c := template[j]
|
||
if (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') || (c >= '0' && c <= '9') || c == '_' {
|
||
j++
|
||
} else {
|
||
break
|
||
}
|
||
}
|
||
|
||
if j == i+1 {
|
||
// 变量名长度为0，保留 $
|
||
result.WriteByte('$')
|
||
i++
|
||
continue
|
||
}
|
||
|
||
varName := template[i+1 : j]
|
||
if v, ok := vc.Get(varName); ok {
|
||
result.WriteString(v)
|
||
} else {
|
||
// 未定义变量，保持原样
|
||
result.WriteByte('$')
|
||
result.WriteString(varName)
|
||
}
|
||
i = j // 跳过变量名
|
||
}
|
||
|
||
return result.String()
|
||
}
|
||
|
||
// ExpandString 展开字符串（静态函数，用于简单场景）
|
||
// 需要提供变量值查找函数
|
||
func ExpandString(template string, lookup func(string) string) string {
|
||
if template == "" {
|
||
return ""
|
||
}
|
||
|
||
// 快速路径
|
||
hasVar := false
|
||
for i := 0; i < len(template); i++ {
|
||
if template[i] == '$' {
|
||
hasVar = true
|
||
break
|
||
}
|
||
}
|
||
if !hasVar {
|
||
return template
|
||
}
|
||
|
||
var result strings.Builder
|
||
result.Grow(len(template) * 2)
|
||
|
||
i := 0
|
||
for i < len(template) {
|
||
if template[i] != '$' {
|
||
result.WriteByte(template[i])
|
||
i++
|
||
continue
|
||
}
|
||
|
||
if i+1 >= len(template) {
|
||
result.WriteByte('$')
|
||
i++
|
||
continue
|
||
}
|
||
|
||
if template[i+1] == '{' {
|
||
end := strings.IndexByte(template[i+2:], '}')
|
||
if end == -1 {
|
||
result.WriteByte('$')
|
||
i++
|
||
continue
|
||
}
|
||
varName := template[i+2 : i+2+end]
|
||
if varName == "" {
|
||
result.WriteByte('$')
|
||
i += 2
|
||
continue
|
||
}
|
||
if v := lookup(varName); v != "" {
|
||
result.WriteString(v)
|
||
} else {
|
||
result.WriteString("${")
|
||
result.WriteString(varName)
|
||
result.WriteByte('}')
|
||
}
|
||
i += 2 + end + 1
|
||
continue
|
||
}
|
||
|
||
j := i + 1
|
||
for j < len(template) {
|
||
c := template[j]
|
||
if (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') || (c >= '0' && c <= '9') || c == '_' {
|
||
j++
|
||
} else {
|
||
break
|
||
}
|
||
}
|
||
|
||
if j == i+1 {
|
||
result.WriteByte('$')
|
||
i++
|
||
continue
|
||
}
|
||
|
||
varName := template[i+1 : j]
|
||
if v := lookup(varName); v != "" {
|
||
result.WriteString(v)
|
||
} else {
|
||
result.WriteByte('$')
|
||
result.WriteString(varName)
|
||
}
|
||
i = j
|
||
}
|
||
|
||
return result.String()
|
||
}
|