Implement wc builtin command

interp/builtin_wc_pentest_test.go

@@ -0,0 +1,237 @@
+// Unless explicitly stated otherwise all files in this repository are licensed
+// under the Apache License Version 2.0.
+// This product includes software developed at Datadog (https://www.datadoghq.com/).
+// Copyright 2026-present Datadog, Inc.
+
+package interp_test
+
+import (
+	"bytes"
+	"context"
+	"errors"
+	"os"
+	"path/filepath"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	"mvdan.cc/sh/v3/syntax"
+
+	"github.com/DataDog/rshell/interp"
+)
+
+func wcRun(t *testing.T, script, dir string) (string, string, int) {
+	t.Helper()
+	return wcRunCtx(context.Background(), t, script, dir)
+}
+
+func wcRunCtx(ctx context.Context, t *testing.T, script, dir string) (string, string, int) {
+	t.Helper()
+	parser := syntax.NewParser()
+	prog, err := parser.Parse(strings.NewReader(script), "")
+	require.NoError(t, err)
+
+	var outBuf, errBuf bytes.Buffer
+	opts := []interp.RunnerOption{
+		interp.StdIO(nil, &outBuf, &errBuf),
+		interp.AllowedPaths([]string{dir}),
+	}
+
+	runner, err := interp.New(opts...)
+	require.NoError(t, err)
+	defer runner.Close()
+
+	if dir != "" {
+		runner.Dir = dir
+	}
+
+	err = runner.Run(ctx, prog)
+	exitCode := 0
+	if err != nil {
+		var es interp.ExitStatus
+		if errors.As(err, &es) {
+			exitCode = int(es)
+		} else if ctx.Err() == nil {
+			t.Fatalf("unexpected error: %v", err)
+		}
+	}
+	return outBuf.String(), errBuf.String(), exitCode
+}
+
+func wcWriteFile(t *testing.T, dir, name, content string) {
+	t.Helper()
+	require.NoError(t, os.WriteFile(filepath.Join(dir, name), []byte(content), 0644))
+}
+
+// --- Flag and argument injection ---
+
+func TestWcPentestUnknownFlags(t *testing.T) {
+	dir := t.TempDir()
+	for _, flag := range []string{"-f", "--follow", "--no-such-flag", "--files0-from=foo"} {
+		_, stderr, code := wcRun(t, "wc "+flag, dir)
+		assert.Equal(t, 1, code, "flag: %s", flag)
+		assert.Contains(t, stderr, "wc:", "flag: %s", flag)
+	}
+}
+
+func TestWcPentestDoubleDashFlagLikeFile(t *testing.T) {
+	dir := t.TempDir()
+	wcWriteFile(t, dir, "-v", "hello\n")
+	stdout, _, code := wcRun(t, "wc -- -v", dir)
+	assert.Equal(t, 0, code)
+	assert.Contains(t, stdout, "-v")
+}
+
+func TestWcPentestMultipleStdin(t *testing.T) {
+	dir := t.TempDir()
+	wcWriteFile(t, dir, "file.txt", "hello\n")
+	stdout, _, code := wcRun(t, "cat file.txt | wc - -", dir)
+	assert.Equal(t, 0, code)
+	assert.Contains(t, stdout, "total")
+}
+
+// --- Path edge cases ---
+
+func TestWcPentestNonexistentFile(t *testing.T) {
+	dir := t.TempDir()
+	stdout, stderr, code := wcRun(t, "wc nonexistent.txt", dir)
+	assert.Equal(t, 1, code)
+	assert.Equal(t, "", stdout)
+	assert.Contains(t, stderr, "wc:")
+}
+
+func TestWcPentestEmptyFilename(t *testing.T) {
+	dir := t.TempDir()
+	stdout, stderr, code := wcRun(t, "wc ''", dir)
+	assert.Equal(t, 1, code)
+	assert.Equal(t, "", stdout)
+	assert.Contains(t, stderr, "wc:")
+}
+
+// --- Special files ---
+
+func TestWcPentestDevNull(t *testing.T) {
+	dir := t.TempDir()
+	wcWriteFile(t, dir, "empty.txt", "")
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+	stdout, _, code := wcRunCtx(ctx, t, "wc empty.txt", dir)
+	assert.Equal(t, 0, code)
+	assert.Contains(t, stdout, "0")
+}
+
+// --- Context cancellation ---
+
+func TestWcPentestContextCancelled(t *testing.T) {
+	dir := t.TempDir()
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+	_, _, _ = wcRunCtx(ctx, t, "wc", dir)
+}
+
+func TestWcPentestContextTimeout(t *testing.T) {
+	dir := t.TempDir()
+	wcWriteFile(t, dir, "file.txt", strings.Repeat("hello\n", 10000))
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+	stdout, _, code := wcRunCtx(ctx, t, "wc file.txt", dir)
+	assert.Equal(t, 0, code)
+	assert.Contains(t, stdout, "10000")
+}
+
+// --- Large input ---
+
+func TestWcPentestLargeFile(t *testing.T) {
+	dir := t.TempDir()
+	content := strings.Repeat("word word word word word\n", 40000)
+	wcWriteFile(t, dir, "large.txt", content)
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+	stdout, _, code := wcRunCtx(ctx, t, "wc -l large.txt", dir)
+	assert.Equal(t, 0, code)
+	assert.Contains(t, stdout, "40000")
+}
+
+// --- Many files (FD leak check) ---
+
+func TestWcPentestManyFiles(t *testing.T) {
+	dir := t.TempDir()
+	var args []string
+	for i := 0; i < 50; i++ {
+		name := filepath.Join(dir, strings.ReplaceAll(filepath.Base(t.Name()), "/", "_")+"_"+string(rune('a'+i%26))+string(rune('0'+i/26))+".txt")
+		require.NoError(t, os.WriteFile(name, []byte("x\n"), 0644))
+		args = append(args, filepath.Base(name))
+	}
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+	stdout, _, code := wcRunCtx(ctx, t, "wc "+strings.Join(args, " "), dir)
+	assert.Equal(t, 0, code)
+	assert.Contains(t, stdout, "total")
+}
+
+// --- Edge case: file with only newlines ---
+
+func TestWcPentestOnlyNewlines(t *testing.T) {
+	dir := t.TempDir()
+	wcWriteFile(t, dir, "file.txt", strings.Repeat("\n", 100))
+	stdout, _, code := wcRun(t, "wc file.txt", dir)
+	assert.Equal(t, 0, code)
+	assert.Contains(t, stdout, "100")
+	assert.Contains(t, stdout, "  0")
+}
+
+// --- Edge case: long line ---
+
+func TestWcPentestLongLine(t *testing.T) {
+	dir := t.TempDir()
+	longLine := strings.Repeat("x", 1024*1024) + "\n"
+	wcWriteFile(t, dir, "file.txt", longLine)
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+	stdout, _, code := wcRunCtx(ctx, t, "wc -L file.txt", dir)
+	assert.Equal(t, 0, code)
+	assert.Contains(t, stdout, "1048576")
+}
+
+// --- Invalid UTF-8 at chunk boundary ---
+
+func TestWcPentestInvalidUTF8AtChunkBoundary(t *testing.T) {
+	dir := t.TempDir()
+	// Build content so that invalid UTF-8 bytes (0xC0 0x80) land at the
+	// exact 32 KiB read boundary. This exercises the carry buffer edge
+	// case where invalid bytes must be handled in-place (not carried).
+	const chunkSize = 32 * 1024
+	padding := strings.Repeat("A", chunkSize-1) // fills up to byte 32767
+	// Place 0xC0 at offset 32767 (last byte of first chunk) and 0x80 at
+	// offset 32768 (first byte of second chunk).
+	content := []byte(padding)
+	content = append(content, 0xC0, 0x80)
+	content = append(content, '\n')
+
+	require.NoError(t, os.WriteFile(filepath.Join(dir, "invalid_utf8.txt"), content, 0644))
+
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+
+	// -c should report exact byte count
+	stdout, _, code := wcRunCtx(ctx, t, "wc -c invalid_utf8.txt", dir)
+	assert.Equal(t, 0, code)
+	assert.Contains(t, stdout, "32770") // chunkSize - 1 + 2 invalid bytes + newline
+
+	// -l should count the newline
+	stdout, _, code = wcRunCtx(ctx, t, "wc -l invalid_utf8.txt", dir)
+	assert.Equal(t, 0, code)
+	assert.Contains(t, stdout, "1")
+}
+
+// --- Flag expansion in loop ---
+
+func TestWcPentestFlagExpansion(t *testing.T) {
+	dir := t.TempDir()
+	wcWriteFile(t, dir, "file.txt", "hello\n")
+	_, stderr, code := wcRun(t, "for flag in --follow; do wc $flag file.txt; done", dir)
+	assert.Equal(t, 1, code)
+	assert.Contains(t, stderr, "wc:")
+}