secret/internal/cli/secrets_size_test.go

package cli

import (
	"bytes"
	"crypto/rand"
	"fmt"
	"io"
	"path/filepath"
	"strings"
	"testing"

	"git.eeqj.de/sneak/secret/internal/secret"
	"git.eeqj.de/sneak/secret/internal/vault"
	"git.eeqj.de/sneak/secret/pkg/agehd"
	"github.com/spf13/afero"
	"github.com/spf13/cobra"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

// TestAddSecretVariousSizes tests adding secrets of various sizes through stdin
func TestAddSecretVariousSizes(t *testing.T) {
	tests := []struct {
		name        string
		size        int
		shouldError bool
		errorMsg    string
	}{
		{
			name:        "1KB secret",
			size:        1024,
			shouldError: false,
		},
		{
			name:        "10KB secret",
			size:        10 * 1024,
			shouldError: false,
		},
		{
			name:        "100KB secret",
			size:        100 * 1024,
			shouldError: false,
		},
		{
			name:        "1MB secret",
			size:        1024 * 1024,
			shouldError: false,
		},
		{
			name:        "10MB secret",
			size:        10 * 1024 * 1024,
			shouldError: false,
		},
		{
			name:        "99MB secret",
			size:        99 * 1024 * 1024,
			shouldError: false,
		},
		{
			name:        "100MB secret minus 1 byte",
			size:        100*1024*1024 - 1,
			shouldError: false,
		},
		{
			name:        "101MB secret - should fail",
			size:        101 * 1024 * 1024,
			shouldError: true,
			errorMsg:    "secret too large: exceeds 100MB limit",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			// Set up test environment
			fs := afero.NewMemMapFs()
			stateDir := "/test/state"

			// Set test mnemonic
			t.Setenv(secret.EnvMnemonic, "abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon about")

			// Create vault
			vaultName := "test-vault"
			_, err := vault.CreateVault(fs, stateDir, vaultName)
			require.NoError(t, err)

			// Set current vault
			currentVaultPath := filepath.Join(stateDir, "currentvault")
			vaultPath := filepath.Join(stateDir, "vaults.d", vaultName)
			err = afero.WriteFile(fs, currentVaultPath, []byte(vaultPath), 0o600)
			require.NoError(t, err)

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

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

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

			// Add newline that will be stripped
			testDataWithNewline := append(testData, '\n')

			// Create fake stdin
			stdin := bytes.NewReader(testDataWithNewline)

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

			// Create CLI instance
			cli := NewCLIInstance()
			cli.fs = fs
			cli.stateDir = stateDir
			cli.cmd = cmd

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

			if tt.shouldError {
				assert.Error(t, err)
				assert.Contains(t, err.Error(), tt.errorMsg)
			} else {
				require.NoError(t, err)

				// Verify the secret was stored correctly
				retrievedValue, err := vlt.GetSecret(secretName)
				require.NoError(t, err)
				assert.Equal(t, testData, retrievedValue, "Retrieved secret should match original (without newline)")
			}
		})
	}
}

// TestImportSecretVariousSizes tests importing secrets of various sizes from files
func TestImportSecretVariousSizes(t *testing.T) {
	tests := []struct {
		name        string
		size        int
		shouldError bool
		errorMsg    string
	}{
		{
			name:        "1KB file",
			size:        1024,
			shouldError: false,
		},
		{
			name:        "10KB file",
			size:        10 * 1024,
			shouldError: false,
		},
		{
			name:        "100KB file",
			size:        100 * 1024,
			shouldError: false,
		},
		{
			name:        "1MB file",
			size:        1024 * 1024,
			shouldError: false,
		},
		{
			name:        "10MB file",
			size:        10 * 1024 * 1024,
			shouldError: false,
		},
		{
			name:        "99MB file",
			size:        99 * 1024 * 1024,
			shouldError: false,
		},
		{
			name:        "100MB file",
			size:        100 * 1024 * 1024,
			shouldError: false,
		},
		{
			name:        "101MB file - should fail",
			size:        101 * 1024 * 1024,
			shouldError: true,
			errorMsg:    "secret file too large: exceeds 100MB limit",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			// Set up test environment
			fs := afero.NewMemMapFs()
			stateDir := "/test/state"

			// Set test mnemonic
			t.Setenv(secret.EnvMnemonic, "abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon about")

			// Create vault
			vaultName := "test-vault"
			_, err := vault.CreateVault(fs, stateDir, vaultName)
			require.NoError(t, err)

			// Set current vault
			currentVaultPath := filepath.Join(stateDir, "currentvault")
			vaultPath := filepath.Join(stateDir, "vaults.d", vaultName)
			err = afero.WriteFile(fs, currentVaultPath, []byte(vaultPath), 0o600)
			require.NoError(t, err)

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

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

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

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

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

			// Create CLI instance
			cli := NewCLIInstance()
			cli.fs = fs
			cli.stateDir = stateDir

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

			if tt.shouldError {
				assert.Error(t, err)
				assert.Contains(t, err.Error(), tt.errorMsg)
			} else {
				require.NoError(t, err)

				// Verify the secret was stored correctly
				retrievedValue, err := vlt.GetSecret(secretName)
				require.NoError(t, err)
				assert.Equal(t, testData, retrievedValue, "Retrieved secret should match original")
			}
		})
	}
}

