fix(shm): 审核修复——arena 并发安全 + ToolCall ring 抢占 + arenaRead 校验

§13.2: arenaAlloc 改 CAS bump(修复并发覆盖);arenaRead 增加 generation + arena 边界校验
§13.3: Reserve 从 Store+Load 改 CAS-free→Reserved 状态机抢占帧(修复并发分配重复);
  增加 RESERVED 中间态;SetReading/SetReady 改 CAS 返回 bool;ReleaseFrame 清零跳过 state
§13.5/13.6: 注入方法接入 resolveText(从 arena 读 SharedRef);cleanerProxy 加 arenaMu 串行化+每次重置
§13.8: ContextPolicy prune topK 改为 maxContextSize-1(不再硬编码 20)
  tool.register 校验 context_policy 只允许 none/prune
This commit is contained in:
JianFeeeee
2026-09-10 17:03:32 +08:00
parent 2bc813be13
commit cf3c99233b
6 changed files with 140 additions and 36 deletions

View File

@ -308,7 +308,11 @@ func (a *Agent) process(input string, stageCtx *sdk.StageContext) (response stri
// ContextPolicy: prune 工具调用后执行上下文裁剪(§13.8) // ContextPolicy: prune 工具调用后执行上下文裁剪(§13.8)
if def := a.stageHost.ToolDef(tc.Name); def != nil && def.ContextPolicy == "prune" { if def := a.stageHost.ToolDef(tc.Name); def != nil && def.ContextPolicy == "prune" {
if a.context != nil { if a.context != nil {
a.context.Prune(result, 20, a.docStore) topK := a.maxContextSize - 1
if topK < 1 {
topK = 1
}
a.context.Prune(result, topK, a.docStore)
} }
} }

View File

@ -154,28 +154,28 @@ func (h *coreHandler) Handle(method string, params json.RawMessage) (interface{}
if err := unmarshal(params, &p); err != nil { if err := unmarshal(params, &p); err != nil {
return nil, err return nil, err
} }
h.sdk.InjectText(p.Source, p.Channel, p.Text) h.sdk.InjectText(p.Source, p.Channel, h.resolveText(p))
return nil, nil return nil, nil
case MethodIOInjectInterrupt: case MethodIOInjectInterrupt:
var p injectParams var p injectParams
if err := unmarshal(params, &p); err != nil { if err := unmarshal(params, &p); err != nil {
return nil, err return nil, err
} }
h.sdk.InjectInterruptText(p.Source, p.Channel, p.Text) h.sdk.InjectInterruptText(p.Source, p.Channel, h.resolveText(p))
return nil, nil return nil, nil
case MethodIOInjectTextNoMem: case MethodIOInjectTextNoMem:
var p injectParams var p injectParams
if err := unmarshal(params, &p); err != nil { if err := unmarshal(params, &p); err != nil {
return nil, err return nil, err
} }
h.sdk.InjectTextNoMemory(p.Source, p.Channel, p.Text) h.sdk.InjectTextNoMemory(p.Source, p.Channel, h.resolveText(p))
return nil, nil return nil, nil
case MethodIOInjectSync: case MethodIOInjectSync:
var p injectParams var p injectParams
if err := unmarshal(params, &p); err != nil { if err := unmarshal(params, &p); err != nil {
return nil, err return nil, err
} }
return map[string]interface{}{"reply": h.sdk.InjectInputSync(p.Source, p.Channel, p.Text)}, nil return map[string]interface{}{"reply": h.sdk.InjectInputSync(p.Source, p.Channel, h.resolveText(p))}, nil
case MethodIOInjectMedia: case MethodIOInjectMedia:
var p injectMediaParams var p injectMediaParams
@ -563,7 +563,7 @@ func (h *coreHandler) Handle(method string, params json.RawMessage) (interface{}
// 否则使用内联 Text。兼容新旧两种协议。 // 否则使用内联 Text。兼容新旧两种协议。
func (h *coreHandler) resolveText(p injectParams) string { func (h *coreHandler) resolveText(p injectParams) string {
if !p.TextRef.IsZero() { if !p.TextRef.IsZero() {
return string(p.TextRef.Slice(h.host.data)) return string(h.host.unified.arenaRead(p.TextRef))
} }
return p.Text return p.Text
} }
@ -610,6 +610,12 @@ func (h *coreHandler) cleanerProxy(scope, name string, enabled bool) (func(strin
return nil, fmt.Errorf("%s %s 声明 Cleaner,但清洗回调通道未就绪", scope, name) return nil, fmt.Errorf("%s %s 声明 Cleaner,但清洗回调通道未就绪", scope, name)
} }
return func(text string) string { return func(text string) string {
// Cleaner 输入和输出共享同一 arena。串行覆盖完整往返,保证结果读完前
// 不被其他调用重置;返回 string 会复制结果,随后即可回收整个临时区。
h.host.arenaMu.Lock()
defer h.host.arenaMu.Unlock()
defer h.host.unified.arenaReset()
ref, err := h.host.unified.arenaWrite([]byte(text)) ref, err := h.host.unified.arenaWrite([]byte(text))
if err != nil { if err != nil {
return text return text
@ -619,6 +625,9 @@ func (h *coreHandler) cleanerProxy(scope, name string, enabled bool) (func(strin
return text return text
} }
result := h.host.unified.arenaRead(resultRef) result := h.host.unified.arenaRead(resultRef)
if result == nil && !resultRef.IsZero() {
return text
}
return string(result) return string(result)
}, nil }, nil
} }
@ -636,6 +645,11 @@ func (h *coreHandler) toolRegister(params json.RawMessage) (interface{}, error)
if p.Name == "" { if p.Name == "" {
return nil, fmt.Errorf("tool.register: 缺少 name") return nil, fmt.Errorf("tool.register: 缺少 name")
} }
switch p.Def.ContextPolicy {
case "", "none", "prune":
default:
return nil, fmt.Errorf("tool.register: context_policy 只允许 none/prune,实际 %q", p.Def.ContextPolicy)
}
p.Def.Plugin = h.name p.Def.Plugin = h.name
// 函数本身不进 JSON;has_cleaner 只声明其存在,实际执行回到插件进程。 // 函数本身不进 JSON;has_cleaner 只声明其存在,实际执行回到插件进程。
cleaner, err := h.cleanerProxy(CleanerScopeTool, p.Name, p.HasCleaner) cleaner, err := h.cleanerProxy(CleanerScopeTool, p.Name, p.HasCleaner)

View File

@ -36,6 +36,7 @@ type Host struct {
evtSubscriber EvtRingSubscriber evtSubscriber EvtRingSubscriber
locks *lockRegistry locks *lockRegistry
arenaMu sync.Mutex // 动态 arena 跨进程调用的生命周期锁
stageMu sync.Mutex stageMu sync.Mutex
coordMu sync.Mutex coordMu sync.Mutex
coord *stageCoordinator coord *stageCoordinator

View File

@ -47,10 +47,11 @@ const (
toolFrameSize = 112 // 单帧总大小 toolFrameSize = 112 // 单帧总大小
// 帧状态 // 帧状态
toolFrameFree uint32 = 0 toolFrameFree uint32 = 0
toolFrameCalling uint32 = 1 // 内核写 input,插件待读 toolFrameReserved uint32 = 1 // 内核已占帧,正在写入元数据
toolFrameReading uint32 = 2 // 插件正在执行 toolFrameCalling uint32 = 2 // 内核写完 input,插件待读
toolFrameReady uint32 = 3 // 插件写 output,内核待读 toolFrameReading uint32 = 3 // 插件正在执行
toolFrameReady uint32 = 4 // 插件写完 output,内核待读
) )
// ToolRing 头偏移(相对区域起始) // ToolRing 头偏移(相对区域起始)
@ -116,10 +117,19 @@ func AttachToolRing(data []byte) (*ToolCallRing, error) {
if magic != toolRingMagic { if magic != toolRingMagic {
return nil, fmt.Errorf("tool ring: 魔数不匹配(0x%x)", magic) return nil, fmt.Errorf("tool ring: 魔数不匹配(0x%x)", magic)
} }
cap := getU32(data[trlOffCap:])
frameSize := getU32(data[trlOffFrameSize:])
version := getU32(data[trlOffVersion:])
if version != toolRingVersion {
return nil, fmt.Errorf("tool ring: 版本不匹配(%d,期望 %d)", version, toolRingVersion)
}
if cap == 0 || frameSize != toolFrameSize || uint64(trlOffFrameBase)+uint64(cap)*uint64(frameSize) > uint64(len(data)) {
return nil, fmt.Errorf("tool ring: 布局非法(cap=%d frameSize=%d total=%d)", cap, frameSize, len(data))
}
return &ToolCallRing{ return &ToolCallRing{
data: data, data: data,
cap: getU32(data[trlOffCap:]), cap: cap,
frameSize: getU32(data[trlOffFrameSize:]), frameSize: frameSize,
framesBase: trlOffFrameBase, framesBase: trlOffFrameBase,
}, nil }, nil
} }
@ -128,13 +138,11 @@ func AttachToolRing(data []byte) (*ToolCallRing, error) {
// //
// 环形扫描:从 writeIdx 开始,绕环一圈找 FREE 帧。全部忙时背压。 // 环形扫描:从 writeIdx 开始,绕环一圈找 FREE 帧。全部忙时背压。
func (r *ToolCallRing) Reserve() (frameIdx uint32, err error) { func (r *ToolCallRing) Reserve() (frameIdx uint32, err error) {
start := r.writeIdx.Load() start := r.writeIdx.Add(1) - 1
for i := uint64(0); i < uint64(r.cap); i++ { for i := uint64(0); i < uint64(r.cap); i++ {
idx := start + i idx := start + i
fi := uint32(idx % uint64(r.cap)) fi := uint32(idx % uint64(r.cap))
state := atomic.LoadUint32(r.statePtr(fi)) if atomic.CompareAndSwapUint32(r.statePtr(fi), toolFrameFree, toolFrameReserved) {
if state == toolFrameFree {
r.writeIdx.Store(idx + 1)
return fi, nil return fi, nil
} }
} }
@ -191,16 +199,20 @@ func (r *ToolCallRing) GetCallingFrame(frameIdx uint32) (name string, inputRef S
return return
} }
// SetReading 将帧状态设为 READING(插件开始执行)。 // SetReading 将 CALLING 帧原子转为 READING(插件开始执行)。
func (r *ToolCallRing) SetReading(frameIdx uint32) { func (r *ToolCallRing) SetReading(frameIdx uint32) bool {
atomic.StoreUint32(r.statePtr(frameIdx), toolFrameReading) return atomic.CompareAndSwapUint32(r.statePtr(frameIdx), toolFrameCalling, toolFrameReading)
} }
// SetReady 将帧状态设为 READY 并写入 output 描述符(插件执行完毕)。 // SetReady 将 READING 帧设为 READY 并写入 output 描述符(插件执行完毕)。
func (r *ToolCallRing) SetReady(frameIdx uint32, outputRef SharedRef) { func (r *ToolCallRing) SetReady(frameIdx uint32, outputRef SharedRef) bool {
if atomic.LoadUint32(r.statePtr(frameIdx)) != toolFrameReading {
return false
}
off := r.frameOff(frameIdx) off := r.frameOff(frameIdx)
packSharedRef(r.data[off+toolFrameOffOutput:], outputRef) packSharedRef(r.data[off+toolFrameOffOutput:], outputRef)
atomic.StoreUint32(r.statePtr(frameIdx), toolFrameReady) atomic.StoreUint32(r.statePtr(frameIdx), toolFrameReady)
return true
} }
// GetReadyFrame 读取帧的 requestID 和 output 描述符(内核消费 READY 帧)。 // GetReadyFrame 读取帧的 requestID 和 output 描述符(内核消费 READY 帧)。
@ -214,8 +226,8 @@ func (r *ToolCallRing) GetReadyFrame(frameIdx uint32) (reqID uint64, outputRef S
// ReleaseFrame 回收帧为 FREE(内核消费完毕后调用)。 // ReleaseFrame 回收帧为 FREE(内核消费完毕后调用)。
func (r *ToolCallRing) ReleaseFrame(frameIdx uint32) { func (r *ToolCallRing) ReleaseFrame(frameIdx uint32) {
off := r.frameOff(frameIdx) off := r.frameOff(frameIdx)
// 清零帧内容 // state 必须最后发布 FREE;否则并发 Reserve 可能在其余字段尚未清零时复用帧。
for i := uint32(0); i < r.frameSize; i++ { for i := uint32(4); i < r.frameSize; i++ {
r.data[off+i] = 0 r.data[off+i] = 0
} }
atomic.StoreUint32(r.statePtr(frameIdx), toolFrameFree) atomic.StoreUint32(r.statePtr(frameIdx), toolFrameFree)

View File

@ -1,6 +1,7 @@
package proc package proc
import ( import (
"sync"
"testing" "testing"
"unsafe" "unsafe"
) )
@ -154,6 +155,56 @@ func TestSharedRef_PackUnpack(t *testing.T) {
} }
} }
func TestToolCallRing_ConcurrentReserveUnique(t *testing.T) {
data := make([]byte, int(trlOffFrameBase+toolRingCap*toolFrameSize))
if err := InitToolRing(data); err != nil {
t.Fatal(err)
}
ring, err := AttachToolRing(data)
if err != nil {
t.Fatal(err)
}
indices := make(chan uint32, toolRingCap)
var wg sync.WaitGroup
for i := uint32(0); i < toolRingCap; i++ {
wg.Add(1)
go func() {
defer wg.Done()
idx, reserveErr := ring.Reserve()
if reserveErr != nil {
t.Errorf("Reserve: %v", reserveErr)
return
}
indices <- idx
}()
}
wg.Wait()
close(indices)
seen := make(map[uint32]bool, toolRingCap)
for idx := range indices {
if seen[idx] {
t.Fatalf("并发 Reserve 重复分配帧 %d", idx)
}
seen[idx] = true
}
if len(seen) != int(toolRingCap) {
t.Fatalf("唯一帧数=%d,期望 %d", len(seen), toolRingCap)
}
}
func TestAttachToolRingRejectsInvalidLayout(t *testing.T) {
data := make([]byte, trlOffFrameBase)
putU32(data[trlOffMagic:], toolRingMagic)
putU32(data[trlOffVersion:], toolRingVersion)
putU32(data[trlOffCap:], toolRingCap)
putU32(data[trlOffFrameSize:], toolFrameSize)
if _, err := AttachToolRing(data); err == nil {
t.Fatal("越界布局应被拒绝")
}
}
func TestToolFrameLayoutAligned(t *testing.T) { func TestToolFrameLayoutAligned(t *testing.T) {
// 确保帧布局字段偏移与内存布局一致(安全断言) // 确保帧布局字段偏移与内存布局一致(安全断言)
var frame toolFrame var frame toolFrame

View File

@ -94,10 +94,11 @@ func (r SharedRef) IsZero() bool {
// Slice 从共享内存中按 SharedRef 切片。data 必须是完整的 mmap 区域。 // Slice 从共享内存中按 SharedRef 切片。data 必须是完整的 mmap 区域。
func (r SharedRef) Slice(data []byte) []byte { func (r SharedRef) Slice(data []byte) []byte {
if r.IsZero() || int(r.Offset)+int(r.Length) > len(data) { end := uint64(r.Offset) + uint64(r.Length)
if r.IsZero() || end > uint64(len(data)) {
return nil return nil
} }
return data[r.Offset : r.Offset+r.Length] return data[r.Offset:uint32(end)]
} }
// unifiedRegion 统一共享内存区域的内核侧视图。 // unifiedRegion 统一共享内存区域的内核侧视图。
@ -171,12 +172,22 @@ func attachUnifiedRegion(data []byte) (*unifiedRegion, error) {
evtOff := getU32(data[sbOffEvtOff:]) evtOff := getU32(data[sbOffEvtOff:])
evtSize := getU32(data[sbOffEvtSize:]) evtSize := getU32(data[sbOffEvtSize:])
capacity := getU32(data[sbOffCapacity:])
arenaOff := getU32(data[sbOffArenaOff:]) arenaOff := getU32(data[sbOffArenaOff:])
arenaCap := getU32(data[sbOffArenaCap:]) arenaCap := getU32(data[sbOffArenaCap:])
if int(evtOff)+int(evtSize) > len(data) { if capacity != uint32(len(data)) {
return nil, fmt.Errorf("unified: capacity 不匹配(header=%d mapped=%d)", capacity, len(data))
}
if uint64(ctxOff)+uint64(ctxSize) > uint64(len(data)) {
return nil, fmt.Errorf("unified: StageContext 越界(off=%d size=%d total=%d)", ctxOff, ctxSize, len(data))
}
if uint64(evtOff)+uint64(evtSize) > uint64(len(data)) {
return nil, fmt.Errorf("unified: EvtRing 越界(off=%d size=%d total=%d)", evtOff, evtSize, len(data)) return nil, fmt.Errorf("unified: EvtRing 越界(off=%d size=%d total=%d)", evtOff, evtSize, len(data))
} }
if uint64(arenaOff)+uint64(arenaCap) > uint64(len(data)) {
return nil, fmt.Errorf("unified: arena 越界(off=%d cap=%d total=%d)", arenaOff, arenaCap, len(data))
}
return &unifiedRegion{ return &unifiedRegion{
data: data, data: data,
@ -248,16 +259,19 @@ func (r *unifiedRegion) arenaAlloc(n int) (uint32, error) {
if n <= 0 { if n <= 0 {
return 0, fmt.Errorf("arena: 非法分配长度 %d", n) return 0, fmt.Errorf("arena: 非法分配长度 %d", n)
} }
used := atomic.LoadUint32(r.arenaUsed) for {
if used == 0 { used := atomic.LoadUint32(r.arenaUsed)
used = 1 if used == 0 {
used = 1
}
end := uint64(used) + uint64(n)
if end > uint64(r.arenaCap) {
return 0, fmt.Errorf("arena: 空间不足(需 %d,剩 %d)", n, uint64(r.arenaCap)-uint64(used))
}
if atomic.CompareAndSwapUint32(r.arenaUsed, used, uint32(end)) {
return used, nil
}
} }
end := uint64(used) + uint64(n)
if end > uint64(r.arenaCap) {
return 0, fmt.Errorf("arena: 空间不足(需 %d,剩 %d)", n, uint64(r.arenaCap)-uint64(used))
}
atomic.StoreUint32(r.arenaUsed, uint32(end))
return used, nil
} }
// arenaWrite 把 b 写入动态 arena 并返回 SharedRef。 // arenaWrite 把 b 写入动态 arena 并返回 SharedRef。
@ -275,8 +289,16 @@ func (r *unifiedRegion) arenaWrite(b []byte) (SharedRef, error) {
return SharedRef{Offset: base, Length: uint32(len(b)), Generation: gen}, nil return SharedRef{Offset: base, Length: uint32(len(b)), Generation: gen}, nil
} }
// arenaRead 按 SharedRef 读取数据。 // arenaRead 按 SharedRef 读取数据。generation 不匹配或引用越出 arena 时拒绝。
func (r *unifiedRegion) arenaRead(ref SharedRef) []byte { func (r *unifiedRegion) arenaRead(ref SharedRef) []byte {
if ref.IsZero() || ref.Generation != uint32(r.generation()) {
return nil
}
end := uint64(ref.Offset) + uint64(ref.Length)
arenaEnd := uint64(r.arenaOff) + uint64(r.arenaCap)
if uint64(ref.Offset) < uint64(r.arenaOff)+1 || end > arenaEnd {
return nil
}
return ref.Slice(r.data) return ref.Slice(r.data)
} }