core/server/sandbox/files_module_test.go

package sandbox

import (
	"context"
	"fmt"
	"io"
	"strings"
	"testing"

	"go.starlark.net/starlark"

	"armature/storage"
)

// memStore is an in-memory ObjectStore for testing.
type memStore struct {
	objects map[string]memObject
}

type memObject struct {
	data        []byte
	contentType string
}

func newMemStore() *memStore {
	return &memStore{objects: make(map[string]memObject)}
}

func (m *memStore) Put(_ context.Context, key string, r io.Reader, size int64, contentType string) error {
	data, err := io.ReadAll(r)
	if err != nil {
		return err
	}
	m.objects[key] = memObject{data: data, contentType: contentType}
	return nil
}

func (m *memStore) Get(_ context.Context, key string) (io.ReadCloser, int64, string, error) {
	obj, ok := m.objects[key]
	if !ok {
		return nil, 0, "", storage.ErrNotFound
	}
	return io.NopCloser(&memReader{data: obj.data}), int64(len(obj.data)), obj.contentType, nil
}

func (m *memStore) Delete(_ context.Context, key string) error {
	delete(m.objects, key)
	return nil
}

func (m *memStore) DeletePrefix(_ context.Context, prefix string) error {
	for k := range m.objects {
		if strings.HasPrefix(k, prefix) {
			delete(m.objects, k)
		}
	}
	return nil
}

func (m *memStore) Exists(_ context.Context, key string) (bool, error) {
	_, ok := m.objects[key]
	return ok, nil
}

func (m *memStore) Healthy(_ context.Context) error { return nil }

func (m *memStore) Stats(_ context.Context) (*storage.StorageStats, error) {
	return &storage.StorageStats{Backend: "mem", Configured: true, Healthy: true}, nil
}

func (m *memStore) Backend() string { return "mem" }

func (m *memStore) List(_ context.Context, prefix string, limit int) ([]storage.ObjectEntry, error) {
	if limit <= 0 {
		limit = 100
	}
	var entries []storage.ObjectEntry
	for k, obj := range m.objects {
		if strings.HasPrefix(k, prefix) {
			entries = append(entries, storage.ObjectEntry{
				Key:         k,
				Size:        int64(len(obj.data)),
				ContentType: obj.contentType,
			})
			if len(entries) >= limit {
				break
			}
		}
	}
	return entries, nil
}

// memReader wraps a byte slice as io.ReadCloser.
type memReader struct {
	data []byte
	pos  int
}

func (r *memReader) Read(p []byte) (int, error) {
	if r.pos >= len(r.data) {
		return 0, io.EOF
	}
	n := copy(p, r.data[r.pos:])
	r.pos += n
	if r.pos >= len(r.data) {
		return n, io.EOF
	}
	return n, nil
}

func (r *memReader) Close() error { return nil }

// ── Test helpers ──────────────────────────────

func execWithFiles(t *testing.T, canWrite bool, script string) (*Result, error) {
	t.Helper()
	ctx := context.Background()
	store := newMemStore()
	mod := BuildFilesModule(ctx, FilesModuleConfig{
		PackageID: "test-pkg",
		CanWrite:  canWrite,
		Store:     store,
	})
	sb := New(DefaultConfig())
	return sb.Exec(ctx, "test.star", script, map[string]starlark.Value{
		"files": mod,
	})
}

func execWithFilesStore(t *testing.T, store *memStore, canWrite bool, script string) (*Result, error) {
	t.Helper()
	ctx := context.Background()
	mod := BuildFilesModule(ctx, FilesModuleConfig{
		PackageID: "test-pkg",
		CanWrite:  canWrite,
		Store:     store,
	})
	sb := New(DefaultConfig())
	return sb.Exec(ctx, "test.star", script, map[string]starlark.Value{
		"files": mod,
	})
}

// ── Put + Get round-trip ──────────────────────

