core/server/sandbox/files_module.go

// Package sandbox — files_module.go
//
// Permissions: files.read, files.write
//
// Starlark API:
//   files.put(name, content, content_type="application/octet-stream", metadata={})
//   result = files.get(name)         → dict or None
//   meta   = files.meta(name)        → dict or None
//   entries = files.list(prefix="", limit=100)
//   files.delete(name)
//   files.delete_prefix(prefix)
//   exists = files.exists(name)
//
// Bridges the ObjectStore (PVC/S3) into Starlark. All keys are scoped
// to ext/{packageID}/ — extensions cannot access each other's files.
// Metadata is stored as companion JSON at ext/{packageID}/_meta/{name}.
package sandbox

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"os"
	"strconv"
	"strings"
	"unicode"

	"go.starlark.net/starlark"
	"go.starlark.net/starlarkstruct"

	"armature/storage"
)

const (
	filesDefaultMaxSize = 50 * 1024 * 1024 // 50 MB
	filesDefaultLimit   = 100
	filesMaxNameLen     = 512
	filesMetaPrefix     = "_meta/"
)

// FilesModuleConfig holds the configuration for the files module.
type FilesModuleConfig struct {
	PackageID string
	CanWrite  bool
	Store     storage.ObjectStore
}

// BuildFilesModule creates the "files" module.
func BuildFilesModule(ctx context.Context, cfg FilesModuleConfig) *starlarkstruct.Module {
	fns := starlark.StringDict{
		"get":    starlark.NewBuiltin("files.get", filesGet(ctx, cfg)),
		"meta":   starlark.NewBuiltin("files.meta", filesMeta(ctx, cfg)),
		"list":   starlark.NewBuiltin("files.list", filesList(ctx, cfg)),
		"exists": starlark.NewBuiltin("files.exists", filesExists(ctx, cfg)),
	}

	if cfg.CanWrite {
		fns["put"] = starlark.NewBuiltin("files.put", filesPut(ctx, cfg))
		fns["delete"] = starlark.NewBuiltin("files.delete", filesDelete(ctx, cfg))
		fns["delete_prefix"] = starlark.NewBuiltin("files.delete_prefix", filesDeletePrefix(ctx, cfg))
	}

	return MakeModule("files", fns)
}

// ── Builtins ──────────────────────────────────

func filesPut(ctx context.Context, cfg FilesModuleConfig) func(*starlark.Thread, *starlark.Builtin, starlark.Tuple, []starlark.Tuple) (starlark.Value, error) {
	return func(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
		var name string
		var content starlark.Value
		var contentType string = "application/octet-stream"
		var metadata *starlark.Dict

		if err := starlark.UnpackArgs(b.Name(), args, kwargs,
			"name", &name,
			"content", &content,
			"content_type?", &contentType,
			"metadata?", &metadata,
		); err != nil {
			return nil, err
		}

		if err := validateFileName(name); err != nil {
			return nil, fmt.Errorf("files.put: %w", err)
		}

		// Extract raw bytes from content (accept String or Bytes).
		var raw []byte
		switch v := content.(type) {
		case starlark.String:
			raw = []byte(string(v))
		case starlark.Bytes:
			raw = []byte(v)
		default:
			return nil, fmt.Errorf("files.put: content must be string or bytes, got %s", content.Type())
		}

		// Enforce size limit.
		maxSize := filesMaxSize()
		if int64(len(raw)) > maxSize {
			return nil, fmt.Errorf("files.put: content size %d exceeds limit %d", len(raw), maxSize)
		}

		// Write main object.
		key := cfg.scopedKey(name)
		if err := cfg.Store.Put(ctx, key, bytes.NewReader(raw), int64(len(raw)), contentType); err != nil {
			return nil, fmt.Errorf("files.put: %w", err)
		}

		// Write metadata companion if provided.
		if metadata != nil && metadata.Len() > 0 {
			metaMap, err := starlarkDictToMap(metadata)
			if err != nil {
				return nil, fmt.Errorf("files.put: metadata: %w", err)
			}
			metaJSON, err := json.Marshal(metaMap)
			if err != nil {
				return nil, fmt.Errorf("files.put: marshal metadata: %w", err)
			}
			metaKey := cfg.metaKey(name)
			if err := cfg.Store.Put(ctx, metaKey, bytes.NewReader(metaJSON), int64(len(metaJSON)), "application/json"); err != nil {
				return nil, fmt.Errorf("files.put: write metadata: %w", err)
			}
		}

		return starlark.True, nil
	}
}