// TestAddSecretBufferGrowth tests that our buffer growth strategy works correctly
func TestAddSecretBufferGrowth(t *testing.T) {
	// Test various sizes that should trigger buffer growth
	sizes := []int{
		1,       // Single byte
		100,     // Small
		4095,    // Just under initial 4KB
		4096,    // Exactly 4KB
		4097,    // Just over 4KB
		8191,    // Just under 8KB (first double)
		8192,    // Exactly 8KB
		8193,    // Just over 8KB
		12288,   // 12KB (should trigger second double)
		16384,   // 16KB
		32768,   // 32KB (after more doublings)
		65536,   // 64KB
		131072,  // 128KB
		524288,  // 512KB
		1048576, // 1MB
		2097152, // 2MB
	}

	for _, size := range sizes {
		t.Run(fmt.Sprintf("size_%d", size), func(t *testing.T) {
			// Set up test environment
			fs := afero.NewMemMapFs()
			stateDir := "/test/state"

			// Set test mnemonic
			t.Setenv(secret.EnvMnemonic, "abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon about")

			// Create vault
			vaultName := "test-vault"
			_, err := vault.CreateVault(fs, stateDir, vaultName)
			require.NoError(t, err)

			// Set current vault
			currentVaultPath := filepath.Join(stateDir, "currentvault")
			vaultPath := filepath.Join(stateDir, "vaults.d", vaultName)
			err = afero.WriteFile(fs, currentVaultPath, []byte(vaultPath), 0o600)
			require.NoError(t, err)

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

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

			// 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 fake stdin without newline
			stdin := bytes.NewReader(testData)

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

			// Create CLI instance
			cli := NewCLIInstance()
			cli.fs = fs
			cli.stateDir = stateDir
			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)
			assert.Equal(t, testData, retrievedValue, "Retrieved secret should match original exactly")
		})
	}
}

// TestAddSecretStreamingBehavior tests that we handle streaming input correctly
func TestAddSecretStreamingBehavior(t *testing.T) {
	// Set up test environment
	fs := afero.NewMemMapFs()
	stateDir := "/test/state"

	// Set test mnemonic
	t.Setenv(secret.EnvMnemonic, "abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon abandon about")

	// Create vault
	vaultName := "test-vault"
	_, err := vault.CreateVault(fs, stateDir, vaultName)
	require.NoError(t, err)

	// Set current vault
	currentVaultPath := filepath.Join(stateDir, "currentvault")
	vaultPath := filepath.Join(stateDir, "vaults.d", vaultName)
	err = afero.WriteFile(fs, currentVaultPath, []byte(vaultPath), 0o600)
	require.NoError(t, err)

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

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

	// 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
	slowReader := &slowReader{
		data:      testData,
		chunkSize: 1000, // Read 1KB at a time
	}

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

	// Create CLI instance
	cli := NewCLIInstance()
	cli.fs = fs
	cli.stateDir = stateDir
	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)
	assert.Equal(t, testData, retrievedValue, "Retrieved secret should match original")
}

// 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) (n int, err error) {
	if r.offset >= len(r.data) {
		return 0, io.EOF
	}

	// Read at most chunkSize bytes
	remaining := len(r.data) - r.offset
	toRead := r.chunkSize
	if toRead > remaining {
		toRead = remaining
	}
	if toRead > len(p) {
		toRead = len(p)
	}

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

	if r.offset >= len(r.data) {
		err = io.EOF
	}

	return n, err
}