feat: restore security package and add regression tests to PR branch
This commit is contained in:
parent
be530df648
commit
edc44f3063
3 changed files with 39 additions and 1 deletions
|
|
@ -9,7 +9,9 @@ import (
|
||||||
"github.com/stretchr/testify/require"
|
"github.com/stretchr/testify/require"
|
||||||
|
|
||||||
"github.com/sipeed/picoclaw/cmd/picoclaw/internal"
|
"github.com/sipeed/picoclaw/cmd/picoclaw/internal"
|
||||||
|
"github.com/sipeed/picoclaw/pkg/agent"
|
||||||
"github.com/sipeed/picoclaw/pkg/config"
|
"github.com/sipeed/picoclaw/pkg/config"
|
||||||
|
"github.com/sipeed/picoclaw/pkg/security"
|
||||||
)
|
)
|
||||||
|
|
||||||
func TestNewPicoclawCommand(t *testing.T) {
|
func TestNewPicoclawCommand(t *testing.T) {
|
||||||
|
|
@ -57,3 +59,24 @@ func TestNewPicoclawCommand(t *testing.T) {
|
||||||
assert.False(t, subcmd.Hidden)
|
assert.False(t, subcmd.Hidden)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
func TestSecurityShieldRegistration_Regression(t *testing.T) {
|
||||||
|
// Initialize security hooks (registers them in pkg/agent)
|
||||||
|
security.Init()
|
||||||
|
|
||||||
|
expectedHooks := []string{
|
||||||
|
"security_policy",
|
||||||
|
"security_canary",
|
||||||
|
"security_behavior",
|
||||||
|
"security_pii",
|
||||||
|
"security_ipia",
|
||||||
|
}
|
||||||
|
|
||||||
|
for _, name := range expectedHooks {
|
||||||
|
t.Run(name, func(t *testing.T) {
|
||||||
|
assert.True(t, agent.IsBuiltinHookRegistered(name),
|
||||||
|
"Builtin hook %q is not registered. This usually means pkg/security was deleted "+
|
||||||
|
"or security.Init() is no longer called in main().", name)
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
|
||||||
|
|
@ -8,8 +8,17 @@ import (
|
||||||
"time"
|
"time"
|
||||||
|
|
||||||
"github.com/sipeed/picoclaw/pkg/config"
|
"github.com/sipeed/picoclaw/pkg/config"
|
||||||
|
"github.com/sipeed/picoclaw/pkg/logger"
|
||||||
)
|
)
|
||||||
|
|
||||||
|
// IsBuiltinHookRegistered returns true if a builtin hook factory is registered with the given name.
|
||||||
|
func IsBuiltinHookRegistered(name string) bool {
|
||||||
|
builtinHookRegistryMu.RLock()
|
||||||
|
defer builtinHookRegistryMu.RUnlock()
|
||||||
|
_, exists := builtinHookRegistry[name]
|
||||||
|
return exists
|
||||||
|
}
|
||||||
|
|
||||||
type hookRuntime struct {
|
type hookRuntime struct {
|
||||||
initOnce sync.Once
|
initOnce sync.Once
|
||||||
mu sync.Mutex
|
mu sync.Mutex
|
||||||
|
|
@ -143,19 +152,25 @@ func (al *AgentLoop) loadConfiguredHooks(ctx context.Context) (err error) {
|
||||||
spec := al.cfg.Hooks.Builtins[name]
|
spec := al.cfg.Hooks.Builtins[name]
|
||||||
factory, ok := lookupBuiltinHook(name)
|
factory, ok := lookupBuiltinHook(name)
|
||||||
if !ok {
|
if !ok {
|
||||||
|
logger.WarnCF("agent", "Builtin hook not registered", map[string]any{"hook": name})
|
||||||
return fmt.Errorf("builtin hook %q is not registered", name)
|
return fmt.Errorf("builtin hook %q is not registered", name)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
logger.DebugCF("agent", "Executing builtin hook factory", map[string]any{"hook": name})
|
||||||
hook, factoryErr := factory(ctx, spec)
|
hook, factoryErr := factory(ctx, spec)
|
||||||
if factoryErr != nil {
|
if factoryErr != nil {
|
||||||
|
logger.ErrorCF("agent", "Builtin hook factory failed", map[string]any{"hook": name, "error": factoryErr.Error()})
|
||||||
return fmt.Errorf("build builtin hook %q: %w", name, factoryErr)
|
return fmt.Errorf("build builtin hook %q: %w", name, factoryErr)
|
||||||
}
|
}
|
||||||
|
logger.DebugCF("agent", "Builtin hook factory finished", map[string]any{"hook": name})
|
||||||
|
|
||||||
if err := al.MountHook(HookRegistration{
|
if err := al.MountHook(HookRegistration{
|
||||||
Name: name,
|
Name: name,
|
||||||
Priority: spec.Priority,
|
Priority: spec.Priority,
|
||||||
Source: HookSourceInProcess,
|
Source: HookSourceInProcess,
|
||||||
Hook: hook,
|
Hook: hook,
|
||||||
}); err != nil {
|
}); err != nil {
|
||||||
|
logger.ErrorCF("agent", "Failed to mount builtin hook", map[string]any{"hook": name, "error": err.Error()})
|
||||||
return fmt.Errorf("mount builtin hook %q: %w", name, err)
|
return fmt.Errorf("mount builtin hook %q: %w", name, err)
|
||||||
}
|
}
|
||||||
mounted = append(mounted, name)
|
mounted = append(mounted, name)
|
||||||
|
|
|
||||||
|
|
@ -55,7 +55,7 @@ func (c *Checker) ApproveTool(ctx context.Context, req *agent.ToolApprovalReques
|
||||||
if c.Config.AllowedTools[req.Tool] {
|
if c.Config.AllowedTools[req.Tool] {
|
||||||
allowed = true
|
allowed = true
|
||||||
} else {
|
} else {
|
||||||
// Check for prefix matches (e.g. "monday" matches "mcp_monday_...")
|
// Check for prefix matches (e.g. "github" matches "mcp_github_...")
|
||||||
// Match logic consistent with ToolRegistry.Filter
|
// Match logic consistent with ToolRegistry.Filter
|
||||||
for w, ok := range c.Config.AllowedTools {
|
for w, ok := range c.Config.AllowedTools {
|
||||||
if !ok {
|
if !ok {
|
||||||
|
|
|
||||||
Loading…
Add table
Reference in a new issue