func filesGet(ctx context.Context, cfg FilesModuleConfig) func(*starlark.Thread, *starlark.Builtin, starlark.Tuple, []starlark.Tuple) (starlark.Value, error) {
	return func(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
		var name string
		if err := starlark.UnpackPositionalArgs(b.Name(), args, kwargs, 1, &name); err != nil {
			return nil, err
		}

		if err := validateFileName(name); err != nil {
			return nil, fmt.Errorf("files.get: %w", err)
		}

		key := cfg.scopedKey(name)
		rc, size, ct, err := cfg.Store.Get(ctx, key)
		if err == storage.ErrNotFound {
			return starlark.None, nil
		}
		if err != nil {
			return nil, fmt.Errorf("files.get: %w", err)
		}
		defer rc.Close()

		data, err := io.ReadAll(rc)
		if err != nil {
			return nil, fmt.Errorf("files.get: read: %w", err)
		}

		// Load metadata companion.
		metaDict := starlark.NewDict(0)
		metaKey := cfg.metaKey(name)
		if metaRC, _, _, metaErr := cfg.Store.Get(ctx, metaKey); metaErr == nil {
			metaBytes, _ := io.ReadAll(metaRC)
			metaRC.Close()
			var m map[string]any
			if json.Unmarshal(metaBytes, &m) == nil {
				metaDict = MapToDict(m)
			}
		}

		result := starlark.NewDict(4)
		result.SetKey(starlark.String("content"), starlark.Bytes(data))
		result.SetKey(starlark.String("content_type"), starlark.String(ct))
		result.SetKey(starlark.String("size"), starlark.MakeInt64(size))
		result.SetKey(starlark.String("metadata"), metaDict)

		return result, nil
	}
}

func filesMeta(ctx context.Context, cfg FilesModuleConfig) func(*starlark.Thread, *starlark.Builtin, starlark.Tuple, []starlark.Tuple) (starlark.Value, error) {
	return func(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
		var name string
		if err := starlark.UnpackPositionalArgs(b.Name(), args, kwargs, 1, &name); err != nil {
			return nil, err
		}

		if err := validateFileName(name); err != nil {
			return nil, fmt.Errorf("files.meta: %w", err)
		}

		// Check the main object exists.
		key := cfg.scopedKey(name)
		exists, err := cfg.Store.Exists(ctx, key)
		if err != nil {
			return nil, fmt.Errorf("files.meta: %w", err)
		}
		if !exists {
			return starlark.None, nil
		}

		// Load metadata companion.
		metaDict := starlark.NewDict(0)
		metaKey := cfg.metaKey(name)
		if metaRC, _, _, metaErr := cfg.Store.Get(ctx, metaKey); metaErr == nil {
			metaBytes, _ := io.ReadAll(metaRC)
			metaRC.Close()
			var m map[string]any
			if json.Unmarshal(metaBytes, &m) == nil {
				metaDict = MapToDict(m)
			}
		}

		return metaDict, nil
	}
}

func filesList(ctx context.Context, cfg FilesModuleConfig) func(*starlark.Thread, *starlark.Builtin, starlark.Tuple, []starlark.Tuple) (starlark.Value, error) {
	return func(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
		var prefix string
		var limit int = filesDefaultLimit

		if err := starlark.UnpackArgs(b.Name(), args, kwargs,
			"prefix?", &prefix,
			"limit?", &limit,
		); err != nil {
			return nil, err
		}

		if limit <= 0 || limit > 1000 {
			limit = filesDefaultLimit
		}

		// Build the full scoped prefix.
		scopedPrefix := cfg.scopedPrefix()
		if prefix != "" {
			// Validate the user-supplied prefix portion.
			if strings.Contains(prefix, "..") {
				return nil, fmt.Errorf("files.list: path traversal not allowed")
			}
			scopedPrefix += prefix
		}

		// Request more than limit to account for _meta/ entries we'll filter out.
		entries, err := cfg.Store.List(ctx, scopedPrefix, limit*2)
		if err != nil {
			return nil, fmt.Errorf("files.list: %w", err)
		}

		// Filter out _meta/ companions and strip the package prefix.
		pkgPrefix := cfg.scopedPrefix()
		var results []starlark.Value
		for _, e := range entries {
			// Strip package prefix to get the extension-relative key.
			relKey := strings.TrimPrefix(e.Key, pkgPrefix)

			// Skip metadata companions.
			if strings.HasPrefix(relKey, filesMetaPrefix) {
				continue
			}

			d := starlark.NewDict(3)
			d.SetKey(starlark.String("name"), starlark.String(relKey))
			d.SetKey(starlark.String("size"), starlark.MakeInt64(e.Size))
			d.SetKey(starlark.String("content_type"), starlark.String(e.ContentType))
			results = append(results, d)

			if len(results) >= limit {
				break
			}
		}

		return starlark.NewList(results), nil
	}
}

