CyberStrikeAI/internal/security/shell_execute_stream_test.go

package security

import (
	"context"
	"errors"
	"io"
	"strings"
	"testing"
	"time"
)

func TestStreamingShell_StreamsStderrBeforeStdoutEOF(t *testing.T) {
	shell := NewStreamingShell()
	cmd := PrepareNonInteractiveShellCommand("echo err-only >&2; exit 1")
	sr, err := shell.ExecuteStreaming(context.Background(), &ShellExecuteRequest{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 TestStreamingShell_SudoFailsFast(t *testing.T) {
	shell := NewStreamingShell()
	// Exercise stderr delivery and failure propagation without relying on the
	// host's sudo policy: CI runners may allow passwordless sudo, even as root.
	// A shell function also prevents this test from invoking the real sudo.
	cmd := PrepareNonInteractiveShellCommand(`
sudo() {
    printf '%s\n' 'sudo: a password is required' >&2
    return 1
}
sudo whoami && printf '%s\n' 'unexpected-command-success'
`)
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	start := time.Now()
	sr, err := shell.ExecuteStreaming(ctx, &ShellExecuteRequest{Command: cmd})
	if err != nil {
		t.Fatalf("ExecuteStreaming: %v", err)
	}
	defer sr.Close()

	var got strings.Builder
	var exitCode *int
	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 resp.ExitCode != nil {
			exitCode = resp.ExitCode
		}
	}
	if ctx.Err() != nil || 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: a password is required") {
		t.Fatalf("expected sudo error text, got: %q", out)
	}
	if strings.Contains(out, "unexpected-command-success") {
		t.Fatalf("command after failed sudo unexpectedly ran: %q", out)
	}
	if exitCode == nil || *exitCode != 1 {
		t.Fatalf("expected exit code 1, got: %v", exitCode)
	}
}

func TestStreamingShell_StderrWhileStdoutBlocks(t *testing.T) {
	shell := NewStreamingShell()
	// 模拟 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, &ShellExecuteRequest{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())
	}
}

// TestStreamingShell_BackgroundJobDoesNotHoldPipe 模拟 cmd & 后继续前台逻辑：重定向后应快速结束。
func TestStreamingShell_BackgroundJobDoesNotHoldPipe(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping shell integration in -short")
	}
	shell := NewStreamingShell()
	cmd := `(sh -c 'printf x; sleep 120') & echo started; sleep 0`
	sr, err := shell.ExecuteStreaming(context.Background(), &ShellExecuteRequest{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())
	}
}