core/server/sandbox/runner.go

// Package sandbox — runner.go
//
// the module set based on granted permissions, and executes it.
//
// network_access config from package manifest.
//
// vault for provider key decryption, and provider.complete module.
//
// at startup. db.read grants query/view/list_tables; db.write adds
// insert/update/delete. If both are granted, db.write wins (superset).
//
// their own permission context. Per-invocation lib cache + cycle detection.
//
// The runner is the bridge between the package/permission system
// and the sandboxed Starlark interpreter.
package sandbox

import (
	"context"
	"database/sql"
	"fmt"
	"log"
	"os"
	"path/filepath"
	"strings"
	"sync/atomic"
	"time"

	"go.starlark.net/starlark"

	starlarkjson "go.starlark.net/lib/json"

	"armature/events"
	"armature/metrics"
	"armature/models"
	"armature/storage"
	"armature/store"
)

// sandboxStats tracks cumulative execution counters for the admin metrics endpoint.
var sandboxStats struct {
	execCount      atomic.Int64
	errorCount     atomic.Int64
	totalDurationNs atomic.Int64
}

// SandboxStats returns cumulative execution counters.
func SandboxStats() (execCount, errorCount uint64, avgDurationMs float64) {
	exec := sandboxStats.execCount.Load()
	errs := sandboxStats.errorCount.Load()
	totalNs := sandboxStats.totalDurationNs.Load()
	if exec > 0 {
		avgDurationMs = float64(totalNs) / float64(exec) / 1e6
	}
	return uint64(exec), uint64(errs), avgDurationMs
}

// RunContext carries per-invocation state that modules need but which
// varies per caller (API route vs filter vs task). Nil is safe — modules
// that need RunContext fields gracefully degrade.
type RunContext struct {
	// UserID is the acting user for provider resolution (BYOK chain).
	UserID string

	// TeamID is the team context for settings cascade resolution.
	// When set, team-scoped package settings are included in the cascade.
	TeamID string
}

// Runner executes Starlark package scripts with permission-gated modules.
type Runner struct {
	sandbox         *Sandbox
	stores          store.Stores
	packagesDir     string
	notifier        NotificationSender // nil = notifications module unavailable
	connResolver    ConnectionResolver  // nil = connections module unavailable
	db              *sql.DB              // nil = db module unavailable
	dbPostgres      bool                 // true = use $N placeholders; false = use ?
	allowPrivateIPs bool
	bus             *events.Bus
	objectStore     storage.ObjectStore  // nil = files module unavailable
	workspaceRoot   string               // empty = workspace module unavailable
	workspaceQuota  int                  // MB, 0 = unlimited
	capabilities    map[string]bool      // detected environment capabilities
}

// NewRunner creates a runner with the given sandbox and dependencies.
func NewRunner(sb *Sandbox, stores store.Stores) *Runner {
	return &Runner{
		sandbox: sb,
		stores:  stores,
	}
}

// SetNotifier attaches the notification sender for the notifications module.
func (r *Runner) SetNotifier(n NotificationSender) {
	r.notifier = n
}

// SetConnectionResolver attaches the connection resolver for the connections module.
func (r *Runner) SetConnectionResolver(cr ConnectionResolver) {
	r.connResolver = cr
}

// SetDB attaches the raw database connection for the db module.
// isPostgres controls whether $N or ? placeholders are used.
// Call this at startup after database.Connect().
func (r *Runner) SetDB(db *sql.DB, isPostgres bool) {
	r.db = db
	r.dbPostgres = isPostgres
}

// SetPackagesDir sets the disk path where package archives are extracted.
// Required for load() support — scripts are read from disk at runtime.
func (r *Runner) SetPackagesDir(dir string) {
	r.packagesDir = dir
}

// SetBus attaches the event bus for the realtime module.
func (r *Runner) SetBus(bus *events.Bus) {
	r.bus = bus
}

// SetObjectStore attaches the blob storage backend for the files module.
func (r *Runner) SetObjectStore(s storage.ObjectStore) {
	r.objectStore = s
}

