secret/internal/cli/secrets_size_test.go

//nolint:testpackage // white-box test of unexported internals
package cli

import (
	"bytes"
	"context"
	"crypto/rand"
	"fmt"
	"io"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"
	"testing"
	"time"

	"github.com/spf13/afero"
	"github.com/spf13/cobra"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"sneak.berlin/go/secret/internal/secret"
	"sneak.berlin/go/secret/internal/vault"
	"sneak.berlin/go/secret/pkg/agehd"
)

// testVaultName is the vault name used by the size tests.
const testVaultName = "test-vault"

// newSizeTestVault creates an in-memory vault unlocked with the test
// mnemonic and returns the filesystem and vault.
//
//nolint:ireturn // afero.Fs is the filesystem abstraction used throughout
func newSizeTestVault(t *testing.T) (afero.Fs, *vault.Vault) {
	t.Helper()

	fs := afero.NewMemMapFs()

	// Create vault, which selects it as the current vault
	_, err := vault.CreateVault(fs, testStateDir, testVaultName,
		testMnemonicBuffer(t), nil)
	require.NoError(t, err)

	// Get vault and set up long-term key
	vlt, err := vault.GetCurrentVault(fs, testStateDir)
	require.NoError(t, err)

	ltIdentity, err := agehd.DeriveIdentity(testMnemonic, 0)
	require.NoError(t, err)
	vlt.Unlock(ltIdentity)

	return fs, vlt
}

// runAddSecretSizeCase adds a secret of the given size through stdin and
// verifies that it is stored.
func runAddSecretSizeCase(t *testing.T, size int) {
	t.Helper()

	fs, vlt := newSizeTestVault(t)

	// Generate test data of specified size
	testData := make([]byte, size)
	_, err := rand.Read(testData)
	require.NoError(t, err)

	// Add newline that will be stripped
	testDataWithNewline := make([]byte, 0, len(testData)+1)
	testDataWithNewline = append(testDataWithNewline, testData...)
	testDataWithNewline = append(testDataWithNewline, '\n')

	// Create command with fake stdin
	cmd := &cobra.Command{}
	cmd.SetIn(bytes.NewReader(testDataWithNewline))

	// Create CLI instance
	cli, err := NewCLIInstance()
	if err != nil {
		t.Fatalf("failed to initialize CLI: %v", err)
	}

	cli.fs = fs
	cli.stateDir = testStateDir
	cli.cmd = cmd

	// Test adding the secret
	secretName := fmt.Sprintf("test-secret-%d", size)
	err = cli.AddSecret(secretName, false)
	require.NoError(t, err)

	// Verify the secret was stored correctly
	retrievedValue, err := vlt.GetSecret(secretName)
	require.NoError(t, err)

	defer retrievedValue.Destroy()

	assert.Equal(t, testData, retrievedValue.Bytes(),
		"Retrieved secret should match original (without newline)")
}

// runImportSecretSizeCase imports a secret file of the given size and
// verifies that it is stored.
func runImportSecretSizeCase(t *testing.T, size int) {
	t.Helper()

	fs, vlt := newSizeTestVault(t)

	// Generate test data of specified size
	testData := make([]byte, size)
	_, err := rand.Read(testData)
	require.NoError(t, err)

	// Write test data to file
	testFile := fmt.Sprintf("/test/secret-%d.bin", size)
	err = afero.WriteFile(fs, testFile, testData, 0o600)
	require.NoError(t, err)

	// Create command
	cmd := &cobra.Command{}

	// Create CLI instance
	cli, err := NewCLIInstance()
	if err != nil {
		t.Fatalf("failed to initialize CLI: %v", err)
	}

	cli.fs = fs
	cli.stateDir = testStateDir

	// Test importing the secret
	secretName := fmt.Sprintf("imported-secret-%d", size)
	err = cli.ImportSecret(cmd, secretName, testFile, false)
	require.NoError(t, err)

	// Verify the secret was stored correctly
	retrievedValue, err := vlt.GetSecret(secretName)
	require.NoError(t, err)

	defer retrievedValue.Destroy()

	assert.Equal(t, testData, retrievedValue.Bytes(),
		"Retrieved secret should match original")
}

