Add files via upload

This commit is contained in:
公明
2026-08-15 01:57:26 +08:00
committed by GitHub
parent f564421b4d
commit 7f8093f8b9
52 changed files with 14151 additions and 0 deletions
File diff suppressed because it is too large Load Diff
+207
View File
@@ -0,0 +1,207 @@
package config
import (
"os"
"path/filepath"
"strings"
"testing"
)
func TestEnsureLocalConfigCreatesFromExample(t *testing.T) {
dir := t.TempDir()
examplePath := filepath.Join(dir, "config.example.yaml")
configPath := filepath.Join(dir, "config.yaml")
example := []byte(`auth:
session_duration_hours: 12
server:
host: 127.0.0.1
port: 8080
`)
if err := os.WriteFile(examplePath, example, 0644); err != nil {
t.Fatalf("write example: %v", err)
}
result, err := EnsureLocalConfig(configPath)
if err != nil {
t.Fatalf("EnsureLocalConfig: %v", err)
}
if !result.Created {
t.Fatal("Created = false, want true")
}
if result.ExamplePath != examplePath {
t.Fatalf("ExamplePath = %q, want %q", result.ExamplePath, examplePath)
}
cfg, err := Load(configPath)
if err != nil {
t.Fatalf("Load generated config: %v", err)
}
if cfg.Auth.SessionDurationHours != 12 {
t.Fatalf("SessionDurationHours = %d, want 12", cfg.Auth.SessionDurationHours)
}
second, err := EnsureLocalConfig(configPath)
if err != nil {
t.Fatalf("EnsureLocalConfig existing: %v", err)
}
if second.Created {
t.Fatal("Created = true for existing config, want false")
}
}
func TestLoadIgnoresLegacyAuthPasswordField(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "config.yaml")
initial := strings.Join([]string{
"auth:",
` password: "legacy-password"`,
" session_duration_hours: 12",
"server:",
" host: 127.0.0.1",
" port: 8080",
"",
}, "\n")
if err := os.WriteFile(path, []byte(initial), 0644); err != nil {
t.Fatalf("write config: %v", err)
}
cfg, err := Load(path)
if err != nil {
t.Fatalf("Load: %v", err)
}
if cfg.Auth.SessionDurationHours != 12 {
t.Fatalf("SessionDurationHours = %d, want 12", cfg.Auth.SessionDurationHours)
}
}
func TestHitlAuditModelEffectiveFallsBackToMainConfig(t *testing.T) {
main := OpenAIConfig{
Provider: "openai",
BaseURL: "https://api.example.com/v1",
APIKey: "main-key",
Model: "large-model",
}
got := (HitlConfig{
AuditModel: OpenAIConfig{Model: "small-reviewer"},
}).AuditModelEffective(main)
if got.Provider != main.Provider || got.BaseURL != main.BaseURL || got.APIKey != main.APIKey {
t.Fatalf("expected provider/base_url/api_key to inherit main config, got %+v", got)
}
if got.Model != "small-reviewer" {
t.Fatalf("expected audit model override, got %q", got.Model)
}
}
func TestLoadUsesAIDefaultChannelAsRuntimeOpenAI(t *testing.T) {
dir := t.TempDir()
path := filepath.Join(dir, "config.yaml")
initial := strings.Join([]string{
"ai:",
" default_channel: deepseek",
" channels:",
" qwen:",
" name: Qwen",
" provider: openai_compatible",
" base_url: https://dashscope.example/v1",
" api_key: qwen-key",
" model: qwen-max",
" deepseek:",
" name: DeepSeek",
" provider: openai_compatible",
" base_url: https://deepseek.example/v1",
" api_key: deepseek-key",
" model: deepseek-chat",
" max_total_tokens: 64000",
"server:",
" host: 127.0.0.1",
" port: 8080",
"",
}, "\n")
if err := os.WriteFile(path, []byte(initial), 0644); err != nil {
t.Fatalf("write config: %v", err)
}
cfg, err := Load(path)
if err != nil {
t.Fatalf("Load: %v", err)
}
if cfg.OpenAI.Model != "deepseek-chat" || cfg.OpenAI.APIKey != "deepseek-key" || cfg.OpenAI.MaxTotalTokens != 64000 {
t.Fatalf("runtime OpenAI config did not follow ai.default_channel: %+v", cfg.OpenAI)
}
oa, id, ok := cfg.ResolveAIChannel("qwen")
if !ok || id != "qwen" || oa.Model != "qwen-max" || oa.APIKey != "qwen-key" {
t.Fatalf("ResolveAIChannel(qwen) = (%+v, %q, %v)", oa, id, ok)
}
}
func TestSummarizationUserIntentLedgerRunesEffective(t *testing.T) {
var zero MultiAgentEinoMiddlewareConfig
if got := zero.SummarizationUserIntentLedgerMaxRunesEffective(); got != DefaultSummarizationUserIntentLedgerMaxRunes {
t.Fatalf("default ledger max runes = %d, want %d", got, DefaultSummarizationUserIntentLedgerMaxRunes)
}
if got := zero.SummarizationUserIntentLedgerEntryMaxRunesEffective(); got != DefaultSummarizationUserIntentLedgerEntryMaxRunes {
t.Fatalf("default ledger entry max runes = %d, want %d", got, DefaultSummarizationUserIntentLedgerEntryMaxRunes)
}
custom := MultiAgentEinoMiddlewareConfig{
SummarizationUserIntentLedgerMaxRunes: 12345,
SummarizationUserIntentLedgerEntryMaxRunes: 2345,
}
if got := custom.SummarizationUserIntentLedgerMaxRunesEffective(); got != 12345 {
t.Fatalf("custom ledger max runes = %d", got)
}
if got := custom.SummarizationUserIntentLedgerEntryMaxRunesEffective(); got != 2345 {
t.Fatalf("custom ledger entry max runes = %d", got)
}
}
func TestSummarizationOutputReserveTokensEffective(t *testing.T) {
var zero MultiAgentEinoMiddlewareConfig
if got := zero.SummarizationOutputReserveTokensEffective(); got != DefaultSummarizationOutputReserveTokens {
t.Fatalf("default output reserve = %d, want %d", got, DefaultSummarizationOutputReserveTokens)
}
custom := MultiAgentEinoMiddlewareConfig{SummarizationOutputReserveTokens: 4096}
if got := custom.SummarizationOutputReserveTokensEffective(); got != 4096 {
t.Fatalf("custom output reserve = %d", got)
}
}
func TestOpenAIOutputLimitValidation(t *testing.T) {
if got := (OpenAIConfig{}).MaxCompletionTokensEffective(); got != DefaultMaxCompletionTokens {
t.Fatalf("max completion default=%d", got)
}
if err := validateOpenAIOutputLimits(OpenAIConfig{MaxCompletionTokens: -1}); err == nil {
t.Fatal("negative completion limit must fail")
}
}
func TestLatestUserMessageRunesEffective(t *testing.T) {
var zero MultiAgentEinoMiddlewareConfig
if got := zero.LatestUserMessageMaxRunesEffective(); got != DefaultLatestUserMessageMaxRunes {
t.Fatalf("default latest user max runes = %d, want %d", got, DefaultLatestUserMessageMaxRunes)
}
if got := zero.LatestUserMessageHeadRunesEffective(); got != DefaultLatestUserMessageHeadRunes {
t.Fatalf("default latest user head runes = %d, want %d", got, DefaultLatestUserMessageHeadRunes)
}
if got := zero.LatestUserMessageTailRunesEffective(); got != DefaultLatestUserMessageTailRunes {
t.Fatalf("default latest user tail runes = %d, want %d", got, DefaultLatestUserMessageTailRunes)
}
custom := MultiAgentEinoMiddlewareConfig{
LatestUserMessageMaxRunes: 100,
LatestUserMessageHeadRunes: 40,
LatestUserMessageTailRunes: 60,
}
if got := custom.LatestUserMessageMaxRunesEffective(); got != 100 {
t.Fatalf("custom latest user max runes = %d", got)
}
if got := custom.LatestUserMessageHeadRunesEffective(); got != 40 {
t.Fatalf("custom latest user head runes = %d", got)
}
if got := custom.LatestUserMessageTailRunesEffective(); got != 60 {
t.Fatalf("custom latest user tail runes = %d", got)
}
}
+66
View File
@@ -0,0 +1,66 @@
package config
import (
"os"
"strings"
)
// expandEnvVar 展开字符串中的 ${VAR} 和 ${VAR:-default} 环境变量引用。
// 与官方 MCP 配置格式一致(Claude Desktop / Cursor / VS Code 均支持此语法)。
func expandEnvVar(s string) string {
var b strings.Builder
i := 0
for i < len(s) {
// 查找 ${
idx := strings.Index(s[i:], "${")
if idx < 0 {
b.WriteString(s[i:])
break
}
b.WriteString(s[i : i+idx])
i += idx + 2 // skip ${
// 查找对应的 }
end := strings.IndexByte(s[i:], '}')
if end < 0 {
// 没有 },原样保留
b.WriteString("${")
continue
}
expr := s[i : i+end]
i += end + 1 // skip }
// 解析 VAR:-default
varName := expr
defaultVal := ""
hasDefault := false
if colonIdx := strings.Index(expr, ":-"); colonIdx >= 0 {
varName = expr[:colonIdx]
defaultVal = expr[colonIdx+2:]
hasDefault = true
}
val := os.Getenv(varName)
if val == "" && hasDefault {
val = defaultVal
}
b.WriteString(val)
}
return b.String()
}
// ExpandConfigEnv 展开 ExternalMCPServerConfig 中所有支持环境变量的字段。
// 展开范围:Command、Args、Env values、URL、Headers values。
func ExpandConfigEnv(cfg *ExternalMCPServerConfig) {
cfg.Command = expandEnvVar(cfg.Command)
for i, arg := range cfg.Args {
cfg.Args[i] = expandEnvVar(arg)
}
for k, v := range cfg.Env {
cfg.Env[k] = expandEnvVar(v)
}
cfg.URL = expandEnvVar(cfg.URL)
for k, v := range cfg.Headers {
cfg.Headers[k] = expandEnvVar(v)
}
}
+81
View File
@@ -0,0 +1,81 @@
package config
import (
"os"
"testing"
)
func TestExpandEnvVar(t *testing.T) {
os.Setenv("TEST_MCP_VAR", "hello")
os.Setenv("TEST_MCP_PATH", "/usr/local/bin")
defer os.Unsetenv("TEST_MCP_VAR")
defer os.Unsetenv("TEST_MCP_PATH")
tests := []struct {
name string
input string
expect string
}{
{"plain string", "no vars here", "no vars here"},
{"empty string", "", ""},
{"simple var", "${TEST_MCP_VAR}", "hello"},
{"var in middle", "prefix-${TEST_MCP_VAR}-suffix", "prefix-hello-suffix"},
{"multiple vars", "${TEST_MCP_PATH}/${TEST_MCP_VAR}", "/usr/local/bin/hello"},
{"missing var empty", "${NONEXISTENT_MCP_VAR_XYZ}", ""},
{"default value used", "${NONEXISTENT_MCP_VAR_XYZ:-fallback}", "fallback"},
{"default not used", "${TEST_MCP_VAR:-unused}", "hello"},
{"default with path", "${NONEXISTENT_MCP_VAR_XYZ:-/tmp/default}", "/tmp/default"},
{"unclosed brace", "${UNCLOSED", "${UNCLOSED"},
{"dollar without brace", "$PLAIN", "$PLAIN"},
{"empty var name", "${}", ""},
{"default empty var", "${:-default}", "default"},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
got := expandEnvVar(tt.input)
if got != tt.expect {
t.Errorf("expandEnvVar(%q) = %q, want %q", tt.input, got, tt.expect)
}
})
}
}
func TestExpandConfigEnv(t *testing.T) {
os.Setenv("TEST_MCP_CMD", "python3")
os.Setenv("TEST_MCP_TOKEN", "secret123")
defer os.Unsetenv("TEST_MCP_CMD")
defer os.Unsetenv("TEST_MCP_TOKEN")
cfg := &ExternalMCPServerConfig{
Command: "${TEST_MCP_CMD}",
Args: []string{"--token", "${TEST_MCP_TOKEN}", "${MISSING:-default_arg}"},
Env: map[string]string{"API_KEY": "${TEST_MCP_TOKEN}", "LEVEL": "${MISSING:-INFO}"},
URL: "https://${MISSING:-example.com}/mcp",
Headers: map[string]string{"Authorization": "Bearer ${TEST_MCP_TOKEN}"},
}
ExpandConfigEnv(cfg)
if cfg.Command != "python3" {
t.Errorf("Command = %q, want %q", cfg.Command, "python3")
}
if cfg.Args[1] != "secret123" {
t.Errorf("Args[1] = %q, want %q", cfg.Args[1], "secret123")
}
if cfg.Args[2] != "default_arg" {
t.Errorf("Args[2] = %q, want %q", cfg.Args[2], "default_arg")
}
if cfg.Env["API_KEY"] != "secret123" {
t.Errorf("Env[API_KEY] = %q, want %q", cfg.Env["API_KEY"], "secret123")
}
if cfg.Env["LEVEL"] != "INFO" {
t.Errorf("Env[LEVEL] = %q, want %q", cfg.Env["LEVEL"], "INFO")
}
if cfg.URL != "https://example.com/mcp" {
t.Errorf("URL = %q, want %q", cfg.URL, "https://example.com/mcp")
}
if cfg.Headers["Authorization"] != "Bearer secret123" {
t.Errorf("Headers[Authorization] = %q, want %q", cfg.Headers["Authorization"], "Bearer secret123")
}
}
+31
View File
@@ -0,0 +1,31 @@
package config
import (
"strings"
"testing"
)
func TestDefaultHitlAuditAgentPromptIncludesPrioritizedRules(t *testing.T) {
prompt := DefaultHitlAuditAgentPrompt()
for _, want := range []string{
"如果同时命中 reject 和 approve,必须 reject",
"修改/重置任意用户或管理员密码",
"修改/创建/删除用户、角色、权限",
"停止、禁用、重启业务服务",
"命中规则:...",
} {
if !strings.Contains(prompt, want) {
t.Fatalf("default approval prompt missing %q", want)
}
}
}
func TestDefaultHitlAuditAgentPromptReviewEditKeepsEditedArguments(t *testing.T) {
prompt := DefaultHitlAuditAgentPromptReviewEdit()
if !strings.Contains(prompt, `"editedArguments":{...}`) {
t.Fatal("review-edit prompt must preserve editedArguments output")
}
if !strings.Contains(prompt, "命中规则:...") {
t.Fatal("review-edit prompt must require a matched rule")
}
}
+69
View File
@@ -0,0 +1,69 @@
package config
import "testing"
func TestValidateWecomConfig(t *testing.T) {
t.Parallel()
tests := []struct {
name string
cfg RobotWecomConfig
wantErr bool
}{
{
name: "disabled without token",
cfg: RobotWecomConfig{Enabled: false, Token: ""},
wantErr: false,
},
{
name: "enabled with token",
cfg: RobotWecomConfig{Enabled: true, Token: "secret"},
wantErr: false,
},
{
name: "enabled without token",
cfg: RobotWecomConfig{Enabled: true, Token: ""},
wantErr: true,
},
{
name: "enabled with whitespace token",
cfg: RobotWecomConfig{Enabled: true, Token: " "},
wantErr: true,
},
}
for _, tt := range tests {
tt := tt
t.Run(tt.name, func(t *testing.T) {
t.Parallel()
err := ValidateWecomConfig(tt.cfg)
if (err != nil) != tt.wantErr {
t.Fatalf("ValidateWecomConfig() error = %v, wantErr %v", err, tt.wantErr)
}
})
}
}
func TestValidateRobotAuthorization(t *testing.T) {
tests := []struct {
name string
cfg RobotAuthorizationConfig
wantErr bool
}{
{name: "default user binding", cfg: RobotAuthorizationConfig{}, wantErr: false},
{name: "explicit user binding", cfg: RobotAuthorizationConfig{Mode: RobotAuthModeUserBinding}, wantErr: false},
{name: "service account", cfg: RobotAuthorizationConfig{Mode: RobotAuthModeServiceAccount, ServiceUserID: "svc-1", AllowedExternalUsers: []string{"t:x|u:y"}}, wantErr: false},
{name: "missing service user", cfg: RobotAuthorizationConfig{Mode: RobotAuthModeServiceAccount, AllowedExternalUsers: []string{"t:x|u:y"}}, wantErr: true},
{name: "admin allowed with exact sender", cfg: RobotAuthorizationConfig{Mode: RobotAuthModeServiceAccount, ServiceUserID: "admin", AllowedExternalUsers: []string{"t:x|u:y"}}, wantErr: false},
{name: "allowlist required", cfg: RobotAuthorizationConfig{Mode: RobotAuthModeServiceAccount, ServiceUserID: "svc-1"}, wantErr: true},
{name: "wildcard forbidden", cfg: RobotAuthorizationConfig{Mode: RobotAuthModeServiceAccount, ServiceUserID: "svc-1", AllowedExternalUsers: []string{"*"}}, wantErr: true},
{name: "unknown mode", cfg: RobotAuthorizationConfig{Mode: "open"}, wantErr: true},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
if err := ValidateRobotAuthorization(tt.cfg, "robots.lark"); (err != nil) != tt.wantErr {
t.Fatalf("ValidateRobotAuthorization() error=%v wantErr=%v", err, tt.wantErr)
}
})
}
}
+60
View File
@@ -0,0 +1,60 @@
package config
import "strings"
// MainWebUIUsesHTTPS 判断主 Web UI 是否以 HTTPS 监听(与 internal/app.prepareMainServerTLS 前置条件一致)。
func MainWebUIUsesHTTPS(s *ServerConfig) bool {
if s == nil {
return false
}
if s.TLSEnabled {
return true
}
if s.TLSAutoSelfSign {
return true
}
cert := strings.TrimSpace(s.TLSCertPath)
key := strings.TrimSpace(s.TLSKeyPath)
return cert != "" && key != ""
}
// ServerHTTPRedirectEnabled 是否在主站启用 HTTPS 时把明文 HTTP 请求重定向到 HTTPS(默认开启)。
func ServerHTTPRedirectEnabled(s *ServerConfig) bool {
if s == nil || !MainWebUIUsesHTTPS(s) {
return false
}
if s.TLSHTTPRedirect == nil {
return true
}
return *s.TLSHTTPRedirect
}
// ApplyDevHTTPSBootstrap 供 --https / 一键脚本使用:强制开启主站 TLS。
// 若已配置 tls_cert_path 与 tls_key_path 则仅用 PEM,不开启自签;否则启用 tls_auto_self_sign(内存证书,仅本地测试)。
func ApplyDevHTTPSBootstrap(cfg *Config) {
if cfg == nil {
return
}
cfg.Server.TLSEnabled = true
cert := strings.TrimSpace(cfg.Server.TLSCertPath)
key := strings.TrimSpace(cfg.Server.TLSKeyPath)
if cert != "" && key != "" {
cfg.Server.TLSAutoSelfSign = false
return
}
cfg.Server.TLSAutoSelfSign = true
}
// ApplyPlainHTTPBootstrap 供 --http / 一键脚本使用:强制主站使用明文 HTTP。
// 它会覆盖配置文件中的 TLS 开关、自签证书以及证书路径,避免 --http 仍被配置中的 HTTPS 选项重新启用。
func ApplyPlainHTTPBootstrap(cfg *Config) {
if cfg == nil {
return
}
cfg.Server.TLSEnabled = false
cfg.Server.TLSAutoSelfSign = false
cfg.Server.TLSCertPath = ""
cfg.Server.TLSKeyPath = ""
disabled := false
cfg.Server.TLSHTTPRedirect = &disabled
}
@@ -0,0 +1,31 @@
package config
import "testing"
func TestApplyPlainHTTPBootstrapDisablesConfiguredTLS(t *testing.T) {
enabled := true
cfg := &Config{
Server: ServerConfig{
TLSEnabled: true,
TLSAutoSelfSign: true,
TLSCertPath: "/tmp/server.crt",
TLSKeyPath: "/tmp/server.key",
TLSHTTPRedirect: &enabled,
},
}
ApplyPlainHTTPBootstrap(cfg)
if MainWebUIUsesHTTPS(&cfg.Server) {
t.Fatal("expected --http bootstrap to disable main web UI HTTPS")
}
if ServerHTTPRedirectEnabled(&cfg.Server) {
t.Fatal("expected --http bootstrap to disable HTTP to HTTPS redirect")
}
if cfg.Server.TLSCertPath != "" || cfg.Server.TLSKeyPath != "" {
t.Fatalf("expected TLS cert paths to be cleared, got cert=%q key=%q", cfg.Server.TLSCertPath, cfg.Server.TLSKeyPath)
}
if cfg.Server.TLSHTTPRedirect == nil || *cfg.Server.TLSHTTPRedirect {
t.Fatal("expected TLSHTTPRedirect to be explicitly disabled")
}
}
+111
View File
@@ -0,0 +1,111 @@
package config
import (
"os"
"path/filepath"
"testing"
)
func TestReloadSecurityToolsFromDir(t *testing.T) {
root := t.TempDir()
toolsDir := filepath.Join(root, "tools")
if err := os.MkdirAll(toolsDir, 0755); err != nil {
t.Fatal(err)
}
configPath := filepath.Join(root, "config.yaml")
if err := os.WriteFile(configPath, []byte(`security:
tools_dir: tools
tools:
- name: inline-only
command: inline-cmd
enabled: true
description: inline tool
`), 0644); err != nil {
t.Fatal(err)
}
writeTool := func(name, command string) {
t.Helper()
content := "name: " + name + "\ncommand: " + command + "\nenabled: true\ndescription: test\n"
if err := os.WriteFile(filepath.Join(toolsDir, name+".yaml"), []byte(content), 0644); err != nil {
t.Fatal(err)
}
}
writeTool("alpha", "alpha-cmd")
cfg := &Config{
Security: SecurityConfig{
ToolsDir: "tools",
Tools: []ToolConfig{
{Name: "stale", Command: "stale-cmd", Enabled: true, Description: "should be removed"},
},
},
}
if err := ReloadSecurityToolsFromDir(cfg, configPath); err != nil {
t.Fatalf("reload: %v", err)
}
if len(cfg.Security.Tools) != 2 {
t.Fatalf("expected 2 tools, got %d", len(cfg.Security.Tools))
}
names := map[string]string{}
for _, tool := range cfg.Security.Tools {
names[tool.Name] = tool.Command
}
if names["alpha"] != "alpha-cmd" {
t.Fatalf("alpha tool missing or wrong command: %#v", names)
}
if names["inline-only"] != "inline-cmd" {
t.Fatalf("inline-only tool missing: %#v", names)
}
if _, ok := names["stale"]; ok {
t.Fatal("stale in-memory tool should not survive reload")
}
writeTool("beta", "beta-cmd")
if err := ReloadSecurityToolsFromDir(cfg, configPath); err != nil {
t.Fatalf("second reload: %v", err)
}
if len(cfg.Security.Tools) != 3 {
t.Fatalf("expected 3 tools after add, got %d", len(cfg.Security.Tools))
}
foundBeta := false
for _, tool := range cfg.Security.Tools {
if tool.Name == "beta" {
foundBeta = true
break
}
}
if !foundBeta {
t.Fatal("beta tool not found after second reload")
}
}
func TestMergeToolsFromDir_DirOverridesInline(t *testing.T) {
root := t.TempDir()
toolsDir := filepath.Join(root, "tools")
if err := os.MkdirAll(toolsDir, 0755); err != nil {
t.Fatal(err)
}
content := "name: shared\ncommand: dir-cmd\nenabled: true\ndescription: from dir\n"
if err := os.WriteFile(filepath.Join(toolsDir, "shared.yaml"), []byte(content), 0644); err != nil {
t.Fatal(err)
}
inline := []ToolConfig{
{Name: "shared", Command: "inline-cmd", Enabled: true, Description: "from inline"},
}
merged, err := MergeToolsFromDir(toolsDir, inline)
if err != nil {
t.Fatal(err)
}
if len(merged) != 1 {
t.Fatalf("expected 1 tool, got %d", len(merged))
}
if merged[0].Command != "dir-cmd" {
t.Fatalf("dir tool should win, got command %q", merged[0].Command)
}
}
+97
View File
@@ -0,0 +1,97 @@
package config
import "strings"
// VisionConfig 独立视觉模型与 analyze_image 工具参数;enabled 时注册 MCP 工具 analyze_image。
type VisionConfig struct {
Enabled bool `yaml:"enabled" json:"enabled"`
APIKey string `yaml:"api_key,omitempty" json:"api_key,omitempty"`
BaseURL string `yaml:"base_url,omitempty" json:"base_url,omitempty"`
Model string `yaml:"model,omitempty" json:"model,omitempty"`
Provider string `yaml:"provider,omitempty" json:"provider,omitempty"`
TimeoutSeconds int `yaml:"timeout_seconds,omitempty" json:"timeout_seconds,omitempty"`
MaxImageBytes int64 `yaml:"max_image_bytes,omitempty" json:"max_image_bytes,omitempty"`
MaxDimension int `yaml:"max_dimension,omitempty" json:"max_dimension,omitempty"`
JPEGQuality int `yaml:"jpeg_quality,omitempty" json:"jpeg_quality,omitempty"`
MaxPayloadBytes int64 `yaml:"max_payload_bytes,omitempty" json:"max_payload_bytes,omitempty"`
SkipPreprocessBelowBytes int64 `yaml:"skip_preprocess_below_bytes,omitempty" json:"skip_preprocess_below_bytes,omitempty"` // 0=始终压缩;默认 2MB 且长边已<=max_dimension 时原图直传
Detail string `yaml:"detail,omitempty" json:"detail,omitempty"` // low | high | auto
}
func (v VisionConfig) TimeoutSecondsEffective() int {
if v.TimeoutSeconds <= 0 {
return 60
}
return v.TimeoutSeconds
}
func (v VisionConfig) MaxImageBytesEffective() int64 {
if v.MaxImageBytes <= 0 {
return 5 * 1024 * 1024
}
return v.MaxImageBytes
}
func (v VisionConfig) MaxDimensionEffective() int {
if v.MaxDimension <= 0 {
return 2048
}
return v.MaxDimension
}
func (v VisionConfig) JPEGQualityEffective() int {
if v.JPEGQuality <= 0 || v.JPEGQuality > 100 {
return 82
}
return v.JPEGQuality
}
func (v VisionConfig) MaxPayloadBytesEffective() int64 {
if v.MaxPayloadBytes <= 0 {
return 512 * 1024
}
return v.MaxPayloadBytes
}
// SkipPreprocessBelowBytesEffective 低于该字节数且长边<=max_dimension、且<=max_payload 时可原图直传;0 表示始终压缩。
func (v VisionConfig) SkipPreprocessBelowBytesEffective() int64 {
if v.SkipPreprocessBelowBytes < 0 {
return 0
}
return v.SkipPreprocessBelowBytes
}
func (v VisionConfig) DetailEffective() string {
d := strings.ToLower(strings.TrimSpace(v.Detail))
switch d {
case "high", "low", "auto":
return d
default:
return "low"
}
}
// OpenAICfgEffective 合并主 openai 配置与 vision 覆盖项,供 VL ChatModel 使用。
// vision.api_key / base_url / provider 留空或省略时,沿用 mainopenai)对应字段;vision.model 必填(由 Ready 校验)。
func (v VisionConfig) OpenAICfgEffective(main OpenAIConfig) OpenAIConfig {
out := main
if k := strings.TrimSpace(v.APIKey); k != "" {
out.APIKey = k
}
if u := strings.TrimSpace(v.BaseURL); u != "" {
out.BaseURL = u
}
if m := strings.TrimSpace(v.Model); m != "" {
out.Model = m
}
if p := strings.TrimSpace(v.Provider); p != "" {
out.Provider = p
}
out.Reasoning.Mode = "off"
return out
}
// Ready 表示已启用且模型名非空。
func (v VisionConfig) Ready() bool {
return v.Enabled && strings.TrimSpace(v.Model) != ""
}
+55
View File
@@ -0,0 +1,55 @@
package config
import "testing"
func TestVisionConfig_OpenAICfgEffective_fallbackToMain(t *testing.T) {
main := OpenAIConfig{
APIKey: "main-key",
BaseURL: "https://main.example/v1",
Model: "main-model",
Provider: "openai",
}
v := VisionConfig{Model: "qwen-vl-max"}
out := v.OpenAICfgEffective(main)
if out.APIKey != main.APIKey || out.BaseURL != main.BaseURL || out.Provider != main.Provider {
t.Fatalf("expected openai fallback, got key=%q url=%q provider=%q", out.APIKey, out.BaseURL, out.Provider)
}
if out.Model != "qwen-vl-max" {
t.Fatalf("model: %s", out.Model)
}
}
func TestVisionConfig_OpenAICfgEffective(t *testing.T) {
main := OpenAIConfig{
APIKey: "main-key",
BaseURL: "https://main.example/v1",
Model: "main-model",
Provider: "openai",
Reasoning: OpenAIReasoningConfig{Mode: "on"},
}
v := VisionConfig{
Model: "vl-model",
APIKey: "vl-key",
BaseURL: "https://vl.example/v1",
Provider: "claude",
}
out := v.OpenAICfgEffective(main)
if out.APIKey != "vl-key" || out.BaseURL != "https://vl.example/v1" || out.Model != "vl-model" {
t.Fatalf("unexpected merge: %+v", out)
}
if out.Provider != "claude" {
t.Fatalf("provider: %s", out.Provider)
}
if out.Reasoning.Mode != "off" {
t.Fatalf("reasoning should be off for vision, got %s", out.Reasoning.Mode)
}
}
func TestVisionConfig_Ready(t *testing.T) {
if (VisionConfig{Enabled: true, Model: "x"}).Ready() != true {
t.Fatal("expected ready")
}
if (VisionConfig{Enabled: true}).Ready() != false {
t.Fatal("expected not ready without model")
}
}
+195
View File
@@ -0,0 +1,195 @@
package builtin
// 内置工具名称常量
// 所有代码中使用内置工具名称的地方都应该使用这些常量,而不是硬编码字符串
const (
// 漏洞管理工具
ToolRecordVulnerability = "record_vulnerability"
ToolListVulnerabilities = "list_vulnerabilities"
ToolGetVulnerability = "get_vulnerability"
// 资产管理工具
ToolCreateAsset = "create_asset"
ToolGetAsset = "get_asset"
ToolQueryAssets = "query_assets"
ToolUpdateAsset = "update_asset"
ToolDeleteAsset = "delete_asset"
ToolCompleteAssetScan = "complete_asset_scan"
// 项目黑板(事实)工具
ToolUpsertProjectFact = "upsert_project_fact"
ToolGetProjectFact = "get_project_fact"
ToolListProjectFacts = "list_project_facts"
ToolSearchProjectFacts = "search_project_facts"
ToolDeprecateProjectFact = "deprecate_project_fact"
ToolRestoreProjectFact = "restore_project_fact"
// 知识库工具
ToolListKnowledgeRiskTypes = "list_knowledge_risk_types"
ToolSearchKnowledgeBase = "search_knowledge_base"
// 视觉分析(本地图片 → VL 模型 → 文本摘要)
ToolAnalyzeImage = "analyze_image"
// 长耗时工具执行控制(后台 execution 查询/等待/取消)
ToolGetToolExecution = "get_tool_execution"
ToolWaitToolExecution = "wait_tool_execution"
ToolCancelToolExecution = "cancel_tool_execution"
// WebShell 助手工具(AI 在 WebShell 管理 - AI 助手 中使用)
ToolWebshellExec = "webshell_exec"
ToolWebshellFileList = "webshell_file_list"
ToolWebshellFileRead = "webshell_file_read"
ToolWebshellFileWrite = "webshell_file_write"
// WebShell 连接管理工具(用于通过 MCP 管理 webshell 连接)
ToolManageWebshellList = "manage_webshell_list"
ToolManageWebshellAdd = "manage_webshell_add"
ToolManageWebshellUpdate = "manage_webshell_update"
ToolManageWebshellDelete = "manage_webshell_delete"
ToolManageWebshellTest = "manage_webshell_test"
// 批量任务队列(与 Web 端批量任务一致,供模型创建/启停/查询队列)
ToolBatchTaskList = "batch_task_list"
ToolBatchTaskGet = "batch_task_get"
ToolBatchTaskCreate = "batch_task_create"
ToolBatchTaskStart = "batch_task_start"
ToolBatchTaskRerun = "batch_task_rerun"
ToolBatchTaskPause = "batch_task_pause"
ToolBatchTaskDelete = "batch_task_delete"
ToolBatchTaskUpdateMetadata = "batch_task_update_metadata"
ToolBatchTaskUpdateSchedule = "batch_task_update_schedule"
ToolBatchTaskScheduleEnabled = "batch_task_schedule_enabled"
ToolBatchTaskAdd = "batch_task_add_task"
ToolBatchTaskUpdate = "batch_task_update_task"
ToolBatchTaskRemove = "batch_task_remove_task"
// C2 工具集(合并同类项,8 个统一工具)
ToolC2Listener = "c2_listener" // 监听器管理(create/start/stop/list/get/update/delete
ToolC2Session = "c2_session" // 会话管理(list/get/set_sleep/kill/delete
ToolC2Task = "c2_task" // 任务下发(统一 task_type 参数)
ToolC2TaskManage = "c2_task_manage" // 任务管理(get_result/wait/list/cancel
ToolC2Payload = "c2_payload" // Payload 生成(oneliner/build
ToolC2Event = "c2_event" // 事件查询
ToolC2Profile = "c2_profile" // Malleable Profile 管理(list/get/create/update/delete
ToolC2File = "c2_file" // 文件管理(list/get_result
)
// IsBuiltinTool 检查工具名称是否是内置工具
func IsBuiltinTool(toolName string) bool {
switch toolName {
case ToolRecordVulnerability,
ToolListVulnerabilities,
ToolGetVulnerability,
ToolCreateAsset,
ToolGetAsset,
ToolQueryAssets,
ToolUpdateAsset,
ToolDeleteAsset,
ToolCompleteAssetScan,
ToolUpsertProjectFact,
ToolGetProjectFact,
ToolListProjectFacts,
ToolSearchProjectFacts,
ToolDeprecateProjectFact,
ToolRestoreProjectFact,
ToolListKnowledgeRiskTypes,
ToolSearchKnowledgeBase,
ToolAnalyzeImage,
ToolGetToolExecution,
ToolWaitToolExecution,
ToolCancelToolExecution,
ToolWebshellExec,
ToolWebshellFileList,
ToolWebshellFileRead,
ToolWebshellFileWrite,
ToolManageWebshellList,
ToolManageWebshellAdd,
ToolManageWebshellUpdate,
ToolManageWebshellDelete,
ToolManageWebshellTest,
ToolBatchTaskList,
ToolBatchTaskGet,
ToolBatchTaskCreate,
ToolBatchTaskStart,
ToolBatchTaskRerun,
ToolBatchTaskPause,
ToolBatchTaskDelete,
ToolBatchTaskUpdateMetadata,
ToolBatchTaskUpdateSchedule,
ToolBatchTaskScheduleEnabled,
ToolBatchTaskAdd,
ToolBatchTaskUpdate,
ToolBatchTaskRemove,
// C2 工具
ToolC2Listener,
ToolC2Session,
ToolC2Task,
ToolC2TaskManage,
ToolC2Payload,
ToolC2Event,
ToolC2Profile,
ToolC2File:
return true
default:
return false
}
}
// GetAllBuiltinTools 返回所有内置工具名称列表
func GetAllBuiltinTools() []string {
return []string{
ToolRecordVulnerability,
ToolListVulnerabilities,
ToolGetVulnerability,
ToolCreateAsset,
ToolGetAsset,
ToolQueryAssets,
ToolUpdateAsset,
ToolDeleteAsset,
ToolCompleteAssetScan,
ToolUpsertProjectFact,
ToolGetProjectFact,
ToolListProjectFacts,
ToolSearchProjectFacts,
ToolDeprecateProjectFact,
ToolRestoreProjectFact,
ToolListKnowledgeRiskTypes,
ToolSearchKnowledgeBase,
ToolAnalyzeImage,
ToolGetToolExecution,
ToolWaitToolExecution,
ToolCancelToolExecution,
ToolWebshellExec,
ToolWebshellFileList,
ToolWebshellFileRead,
ToolWebshellFileWrite,
ToolManageWebshellList,
ToolManageWebshellAdd,
ToolManageWebshellUpdate,
ToolManageWebshellDelete,
ToolManageWebshellTest,
ToolBatchTaskList,
ToolBatchTaskGet,
ToolBatchTaskCreate,
ToolBatchTaskStart,
ToolBatchTaskRerun,
ToolBatchTaskPause,
ToolBatchTaskDelete,
ToolBatchTaskUpdateMetadata,
ToolBatchTaskUpdateSchedule,
ToolBatchTaskScheduleEnabled,
ToolBatchTaskAdd,
ToolBatchTaskUpdate,
ToolBatchTaskRemove,
// C2 工具
ToolC2Listener,
ToolC2Session,
ToolC2Task,
ToolC2TaskManage,
ToolC2Payload,
ToolC2Event,
ToolC2Profile,
ToolC2File,
}
}
+475
View File
@@ -0,0 +1,475 @@
// Package mcp 外部 MCP 客户端 - 基于官方 go-sdk 实现,保证协议兼容性
package mcp
import (
"context"
"encoding/json"
"fmt"
"net/http"
"os"
"os/exec"
"strings"
"sync"
"time"
"cyberstrike-ai/internal/config"
"github.com/modelcontextprotocol/go-sdk/mcp"
"go.uber.org/zap"
)
const (
clientName = "CyberStrikeAI"
clientVersion = "1.0.0"
)
// sdkClient 基于官方 MCP Go SDK 的外部 MCP 客户端,实现 ExternalMCPClient 接口
type sdkClient struct {
session *mcp.ClientSession
client *mcp.Client
logger *zap.Logger
mu sync.RWMutex
status string // "disconnected", "connecting", "connected", "error"
}
// newSDKClientFromSession 用已连接成功的 session 构造(供 createSDKClient 内部使用)
func newSDKClientFromSession(session *mcp.ClientSession, client *mcp.Client, logger *zap.Logger) *sdkClient {
return &sdkClient{
session: session,
client: client,
logger: logger,
status: "connected",
}
}
// lazySDKClient 延迟连接:Initialize() 时才调用官方 SDK 建立连接,对外实现 ExternalMCPClient
type lazySDKClient struct {
serverCfg config.ExternalMCPServerConfig
logger *zap.Logger
sessionCancel context.CancelFunc
inner ExternalMCPClient // connected SDK client
mu sync.RWMutex
status string
}
func newLazySDKClient(serverCfg config.ExternalMCPServerConfig, logger *zap.Logger) *lazySDKClient {
return &lazySDKClient{
serverCfg: serverCfg,
logger: logger,
status: "connecting",
}
}
func (c *lazySDKClient) setStatus(s string) {
c.mu.Lock()
defer c.mu.Unlock()
c.status = s
}
func (c *lazySDKClient) GetStatus() string {
c.mu.RLock()
defer c.mu.RUnlock()
if c.inner != nil {
return c.inner.GetStatus()
}
return c.status
}
func (c *lazySDKClient) IsConnected() bool {
c.mu.RLock()
inner := c.inner
c.mu.RUnlock()
if inner != nil {
return inner.IsConnected()
}
return false
}
func (c *lazySDKClient) Initialize(ctx context.Context) error {
c.mu.Lock()
if c.inner != nil {
c.mu.Unlock()
return nil
}
c.mu.Unlock()
sessionCtx, sessionCancel := context.WithCancel(context.Background())
type connectResult struct {
inner ExternalMCPClient
err error
}
resultCh := make(chan connectResult)
abandoned := make(chan struct{})
go func() {
inner, err := createSDKClient(sessionCtx, c.serverCfg, c.logger)
select {
case resultCh <- connectResult{inner: inner, err: err}:
case <-abandoned:
if inner != nil {
_ = inner.Close()
}
sessionCancel()
}
}()
var result connectResult
select {
case result = <-resultCh:
case <-ctx.Done():
close(abandoned)
sessionCancel()
c.setStatus("error")
return ctx.Err()
}
if err := ctx.Err(); err != nil {
sessionCancel()
if result.inner != nil {
_ = result.inner.Close()
}
c.setStatus("error")
return err
}
if result.err != nil {
sessionCancel()
c.setStatus("error")
return result.err
}
c.mu.Lock()
if c.inner != nil {
c.mu.Unlock()
sessionCancel()
if result.inner != nil {
_ = result.inner.Close()
}
return nil
}
c.inner = result.inner
c.sessionCancel = sessionCancel
c.mu.Unlock()
c.setStatus("connected")
return nil
}
func (c *lazySDKClient) ListTools(ctx context.Context) ([]Tool, error) {
c.mu.RLock()
inner := c.inner
c.mu.RUnlock()
if inner == nil {
return nil, fmt.Errorf("未连接")
}
return inner.ListTools(ctx)
}
func (c *lazySDKClient) CallTool(ctx context.Context, name string, args map[string]interface{}) (*ToolResult, error) {
c.mu.RLock()
inner := c.inner
c.mu.RUnlock()
if inner == nil {
return nil, fmt.Errorf("未连接")
}
return inner.CallTool(ctx, name, args)
}
func (c *lazySDKClient) Close() error {
c.mu.Lock()
inner := c.inner
sessionCancel := c.sessionCancel
c.inner = nil
c.sessionCancel = nil
c.mu.Unlock()
c.setStatus("disconnected")
if sessionCancel != nil {
sessionCancel()
}
if inner != nil {
return inner.Close()
}
return nil
}
// markDisconnected 在检测到传输层断连时关闭底层 session,避免 IsConnected 仍返回 true。
func (c *lazySDKClient) markDisconnected() {
c.mu.Lock()
inner := c.inner
sessionCancel := c.sessionCancel
c.inner = nil
c.sessionCancel = nil
c.mu.Unlock()
if sessionCancel != nil {
sessionCancel()
}
if inner != nil {
_ = inner.Close()
}
c.setStatus("disconnected")
}
func (c *sdkClient) setStatus(s string) {
c.mu.Lock()
defer c.mu.Unlock()
c.status = s
}
func (c *sdkClient) GetStatus() string {
c.mu.RLock()
defer c.mu.RUnlock()
return c.status
}
func (c *sdkClient) IsConnected() bool {
return c.GetStatus() == "connected"
}
func (c *sdkClient) Initialize(ctx context.Context) error {
// sdkClient 由 createSDKClient 在 Connect 成功后才创建,因此 Initialize 时已经连接
// 此方法仅用于满足 ExternalMCPClient 接口,实际连接在 createSDKClient 中完成
return nil
}
func (c *sdkClient) ListTools(ctx context.Context) ([]Tool, error) {
if c.session == nil {
return nil, fmt.Errorf("未连接")
}
res, err := c.session.ListTools(ctx, nil)
if err != nil {
return nil, err
}
if res == nil {
return nil, nil
}
return sdkToolsToOur(res.Tools), nil
}
func (c *sdkClient) CallTool(ctx context.Context, name string, args map[string]interface{}) (*ToolResult, error) {
if c.session == nil {
return nil, fmt.Errorf("未连接")
}
params := &mcp.CallToolParams{
Name: name,
Arguments: args,
}
res, err := c.session.CallTool(ctx, params)
if err != nil {
return nil, err
}
return sdkCallToolResultToOurs(res), nil
}
func (c *sdkClient) Close() error {
c.setStatus("disconnected")
if c.session != nil {
err := c.session.Close()
c.session = nil
return err
}
return nil
}
// sdkToolsToOur 将 SDK 的 []*mcp.Tool 转为我们的 []Tool
func sdkToolsToOur(tools []*mcp.Tool) []Tool {
if len(tools) == 0 {
return nil
}
out := make([]Tool, 0, len(tools))
for _, t := range tools {
if t == nil {
continue
}
schema := make(map[string]interface{})
if t.InputSchema != nil {
// SDK InputSchema 可能为 *jsonschema.Schema 或 map,统一转为 map
if m, ok := t.InputSchema.(map[string]interface{}); ok {
schema = m
} else {
_ = json.Unmarshal(mustJSON(t.InputSchema), &schema)
}
}
desc := t.Description
shortDesc := desc
if t.Annotations != nil && t.Annotations.Title != "" {
shortDesc = t.Annotations.Title
}
out = append(out, Tool{
Name: t.Name,
Description: desc,
ShortDescription: shortDesc,
InputSchema: schema,
})
}
return out
}
// sdkCallToolResultToOurs 将 SDK 的 *mcp.CallToolResult 转为我们的 *ToolResult
func sdkCallToolResultToOurs(res *mcp.CallToolResult) *ToolResult {
if res == nil {
return &ToolResult{Content: []Content{}}
}
content := sdkContentToOurs(res.Content)
return &ToolResult{
Content: content,
IsError: res.IsError,
}
}
func sdkContentToOurs(list []mcp.Content) []Content {
if len(list) == 0 {
return nil
}
out := make([]Content, 0, len(list))
for _, c := range list {
switch v := c.(type) {
case *mcp.TextContent:
out = append(out, Content{Type: "text", Text: v.Text})
default:
out = append(out, Content{Type: "text", Text: fmt.Sprintf("%v", c)})
}
}
return out
}
func mustJSON(v interface{}) []byte {
b, _ := json.Marshal(v)
return b
}
// createSDKClient 根据配置创建并连接外部 MCP 客户端(使用官方 SDK),返回实现 ExternalMCPClient 的 *sdkClient
// 若连接失败返回 (nil, error)。ctx 用于连接超时与取消。
func createSDKClient(ctx context.Context, serverCfg config.ExternalMCPServerConfig, logger *zap.Logger) (ExternalMCPClient, error) {
timeout := time.Duration(serverCfg.Timeout) * time.Second
if timeout <= 0 {
timeout = 30 * time.Second
}
transport := serverCfg.GetTransportType()
if transport == "" {
return nil, fmt.Errorf("配置缺少 command 或 url,且未指定 type/transport")
}
// 构造 ClientOptionsKeepAlive 心跳
var clientOpts *mcp.ClientOptions
if serverCfg.KeepAlive > 0 {
clientOpts = &mcp.ClientOptions{
KeepAlive: time.Duration(serverCfg.KeepAlive) * time.Second,
}
}
client := mcp.NewClient(&mcp.Implementation{
Name: clientName,
Version: clientVersion,
}, clientOpts)
var t mcp.Transport
switch transport {
case "stdio":
if serverCfg.Command == "" {
return nil, fmt.Errorf("stdio 模式需要配置 command")
}
// 必须用 exec.Command 而非 CommandContextdoConnect 返回后 ctx 会被 cancel
// 若用 CommandContext(ctx) 会立刻杀掉子进程,导致 ListTools 等后续请求失败、显示 0 工具
cmd := exec.Command(serverCfg.Command, serverCfg.Args...)
if len(serverCfg.Env) > 0 {
cmd.Env = append(cmd.Env, envMapToSlice(serverCfg.Env)...)
}
ct := &mcp.CommandTransport{Command: cmd}
if serverCfg.TerminateDuration > 0 {
ct.TerminateDuration = time.Duration(serverCfg.TerminateDuration) * time.Second
}
t = ct
case "sse":
if serverCfg.URL == "" {
return nil, fmt.Errorf("sse 模式需要配置 url")
}
// SSE 是长连接(GET 流持续打开),不能设置 http.Client.Timeout(会在超时后杀掉整个连接导致 EOF)。
// 超时由每次 ListTools/CallTool 的 context 单独控制。
httpClient := httpClientForLongLived(serverCfg.Headers)
t = &mcp.SSEClientTransport{
Endpoint: serverCfg.URL,
HTTPClient: httpClient,
}
case "http":
if serverCfg.URL == "" {
return nil, fmt.Errorf("http 模式需要配置 url")
}
httpClient := httpClientWithTimeoutAndHeaders(timeout, serverCfg.Headers)
st := &mcp.StreamableClientTransport{
Endpoint: serverCfg.URL,
HTTPClient: httpClient,
}
if serverCfg.MaxRetries > 0 {
st.MaxRetries = serverCfg.MaxRetries
}
t = st
default:
return nil, fmt.Errorf("不支持的传输模式: %s(支持: stdio, sse, http", transport)
}
session, err := client.Connect(ctx, t, nil)
if err != nil {
return nil, fmt.Errorf("连接失败: %w", err)
}
return newSDKClientFromSession(session, client, logger), nil
}
func envMapToSlice(env map[string]string) []string {
m := make(map[string]string)
for _, s := range os.Environ() {
if i := strings.IndexByte(s, '='); i > 0 {
m[s[:i]] = s[i+1:]
}
}
for k, v := range env {
m[k] = v
}
out := make([]string, 0, len(m))
for k, v := range m {
out = append(out, k+"="+v)
}
return out
}
func httpClientWithTimeoutAndHeaders(timeout time.Duration, headers map[string]string) *http.Client {
transport := http.DefaultTransport
if len(headers) > 0 {
transport = &headerRoundTripper{
headers: headers,
base: http.DefaultTransport,
}
}
return &http.Client{
Timeout: timeout,
Transport: transport,
}
}
// httpClientForLongLived 创建不设超时的 HTTP 客户端,用于 SSE 等长连接传输。
// SSE 的 GET 流会持续打开,http.Client.Timeout 会在超时后强制关闭连接导致 EOF。
// 超时由调用方通过 context 控制。
func httpClientForLongLived(headers map[string]string) *http.Client {
transport := http.DefaultTransport
if len(headers) > 0 {
transport = &headerRoundTripper{
headers: headers,
base: http.DefaultTransport,
}
}
return &http.Client{
Transport: transport,
// 不设 TimeoutSSE 长连接的超时由 per-request context 控制
}
}
type headerRoundTripper struct {
headers map[string]string
base http.RoundTripper
}
func (h *headerRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
for k, v := range h.headers {
req.Header.Set(k, v)
}
return h.base.RoundTrip(req)
}
+192
View File
@@ -0,0 +1,192 @@
package mcp
import (
"context"
"errors"
"io"
"strings"
"time"
"go.uber.org/zap"
)
const (
// externalReconnectMinInterval 两次自动重连之间的最短间隔
externalReconnectMinInterval = 30 * time.Second
// externalReconnectMaxBackoff 指数退避上限
externalReconnectMaxBackoff = 5 * time.Minute
)
// isConnectionDeadError 判断错误是否表示底层传输已断开(而非调用方主动取消或超时)。
func isConnectionDeadError(err error) bool {
if err == nil {
return false
}
if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
return false
}
if errors.Is(err, io.EOF) {
return true
}
s := strings.ToLower(err.Error())
return strings.Contains(s, "eof") ||
strings.Contains(s, "client is closing") ||
strings.Contains(s, "connection closed") ||
strings.Contains(s, "connection reset") ||
strings.Contains(s, "broken pipe")
}
// handleConnectionDead 在 ListTools/CallTool 等操作失败且判定为断连时,标记客户端并调度重连。
func (m *ExternalMCPManager) handleConnectionDead(name string, client ExternalMCPClient, err error) {
if !isConnectionDeadError(err) {
return
}
m.logger.Warn("检测到外部MCP连接已断开,将尝试自动重连",
zap.String("name", name),
zap.Error(err),
)
m.markClientDisconnected(name, client, err)
m.scheduleReconnect(name)
}
func (m *ExternalMCPManager) markClientDisconnected(name string, client ExternalMCPClient, err error) {
if lazy, ok := client.(*lazySDKClient); ok {
lazy.markDisconnected()
}
m.mu.Lock()
if err != nil {
m.errors[name] = "连接已断开: " + err.Error()
}
m.mu.Unlock()
m.toolCountsMu.Lock()
m.toolCounts[name] = 0
m.toolCountsMu.Unlock()
}
func (m *ExternalMCPManager) onClientConnected(name string) {
m.clearReconnectState(name)
}
func (m *ExternalMCPManager) clearReconnectState(name string) {
m.reconnectMu.Lock()
delete(m.reconnectAttempts, name)
delete(m.reconnectLastTry, name)
delete(m.reconnecting, name)
m.reconnectMu.Unlock()
}
func (m *ExternalMCPManager) reconnectBackoff(attempts int) time.Duration {
if attempts <= 0 {
return 0
}
d := externalReconnectMinInterval
for i := 1; i < attempts && d < externalReconnectMaxBackoff; i++ {
d *= 2
}
if d > externalReconnectMaxBackoff {
return externalReconnectMaxBackoff
}
return d
}
func (m *ExternalMCPManager) scheduleReconnect(name string) {
m.mu.RLock()
cfg, exists := m.configs[name]
enabled := exists && m.isEnabled(cfg)
m.mu.RUnlock()
if !enabled {
return
}
go m.tryReconnect(name)
}
func (m *ExternalMCPManager) tryReconnect(name string) {
m.reconnectMu.Lock()
if m.reconnecting[name] {
m.reconnectMu.Unlock()
return
}
attempts := m.reconnectAttempts[name]
if wait := m.reconnectBackoff(attempts); wait > 0 {
if last, ok := m.reconnectLastTry[name]; ok {
if elapsed := time.Since(last); elapsed < wait {
remaining := wait - elapsed
m.reconnectMu.Unlock()
m.scheduleReconnectAfter(name, remaining)
return
}
}
}
m.reconnecting[name] = true
m.reconnectMu.Unlock()
defer func() {
m.reconnectMu.Lock()
delete(m.reconnecting, name)
m.reconnectMu.Unlock()
}()
m.mu.RLock()
cfg, exists := m.configs[name]
enabled := exists && m.isEnabled(cfg)
client, hasClient := m.clients[name]
connecting := hasClient && client.GetStatus() == "connecting"
m.mu.RUnlock()
if !enabled {
m.logger.Debug("跳过自动重连(外部MCP已停用)", zap.String("name", name))
return
}
if connecting {
m.logger.Debug("跳过自动重连(连接正在进行中)", zap.String("name", name))
return
}
m.reconnectMu.Lock()
m.reconnectLastTry[name] = time.Now()
m.reconnectAttempts[name] = attempts + 1
attemptNum := m.reconnectAttempts[name]
m.reconnectMu.Unlock()
m.logger.Info("正在自动重连外部MCP",
zap.String("name", name),
zap.Int("attempt", attemptNum),
)
if err := m.startClient(name, true); err != nil {
m.logger.Warn("自动重连外部MCP失败",
zap.String("name", name),
zap.Error(err),
)
}
}
// scheduleReconnectAfterFailure 在自动重连失败后,按当前退避间隔预约下一次重试。
func (m *ExternalMCPManager) scheduleReconnectAfterFailure(name string) {
m.mu.RLock()
cfg, exists := m.configs[name]
enabled := exists && m.isEnabled(cfg)
m.mu.RUnlock()
if !enabled {
return
}
m.reconnectMu.Lock()
wait := m.reconnectBackoff(m.reconnectAttempts[name])
m.reconnectMu.Unlock()
m.logger.Info("自动重连失败,将按退避间隔再次尝试",
zap.String("name", name),
zap.Duration("after", wait),
)
m.scheduleReconnectAfter(name, wait)
}
// scheduleReconnectAfter 在 delay 后触发 tryReconnectdelay<=0 时立即执行)。
func (m *ExternalMCPManager) scheduleReconnectAfter(name string, delay time.Duration) {
if delay <= 0 {
go m.tryReconnect(name)
return
}
time.AfterFunc(delay, func() {
m.tryReconnect(name)
})
}
+215
View File
@@ -0,0 +1,215 @@
package mcp
import (
"context"
"errors"
"fmt"
"io"
"testing"
"time"
"cyberstrike-ai/internal/config"
"go.uber.org/zap"
)
func TestIsConnectionDeadError(t *testing.T) {
t.Parallel()
cases := []struct {
name string
err error
want bool
}{
{"nil", nil, false},
{"eof", io.EOF, true},
{"wrapped eof", fmt.Errorf("connection closed: %w", io.EOF), true},
{"client closing", errors.New(`calling "tools/list": client is closing: EOF`), true},
{"connection reset", errors.New("read tcp: connection reset by peer"), true},
{"canceled", context.Canceled, false},
{"deadline", context.DeadlineExceeded, false},
{"other", errors.New("invalid params"), false},
}
for _, tc := range cases {
tc := tc
t.Run(tc.name, func(t *testing.T) {
t.Parallel()
if got := isConnectionDeadError(tc.err); got != tc.want {
t.Fatalf("isConnectionDeadError(%v) = %v, want %v", tc.err, got, tc.want)
}
})
}
}
func TestLazySDKClient_MarkDisconnected(t *testing.T) {
c := &lazySDKClient{status: "connected"}
c.inner = &sdkClient{status: "connected"}
c.markDisconnected()
if c.IsConnected() {
t.Fatal("expected disconnected after markDisconnected")
}
if c.GetStatus() != "disconnected" {
t.Fatalf("expected status disconnected, got %s", c.GetStatus())
}
}
func TestHandleConnectionDead_MarksLazyClientDisconnected(t *testing.T) {
logger := zap.NewNop()
m := NewExternalMCPManager(logger)
name := "dead-mcp"
cfg := config.ExternalMCPServerConfig{
Type: "http",
URL: "http://example.com/mcp",
ExternalMCPEnable: true,
}
m.mu.Lock()
m.configs[name] = cfg
client := newLazySDKClient(cfg, logger)
client.inner = &sdkClient{status: "connected"}
client.status = "connected"
m.clients[name] = client
m.mu.Unlock()
deadErr := errors.New(`connection closed: calling "tools/list": client is closing: EOF`)
m.handleConnectionDead(name, client, deadErr)
if client.IsConnected() {
t.Fatal("expected disconnected after handleConnectionDead")
}
if m.GetError(name) == "" {
t.Fatal("expected error message to be recorded")
}
counts := m.GetToolCounts()
if counts[name] != 0 {
t.Fatalf("expected tool count 0 after disconnect, got %d", counts[name])
}
}
func TestReconnectBackoff(t *testing.T) {
t.Parallel()
if d := (&ExternalMCPManager{}).reconnectBackoff(0); d != 0 {
t.Fatalf("attempt 0: got %v", d)
}
if d := (&ExternalMCPManager{}).reconnectBackoff(1); d != externalReconnectMinInterval {
t.Fatalf("attempt 1: got %v", d)
}
if d := (&ExternalMCPManager{}).reconnectBackoff(10); d != externalReconnectMaxBackoff {
t.Fatalf("attempt 10: got %v, want cap %v", d, externalReconnectMaxBackoff)
}
}
func TestTryReconnect_RateLimited(t *testing.T) {
logger := zap.NewNop()
m := NewExternalMCPManager(logger)
name := "rate-limited"
m.reconnectMu.Lock()
m.reconnectLastTry[name] = time.Now()
m.reconnectAttempts[name] = 2
m.reconnectMu.Unlock()
m.tryReconnect(name)
m.reconnectMu.Lock()
attempts := m.reconnectAttempts[name]
m.reconnectMu.Unlock()
if attempts != 2 {
t.Fatalf("rate limited reconnect should not increment attempts, got %d", attempts)
}
}
func TestTryReconnect_SkipsWhenDisabled(t *testing.T) {
logger := zap.NewNop()
m := NewExternalMCPManager(logger)
name := "disabled-mcp"
m.mu.Lock()
m.configs[name] = config.ExternalMCPServerConfig{
Type: "http",
URL: "http://example.com/mcp",
ExternalMCPEnable: false,
}
m.mu.Unlock()
m.tryReconnect(name)
m.reconnectMu.Lock()
attempts := m.reconnectAttempts[name]
m.reconnectMu.Unlock()
if attempts != 0 {
t.Fatalf("disabled MCP should not increment reconnect attempts, got %d", attempts)
}
}
func TestTryReconnect_SkipsWhenConnecting(t *testing.T) {
logger := zap.NewNop()
m := NewExternalMCPManager(logger)
name := "connecting-mcp"
cfg := config.ExternalMCPServerConfig{
Type: "http",
URL: "http://example.com/mcp",
ExternalMCPEnable: true,
}
client := newLazySDKClient(cfg, logger)
client.setStatus("connecting")
m.mu.Lock()
m.configs[name] = cfg
m.clients[name] = client
m.mu.Unlock()
m.tryReconnect(name)
m.reconnectMu.Lock()
attempts := m.reconnectAttempts[name]
m.reconnectMu.Unlock()
if attempts != 0 {
t.Fatalf("connecting MCP should not increment reconnect attempts, got %d", attempts)
}
}
func TestStartClientAutoReconnect_SkipsWhenDisabled(t *testing.T) {
logger := zap.NewNop()
m := NewExternalMCPManager(logger)
m.stopRefresh = make(chan struct{})
name := "stopped"
m.mu.Lock()
m.configs[name] = config.ExternalMCPServerConfig{
Type: "http",
URL: "http://example.com/mcp",
ExternalMCPEnable: false,
}
m.mu.Unlock()
if err := m.startClient(name, true); err != nil {
t.Fatalf("startClient: %v", err)
}
m.mu.RLock()
cfg := m.configs[name]
_, hasClient := m.clients[name]
m.mu.RUnlock()
if cfg.ExternalMCPEnable {
t.Fatal("auto reconnect should not enable stopped MCP")
}
if hasClient {
t.Fatal("auto reconnect should not create client when disabled")
}
}
func TestOnClientConnected_ClearsReconnectState(t *testing.T) {
m := &ExternalMCPManager{
reconnectAttempts: map[string]int{"x": 3},
reconnectLastTry: map[string]time.Time{"x": time.Now()},
reconnecting: map[string]bool{"x": true},
}
m.onClientConnected("x")
m.reconnectMu.Lock()
defer m.reconnectMu.Unlock()
if len(m.reconnectAttempts) != 0 || len(m.reconnectLastTry) != 0 || len(m.reconnecting) != 0 {
t.Fatal("expected reconnect state cleared")
}
}
+296
View File
@@ -0,0 +1,296 @@
package mcp
import (
"context"
"encoding/json"
"errors"
"fmt"
"strconv"
"strings"
"time"
"cyberstrike-ai/internal/mcp/builtin"
)
const (
defaultExecutionWaitTimeout = 60 * time.Second
maxExecutionWaitTimeout = 10 * time.Minute
defaultPartialPreviewBytes = 4096
maxPartialPreviewBytes = 64 * 1024
)
// RegisterExecutionControlTools exposes execution handle operations to Eino as
// ordinary MCP tools. This keeps the agent loop native: the model calls a tool,
// receives a bounded result, and may call wait_tool_execution again if needed.
func RegisterExecutionControlTools(server *Server, external *ExternalMCPManager) {
if server == nil {
return
}
server.RegisterTool(Tool{
Name: builtin.ToolGetToolExecution,
Description: "查询后台工具 execution 的当前状态、结果和错误。用于外部 MCP 工具等待超时后,凭 execution_id 继续查看进度。",
ShortDescription: "查询后台工具执行状态",
InputSchema: map[string]interface{}{
"type": "object",
"properties": map[string]interface{}{
"execution_id": map[string]interface{}{"type": "string", "description": "工具执行 ID"},
"include_partial_output": map[string]interface{}{"type": "boolean", "description": "是否返回运行中已产生输出的尾部预览,默认 true"},
"partial_output_max_bytes": map[string]interface{}{"type": "number", "description": "partial_output 最多返回字节数,默认 4096,最大 65536"},
},
"required": []string{"execution_id"},
},
}, func(ctx context.Context, args map[string]interface{}) (*ToolResult, error) {
id := stringArg(args, "execution_id")
if id == "" {
return textToolResult("execution_id 必填", true), nil
}
exec := lookupToolExecution(server, external, id)
if exec == nil {
return textToolResult("未找到该 execution_id: "+id, true), nil
}
return textToolResult(formatExecutionForModel(exec, executionFormatOptionsFromArgs(args)), false), nil
})
server.RegisterTool(Tool{
Name: builtin.ToolWaitToolExecution,
Description: "继续等待一个后台工具 execution 完成。每次等待都有 timeout_seconds 上限;若仍未完成,会返回当前状态,模型可稍后再次调用。",
ShortDescription: "有界等待后台工具执行",
InputSchema: map[string]interface{}{
"type": "object",
"properties": map[string]interface{}{
"execution_id": map[string]interface{}{"type": "string", "description": "工具执行 ID"},
"timeout_seconds": map[string]interface{}{"type": "number", "description": "本次最多等待秒数,默认 60,最大 600"},
"include_partial_output": map[string]interface{}{"type": "boolean", "description": "是否返回运行中已产生输出的尾部预览,默认 true"},
"partial_output_max_bytes": map[string]interface{}{"type": "number", "description": "partial_output 最多返回字节数,默认 4096,最大 65536"},
},
"required": []string{"execution_id"},
},
}, func(ctx context.Context, args map[string]interface{}) (*ToolResult, error) {
id := stringArg(args, "execution_id")
if id == "" {
return textToolResult("execution_id 必填", true), nil
}
wait := durationSecondsArg(args, "timeout_seconds", defaultExecutionWaitTimeout, maxExecutionWaitTimeout)
snap, err := waitToolExecutionSnapshot(ctx, server, external, id, wait)
if err != nil && !errors.Is(err, ErrExecutionWaitTimeout) {
return textToolResult("等待 execution 失败: "+err.Error(), true), nil
}
if snap == nil || snap.Execution == nil {
return textToolResult("未找到该 execution_id: "+id, true), nil
}
body := formatExecutionForModel(snap.Execution, executionFormatOptionsFromArgs(args))
if errors.Is(err, ErrExecutionWaitTimeout) {
body += "\n\n本次等待已到达 timeout_seconds,上述 execution 仍未完成。可继续等待、取消,或采用其他步骤。"
}
return textToolResult(body, false), nil
})
server.RegisterTool(Tool{
Name: builtin.ToolCancelToolExecution,
Description: "取消一个后台工具 execution。用于外部 MCP 工具长时间运行、误调用或用户要求停止时。",
ShortDescription: "取消后台工具执行",
InputSchema: map[string]interface{}{
"type": "object",
"properties": map[string]interface{}{
"execution_id": map[string]interface{}{"type": "string", "description": "工具执行 ID"},
"reason": map[string]interface{}{"type": "string", "description": "取消原因,可选,会写入终止说明"},
},
"required": []string{"execution_id"},
},
}, func(ctx context.Context, args map[string]interface{}) (*ToolResult, error) {
id := stringArg(args, "execution_id")
if id == "" {
return textToolResult("execution_id 必填", true), nil
}
reason := stringArg(args, "reason")
if server.CancelToolExecutionWithNote(id, reason) {
return textToolResult("已请求取消内部工具 execution: "+id, false), nil
}
if external != nil && external.CancelToolExecutionWithNote(id, reason) {
return textToolResult("已请求取消外部 MCP execution: "+id, false), nil
}
return textToolResult("未找到进行中的 execution,或该 execution 已结束: "+id, true), nil
})
}
func waitToolExecutionSnapshot(ctx context.Context, server *Server, external *ExternalMCPManager, id string, wait time.Duration) (*ExecutionSnapshot, error) {
if server != nil && server.executionService != nil && server.executionService.getEntry(id) != nil {
return server.executionService.Wait(ctx, id, wait)
}
if external != nil && external.executionService != nil && external.executionService.getEntry(id) != nil {
return external.executionService.Wait(ctx, id, wait)
}
if server != nil && server.executionService != nil {
if snap, err := server.executionService.Get(id); err == nil {
return snap, nil
}
}
if external != nil && external.executionService != nil {
return external.executionService.Get(id)
}
exec := lookupToolExecution(server, external, id)
if exec == nil {
return nil, fmt.Errorf("execution not found: %s", id)
}
return &ExecutionSnapshot{Execution: exec}, nil
}
func lookupToolExecution(server *Server, external *ExternalMCPManager, id string) *ToolExecution {
if server != nil {
if exec, ok := server.GetExecution(id); ok && exec != nil {
return exec
}
}
if external != nil {
if exec, ok := external.GetExecution(id); ok && exec != nil {
return exec
}
}
return nil
}
type executionFormatOptions struct {
includePartialOutput bool
partialMaxBytes int
}
func executionFormatOptionsFromArgs(args map[string]interface{}) executionFormatOptions {
includePartial := true
if raw, ok := args["include_partial_output"]; ok {
if b, ok := raw.(bool); ok {
includePartial = b
} else if s := strings.TrimSpace(fmt.Sprint(raw)); s != "" {
includePartial = strings.EqualFold(s, "true") || s == "1" || strings.EqualFold(s, "yes")
}
}
maxBytes := intArg(args, "partial_output_max_bytes", defaultPartialPreviewBytes, maxPartialPreviewBytes)
return executionFormatOptions{includePartialOutput: includePartial, partialMaxBytes: maxBytes}
}
func formatExecutionForModel(exec *ToolExecution, opts executionFormatOptions) string {
if exec == nil {
return "execution: null"
}
payload := map[string]interface{}{
"execution_id": exec.ID,
"tool": exec.ToolName,
"status": exec.Status,
"started_at": exec.StartTime.Format(time.RFC3339),
}
if exec.EndTime != nil {
payload["ended_at"] = exec.EndTime.Format(time.RFC3339)
}
if exec.Duration > 0 {
payload["duration"] = exec.Duration.String()
}
if exec.Error != "" {
payload["error"] = exec.Error
}
if exec.Result != nil {
payload["result"] = ToolResultPlainText(exec.Result)
payload["is_error"] = exec.Result.IsError
}
if opts.includePartialOutput && exec.PartialOutput != "" {
partial := tailStringBytes(exec.PartialOutput, opts.partialMaxBytes)
payload["partial_output"] = partial
payload["partial_output_bytes"] = exec.PartialOutputBytes
payload["partial_output_truncated"] = exec.PartialOutputTruncated || len([]byte(partial)) < len([]byte(exec.PartialOutput))
if exec.PartialOutputUpdatedAt != nil {
payload["partial_output_updated_at"] = exec.PartialOutputUpdatedAt.Format(time.RFC3339)
}
}
b, err := json.MarshalIndent(payload, "", " ")
if err != nil {
return fmt.Sprintf("execution_id: %s\nstatus: %s\nerror: %s", exec.ID, exec.Status, exec.Error)
}
return string(b)
}
func tailStringBytes(s string, maxBytes int) string {
if maxBytes <= 0 {
maxBytes = defaultPartialPreviewBytes
}
b := []byte(s)
if len(b) <= maxBytes {
return s
}
return string(b[len(b)-maxBytes:])
}
func textToolResult(text string, isErr bool) *ToolResult {
return &ToolResult{Content: []Content{{Type: "text", Text: text}}, IsError: isErr}
}
func stringArg(args map[string]interface{}, key string) string {
if args == nil {
return ""
}
raw, ok := args[key]
if !ok || raw == nil {
return ""
}
switch v := raw.(type) {
case string:
return strings.TrimSpace(v)
default:
return strings.TrimSpace(fmt.Sprint(v))
}
}
func durationSecondsArg(args map[string]interface{}, key string, def, max time.Duration) time.Duration {
if args == nil {
return def
}
var seconds float64
switch v := args[key].(type) {
case int:
seconds = float64(v)
case int64:
seconds = float64(v)
case float64:
seconds = v
case json.Number:
f, _ := v.Float64()
seconds = f
case string:
f, _ := strconv.ParseFloat(strings.TrimSpace(v), 64)
seconds = f
}
if seconds <= 0 {
return def
}
d := time.Duration(seconds * float64(time.Second))
if max > 0 && d > max {
return max
}
return d
}
func intArg(args map[string]interface{}, key string, def, max int) int {
if args == nil {
return def
}
var n int
switch v := args[key].(type) {
case int:
n = v
case int64:
n = int(v)
case float64:
n = int(v)
case json.Number:
i, _ := v.Int64()
n = int(i)
case string:
i, _ := strconv.Atoi(strings.TrimSpace(v))
n = i
}
if n <= 0 {
return def
}
if max > 0 && n > max {
return max
}
return n
}
+625
View File
@@ -0,0 +1,625 @@
package mcp
import (
"context"
"errors"
"fmt"
"strings"
"sync"
"time"
"cyberstrike-ai/internal/authctx"
"github.com/google/uuid"
"go.uber.org/zap"
)
const (
ToolExecutionStatusQueued = "queued"
ToolExecutionStatusRunning = "running"
ToolExecutionStatusCompleted = "completed"
ToolExecutionStatusFailed = "failed"
ToolExecutionStatusCancelled = "cancelled"
ToolExecutionStatusHardTimeout = "hard_timeout"
ToolExecutionStatusOrphaned = "orphaned"
)
var ErrExecutionWaitTimeout = errors.New("tool execution wait timeout")
// ExecutionRunFunc is the blocking operation owned by a worker.
type ExecutionRunFunc func(context.Context) (*ToolResult, error)
type ExecutionPreRunFunc func(context.Context, *ToolExecution) (func(), error)
// ExecutionDoneFunc observes the final persisted state. It is invoked once,
// including for late completions after an agent has stopped waiting.
type ExecutionDoneFunc func(*ToolExecution)
type ExecutionRequest struct {
ID string
ToolName string
Arguments map[string]interface{}
ConversationID string
OwnerUserID string
HardTimeout time.Duration
PreRun ExecutionPreRunFunc
Run ExecutionRunFunc
OnDone ExecutionDoneFunc
}
type ExecutionHandle struct {
ID string
}
type ExecutionSnapshot struct {
Execution *ToolExecution
}
type executionEntry struct {
exec *ToolExecution
cancel context.CancelFunc
done chan struct{}
preRun ExecutionPreRunFunc
run ExecutionRunFunc
result *ToolResult
err error
}
// ExecutionService keeps Eino-facing tool calls synchronous while moving the
// untrusted blocking work into cancellable workers with explicit execution IDs.
type ExecutionService struct {
storage MonitorStorage
logger *zap.Logger
mu sync.Mutex
entries map[string]*executionEntry
abortUserNotes map[string]string
maxInMemory int
resultMaxBytes int
spillRootDir string
}
func NewExecutionService(storage MonitorStorage, logger *zap.Logger) *ExecutionService {
if logger == nil {
logger = zap.NewNop()
}
return &ExecutionService{
storage: storage,
logger: logger,
entries: make(map[string]*executionEntry),
abortUserNotes: make(map[string]string),
maxInMemory: 1000,
resultMaxBytes: DefaultToolResultMaxBytes,
}
}
func (s *ExecutionService) ConfigureToolResultMaxBytes(maxBytes int) {
if s == nil {
return
}
s.mu.Lock()
defer s.mu.Unlock()
s.resultMaxBytes = maxBytes
}
// ConfigureToolResultSpillRoot sets the reduction-compatible root used when
// oversized tool results are spilled to local files (empty → tmp/reduction).
func (s *ExecutionService) ConfigureToolResultSpillRoot(rootDir string) {
if s == nil {
return
}
s.mu.Lock()
defer s.mu.Unlock()
s.spillRootDir = strings.TrimSpace(rootDir)
}
func (s *ExecutionService) Submit(ctx context.Context, req ExecutionRequest) (*ExecutionHandle, error) {
if s == nil {
return nil, fmt.Errorf("execution service is nil")
}
if req.Run == nil {
return nil, fmt.Errorf("execution run func is nil")
}
id := strings.TrimSpace(req.ID)
if id == "" {
id = uuid.New().String()
}
start := time.Now()
exec := &ToolExecution{
ID: id,
ToolName: strings.TrimSpace(req.ToolName),
Arguments: cloneArgsMap(req.Arguments),
Status: ToolExecutionStatusQueued,
StartTime: start,
ConversationID: strings.TrimSpace(req.ConversationID),
OwnerUserID: strings.TrimSpace(req.OwnerUserID),
}
if exec.ConversationID == "" {
exec.ConversationID = MCPConversationIDFromContext(ctx)
}
if exec.OwnerUserID == "" {
if principal, ok := authctx.PrincipalFromContext(ctx); ok {
exec.OwnerUserID = principal.UserID
}
}
runCtx := detachedExecutionContext(ctx)
var cancel context.CancelFunc
if req.HardTimeout > 0 {
runCtx, cancel = context.WithTimeout(runCtx, req.HardTimeout)
} else {
runCtx, cancel = context.WithCancel(runCtx)
}
entry := &executionEntry{exec: exec, cancel: cancel, done: make(chan struct{}), preRun: req.PreRun, run: req.Run}
s.mu.Lock()
if _, exists := s.entries[id]; exists {
s.mu.Unlock()
cancel()
return nil, fmt.Errorf("execution already exists: %s", id)
}
s.entries[id] = entry
s.cleanupOldEntriesLocked()
s.mu.Unlock()
if s.storage != nil {
if err := s.storage.SaveToolExecution(exec); err != nil {
s.logger.Warn("保存执行记录到数据库失败", zap.Error(err), zap.String("executionId", id))
}
}
notifyToolRunBegin(ctx, id)
go s.runWorker(runCtx, entry, req.OnDone)
return &ExecutionHandle{ID: id}, nil
}
func (s *ExecutionService) runWorker(ctx context.Context, entry *executionEntry, onDone ExecutionDoneFunc) {
id := entry.exec.ID
ctx = WithMCPExecutionID(ctx, id)
if conv := strings.TrimSpace(entry.exec.ConversationID); conv != "" {
ctx = WithMCPConversationID(ctx, conv)
}
var release func()
defer func() {
if release != nil {
release()
}
entry.cancel()
notifyToolRunEnd(ctx, id)
close(entry.done)
}()
if entry.preRun != nil {
var preErr error
release, preErr = entry.preRun(ctx, cloneToolExecution(entry.exec))
if preErr != nil {
s.finishEntry(ctx, entry, nil, preErr, onDone)
return
}
}
s.markEntryRunning(entry)
result, err := entryResultRecover(ctx, entry.exec.ToolName, s.logger, func() (*ToolResult, error) {
return nilSafeRun(ctx, entry)
})
s.finishEntry(ctx, entry, result, err, onDone)
}
func (s *ExecutionService) markEntryRunning(entry *executionEntry) {
if s == nil || entry == nil || entry.exec == nil {
return
}
s.mu.Lock()
if !isExecutionTerminal(entry.exec.Status) {
entry.exec.Status = ToolExecutionStatusRunning
}
runningExec := cloneToolExecution(entry.exec)
s.mu.Unlock()
if s.storage != nil {
if err := s.storage.SaveToolExecution(runningExec); err != nil {
s.logger.Warn("保存执行记录到数据库失败", zap.Error(err), zap.String("executionId", runningExec.ID))
}
}
}
func (s *ExecutionService) finishEntry(ctx context.Context, entry *executionEntry, result *ToolResult, err error, onDone ExecutionDoneFunc) {
id := entry.exec.ID
cancelledWithUserNote := s.applyAbortUserNoteToCancelledToolResult(id, &result, &err)
now := time.Now()
s.mu.Lock()
spill := ToolResultSpillConfig{
RootDir: s.spillRootDir,
ConversationID: entry.exec.ConversationID,
ExecutionID: id,
}
if ctx != nil {
if pid := MCPProjectIDFromContext(ctx); pid != "" {
spill.ProjectID = pid
}
if conv := MCPConversationIDFromContext(ctx); conv != "" {
spill.ConversationID = conv
}
}
result = NormalizeToolResultForStorageWithSpill(result, s.resultMaxBytes, spill)
entry.result = result
entry.err = err
entry.exec.EndTime = &now
entry.exec.Duration = now.Sub(entry.exec.StartTime)
if err != nil {
switch {
case errors.Is(err, context.DeadlineExceeded):
entry.exec.Status = ToolExecutionStatusHardTimeout
entry.exec.Error = "工具执行超过硬超时限制"
case errors.Is(err, context.Canceled):
entry.exec.Status = ToolExecutionStatusCancelled
entry.exec.Error = "已手动终止或任务已取消"
default:
entry.exec.Status = ToolExecutionStatusFailed
entry.exec.Error = err.Error()
}
} else if result != nil && result.IsError {
if cancelledWithUserNote {
entry.exec.Status = ToolExecutionStatusCancelled
entry.exec.Error = ""
} else if isBackgroundWaitToolResult(result) {
entry.exec.Status = ToolExecutionStatusCompleted
entry.exec.Error = ""
} else {
entry.exec.Status = ToolExecutionStatusFailed
entry.exec.Error = firstToolResultText(result, "工具执行返回错误结果")
}
entry.exec.Result = result
} else {
entry.exec.Status = ToolExecutionStatusCompleted
if result == nil {
result = &ToolResult{Content: []Content{{Type: "text", Text: "工具执行完成,但未返回结果"}}}
entry.result = result
}
entry.exec.Result = result
}
finalExec := cloneToolExecution(entry.exec)
s.mu.Unlock()
if s.storage != nil {
if saveErr := s.storage.SaveToolExecution(finalExec); saveErr != nil {
s.logger.Warn("保存执行记录到数据库失败", zap.Error(saveErr), zap.String("executionId", id))
}
}
if onDone != nil {
onDone(finalExec)
}
}
func nilSafeRun(ctx context.Context, entry *executionEntry) (*ToolResult, error) {
if entry == nil {
return nil, fmt.Errorf("execution entry is nil")
}
if entry.run == nil {
return nil, fmt.Errorf("execution run func not wired")
}
return entry.run(ctx)
}
func entryResultRecover(ctx context.Context, toolName string, logger *zap.Logger, fn func() (*ToolResult, error)) (res *ToolResult, err error) {
defer func() {
if r := recover(); r != nil {
if logger != nil {
logger.Error("tool execution worker panic recovered", zap.Any("recover", r), zap.String("toolName", toolName), zap.Stack("stack"))
}
err = fmt.Errorf("tool execution panic: %v", r)
}
}()
return fn()
}
func (s *ExecutionService) Wait(ctx context.Context, executionID string, timeout time.Duration) (*ExecutionSnapshot, error) {
entry := s.getEntry(executionID)
if entry == nil {
return s.getPersistedSnapshot(executionID)
}
if isExecutionTerminal(entry.exec.Status) {
return &ExecutionSnapshot{Execution: cloneToolExecution(entry.exec)}, nil
}
var timeoutCh <-chan time.Time
var timer *time.Timer
if timeout > 0 {
timer = time.NewTimer(timeout)
timeoutCh = timer.C
defer timer.Stop()
}
select {
case <-entry.done:
return &ExecutionSnapshot{Execution: cloneToolExecution(entry.exec)}, nil
case <-timeoutCh:
return &ExecutionSnapshot{Execution: cloneToolExecution(entry.exec)}, ErrExecutionWaitTimeout
case <-ctxDone(ctx):
return &ExecutionSnapshot{Execution: cloneToolExecution(entry.exec)}, ctx.Err()
}
}
func (s *ExecutionService) Get(executionID string) (*ExecutionSnapshot, error) {
entry := s.getEntry(executionID)
if entry != nil {
return &ExecutionSnapshot{Execution: cloneToolExecution(entry.exec)}, nil
}
return s.getPersistedSnapshot(executionID)
}
func (s *ExecutionService) AppendPartialOutput(executionID, chunk string) bool {
id := strings.TrimSpace(executionID)
if s == nil || id == "" || chunk == "" {
return false
}
now := time.Now()
s.mu.Lock()
defer s.mu.Unlock()
entry := s.entries[id]
if entry == nil || entry.exec == nil {
return false
}
appendPartialOutput(entry.exec, chunk, defaultPartialOutputMaxBytes, now)
return true
}
func (s *ExecutionService) Cancel(executionID, note string) bool {
id := strings.TrimSpace(executionID)
if id == "" || s == nil {
return false
}
s.mu.Lock()
entry := s.entries[id]
if entry == nil || isExecutionTerminal(entry.exec.Status) {
s.mu.Unlock()
return false
}
if strings.TrimSpace(note) != "" {
s.abortUserNotes[id] = strings.TrimSpace(note)
}
cancel := entry.cancel
s.mu.Unlock()
if cancel != nil {
cancel()
}
return true
}
func (s *ExecutionService) ActiveRunningExecutionIDs() map[string]struct{} {
if s == nil {
return nil
}
s.mu.Lock()
defer s.mu.Unlock()
out := make(map[string]struct{})
for id, entry := range s.entries {
if entry != nil && entry.exec != nil && !isExecutionTerminal(entry.exec.Status) {
out[id] = struct{}{}
}
}
if len(out) == 0 {
return nil
}
return out
}
func (s *ExecutionService) CancelAll(note string) {
if s == nil {
return
}
s.mu.Lock()
cancels := make([]context.CancelFunc, 0, len(s.entries))
for id, entry := range s.entries {
if entry == nil || isExecutionTerminal(entry.exec.Status) {
continue
}
if strings.TrimSpace(note) != "" {
s.abortUserNotes[id] = strings.TrimSpace(note)
}
if entry.cancel != nil {
cancels = append(cancels, entry.cancel)
}
}
s.mu.Unlock()
for _, cancel := range cancels {
cancel()
}
}
func (s *ExecutionService) getEntry(executionID string) *executionEntry {
if s == nil {
return nil
}
id := strings.TrimSpace(executionID)
if id == "" {
return nil
}
s.mu.Lock()
defer s.mu.Unlock()
return s.entries[id]
}
func (s *ExecutionService) getPersistedSnapshot(executionID string) (*ExecutionSnapshot, error) {
id := strings.TrimSpace(executionID)
if id == "" {
return nil, fmt.Errorf("execution_id is required")
}
if s != nil && s.storage != nil {
exec, err := s.storage.GetToolExecution(id)
if err == nil && exec != nil {
return &ExecutionSnapshot{Execution: exec}, nil
}
if err != nil {
return nil, err
}
}
return nil, fmt.Errorf("execution not found: %s", id)
}
func (s *ExecutionService) applyAbortUserNoteToCancelledToolResult(executionID string, result **ToolResult, err *error) (cancelledWithUserNote bool) {
note := strings.TrimSpace(s.takeAbortUserNote(executionID))
if note == "" {
return false
}
hasErr := err != nil && *err != nil
hasRes := result != nil && *result != nil
if !hasErr && !hasRes {
return false
}
partial := ""
if hasRes {
partial = ToolResultPlainText(*result)
}
if partial == "" && hasErr {
partial = (*err).Error()
}
merged := MergePartialToolOutputAndAbortNote(partial, note)
if err != nil {
*err = nil
}
if result != nil {
*result = &ToolResult{Content: []Content{{Type: "text", Text: merged}}, IsError: true}
}
return true
}
func (s *ExecutionService) takeAbortUserNote(id string) string {
s.mu.Lock()
defer s.mu.Unlock()
note := s.abortUserNotes[id]
delete(s.abortUserNotes, id)
return note
}
func (s *ExecutionService) cleanupOldEntriesLocked() {
if s.maxInMemory <= 0 || len(s.entries) <= s.maxInMemory {
return
}
type oldEntry struct {
id string
startTime time.Time
}
var terminal []oldEntry
for id, entry := range s.entries {
if entry != nil && entry.exec != nil && isExecutionTerminal(entry.exec.Status) {
terminal = append(terminal, oldEntry{id: id, startTime: entry.exec.StartTime})
}
}
for len(s.entries) > s.maxInMemory && len(terminal) > 0 {
oldest := 0
for i := 1; i < len(terminal); i++ {
if terminal[i].startTime.Before(terminal[oldest].startTime) {
oldest = i
}
}
delete(s.entries, terminal[oldest].id)
terminal = append(terminal[:oldest], terminal[oldest+1:]...)
}
}
func firstToolResultText(result *ToolResult, fallback string) string {
if result != nil {
for _, c := range result.Content {
if strings.TrimSpace(c.Text) != "" {
return c.Text
}
}
}
return fallback
}
func isBackgroundWaitToolResult(result *ToolResult) bool {
text := strings.ToLower(strings.TrimSpace(ToolResultPlainText(result)))
if text == "" {
return false
}
hasExecutionID := strings.Contains(text, "execution_id:") || strings.Contains(text, `"execution_id"`)
hasRunningStatus := strings.Contains(text, "status: running") || strings.Contains(text, "status: queued") ||
strings.Contains(text, `"status": "running"`) || strings.Contains(text, `"status":"running"`) ||
strings.Contains(text, `"status": "queued"`) || strings.Contains(text, `"status":"queued"`)
hasSoftWaitSignal := strings.Contains(text, "工具已提交到后台执行") ||
strings.Contains(text, "本次等待已到达") ||
strings.Contains(text, "wait_timeout:") ||
strings.Contains(text, "background execution") ||
strings.Contains(text, "still running") ||
strings.Contains(text, "仍未完成")
return hasExecutionID && hasRunningStatus && hasSoftWaitSignal
}
func isExecutionTerminal(status string) bool {
switch strings.TrimSpace(strings.ToLower(status)) {
case ToolExecutionStatusCompleted, ToolExecutionStatusFailed, ToolExecutionStatusCancelled, ToolExecutionStatusHardTimeout, ToolExecutionStatusOrphaned:
return true
default:
return false
}
}
func ctxDone(ctx context.Context) <-chan struct{} {
if ctx == nil {
return nil
}
return ctx.Done()
}
func detachedExecutionContext(ctx context.Context) context.Context {
if ctx == nil {
return context.Background()
}
return context.WithoutCancel(ctx)
}
func cloneArgsMap(in map[string]interface{}) map[string]interface{} {
if in == nil {
return map[string]interface{}{}
}
out := make(map[string]interface{}, len(in))
for k, v := range in {
out[k] = v
}
return out
}
func cloneToolExecution(in *ToolExecution) *ToolExecution {
if in == nil {
return nil
}
out := *in
out.Arguments = cloneArgsMap(in.Arguments)
if in.Result != nil {
res := *in.Result
if in.Result.Content != nil {
res.Content = append([]Content(nil), in.Result.Content...)
}
out.Result = &res
}
if in.EndTime != nil {
t := *in.EndTime
out.EndTime = &t
}
if in.PartialOutputUpdatedAt != nil {
t := *in.PartialOutputUpdatedAt
out.PartialOutputUpdatedAt = &t
}
return &out
}
func appendPartialOutput(exec *ToolExecution, chunk string, maxBytes int, updatedAt time.Time) {
if exec == nil || chunk == "" {
return
}
if maxBytes <= 0 {
maxBytes = defaultPartialOutputMaxBytes
}
exec.PartialOutputBytes += int64(len([]byte(chunk)))
combined := exec.PartialOutput + chunk
if len([]byte(combined)) > maxBytes {
b := []byte(combined)
combined = string(b[len(b)-maxBytes:])
exec.PartialOutputTruncated = true
}
exec.PartialOutput = combined
t := updatedAt
exec.PartialOutputUpdatedAt = &t
}
+41
View File
@@ -0,0 +1,41 @@
package mcp
import (
"context"
"testing"
)
func TestExecutionServiceBackgroundWaitResultCompletesWaitTool(t *testing.T) {
service := NewExecutionService(nil, nil)
handle, err := service.Submit(context.Background(), ExecutionRequest{
ToolName: "wait_tool_execution",
Run: func(context.Context) (*ToolResult, error) {
return &ToolResult{
Content: []Content{{Type: "text", Text: `{
"execution_id": "3eaaa391-050b-4be1-a870-48a855923cb7",
"tool": "exec",
"status": "running"
}
本次等待已到达 timeout_seconds,上述 execution 仍未完成。可继续等待、取消,或采用其他步骤。`}},
IsError: true,
}, nil
},
})
if err != nil {
t.Fatalf("Submit: %v", err)
}
snap, err := service.Wait(context.Background(), handle.ID, 0)
if err != nil {
t.Fatalf("Wait: %v", err)
}
if snap == nil || snap.Execution == nil {
t.Fatal("missing execution snapshot")
}
if snap.Execution.Status != ToolExecutionStatusCompleted {
t.Fatalf("status = %q, want %q", snap.Execution.Status, ToolExecutionStatusCompleted)
}
if snap.Execution.Result == nil || !snap.Execution.Result.IsError {
t.Fatal("model-facing result should remain IsError")
}
}
File diff suppressed because it is too large Load Diff
+230
View File
@@ -0,0 +1,230 @@
package mcp
import (
"context"
"errors"
"strings"
"sync/atomic"
"testing"
"time"
"go.uber.org/zap"
)
type blockingExternalMCPClient struct {
started chan struct{}
calls chan string
release chan struct{}
result *ToolResult
count atomic.Int32
}
func newBlockingExternalMCPClient(resultText string) *blockingExternalMCPClient {
return &blockingExternalMCPClient{
started: make(chan struct{}),
calls: make(chan string, 8),
release: make(chan struct{}),
result: &ToolResult{Content: []Content{{Type: "text", Text: resultText}}},
}
}
func (c *blockingExternalMCPClient) Initialize(ctx context.Context) error { return nil }
func (c *blockingExternalMCPClient) ListTools(ctx context.Context) ([]Tool, error) {
return []Tool{{Name: "slow_tool"}}, nil
}
func (c *blockingExternalMCPClient) CallTool(ctx context.Context, name string, args map[string]interface{}) (*ToolResult, error) {
c.count.Add(1)
select {
case c.calls <- name:
default:
}
select {
case <-c.started:
default:
close(c.started)
}
select {
case <-c.release:
return c.result, nil
case <-ctx.Done():
return nil, ctx.Err()
}
}
func (c *blockingExternalMCPClient) Close() error { return nil }
func (c *blockingExternalMCPClient) IsConnected() bool { return true }
func (c *blockingExternalMCPClient) GetStatus() string { return "connected" }
type failingExternalMCPClient struct{}
func (c *failingExternalMCPClient) Initialize(ctx context.Context) error { return nil }
func (c *failingExternalMCPClient) ListTools(ctx context.Context) ([]Tool, error) {
return []Tool{{Name: "fail_tool"}}, nil
}
func (c *failingExternalMCPClient) CallTool(ctx context.Context, name string, args map[string]interface{}) (*ToolResult, error) {
return nil, errors.New("boom")
}
func (c *failingExternalMCPClient) Close() error { return nil }
func (c *failingExternalMCPClient) IsConnected() bool { return true }
func (c *failingExternalMCPClient) GetStatus() string { return "connected" }
func TestExternalMCPManager_CallToolBoundedWaitThenContinue(t *testing.T) {
manager := NewExternalMCPManager(zap.NewNop())
manager.ConfigureToolWaitTimeoutSeconds(1)
manager.toolWaitTimeout = 10 * time.Millisecond
client := newBlockingExternalMCPClient("slow result ready")
manager.clients["lab"] = client
callCtx, callCancel := context.WithCancel(context.Background())
result, executionID, err := manager.CallTool(callCtx, "lab::slow_tool", map[string]interface{}{"target": "example"})
if err != nil {
t.Fatalf("CallTool returned error: %v", err)
}
if executionID == "" {
t.Fatal("expected execution id")
}
if result == nil || !result.IsError {
t.Fatalf("expected soft timeout tool result, got %#v", result)
}
text := ToolResultPlainText(result)
if !strings.Contains(text, executionID) || !strings.Contains(text, "wait_tool_execution") {
t.Fatalf("timeout result should include execution id and wait guidance, got %q", text)
}
select {
case <-client.started:
default:
t.Fatal("worker did not start")
}
callCancel()
close(client.release)
snapshot, err := manager.executionService.Wait(context.Background(), executionID, time.Second)
if err != nil {
t.Fatalf("Wait returned error: %v", err)
}
if snapshot == nil || snapshot.Execution == nil {
t.Fatal("expected execution snapshot")
}
if snapshot.Execution.Status != ToolExecutionStatusCompleted {
t.Fatalf("status = %q, want completed", snapshot.Execution.Status)
}
if got := ToolResultPlainText(snapshot.Execution.Result); got != "slow result ready" {
t.Fatalf("result = %q, want slow result ready", got)
}
}
func TestExecutionControlWaitToolReturnsCompletedResult(t *testing.T) {
manager := NewExternalMCPManager(zap.NewNop())
manager.toolWaitTimeout = 10 * time.Millisecond
client := newBlockingExternalMCPClient("control wait result")
manager.clients["lab"] = client
result, executionID, err := manager.CallTool(context.Background(), "lab::slow_tool", nil)
if err != nil {
t.Fatalf("CallTool returned error: %v", err)
}
if result == nil || !result.IsError || executionID == "" {
t.Fatalf("expected soft timeout and execution id, got result=%#v id=%q", result, executionID)
}
server := NewServer(zap.NewNop())
RegisterExecutionControlTools(server, manager)
close(client.release)
waitResult, _, err := server.CallTool(context.Background(), "wait_tool_execution", map[string]interface{}{
"execution_id": executionID,
"timeout_seconds": 1,
})
if err != nil {
t.Fatalf("wait_tool_execution returned error: %v", err)
}
if waitResult == nil || waitResult.IsError {
t.Fatalf("expected successful wait result, got %#v", waitResult)
}
body := ToolResultPlainText(waitResult)
if !strings.Contains(body, `"status": "completed"`) || !strings.Contains(body, "control wait result") {
t.Fatalf("wait result body missing completed status/result: %s", body)
}
}
func TestExternalMCPManager_PerServerConcurrencyLimitsWorkers(t *testing.T) {
manager := NewExternalMCPManager(zap.NewNop())
manager.toolWaitTimeout = 10 * time.Millisecond
manager.ConfigureResilience(ExternalMCPResilienceConfig{
MaxConcurrentPerServer: 1,
MaxConcurrentTotal: 4,
CircuitFailureThreshold: -1,
CircuitCooldown: time.Second,
})
client := newBlockingExternalMCPClient("ok")
manager.clients["lab"] = client
done1 := make(chan struct{})
go func() {
_, _, _ = manager.CallTool(context.Background(), "lab::slow_tool", nil)
close(done1)
}()
select {
case <-client.calls:
case <-time.After(time.Second):
t.Fatal("first worker did not enter client")
}
type callOutcome struct {
executionID string
err error
}
done2 := make(chan callOutcome, 1)
go func() {
_, executionID, err := manager.CallTool(context.Background(), "lab::slow_tool", nil)
done2 <- callOutcome{executionID: executionID, err: err}
}()
select {
case <-client.calls:
t.Fatal("second worker entered client before per-server slot was released")
case <-time.After(50 * time.Millisecond):
}
var second callOutcome
select {
case second = <-done2:
case <-time.After(time.Second):
t.Fatal("second call did not return after bounded wait")
}
if second.err != nil || second.executionID == "" {
t.Fatalf("second call should return queued execution id after bounded wait, id=%q err=%v", second.executionID, second.err)
}
snapshot, err := manager.executionService.Get(second.executionID)
if err != nil {
t.Fatalf("Get queued execution: %v", err)
}
if snapshot == nil || snapshot.Execution == nil || snapshot.Execution.Status != ToolExecutionStatusQueued {
t.Fatalf("second execution status = %#v, want queued", snapshot)
}
close(client.release)
select {
case <-client.calls:
case <-time.After(time.Second):
t.Fatal("second worker did not enter client after slot release")
}
<-done1
}
func TestExternalMCPManager_CircuitBreakerOpensAfterFailures(t *testing.T) {
manager := NewExternalMCPManager(zap.NewNop())
manager.ConfigureResilience(ExternalMCPResilienceConfig{
MaxConcurrentPerServer: 2,
MaxConcurrentTotal: 4,
CircuitFailureThreshold: 1,
CircuitCooldown: time.Minute,
})
manager.clients["lab"] = &failingExternalMCPClient{}
_, _, err := manager.CallTool(context.Background(), "lab::fail_tool", nil)
if err == nil || !strings.Contains(err.Error(), "boom") {
t.Fatalf("expected first call to fail with client error, got %v", err)
}
_, _, err = manager.CallTool(context.Background(), "lab::fail_tool", nil)
if err == nil || !strings.Contains(err.Error(), "熔断") {
t.Fatalf("expected circuit breaker rejection, got %v", err)
}
}
+261
View File
@@ -0,0 +1,261 @@
package mcp
import (
"context"
"errors"
"strings"
"testing"
"time"
"cyberstrike-ai/internal/authctx"
"cyberstrike-ai/internal/config"
"go.uber.org/zap"
)
func TestExternalManagerEnforcesConfiguredAuthorizer(t *testing.T) {
manager := NewExternalMCPManager(zap.NewNop())
t.Cleanup(manager.StopAll)
manager.SetToolAuthorizer(func(context.Context, string, map[string]interface{}) error {
return errors.New("denied by policy")
})
ctx := authctx.WithPrincipal(context.Background(), authctx.NewPrincipal("u1", "user", "assigned", map[string]bool{"agent:execute": true}))
_, executionID, err := manager.CallTool(ctx, "server::tool", map[string]interface{}{})
if err == nil || !strings.Contains(err.Error(), "authorization denied") {
t.Fatalf("external call bypassed authorizer: %v", err)
}
if executionID == "" {
t.Fatal("denied external call should still return an execution id")
}
execution, ok := manager.GetExecution(executionID)
if !ok || execution == nil {
t.Fatalf("missing denied external execution %q", executionID)
}
if execution.Status != ToolExecutionStatusFailed || !strings.Contains(execution.Error, "denied by policy") {
t.Fatalf("denied external execution = %#v, want failed with policy error", execution)
}
}
func TestExternalMCPManager_AddOrUpdateConfig(t *testing.T) {
logger := zap.NewNop()
manager := NewExternalMCPManager(logger)
// 测试添加stdio配置
stdioCfg := config.ExternalMCPServerConfig{
Command: "python3",
Args: []string{"/path/to/script.py"},
Description: "Test stdio MCP",
Timeout: 30,
ExternalMCPEnable: true,
}
err := manager.AddOrUpdateConfig("test-stdio", stdioCfg)
if err != nil {
t.Fatalf("添加stdio配置失败: %v", err)
}
// 测试添加HTTP配置
httpCfg := config.ExternalMCPServerConfig{
Type: "http",
URL: "http://127.0.0.1:8081/mcp",
Description: "Test HTTP MCP",
Timeout: 30,
ExternalMCPEnable: false,
}
err = manager.AddOrUpdateConfig("test-http", httpCfg)
if err != nil {
t.Fatalf("添加HTTP配置失败: %v", err)
}
// 验证配置已保存
configs := manager.GetConfigs()
if len(configs) != 2 {
t.Fatalf("期望2个配置,实际%d个", len(configs))
}
if configs["test-stdio"].Command != stdioCfg.Command {
t.Errorf("stdio配置命令不匹配")
}
if configs["test-http"].URL != httpCfg.URL {
t.Errorf("HTTP配置URL不匹配")
}
}
func TestExternalMCPManager_RemoveConfig(t *testing.T) {
logger := zap.NewNop()
manager := NewExternalMCPManager(logger)
cfg := config.ExternalMCPServerConfig{
Command: "python3",
ExternalMCPEnable: false,
}
manager.AddOrUpdateConfig("test-remove", cfg)
// 移除配置
err := manager.RemoveConfig("test-remove")
if err != nil {
t.Fatalf("移除配置失败: %v", err)
}
configs := manager.GetConfigs()
if _, exists := configs["test-remove"]; exists {
t.Error("配置应该已被移除")
}
}
func TestExternalMCPManager_GetStats(t *testing.T) {
logger := zap.NewNop()
manager := NewExternalMCPManager(logger)
// 添加多个配置
manager.AddOrUpdateConfig("enabled1", config.ExternalMCPServerConfig{
Command: "python3",
ExternalMCPEnable: true,
})
manager.AddOrUpdateConfig("enabled2", config.ExternalMCPServerConfig{
URL: "http://127.0.0.1:8081/mcp",
ExternalMCPEnable: true,
})
manager.AddOrUpdateConfig("disabled1", config.ExternalMCPServerConfig{
Command: "python3",
ExternalMCPEnable: false,
})
stats := manager.GetStats()
if stats["total"].(int) != 3 {
t.Errorf("期望总数3,实际%d", stats["total"])
}
if stats["enabled"].(int) != 2 {
t.Errorf("期望启用数2,实际%d", stats["enabled"])
}
if stats["disabled"].(int) != 1 {
t.Errorf("期望停用数1,实际%d", stats["disabled"])
}
}
func TestExternalMCPManager_LoadConfigs(t *testing.T) {
logger := zap.NewNop()
manager := NewExternalMCPManager(logger)
externalMCPConfig := config.ExternalMCPConfig{
Servers: map[string]config.ExternalMCPServerConfig{
"loaded1": {
Command: "python3",
ExternalMCPEnable: true,
},
"loaded2": {
URL: "http://127.0.0.1:8081/mcp",
ExternalMCPEnable: false,
},
},
}
manager.LoadConfigs(&externalMCPConfig)
configs := manager.GetConfigs()
if len(configs) != 2 {
t.Fatalf("期望2个配置,实际%d个", len(configs))
}
if configs["loaded1"].Command != "python3" {
t.Error("配置1加载失败")
}
if configs["loaded2"].URL != "http://127.0.0.1:8081/mcp" {
t.Error("配置2加载失败")
}
}
// TestLazySDKClient_InitializeFails 验证无效配置时 SDK 客户端 Initialize 失败并设置 error 状态
func TestLazySDKClient_InitializeFails(t *testing.T) {
logger := zap.NewNop()
// 使用不存在的 HTTP 地址,Initialize 应失败
cfg := config.ExternalMCPServerConfig{
Type: "http",
URL: "http://127.0.0.1:19999/nonexistent",
Timeout: 2,
}
c := newLazySDKClient(cfg, logger)
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
err := c.Initialize(ctx)
if err == nil {
t.Fatal("expected error when connecting to invalid server")
}
if c.GetStatus() != "error" {
t.Errorf("expected status error, got %s", c.GetStatus())
}
c.Close()
}
func TestExternalMCPManager_StartStopClient(t *testing.T) {
logger := zap.NewNop()
manager := NewExternalMCPManager(logger)
// 添加一个禁用的配置
cfg := config.ExternalMCPServerConfig{
Command: "python3",
ExternalMCPEnable: false,
}
manager.AddOrUpdateConfig("test-start-stop", cfg)
// 尝试启动(可能会失败,因为没有真实的服务器)
err := manager.StartClient("test-start-stop")
if err != nil {
t.Logf("启动失败(可能是没有服务器): %v", err)
}
// 停止
err = manager.StopClient("test-start-stop")
if err != nil {
t.Fatalf("停止失败: %v", err)
}
// 验证配置已更新为禁用
configs := manager.GetConfigs()
if configs["test-start-stop"].ExternalMCPEnable {
t.Error("配置应该已被禁用")
}
}
func TestExternalMCPManager_CallTool(t *testing.T) {
logger := zap.NewNop()
manager := NewExternalMCPManager(logger)
// 测试调用不存在的工具
_, _, err := manager.CallTool(context.Background(), "nonexistent::tool", map[string]interface{}{})
if err == nil {
t.Error("应该返回错误")
}
// 测试无效的工具名称格式
_, _, err = manager.CallTool(context.Background(), "invalid-tool-name", map[string]interface{}{})
if err == nil {
t.Error("应该返回错误(无效格式)")
}
}
func TestExternalMCPManager_GetAllTools(t *testing.T) {
logger := zap.NewNop()
manager := NewExternalMCPManager(logger)
ctx := context.Background()
tools, err := manager.GetAllTools(ctx)
if err != nil {
t.Fatalf("获取工具列表失败: %v", err)
}
// 如果没有连接的客户端,应该返回空列表
if len(tools) != 0 {
t.Logf("获取到%d个工具", len(tools))
}
}
+147
View File
@@ -0,0 +1,147 @@
package mcp
import (
"context"
"strings"
)
// ToolRunRegistry 在工具开始/结束时登记当前 executionId,供对话页「仅终止当前工具」与监控页共用取消逻辑。
type ToolRunRegistry interface {
RegisterRunningTool(conversationID, executionID string)
UnregisterRunningTool(conversationID, executionID string)
}
// EinoExecuteRunRegistry 登记进行中的 Eino filesystem execute,供「中断并继续」终止 amass 等长命令。
type EinoExecuteRunRegistry interface {
RegisterActiveEinoExecute(conversationID string, cancel context.CancelFunc)
UnregisterActiveEinoExecute(conversationID string)
AbortActiveEinoExecute(conversationID, note string) bool
TakeEinoExecuteAbortNote(conversationID string) string
}
type toolRunRegistryCtxKey struct{}
type einoExecuteRunRegistryCtxKey struct{}
type mcpConversationIDCtxKey struct{}
type mcpExecutionIDCtxKey struct{}
type mcpProjectIDCtxKey struct{}
// WithToolRunRegistry 将登记器注入 ctxEino / 原生 Agent 任务 ctx)。
func WithToolRunRegistry(ctx context.Context, reg ToolRunRegistry) context.Context {
if ctx == nil || reg == nil {
return ctx
}
return context.WithValue(ctx, toolRunRegistryCtxKey{}, reg)
}
// ToolRunRegistryFromContext 取出登记器(无则 nil)。
func ToolRunRegistryFromContext(ctx context.Context) ToolRunRegistry {
if ctx == nil {
return nil
}
v, _ := ctx.Value(toolRunRegistryCtxKey{}).(ToolRunRegistry)
return v
}
// WithEinoExecuteRunRegistry 将 Eino execute 取消登记器注入 ctx。
func WithEinoExecuteRunRegistry(ctx context.Context, reg EinoExecuteRunRegistry) context.Context {
if ctx == nil || reg == nil {
return ctx
}
return context.WithValue(ctx, einoExecuteRunRegistryCtxKey{}, reg)
}
// EinoExecuteRunRegistryFromContext 取出 Eino execute 登记器(无则 nil)。
func EinoExecuteRunRegistryFromContext(ctx context.Context) EinoExecuteRunRegistry {
if ctx == nil {
return nil
}
v, _ := ctx.Value(einoExecuteRunRegistryCtxKey{}).(EinoExecuteRunRegistry)
return v
}
// WithMCPConversationID 将对话 ID 注入 ctx,供 CallTool 内与 executionId 关联。
func WithMCPConversationID(ctx context.Context, conversationID string) context.Context {
if ctx == nil {
return nil
}
id := strings.TrimSpace(conversationID)
if id == "" {
return ctx
}
return context.WithValue(ctx, mcpConversationIDCtxKey{}, id)
}
// MCPConversationIDFromContext 读取对话 ID。
func MCPConversationIDFromContext(ctx context.Context) string {
if ctx == nil {
return ""
}
v, _ := ctx.Value(mcpConversationIDCtxKey{}).(string)
return v
}
// WithMCPExecutionID 将当前工具 executionId 注入 ctx,供超长输出落盘文件名对齐。
func WithMCPExecutionID(ctx context.Context, executionID string) context.Context {
if ctx == nil {
return nil
}
id := strings.TrimSpace(executionID)
if id == "" {
return ctx
}
return context.WithValue(ctx, mcpExecutionIDCtxKey{}, id)
}
// MCPExecutionIDFromContext 读取当前工具 executionId。
func MCPExecutionIDFromContext(ctx context.Context) string {
if ctx == nil {
return ""
}
v, _ := ctx.Value(mcpExecutionIDCtxKey{}).(string)
return v
}
// WithMCPProjectID 将项目 ID 注入 ctx,供 reduction/trunc 落盘路径与项目隔离对齐。
func WithMCPProjectID(ctx context.Context, projectID string) context.Context {
if ctx == nil {
return nil
}
id := strings.TrimSpace(projectID)
if id == "" {
return ctx
}
return context.WithValue(ctx, mcpProjectIDCtxKey{}, id)
}
// MCPProjectIDFromContext 读取项目 ID。
func MCPProjectIDFromContext(ctx context.Context) string {
if ctx == nil {
return ""
}
v, _ := ctx.Value(mcpProjectIDCtxKey{}).(string)
return v
}
func notifyToolRunBegin(ctx context.Context, executionID string) {
reg := ToolRunRegistryFromContext(ctx)
if reg == nil {
return
}
conv := MCPConversationIDFromContext(ctx)
if conv == "" || strings.TrimSpace(executionID) == "" {
return
}
reg.RegisterRunningTool(conv, executionID)
}
func notifyToolRunEnd(ctx context.Context, executionID string) {
reg := ToolRunRegistryFromContext(ctx)
if reg == nil {
return
}
conv := MCPConversationIDFromContext(ctx)
if conv == "" || strings.TrimSpace(executionID) == "" {
return
}
reg.UnregisterRunningTool(conv, executionID)
}
File diff suppressed because it is too large Load Diff
+231
View File
@@ -0,0 +1,231 @@
package mcp
import (
"context"
"errors"
"strings"
"testing"
"time"
"cyberstrike-ai/internal/authctx"
"go.uber.org/zap"
)
func TestToolAuthorizerIsUniversalAndExecutionKeepsOwner(t *testing.T) {
server := NewServer(zap.NewNop())
server.RegisterTool(Tool{Name: "echo", InputSchema: map[string]interface{}{"type": "object"}}, func(ctx context.Context, args map[string]interface{}) (*ToolResult, error) {
return &ToolResult{Content: []Content{{Type: "text", Text: "ok"}}}, nil
})
server.SetToolAuthorizer(func(ctx context.Context, toolName string, args map[string]interface{}) error {
if _, ok := authctx.PrincipalFromContext(ctx); !ok {
return errors.New("principal required")
}
return nil
})
_, deniedExecutionID, err := server.CallTool(context.Background(), "echo", nil)
if err == nil {
t.Fatal("tool call without principal was allowed")
}
if deniedExecutionID == "" {
t.Fatal("denied tool call should still return an execution id")
}
deniedExecution, ok := server.GetExecution(deniedExecutionID)
if !ok || deniedExecution == nil {
t.Fatalf("missing denied execution %q", deniedExecutionID)
}
if deniedExecution.Status != ToolExecutionStatusFailed || !strings.Contains(deniedExecution.Error, "principal required") {
t.Fatalf("denied execution = %#v, want failed with authorization error", deniedExecution)
}
ctx := authctx.WithPrincipal(context.Background(), authctx.NewPrincipal("u1", "user", "assigned", map[string]bool{"mcp:execute": true}))
_, executionID, err := server.CallTool(ctx, "echo", nil)
if err != nil {
t.Fatal(err)
}
execution, ok := server.GetExecution(executionID)
if !ok || execution.OwnerUserID != "u1" {
t.Fatalf("execution owner = %#v, want u1", execution)
}
}
func TestServerCallToolBoundedWaitForInternalTool(t *testing.T) {
server := NewServer(zap.NewNop())
server.toolWaitTimeout = 10 * time.Millisecond
release := make(chan struct{})
started := make(chan struct{})
server.RegisterTool(Tool{Name: "slow", InputSchema: map[string]interface{}{"type": "object"}}, func(ctx context.Context, args map[string]interface{}) (*ToolResult, error) {
close(started)
select {
case <-release:
return &ToolResult{Content: []Content{{Type: "text", Text: "internal done"}}}, nil
case <-ctx.Done():
return nil, ctx.Err()
}
})
callCtx, callCancel := context.WithCancel(context.Background())
result, executionID, err := server.CallTool(callCtx, "slow", nil)
if err != nil {
t.Fatalf("CallTool returned error: %v", err)
}
if executionID == "" || result == nil || !result.IsError {
t.Fatalf("expected soft timeout with execution id, result=%#v id=%q", result, executionID)
}
if text := ToolResultPlainText(result); !strings.Contains(text, executionID) || !strings.Contains(text, "wait_tool_execution") {
t.Fatalf("timeout result missing execution guidance: %q", text)
}
select {
case <-started:
default:
t.Fatal("internal worker did not start")
}
callCancel()
close(release)
snapshot, err := server.executionService.Wait(context.Background(), executionID, time.Second)
if err != nil {
t.Fatalf("wait internal execution: %v", err)
}
if snapshot == nil || snapshot.Execution == nil || snapshot.Execution.Status != ToolExecutionStatusCompleted {
t.Fatalf("snapshot = %#v, want completed", snapshot)
}
if got := ToolResultPlainText(snapshot.Execution.Result); got != "internal done" {
t.Fatalf("result = %q, want internal done", got)
}
}
func TestWaitToolExecutionWaitsForInternalActiveExecution(t *testing.T) {
server := NewServer(zap.NewNop())
server.toolWaitTimeout = 10 * time.Millisecond
release := make(chan struct{})
server.RegisterTool(Tool{Name: "slow", InputSchema: map[string]interface{}{"type": "object"}}, func(ctx context.Context, args map[string]interface{}) (*ToolResult, error) {
select {
case <-release:
return &ToolResult{Content: []Content{{Type: "text", Text: "wait saw completion"}}}, nil
case <-ctx.Done():
return nil, ctx.Err()
}
})
RegisterExecutionControlTools(server, nil)
result, executionID, err := server.CallTool(context.Background(), "slow", nil)
if err != nil {
t.Fatalf("CallTool returned error: %v", err)
}
if result == nil || !result.IsError || executionID == "" {
t.Fatalf("expected initial bounded wait timeout, result=%#v id=%q", result, executionID)
}
done := make(chan *ToolResult, 1)
errCh := make(chan error, 1)
go func() {
waitResult, _, waitErr := server.CallTool(context.Background(), "wait_tool_execution", map[string]interface{}{
"execution_id": executionID,
"timeout_seconds": 1,
})
if waitErr != nil {
errCh <- waitErr
return
}
done <- waitResult
}()
select {
case <-done:
t.Fatal("wait_tool_execution returned before target execution completed")
case err := <-errCh:
t.Fatalf("wait_tool_execution errored before release: %v", err)
case <-time.After(50 * time.Millisecond):
}
close(release)
select {
case err := <-errCh:
t.Fatalf("wait_tool_execution returned error: %v", err)
case waitResult := <-done:
if waitResult == nil || waitResult.IsError {
t.Fatalf("expected successful wait result, got %#v", waitResult)
}
if body := ToolResultPlainText(waitResult); !strings.Contains(body, "wait saw completion") || !strings.Contains(body, `"status": "completed"`) {
t.Fatalf("wait result missing completed target: %s", body)
}
case <-time.After(time.Second):
t.Fatal("wait_tool_execution did not return after target completion")
}
}
func TestWaitToolExecutionTimeoutIsObservationNotFailure(t *testing.T) {
server := NewServer(zap.NewNop())
server.toolWaitTimeout = 10 * time.Millisecond
release := make(chan struct{})
server.RegisterTool(Tool{Name: "slow_observed", InputSchema: map[string]interface{}{"type": "object"}}, func(ctx context.Context, args map[string]interface{}) (*ToolResult, error) {
<-release
return &ToolResult{Content: []Content{{Type: "text", Text: "done"}}}, nil
})
RegisterExecutionControlTools(server, nil)
result, executionID, err := server.CallTool(context.Background(), "slow_observed", nil)
if err != nil {
t.Fatalf("CallTool returned error: %v", err)
}
if result == nil || !result.IsError || executionID == "" {
t.Fatalf("expected initial bounded wait timeout, result=%#v id=%q", result, executionID)
}
waitResult, _, err := server.CallTool(context.Background(), "wait_tool_execution", map[string]interface{}{
"execution_id": executionID,
"timeout_seconds": 0.01,
})
if err != nil {
t.Fatalf("wait_tool_execution returned error: %v", err)
}
if waitResult == nil {
t.Fatal("missing wait result")
}
if waitResult.IsError {
t.Fatalf("wait timeout should be a successful observation, got %#v", waitResult)
}
body := ToolResultPlainText(waitResult)
if !strings.Contains(body, `"status": "running"`) || !strings.Contains(body, "本次等待已到达") {
t.Fatalf("wait timeout body missing running status/guidance: %s", body)
}
close(release)
}
func TestGetToolExecutionIncludesBoundedPartialOutput(t *testing.T) {
server := NewServer(zap.NewNop())
RegisterExecutionControlTools(server, nil)
executionID := server.BeginToolExecution(context.Background(), "execute", map[string]interface{}{"command": "demo"})
if executionID == "" {
t.Fatal("missing execution id")
}
server.AppendToolExecutionPartialOutput(executionID, "first\n")
server.AppendToolExecutionPartialOutput(executionID, strings.Repeat("x", 32))
result, _, err := server.CallTool(context.Background(), "get_tool_execution", map[string]interface{}{
"execution_id": executionID,
"partial_output_max_bytes": 8,
})
if err != nil {
t.Fatalf("get_tool_execution: %v", err)
}
body := ToolResultPlainText(result)
if !strings.Contains(body, `"partial_output": "xxxxxxxx"`) {
t.Fatalf("missing bounded partial output: %s", body)
}
if !strings.Contains(body, `"partial_output_bytes": 38`) {
t.Fatalf("missing partial byte count: %s", body)
}
result, _, err = server.CallTool(context.Background(), "get_tool_execution", map[string]interface{}{
"execution_id": executionID,
"include_partial_output": false,
})
if err != nil {
t.Fatalf("get_tool_execution without partial: %v", err)
}
if body := ToolResultPlainText(result); strings.Contains(body, "partial_output") {
t.Fatalf("partial output should be omitted: %s", body)
}
}
+65
View File
@@ -0,0 +1,65 @@
package mcp
import "cyberstrike-ai/internal/tooloutput"
const DefaultToolResultMaxBytes = 12000
// ToolResultSpillConfig controls where oversized tool results are written on disk
// before the in-memory/DB/agent-facing payload is truncated.
type ToolResultSpillConfig struct {
RootDir string
ProjectID string
ConversationID string
ExecutionID string
}
// NormalizeToolResultForStorage returns the canonical result used by both the
// agent-facing response and monitor persistence. When maxBytes is exceeded the
// full text is spilled under the reduction cache tree and replaced with a
// <persisted-output> notice that includes the file path.
func NormalizeToolResultForStorage(result *ToolResult, maxBytes int) *ToolResult {
return NormalizeToolResultForStorageWithSpill(result, maxBytes, ToolResultSpillConfig{})
}
// NormalizeToolResultForStorageWithSpill is NormalizeToolResultForStorage with
// an explicit spill location (conversation/execution scoped).
func NormalizeToolResultForStorageWithSpill(result *ToolResult, maxBytes int, spill ToolResultSpillConfig) *ToolResult {
if result == nil {
return nil
}
out := cloneToolResult(result)
if maxBytes <= 0 {
return out
}
total := 0
for _, c := range out.Content {
if c.Type == "text" {
total += len(c.Text)
}
}
if total <= maxBytes {
return out
}
full := ToolResultPlainText(out)
bound := tooloutput.BoundWithSpill(full, maxBytes, tooloutput.SpillOpts{
RootDir: spill.RootDir,
ProjectID: spill.ProjectID,
ConversationID: spill.ConversationID,
ExecutionID: spill.ExecutionID,
})
out.Content = []Content{{Type: "text", Text: bound}}
return out
}
func cloneToolResult(in *ToolResult) *ToolResult {
if in == nil {
return nil
}
out := *in
if in.Content != nil {
out.Content = append([]Content(nil), in.Content...)
}
return &out
}
+158
View File
@@ -0,0 +1,158 @@
package mcp
import (
"context"
"os"
"path/filepath"
"strings"
"testing"
"time"
"go.uber.org/zap"
)
type inMemoryMonitorStorage struct {
executions map[string]*ToolExecution
}
func newInMemoryMonitorStorage() *inMemoryMonitorStorage {
return &inMemoryMonitorStorage{executions: map[string]*ToolExecution{}}
}
func (s *inMemoryMonitorStorage) SaveToolExecution(exec *ToolExecution) error {
if exec != nil {
s.executions[exec.ID] = cloneToolExecution(exec)
}
return nil
}
func (s *inMemoryMonitorStorage) UpdateToolExecutionResult(id string, result *ToolResult) error {
exec := s.executions[id]
if exec == nil {
exec = &ToolExecution{ID: id}
s.executions[id] = exec
}
exec.Result = cloneToolResult(result)
return nil
}
func (s *inMemoryMonitorStorage) LoadToolExecutions() ([]*ToolExecution, error) {
out := make([]*ToolExecution, 0, len(s.executions))
for _, exec := range s.executions {
out = append(out, cloneToolExecution(exec))
}
return out, nil
}
func (s *inMemoryMonitorStorage) GetToolExecution(id string) (*ToolExecution, error) {
if exec := s.executions[id]; exec != nil {
return cloneToolExecution(exec), nil
}
return nil, nil
}
func (s *inMemoryMonitorStorage) SaveToolStats(string, *ToolStats) error { return nil }
func (s *inMemoryMonitorStorage) LoadToolStats() (map[string]*ToolStats, error) {
return map[string]*ToolStats{}, nil
}
func (s *inMemoryMonitorStorage) UpdateToolStats(string, int, int, int, *time.Time) error {
return nil
}
func TestServerCallToolStoresAndReturnsSameGuardedResult(t *testing.T) {
storage := newInMemoryMonitorStorage()
server := NewServerWithStorage(zap.NewNop(), storage)
server.ConfigureToolWaitTimeoutSeconds(0)
server.ConfigureToolResultMaxBytes(400)
spillRoot := t.TempDir()
server.ConfigureToolResultSpillRoot(spillRoot)
server.RegisterTool(Tool{Name: "big", InputSchema: map[string]interface{}{"type": "object"}}, func(ctx context.Context, args map[string]interface{}) (*ToolResult, error) {
return &ToolResult{Content: []Content{{Type: "text", Text: strings.Repeat("x", 800)}}}, nil
})
ctx := WithMCPConversationID(context.Background(), "conv-spill")
result, executionID, err := server.CallTool(ctx, "big", nil)
if err != nil {
t.Fatalf("CallTool: %v", err)
}
if executionID == "" {
t.Fatal("missing execution id")
}
returned := ToolResultPlainText(result)
if !strings.Contains(returned, "<persisted-output>") || !strings.Contains(returned, "Full output saved to:") {
t.Fatalf("returned result was not spilled: %q", returned)
}
if len(returned) > 400 {
t.Fatalf("returned result exceeded hard limit: len=%d text=%q", len(returned), returned)
}
spillPath := filepath.Join(spillRoot, "conversations", "conv-spill", "trunc", executionID)
abs, err := filepath.Abs(spillPath)
if err != nil {
t.Fatal(err)
}
if !strings.Contains(returned, abs) {
t.Fatalf("missing spill path %q in %q", abs, returned)
}
body, err := os.ReadFile(abs)
if err != nil {
t.Fatalf("read spill file: %v", err)
}
if string(body) != strings.Repeat("x", 800) {
t.Fatalf("spill body mismatch: len=%d", len(body))
}
inMem, ok := server.GetExecution(executionID)
if !ok || inMem == nil || inMem.Result == nil {
t.Fatalf("missing in-memory execution: %#v", inMem)
}
stored := storage.executions[executionID]
if stored == nil || stored.Result == nil {
t.Fatalf("missing stored execution: %#v", stored)
}
if ToolResultPlainText(inMem.Result) != returned {
t.Fatalf("in-memory result != returned\nmem=%q\nret=%q", ToolResultPlainText(inMem.Result), returned)
}
if ToolResultPlainText(stored.Result) != returned {
t.Fatalf("stored result != returned\nstored=%q\nret=%q", ToolResultPlainText(stored.Result), returned)
}
}
func TestExecutionServiceStoresGuardedResult(t *testing.T) {
service := NewExecutionService(nil, zap.NewNop())
service.ConfigureToolResultMaxBytes(400)
spillRoot := t.TempDir()
service.ConfigureToolResultSpillRoot(spillRoot)
handle, err := service.Submit(context.Background(), ExecutionRequest{
ToolName: "big",
ConversationID: "svc-conv",
Run: func(context.Context) (*ToolResult, error) {
return &ToolResult{Content: []Content{{Type: "text", Text: strings.Repeat("a", 800)}}}, nil
},
})
if err != nil {
t.Fatalf("Submit: %v", err)
}
snap, err := service.Wait(context.Background(), handle.ID, time.Second)
if err != nil {
t.Fatalf("Wait: %v", err)
}
got := ToolResultPlainText(snap.Execution.Result)
if !strings.Contains(got, "<persisted-output>") {
t.Fatalf("service result was not spilled: %q", got)
}
if len(got) > 400 {
t.Fatalf("service result exceeded hard limit: len=%d text=%q", len(got), got)
}
path := filepath.Join(spillRoot, "conversations", "svc-conv", "trunc", handle.ID)
abs, _ := filepath.Abs(path)
body, err := os.ReadFile(abs)
if err != nil {
t.Fatalf("read spill: %v", err)
}
if string(body) != strings.Repeat("a", 800) {
t.Fatalf("unexpected spill body len=%d", len(body))
}
}
+338
View File
@@ -0,0 +1,338 @@
package mcp
import (
"context"
"encoding/json"
"fmt"
"strings"
"time"
)
// ExternalMCPClient 外部 MCP 客户端接口(由 client_sdk.go 基于官方 SDK 实现)
type ExternalMCPClient interface {
Initialize(ctx context.Context) error
ListTools(ctx context.Context) ([]Tool, error)
CallTool(ctx context.Context, name string, args map[string]interface{}) (*ToolResult, error)
Close() error
IsConnected() bool
GetStatus() string
}
// MCP消息类型
const (
MessageTypeRequest = "request"
MessageTypeResponse = "response"
MessageTypeError = "error"
MessageTypeNotify = "notify"
)
// MCP协议版本
const ProtocolVersion = "2024-11-05"
// MessageID 表示JSON-RPC 2.0的id字段,可以是字符串、数字或null
type MessageID struct {
value interface{}
}
// UnmarshalJSON 自定义反序列化,支持字符串、数字和null
func (m *MessageID) UnmarshalJSON(data []byte) error {
// 尝试解析为null
if string(data) == "null" {
m.value = nil
return nil
}
// 尝试解析为字符串
var str string
if err := json.Unmarshal(data, &str); err == nil {
m.value = str
return nil
}
// 尝试解析为数字
var num json.Number
if err := json.Unmarshal(data, &num); err == nil {
m.value = num
return nil
}
return fmt.Errorf("invalid id type")
}
// MarshalJSON 自定义序列化
func (m MessageID) MarshalJSON() ([]byte, error) {
if m.value == nil {
return []byte("null"), nil
}
return json.Marshal(m.value)
}
// String 返回字符串表示
func (m MessageID) String() string {
if m.value == nil {
return ""
}
return fmt.Sprintf("%v", m.value)
}
// Value 返回原始值
func (m MessageID) Value() interface{} {
return m.value
}
// Message 表示MCP消息(符合JSON-RPC 2.0规范)
type Message struct {
ID MessageID `json:"id,omitempty"`
Type string `json:"-"` // 内部使用,不序列化到JSON
Method string `json:"method,omitempty"`
Params json.RawMessage `json:"params,omitempty"`
Result json.RawMessage `json:"result,omitempty"`
Error *Error `json:"error,omitempty"`
Version string `json:"jsonrpc,omitempty"` // JSON-RPC 2.0 版本标识
}
// Error 表示MCP错误
type Error struct {
Code int `json:"code"`
Message string `json:"message"`
Data interface{} `json:"data,omitempty"`
}
// Tool 表示MCP工具定义
type Tool struct {
Name string `json:"name"`
Description string `json:"description"` // 详细描述
ShortDescription string `json:"shortDescription,omitempty"` // 简短描述(用于工具列表,减少token消耗)
InputSchema map[string]interface{} `json:"inputSchema"`
}
// ToolCall 表示工具调用
type ToolCall struct {
Name string `json:"name"`
Arguments map[string]interface{} `json:"arguments"`
}
// ToolResult 表示工具执行结果
type ToolResult struct {
Content []Content `json:"content"`
IsError bool `json:"isError,omitempty"`
}
// Content 表示内容
type Content struct {
Type string `json:"type"`
Text string `json:"text"`
}
// InitializeRequest 初始化请求
type InitializeRequest struct {
ProtocolVersion string `json:"protocolVersion"`
Capabilities map[string]interface{} `json:"capabilities"`
ClientInfo ClientInfo `json:"clientInfo"`
}
// ClientInfo 客户端信息
type ClientInfo struct {
Name string `json:"name"`
Version string `json:"version"`
}
// InitializeResponse 初始化响应
type InitializeResponse struct {
ProtocolVersion string `json:"protocolVersion"`
Capabilities ServerCapabilities `json:"capabilities"`
ServerInfo ServerInfo `json:"serverInfo"`
}
// ServerCapabilities 服务器能力
type ServerCapabilities struct {
Tools map[string]interface{} `json:"tools,omitempty"`
Prompts map[string]interface{} `json:"prompts,omitempty"`
Resources map[string]interface{} `json:"resources,omitempty"`
Sampling map[string]interface{} `json:"sampling,omitempty"`
}
// ServerInfo 服务器信息
type ServerInfo struct {
Name string `json:"name"`
Version string `json:"version"`
}
// ListToolsRequest 列出工具请求
type ListToolsRequest struct{}
// ListToolsResponse 列出工具响应
type ListToolsResponse struct {
Tools []Tool `json:"tools"`
}
// ListPromptsResponse 列出提示词响应
type ListPromptsResponse struct {
Prompts []Prompt `json:"prompts"`
}
// ListResourcesResponse 列出资源响应
type ListResourcesResponse struct {
Resources []Resource `json:"resources"`
}
// CallToolRequest 调用工具请求
type CallToolRequest struct {
Name string `json:"name"`
Arguments map[string]interface{} `json:"arguments"`
}
// CallToolResponse 调用工具响应
type CallToolResponse struct {
Content []Content `json:"content"`
IsError bool `json:"isError,omitempty"`
}
// ToolExecution 工具执行记录
type ToolExecution struct {
ID string `json:"id"`
ToolName string `json:"toolName"`
Arguments map[string]interface{} `json:"arguments"`
Status string `json:"status"` // pending, running, completed, failed, cancelled
Result *ToolResult `json:"result,omitempty"`
Error string `json:"error,omitempty"`
StartTime time.Time `json:"startTime"`
EndTime *time.Time `json:"endTime,omitempty"`
Duration time.Duration `json:"duration,omitempty"`
// PartialOutput is a bounded tail preview of output produced by a running tool.
// It is intentionally separate from Result, which remains the final canonical tool result.
PartialOutput string `json:"partialOutput,omitempty"`
PartialOutputBytes int64 `json:"partialOutputBytes,omitempty"`
PartialOutputTruncated bool `json:"partialOutputTruncated,omitempty"`
PartialOutputUpdatedAt *time.Time `json:"partialOutputUpdatedAt,omitempty"`
// ConversationID 仅 API 展示用(进行中的 Agent 任务),不写入 tool_executions 表。
ConversationID string `json:"conversationId,omitempty"`
OwnerUserID string `json:"-"`
}
// ToolStats 工具统计信息
type ToolStats struct {
ToolName string `json:"toolName"`
TotalCalls int `json:"totalCalls"`
SuccessCalls int `json:"successCalls"`
FailedCalls int `json:"failedCalls"`
LastCallTime *time.Time `json:"lastCallTime,omitempty"`
}
// Prompt 提示词模板
type Prompt struct {
Name string `json:"name"`
Description string `json:"description,omitempty"`
Arguments []PromptArgument `json:"arguments,omitempty"`
}
// PromptArgument 提示词参数
type PromptArgument struct {
Name string `json:"name"`
Description string `json:"description,omitempty"`
Required bool `json:"required,omitempty"`
}
// GetPromptRequest 获取提示词请求
type GetPromptRequest struct {
Name string `json:"name"`
Arguments map[string]interface{} `json:"arguments,omitempty"`
}
// GetPromptResponse 获取提示词响应
type GetPromptResponse struct {
Messages []PromptMessage `json:"messages"`
}
// PromptMessage 提示词消息
type PromptMessage struct {
Role string `json:"role"`
Content string `json:"content"`
}
// Resource 资源
type Resource struct {
URI string `json:"uri"`
Name string `json:"name"`
Description string `json:"description,omitempty"`
MimeType string `json:"mimeType,omitempty"`
}
// ReadResourceRequest 读取资源请求
type ReadResourceRequest struct {
URI string `json:"uri"`
}
// ReadResourceResponse 读取资源响应
type ReadResourceResponse struct {
Contents []ResourceContent `json:"contents"`
}
// ResourceContent 资源内容
type ResourceContent struct {
URI string `json:"uri"`
MimeType string `json:"mimeType,omitempty"`
Text string `json:"text,omitempty"`
Blob string `json:"blob,omitempty"`
}
// SamplingRequest 采样请求
type SamplingRequest struct {
Messages []SamplingMessage `json:"messages"`
Model string `json:"model,omitempty"`
MaxTokens int `json:"maxTokens,omitempty"`
Temperature float64 `json:"temperature,omitempty"`
TopP float64 `json:"topP,omitempty"`
}
// SamplingMessage 采样消息
type SamplingMessage struct {
Role string `json:"role"`
Content string `json:"content"`
}
// SamplingResponse 采样响应
type SamplingResponse struct {
Content []SamplingContent `json:"content"`
Model string `json:"model,omitempty"`
StopReason string `json:"stopReason,omitempty"`
}
// SamplingContent 采样内容
type SamplingContent struct {
Type string `json:"type"`
Text string `json:"text,omitempty"`
}
// ToolResultPlainText 拼接工具结果中的文本(手动终止时作为「工具原始输出」)。
func ToolResultPlainText(r *ToolResult) string {
if r == nil || len(r.Content) == 0 {
return ""
}
var b strings.Builder
for _, c := range r.Content {
b.WriteString(c.Text)
}
return strings.TrimSpace(b.String())
}
// AbortNoteBannerForModel 标出后续文本来自「用户手动终止工具时在弹窗中填写」,避免与 stdout/stderr 混淆。
const AbortNoteBannerForModel = "---\n" +
"【用户终止说明|USER INTERRUPT NOTE】\n" +
"(以下由操作者填写,用于指示模型如何继续;不是工具原始输出。)\n" +
"Written by the operator when stopping this tool; not raw tool output.\n" +
"---"
// MergePartialToolOutputAndAbortNote 格式:工具原始输出 + 醒目标题 + 用户终止说明(无说明则原样返回 partial)。
func MergePartialToolOutputAndAbortNote(partial, userNote string) string {
partial = strings.TrimSpace(partial)
userNote = strings.TrimSpace(userNote)
if userNote == "" {
return partial
}
section := AbortNoteBannerForModel + "\n" + userNote
if partial == "" {
return section
}
return partial + "\n\n" + section
}
+266
View File
@@ -0,0 +1,266 @@
package security
import (
"database/sql"
"errors"
"strings"
"sync"
"time"
"cyberstrike-ai/internal/database"
"github.com/google/uuid"
)
// Predefined errors for authentication operations.
var (
ErrInvalidPassword = errors.New("invalid password")
)
// Session represents an authenticated user session.
type Session struct {
Token string
ExpiresAt time.Time
UserID string
Username string
DisplayName string
Roles []string
Permissions map[string]bool
PermissionScopes map[string]string
Scope string
}
// AuthManager manages password-based authentication and session lifecycle.
type AuthManager struct {
sessionDuration time.Duration
db *database.DB
mu sync.RWMutex
sessions map[string]Session
}
// NewAuthManager creates a new AuthManager instance.
func NewAuthManager(sessionDurationHours int) *AuthManager {
if sessionDurationHours <= 0 {
sessionDurationHours = 12
}
return &AuthManager{
sessionDuration: time.Duration(sessionDurationHours) * time.Hour,
sessions: make(map[string]Session),
}
}
// AttachRBACStore enables multi-user RBAC authentication. When no users exist yet,
// it bootstraps the built-in admin account and returns the generated initial password.
func (a *AuthManager) AttachRBACStore(db *database.DB) (generatedAdminPassword string, err error) {
if db == nil {
return "", errors.New("database is required for authentication")
}
needsAdminPassword, err := db.RBACNeedsAdminPassword()
if err != nil {
return "", err
}
adminPasswordHash := ""
if needsAdminPassword {
generatedAdminPassword, err = GenerateStrongPassword(24)
if err != nil {
return "", err
}
adminPasswordHash, err = HashPassword(generatedAdminPassword)
if err != nil {
return "", err
}
}
if err := db.BootstrapRBAC(adminPasswordHash, PermissionCatalog); err != nil {
return "", err
}
a.mu.Lock()
a.db = db
a.mu.Unlock()
return generatedAdminPassword, nil
}
// Authenticate validates the password and creates a new session.
func (a *AuthManager) Authenticate(username, password string) (string, time.Time, error) {
session, err := a.authenticateSession(username, password)
if err != nil {
return "", time.Time{}, err
}
a.mu.Lock()
a.sessions[session.Token] = session
a.mu.Unlock()
return session.Token, session.ExpiresAt, nil
}
func (a *AuthManager) authenticateSession(username, password string) (Session, error) {
token := uuid.NewString()
expiresAt := time.Now().Add(a.sessionDuration)
a.mu.RLock()
db := a.db
a.mu.RUnlock()
if db == nil {
return Session{}, errors.New("authentication store is not configured")
}
username = strings.TrimSpace(strings.ToLower(username))
if username == "" {
username = "admin"
}
user, err := db.GetRBACUserByUsername(username)
if err != nil {
if err == sql.ErrNoRows {
return Session{}, ErrInvalidPassword
}
return Session{}, err
}
if !user.Enabled || !VerifyPasswordHash(password, user.PasswordHash) {
return Session{}, ErrInvalidPassword
}
access, err := db.ResolveRBACAccess(user.ID)
if err != nil {
return Session{}, err
}
roleIDs := make([]string, 0, len(access.Roles))
for _, role := range access.Roles {
roleIDs = append(roleIDs, role.ID)
}
return Session{
Token: token,
ExpiresAt: expiresAt,
UserID: user.ID,
Username: user.Username,
DisplayName: user.DisplayName,
Roles: roleIDs,
Permissions: access.Permissions,
PermissionScopes: access.PermissionScopes,
Scope: access.Scope,
}, nil
}
func (s Session) ScopeFor(permission string) string {
if scope := strings.TrimSpace(s.PermissionScopes[strings.TrimSpace(permission)]); scope != "" {
return scope
}
return strings.TrimSpace(s.Scope)
}
// ValidateToken checks whether the provided token is still valid.
func (a *AuthManager) ValidateToken(token string) (Session, bool) {
if strings.TrimSpace(token) == "" {
return Session{}, false
}
a.mu.RLock()
session, ok := a.sessions[token]
a.mu.RUnlock()
if !ok {
return Session{}, false
}
if time.Now().After(session.ExpiresAt) {
a.mu.Lock()
delete(a.sessions, token)
a.mu.Unlock()
return Session{}, false
}
return session, true
}
// CheckPassword verifies whether the provided password matches the current password.
func (a *AuthManager) CheckPassword(password string) bool {
return a.CheckUserPassword("admin", password)
}
// CheckUserPassword verifies whether the provided password matches a user.
func (a *AuthManager) CheckUserPassword(username, password string) bool {
a.mu.RLock()
db := a.db
a.mu.RUnlock()
if db == nil {
return false
}
user, err := db.GetRBACUserByUsername(username)
if err != nil {
return false
}
return VerifyPasswordHash(password, user.PasswordHash)
}
func (a *AuthManager) UpdateUserPassword(userID, password string) error {
password = strings.TrimSpace(password)
if password == "" {
return errors.New("auth password must be configured")
}
hash, err := HashPassword(password)
if err != nil {
return err
}
a.mu.RLock()
db := a.db
a.mu.RUnlock()
if db == nil {
return errors.New("authentication store is not configured")
}
if err := db.UpdateRBACUserPassword(userID, hash); err != nil {
return err
}
a.mu.Lock()
for token, session := range a.sessions {
if session.UserID == userID {
delete(a.sessions, token)
}
}
a.mu.Unlock()
return nil
}
// RevokeToken invalidates the specified token.
func (a *AuthManager) RevokeToken(token string) {
if strings.TrimSpace(token) == "" {
return
}
a.mu.Lock()
delete(a.sessions, token)
a.mu.Unlock()
}
func (a *AuthManager) RevokeUserSessions(userID string) {
userID = strings.TrimSpace(userID)
if userID == "" {
return
}
a.mu.Lock()
for token, session := range a.sessions {
if session.UserID == userID {
delete(a.sessions, token)
}
}
a.mu.Unlock()
}
func (a *AuthManager) RevokeAllSessions() {
a.mu.Lock()
a.sessions = make(map[string]Session)
a.mu.Unlock()
}
// SessionDurationHours returns the configured session duration in hours.
func (a *AuthManager) SessionDurationHours() int {
return int(a.sessionDuration / time.Hour)
}
func allPermissions() map[string]bool {
out := make(map[string]bool, len(PermissionCatalog))
for key := range PermissionCatalog {
out[key] = true
}
return out
}
@@ -0,0 +1,38 @@
package security
import (
"path/filepath"
"testing"
"cyberstrike-ai/internal/database"
"go.uber.org/zap"
)
func TestAttachRBACStoreBootstrapsAdminPassword(t *testing.T) {
db, err := database.NewDB(filepath.Join(t.TempDir(), "auth-bootstrap.db"), zap.NewNop())
if err != nil {
t.Fatalf("NewDB: %v", err)
}
t.Cleanup(func() { _ = db.Close() })
manager := NewAuthManager(12)
generated, err := manager.AttachRBACStore(db)
if err != nil {
t.Fatalf("AttachRBACStore: %v", err)
}
if generated == "" {
t.Fatal("expected generated admin password on first bootstrap")
}
if !manager.CheckUserPassword("admin", generated) {
t.Fatal("generated password should authenticate admin")
}
second, err := manager.AttachRBACStore(db)
if err != nil {
t.Fatalf("AttachRBACStore second call: %v", err)
}
if second != "" {
t.Fatalf("expected no password on second bootstrap, got %q", second)
}
}
+94
View File
@@ -0,0 +1,94 @@
package security
import (
"net/http"
"net/http/httptest"
"path/filepath"
"testing"
"cyberstrike-ai/internal/authctx"
"cyberstrike-ai/internal/database"
"github.com/gin-gonic/gin"
"go.uber.org/zap"
)
func TestAuthManagerAuthenticatesCreatedRBACUser(t *testing.T) {
db, err := database.NewDB(filepath.Join(t.TempDir(), "auth-rbac.db"), zap.NewNop())
if err != nil {
t.Fatalf("NewDB: %v", err)
}
t.Cleanup(func() { _ = db.Close() })
manager := NewAuthManager(12)
if _, err := manager.AttachRBACStore(db); err != nil {
t.Fatalf("AttachRBACStore: %v", err)
}
hash, err := HashPassword("operator-secret")
if err != nil {
t.Fatalf("HashPassword: %v", err)
}
user, err := db.CreateRBACUser("operator1", "Operator One", hash, true, []string{database.RBACSystemRoleViewer})
if err != nil {
t.Fatalf("CreateRBACUser: %v", err)
}
token, _, err := manager.Authenticate("operator1", "operator-secret")
if err != nil {
t.Fatalf("Authenticate created user: %v", err)
}
session, ok := manager.ValidateToken(token)
if !ok {
t.Fatalf("expected created user session to validate")
}
if session.UserID != user.ID || session.Username != "operator1" {
t.Fatalf("session user = %s/%s, want %s/operator1", session.UserID, session.Username, user.ID)
}
if !session.Permissions["auth:self"] || !session.Permissions["chat:read"] {
t.Fatalf("expected viewer permissions in session, got %#v", session.Permissions)
}
if _, _, err := manager.Authenticate("", "operator-secret"); err == nil {
t.Fatalf("empty username must not authenticate non-admin user")
}
router := gin.New()
router.Use(AuthMiddleware(manager))
router.GET("/principal", func(c *gin.Context) {
principal, ok := authctx.PrincipalFromContext(c.Request.Context())
if !ok || principal.UserID != user.ID || !principal.HasPermission("chat:read") || principal.ScopeFor("chat:read") != database.RBACScopeAssigned {
c.Status(http.StatusInternalServerError)
return
}
c.Status(http.StatusNoContent)
})
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/principal", nil)
req.Header.Set("Authorization", "Bearer "+token)
router.ServeHTTP(w, req)
if w.Code != http.StatusNoContent {
t.Fatalf("principal propagation status = %d", w.Code)
}
}
func TestQueryTokenOnlyAllowedForSSEAndWebSocketGET(t *testing.T) {
requestToken := func(method, accept, upgrade string) string {
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest(method, "/api/test?token=secret", nil)
c.Request.Header.Set("Accept", accept)
c.Request.Header.Set("Upgrade", upgrade)
return extractTokenFromRequest(c)
}
if got := requestToken(http.MethodGet, "application/json", ""); got != "" {
t.Fatalf("ordinary GET accepted query token %q", got)
}
if got := requestToken(http.MethodPost, "text/event-stream", ""); got != "" {
t.Fatalf("POST accepted query token %q", got)
}
if got := requestToken(http.MethodGet, "text/event-stream", ""); got != "secret" {
t.Fatalf("SSE token = %q", got)
}
if got := requestToken(http.MethodGet, "", "websocket"); got != "secret" {
t.Fatalf("WebSocket token = %q", got)
}
}
+151
View File
@@ -0,0 +1,151 @@
package security
import (
"net/http"
"strings"
"cyberstrike-ai/internal/authctx"
"cyberstrike-ai/internal/database"
"github.com/gin-gonic/gin"
)
const (
ContextAuthTokenKey = "authToken"
ContextSessionExpiry = "authSessionExpiry"
ContextUserIDKey = "authUserID"
ContextUsernameKey = "authUsername"
ContextUserScopeKey = "authUserScope"
ContextSessionKey = "authSession"
)
// AuthMiddleware enforces authentication on protected routes.
func AuthMiddleware(manager *AuthManager) gin.HandlerFunc {
return func(c *gin.Context) {
token := extractTokenFromRequest(c)
session, ok := manager.ValidateToken(token)
if !ok {
c.AbortWithStatusJSON(http.StatusUnauthorized, gin.H{
"error": "未授权访问,请先登录",
})
return
}
c.Set(ContextAuthTokenKey, session.Token)
c.Set(ContextSessionExpiry, session.ExpiresAt)
c.Set(ContextUserIDKey, session.UserID)
c.Set(ContextUsernameKey, session.Username)
c.Set(ContextUserScopeKey, session.Scope)
c.Set(ContextSessionKey, session)
// Gin context values do not survive into Agent/MCP/background contexts.
// Attach an immutable principal to the request context as the canonical
// identity for every downstream execution layer.
principal := authctx.NewPrincipalWithScopes(session.UserID, session.Username, session.Scope, session.Permissions, session.PermissionScopes)
c.Request = c.Request.WithContext(authctx.WithPrincipal(c.Request.Context(), principal))
c.Next()
}
}
func RequirePermission(permission string) gin.HandlerFunc {
permission = strings.TrimSpace(permission)
return func(c *gin.Context) {
if permission == "" || SessionHasPermission(c, permission) {
c.Next()
return
}
c.AbortWithStatusJSON(http.StatusForbidden, gin.H{
"error": "权限不足",
"permission": permission,
})
}
}
func RequireAnyPermission(permissions ...string) gin.HandlerFunc {
return func(c *gin.Context) {
for _, permission := range permissions {
if SessionHasPermission(c, permission) {
c.Next()
return
}
}
c.AbortWithStatusJSON(http.StatusForbidden, gin.H{
"error": "权限不足",
"permissions": permissions,
})
}
}
func RequireResourcePermission(db *database.DB, permission, resourceType, paramName string) gin.HandlerFunc {
return func(c *gin.Context) {
if !SessionHasPermission(c, permission) {
c.AbortWithStatusJSON(http.StatusForbidden, gin.H{
"error": "权限不足",
"permission": permission,
})
return
}
if db == nil {
c.AbortWithStatusJSON(http.StatusServiceUnavailable, gin.H{"error": "资源鉴权服务不可用"})
return
}
resourceID := strings.TrimSpace(c.Param(paramName))
if resourceID == "" {
c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": "资源 ID 不能为空"})
return
}
session, ok := CurrentSession(c)
if !ok || !db.UserCanAccessResource(session.UserID, session.Scope, resourceType, resourceID) {
c.AbortWithStatusJSON(http.StatusForbidden, gin.H{
"error": "无权访问该资源",
"resource_type": resourceType,
"resource_id": resourceID,
})
return
}
c.Next()
}
}
func CurrentSession(c *gin.Context) (Session, bool) {
if c == nil {
return Session{}, false
}
v, ok := c.Get(ContextSessionKey)
if !ok {
return Session{}, false
}
session, ok := v.(Session)
return session, ok
}
func SessionHasPermission(c *gin.Context, permission string) bool {
session, ok := CurrentSession(c)
if !ok {
return false
}
return session.Permissions[permission]
}
func extractTokenFromRequest(c *gin.Context) string {
authHeader := c.GetHeader("Authorization")
if authHeader != "" {
if len(authHeader) > 7 && strings.EqualFold(authHeader[0:7], "Bearer ") {
return strings.TrimSpace(authHeader[7:])
}
return strings.TrimSpace(authHeader)
}
if token := c.Query("token"); token != "" && c.Request.Method == http.MethodGet {
acceptsSSE := strings.Contains(strings.ToLower(c.GetHeader("Accept")), "text/event-stream")
upgradesWebSocket := strings.EqualFold(strings.TrimSpace(c.GetHeader("Upgrade")), "websocket")
if acceptsSSE || upgradesWebSocket {
return strings.TrimSpace(token)
}
}
if cookie, err := c.Cookie("auth_token"); err == nil {
return strings.TrimSpace(cookie)
}
return ""
}
@@ -0,0 +1,56 @@
package security
import (
"errors"
"fmt"
"os/exec"
"strings"
)
// FormatCommandFailureResult 与 exec 工具 ToolResult 文案一致(不含 ToolErrorPrefix)。
func FormatCommandFailureResult(exitCode int, output string) string {
output = strings.TrimSpace(output)
errMsg := fmt.Sprintf("exit status %d", exitCode)
if output == "" {
return fmt.Sprintf("命令执行失败: %s", errMsg)
}
if strings.HasPrefix(output, "命令执行失败:") {
return output
}
return fmt.Sprintf("命令执行失败: %s\n输出: %s", errMsg, output)
}
// FormatCommandFailureFromErr 根据 exec/execute 返回的 error 生成统一失败文案(IsError 正文)。
func FormatCommandFailureFromErr(err error, output string) string {
if err == nil {
return strings.TrimSpace(output)
}
var exitError *exec.ExitError
if errors.As(err, &exitError) {
return FormatCommandFailureResult(exitError.ExitCode(), output)
}
output = strings.TrimSpace(output)
if output == "" {
return fmt.Sprintf("命令执行失败: %v", err)
}
if strings.HasPrefix(output, "命令执行失败:") {
return output
}
return fmt.Sprintf("命令执行失败: %v\n输出: %s", err, output)
}
// ExecuteFailureStatusLine 流式 execute 结束时追加的单行状态(输出正文已在流中推送过)。
func ExecuteFailureStatusLine(exitCode int) string {
return fmt.Sprintf("\n命令执行失败: exit status %d", exitCode)
}
// IsCommandFailureResult 判断工具结果正文是否表示命令非零退出(用于 execute / exec 对齐 isError)。
func IsCommandFailureResult(content string) bool {
return strings.Contains(content, "命令执行失败:")
}
// IsLegacyShellExitNoise 过滤旧版 shell 流中冗余的 exit code 行。
func IsLegacyShellExitNoise(s string) bool {
trimmed := strings.TrimSpace(s)
return strings.HasPrefix(trimmed, "command exited with non-zero code ")
}
@@ -0,0 +1,54 @@
package security
import (
"errors"
"os/exec"
"strings"
"testing"
)
func TestFormatCommandFailureResult(t *testing.T) {
got := FormatCommandFailureResult(1, "sudo: password required")
want := "命令执行失败: exit status 1\n输出: sudo: password required"
if got != want {
t.Fatalf("got %q want %q", got, want)
}
if FormatCommandFailureResult(2, "") != "命令执行失败: exit status 2" {
t.Fatal("empty output format")
}
if FormatCommandFailureResult(1, "命令执行失败: exit status 1") != "命令执行失败: exit status 1" {
t.Fatal("should not double-wrap")
}
}
func TestIsCommandFailureResult(t *testing.T) {
if !IsCommandFailureResult("sudo: err\n命令执行失败: exit status 1") {
t.Fatal("expected true")
}
if IsCommandFailureResult("sudo: err only") {
t.Fatal("expected false")
}
}
func TestFormatCommandFailureFromErr(t *testing.T) {
cmd := exec.Command("sh", "-c", "exit 42")
err := cmd.Run()
got := FormatCommandFailureFromErr(err, "oops")
if got != "命令执行失败: exit status 42\n输出: oops" {
t.Fatalf("got %q", got)
}
timeoutErr := errors.New("shell inactivity timeout (300s)")
got2 := FormatCommandFailureFromErr(timeoutErr, "already timed out")
if !strings.Contains(got2, "shell inactivity timeout") || !strings.Contains(got2, "already timed out") {
t.Fatalf("got %q", got2)
}
}
func TestIsLegacyShellExitNoise(t *testing.T) {
if !IsLegacyShellExitNoise("command exited with non-zero code 1\n") {
t.Fatal("expected legacy noise")
}
if IsLegacyShellExitNoise("sudo: failed") {
t.Fatal("unexpected noise")
}
}
File diff suppressed because it is too large Load Diff
+282
View File
@@ -0,0 +1,282 @@
package security
import (
"context"
"os/exec"
"runtime"
"strings"
"testing"
"time"
"cyberstrike-ai/internal/config"
"cyberstrike-ai/internal/mcp"
"go.uber.org/zap"
)
// setupTestExecutor 创建测试用的执行器
func setupTestExecutor(t *testing.T) (*Executor, *mcp.Server) {
logger := zap.NewNop()
mcpServer := mcp.NewServer(logger)
cfg := &config.SecurityConfig{
Tools: []config.ToolConfig{},
}
executor := NewExecutor(cfg, mcpServer, logger)
return executor, mcpServer
}
func TestExecutor_ExecuteInternalTool_UnknownTool(t *testing.T) {
executor, _ := setupTestExecutor(t)
ctx := context.Background()
args := map[string]interface{}{
"test": "value",
}
// 测试未知的内部工具类型
toolResult, err := executor.executeInternalTool(ctx, "unknown_tool", "internal:unknown_tool", args)
if err != nil {
t.Fatalf("执行内部工具失败: %v", err)
}
if !toolResult.IsError {
t.Fatal("未知的工具类型应该返回错误")
}
if !strings.Contains(toolResult.Content[0].Text, "未知的内部工具类型") {
t.Errorf("错误消息应该包含'未知的内部工具类型'")
}
}
func TestExecuteSystemCommand_BackgroundDoesNotBlockOnChildStdout(t *testing.T) {
executor, _ := setupTestExecutor(t)
// 子进程先向 stdout 写无换行字符再长时间 sleep;若与 echo $pid 共享管道且未重定向子进程 stdout,
// ReadString('\n') 会阻塞到子进程退出。后台包装须将子进程标准流与 PID 行分离。
ctx, cancel := context.WithTimeout(context.Background(), 4*time.Second)
defer cancel()
args := map[string]interface{}{
"command": `(sh -c 'printf x; sleep 120') &`,
"shell": "sh",
}
res, err := executor.executeSystemCommand(ctx, args)
if err != nil {
t.Fatalf("executeSystemCommand: %v", err)
}
if res == nil || res.IsError {
t.Fatalf("expected success, got %+v", res)
}
txt := res.Content[0].Text
if !strings.Contains(txt, "后台命令已启动") {
t.Fatalf("unexpected body: %q", txt)
}
}
func TestExecToolSoftWaitExposesPartialOutput(t *testing.T) {
executor, server := setupTestExecutor(t)
server.ConfigureToolWaitTimeoutSeconds(1)
mcp.RegisterExecutionControlTools(server, nil)
server.RegisterTool(mcp.Tool{Name: "exec", InputSchema: map[string]interface{}{"type": "object"}}, func(ctx context.Context, args map[string]interface{}) (*mcp.ToolResult, error) {
return executor.ExecuteTool(ctx, "exec", args)
})
result, executionID, err := server.CallTool(context.Background(), "exec", map[string]interface{}{
"command": "for i in 1 2 3 4; do echo partial-$i; sleep 0.3; done; sleep 5",
"shell": "sh",
})
if err != nil {
t.Fatalf("CallTool exec: %v", err)
}
if executionID == "" || result == nil || !result.IsError {
t.Fatalf("expected soft wait timeout, id=%q result=%#v", executionID, result)
}
status, _, err := server.CallTool(context.Background(), "get_tool_execution", map[string]interface{}{
"execution_id": executionID,
"include_partial_output": true,
"partial_output_max_bytes": 4096,
})
if err != nil {
t.Fatalf("get_tool_execution: %v", err)
}
body := mcp.ToolResultPlainText(status)
if !strings.Contains(body, `"status": "running"`) {
t.Fatalf("expected running execution, got: %s", body)
}
if !strings.Contains(body, "partial-") || !strings.Contains(body, "partial_output") {
t.Fatalf("expected partial output in execution status, got: %s", body)
}
server.CancelToolExecution(executionID)
}
func TestExecuteSystemCommand_FailureFormat(t *testing.T) {
executor, _ := setupTestExecutor(t)
res, err := executor.executeSystemCommand(context.Background(), map[string]interface{}{
"command": "echo fail-msg >&2; exit 7",
"shell": "sh",
})
if err != nil {
t.Fatalf("executeSystemCommand: %v", err)
}
if res == nil || !res.IsError {
t.Fatalf("expected IsError, got %+v", res)
}
text := res.Content[0].Text
if text != FormatCommandFailureResult(7, "fail-msg\n") && text != FormatCommandFailureResult(7, "fail-msg") {
t.Fatalf("unexpected failure text: %q", text)
}
if !strings.Contains(text, "exit status 7") || !strings.Contains(text, "fail-msg") {
t.Fatalf("unexpected failure text: %q", text)
}
}
func TestExecuteSystemCommand_OutputIsSourceLimited(t *testing.T) {
executor, _ := setupTestExecutor(t)
spillRoot := t.TempDir()
executor.SetToolOutputMaxBytes(200)
executor.SetToolOutputSpillRoot(spillRoot)
ctx := mcp.WithMCPConversationID(context.Background(), "exec-spill")
res, err := executor.executeSystemCommand(ctx, map[string]interface{}{
"command": "i=0; while [ $i -lt 2000 ]; do printf 0123456789; i=$((i+1)); done",
"shell": "sh",
})
if err != nil {
t.Fatalf("executeSystemCommand: %v", err)
}
if res == nil || res.IsError {
t.Fatalf("expected success, got %+v", res)
}
text := res.Content[0].Text
if !strings.Contains(text, "<persisted-output>") || !strings.Contains(text, "Full output saved to:") {
t.Fatalf("missing persisted-output notice: %q", text)
}
if len(text) > 200 {
t.Fatalf("output exceeded hard limit: len=%d text=%q", len(text), text)
}
if strings.Contains(text, strings.Repeat("0123456789", 20)) {
t.Fatalf("output kept too much data: len=%d", len(text))
}
}
func TestExecuteSystemCommand_StreamingOutputIsSourceLimited(t *testing.T) {
executor, _ := setupTestExecutor(t)
spillRoot := t.TempDir()
executor.SetToolOutputMaxBytes(200)
executor.SetToolOutputSpillRoot(spillRoot)
var streamed strings.Builder
ctx := context.WithValue(context.Background(), ToolOutputCallbackCtxKey, ToolOutputCallback(func(chunk string) {
streamed.WriteString(chunk)
}))
ctx = mcp.WithMCPConversationID(ctx, "exec-stream-spill")
res, err := executor.executeSystemCommand(ctx, map[string]interface{}{
"command": "i=0; while [ $i -lt 2000 ]; do printf abcdefghij; i=$((i+1)); done",
"shell": "sh",
})
if err != nil {
t.Fatalf("executeSystemCommand: %v", err)
}
text := res.Content[0].Text
if !strings.Contains(text, "<persisted-output>") {
t.Fatalf("missing persisted-output notice: %q", text)
}
if len(text) > 200 {
t.Fatalf("returned output exceeded hard limit: len=%d text=%q", len(text), text)
}
// SSE only streams the bounded prefix; final agent-facing body is the spill notice.
if len(streamed.String()) > 200 {
t.Fatalf("streamed prefix exceeded hard limit: len=%d", len(streamed.String()))
}
if streamed.Len() == 0 {
t.Fatal("expected some streamed prefix before truncation")
}
if strings.Contains(text, strings.Repeat("abcdefghij", 50)) {
t.Fatalf("returned output kept too much raw data: len=%d", len(text))
}
}
func TestBuildCommandArgs_NmapSkipsEmptyOptionalFlags(t *testing.T) {
pos1 := 1
executor, _ := setupTestExecutor(t)
toolConfig := &config.ToolConfig{
Name: "nmap",
Command: "nmap",
Args: []string{"-sT", "-sV", "-sC"},
Parameters: []config.ParameterConfig{
{Name: "target", Type: "string", Required: true, Position: &pos1, Format: "positional"},
{Name: "ports", Type: "string", Flag: "-p", Format: "flag"},
{Name: "timing", Type: "string", Template: "-T{value}", Format: "template"},
{Name: "nse_scripts", Type: "string", Flag: "--script", Format: "flag"},
{Name: "os_detection", Type: "bool", Flag: "-O", Format: "flag", Default: false},
{Name: "aggressive", Type: "bool", Flag: "-A", Format: "flag", Default: false},
{Name: "scan_type", Type: "string", Format: "template", Template: "{value}"},
{Name: "additional_args", Type: "string", Format: "positional"},
},
}
args := map[string]interface{}{
"target": "110.52.223.114",
"ports": "21, 22, 80, 443",
"timing": "4",
"nse_scripts": "",
"scan_type": "",
"os_detection": false,
"aggressive": false,
"additional_args": "-Pn",
}
cmdArgs := executor.buildCommandArgs("nmap", toolConfig, args)
joined := strings.Join(cmdArgs, " ")
if strings.Contains(joined, "--script") {
t.Fatalf("empty nse_scripts must not emit --script, got: %v", cmdArgs)
}
if !strings.Contains(joined, "110.52.223.114") {
t.Fatalf("target missing from args: %v", cmdArgs)
}
// target 应出现在 -Pn 之前,避免被误当作 --script 的参数
pnIdx := indexOf(cmdArgs, "-Pn")
targetIdx := indexOf(cmdArgs, "110.52.223.114")
if pnIdx < 0 || targetIdx < 0 || targetIdx >= pnIdx {
t.Fatalf("expected target before -Pn, got: %v", cmdArgs)
}
}
func indexOf(slice []string, s string) int {
for i, v := range slice {
if v == s {
return i
}
}
return -1
}
// TestCombinedOutputCancellable_ContextCancelKillsTree 验证 ctx 取消时能在数秒内结束(杀进程组,非挂死)。
func TestCombinedOutputCancellable_ContextCancelKillsTree(t *testing.T) {
if runtime.GOOS == "windows" {
t.Skip("unix process group kill")
}
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
cmd := exec.CommandContext(ctx, "sh", "-c", "sleep 300")
ConfigureShellCmdForAgentExecute(cmd)
done := make(chan error, 1)
go func() {
_, err := combinedOutputCancellable(ctx, cmd)
done <- err
}()
time.Sleep(150 * time.Millisecond)
cancel()
select {
case err := <-done:
if err == nil {
t.Fatal("expected context cancel error")
}
case <-time.After(5 * time.Second):
t.Fatal("combinedOutputCancellable did not return within 5s after context cancel")
}
}
+24
View File
@@ -0,0 +1,24 @@
package security
import (
"crypto/rand"
"encoding/base64"
)
// GenerateStrongPassword returns a URL-safe random password of the given length.
func GenerateStrongPassword(length int) (string, error) {
if length <= 0 {
length = 24
}
randomBytes := make([]byte, length)
if _, err := rand.Read(randomBytes); err != nil {
return "", err
}
password := base64.RawURLEncoding.EncodeToString(randomBytes)
if len(password) > length {
password = password[:length]
}
return password, nil
}
+41
View File
@@ -0,0 +1,41 @@
//go:build !windows
package security
import (
"os/exec"
"syscall"
)
// prepareShellCmdSession 让 shell 子进程在独立会话中运行,便于超时/取消时整组 SIGKILL(含子进程)。
func prepareShellCmdSession(cmd *exec.Cmd) error {
if cmd == nil {
return nil
}
if cmd.SysProcAttr == nil {
cmd.SysProcAttr = &syscall.SysProcAttr{}
}
cmd.SysProcAttr.Setsid = true
return nil
}
// terminateProcessGroup 对 rootPID 对应进程组发 SIGKILLrootPID 为 0 时回退到 cmd.Process.Pid。
func terminateProcessGroup(rootPID int, cmd *exec.Cmd) {
pid := rootPID
if pid <= 0 && cmd != nil && cmd.Process != nil {
pid = cmd.Process.Pid
}
if pid <= 0 {
return
}
if err := syscall.Kill(-pid, syscall.SIGKILL); err != nil {
if cmd != nil && cmd.Process != nil {
_ = cmd.Process.Kill()
}
}
}
// terminateCmdTree 尽力终止 cmd 及其进程组(Unix 下 Setsid 后 PGID == 首进程 PID)。
func terminateCmdTree(cmd *exec.Cmd) {
terminateProcessGroup(0, cmd)
}
+43
View File
@@ -0,0 +1,43 @@
//go:build windows
package security
import (
"os/exec"
"strconv"
"syscall"
)
func prepareShellCmdSession(cmd *exec.Cmd) error {
if cmd == nil {
return nil
}
// 独立进程组,便于 taskkill /T 终止整棵子进程树。
if cmd.SysProcAttr == nil {
cmd.SysProcAttr = &syscall.SysProcAttr{}
}
cmd.SysProcAttr.CreationFlags = syscall.CREATE_NEW_PROCESS_GROUP
return nil
}
// terminateProcessGroup 使用 taskkill /F /T 终止进程及其子进程;rootPID 为 0 时回退到 cmd.Process.Pid。
func terminateProcessGroup(rootPID int, cmd *exec.Cmd) {
pid := rootPID
if pid <= 0 && cmd != nil && cmd.Process != nil {
pid = cmd.Process.Pid
}
if pid <= 0 {
return
}
tk := exec.Command("taskkill", "/F", "/T", "/PID", strconv.Itoa(pid))
if err := tk.Run(); err != nil {
if cmd != nil && cmd.Process != nil {
_ = cmd.Process.Kill()
}
}
}
// terminateCmdTree 使用 taskkill /F /T 终止进程及其子进程(Windows 上 Process.Kill 无法保证杀掉 python 等孙进程)。
func terminateCmdTree(cmd *exec.Cmd) {
terminateProcessGroup(0, cmd)
}
+81
View File
@@ -0,0 +1,81 @@
package security
import (
"net/http"
"sync"
"time"
"github.com/gin-gonic/gin"
)
// rateLimitEntry 记录某个 IP 的请求窗口信息
type rateLimitEntry struct {
count int
windowAt time.Time
}
// RateLimiter 基于 IP 的滑动窗口速率限制器
type RateLimiter struct {
mu sync.Mutex
entries map[string]*rateLimitEntry
limit int // 窗口内允许的最大请求数
window time.Duration // 窗口时长
}
// NewRateLimiter 创建速率限制器
func NewRateLimiter(limit int, window time.Duration) *RateLimiter {
rl := &RateLimiter{
entries: make(map[string]*rateLimitEntry),
limit: limit,
window: window,
}
// 后台定期清理过期条目,防止内存泄漏
go rl.cleanup()
return rl
}
// cleanup 每分钟清理一次过期条目
func (rl *RateLimiter) cleanup() {
ticker := time.NewTicker(1 * time.Minute)
defer ticker.Stop()
for range ticker.C {
rl.mu.Lock()
now := time.Now()
for ip, entry := range rl.entries {
if now.Sub(entry.windowAt) > rl.window {
delete(rl.entries, ip)
}
}
rl.mu.Unlock()
}
}
// allow 检查指定 IP 是否允许通过
func (rl *RateLimiter) allow(ip string) bool {
rl.mu.Lock()
defer rl.mu.Unlock()
now := time.Now()
entry, ok := rl.entries[ip]
if !ok || now.Sub(entry.windowAt) > rl.window {
rl.entries[ip] = &rateLimitEntry{count: 1, windowAt: now}
return true
}
entry.count++
return entry.count <= rl.limit
}
// RateLimitMiddleware 返回 Gin 中间件,对超限请求返回 429
func RateLimitMiddleware(rl *RateLimiter) gin.HandlerFunc {
return func(c *gin.Context) {
ip := c.ClientIP()
if !rl.allow(ip) {
c.AbortWithStatusJSON(http.StatusTooManyRequests, gin.H{
"error": "rate limit exceeded, please try again later",
})
return
}
c.Next()
}
}
+119
View File
@@ -0,0 +1,119 @@
package security
import (
"crypto/sha256"
"crypto/subtle"
"encoding/hex"
"fmt"
"strings"
"golang.org/x/crypto/bcrypt"
)
// Platform permissions use module:action naming. They are intentionally
// separate from AI testing roles under roles/.
var PermissionCatalog = map[string]string{
"auth:self": "Manage own session and password",
"dashboard:read": "View dashboard summaries",
"chat:read": "View conversations",
"chat:write": "Create and update conversations",
"chat:delete": "Delete conversations and turns",
"agent:execute": "Run AI agents and workflows",
"agent:local-execute": "Use local filesystem, shell, and configured command tools from an agent",
"hitl:read": "View HITL queues and logs",
"hitl:write": "Approve, dismiss, and configure HITL",
"tasks:read": "View task queues",
"tasks:write": "Create and run task queues",
"tasks:delete": "Delete task queues",
"project:read": "View projects and project facts",
"project:write": "Create and update projects and facts",
"project:delete": "Delete projects and facts",
"vulnerability:read": "View vulnerabilities",
"vulnerability:write": "Create and update vulnerabilities",
"vulnerability:delete": "Delete vulnerabilities",
"asset:read": "View managed assets and asset summaries",
"asset:write": "Create, import, and update assets",
"asset:delete": "Delete managed assets",
"webshell:read": "View WebShell connections",
"webshell:write": "Manage and use WebShell connections",
"webshell:delete": "Delete WebShell connections",
"c2:read": "View C2 listeners, sessions, tasks, events, and profiles",
"c2:write": "Operate C2 listeners, sessions, tasks, payloads, files, and profiles",
"c2:delete": "Delete C2 objects",
"mcp:read": "View MCP status and external MCP configuration",
"mcp:execute": "Invoke the authenticated MCP endpoint",
"mcp:external:execute": "Invoke tools exposed by configured external MCP servers",
"mcp:write": "Manage external MCP server configuration and lifecycle",
"knowledge:read": "View knowledge base and retrieval logs",
"knowledge:write": "Create, update, index, and scan knowledge base",
"knowledge:delete": "Delete knowledge items and retrieval logs",
"skills:read": "View skills and skill stats",
"skills:write": "Create and update skills",
"skills:delete": "Delete skills and stats",
"agents:read": "View markdown agents",
"agents:write": "Create and update markdown agents",
"agents:delete": "Delete markdown agents",
"roles:read": "View AI testing roles",
"roles:write": "Create and update AI testing roles",
"roles:delete": "Delete AI testing roles",
"workflow:read": "View workflow definitions and runs",
"workflow:execute": "Validate, dry-run, and resume authorized workflow runs",
"workflow:write": "Create and update workflow definitions",
"workflow:delete": "Delete workflows",
"config:read": "View system configuration",
"config:write": "Update and apply system configuration",
"terminal:execute": "Run terminal commands",
"audit:read": "View and export audit logs",
"audit:delete": "Delete audit logs",
"rbac:read": "View users, platform roles, permissions, and assignments",
"rbac:write": "Manage users, platform roles, permissions, and assignments",
"notification:read": "View notifications",
"notification:write": "Mark notifications as read",
"robot:read": "View robot binding status",
"robot:write": "Manage robot bindings and test robot callbacks",
"files:read": "View chat uploads",
"files:write": "Upload, edit, and rename chat files",
"files:delete": "Delete chat files",
"attackchain:read": "View attack chains",
"attackchain:write": "Regenerate attack chains",
"fofa:execute": "Run FOFA searches and query parsing",
"openapi:read": "Read OpenAPI aggregation results",
"group:read": "View conversation groups",
"group:write": "Create and update conversation groups",
"group:delete": "Delete conversation groups",
"monitor:read": "View execution monitor",
"monitor:write": "Cancel monitor executions",
"monitor:delete": "Delete monitor executions",
}
func HashPassword(password string) (string, error) {
password = strings.TrimSpace(password)
if password == "" {
return "", fmt.Errorf("password is empty")
}
hash, err := bcrypt.GenerateFromPassword([]byte(password), bcrypt.DefaultCost)
if err != nil {
return "", err
}
return string(hash), nil
}
func VerifyPasswordHash(password, encoded string) bool {
if strings.HasPrefix(encoded, "$2a$") || strings.HasPrefix(encoded, "$2b$") || strings.HasPrefix(encoded, "$2y$") {
return bcrypt.CompareHashAndPassword([]byte(encoded), []byte(strings.TrimSpace(password))) == nil
}
parts := strings.Split(encoded, "$")
if len(parts) != 3 || parts[0] != "sha256" {
return false
}
salt, err := hex.DecodeString(parts[1])
if err != nil {
return false
}
expected, err := hex.DecodeString(parts[2])
if err != nil {
return false
}
sum := sha256.Sum256(append(salt, []byte(strings.TrimSpace(password))...))
return subtle.ConstantTimeCompare(sum[:], expected) == 1
}
+282
View File
@@ -0,0 +1,282 @@
package security
import (
"net/http"
"strings"
"cyberstrike-ai/internal/database"
"github.com/gin-gonic/gin"
)
// RBACMiddleware maps protected API routes to platform permissions. It keeps
// enforcement centralized so route declarations stay readable.
func RBACMiddleware(db *database.DB) gin.HandlerFunc {
return RBACMiddlewareWithDenyHook(db, nil)
}
type RBACDenyHook func(c *gin.Context, reason, permission string)
func RBACMiddlewareWithDenyHook(db *database.DB, denyHook RBACDenyHook) gin.HandlerFunc {
return func(c *gin.Context) {
permission := permissionForRequest(c.Request.Method, c.FullPath())
if permission == "" {
if denyHook != nil {
denyHook(c, "unmapped_route", "")
}
c.AbortWithStatusJSON(http.StatusForbidden, gin.H{
"error": "未配置访问权限",
})
return
}
permission, allowed := sessionHasRoutePermission(c, c.Request.Method, c.FullPath())
if !allowed {
if denyHook != nil {
denyHook(c, "permission_denied", permission)
}
c.AbortWithStatusJSON(http.StatusForbidden, gin.H{
"error": "权限不足",
"permission": permission,
})
return
}
// Bind the scope of the permission authorizing this request. Scope is
// permission-specific; using the user's broadest role scope here would
// let an unrelated global read role widen a write permission.
session, _ := CurrentSession(c)
session.Scope = session.ScopeFor(permission)
c.Set(ContextSessionKey, session)
c.Set(ContextUserScopeKey, session.Scope)
if db != nil && !resourceAllowed(c, db) {
if denyHook != nil {
denyHook(c, "resource_denied", permission)
}
c.AbortWithStatusJSON(http.StatusForbidden, gin.H{"error": "无权访问该资源"})
return
}
c.Next()
}
}
func sessionHasRoutePermission(c *gin.Context, method, fullPath string) (string, bool) {
path := strings.TrimPrefix(fullPath, "/api")
if alts := permissionAlternativesForRequest(method, path); len(alts) > 0 {
for _, permission := range alts {
if SessionHasPermission(c, permission) {
return permission, true
}
}
return alts[0], false
}
permission := permissionForRequest(method, fullPath)
if permission == "" {
return "", false
}
return permission, SessionHasPermission(c, permission)
}
func permissionAlternativesForRequest(method, path string) []string {
if method != http.MethodGet && method != http.MethodHead {
return nil
}
switch {
case strings.HasPrefix(path, "/config/tools"):
// MCP 管理页只需 mcp:read;系统设置页仍可用 config:read 访问同一接口。
return []string{"mcp:read", "config:read"}
default:
return nil
}
}
func permissionForRequest(method, fullPath string) string {
path := strings.TrimPrefix(fullPath, "/api")
switch {
case path == "/rbac/me":
return "auth:self"
case path == "/rbac/resources":
// The picker enumerates resource names and IDs and is only needed by
// administrators who can actually create assignments.
return "rbac:write"
case strings.HasPrefix(path, "/rbac"):
if method == http.MethodGet {
return "rbac:read"
}
return "rbac:write"
case strings.HasPrefix(path, "/robot/wechat/status"):
return "robot:read"
case strings.HasPrefix(path, "/robot"):
return "robot:write"
case strings.HasPrefix(path, "/eino-agent"), strings.HasPrefix(path, "/multi-agent"):
if strings.Contains(path, "/markdown-agents") {
return crudPermission(method, "agents")
}
return "agent:execute"
case strings.HasPrefix(path, "/hitl"):
if method == http.MethodGet || method == http.MethodHead {
return "hitl:read"
}
return "hitl:write"
case strings.HasPrefix(path, "/agent-loop"), strings.HasPrefix(path, "/batch-tasks"):
return crudPermission(method, "tasks")
case strings.HasPrefix(path, "/conversations"), strings.HasPrefix(path, "/messages"), strings.HasPrefix(path, "/process-details"):
return crudPermission(method, "chat")
case strings.HasPrefix(path, "/groups"):
return crudPermission(method, "group")
case strings.HasPrefix(path, "/monitor"):
return crudPermission(method, "monitor")
case strings.HasPrefix(path, "/notifications"):
if method == http.MethodGet {
return "notification:read"
}
return "notification:write"
case strings.HasPrefix(path, "/config"):
return crudPermission(method, "config")
case strings.HasPrefix(path, "/terminal"):
return "terminal:execute"
case strings.HasPrefix(path, "/audit"):
return crudPermission(method, "audit")
case path == "/mcp":
return "mcp:execute"
case strings.HasPrefix(path, "/external-mcp"):
if method == http.MethodGet || method == http.MethodHead {
return "mcp:read"
}
return "mcp:write"
case strings.HasPrefix(path, "/attack-chain"):
return crudPermission(method, "attackchain")
case strings.HasPrefix(path, "/knowledge"):
if path == "/knowledge/search" {
return "knowledge:read"
}
return crudPermission(method, "knowledge")
case strings.HasPrefix(path, "/vulnerabilities"):
return crudPermission(method, "vulnerability")
case path == "/assets/batch-delete", path == "/assets/merge":
return "asset:delete"
case strings.HasPrefix(path, "/assets"):
return crudPermission(method, "asset")
case strings.HasPrefix(path, "/vulnerability-alerts"):
// This endpoint only changes the authenticated user's own preference.
return "vulnerability:read"
case strings.HasPrefix(path, "/projects"):
return crudPermission(method, "project")
case strings.HasPrefix(path, "/webshell"):
return crudPermission(method, "webshell")
case strings.HasPrefix(path, "/c2"):
return crudPermission(method, "c2")
case strings.HasPrefix(path, "/chat-uploads"):
return crudPermission(method, "files")
case strings.HasPrefix(path, "/roles"):
return crudPermission(method, "roles")
case path == "/workflows/:id/package":
return "workflow:read"
case strings.HasPrefix(path, "/workflow-package-inspections"), strings.HasPrefix(path, "/workflow-package-imports"):
return "workflow:write"
case path == "/workflows/generate-draft":
return "workflow:write"
case strings.HasPrefix(path, "/workflows"):
if path == "/workflows/validate" || path == "/workflows/dry-run" || strings.HasSuffix(path, "/resume") {
return "workflow:execute"
}
return crudPermission(method, "workflow")
case strings.HasPrefix(path, "/skills"):
return crudPermission(method, "skills")
case strings.HasPrefix(path, "/openapi"):
return "openapi:read"
case strings.HasPrefix(path, "/fofa"):
return "fofa:execute"
default:
return ""
}
}
func crudPermission(method, module string) string {
switch method {
case http.MethodGet, http.MethodHead:
return module + ":read"
case http.MethodDelete:
return module + ":delete"
default:
return module + ":write"
}
}
func resourceAllowed(c *gin.Context, db *database.DB) bool {
session, ok := CurrentSession(c)
if !ok || session.Scope == database.RBACScopeAll {
return ok
}
path := strings.TrimPrefix(c.FullPath(), "/api")
switch {
case path == "/monitor/stats", path == "/monitor/calls-timeline":
// These APIs currently operate on process-global state. Until every MCP
// invocation and persisted execution record carries an immutable owner,
// allowing an assigned/own-scoped session would be a cross-user bypass.
return session.Scope == database.RBACScopeAll
case strings.HasPrefix(path, "/c2/profiles") && c.Request.Method != http.MethodGet:
return session.Scope == database.RBACScopeAll
case (strings.HasPrefix(path, "/hitl/tool-whitelist") || strings.HasPrefix(path, "/hitl/default-reviewer") || strings.HasPrefix(path, "/hitl/audit-strategy")) && c.Request.Method != http.MethodGet:
return session.Scope == database.RBACScopeAll
case isMutationMethod(c.Request.Method) && isProcessGlobalMutationPath(path):
// These definitions/configurations are shared by every user and do not
// carry owners. A module write permission with assigned/own scope must
// not silently become a process-global administrative capability.
return session.Scope == database.RBACScopeAll
case strings.HasPrefix(path, "/projects/:id"):
return db.UserCanAccessResource(session.UserID, session.Scope, "project", c.Param("id"))
case strings.HasPrefix(path, "/conversations/:id"):
return db.UserCanAccessResource(session.UserID, session.Scope, "conversation", c.Param("id"))
case strings.HasPrefix(path, "/messages/:id/process-details"):
return db.UserCanAccessMessage(session.UserID, session.Scope, c.Param("id"))
case strings.HasPrefix(path, "/process-details/:id"):
return db.UserCanAccessProcessDetail(session.UserID, session.Scope, c.Param("id"))
case strings.HasPrefix(path, "/attack-chain/:conversationId"):
return db.UserCanAccessResource(session.UserID, session.Scope, "conversation", c.Param("conversationId"))
case strings.HasPrefix(path, "/webshell/connections/:id"):
return db.UserCanAccessResource(session.UserID, session.Scope, "webshell", c.Param("id"))
case strings.HasPrefix(path, "/batch-tasks/:queueId"):
return db.UserCanAccessResource(session.UserID, session.Scope, "batch_task", c.Param("queueId"))
case strings.HasPrefix(path, "/vulnerabilities/:id"):
return db.UserCanAccessResource(session.UserID, session.Scope, "vulnerability", c.Param("id"))
case strings.HasPrefix(path, "/assets/:id"):
return db.UserCanAccessResource(session.UserID, session.Scope, "asset", c.Param("id"))
case strings.HasPrefix(path, "/c2/listeners/:id"):
return db.UserCanAccessResource(session.UserID, session.Scope, "c2_listener", c.Param("id"))
case strings.HasPrefix(path, "/c2/sessions/:id"):
return db.UserCanAccessResource(session.UserID, session.Scope, "c2_session", c.Param("id"))
case strings.HasPrefix(path, "/c2/tasks/:id"):
return db.UserCanAccessResource(session.UserID, session.Scope, "c2_task", c.Param("id"))
default:
return true
}
}
func isMutationMethod(method string) bool {
switch method {
case http.MethodPost, http.MethodPut, http.MethodPatch, http.MethodDelete:
return true
default:
return false
}
}
func isProcessGlobalMutationPath(path string) bool {
if strings.HasPrefix(path, "/roles") || strings.HasPrefix(path, "/skills") ||
strings.HasPrefix(path, "/external-mcp") || strings.HasPrefix(path, "/robot") {
return true
}
if strings.HasPrefix(path, "/workflows") {
// Workflow runs inherit conversation access; definitions are global.
return !strings.HasPrefix(path, "/workflows/runs/") && path != "/workflows/validate" && path != "/workflows/dry-run" && path != "/workflows/generate-draft"
}
if strings.HasPrefix(path, "/workflow-package-inspections") || strings.HasPrefix(path, "/workflow-package-imports") {
return true
}
if strings.HasPrefix(path, "/knowledge") {
return path != "/knowledge/search"
}
if strings.HasPrefix(path, "/eino-agent/markdown-agents") || strings.HasPrefix(path, "/multi-agent/markdown-agents") {
return true
}
return false
}
+269
View File
@@ -0,0 +1,269 @@
package security
import (
"net/http"
"net/http/httptest"
"testing"
"cyberstrike-ai/internal/database"
"github.com/gin-gonic/gin"
)
func TestRBACMiddlewareUsesMatchedFullPath(t *testing.T) {
gin.SetMode(gin.TestMode)
router := gin.New()
router.Use(func(c *gin.Context) {
c.Set(ContextSessionKey, Session{
UserID: "u1",
Username: "operator",
Permissions: map[string]bool{"project:read": true},
Scope: database.RBACScopeAll,
})
c.Next()
})
router.Use(RBACMiddleware(nil))
router.GET("/api/projects/:id", func(c *gin.Context) {
c.JSON(http.StatusOK, gin.H{"ok": true})
})
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/api/projects/p1", nil)
router.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("status = %d, want %d: %s", w.Code, http.StatusOK, w.Body.String())
}
}
func TestRBACMiddlewareRejectsMissingPermission(t *testing.T) {
gin.SetMode(gin.TestMode)
router := gin.New()
router.Use(func(c *gin.Context) {
c.Set(ContextSessionKey, Session{
UserID: "u1",
Username: "viewer",
Permissions: map[string]bool{"project:read": true},
Scope: database.RBACScopeAll,
})
c.Next()
})
router.Use(RBACMiddleware(nil))
router.POST("/api/projects", func(c *gin.Context) {
c.JSON(http.StatusOK, gin.H{"ok": true})
})
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodPost, "/api/projects", nil)
router.ServeHTTP(w, req)
if w.Code != http.StatusForbidden {
t.Fatalf("status = %d, want %d", w.Code, http.StatusForbidden)
}
}
func TestRBACMiddlewareRejectsUnmappedProtectedRoute(t *testing.T) {
gin.SetMode(gin.TestMode)
router := gin.New()
router.Use(func(c *gin.Context) {
c.Set(ContextSessionKey, Session{
UserID: "u1",
Username: "admin",
Permissions: allPermissions(),
Scope: database.RBACScopeAll,
})
c.Next()
})
router.Use(RBACMiddleware(nil))
router.GET("/api/new-module", func(c *gin.Context) {
c.JSON(http.StatusOK, gin.H{"ok": true})
})
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/api/new-module", nil)
router.ServeHTTP(w, req)
if w.Code != http.StatusForbidden {
t.Fatalf("status = %d, want %d", w.Code, http.StatusForbidden)
}
}
func TestRBACMiddlewareMapsOpenAPISpec(t *testing.T) {
gin.SetMode(gin.TestMode)
router := gin.New()
router.Use(func(c *gin.Context) {
c.Set(ContextSessionKey, Session{
UserID: "u1",
Username: "viewer",
Permissions: map[string]bool{"openapi:read": true},
Scope: database.RBACScopeAll,
})
c.Next()
})
router.Use(RBACMiddleware(nil))
router.GET("/api/openapi/spec", func(c *gin.Context) {
c.JSON(http.StatusOK, gin.H{"ok": true})
})
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/api/openapi/spec", nil)
router.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("status = %d, want %d: %s", w.Code, http.StatusOK, w.Body.String())
}
}
func TestRBACResourcePickerRequiresWritePermission(t *testing.T) {
if got := permissionForRequest(http.MethodGet, "/api/rbac/resources"); got != "rbac:write" {
t.Fatalf("picker permission = %q, want rbac:write", got)
}
if got := permissionForRequest(http.MethodGet, "/api/rbac/resource-assignments"); got != "rbac:read" {
t.Fatalf("assignment list permission = %q, want rbac:read", got)
}
}
func TestMCPInvocationPermissionIsSeparateFromMCPAdministration(t *testing.T) {
if got := permissionForRequest(http.MethodPost, "/api/mcp"); got != "mcp:execute" {
t.Fatalf("MCP invocation permission = %q, want mcp:execute", got)
}
if got := permissionForRequest(http.MethodPut, "/api/external-mcp/example"); got != "mcp:write" {
t.Fatalf("external MCP admin permission = %q, want mcp:write", got)
}
}
func TestConfigToolsReadAllowsMCPReadWithoutConfigRead(t *testing.T) {
gin.SetMode(gin.TestMode)
router := gin.New()
router.Use(func(c *gin.Context) {
c.Set(ContextSessionKey, Session{
UserID: "viewer",
Username: "viewer",
Permissions: map[string]bool{"mcp:read": true},
Scope: database.RBACScopeAssigned,
})
c.Next()
})
router.Use(RBACMiddleware(nil))
router.GET("/api/config/tools", func(c *gin.Context) {
c.JSON(http.StatusOK, gin.H{"tools": []any{}})
})
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, "/api/config/tools", nil)
router.ServeHTTP(w, req)
if w.Code != http.StatusOK {
t.Fatalf("status = %d, want %d: %s", w.Code, http.StatusOK, w.Body.String())
}
}
func TestWorkflowRunPermissionIsSeparateFromDefinitionManagement(t *testing.T) {
if got := permissionForRequest(http.MethodPost, "/api/workflows/runs/run-1/resume"); got != "workflow:execute" {
t.Fatalf("resume permission = %q, want workflow:execute", got)
}
if got := permissionForRequest(http.MethodPost, "/api/workflows/generate-draft"); got != "workflow:write" {
t.Fatalf("generate draft permission = %q, want workflow:write", got)
}
if got := permissionForRequest(http.MethodPut, "/api/workflows/workflow-1"); got != "workflow:write" {
t.Fatalf("definition permission = %q, want workflow:write", got)
}
if isProcessGlobalMutationPath("/workflows/generate-draft") {
t.Fatalf("generate draft should not be treated as a process-global mutation")
}
}
func TestRBACDenyHookReceivesDeniedDecision(t *testing.T) {
gin.SetMode(gin.TestMode)
called := false
router := gin.New()
router.Use(func(c *gin.Context) {
c.Set(ContextSessionKey, Session{UserID: "viewer", Permissions: map[string]bool{"project:read": true}, Scope: database.RBACScopeAssigned})
c.Next()
})
router.Use(RBACMiddlewareWithDenyHook(nil, func(_ *gin.Context, reason, permission string) {
called = reason == "permission_denied" && permission == "project:write"
}))
router.POST("/api/projects", func(c *gin.Context) { c.Status(http.StatusNoContent) })
w := httptest.NewRecorder()
router.ServeHTTP(w, httptest.NewRequest(http.MethodPost, "/api/projects", nil))
if w.Code != http.StatusForbidden || !called {
t.Fatalf("denial = status %d, hook called %v", w.Code, called)
}
}
func TestRBACMiddlewareBindsPermissionSpecificScope(t *testing.T) {
gin.SetMode(gin.TestMode)
router := gin.New()
router.Use(func(c *gin.Context) {
c.Set(ContextSessionKey, Session{
UserID: "mixed", Scope: database.RBACScopeAll,
Permissions: map[string]bool{"project:read": true, "project:write": true},
PermissionScopes: map[string]string{"project:read": database.RBACScopeAll, "project:write": database.RBACScopeOwn},
})
c.Next()
})
router.Use(RBACMiddleware(nil))
handler := func(c *gin.Context) {
session, _ := CurrentSession(c)
c.String(http.StatusOK, session.Scope)
}
router.GET("/api/projects/:id", handler)
router.PUT("/api/projects/:id", handler)
for _, tc := range []struct{ method, want string }{
{http.MethodGet, database.RBACScopeAll},
{http.MethodPut, database.RBACScopeOwn},
} {
w := httptest.NewRecorder()
router.ServeHTTP(w, httptest.NewRequest(tc.method, "/api/projects/p1", nil))
if w.Code != http.StatusOK || w.Body.String() != tc.want {
t.Fatalf("%s scope response = %d/%q, want 200/%q", tc.method, w.Code, w.Body.String(), tc.want)
}
}
}
func TestRBACMiddlewareRejectsAssignedScopeForGlobalMonitorAggregates(t *testing.T) {
gin.SetMode(gin.TestMode)
for _, tc := range []struct {
method string
path string
permission string
}{
{method: http.MethodGet, path: "/api/monitor/stats", permission: "monitor:read"},
} {
t.Run(tc.path, func(t *testing.T) {
router := gin.New()
router.Use(func(c *gin.Context) {
c.Set(ContextSessionKey, Session{
UserID: "assigned-user", Permissions: map[string]bool{tc.permission: true}, Scope: database.RBACScopeAssigned,
})
c.Next()
})
router.Use(RBACMiddleware(&database.DB{}))
router.Handle(tc.method, tc.path, func(c *gin.Context) { c.Status(http.StatusOK) })
w := httptest.NewRecorder()
router.ServeHTTP(w, httptest.NewRequest(tc.method, tc.path, nil))
if w.Code != http.StatusForbidden {
t.Fatalf("status = %d, want %d", w.Code, http.StatusForbidden)
}
})
}
}
func TestAssignedScopeCannotMutateProcessGlobalAssets(t *testing.T) {
gin.SetMode(gin.TestMode)
for _, path := range []string{"/api/roles/demo", "/api/skills/demo", "/api/external-mcp/demo", "/api/workflows/demo", "/api/knowledge/items/demo"} {
t.Run(path, func(t *testing.T) {
permission := permissionForRequest(http.MethodPut, path)
router := gin.New()
router.Use(func(c *gin.Context) {
c.Set(ContextSessionKey, Session{UserID: "operator", Scope: database.RBACScopeAssigned, Permissions: map[string]bool{permission: true}, PermissionScopes: map[string]string{permission: database.RBACScopeAssigned}})
c.Next()
})
router.Use(RBACMiddleware(&database.DB{}))
router.PUT(path, func(c *gin.Context) { c.Status(http.StatusNoContent) })
w := httptest.NewRecorder()
router.ServeHTTP(w, httptest.NewRequest(http.MethodPut, path, nil))
if w.Code != http.StatusForbidden {
t.Fatalf("global mutation status = %d, want 403", w.Code)
}
})
}
}
+60
View File
@@ -0,0 +1,60 @@
package security
import (
"go/ast"
"go/parser"
"go/token"
"net/http"
"path/filepath"
"strconv"
"testing"
)
func TestEveryProtectedRouteHasCatalogPermission(t *testing.T) {
file, err := parser.ParseFile(token.NewFileSet(), filepath.Join("..", "app", "app.go"), nil, 0)
if err != nil {
t.Fatal(err)
}
methods := map[string]string{
"GET": http.MethodGet, "POST": http.MethodPost, "PUT": http.MethodPut,
"PATCH": http.MethodPatch, "DELETE": http.MethodDelete,
}
prefixes := map[string]string{"protected": "", "c2Routes": "/c2", "knowledgeRoutes": "/knowledge"}
found := 0
ast.Inspect(file, func(node ast.Node) bool {
call, ok := node.(*ast.CallExpr)
if !ok || len(call.Args) == 0 {
return true
}
sel, ok := call.Fun.(*ast.SelectorExpr)
if !ok {
return true
}
ident, ok := sel.X.(*ast.Ident)
if !ok {
return true
}
prefix, protected := prefixes[ident.Name]
method, routeMethod := methods[sel.Sel.Name]
literal, literalPath := call.Args[0].(*ast.BasicLit)
if !protected || !routeMethod || !literalPath || literal.Kind != token.STRING {
return true
}
path, err := strconv.Unquote(literal.Value)
if err != nil {
t.Errorf("invalid route literal %s", literal.Value)
return true
}
found++
permission := permissionForRequest(method, "/api"+prefix+path)
if permission == "" {
t.Errorf("unmapped protected route: %s %s%s", method, prefix, path)
} else if _, ok := PermissionCatalog[permission]; !ok {
t.Errorf("route %s %s%s maps to unknown permission %q", method, prefix, path, permission)
}
return true
})
if found < 100 {
t.Fatalf("route inventory unexpectedly small: %d", found)
}
}
+111
View File
@@ -0,0 +1,111 @@
package security
import "strings"
const backgroundJobStdioRedirect = " </dev/null >/dev/null 2>&1"
// findStandaloneAmpersandPositions 返回不在引号内的独立 & 下标(排除 &&)。
func findStandaloneAmpersandPositions(command string) []int {
command = strings.TrimSpace(command)
if command == "" {
return nil
}
var positions []int
inSingleQuote := false
inDoubleQuote := false
escaped := false
for i := 0; i < len(command); i++ {
r := command[i]
if escaped {
escaped = false
continue
}
if r == '\\' {
escaped = true
continue
}
if r == '\'' && !inDoubleQuote {
inSingleQuote = !inSingleQuote
continue
}
if r == '"' && !inSingleQuote {
inDoubleQuote = !inDoubleQuote
continue
}
if r != '&' || inSingleQuote || inDoubleQuote {
continue
}
if i+1 < len(command) && command[i+1] == '&' {
continue
}
if i > 0 && command[i-1] == '&' {
continue
}
isStandalone := i == 0
if !isStandalone {
prev := command[i-1]
isStandalone = prev == ' ' || prev == '\t' || prev == '\n' || prev == '\r'
}
if !isStandalone {
continue
}
if i == len(command)-1 {
positions = append(positions, i)
continue
}
next := command[i+1]
if next == ' ' || next == '\t' || next == '\n' || next == '\r' {
positions = append(positions, i)
}
}
return positions
}
func segmentHasStdioRedirect(segment string) bool {
lower := strings.ToLower(strings.TrimSpace(segment))
if lower == "" {
return false
}
if strings.Contains(lower, ">/dev/null") || strings.Contains(lower, "2>/dev/null") {
return true
}
if strings.Contains(lower, "&>") || strings.Contains(lower, "&>>") {
return true
}
if strings.Contains(lower, "2>&1") && strings.Contains(lower, "/dev/null") {
return true
}
return false
}
// RedirectBackgroundJobStdio 为每个独立 & 前的后台段注入 </dev/null >/dev/null 2>&1
// 避免后台子进程占用 execute/exec 管道导致挂死。
func RedirectBackgroundJobStdio(command string) string {
positions := findStandaloneAmpersandPositions(command)
if len(positions) == 0 {
return command
}
out := command
for j := len(positions) - 1; j >= 0; j-- {
i := positions[j]
before := out[:i]
after := out[i:]
trimmed := strings.TrimRight(before, " \t\r\n")
if segmentHasStdioRedirect(trimmed) {
continue
}
trailing := before[len(trimmed):]
out = trimmed + backgroundJobStdioRedirect + trailing + after
}
return out
}
// PrepareShellCommandForExecute 组合 execute/exec 用的非交互包装与后台 IO 重定向。
// 须先注入 exec </dev/null,再改写 & 后台段,否则段内 </dev/null 会使 stdin 重定向被误判为已存在。
func PrepareShellCommandForExecute(shellCommand string) string {
return RedirectBackgroundJobStdio(PrepareNonInteractiveShellCommand(shellCommand))
}
@@ -0,0 +1,64 @@
package security
import (
"strings"
"testing"
)
func TestRedirectBackgroundJobStdio_mixedCommand(t *testing.T) {
in := "java -jar app.jar & JRMP_PID=$!; echo started"
out := RedirectBackgroundJobStdio(in)
if !strings.Contains(out, "java -jar app.jar </dev/null >/dev/null 2>&1 &") {
t.Fatalf("expected redirect before &: %q", out)
}
if !strings.Contains(out, "echo started") {
t.Fatalf("foreground tail preserved: %q", out)
}
}
func TestRedirectBackgroundJobStdio_trailingOnly(t *testing.T) {
in := "sleep 120 &"
out := RedirectBackgroundJobStdio(in)
want := "sleep 120 </dev/null >/dev/null 2>&1 &"
if strings.TrimSpace(out) != want {
t.Fatalf("got %q want %q", out, want)
}
}
func TestRedirectBackgroundJobStdio_skipsAlreadyRedirected(t *testing.T) {
in := "sleep 1 >/dev/null 2>&1 & echo ok"
out := RedirectBackgroundJobStdio(in)
if out != in {
t.Fatalf("should not double-redirect: %q", out)
}
}
func TestRedirectBackgroundJobStdio_skipsAndAnd(t *testing.T) {
in := "test -f /etc/passwd && echo ok"
out := RedirectBackgroundJobStdio(in)
if out != in {
t.Fatalf("&& must not be treated as background &: %q", out)
}
}
func TestPrepareShellCommandForExecute(t *testing.T) {
out := PrepareShellCommandForExecute("java -jar x & echo hi")
if !strings.Contains(out, "exec </dev/null") {
t.Fatalf("missing stdin redirect: %q", out)
}
if !strings.Contains(out, "GIT_PAGER=cat") {
t.Fatalf("missing pager export: %q", out)
}
if !strings.Contains(out, "java -jar x </dev/null >/dev/null 2>&1 &") {
t.Fatalf("missing background redirect: %q", out)
}
}
func TestIsBackgroundShellCommand_usesSharedParser(t *testing.T) {
if !IsBackgroundShellCommand("sleep 1 &") {
t.Fatal("trailing & should be background")
}
if IsBackgroundShellCommand("sleep 1 & echo hi") {
t.Fatal("mixed should not be fully background")
}
}
+211
View File
@@ -0,0 +1,211 @@
package security
import (
"context"
"errors"
"fmt"
"io"
"os/exec"
"sync"
"github.com/cloudwego/eino/adk/filesystem"
"github.com/cloudwego/eino/schema"
)
// ConfigureShellCmdForAgentExecute 与 exec 工具一致:非交互 stdin、pager/TERM 环境、独立进程组。
func ConfigureShellCmdForAgentExecute(cmd *exec.Cmd) {
if cmd == nil {
return
}
applyDefaultTerminalEnv(cmd)
attachNonInteractiveStdin(cmd)
_ = prepareShellCmdSession(cmd)
}
// TerminateShellCmdTree 尽力终止 shell 及其子进程组(与 exec/execute 超时取消一致)。
func TerminateShellCmdTree(cmd *exec.Cmd) {
terminateCmdTree(cmd)
}
// TerminateShellCmdSession 使用 Start 时缓存的进程组 ID 终止(shell 已退出时仍有效)。
func TerminateShellCmdSession(session *ShellSession) {
TerminateShellSession(session)
}
// EinoStreamingShell 为 Eino ADK execute 工具提供流式 shell,行为与 exec 对齐:
// 并发读取 stdout/stderr(定长块,非按行),避免官方 local.ExecuteStreaming 先排空 stdout
// 导致 stderr 错误(如 sudo 密码提示)长时间不可见、UI 一直显示「执行中」。
type EinoStreamingShell struct{}
// NewEinoStreamingShell 创建 execute 流式 shell 实现。
func NewEinoStreamingShell() *EinoStreamingShell {
return &EinoStreamingShell{}
}
// ExecuteStreaming 实现 filesystem.StreamingShell。
func (s *EinoStreamingShell) ExecuteStreaming(ctx context.Context, input *filesystem.ExecuteRequest) (*schema.StreamReader[*filesystem.ExecuteResponse], error) {
if input == nil || input.Command == "" {
return nil, fmt.Errorf("command is required")
}
sr, w := schema.Pipe[*filesystem.ExecuteResponse](100)
if input.RunInBackendGround {
go runShellInBackground(ctx, input.Command, w)
return sr, nil
}
go streamShellForeground(ctx, input.Command, w)
return sr, nil
}
func runShellInBackground(ctx context.Context, command string, w *schema.StreamWriter[*filesystem.ExecuteResponse]) {
defer w.Close()
command = PrepareShellCommandForExecute(command)
cmd := exec.CommandContext(ctx, "/bin/sh", "-c", command)
applyDefaultTerminalEnv(cmd)
attachNonInteractiveStdin(cmd)
stdout, err := cmd.StdoutPipe()
if err != nil {
_ = w.Send(nil, fmt.Errorf("failed to create stdout pipe: %w", err))
return
}
stderr, err := cmd.StderrPipe()
if err != nil {
_ = stdout.Close()
_ = w.Send(nil, fmt.Errorf("failed to create stderr pipe: %w", err))
return
}
session, err := StartShellSession(cmd)
if err != nil {
_ = stdout.Close()
_ = stderr.Close()
_ = w.Send(nil, fmt.Errorf("failed to start command: %w", err))
return
}
done := make(chan struct{})
go func() {
drainShellPipes(stdout, stderr)
_ = session.Wait()
close(done)
}()
select {
case <-done:
case <-ctx.Done():
TerminateShellCmdSession(session)
}
exitCode := 0
_ = w.Send(&filesystem.ExecuteResponse{
Output: "command started in background\n",
ExitCode: &exitCode,
}, nil)
}
func drainShellPipes(stdout, stderr io.Reader) {
var wg sync.WaitGroup
wg.Add(2)
go func() {
defer wg.Done()
_, _ = io.Copy(io.Discard, stdout)
}()
go func() {
defer wg.Done()
_, _ = io.Copy(io.Discard, stderr)
}()
wg.Wait()
}
func streamShellForeground(ctx context.Context, command string, w *schema.StreamWriter[*filesystem.ExecuteResponse]) {
defer w.Close()
command = PrepareShellCommandForExecute(command)
cmd := exec.CommandContext(ctx, "/bin/sh", "-c", command)
applyDefaultTerminalEnv(cmd)
attachNonInteractiveStdin(cmd)
stdoutPipe, err := cmd.StdoutPipe()
if err != nil {
_ = w.Send(nil, fmt.Errorf("failed to create stdout pipe: %w", err))
return
}
stderrPipe, err := cmd.StderrPipe()
if err != nil {
_ = stdoutPipe.Close()
_ = w.Send(nil, fmt.Errorf("failed to create stderr pipe: %w", err))
return
}
session, err := StartShellSession(cmd)
if err != nil {
_ = stdoutPipe.Close()
_ = stderrPipe.Close()
_ = w.Send(nil, fmt.Errorf("failed to start command: %w", err))
return
}
stopWatch := make(chan struct{})
go func() {
select {
case <-ctx.Done():
TerminateShellCmdSession(session)
case <-stopWatch:
}
}()
defer close(stopWatch)
chunks := make(chan string, 64)
var wg sync.WaitGroup
readFn := func(r io.Reader) {
defer wg.Done()
buf := make([]byte, 8192)
for {
n, readErr := r.Read(buf)
if n > 0 {
chunks <- string(buf[:n])
}
if readErr != nil {
return
}
}
}
wg.Add(2)
go readFn(stdoutPipe)
go readFn(stderrPipe)
go func() {
wg.Wait()
close(chunks)
}()
hadOutput := false
for chunk := range chunks {
if chunk == "" {
continue
}
hadOutput = true
if w.Send(&filesystem.ExecuteResponse{Output: chunk}, nil) {
TerminateShellCmdSession(session)
return
}
}
waitErr := session.Wait()
if waitErr == nil {
exitCode := 0
_ = w.Send(&filesystem.ExecuteResponse{ExitCode: &exitCode}, nil)
return
}
var exitError *exec.ExitError
if errors.As(waitErr, &exitError) {
exitCode := exitError.ExitCode()
resp := &filesystem.ExecuteResponse{ExitCode: &exitCode}
if !hadOutput {
resp.Output = FormatCommandFailureResult(exitCode, "")
}
_ = w.Send(resp, nil)
return
}
_ = w.Send(nil, fmt.Errorf("command failed: %w", waitErr))
}
@@ -0,0 +1,152 @@
package security
import (
"context"
"errors"
"io"
"strings"
"testing"
"time"
"github.com/cloudwego/eino/adk/filesystem"
)
func TestEinoStreamingShell_StreamsStderrBeforeStdoutEOF(t *testing.T) {
shell := NewEinoStreamingShell()
cmd := PrepareNonInteractiveShellCommand("echo err-only >&2; exit 1")
sr, err := shell.ExecuteStreaming(context.Background(), &filesystem.ExecuteRequest{Command: cmd})
if err != nil {
t.Fatalf("ExecuteStreaming: %v", err)
}
defer sr.Close()
start := time.Now()
var got strings.Builder
for {
resp, rerr := sr.Recv()
if errors.Is(rerr, io.EOF) {
break
}
if rerr != nil {
t.Fatalf("recv: %v", rerr)
}
if resp != nil && resp.Output != "" {
got.WriteString(resp.Output)
}
}
if time.Since(start) > 3*time.Second {
t.Fatalf("expected fast completion, took %v", time.Since(start))
}
if !strings.Contains(got.String(), "err-only") {
t.Fatalf("expected stderr in output, got: %q", got.String())
}
}
func TestEinoStreamingShell_SudoFailsFast(t *testing.T) {
shell := NewEinoStreamingShell()
cmd := PrepareNonInteractiveShellCommand("sudo whoami && sudo cat /etc/os-release")
sr, err := shell.ExecuteStreaming(context.Background(), &filesystem.ExecuteRequest{Command: cmd})
if err != nil {
t.Fatalf("ExecuteStreaming: %v", err)
}
defer sr.Close()
start := time.Now()
var got strings.Builder
for {
resp, rerr := sr.Recv()
if errors.Is(rerr, io.EOF) {
break
}
if rerr != nil {
t.Fatalf("recv: %v", rerr)
}
if resp == nil {
continue
}
got.WriteString(resp.Output)
}
if time.Since(start) > 5*time.Second {
t.Fatalf("sudo should fail quickly, took %v output=%q", time.Since(start), got.String())
}
out := got.String()
if strings.Contains(out, "command exited with non-zero code") {
t.Fatalf("legacy exit line present: %q", out)
}
if !strings.Contains(out, "sudo") && !strings.Contains(out, "password") && !strings.Contains(out, "terminal") {
t.Fatalf("expected sudo error text, got: %q", out)
}
}
func TestEinoStreamingShell_StderrWhileStdoutBlocks(t *testing.T) {
shell := NewEinoStreamingShell()
// 模拟 sudostderr 先有输出,stdout 侧进程仍挂起;旧 eino local 在首包 stderr 前不会向流写任何内容。
cmd := PrepareNonInteractiveShellCommand(`echo "password prompt" >&2; sleep 30`)
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
defer cancel()
sr, err := shell.ExecuteStreaming(ctx, &filesystem.ExecuteRequest{Command: cmd})
if err != nil {
t.Fatalf("ExecuteStreaming: %v", err)
}
defer sr.Close()
start := time.Now()
var got strings.Builder
for {
resp, rerr := sr.Recv()
if errors.Is(rerr, io.EOF) {
break
}
if rerr != nil {
break
}
if resp != nil && resp.Output != "" {
got.WriteString(resp.Output)
if strings.Contains(got.String(), "password prompt") {
break
}
}
}
if time.Since(start) > 1500*time.Millisecond {
t.Fatalf("expected stderr promptly, took %v output=%q", time.Since(start), got.String())
}
if !strings.Contains(got.String(), "password prompt") {
t.Fatalf("expected early stderr, got: %q", got.String())
}
}
// TestEinoStreamingShell_BackgroundJobDoesNotHoldPipe 模拟 cmd & 后继续前台逻辑:重定向后应快速结束。
func TestEinoStreamingShell_BackgroundJobDoesNotHoldPipe(t *testing.T) {
if testing.Short() {
t.Skip("skipping shell integration in -short")
}
shell := NewEinoStreamingShell()
cmd := `(sh -c 'printf x; sleep 120') & echo started; sleep 0`
sr, err := shell.ExecuteStreaming(context.Background(), &filesystem.ExecuteRequest{Command: cmd})
if err != nil {
t.Fatalf("ExecuteStreaming: %v", err)
}
defer sr.Close()
start := time.Now()
var got strings.Builder
for {
resp, rerr := sr.Recv()
if errors.Is(rerr, io.EOF) {
break
}
if rerr != nil {
t.Fatalf("recv: %v", rerr)
}
if resp != nil && resp.Output != "" {
got.WriteString(resp.Output)
}
}
if time.Since(start) > 3*time.Second {
t.Fatalf("expected fast completion, took %v output=%q", time.Since(start), got.String())
}
if !strings.Contains(got.String(), "started") {
t.Fatalf("expected foreground echo, got: %q", got.String())
}
}
+163
View File
@@ -0,0 +1,163 @@
package security
import (
"fmt"
"os"
"os/exec"
"strings"
"sync"
"time"
)
// ShellNoOutputTimeoutMessage 长时间无新 stdout/stderr 时的提示(软失败,模型可见)。
func ShellNoOutputTimeoutMessage(idleSec int) string {
return fmt.Sprintf(`命令已终止超过 %d 秒没有新的输出疑似在等待交互输入或已挂起
长时静默任务请使用末尾 & 后台运行或增大 agent.shell_no_output_timeout_seconds-1=关闭此检测
Command terminated: no new output for %d seconds (possible interactive wait or hung process).`, idleSec, idleSec)
}
// ShellInactivityWatch 在 noOutputSec 内无任何新输出时向 expired 发送信号;每次 Bump 重置计时。
// 与「仅有首包输出就永久取消计时」不同,可兜住 sudo 打印 Password 提示后继续挂起等情况。
type ShellInactivityWatch struct {
Sec int
mu sync.Mutex
timer *time.Timer
Expired chan struct{}
}
func NewShellInactivityWatch(noOutputSec int) *ShellInactivityWatch {
sec := ResolveShellNoOutputTimeoutSeconds(noOutputSec)
if sec <= 0 {
return nil
}
w := &ShellInactivityWatch{
Sec: sec,
Expired: make(chan struct{}, 1),
}
w.Bump()
return w
}
func (w *ShellInactivityWatch) Bump() {
if w == nil || w.Sec <= 0 {
return
}
w.mu.Lock()
defer w.mu.Unlock()
if w.timer != nil {
w.timer.Stop()
}
w.timer = time.AfterFunc(time.Duration(w.Sec)*time.Second, func() {
select {
case w.Expired <- struct{}{}:
default:
}
})
}
func (w *ShellInactivityWatch) Stop() {
if w == nil {
return
}
w.mu.Lock()
defer w.mu.Unlock()
if w.timer != nil {
w.timer.Stop()
w.timer = nil
}
}
// ResolveShellNoOutputTimeoutSeconds0=默认 3005 分钟);-1=关闭;>0=自定义。
func ResolveShellNoOutputTimeoutSeconds(sec int) int {
if sec < 0 {
return 0
}
if sec == 0 {
return 300
}
return sec
}
// PrependNonInteractiveShellExports 为 sh -c 注入通用非交互环境(pager 等),不维护命令黑名单。
func PrependNonInteractiveShellExports(shellCommand string) string {
if strings.TrimSpace(shellCommand) == "" {
return shellCommand
}
upper := strings.ToUpper(shellCommand)
var pairs []string
add := func(key, val string) {
if strings.Contains(upper, strings.ToUpper(key)) {
return
}
pairs = append(pairs, key+"="+val)
}
add("GIT_PAGER", "cat")
add("PAGER", "cat")
add("SYSTEMD_PAGER", "cat")
add("DEBIAN_FRONTEND", "noninteractive")
if len(pairs) == 0 {
return shellCommand
}
return "export " + strings.Join(pairs, " ") + "\n" + shellCommand
}
// PrependNonInteractiveStdinRedirect 为 sh -c 关闭 stdin(与 attachNonInteractiveStdin 等价),
// 使 read/input()/sudo -S 等从 stdin 读取的程序快速失败而非挂起。已含 </dev/null 时不重复注入。
func PrependNonInteractiveStdinRedirect(shellCommand string) string {
if strings.TrimSpace(shellCommand) == "" {
return shellCommand
}
lower := strings.ToLower(shellCommand)
if strings.Contains(lower, "</dev/null") || strings.Contains(lower, "0</dev/null") {
return shellCommand
}
return "exec </dev/null\n" + shellCommand
}
// PrepareNonInteractiveShellCommand 组合非交互包装:stdin 关闭 + pager 等环境变量(零名单)。
func PrepareNonInteractiveShellCommand(shellCommand string) string {
return PrependNonInteractiveStdinRedirect(PrependNonInteractiveShellExports(shellCommand))
}
// ApplyNonInteractivePagerEnv 为 exec.Cmd 补齐与 PrependNonInteractiveShellExports 一致的环境变量。
func ApplyNonInteractivePagerEnv(cmdEnv []string) []string {
if cmdEnv == nil {
cmdEnv = []string{}
}
has := func(k string) bool {
prefix := k + "="
for _, e := range cmdEnv {
if strings.HasPrefix(e, prefix) {
return true
}
}
return false
}
if !has("GIT_PAGER") {
cmdEnv = append(cmdEnv, "GIT_PAGER=cat")
}
if !has("PAGER") {
cmdEnv = append(cmdEnv, "PAGER=cat")
}
if !has("SYSTEMD_PAGER") {
cmdEnv = append(cmdEnv, "SYSTEMD_PAGER=cat")
}
if !has("DEBIAN_FRONTEND") {
cmdEnv = append(cmdEnv, "DEBIAN_FRONTEND=noninteractive")
}
return cmdEnv
}
// attachNonInteractiveStdin 关闭交互式 stdin,使部分命令快速失败而非等待输入。
func attachNonInteractiveStdin(cmd *exec.Cmd) {
if cmd == nil || cmd.Stdin != nil {
return
}
f, err := os.Open(os.DevNull)
if err != nil {
return
}
cmd.Stdin = f
}
@@ -0,0 +1,128 @@
package security
import (
"context"
"os"
"os/exec"
"strings"
"testing"
"time"
)
func TestPrependNonInteractiveShellExports(t *testing.T) {
out := PrependNonInteractiveShellExports("echo hi")
if !strings.Contains(out, "GIT_PAGER=cat") || !strings.Contains(out, "PAGER=cat") {
t.Fatalf("missing pager exports: %q", out)
}
if !strings.HasSuffix(strings.TrimSpace(out), "echo hi") {
t.Fatalf("command suffix lost: %q", out)
}
skip := PrependNonInteractiveShellExports("GIT_PAGER=less echo hi")
if strings.Contains(skip, "export GIT_PAGER=cat") {
t.Fatalf("should not override existing GIT_PAGER: %q", skip)
}
}
func TestPrependNonInteractiveStdinRedirect(t *testing.T) {
out := PrependNonInteractiveStdinRedirect("echo hi")
if !strings.HasPrefix(out, "exec </dev/null") {
t.Fatalf("missing stdin redirect: %q", out)
}
if !strings.HasSuffix(strings.TrimSpace(out), "echo hi") {
t.Fatalf("command suffix lost: %q", out)
}
skip := PrependNonInteractiveStdinRedirect("cmd </dev/null")
if strings.HasPrefix(skip, "exec </dev/null") {
t.Fatalf("should not double redirect: %q", skip)
}
}
func TestPrepareNonInteractiveShellCommand(t *testing.T) {
out := PrepareNonInteractiveShellCommand("echo hi")
if !strings.Contains(out, "exec </dev/null") {
t.Fatalf("missing stdin redirect: %q", out)
}
if !strings.Contains(out, "GIT_PAGER=cat") {
t.Fatalf("missing pager export: %q", out)
}
}
func TestNewShellInactivityWatch(t *testing.T) {
w := NewShellInactivityWatch(1)
if w == nil {
t.Fatal("expected watch")
}
w.Bump()
select {
case <-w.Expired:
case <-time.After(3 * time.Second):
t.Fatal("expected inactivity fire within 3s")
}
}
func TestResolveShellNoOutputTimeoutSeconds(t *testing.T) {
if ResolveShellNoOutputTimeoutSeconds(0) != 300 {
t.Fatal("zero should default to 300")
}
if ResolveShellNoOutputTimeoutSeconds(-1) != 0 {
t.Fatal("-1 should disable")
}
if ResolveShellNoOutputTimeoutSeconds(30) != 30 {
t.Fatal("explicit value")
}
}
// TestNonInteractiveStdinReadExitsQuickly 验证 exec </dev/null + attachNonInteractiveStdin 时 read 立即 EOF,不挂起。
func TestNonInteractiveStdinReadExitsQuickly(t *testing.T) {
if testing.Short() {
t.Skip("skipping shell integration in -short")
}
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
cmd := exec.CommandContext(ctx, "sh", "-c", PrepareNonInteractiveShellCommand(`read x; echo "x=<$x>"`))
attachNonInteractiveStdin(cmd)
start := time.Now()
out, err := cmd.CombinedOutput()
elapsed := time.Since(start)
if elapsed > 2*time.Second {
t.Fatalf("read with closed stdin took %v, want <2s", elapsed)
}
if err != nil {
t.Fatalf("unexpected error: %v output=%q", err, out)
}
if !strings.Contains(string(out), "x=<>") {
t.Fatalf("unexpected output: %q", out)
}
}
// TestNonInteractiveStdinReadBlocksWithoutRedirect 对照:stdin 为永不写入的管道时 read 会挂起。
func TestNonInteractiveStdinReadBlocksWithoutRedirect(t *testing.T) {
if testing.Short() {
t.Skip("skipping shell integration in -short")
}
r, w, err := os.Pipe()
if err != nil {
t.Fatal(err)
}
defer r.Close()
// 保持 w 打开且不写数据,模拟「等待用户输入」
cmd := exec.Command("sh", "-c", `read x; echo done`)
cmd.Stdin = r
done := make(chan error, 1)
go func() { done <- cmd.Run() }()
select {
case err := <-done:
t.Fatalf("expected hang, but command finished: %v", err)
case <-time.After(500 * time.Millisecond):
if cmd.Process != nil {
_ = cmd.Process.Kill()
}
_ = w.Close()
<-done // 等待 goroutine 退出
}
}
+47
View File
@@ -0,0 +1,47 @@
package security
import "os/exec"
// ShellSession 在 Start 时记录根 shell 的进程组 ID,取消/超时时可杀整组(即使 cmd.Process 已失效)。
type ShellSession struct {
Cmd *exec.Cmd
rootPID int
}
// StartShellSession 配置独立进程组并启动 shell,缓存 rootPIDUnix 下即 PGID)。
func StartShellSession(cmd *exec.Cmd) (*ShellSession, error) {
if err := prepareShellCmdSession(cmd); err != nil {
return nil, err
}
if err := cmd.Start(); err != nil {
return nil, err
}
pid := 0
if cmd.Process != nil {
pid = cmd.Process.Pid
}
return &ShellSession{Cmd: cmd, rootPID: pid}, nil
}
// Wait 等待 shell 退出。
func (s *ShellSession) Wait() error {
if s == nil || s.Cmd == nil {
return nil
}
return s.Cmd.Wait()
}
// Terminate 终止 shell 及其进程组。
func (s *ShellSession) Terminate() {
if s == nil {
return
}
terminateProcessGroup(s.rootPID, s.Cmd)
}
// TerminateShellSession 终止由 StartShellSession 启动的会话。
func TerminateShellSession(session *ShellSession) {
if session != nil {
session.Terminate()
}
}
+65
View File
@@ -0,0 +1,65 @@
package security
import (
"context"
"os/exec"
"runtime"
"testing"
"time"
)
func TestShellSession_TerminateUsesCachedRootPID(t *testing.T) {
if runtime.GOOS == "windows" {
t.Skip("unix process group kill")
}
cmd := exec.Command("sh", "-c", "sleep 300")
ConfigureShellCmdForAgentExecute(cmd)
session, err := StartShellSession(cmd)
if err != nil {
t.Fatalf("StartShellSession: %v", err)
}
time.Sleep(100 * time.Millisecond)
session.Terminate()
done := make(chan error, 1)
go func() { done <- session.Wait() }()
select {
case <-done:
case <-time.After(5 * time.Second):
t.Fatal("session did not finish within 5s after Terminate")
}
}
func TestShellSession_TerminateAfterContextCancel(t *testing.T) {
if runtime.GOOS == "windows" {
t.Skip("unix process group kill")
}
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
cmd := exec.CommandContext(ctx, "sh", "-c", "sleep 300")
ConfigureShellCmdForAgentExecute(cmd)
session, err := StartShellSession(cmd)
if err != nil {
t.Fatalf("StartShellSession: %v", err)
}
time.Sleep(100 * time.Millisecond)
cancel()
TerminateShellCmdSession(session)
done := make(chan error, 1)
go func() { done <- session.Wait() }()
select {
case <-done:
case <-time.After(5 * time.Second):
t.Fatal("session did not finish within 5s after cancel+terminate")
}
}
@@ -0,0 +1,20 @@
package security
import (
"net/http"
"testing"
)
func TestWorkflowPackageRoutesHaveExplicitWorkflowPermissions(t *testing.T) {
if got := permissionForRequest(http.MethodGet, "/api/workflows/:id/package"); got != "workflow:read" {
t.Fatalf("export permission=%q", got)
}
for _, path := range []string{"/api/workflow-package-inspections", "/api/workflow-package-inspections/:inspectionId", "/api/workflow-package-imports", "/api/workflow-package-imports/:importId"} {
if got := permissionForRequest(http.MethodGet, path); got != "workflow:write" {
t.Fatalf("%s permission=%q", path, got)
}
}
if !isProcessGlobalMutationPath("/workflow-package-imports") || !isProcessGlobalMutationPath("/workflow-package-inspections") {
t.Fatal("package mutations must require all-resource scope")
}
}