mirror of
https://github.com/Ed1s0nZ/CyberStrikeAI.git
synced 2026-08-14 23:20:25 +02:00
Add files via upload
This commit is contained in:
File diff suppressed because it is too large
Load Diff
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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 留空或省略时,沿用 main(openai)对应字段;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) != ""
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -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,
|
||||
}
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
|
||||
// 构造 ClientOptions:KeepAlive 心跳
|
||||
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 而非 CommandContext:doConnect 返回后 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,
|
||||
// 不设 Timeout,SSE 长连接的超时由 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)
|
||||
}
|
||||
@@ -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 后触发 tryReconnect(delay<=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)
|
||||
})
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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))
|
||||
}
|
||||
}
|
||||
@@ -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 将登记器注入 ctx(Eino / 原生 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
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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))
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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 对应进程组发 SIGKILL;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
|
||||
}
|
||||
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)
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
@@ -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()
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -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()
|
||||
// 模拟 sudo:stderr 先有输出,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())
|
||||
}
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
}
|
||||
|
||||
// ResolveShellNoOutputTimeoutSeconds:0=默认 300(5 分钟);-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 退出
|
||||
}
|
||||
}
|
||||
@@ -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,缓存 rootPID(Unix 下即 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()
|
||||
}
|
||||
}
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user