// TestAddSecretVariousSizes tests adding secrets of various sizes through stdin
//
//nolint:paralleltest // in parallel, size tests could exceed the memlock limit
func TestAddSecretVariousSizes(t *testing.T) {
	tests := []struct {
		name string
		size int
	}{
		{
			name: "1KB secret",
			size: 1024,
		},
		{
			name: "10KB secret",
			size: 10 * 1024,
		},
		{
			name: "100KB secret",
			size: 100 * 1024,
		},
		{
			name: "1MB secret",
			size: 1024 * 1024,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			runAddSecretSizeCase(t, tt.size)
		})
	}
}

// TestImportSecretVariousSizes tests importing secrets of various sizes from files
//
//nolint:paralleltest // in parallel, size tests could exceed the memlock limit
func TestImportSecretVariousSizes(t *testing.T) {
	tests := []struct {
		name string
		size int
	}{
		{
			name: "1KB file",
			size: 1024,
		},
		{
			name: "10KB file",
			size: 10 * 1024,
		},
		{
			name: "100KB file",
			size: 100 * 1024,
		},
		{
			name: "1MB file",
			size: 1024 * 1024,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			runImportSecretSizeCase(t, tt.size)
		})
	}
}

// TestAddSecretAroundPageSize tests that secret add, which reads its input
// one memory page (usually 4 KiB) at a time, stores values just under, at and
// just over one and two pages unchanged.
//
//nolint:paralleltest // in parallel, size tests could exceed the memlock limit
func TestAddSecretAroundPageSize(t *testing.T) {
	// Sizes around one and two pages of 4 KiB, then larger ones up to 1 MiB
	sizes := []int{
		1,       // Single byte
		100,     // Small
		4095,    // Just under 4KB
		4096,    // Exactly 4KB
		4097,    // Just over 4KB
		8191,    // Just under 8KB
		8192,    // Exactly 8KB
		8193,    // Just over 8KB
		12288,   // 12KB
		16384,   // 16KB
		32768,   // 32KB
		65536,   // 64KB
		131072,  // 128KB
		524288,  // 512KB
		1048576, // 1MB
	}

	for _, size := range sizes {
		t.Run(fmt.Sprintf("size_%d", size), func(t *testing.T) {
			fs, vlt := newSizeTestVault(t)

			// Create test data of exactly the specified size
			// Use a pattern that's easy to verify
			testData := make([]byte, size)
			for i := range testData {
				testData[i] = byte(i % 256)
			}

			// Create command with fake stdin (no newline)
			cmd := &cobra.Command{}
			cmd.SetIn(bytes.NewReader(testData))

			// Create CLI instance
			cli, err := NewCLIInstance()
			if err != nil {
				t.Fatalf("failed to initialize CLI: %v", err)
			}

			cli.fs = fs
			cli.stateDir = testStateDir
			cli.cmd = cmd

			// Test adding the secret
			secretName := fmt.Sprintf("buffer-test-%d", size)
			err = cli.AddSecret(secretName, false)
			require.NoError(t, err)

			// Verify the secret was stored correctly
			retrievedValue, err := vlt.GetSecret(secretName)
			require.NoError(t, err)

			defer retrievedValue.Destroy()

			assert.Equal(t, testData, retrievedValue.Bytes(),
				"Retrieved secret should match original exactly")
		})
	}
}

