core/server/sandbox/workspace_module.go

// Package sandbox — workspace_module.go
//
// Permission: workspace.manage
//
// Starlark API:
//   path = workspace.create(name)   → absolute path string (idempotent)
//   path = workspace.path(name)     → path string or None
//   names = workspace.list()        → list of workspace name strings
//   workspace.delete(name)          → True (idempotent, destructive)
//   bytes = workspace.usage(name)   → int (disk usage in bytes)
//
// Managed disk directories for extensions that need a real filesystem
// (git, compilers, media tools). All directories are scoped to
// {WORKSPACE_ROOT}/{packageID}/{name}/ — extensions cannot access each
// other's workspaces. Optional quota enforcement via WORKSPACE_QUOTA_MB.
package sandbox

import (
	"context"
	"fmt"
	"io/fs"
	"os"
	"path/filepath"
	"regexp"
	"strings"
	"time"

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

var workspaceNameRe = regexp.MustCompile(`^[a-z][a-z0-9_]{0,62}$`)

// WorkspaceModuleConfig holds the configuration for the workspace module.
type WorkspaceModuleConfig struct {
	PackageID     string
	WorkspaceRoot string
	QuotaMB       int // 0 = unlimited
}

// BuildWorkspaceModule creates the "workspace" module.
func BuildWorkspaceModule(ctx context.Context, cfg WorkspaceModuleConfig) *starlarkstruct.Module {
	fns := starlark.StringDict{
		"create": starlark.NewBuiltin("workspace.create", workspaceCreate(ctx, cfg)),
		"path":   starlark.NewBuiltin("workspace.path", workspacePath(ctx, cfg)),
		"list":   starlark.NewBuiltin("workspace.list", workspaceList(ctx, cfg)),
		"delete": starlark.NewBuiltin("workspace.delete", workspaceDelete(ctx, cfg)),
		"usage":  starlark.NewBuiltin("workspace.usage", workspaceUsage(ctx, cfg)),
	}
	return MakeModule("workspace", fns)
}

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

func workspaceCreate(ctx context.Context, cfg WorkspaceModuleConfig) 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
		}

		dir, err := cfg.safePath(name)
		if err != nil {
			return nil, fmt.Errorf("workspace.create: %w", err)
		}

		// Enforce quota before creating.
		if cfg.QuotaMB > 0 {
			quotaCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
			defer cancel()
			used, err := dirUsage(quotaCtx, cfg.packageDir())
			if err != nil && !os.IsNotExist(err) {
				return nil, fmt.Errorf("workspace.create: quota check: %w", err)
			}
			limit := int64(cfg.QuotaMB) * 1024 * 1024
			if used >= limit {
				return nil, fmt.Errorf("workspace.create: quota exceeded (%d bytes used, limit %d)", used, limit)
			}
		}

		if err := os.MkdirAll(dir, 0755); err != nil {
			return nil, fmt.Errorf("workspace.create: %w", err)
		}

		return starlark.String(dir), nil
	}
}

func workspacePath(ctx context.Context, cfg WorkspaceModuleConfig) 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
		}

		dir, err := cfg.safePath(name)
		if err != nil {
			return nil, fmt.Errorf("workspace.path: %w", err)
		}

		info, err := os.Stat(dir)
		if err != nil || !info.IsDir() {
			return starlark.None, nil
		}

		return starlark.String(dir), nil
	}
}

func workspaceList(ctx context.Context, cfg WorkspaceModuleConfig) 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) {
		if err := starlark.UnpackPositionalArgs(b.Name(), args, kwargs, 0); err != nil {
			return nil, err
		}

		entries, err := os.ReadDir(cfg.packageDir())
		if os.IsNotExist(err) {
			return starlark.NewList(nil), nil
		}
		if err != nil {
			return nil, fmt.Errorf("workspace.list: %w", err)
		}

		var names []starlark.Value
		for _, e := range entries {
			if e.IsDir() && workspaceNameRe.MatchString(e.Name()) {
				names = append(names, starlark.String(e.Name()))
			}
		}

		return starlark.NewList(names), nil
	}
}