func TestFilesPut_StringContent(t *testing.T) {
	_, err := execWithFiles(t, true, `
files.put("hello.txt", "hello world", content_type="text/plain")
result = files.get("hello.txt")
`)
	if err != nil {
		t.Fatalf("exec error: %v", err)
	}
}

func TestFilesPut_BytesContent(t *testing.T) {
	result, err := execWithFiles(t, true, `
files.put("bin.dat", b"\x00\x01\x02\xff")
got = files.get("bin.dat")
size = got["size"]
ct = got["content_type"]
`)
	if err != nil {
		t.Fatalf("exec error: %v", err)
	}

	size, _ := starlark.AsInt32(result.Globals["size"])
	if size != 4 {
		t.Errorf("size = %d, want 4", size)
	}
	ct := string(result.Globals["ct"].(starlark.String))
	if ct != "application/octet-stream" {
		t.Errorf("content_type = %q, want application/octet-stream", ct)
	}
}

func TestFilesPut_WithMetadata(t *testing.T) {
	result, err := execWithFiles(t, true, `
files.put("doc.pdf", "pdf content", content_type="application/pdf", metadata={"author": "test", "pages": 5})
meta = files.meta("doc.pdf")
author = meta["author"]
pages = meta["pages"]
`)
	if err != nil {
		t.Fatalf("exec error: %v", err)
	}

	author := string(result.Globals["author"].(starlark.String))
	if author != "test" {
		t.Errorf("author = %q, want %q", author, "test")
	}
	pages, _ := starlark.AsInt32(result.Globals["pages"])
	if pages != 5 {
		t.Errorf("pages = %d, want 5", pages)
	}
}

func TestFilesPut_SizeLimit(t *testing.T) {
	// Set a small limit for testing.
	t.Setenv("EXT_FILES_MAX_SIZE", "10")
	_, err := execWithFiles(t, true, `
files.put("big.bin", "this content is longer than ten bytes")
`)
	if err == nil {
		t.Fatal("expected size limit error")
	}
	if !strings.Contains(err.Error(), "exceeds limit") {
		t.Errorf("unexpected error: %v", err)
	}
}

// ── Get ───────────────────────────────────────

func TestFilesGet_NotFound(t *testing.T) {
	result, err := execWithFiles(t, false, `
got = files.get("nonexistent.txt")
`)
	if err != nil {
		t.Fatalf("exec error: %v", err)
	}
	if result.Globals["got"] != starlark.None {
		t.Errorf("expected None, got %v", result.Globals["got"])
	}
}

func TestFilesGet_ReturnsAllFields(t *testing.T) {
	result, err := execWithFiles(t, true, `
files.put("test.txt", "data", content_type="text/plain", metadata={"k": "v"})
got = files.get("test.txt")
has_content = "content" in got
has_ct = "content_type" in got
has_size = "size" in got
has_meta = "metadata" in got
`)
	if err != nil {
		t.Fatalf("exec error: %v", err)
	}
	for _, key := range []string{"has_content", "has_ct", "has_size", "has_meta"} {
		if result.Globals[key] != starlark.True {
			t.Errorf("%s = False, want True", key)
		}
	}
}

// ── Meta ──────────────────────────────────────

func TestFilesMeta_NotFound(t *testing.T) {
	result, err := execWithFiles(t, false, `
got = files.meta("nonexistent.txt")
`)
	if err != nil {
		t.Fatalf("exec error: %v", err)
	}
	if result.Globals["got"] != starlark.None {
		t.Errorf("expected None, got %v", result.Globals["got"])
	}
}

func TestFilesMeta_NoCompanion(t *testing.T) {
	result, err := execWithFiles(t, true, `
files.put("plain.txt", "data")
meta = files.meta("plain.txt")
count = len(meta)
`)
	if err != nil {
		t.Fatalf("exec error: %v", err)
	}
	count, _ := starlark.AsInt32(result.Globals["count"])
	if count != 0 {
		t.Errorf("meta length = %d, want 0", count)
	}
}

// ── List ──────────────────────────────────────

