From 7d53cc3dea596f973ed150f5a224a6e3589ba9af Mon Sep 17 00:00:00 2001 From: xfy Date: Sun, 12 Apr 2026 13:03:42 +0800 Subject: [PATCH] =?UTF-8?q?feat(lua):=20=E6=94=B9=E8=BF=9B=E5=AD=90?= =?UTF-8?q?=E8=AF=B7=E6=B1=82=E7=88=B6=E8=AF=B7=E6=B1=82=E6=95=B0=E6=8D=AE?= =?UTF-8?q?=E7=BB=A7=E6=89=BF?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ngx.location.capture 改进： - 子请求继承父请求的 headers 和 query args - 支持通过 opts.args 覆盖查询参数 - 通过 LState Context 传递 RequestCtx - 添加测试验证父请求数据继承 Co-Authored-By: Claude Opus 4.6 --- internal/lua/api_location.go | 90 +++++++++++++++++++++++-------- internal/lua/api_location_test.go | 30 +++++++++++ 2 files changed, 97 insertions(+), 23 deletions(-) diff --git a/internal/lua/api_location.go b/internal/lua/api_location.go index 12de255..f8fae2b 100644 --- a/internal/lua/api_location.go +++ b/internal/lua/api_location.go @@ -54,11 +54,13 @@ func (m *LocationManager) Capture(parentCtx *fasthttp.RequestCtx, location strin // 创建子请求上下文（不设置 Conn） subCtx := &fasthttp.RequestCtx{} - // 复制父请求作为基础 + // 复制父请求作为基础（深拷贝） parentCtx.Request.CopyTo(&subCtx.Request) - // 设置子请求的 URI - subCtx.Request.SetRequestURI(location) + // 设置子请求的路径，保留父请求的查询参数 + // 解析 location，分离路径和查询参数 + uri := subCtx.URI() + uri.SetPath(location) // 应用选项 if opts != nil { @@ -73,6 +75,13 @@ func (m *LocationManager) Capture(parentCtx *fasthttp.RequestCtx, location strin subCtx.Request.Header.Set(k, v) } } + // 如果选项中显式指定了 args，则覆盖父请求的查询参数 + if args, ok := opts["args"].(map[string]string); ok { + uri.QueryArgs().Reset() + for k, v := range args { + uri.QueryArgs().Add(k, v) + } + } } // 执行 handler @@ -93,6 +102,18 @@ func (m *LocationManager) Capture(parentCtx *fasthttp.RequestCtx, location strin return result, nil } +// getRequestCtx 从当前 Lua 协程的 UserData 中获取 RequestCtx +// 通过协程关联的 RequestCtx 实现子请求对父请求数据的访问 +func getRequestCtx(L *glua.LState) *fasthttp.RequestCtx { + // 获取当前协程的上下文（在创建协程时通过 SetContext 设置） + if ctx := L.Context(); ctx != nil { + if reqCtx, ok := ctx.(*fasthttp.RequestCtx); ok { + return reqCtx + } + } + return nil +} + // RegisterLocationAPI 注册 ngx.location API func RegisterLocationAPI(L *glua.LState, manager *LocationManager, ngx *glua.LTable) { // 创建 ngx.location 表 @@ -129,31 +150,36 @@ func RegisterLocationAPI(L *glua.LState, manager *LocationManager, ngx *glua.LTa // 创建结果表 result := L.NewTable() - // 尝试执行子请求 - // 注意：由于无法直接获取 RequestCtx，这里使用模拟的上下文 - // 在完整实现中，应该通过 coroutine 传递 RequestCtx - if manager != nil { - // 创建模拟请求上下文用于子请求执行 + if manager == nil { + // manager 未初始化 + L.SetField(result, "status", glua.LNumber(500)) + L.SetField(result, "body", glua.LString("location manager not initialized")) + L.Push(result) + return 1 + } + + // 获取父请求上下文（从当前协程） + parentCtx := getRequestCtx(L) + if parentCtx == nil { + // 没有父请求上下文，使用模拟上下文 mockCtx := &fasthttp.RequestCtx{} mockCtx.Request.SetRequestURI(uri) + parentCtx = mockCtx + } - captureResult, err := manager.Capture(mockCtx, uri, opts) - if err == nil && captureResult != nil { - L.SetField(result, "status", glua.LNumber(captureResult.Status)) - L.SetField(result, "body", glua.LString(string(captureResult.Body))) + // 执行子请求，传递父请求上下文用于数据复制 + captureResult, err := manager.Capture(parentCtx, uri, opts) + if err == nil && captureResult != nil { + L.SetField(result, "status", glua.LNumber(captureResult.Status)) + L.SetField(result, "body", glua.LString(string(captureResult.Body))) - // 设置 headers - headersTable := headersToLuaTable(L, captureResult.Headers) - L.SetField(result, "headers", headersTable) - } else { - // 执行失败 - L.SetField(result, "status", glua.LNumber(500)) - L.SetField(result, "body", glua.LString("subrequest failed")) - } + // 设置 headers + headersTable := headersToLuaTable(L, captureResult.Headers) + L.SetField(result, "headers", headersTable) } else { - // manager 未初始化 - L.SetField(result, "status", glua.LNumber(404)) - L.SetField(result, "body", glua.LString("location manager not initialized")) + // 执行失败 + L.SetField(result, "status", glua.LNumber(500)) + L.SetField(result, "body", glua.LString("subrequest failed: "+err.Error())) } L.Push(result) @@ -172,3 +198,21 @@ func headersToLuaTable(L *glua.LState, headers map[string]string) *glua.LTable { } return table } + +// RegisterSchedulerUnsafeLocationAPI 为 Scheduler LState 注册不安全的 ngx.location API +// 这些 API 在 scheduler 模式下会返回错误 +func RegisterSchedulerUnsafeLocationAPI(L *glua.LState, ngx *glua.LTable) { + // 创建 ngx.location 表 + location := L.NewTable() + + // ngx.location.capture 在 scheduler 模式下不可用 + L.SetField(location, "capture", L.NewFunction(luaSchedulerUnsafeLocation)) + + L.SetField(ngx, "location", location) +} + +// luaSchedulerUnsafeLocation 返回 scheduler 模式下不可用的错误 +func luaSchedulerUnsafeLocation(L *glua.LState) int { + L.RaiseError("API ngx.location.capture not available in timer callback context") + return 0 +} diff --git a/internal/lua/api_location_test.go b/internal/lua/api_location_test.go index b8148c4..c98c3bc 100644 --- a/internal/lua/api_location_test.go +++ b/internal/lua/api_location_test.go @@ -116,3 +116,33 @@ func TestLocationLuaAPI(t *testing.T) { `) require.NoError(t, err) } + +// TestLocationCaptureWithParentRequest 测试子请求能够访问父请求数据 +func TestLocationCaptureWithParentRequest(t *testing.T) { + engine, err := NewEngine(DefaultConfig()) + require.NoError(t, err) + defer engine.Close() + + // 注册 location handler，它会检查父请求的头信息 + engine.LocationManager().Register("/api/sub", func(ctx *fasthttp.RequestCtx) { + // 检查是否继承了父请求的头 + parentHeader := ctx.Request.Header.Peek("X-Parent-Header") + + ctx.SetStatusCode(200) + ctx.SetBodyString("parent_header: " + string(parentHeader)) + }) + + // 创建父请求上下文 + parentCtx := &fasthttp.RequestCtx{} + parentCtx.Request.SetRequestURI("/parent") + parentCtx.Request.Header.Set("X-Parent-Header", "header_value") + + // 使用 Capture 进行子请求 + result, err := engine.LocationManager().Capture(parentCtx, "/api/sub", nil) + require.NoError(t, err) + require.NotNil(t, result) + + // 验证子请求能够访问父请求数据（Headers 被继承） + assert.Equal(t, 200, result.Status) + assert.Contains(t, string(result.Body), "parent_header: header_value") +}