func workspaceDelete(ctx context.Context, cfg WorkspaceModuleConfig) 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
		}

		dir, err := cfg.safePath(name)
		if err != nil {
			return nil, fmt.Errorf("workspace.delete: %w", err)
		}

		if err := os.RemoveAll(dir); err != nil {
			return nil, fmt.Errorf("workspace.delete: %w", err)
		}

		return starlark.True, nil
	}
}

func workspaceUsage(ctx context.Context, cfg WorkspaceModuleConfig) 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
		}

		dir, err := cfg.safePath(name)
		if err != nil {
			return nil, fmt.Errorf("workspace.usage: %w", err)
		}

		info, statErr := os.Stat(dir)
		if statErr != nil || !info.IsDir() {
			return nil, fmt.Errorf("workspace.usage: workspace %q does not exist", name)
		}

		usageCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
		defer cancel()

		total, err := dirUsage(usageCtx, dir)
		if err != nil {
			return nil, fmt.Errorf("workspace.usage: %w", err)
		}

		return starlark.MakeInt64(total), nil
	}
}

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

// packageDir returns the root directory for this extension's workspaces.
func (cfg WorkspaceModuleConfig) packageDir() string {
	return filepath.Join(cfg.WorkspaceRoot, cfg.PackageID)
}

// safePath validates a workspace name and returns the absolute directory
// path, guarding against path traversal and symlink escape.
func (cfg WorkspaceModuleConfig) safePath(name string) (string, error) {
	if err := validateWorkspaceName(name); err != nil {
		return "", err
	}

	dir := filepath.Join(cfg.WorkspaceRoot, cfg.PackageID, name)
	cleaned := filepath.Clean(dir)

	// Ensure the cleaned path is under the package directory.
	pkgDir := filepath.Clean(cfg.packageDir())
	if !strings.HasPrefix(cleaned, pkgDir+string(filepath.Separator)) {
		return "", fmt.Errorf("path escapes workspace root")
	}

	// If the parent directory exists, resolve symlinks to detect escape.
	parent := filepath.Dir(cleaned)
	if _, err := os.Stat(parent); err == nil {
		resolved, err := filepath.EvalSymlinks(parent)
		if err != nil {
			return "", fmt.Errorf("symlink resolution failed: %w", err)
		}
		resolvedRoot := filepath.Clean(cfg.WorkspaceRoot)
		if realRoot, err := filepath.EvalSymlinks(resolvedRoot); err == nil {
			resolvedRoot = realRoot
		}
		if !strings.HasPrefix(resolved, resolvedRoot) {
			return "", fmt.Errorf("path escapes workspace root via symlink")
		}
	}

	return cleaned, nil
}

// validateWorkspaceName checks that a workspace name matches the allowed pattern.
func validateWorkspaceName(name string) error {
	if name == "" {
		return fmt.Errorf("workspace name is required")
	}
	if !workspaceNameRe.MatchString(name) {
		return fmt.Errorf("workspace name %q must match %s", name, workspaceNameRe.String())
	}
	return nil
}

// dirUsage walks a directory tree and returns the total size in bytes.
// Respects context cancellation for timeout enforcement.
func dirUsage(ctx context.Context, root string) (int64, error) {
	var total int64
	err := filepath.WalkDir(root, func(path string, d fs.DirEntry, err error) error {
		if err != nil {
			return err
		}
		if ctx.Err() != nil {
			return ctx.Err()
		}
		if !d.IsDir() {
			info, err := d.Info()
			if err != nil {
				return nil // skip unreadable files
			}
			total += info.Size()
		}
		return nil
	})
	if os.IsNotExist(err) {
		return 0, nil
	}
	return total, err
}