264 lines
7.7 KiB
Go
264 lines
7.7 KiB
Go
// Package lua 提供 ngx.var API 实现。
|
||
//
|
||
// 该文件实现 nginx 变量访问的 Lua API，兼容 OpenResty/ngx_lua 语义。
|
||
// 支持：
|
||
// - 标准 nginx 变量：request_method、request_uri、uri、query_string 等
|
||
// - 请求头变量：http_host、http_user_agent、http_content_type 等
|
||
// - 客户端信息：remote_addr、server_addr 等
|
||
// - arg_ 前缀变量：arg_name 用于访问查询参数
|
||
// - http_ 前缀变量：http_xxx 用于访问任意请求头
|
||
// - 自定义变量存储：支持通过 store map 读写自定义变量
|
||
//
|
||
// 实现说明：
|
||
// - 通过元表（__index/__newindex）实现动态变量读写
|
||
// - 读取优先级：自定义变量 > fasthttp 内置变量
|
||
// - 写入始终存储到自定义变量 store 中
|
||
//
|
||
// 作者：xfy
|
||
package lua
|
||
|
||
import (
|
||
"strconv"
|
||
"strings"
|
||
|
||
"github.com/valyala/fasthttp"
|
||
glua "github.com/yuin/gopher-lua"
|
||
)
|
||
|
||
// argPrefix 是 arg_ 变量的前缀，用于获取查询参数
|
||
const argPrefix = "arg_"
|
||
|
||
// ngxVarAPI 封装 nginx 变量访问 API。
|
||
//
|
||
// 支持读取 fasthttp 请求中的标准 nginx 变量，
|
||
// 以及读写自定义变量（存储于 store map 中）。
|
||
type ngxVarAPI struct {
|
||
// ctx 关联的 fasthttp 请求上下文
|
||
ctx *fasthttp.RequestCtx
|
||
|
||
// store 自定义变量存储，支持 Lua 脚本读写自定义变量
|
||
store map[string]string
|
||
}
|
||
|
||
// newNgxVarAPI 创建 ngx.var API 实例。
|
||
//
|
||
// 参数：
|
||
// - ctx: fasthttp 请求上下文
|
||
//
|
||
// 返回值：
|
||
// - *ngxVarAPI: 初始化的 API 实例
|
||
func newNgxVarAPI(ctx *fasthttp.RequestCtx) *ngxVarAPI {
|
||
return &ngxVarAPI{
|
||
ctx: ctx,
|
||
store: make(map[string]string),
|
||
}
|
||
}
|
||
|
||
// RegisterNgxVarAPI 在 Lua 状态机中注册 ngx.var API
|
||
// 使用元表实现动态读写：ngx.var.key 和 ngx.var[key]
|
||
func RegisterNgxVarAPI(L *glua.LState, api *ngxVarAPI, ngxTable *glua.LTable) {
|
||
// 检查 ngx 表是否已存在 var 字段，避免重复写入
|
||
if existing := ngxTable.RawGetString("var"); existing != glua.LNil {
|
||
return
|
||
}
|
||
|
||
// 创建 ngx.var 表（使用元表实现动态访问）
|
||
ngxVar := L.NewTable()
|
||
|
||
// 创建元表
|
||
mt := L.NewTable()
|
||
|
||
// __index 元方法：读取变量
|
||
mt.RawSetString("__index", L.NewFunction(api.luaVarIndex))
|
||
|
||
// __newindex 元方法：设置变量
|
||
mt.RawSetString("__newindex", L.NewFunction(api.luaVarNewIndex))
|
||
|
||
// 设置元表
|
||
L.SetMetatable(ngxVar, mt)
|
||
|
||
// 将 ngx.var 添加到 ngx 表
|
||
ngxTable.RawSetString("var", ngxVar)
|
||
}
|
||
|
||
// luaVarIndex 实现 ngx.var[key] 读取
|
||
// Lua 调用: local value = ngx.var.key 或 ngx.var[key]
|
||
func (api *ngxVarAPI) luaVarIndex(L *glua.LState) int {
|
||
// 第一个参数是表本身（ngx.var）
|
||
// 第二个参数是键名
|
||
key := L.CheckString(2)
|
||
|
||
// 1. 先查自定义变量存储
|
||
if value, ok := api.store[key]; ok {
|
||
L.Push(glua.LString(value))
|
||
return 1
|
||
}
|
||
|
||
// 2. 从 fasthttp RequestCtx 获取变量
|
||
// 某些变量需要返回数值类型（如 request_length）
|
||
lv := api.getVariableLua(key)
|
||
if lv != nil {
|
||
L.Push(lv)
|
||
return 1
|
||
}
|
||
|
||
// 3. 未找到变量，返回 nil
|
||
L.Push(glua.LNil)
|
||
return 1
|
||
}
|
||
|
||
// luaVarNewIndex 实现 ngx.var[key] = value 写入
|
||
// Lua 调用: ngx.var.key = value 或 ngx.var[key] = value
|
||
// 注意：Lua 的 nil 会被转换为空字符串存储
|
||
func (api *ngxVarAPI) luaVarNewIndex(L *glua.LState) int {
|
||
// 第一个参数是表本身（ngx.var）
|
||
// 第二个参数是键名
|
||
// 第三个参数是值
|
||
key := L.CheckString(2)
|
||
value := L.OptString(3, "")
|
||
|
||
// 存储到自定义变量存储（nil 会转换为空字符串）
|
||
api.store[key] = value
|
||
|
||
return 0
|
||
}
|
||
|
||
// getVariableLua 从 fasthttp RequestCtx 获取变量值，返回 Lua 类型
|
||
// 支持常见的 nginx 变量，某些变量返回数值类型
|
||
func (api *ngxVarAPI) getVariableLua(name string) glua.LValue {
|
||
if api.ctx == nil {
|
||
return nil
|
||
}
|
||
|
||
switch name {
|
||
// HTTP 请求相关 - 数值类型
|
||
case "request_length":
|
||
return glua.LNumber(api.ctx.Request.Header.ContentLength())
|
||
|
||
// HTTP 请求相关 - 字符串类型
|
||
case "request_method":
|
||
return glua.LString(string(api.ctx.Method()))
|
||
case "request_uri":
|
||
return glua.LString(string(api.ctx.RequestURI()))
|
||
case "uri":
|
||
return glua.LString(string(api.ctx.URI().Path()))
|
||
case "document_uri":
|
||
return glua.LString(string(api.ctx.URI().Path()))
|
||
case "query_string", "args":
|
||
return glua.LString(string(api.ctx.URI().QueryString()))
|
||
case "server_protocol", "protocol":
|
||
return glua.LString(string(api.ctx.Request.Header.Protocol()))
|
||
case "scheme":
|
||
return glua.LString(string(api.ctx.URI().Scheme()))
|
||
case "request_time":
|
||
// 简化实现，返回空字符串
|
||
return glua.LString("")
|
||
|
||
// 请求头相关
|
||
case "http_host":
|
||
return glua.LString(string(api.ctx.Host()))
|
||
case "http_user_agent", "http_user-agent":
|
||
return glua.LString(string(api.ctx.UserAgent()))
|
||
case "http_referer":
|
||
return glua.LString(string(api.ctx.Referer()))
|
||
case "http_accept":
|
||
return glua.LString(string(api.ctx.Request.Header.Peek("Accept")))
|
||
case "http_accept_encoding", "http_accept-encoding":
|
||
return glua.LString(string(api.ctx.Request.Header.Peek("Accept-Encoding")))
|
||
case "http_accept_language", "http_accept-language":
|
||
return glua.LString(string(api.ctx.Request.Header.Peek("Accept-Language")))
|
||
case "http_connection":
|
||
return glua.LString(string(api.ctx.Request.Header.Peek("Connection")))
|
||
case "http_content_type", "http_content-type":
|
||
return glua.LString(string(api.ctx.Request.Header.ContentType()))
|
||
case "http_content_length", "http_content-length":
|
||
return glua.LString(string(api.ctx.Request.Header.Peek("Content-Length")))
|
||
|
||
// 客户端信息
|
||
case "remote_addr":
|
||
return glua.LString(api.ctx.RemoteAddr().String())
|
||
case "remote_port":
|
||
addr := api.ctx.RemoteAddr()
|
||
if addr != nil {
|
||
addrStr := addr.String()
|
||
// IPv6 地址格式: [ip]:port，IPv4 格式: ip:port
|
||
// 统一使用 LastIndex(":") 可正确处理两种格式
|
||
if idx := strings.LastIndex(addrStr, ":"); idx >= 0 {
|
||
return glua.LString(addrStr[idx+1:])
|
||
}
|
||
}
|
||
return glua.LString("")
|
||
case "binary_remote_addr":
|
||
return glua.LString("")
|
||
|
||
// 服务器信息
|
||
case "server_addr":
|
||
addr := api.ctx.LocalAddr()
|
||
if addr != nil {
|
||
return glua.LString(addr.String())
|
||
}
|
||
return glua.LString("")
|
||
case "server_port":
|
||
addr := api.ctx.LocalAddr()
|
||
if addr != nil {
|
||
addrStr := addr.String()
|
||
if idx := strings.LastIndex(addrStr, ":"); idx >= 0 {
|
||
return glua.LString(addrStr[idx+1:])
|
||
}
|
||
}
|
||
return glua.LString("")
|
||
case "server_name":
|
||
return glua.LString(string(api.ctx.Host()))
|
||
|
||
// URI 参数
|
||
case argPrefix:
|
||
// 获取所有参数
|
||
return glua.LString(string(api.ctx.URI().QueryString()))
|
||
default:
|
||
// 检查是否是 arg_ 开头的参数
|
||
if len(name) > len(argPrefix) && name[:len(argPrefix)] == argPrefix {
|
||
paramName := name[4:]
|
||
return glua.LString(string(api.ctx.QueryArgs().Peek(paramName)))
|
||
}
|
||
// 检查是否是 http_ 开头的请求头
|
||
if len(name) > 5 && name[:5] == "http_" {
|
||
headerName := name[5:]
|
||
return glua.LString(string(api.ctx.Request.Header.Peek(headerName)))
|
||
}
|
||
return nil
|
||
}
|
||
}
|
||
|
||
// getVariable 从 fasthttp RequestCtx 获取变量值（字符串形式）
|
||
// 用于 Go 层调用，内部复用 getVariableLua 实现
|
||
func (api *ngxVarAPI) getVariable(name string) string {
|
||
lv := api.getVariableLua(name)
|
||
if lv == nil {
|
||
return ""
|
||
}
|
||
switch v := lv.(type) {
|
||
case glua.LString:
|
||
return string(v)
|
||
case glua.LNumber:
|
||
return strconv.Itoa(int(v))
|
||
default:
|
||
return lv.String()
|
||
}
|
||
}
|
||
|
||
// SetVariable 设置自定义变量（Go 层调用）
|
||
func (api *ngxVarAPI) SetVariable(name, value string) {
|
||
api.store[name] = value
|
||
}
|
||
|
||
// GetVariable 获取变量值（Go 层调用）
|
||
func (api *ngxVarAPI) GetVariable(name string) (string, bool) {
|
||
// 先查自定义变量
|
||
if value, ok := api.store[name]; ok {
|
||
return value, true
|
||
}
|
||
// 再查 fasthttp 变量
|
||
value := api.getVariable(name)
|
||
return value, value != ""
|
||
}
|