// SetWorkspaceRoot sets the base directory for extension workspaces.
// quotaMB is the per-extension quota in MB (0 = unlimited).
func (r *Runner) SetWorkspaceRoot(root string, quotaMB int) {
	r.workspaceRoot = root
	r.workspaceQuota = quotaMB
}

// SetCapabilities stores the detected environment capabilities map.
// Passed to the settings module so extensions can call has_capability().
func (r *Runner) SetCapabilities(caps map[string]bool) {
	r.capabilities = caps
}

// SetAllowPrivateIPs disables the SSRF check that blocks connections to
// private/loopback IPs. For self-hosted environments where extensions
// reach internal services. Controlled by EXT_ALLOW_PRIVATE_IPS env var.
func (r *Runner) SetAllowPrivateIPs(allow bool) {
	r.allowPrivateIPs = allow
}

// ExecPackage loads a package's script from disk (primary) or manifest
// (legacy) and executes it with modules gated by granted permissions.
//
//
// The RunContext carries per-invocation state (user_id for provider
// resolution). Pass nil if no per-invocation context is needed.
//
// Returns the script's globals (which may contain callable entry points
// like on_pre_completion, on_run, etc.) and any captured print output.
func (r *Runner) ExecPackage(ctx context.Context, pkg *store.PackageRegistration, rc *RunContext) (*Result, error) {
	// Only active starlark packages can execute
	if pkg.Status != models.PackageStatusActive {
		return nil, fmt.Errorf("package %q is %s, not active", pkg.ID, pkg.Status)
	}
	if pkg.Tier != models.ExtTierStarlark {
		return nil, fmt.Errorf("package %q is tier %s, not starlark", pkg.ID, pkg.Tier)
	}

	// Load script from disk (primary) or manifest (legacy)
	script, err := r.loadScript(pkg)
	if err != nil {
		return nil, err
	}

	lc := newLibContext()

	// Build modules based on granted permissions
	modules, err := r.buildModulesWithLibCtx(ctx, pkg.ID, pkg.Manifest, rc, lc)
	if err != nil {
		return nil, fmt.Errorf("failed to build modules for %q: %w", pkg.ID, err)
	}

	// Build package-scoped load callback
	loader := r.packageLoader(pkg.ID, modules)

	log.Printf("  🔧 runner: exec %s (%d modules granted)", pkg.ID, len(modules))

	start := time.Now()
	result, err := r.sandbox.ExecWithLoader(ctx, pkg.ID+"/script.star", script, modules, loader)
	duration := time.Since(start)

	sandboxStats.execCount.Add(1)
	sandboxStats.totalDurationNs.Add(int64(duration))
	status := "success"
	if err != nil {
		sandboxStats.errorCount.Add(1)
		status = "error"
	}
	metrics.SandboxExecutionsTotal.WithLabelValues("exec", status).Inc()

	return result, err
}

// loadScript reads the entry point script from disk (primary path)
// or falls back to the legacy _starlark_script manifest field.
func (r *Runner) loadScript(pkg *store.PackageRegistration) (string, error) {
	// Primary: read from disk
	if r.packagesDir != "" {
		entryPoint := "script.star"
		if ep, ok := pkg.Manifest["entry_point"].(string); ok && ep != "" {
			entryPoint = ep
		}
		path := filepath.Join(r.packagesDir, pkg.ID, entryPoint)
		data, err := os.ReadFile(path)
		if err == nil && len(data) > 0 {
			return string(data), nil
		}
		// Fall through to legacy
	}

	// Legacy: inline in manifest
	script, ok := pkg.Manifest["_starlark_script"].(string)
	if ok && script != "" {
		return script, nil
	}

	return "", fmt.Errorf("package %q: no script.star on disk and no _starlark_script in manifest", pkg.ID)
}

// loadEntry tracks cache and cycle detection for the package loader.
type loadEntry struct {
	loading bool
	globals starlark.StringDict
}