func TestFilesList_PrefixFilter(t *testing.T) {
	result, err := execWithFiles(t, true, `
files.put("images/a.png", "a")
files.put("images/b.png", "b")
files.put("docs/c.pdf", "c")
all_entries = files.list()
img_entries = files.list(prefix="images/")
all_count = len(all_entries)
img_count = len(img_entries)
`)
	if err != nil {
		t.Fatalf("exec error: %v", err)
	}

	allCount, _ := starlark.AsInt32(result.Globals["all_count"])
	imgCount, _ := starlark.AsInt32(result.Globals["img_count"])

	if allCount != 3 {
		t.Errorf("all_count = %d, want 3", allCount)
	}
	if imgCount != 2 {
		t.Errorf("img_count = %d, want 2", imgCount)
	}
}

func TestFilesList_ExcludesMeta(t *testing.T) {
	result, err := execWithFiles(t, true, `
files.put("file.txt", "data", metadata={"key": "val"})
entries = files.list()
count = len(entries)
`)
	if err != nil {
		t.Fatalf("exec error: %v", err)
	}

	count, _ := starlark.AsInt32(result.Globals["count"])
	if count != 1 {
		t.Errorf("count = %d, want 1 (should exclude _meta/ companion)", count)
	}
}

// ── Delete ────────────────────────────────────

func TestFilesDelete_Idempotent(t *testing.T) {
	_, err := execWithFiles(t, true, `
files.put("temp.txt", "temp")
files.delete("temp.txt")
files.delete("temp.txt")
exists = files.exists("temp.txt")
`)
	if err != nil {
		t.Fatalf("exec error: %v", err)
	}
}

func TestFilesDeletePrefix(t *testing.T) {
	result, err := execWithFiles(t, true, `
files.put("uploads/a.txt", "a")
files.put("uploads/b.txt", "b")
files.put("keep.txt", "keep")
files.delete_prefix("uploads/")
a_exists = files.exists("uploads/a.txt")
b_exists = files.exists("uploads/b.txt")
keep_exists = files.exists("keep.txt")
`)
	if err != nil {
		t.Fatalf("exec error: %v", err)
	}

	if result.Globals["a_exists"] != starlark.False {
		t.Error("uploads/a.txt should be deleted")
	}
	if result.Globals["b_exists"] != starlark.False {
		t.Error("uploads/b.txt should be deleted")
	}
	if result.Globals["keep_exists"] != starlark.True {
		t.Error("keep.txt should be preserved")
	}
}

// ── Exists ────────────────────────────────────

func TestFilesExists(t *testing.T) {
	result, err := execWithFiles(t, true, `
files.put("exists.txt", "yes")
yes = files.exists("exists.txt")
no = files.exists("nope.txt")
`)
	if err != nil {
		t.Fatalf("exec error: %v", err)
	}

	if result.Globals["yes"] != starlark.True {
		t.Error("exists should be True")
	}
	if result.Globals["no"] != starlark.False {
		t.Error("nope should be False")
	}
}

// ── Name Validation ──────────────────────────

func TestFilesNameValidation(t *testing.T) {
	badNames := []string{
		"../escape",
		"/absolute",
		"_meta/reserved",
		"",
	}

	for _, name := range badNames {
		t.Run(name, func(t *testing.T) {
			script := fmt.Sprintf(`files.exists(%q)`, name)
			_, err := execWithFiles(t, false, script)
			if err == nil {
				t.Errorf("expected error for name %q", name)
			}
		})
	}
}

// ── Write Permission Gate ────────────────────

func TestFilesWritePermission(t *testing.T) {
	// Read-only mode: put/delete/delete_prefix should not be available.
	_, err := execWithFiles(t, false, `
files.put("test.txt", "data")
`)
	if err == nil {
		t.Fatal("expected error: put should not be available in read-only mode")
	}
	if !strings.Contains(err.Error(), "no .put field") && !strings.Contains(err.Error(), "has no .put") {
		t.Errorf("unexpected error: %v", err)
	}
}