diff --git a/go.mod b/go.mod index ee66dac2..aade70d6 100644 --- a/go.mod +++ b/go.mod @@ -9,6 +9,7 @@ tool ( require ( connectrpc.com/connect v1.20.0 + github.com/Microsoft/go-winio v0.6.2 github.com/alecthomas/chroma/v2 v2.14.0 github.com/carapace-sh/carapace v1.11.6 github.com/chainreactors/crtm v0.0.3-0.20260618163257-073207497076 diff --git a/go.sum b/go.sum index 19defe2a..0ba0f39b 100644 --- a/go.sum +++ b/go.sum @@ -70,6 +70,8 @@ github.com/M09ic/go-ntlmssp v0.0.0-20230312133735-dcccd454dfe0 h1:9Y+BdzDIHfpKy0 github.com/M09ic/go-ntlmssp v0.0.0-20230312133735-dcccd454dfe0/go.mod h1:yMNEF6ulbFipt3CakMhcmcNVACshPRG4Ap4l00V+mMs= github.com/Masterminds/semver/v3 v3.4.0 h1:Zog+i5UMtVoCU8oKka5P7i9q9HgrJeGzI9SA1Xbatp0= github.com/Masterminds/semver/v3 v3.4.0/go.mod h1:4V+yj/TJE1HU9XfppCwVMZq3I84lprf4nC11bSS5beM= +github.com/Microsoft/go-winio v0.6.2 h1:F2VQgta7ecxGYO8k3ZZz3RS8fVIXVxONVUPlNERoyfY= +github.com/Microsoft/go-winio v0.6.2/go.mod h1:yd8OoFMLzJbo9gZq8j5qaps8bJ9aShtEA8Ipt1oGCvU= github.com/Mzack9999/gcache v0.0.0-20230410081825-519e28eab057 h1:KFac3SiGbId8ub47e7kd2PLZeACxc1LkiiNoDOFRClE= github.com/Mzack9999/gcache v0.0.0-20230410081825-519e28eab057/go.mod h1:iLB2pivrPICvLOuROKmlqURtFIEsoJZaMidQfCG1+D4= github.com/Mzack9999/go-http-digest-auth-client v0.6.1-0.20220414142836-eb8883508809 h1:ZbFL+BDfBqegi+/Ssh7im5+aQfBRx6it+kHnC7jaDU8= diff --git a/pkg/commands/bash.go b/pkg/commands/bash.go index a99e6153..959e8362 100644 --- a/pkg/commands/bash.go +++ b/pkg/commands/bash.go @@ -8,6 +8,7 @@ import ( "os/exec" "sort" "strings" + "sync" "time" "github.com/chainreactors/aiscan/agent/inbox" @@ -44,6 +45,10 @@ type BashTool struct { tasks *tmux.Manager commandNames func() []string resolveCommand func(string) (Command, bool) + shellRegistry *CommandRegistry + adapterMu sync.Mutex + shellAdapter *shellCommandAdapter + closeOnce sync.Once } func NewBashTool(workDir string, timeout int) *BashTool { @@ -60,7 +65,47 @@ func (t *BashTool) SetCommandResolver(fn func(string) (Command, bool)) { t.resolveCommand = fn } func (t *BashTool) Name() string { return "bash" } -func (t *BashTool) Close() { t.tasks.Shutdown() } +func (t *BashTool) Close() { + t.closeOnce.Do(func() { + t.adapterMu.Lock() + adapter := t.shellAdapter + if adapter != nil { + adapter.shutdown() + } + t.adapterMu.Unlock() + t.tasks.Shutdown() + if adapter != nil { + adapter.cleanup() + } + }) +} + +func (t *BashTool) attachShellCommands(registry *CommandRegistry) { + t.shellRegistry = registry +} + +func (t *BashTool) ensureShellCommands() (*shellCommandAdapter, error) { + if t.shellRegistry == nil { + return nil, nil + } + t.adapterMu.Lock() + defer t.adapterMu.Unlock() + if t.shellAdapter == nil { + adapter, err := newShellCommandAdapter(t.shellRegistry) + if err != nil { + return nil, err + } + t.shellAdapter = adapter + } + if t.commandNames != nil { + if err := t.shellAdapter.syncAliases(t.commandNames()); err != nil { + t.shellAdapter.close() + t.shellAdapter = nil + return nil, err + } + } + return t.shellAdapter, nil +} func (t *BashTool) WithScannerProxy(proxy string) *BashTool { t.scannerProxy = proxy @@ -208,9 +253,39 @@ func (t *BashTool) Start(ctx context.Context, command string, options BashExecOp if workDir == "" { workDir = t.workDir } - env := t.runEnv(options.Env) left, right, hasPipe := splitPipeline(command) leftToken := firstCommandToken(left) + if !hasPipe { + if cmd, ok := t.resolve(leftToken); ok { + if tokens, err := SplitCommandLine(left); err == nil { + if args, syntaxErr := stripShellSyntax(tokens[1:]); syntaxErr == nil { + args = normalizeNoColor(cmd.Name, args) + return t.startBuiltin(ctx, cmd, args, timeout, workDir, t.runEnv(options.Env, nil, ""), options) + } + } + } + } + adapter, err := t.ensureShellCommands() + if err != nil { + return nil, err + } + if adapter != nil { + contextID := adapter.retainContext(ctx) + env := t.runEnv(options.Env, adapter, contextID) + execution := newExecution(t.tasks, command, nil, workDir, env) + info, err := t.tasks.Create(workDir, command, options.Name, timeout, env, "") + if err != nil { + adapter.releaseContext(contextID) + return nil, err + } + execution.bind(info) + go func() { + <-t.tasks.Done(execution.ID) + adapter.releaseContext(contextID) + }() + return execution, nil + } + env := t.runEnv(options.Env, nil, "") if cmd, ok := t.resolve(leftToken); ok { tokens, err := SplitCommandLine(left) if err != nil { @@ -447,7 +522,7 @@ func (t *BashTool) collectResult(execution *Execution) *coretool.Result { return result } -func (t *BashTool) runEnv(overrides map[string]string) []string { +func (t *BashTool) runEnv(overrides map[string]string, adapter *shellCommandAdapter, shellContextID string) []string { values := make(map[string]string) for _, item := range t.proxyEnv() { if key, value, ok := strings.Cut(item, "="); ok { @@ -457,6 +532,18 @@ func (t *BashTool) runEnv(overrides map[string]string) []string { for key, value := range overrides { values[key] = value } + if adapter != nil && shellContextID != "" { + for _, item := range adapter.environment(shellContextID) { + if key, value, ok := strings.Cut(item, "="); ok { + values[key] = value + } + } + path := values["PATH"] + if path == "" { + path = os.Getenv("PATH") + } + values["PATH"] = adapter.runtimeDir + string(os.PathListSeparator) + path + } keys := make([]string, 0, len(values)) for key := range values { keys = append(keys, key) diff --git a/pkg/commands/bash_test.go b/pkg/commands/bash_test.go index a85ba5d1..0a990f72 100644 --- a/pkg/commands/bash_test.go +++ b/pkg/commands/bash_test.go @@ -479,7 +479,9 @@ func TestBashExecOptionsAreIsolatedAcrossConcurrentCalls(t *testing.T) { wg.Add(1) go func(i int) { defer wg.Done() - results[i], errs[i] = bash.RunForeground(context.Background(), `printf '%s\n' "$AISCAN_RUN_VALUE"; pwd`, BashExecOptions{ + // Keep the short-lived shell alive until the PTY reader is scheduled; + // this test exercises concurrent option isolation, not PTY drain timing. + results[i], errs[i] = bash.RunForeground(context.Background(), `printf '%s\n' "$AISCAN_RUN_VALUE"; pwd; sleep 0.05`, BashExecOptions{ WorkDir: dirs[i], Env: map[string]string{"AISCAN_RUN_VALUE": fmt.Sprintf("value-%d", i)}, OnOutput: func(data []byte) { diff --git a/pkg/commands/register.go b/pkg/commands/register.go index 367fddc2..4158edc4 100644 --- a/pkg/commands/register.go +++ b/pkg/commands/register.go @@ -33,6 +33,7 @@ func init() { bash.SetCommandNames(reg.Names) bash.SetCommandResolver(reg.Get) reg.RegisterTool(bash) + bash.attachShellCommands(reg) tmuxCmd := NewTmuxCommand(bash) reg.Register(tmuxCmd, "core") diff --git a/pkg/commands/register_test.go b/pkg/commands/register_test.go index 1b625953..d8de8900 100644 --- a/pkg/commands/register_test.go +++ b/pkg/commands/register_test.go @@ -29,3 +29,20 @@ func TestNativeListToolIsRunnerOnly(t *testing.T) { t.Fatal("runner mode must expose the native ls tool") } } + +func TestCoreFactoryAttachesRegisteredCommandsToShell(t *testing.T) { + registry := NewRegistry() + BuildPlan(capability.Select(capability.Options{Groups: []string{"core"}}), &Deps{WorkDir: t.TempDir()}, registry) + defer closeRegistryTools(registry) + tool, ok := registry.GetTool("bash") + if !ok { + t.Fatal("core factory did not register bash") + } + bash := tool.(*BashTool) + if bash.shellRegistry != registry { + t.Fatal("registered commands were not attached to bash") + } + if bash.shellAdapter != nil { + t.Fatal("shell adapter must remain lazy until a real shell is needed") + } +} diff --git a/pkg/commands/shell_command_adapter.go b/pkg/commands/shell_command_adapter.go new file mode 100644 index 00000000..c7d0c754 --- /dev/null +++ b/pkg/commands/shell_command_adapter.go @@ -0,0 +1,547 @@ +package commands + +import ( + "bufio" + "context" + "encoding/binary" + "encoding/json" + "errors" + "fmt" + "io" + "net" + "os" + "path/filepath" + "sort" + "strconv" + "strings" + "sync" + "time" + + coretool "github.com/chainreactors/aiscan/core/tool" +) + +const ( + shellCommandAdapterMarkerEnv = "AISCAN_SHELL_COMMAND" + shellCommandAdapterEndpointEnv = "AISCAN_SHELL_COMMAND_ENDPOINT" + shellCommandAdapterCommandEnv = "AISCAN_SHELL_COMMAND_COMMAND" + shellCommandAdapterExecutableEnv = "AISCAN_SHELL_COMMAND_EXECUTABLE" + shellCommandAdapterContextEnv = "AISCAN_SHELL_COMMAND_CONTEXT" + + shellCommandAdapterProtocolVersion = 1 + shellCommandAdapterChunkSize = 32 << 10 + shellCommandAdapterMaxFrameSize = 1 << 20 + shellCommandAdapterDialTimeout = 5 * time.Second +) + +func init() { + if code, ok := runShellCommandProxyIfRequested(); ok { + os.Exit(code) + } +} + +type shellCommandAdapterFrame struct { + Type string `json:"type"` + Version int `json:"version,omitempty"` + Command string `json:"command,omitempty"` + Args []string `json:"args,omitempty"` + Dir string `json:"dir,omitempty"` + ContextID string `json:"context_id,omitempty"` + Data []byte `json:"data,omitempty"` + ExitCode int `json:"exit_code,omitempty"` +} + +type shellCommandAdapter struct { + registry *CommandRegistry + executable string + runtimeDir string + endpoint string + listener net.Listener + cancel context.CancelFunc + + mu sync.Mutex + aliases map[string]string + contexts map[string]context.Context + nextContext uint64 + connections map[net.Conn]struct{} + wg sync.WaitGroup + shutdownOnce sync.Once + cleanupOnce sync.Once +} + +func newShellCommandAdapter(registry *CommandRegistry) (*shellCommandAdapter, error) { + if registry == nil { + return nil, fmt.Errorf("shell command adapter requires a registry") + } + executable, err := os.Executable() + if err != nil { + return nil, fmt.Errorf("resolve shell command executable: %w", err) + } + executable, err = filepath.Abs(executable) + if err != nil { + return nil, fmt.Errorf("resolve shell command executable path: %w", err) + } + if err := cleanupStaleShellCommandAdapterRuntime(); err != nil { + return nil, err + } + + root := shellCommandAdapterRuntimeRoot() + if err := os.MkdirAll(root, 0o700); err != nil { + return nil, fmt.Errorf("create shell command runtime root: %w", err) + } + _ = os.Chmod(root, 0o700) + runtimeDir, err := os.MkdirTemp(root, strconv.Itoa(os.Getpid())+"-") + if err != nil { + return nil, fmt.Errorf("create shell command runtime: %w", err) + } + cleanup := func() { _ = os.RemoveAll(runtimeDir) } + _ = os.Chmod(runtimeDir, 0o700) + endpoint := shellCommandAdapterEndpoint(runtimeDir) + listener, err := listenShellCommandAdapter(endpoint) + if err != nil { + cleanup() + return nil, fmt.Errorf("listen for shell commands: %w", err) + } + ctx, cancel := context.WithCancel(context.Background()) + adapter := &shellCommandAdapter{ + registry: registry, executable: executable, runtimeDir: runtimeDir, + endpoint: endpoint, listener: listener, cancel: cancel, + aliases: make(map[string]string), contexts: make(map[string]context.Context), + connections: make(map[net.Conn]struct{}), + } + adapter.wg.Add(1) + go adapter.accept(ctx) + return adapter, nil +} + +func shellCommandAdapterRuntimeRoot() string { + return filepath.Join(os.TempDir(), "aiscan-shell-commands") +} + +func cleanupStaleShellCommandAdapterRuntime() error { + root := shellCommandAdapterRuntimeRoot() + entries, err := os.ReadDir(root) + if errors.Is(err, os.ErrNotExist) { + return nil + } + if err != nil { + return fmt.Errorf("read shell command runtime root: %w", err) + } + for _, entry := range entries { + if !entry.IsDir() { + continue + } + pidText, _, ok := strings.Cut(entry.Name(), "-") + pid, parseErr := strconv.Atoi(pidText) + if !ok || parseErr != nil || pid <= 0 { + continue + } + if !shellCommandAdapterProcessAlive(pid) { + _ = os.RemoveAll(filepath.Join(root, entry.Name())) + } + } + return nil +} + +func (b *shellCommandAdapter) accept(ctx context.Context) { + defer b.wg.Done() + for { + conn, err := b.listener.Accept() + if err != nil { + if ctx.Err() != nil { + return + } + continue + } + b.mu.Lock() + b.connections[conn] = struct{}{} + b.wg.Add(1) + b.mu.Unlock() + go b.handle(ctx, conn) + } +} + +func (b *shellCommandAdapter) handle(parent context.Context, conn net.Conn) { + defer b.wg.Done() + defer func() { + _ = conn.Close() + b.mu.Lock() + delete(b.connections, conn) + b.mu.Unlock() + }() + + reader := bufio.NewReader(conn) + header, err := readShellCommandAdapterFrame(reader) + if err != nil { + return + } + writer := &shellCommandAdapterFrameWriter{writer: conn} + if header.Type != "request" || header.Version != shellCommandAdapterProtocolVersion { + _, _ = io.WriteString(&shellCommandAdapterStreamWriter{writer: writer, frameType: "stderr"}, "unsupported shell command protocol\n") + _ = writer.write(shellCommandAdapterFrame{Type: "final", ExitCode: 125}) + return + } + command, ok := b.registry.Get(header.Command) + if !ok || command.Run == nil { + message := "unknown in-memory command: " + header.Command + _, _ = io.WriteString(&shellCommandAdapterStreamWriter{writer: writer, frameType: "stderr"}, message+"\n") + _ = writer.write(shellCommandAdapterFrame{Type: "final", ExitCode: 127}) + return + } + baseCtx, ok := b.context(header.ContextID) + if !ok { + _, _ = io.WriteString(&shellCommandAdapterStreamWriter{writer: writer, frameType: "stderr"}, "shell command context expired\n") + _ = writer.write(shellCommandAdapterFrame{Type: "final", ExitCode: 125}) + return + } + + ctx, cancel := context.WithCancel(baseCtx) + defer cancel() + stopParent := context.AfterFunc(parent, cancel) + defer stopParent() + stdinReader, stdinWriter := io.Pipe() + defer stdinReader.Close() + readDone := make(chan error, 1) + go func() { + readDone <- readShellCommandAdapterInput(reader, stdinWriter, cancel) + }() + + stdout := &shellCommandAdapterStreamWriter{writer: writer, frameType: "stdout"} + stderr := &shellCommandAdapterStreamWriter{writer: writer, frameType: "stderr"} + execution := newExecution(nil, command.Name, normalizeNoColor(command.Name, header.Args), header.Dir, nil) + execution.ID = coretool.InvocationFromContext(ctx).CallID + execution.StartedAt = time.Now() + execution.setIO(stdinReader, stdout, stderr) + _, runErr := command.Run(ctx, execution) + execution.EndedAt = time.Now() + _ = stdinReader.Close() + cancel() + + exitCode := shellCommandAdapterExitCode(runErr) + if runErr != nil && !errors.Is(runErr, context.Canceled) && !errors.Is(runErr, context.DeadlineExceeded) { + _, _ = io.WriteString(stderr, runErr.Error()+"\n") + } + _ = writer.write(shellCommandAdapterFrame{Type: "final", ExitCode: exitCode}) + _ = conn.Close() + select { + case <-readDone: + case <-time.After(time.Second): + } +} + +func readShellCommandAdapterInput(reader *bufio.Reader, stdin *io.PipeWriter, cancel context.CancelFunc) error { + stdinOpen := true + closeStdin := func() { + if stdinOpen { + stdinOpen = false + _ = stdin.Close() + } + } + defer closeStdin() + for { + frame, err := readShellCommandAdapterFrame(reader) + if err != nil { + cancel() + return err + } + switch frame.Type { + case "stdin": + if stdinOpen && len(frame.Data) > 0 { + if _, err := stdin.Write(frame.Data); err != nil { + return err + } + } + case "stdin_eof": + closeStdin() + case "cancel": + cancel() + return context.Canceled + default: + cancel() + return fmt.Errorf("unexpected shell command frame %q", frame.Type) + } + } +} + +type shellCommandAdapterFrameWriter struct { + mu sync.Mutex + writer io.Writer +} + +func (w *shellCommandAdapterFrameWriter) write(frame shellCommandAdapterFrame) error { + w.mu.Lock() + defer w.mu.Unlock() + return writeShellCommandAdapterFrame(w.writer, frame) +} + +type shellCommandAdapterStreamWriter struct { + writer *shellCommandAdapterFrameWriter + frameType string +} + +func (w *shellCommandAdapterStreamWriter) Write(data []byte) (int, error) { + for offset := 0; offset < len(data); { + end := offset + shellCommandAdapterChunkSize + if end > len(data) { + end = len(data) + } + if err := w.writer.write(shellCommandAdapterFrame{Type: w.frameType, Data: data[offset:end]}); err != nil { + return offset, err + } + offset = end + } + return len(data), nil +} + +func shellCommandAdapterExitCode(err error) int { + if err == nil { + return 0 + } + var exitCoder interface{ ExitCode() int } + if errors.As(err, &exitCoder) { + code := exitCoder.ExitCode() + if code >= 0 && code <= 255 { + return code + } + } + if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { + return 130 + } + return 1 +} + +func writeShellCommandAdapterFrame(writer io.Writer, frame shellCommandAdapterFrame) error { + data, err := json.Marshal(frame) + if err != nil { + return err + } + if len(data) > shellCommandAdapterMaxFrameSize { + return fmt.Errorf("shell command frame too large: %d", len(data)) + } + var size [4]byte + binary.BigEndian.PutUint32(size[:], uint32(len(data))) + if err := writeShellCommandAdapterBytes(writer, size[:]); err != nil { + return err + } + return writeShellCommandAdapterBytes(writer, data) +} + +func writeShellCommandAdapterBytes(writer io.Writer, data []byte) error { + for len(data) > 0 { + n, err := writer.Write(data) + if err != nil { + return err + } + if n <= 0 { + return io.ErrShortWrite + } + data = data[n:] + } + return nil +} + +func readShellCommandAdapterFrame(reader io.Reader) (shellCommandAdapterFrame, error) { + var size [4]byte + if _, err := io.ReadFull(reader, size[:]); err != nil { + return shellCommandAdapterFrame{}, err + } + length := binary.BigEndian.Uint32(size[:]) + if length == 0 || length > shellCommandAdapterMaxFrameSize { + return shellCommandAdapterFrame{}, fmt.Errorf("invalid shell command frame size: %d", length) + } + data := make([]byte, int(length)) + if _, err := io.ReadFull(reader, data); err != nil { + return shellCommandAdapterFrame{}, err + } + var frame shellCommandAdapterFrame + if err := json.Unmarshal(data, &frame); err != nil { + return shellCommandAdapterFrame{}, err + } + return frame, nil +} + +func (b *shellCommandAdapter) syncAliases(names []string) error { + b.mu.Lock() + defer b.mu.Unlock() + sorted := append([]string(nil), names...) + sort.Strings(sorted) + for _, name := range sorted { + if !validShellCommandAdapterName(name) { + continue + } + if _, ok := b.aliases[name]; ok { + continue + } + path, err := createShellCommandAdapterAlias(b.executable, b.runtimeDir, name) + if err != nil { + return fmt.Errorf("create shell command alias %s: %w", name, err) + } + b.aliases[name] = path + } + return nil +} + +func validShellCommandAdapterName(name string) bool { + if name == "" || name == "." || name == ".." { + return false + } + for _, r := range name { + if (r >= 'a' && r <= 'z') || (r >= 'A' && r <= 'Z') || (r >= '0' && r <= '9') || r == '-' || r == '_' || r == '.' { + continue + } + return false + } + return true +} + +func (b *shellCommandAdapter) retainContext(ctx context.Context) string { + if ctx == nil { + ctx = context.Background() + } + b.mu.Lock() + defer b.mu.Unlock() + b.nextContext++ + id := strconv.FormatUint(b.nextContext, 10) + b.contexts[id] = ctx + return id +} + +func (b *shellCommandAdapter) context(id string) (context.Context, bool) { + b.mu.Lock() + defer b.mu.Unlock() + ctx, ok := b.contexts[id] + return ctx, ok +} + +func (b *shellCommandAdapter) releaseContext(id string) { + b.mu.Lock() + delete(b.contexts, id) + b.mu.Unlock() +} + +func (b *shellCommandAdapter) environment(contextID string) []string { + return []string{ + shellCommandAdapterMarkerEnv + "=1", + shellCommandAdapterEndpointEnv + "=" + b.endpoint, + shellCommandAdapterExecutableEnv + "=" + b.executable, + shellCommandAdapterContextEnv + "=" + contextID, + } +} + +func (b *shellCommandAdapter) shutdown() { + b.shutdownOnce.Do(func() { + b.cancel() + _ = b.listener.Close() + b.mu.Lock() + for conn := range b.connections { + _ = conn.Close() + } + b.mu.Unlock() + b.wg.Wait() + }) +} + +func (b *shellCommandAdapter) cleanup() { + b.cleanupOnce.Do(func() { + deadline := time.Now().Add(2 * time.Second) + for { + if err := os.RemoveAll(b.runtimeDir); err == nil { + break + } + if time.Now().After(deadline) { + break + } + time.Sleep(50 * time.Millisecond) + } + _ = os.Remove(shellCommandAdapterRuntimeRoot()) + }) +} + +func (b *shellCommandAdapter) close() { + b.shutdown() + b.cleanup() +} + +// runShellCommandProxyIfRequested dispatches the process-local PATH shim. The +// marker and command variables are only injected into shim children, so normal +// AIScan and embedding-host startup is unchanged. +func runShellCommandProxyIfRequested() (code int, ok bool) { + command := strings.TrimSpace(os.Getenv(shellCommandAdapterCommandEnv)) + if os.Getenv(shellCommandAdapterMarkerEnv) != "1" || command == "" { + return 0, false + } + if !validShellCommandAdapterName(command) { + fmt.Fprintln(os.Stderr, "invalid shell command alias") + return 126, true + } + return runShellCommandAdapterProxy(command, os.Args[1:]), true +} + +func runShellCommandAdapterProxy(command string, args []string) int { + endpoint := os.Getenv(shellCommandAdapterEndpointEnv) + if endpoint == "" { + fmt.Fprintln(os.Stderr, "AIScan shell command environment is incomplete") + return 126 + } + ctx, cancel := context.WithTimeout(context.Background(), shellCommandAdapterDialTimeout) + conn, err := dialShellCommandAdapter(ctx, endpoint) + cancel() + if err != nil { + fmt.Fprintf(os.Stderr, "connect AIScan shell command adapter: %s\n", err) + return 125 + } + defer conn.Close() + cwd, _ := os.Getwd() + header := shellCommandAdapterFrame{ + Type: "request", Version: shellCommandAdapterProtocolVersion, + Command: command, Args: append([]string(nil), args...), Dir: cwd, + ContextID: os.Getenv(shellCommandAdapterContextEnv), + } + writer := &shellCommandAdapterFrameWriter{writer: conn} + if err := writer.write(header); err != nil { + fmt.Fprintf(os.Stderr, "start AIScan shell command request: %s\n", err) + return 125 + } + go streamShellCommandAdapterStdin(writer) + + reader := bufio.NewReader(conn) + for { + frame, err := readShellCommandAdapterFrame(reader) + if err != nil { + fmt.Fprintf(os.Stderr, "read AIScan shell command response: %s\n", err) + return 125 + } + switch frame.Type { + case "stdout": + _, _ = os.Stdout.Write(frame.Data) + case "stderr": + _, _ = os.Stderr.Write(frame.Data) + case "final": + flushShellCommandAdapterProxyOutput() + return frame.ExitCode + default: + fmt.Fprintf(os.Stderr, "unexpected AIScan shell command response %q\n", frame.Type) + return 125 + } + } +} + +func streamShellCommandAdapterStdin(writer *shellCommandAdapterFrameWriter) { + info, err := os.Stdin.Stat() + if err != nil || info.Mode()&os.ModeCharDevice != 0 { + _ = writer.write(shellCommandAdapterFrame{Type: "stdin_eof"}) + return + } + buffer := make([]byte, shellCommandAdapterChunkSize) + for { + n, readErr := os.Stdin.Read(buffer) + if n > 0 { + if err := writer.write(shellCommandAdapterFrame{Type: "stdin", Data: buffer[:n]}); err != nil { + return + } + } + if readErr != nil { + _ = writer.write(shellCommandAdapterFrame{Type: "stdin_eof"}) + return + } + } +} diff --git a/pkg/commands/shell_command_adapter_test.go b/pkg/commands/shell_command_adapter_test.go new file mode 100644 index 00000000..e2499bda --- /dev/null +++ b/pkg/commands/shell_command_adapter_test.go @@ -0,0 +1,247 @@ +package commands + +import ( + "bytes" + "context" + "errors" + "fmt" + "io" + "os" + "path/filepath" + "strings" + "sync" + "testing" + "time" + + coretool "github.com/chainreactors/aiscan/core/tool" +) + +type adapterTestExitError struct{ code int } + +func (e adapterTestExitError) Error() string { return fmt.Sprintf("adapter test exit %d", e.code) } +func (e adapterTestExitError) ExitCode() int { return e.code } + +type adapterTestCommands struct { + started chan struct{} + canceled chan struct{} + once sync.Once +} + +func newAdapterTestBash(t *testing.T) (*BashTool, *CommandRegistry, *adapterTestCommands) { + t.Helper() + state := &adapterTestCommands{started: make(chan struct{}), canceled: make(chan struct{})} + registry := NewRegistry() + registry.Register(Command{Name: "memory_echo", Run: func(_ context.Context, execution *Execution) (any, error) { + fmt.Fprintln(execution.Stdout, strings.Join(execution.Args, " ")) + return nil, nil + }}, "test") + registry.Register(Command{Name: "memory_upper", Run: func(_ context.Context, execution *Execution) (any, error) { + data, err := io.ReadAll(execution.Stdin) + if err != nil { + return nil, err + } + _, err = execution.Stdout.Write(bytes.ToUpper(data)) + return nil, err + }}, "test") + registry.Register(Command{Name: "memory_fail", Run: func(context.Context, *Execution) (any, error) { + return nil, adapterTestExitError{code: 7} + }}, "test") + registry.Register(Command{Name: "memory_context", Run: func(ctx context.Context, execution *Execution) (any, error) { + invocation := coretool.InvocationFromContext(ctx) + fmt.Fprintf(execution.Stdout, "dir=%s call=%s session=%s turn=%s emitter=%s\n", + execution.Dir, invocation.CallID, invocation.SessionID, invocation.TurnID, invocation.Emitter) + return nil, nil + }}, "test") + registry.Register(Command{Name: "memory_wait", Run: func(ctx context.Context, _ *Execution) (any, error) { + state.once.Do(func() { close(state.started) }) + <-ctx.Done() + close(state.canceled) + return nil, ctx.Err() + }}, "test") + + bash := NewBashTool(t.TempDir(), 10) + bash.SetCommandNames(registry.Names) + bash.SetCommandResolver(registry.Get) + bash.attachShellCommands(registry) + t.Cleanup(bash.Close) + return bash, registry, state +} + +func runAdapterCommand(t *testing.T, bash *BashTool, ctx context.Context, command string, workDir string) (*Execution, string) { + t.Helper() + var output strings.Builder + execution, err := bash.RunForeground(ctx, command, BashExecOptions{ + WorkDir: workDir, + OnOutput: func(data []byte) { + _, _ = output.Write(data) + }, + }) + if err != nil { + t.Fatalf("RunForeground(%q): %v", command, err) + } + return execution, output.String() +} + +func TestShellCommandAdapterIsLazy(t *testing.T) { + bash, _, _ := newAdapterTestBash(t) + execution, output := runAdapterCommand(t, bash, context.Background(), "memory_echo direct", t.TempDir()) + if execution.ExitCode != 0 || !strings.Contains(output, "direct") { + t.Fatalf("direct command exit=%d output=%q", execution.ExitCode, output) + } + if bash.shellAdapter != nil { + t.Fatal("simple registered command allocated shell runtime state") + } +} + +func TestShellCommandMarkerDoesNotHijackNormalChildProcess(t *testing.T) { + t.Setenv(shellCommandAdapterMarkerEnv, "1") + t.Setenv(shellCommandAdapterCommandEnv, "") + if code, ok := runShellCommandProxyIfRequested(); ok { + t.Fatalf("marker-only child was treated as proxy with code %d", code) + } +} + +func TestShellCommandComposition(t *testing.T) { + bash, _, _ := newAdapterTestBash(t) + workDir := t.TempDir() + + execution, output := runAdapterCommand(t, bash, context.Background(), "memory_echo one && memory_echo two", workDir) + if execution.ExitCode != 0 || !strings.Contains(output, "one") || !strings.Contains(output, "two") { + t.Fatalf("and composition exit=%d output=%q", execution.ExitCode, output) + } + + execution, output = runAdapterCommand(t, bash, context.Background(), "memory_fail || memory_echo recovered", workDir) + if execution.ExitCode != 0 || !strings.Contains(output, "recovered") { + t.Fatalf("or composition exit=%d output=%q", execution.ExitCode, output) + } + + execution, output = runAdapterCommand(t, bash, context.Background(), "memory_echo hello | memory_upper", workDir) + if execution.ExitCode != 0 || !strings.Contains(output, "HELLO") { + t.Fatalf("pipeline exit=%d output=%q", execution.ExitCode, output) + } + + execution, _ = runAdapterCommand(t, bash, context.Background(), "memory_fail && memory_echo unreachable", workDir) + if execution.ExitCode != 7 { + t.Fatalf("short-circuit exit code = %d, want 7", execution.ExitCode) + } +} + +func TestShellCommandRedirectionAndInvocationContext(t *testing.T) { + bash, _, _ := newAdapterTestBash(t) + workDir := t.TempDir() + ctx := coretool.ContextWithInvocation(context.Background(), coretool.Invocation{ + WorkDir: workDir, CallID: "call-1", SessionID: "session-1", TurnID: "turn-1", Emitter: "runner", + }) + + execution, _ := runAdapterCommand(t, bash, ctx, "memory_context > adapter-context.txt", workDir) + if execution.ExitCode != 0 { + t.Fatalf("redirection exit code = %d", execution.ExitCode) + } + data, err := os.ReadFile(filepath.Join(workDir, "adapter-context.txt")) + if err != nil { + t.Fatalf("read redirected output: %v", err) + } + text := string(data) + for _, expected := range []string{"dir=" + workDir, "call=call-1", "session=session-1", "turn=turn-1", "emitter=runner"} { + if !strings.Contains(text, expected) { + t.Fatalf("redirected context missing %q: %q", expected, text) + } + } +} + +func TestShellCommandAdapterFindsCommandsRegisteredLater(t *testing.T) { + bash, registry, _ := newAdapterTestBash(t) + registry.Register(Command{Name: "scan", Run: func(_ context.Context, execution *Execution) (any, error) { + fmt.Fprintln(execution.Stdout, strings.Join(execution.Args, " ")) + return nil, nil + }}, "test") + command := "scan -i http://127.0.0.1:1 --timeout 1 --no-color && memory_echo ready" + execution, output := runAdapterCommand(t, bash, context.Background(), command, t.TempDir()) + if execution.ExitCode != 0 || !strings.Contains(output, "http://127.0.0.1:1") || !strings.Contains(output, "ready") { + t.Fatalf("late alias exit=%d output=%q", execution.ExitCode, output) + } +} + +func TestShellCommandCloseCancelsCallsAndRemovesRuntime(t *testing.T) { + bash, _, state := newAdapterTestBash(t) + execution, err := bash.Start(context.Background(), "memory_wait && memory_echo unreachable", BashExecOptions{WorkDir: t.TempDir()}) + if err != nil { + t.Fatalf("Start memory_wait: %v", err) + } + runtimeDir := bash.shellAdapter.runtimeDir + if _, err := os.Stat(runtimeDir); err != nil { + t.Fatalf("runtime directory before close: %v", err) + } + select { + case <-state.started: + case <-time.After(5 * time.Second): + t.Fatal("shell command did not start") + } + bash.Close() + bash.Close() + select { + case <-state.canceled: + case <-time.After(5 * time.Second): + t.Fatal("shell command was not canceled") + } + if err := execution.Wait(context.Background()); err != nil && !errors.Is(err, context.Canceled) { + t.Fatalf("wait after close: %v", err) + } + if _, err := os.Stat(runtimeDir); !errors.Is(err, os.ErrNotExist) { + t.Fatalf("runtime directory still exists after close: %v", err) + } +} + +func TestShellCommandDisconnectCancelsRunningCommand(t *testing.T) { + bash, _, state := newAdapterTestBash(t) + adapter, err := bash.ensureShellCommands() + if err != nil { + t.Fatalf("ensure shell commands: %v", err) + } + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + conn, err := dialShellCommandAdapter(ctx, adapter.endpoint) + if err != nil { + t.Fatalf("dial shell command adapter: %v", err) + } + contextID := adapter.retainContext(context.Background()) + defer adapter.releaseContext(contextID) + writer := &shellCommandAdapterFrameWriter{writer: conn} + if err := writer.write(shellCommandAdapterFrame{ + Type: "request", Version: shellCommandAdapterProtocolVersion, + Command: "memory_wait", Dir: t.TempDir(), ContextID: contextID, + }); err != nil { + t.Fatalf("write request: %v", err) + } + if err := writer.write(shellCommandAdapterFrame{Type: "stdin_eof"}); err != nil { + t.Fatalf("write stdin eof: %v", err) + } + select { + case <-state.started: + case <-ctx.Done(): + t.Fatal("shell command did not start") + } + _ = conn.Close() + select { + case <-state.canceled: + case <-ctx.Done(): + t.Fatal("disconnect did not cancel shell command") + } +} + +func TestShellCommandStartupReclaimsOwnedStaleRuntime(t *testing.T) { + root := shellCommandAdapterRuntimeRoot() + if err := os.MkdirAll(root, 0o700); err != nil { + t.Fatal(err) + } + stale, err := os.MkdirTemp(root, "1073741824-stale-") + if err != nil { + t.Fatal(err) + } + if err := cleanupStaleShellCommandAdapterRuntime(); err != nil { + t.Fatalf("cleanup stale runtime: %v", err) + } + if _, err := os.Stat(stale); !errors.Is(err, os.ErrNotExist) { + t.Fatalf("stale runtime still exists: %v", err) + } +} diff --git a/pkg/commands/shell_command_adapter_unix.go b/pkg/commands/shell_command_adapter_unix.go new file mode 100644 index 00000000..9b61a9c7 --- /dev/null +++ b/pkg/commands/shell_command_adapter_unix.go @@ -0,0 +1,57 @@ +//go:build !windows + +package commands + +import ( + "context" + "errors" + "net" + "os" + "path/filepath" + "syscall" +) + +func shellCommandAdapterEndpoint(runtimeDir string) string { + return filepath.Join(runtimeDir, "commands.sock") +} + +func listenShellCommandAdapter(endpoint string) (net.Listener, error) { + _ = os.Remove(endpoint) + listener, err := net.Listen("unix", endpoint) + if err != nil { + return nil, err + } + if err := os.Chmod(endpoint, 0o600); err != nil { + _ = listener.Close() + return nil, err + } + return listener, nil +} + +func dialShellCommandAdapter(ctx context.Context, endpoint string) (net.Conn, error) { + var dialer net.Dialer + return dialer.DialContext(ctx, "unix", endpoint) +} + +func createShellCommandAdapterAlias(executable, runtimeDir, name string) (string, error) { + _ = executable + path := filepath.Join(runtimeDir, name) + if _, err := os.Stat(path); err == nil { + return path, nil + } + content := "#!/bin/sh\n" + shellCommandAdapterCommandEnv + "='" + name + "' exec \"$" + shellCommandAdapterExecutableEnv + "\" \"$@\"\n" + if err := os.WriteFile(path, []byte(content), 0o700); err != nil { + return "", err + } + return path, nil +} + +func shellCommandAdapterProcessAlive(pid int) bool { + err := syscall.Kill(pid, 0) + return err == nil || errors.Is(err, syscall.EPERM) +} + +func flushShellCommandAdapterProxyOutput() { + _ = os.Stdout.Sync() + _ = os.Stderr.Sync() +} diff --git a/pkg/commands/shell_command_adapter_windows.go b/pkg/commands/shell_command_adapter_windows.go new file mode 100644 index 00000000..6b2dce86 --- /dev/null +++ b/pkg/commands/shell_command_adapter_windows.go @@ -0,0 +1,73 @@ +//go:build windows + +package commands + +import ( + "context" + "fmt" + "net" + "os" + "path/filepath" + "time" + + "github.com/Microsoft/go-winio" + "golang.org/x/sys/windows" +) + +func shellCommandAdapterEndpoint(runtimeDir string) string { + return `\\.\pipe\aiscan-shell-commands-` + fmt.Sprintf("%d-%s", os.Getpid(), filepath.Base(runtimeDir)) +} + +func listenShellCommandAdapter(endpoint string) (net.Listener, error) { + user, err := windows.GetCurrentProcessToken().GetTokenUser() + if err != nil { + return nil, err + } + return winio.ListenPipe(endpoint, &winio.PipeConfig{ + SecurityDescriptor: "D:P(A;;GA;;;" + user.User.Sid.String() + ")", + InputBufferSize: shellCommandAdapterChunkSize, + OutputBufferSize: shellCommandAdapterChunkSize, + }) +} + +func dialShellCommandAdapter(ctx context.Context, endpoint string) (net.Conn, error) { + return winio.DialPipeContext(ctx, endpoint) +} + +func createShellCommandAdapterAlias(executable, runtimeDir, name string) (string, error) { + _ = executable + path := filepath.Join(runtimeDir, name+".cmd") + if _, err := os.Stat(path); err == nil { + return path, nil + } + content := "@echo off\r\n" + + "set \"" + shellCommandAdapterCommandEnv + "=" + name + "\"\r\n" + + "\"%" + shellCommandAdapterExecutableEnv + "%\" %*\r\n" + + "exit /b %errorlevel%\r\n" + if err := os.WriteFile(path, []byte(content), 0o600); err != nil { + return "", err + } + return path, nil +} + +func shellCommandAdapterProcessAlive(pid int) bool { + handle, err := windows.OpenProcess(windows.PROCESS_QUERY_LIMITED_INFORMATION, false, uint32(pid)) + if err != nil { + return err == windows.ERROR_ACCESS_DENIED + } + defer windows.CloseHandle(handle) + var exitCode uint32 + if err := windows.GetExitCodeProcess(handle, &exitCode); err != nil { + return true + } + return exitCode == 259 // STILL_ACTIVE +} + +func flushShellCommandAdapterProxyOutput() { + _ = os.Stdout.Sync() + _ = os.Stderr.Sync() + // ConPTY can report the proxy process exit before consuming its final + // console write. A short drain window prevents the last command in a chain + // from losing output when cmd.exe exits immediately afterwards. + time.Sleep(15 * time.Millisecond) +}