// packageLoader builds a LoadFunc scoped to a package's directory.
// Returns nil if no packages directory is configured.
func (r *Runner) packageLoader(pkgID string, modules map[string]starlark.Value) LoadFunc {
	if r.packagesDir == "" {
		return nil // no disk = no load support
	}

	pkgDir := filepath.Join(r.packagesDir, pkgID)
	cache := make(map[string]*loadEntry) // dedup + cycle detection

	return func(thread *starlark.Thread, module string) (starlark.StringDict, error) {
		// Security: reject path traversal
		if strings.Contains(module, "..") || filepath.IsAbs(module) {
			return nil, fmt.Errorf("load: path traversal not allowed: %q", module)
		}

		// Resolve to package directory
		resolved := filepath.Join(pkgDir, module)
		if !strings.HasPrefix(filepath.Clean(resolved), filepath.Clean(pkgDir)) {
			return nil, fmt.Errorf("load: path escapes package directory: %q", module)
		}

		// Must be a .star file
		if !strings.HasSuffix(resolved, ".star") {
			return nil, fmt.Errorf("load: only .star files can be loaded: %q", module)
		}

		// Cache / cycle detection
		if entry, ok := cache[module]; ok {
			if entry.loading {
				return nil, fmt.Errorf("load: circular dependency: %q", module)
			}
			return entry.globals, nil
		}

		entry := &loadEntry{loading: true}
		cache[module] = entry

		// Read and execute
		data, err := os.ReadFile(resolved)
		if err != nil {
			return nil, fmt.Errorf("load: %q not found in package %q", module, pkgID)
		}

		// Build predeclared with same modules as entry point
		predeclared := make(starlark.StringDict, len(modules)+2)
		for k, v := range modules {
			predeclared[k] = v
		}
		// json is always available (added by ExecWithLoader for entry point;
		// sub-modules need it too)
		predeclared["json"] = starlarkjson.Module

		globals, err := starlark.ExecFile(thread, module, string(data), predeclared)
		if err != nil {
			return nil, fmt.Errorf("load: error in %q: %w", module, err)
		}

		entry.globals = globals
		entry.loading = false
		return globals, nil
	}
}

// CallEntryPoint executes a package script and calls a named function.
// This is the standard pattern for event-driven extensions:
//  1. Exec the script (defines functions)
//  2. Find the named entry point in globals
//  3. Call it with the provided arguments
//
// The RunContext carries per-invocation state. Pass nil if not needed.
//
// Returns the function's return value and captured output.
func (r *Runner) CallEntryPoint(ctx context.Context, pkg *store.PackageRegistration, entryPoint string, args starlark.Tuple, kwargs []starlark.Tuple, rc *RunContext) (starlark.Value, string, error) {
	result, err := r.ExecPackage(ctx, pkg, rc)
	if err != nil {
		return nil, "", err
	}

	fn, ok := result.Globals[entryPoint]
	if !ok {
		return nil, result.Output, fmt.Errorf("package %q has no %s function", pkg.ID, entryPoint)
	}

	callable, ok := fn.(starlark.Callable)
	if !ok {
		return nil, result.Output, fmt.Errorf("package %q: %s is not callable", pkg.ID, entryPoint)
	}

	start := time.Now()
	val, callOutput, callErr := r.sandbox.Call(ctx, callable, args, kwargs)
	duration := time.Since(start)

	sandboxStats.execCount.Add(1)
	sandboxStats.totalDurationNs.Add(int64(duration))
	status := "success"
	if callErr != nil {
		sandboxStats.errorCount.Add(1)
		status = "error"
	}
	metrics.SandboxExecutionsTotal.WithLabelValues(entryPoint, status).Inc()

	return val, result.Output + callOutput, callErr
}

// buildModules assembles the module map based on granted permissions.
// Delegates to buildModulesWithLibCtx with a nil lib context (no lib.load support).
func (r *Runner) buildModules(ctx context.Context, packageID string, manifest map[string]any, rc *RunContext) (map[string]starlark.Value, error) {
	return r.buildModulesWithLibCtx(ctx, packageID, manifest, rc, nil)
}