// TestAddSecretStreamingBehavior tests that we handle streaming input correctly
func TestAddSecretStreamingBehavior(t *testing.T) {
	t.Parallel()

	fs, vlt := newSizeTestVault(t)

	// Create a custom reader that simulates slow streaming input
	// This will help verify our buffer handling works correctly with partial reads
	testData := []byte(strings.Repeat("Hello, World! ", 1000)) // ~14KB
	streamingStdin := &slowReader{
		data:      testData,
		chunkSize: 1000, // Read 1KB at a time
	}

	// Create command with slow reader as stdin
	cmd := &cobra.Command{}
	cmd.SetIn(streamingStdin)

	// Create CLI instance
	cli, err := NewCLIInstance()
	if err != nil {
		t.Fatalf("failed to initialize CLI: %v", err)
	}

	cli.fs = fs
	cli.stateDir = testStateDir
	cli.cmd = cmd

	// Test adding the secret
	err = cli.AddSecret("streaming-test", false)
	require.NoError(t, err)

	// Verify the secret was stored correctly
	retrievedValue, err := vlt.GetSecret("streaming-test")
	require.NoError(t, err)

	defer retrievedValue.Destroy()

	assert.Equal(t, testData, retrievedValue.Bytes(),
		"Retrieved secret should match original")
}

// newBinaryTestStateDir returns a new state directory holding a vault made
// from the test mnemonic, for the secret binary to use.
func newBinaryTestStateDir(t *testing.T) string {
	t.Helper()

	stateDir := t.TempDir()
	_, err := vault.CreateVault(afero.NewOsFs(), stateDir, "default",
		testMnemonicBuffer(t), nil)
	require.NoError(t, err)

	return stateDir
}

// runAtLockedMemoryLimit runs the secret binary with args, stdin, the state
// directory stateDir and the test mnemonic, its locked-memory limit (ulimit
// -l) lowered to limitKiB, and returns what it printed.
func runAtLockedMemoryLimit(
	t *testing.T, stateDir string, limitKiB int, stdin io.Reader,
	args ...string,
) (string, error) {
	t.Helper()

	const commandWait = 10 * time.Second

	secretPath, err := filepath.Abs("../../secret")
	require.NoError(t, err)

	ctx, cancel := context.WithTimeout(t.Context(), commandWait)
	defer cancel()

	// The shell lowers the limit for itself and the binary only.
	script := fmt.Sprintf(`ulimit -l %d && exec "$0" "$@"`, limitKiB)

	//nolint:gosec // G204: runs the freshly built secret binary
	cmd := exec.CommandContext(ctx, "sh",
		append([]string{"-c", script, secretPath}, args...)...)
	cmd.Env = []string{
		secret.EnvStateDir + "=" + stateDir,
		secret.EnvMnemonic + "=" + testMnemonic,
	}
	cmd.Stdin = stdin

	output, err := cmd.CombinedOutput()

	return string(output), err
}

// requireRefusedAtLockedMemoryLimit runs the secret binary as
// runAtLockedMemoryLimit does, and requires it to exit with status 1 and
// print wantError, with no panic.
func requireRefusedAtLockedMemoryLimit(
	t *testing.T, stateDir string, limitKiB int, wantError string,
	stdin io.Reader, args ...string,
) {
	t.Helper()

	output, err := runAtLockedMemoryLimit(t, stateDir, limitKiB, stdin, args...)

	var exitErr *exec.ExitError

	require.ErrorAs(t, err, &exitErr, "%v: %s", args, output)
	assert.Equal(t, 1, exitErr.ExitCode(), "%v: %s", args, output)
	assert.Contains(t, output, wantError, args)
	assert.NotContains(t, output, "panic", args)
}