func filesDelete(ctx context.Context, cfg FilesModuleConfig) func(*starlark.Thread, *starlark.Builtin, starlark.Tuple, []starlark.Tuple) (starlark.Value, error) {
	return func(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
		var name string
		if err := starlark.UnpackPositionalArgs(b.Name(), args, kwargs, 1, &name); err != nil {
			return nil, err
		}

		if err := validateFileName(name); err != nil {
			return nil, fmt.Errorf("files.delete: %w", err)
		}

		// Delete main object.
		key := cfg.scopedKey(name)
		if err := cfg.Store.Delete(ctx, key); err != nil {
			return nil, fmt.Errorf("files.delete: %w", err)
		}

		// Delete metadata companion (best-effort, idempotent).
		metaKey := cfg.metaKey(name)
		cfg.Store.Delete(ctx, metaKey)

		return starlark.True, nil
	}
}

func filesDeletePrefix(ctx context.Context, cfg FilesModuleConfig) func(*starlark.Thread, *starlark.Builtin, starlark.Tuple, []starlark.Tuple) (starlark.Value, error) {
	return func(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
		var prefix string
		if err := starlark.UnpackPositionalArgs(b.Name(), args, kwargs, 1, &prefix); err != nil {
			return nil, err
		}

		if strings.Contains(prefix, "..") {
			return nil, fmt.Errorf("files.delete_prefix: path traversal not allowed")
		}

		scopedPrefix := cfg.scopedPrefix() + prefix
		if err := cfg.Store.DeletePrefix(ctx, scopedPrefix); err != nil {
			return nil, fmt.Errorf("files.delete_prefix: %w", err)
		}

		// Also delete metadata companions under this prefix.
		metaPrefix := cfg.scopedPrefix() + filesMetaPrefix + prefix
		cfg.Store.DeletePrefix(ctx, metaPrefix)

		return starlark.True, nil
	}
}

func filesExists(ctx context.Context, cfg FilesModuleConfig) func(*starlark.Thread, *starlark.Builtin, starlark.Tuple, []starlark.Tuple) (starlark.Value, error) {
	return func(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
		var name string
		if err := starlark.UnpackPositionalArgs(b.Name(), args, kwargs, 1, &name); err != nil {
			return nil, err
		}

		if err := validateFileName(name); err != nil {
			return nil, fmt.Errorf("files.exists: %w", err)
		}

		key := cfg.scopedKey(name)
		exists, err := cfg.Store.Exists(ctx, key)
		if err != nil {
			return nil, fmt.Errorf("files.exists: %w", err)
		}

		return starlark.Bool(exists), nil
	}
}

// ── Helpers ──────────────────────────────────

// scopedKey returns the full ObjectStore key for a file name.
func (cfg FilesModuleConfig) scopedKey(name string) string {
	return fmt.Sprintf("ext/%s/%s", cfg.PackageID, name)
}

// metaKey returns the ObjectStore key for a file's metadata companion.
func (cfg FilesModuleConfig) metaKey(name string) string {
	return fmt.Sprintf("ext/%s/%s%s", cfg.PackageID, filesMetaPrefix, name)
}

// scopedPrefix returns the ObjectStore key prefix for this package.
func (cfg FilesModuleConfig) scopedPrefix() string {
	return fmt.Sprintf("ext/%s/", cfg.PackageID)
}

// validateFileName checks that a file name is safe.
func validateFileName(name string) error {
	if name == "" {
		return fmt.Errorf("file name is required")
	}
	if len(name) > filesMaxNameLen {
		return fmt.Errorf("file name exceeds %d characters", filesMaxNameLen)
	}
	if strings.HasPrefix(name, "/") || strings.HasPrefix(name, "\\") {
		return fmt.Errorf("absolute paths not allowed: %q", name)
	}
	if strings.Contains(name, "..") {
		return fmt.Errorf("path traversal not allowed: %q", name)
	}
	if strings.HasPrefix(name, filesMetaPrefix) {
		return fmt.Errorf("names starting with %q are reserved", filesMetaPrefix)
	}
	for _, r := range name {
		if unicode.IsControl(r) {
			return fmt.Errorf("control characters not allowed in file name")
		}
	}
	return nil
}

// filesMaxSize returns the maximum file size from env or default.
func filesMaxSize() int64 {
	if v := os.Getenv("EXT_FILES_MAX_SIZE"); v != "" {
		if n, err := strconv.ParseInt(v, 10, 64); err == nil && n > 0 {
			return n
		}
	}
	return filesDefaultMaxSize
}

// starlarkDictToMap converts a *starlark.Dict to map[string]any, validating
// that all keys are strings.  Returns an error if a non-string key is found.
func starlarkDictToMap(d *starlark.Dict) (map[string]any, error) {
	for _, item := range d.Items() {
		if _, ok := starlark.AsString(item[0]); !ok {
			return nil, fmt.Errorf("metadata keys must be strings, got %s", item[0].Type())
		}
	}
	return DictToMap(d), nil
}