// buildModulesWithLibCtx assembles the module map based on granted permissions.
// The manifest is passed through so modules can read their config
// (e.g., http module reads network_access, provider reads requires_provider).
// The RunContext carries per-invocation state for user-scoped modules.
// The libContext enables lib.load() caching and cycle detection.
func (r *Runner) buildModulesWithLibCtx(ctx context.Context, packageID string, manifest map[string]any, rc *RunContext, lc *libContext) (map[string]starlark.Value, error) {
	if r.stores.ExtPermissions == nil {
		return nil, nil
	}

	granted, err := r.stores.ExtPermissions.GrantedForPackage(ctx, packageID)
	if err != nil {
		return nil, err
	}

	modules := make(map[string]starlark.Value)

	// Track tiered permission levels: 0=none, 1=read, 2=write
	dbLevel := 0
	filesLevel := 0
	hasBatchExec := false

	for _, perm := range granted {
		switch perm {
		case models.ExtPermSecretsRead:
			modules["secrets"] = BuildSecretsModule(ctx, r.stores, packageID)

		case models.ExtPermNotificationsSend:
			if r.notifier != nil {
				modules["notifications"] = BuildNotificationsModule(ctx, r.notifier, packageID)
			}

		case models.ExtPermAPIHTTP:
			httpCfg := ParseNetworkAccess(manifest)
			httpCfg.AllowPrivateIPs = r.allowPrivateIPs
			modules["http"] = BuildHTTPModule(ctx, httpCfg)

		case models.ExtPermDBRead:
			if dbLevel < 1 {
				dbLevel = 1
			}

		case models.ExtPermDBWrite:
			dbLevel = 2

		case models.ExtPermFilesRead:
			if filesLevel < 1 {
				filesLevel = 1
			}

		case models.ExtPermFilesWrite:
			filesLevel = 2

		case models.ExtPermWorkflowAccess:
			modules["workflow"] = BuildWorkflowModule(ctx, r.stores)

		case models.ExtPermConnectionsRead:
			if r.connResolver != nil && rc != nil && rc.UserID != "" {
				modules["connections"] = BuildConnectionsModule(ctx, r.connResolver, rc.UserID)
			}

		case models.ExtPermRealtimePublish:
			if r.bus != nil {
				modules["realtime"] = BuildRealtimeModule(ctx, r.bus, packageID)
			}

		case models.ExtPermWorkspaceManage:
			if r.workspaceRoot != "" {
				modules["workspace"] = BuildWorkspaceModule(ctx, WorkspaceModuleConfig{
					PackageID:     packageID,
					WorkspaceRoot: r.workspaceRoot,
					QuotaMB:       r.workspaceQuota,
				})
			}

		case models.ExtPermBatchExec:
			hasBatchExec = true
		}
	}

	// Wire db module at the highest granted level.
	if dbLevel > 0 && r.db != nil {
		modules["db"] = BuildDBModule(ctx, DBModuleConfig{
			PackageID:   packageID,
			CanWrite:    dbLevel == 2,
			DB:          r.db,
			IsPostgres:  r.dbPostgres,
			HasPgvector: r.capabilities["pgvector"],
		})
	}

	// Wire files module at the highest granted level.
	if filesLevel > 0 && r.objectStore != nil {
		modules["files"] = BuildFilesModule(ctx, FilesModuleConfig{
			PackageID: packageID,
			CanWrite:  filesLevel == 2,
			Store:     r.objectStore,
		})
	}

	// A package reads its own admin + team + user settings.
	userID := ""
	teamID := ""
	if rc != nil {
		userID = rc.UserID
		teamID = rc.TeamID
	}
	modules["settings"] = BuildSettingsModule(ctx, r.stores, packageID, userID, teamID, r.capabilities)

	// Always available — read-only check against kernel permission data
	modules["permissions"] = BuildPermissionsModule(ctx, r.stores)

	// Always available — read-only team role queries
	modules["teams"] = BuildTeamsModule(ctx, r.stores)

	// Allows any starlark package to load declared library dependencies.
	if lc != nil {
		modules["lib"] = BuildLibModule(ctx, r, packageID, rc, lc)
	}

	// Concurrent execution primitive — wired after all other modules
	// so the runner reference can construct per-branch module sets if
	// needed in future (Option B).
	if hasBatchExec {
		modules["batch"] = BuildBatchModule(ctx, r, packageID, manifest, rc, lc)
	}

	return modules, nil
}