// TestSecretsAtLockedMemoryLimit runs the secret binary with the
// locked-memory limit (ulimit -l) lowered to 1 MiB, where memguard would crash
// the process with a panic instead of returning an error. add must store the
// largest value that fits, which takes holding it in locked memory only once.
// add, import and generate secret must refuse a value one byte larger with
// exit status 1 and an error saying why.
func TestSecretsAtLockedMemoryLimit(t *testing.T) {
	t.Parallel()

	const limitKiB = 1024

	fits := limitKiB*kibibyte - lockedMemoryReservePages*os.Getpagesize()
	limitError := fmt.Sprintf("locked-memory limit of %d KiB", limitKiB)
	stateDir := newBinaryTestStateDir(t)

	tooBigFile := filepath.Join(t.TempDir(), "too-big")
	require.NoError(t, os.WriteFile(tooBigFile, make([]byte, fits+1), 0o600))

	output, err := runAtLockedMemoryLimit(t, stateDir, limitKiB,
		bytes.NewReader(make([]byte, fits)), "add", "fits")
	require.NoError(t, err, "the largest value that fits: %s", output)

	requireRefusedAtLockedMemoryLimit(t, stateDir, limitKiB, limitError,
		bytes.NewReader(make([]byte, fits+1)), "add", "too-big")
	requireRefusedAtLockedMemoryLimit(t, stateDir, limitKiB, limitError,
		nil, "import", "too-big", "--source", tooBigFile)
	requireRefusedAtLockedMemoryLimit(t, stateDir, limitKiB, limitError,
		nil, "generate", "secret", "too-big", "--length", strconv.Itoa(fits+1))
}

// TestSecretsUnderSmallLockedMemoryLimits runs add, import and generate
// secret with a short value under two small locked-memory limits (ulimit
// -l). Under 20 KiB, too little for the buffers they need besides the value,
// each must refuse it with exit status 1 and an error naming the limit, where
// memguard would panic or hang. Under 64 KiB, the default of older Linux
// kernels, each must store it.
func TestSecretsUnderSmallLockedMemoryLimits(t *testing.T) {
	t.Parallel()

	const (
		tooSmallLimitKiB   = 20
		oldDefaultLimitKiB = 64
		shortValue         = "hunter2"
	)

	tooSmallLimitError := fmt.Sprintf("locked-memory limit of %d KiB",
		tooSmallLimitKiB)
	stateDir := newBinaryTestStateDir(t)

	shortFile := filepath.Join(t.TempDir(), "short")
	require.NoError(t, os.WriteFile(shortFile, []byte(shortValue), 0o600))

	requireRefusedAtLockedMemoryLimit(t, stateDir, tooSmallLimitKiB,
		tooSmallLimitError, strings.NewReader(shortValue), "add", "added")
	requireRefusedAtLockedMemoryLimit(t, stateDir, tooSmallLimitKiB,
		tooSmallLimitError, nil, "import", "imported", "--source", shortFile)
	requireRefusedAtLockedMemoryLimit(t, stateDir, tooSmallLimitKiB,
		tooSmallLimitError, nil, "generate", "secret", "generated")

	output, err := runAtLockedMemoryLimit(t, stateDir, oldDefaultLimitKiB,
		strings.NewReader(shortValue), "add", "added")
	require.NoError(t, err, "add: %s", output)

	output, err = runAtLockedMemoryLimit(t, stateDir, oldDefaultLimitKiB,
		nil, "import", "imported", "--source", shortFile)
	require.NoError(t, err, "import: %s", output)

	output, err = runAtLockedMemoryLimit(t, stateDir, oldDefaultLimitKiB,
		nil, "generate", "secret", "generated")
	require.NoError(t, err, "generate secret: %s", output)
}

// slowReader simulates a reader that returns data in small chunks
type slowReader struct {
	data      []byte
	offset    int
	chunkSize int
}

func (r *slowReader) Read(p []byte) (int, error) {
	if r.offset >= len(r.data) {
		return 0, io.EOF
	}

	// Read at most chunkSize bytes, bounded by the remaining data and
	// the destination buffer
	remaining := len(r.data) - r.offset
	toRead := min(r.chunkSize, remaining, len(p))

	n := copy(p, r.data[r.offset:r.offset+toRead])
	r.offset += n

	if r.offset >= len(r.data) {
		return n, io.EOF
	}

	return n, nil
}