browser/plugins/examples/scum-server-plugin/companion/dispatcher.go

package companion

import (
	"context"
	"fmt"
	"strings"
	"sync"
	"time"
)

// SafeAdapter is intentionally narrow: it receives typed values only and has
// no raw RCON, SQL, host-path, credential, or shell access.
type SafeAdapter interface {
	ReadConfiguration(context.Context) (map[string]any, error)
	PatchConfiguration(context.Context, map[string]any) (map[string]any, error)
	Diagnostics(context.Context) (map[string]any, error)
	PatchGameState(context.Context, map[string]any) (map[string]any, error)
	DeliverReward(context.Context, map[string]any) (map[string]any, error)
	NotifyPlayer(context.Context, map[string]any) (map[string]any, error)
}

type HandlerAvailability struct {
	ServerVersion string
	Capabilities  map[string]bool
	Approved      bool
}
type CommandHandler func(context.Context, map[string]any) (map[string]any, error)
type HandlerRegistry struct {
	availability HandlerAvailability
	handlers     map[string]CommandHandler
	mu           sync.Mutex
	completed    map[string]CommandResult
}

func NewHandlerRegistry(availability HandlerAvailability, adapter SafeAdapter) *HandlerRegistry {
	registry := &HandlerRegistry{availability: availability, handlers: map[string]CommandHandler{}, completed: map[string]CommandResult{}}
	if adapter == nil {
		return registry
	}
	registry.handlers["config.read"] = func(ctx context.Context, _ map[string]any) (map[string]any, error) {
		return adapter.ReadConfiguration(ctx)
	}
	registry.handlers["config.patch"] = func(ctx context.Context, payload map[string]any) (map[string]any, error) {
		return adapter.PatchConfiguration(ctx, payload)
	}
	registry.handlers["companion.diagnostics"] = func(ctx context.Context, _ map[string]any) (map[string]any, error) { return adapter.Diagnostics(ctx) }
	registry.handlers["game-state.patch"] = func(ctx context.Context, payload map[string]any) (map[string]any, error) {
		return adapter.PatchGameState(ctx, payload)
	}
	registry.handlers["reward.deliver"] = func(ctx context.Context, payload map[string]any) (map[string]any, error) {
		return adapter.DeliverReward(ctx, payload)
	}
	registry.handlers["player.notify"] = func(ctx context.Context, payload map[string]any) (map[string]any, error) {
		return adapter.NotifyPlayer(ctx, payload)
	}
	return registry
}

func (registry *HandlerRegistry) Execute(ctx context.Context, command ClaimedCommand) (CommandResult, error) {
	registry.mu.Lock()
	cached, done := registry.completed[command.ID]
	registry.mu.Unlock()
	if done {
		return cached, nil
	}
	if err := validateDeclaredCommandAt(command, time.Now); err != nil {
		return unsupportedResult("validation-failed"), nil
	}
	if !registry.availability.Approved || !registry.availability.Capabilities[command.CommandType] {
		return unsupportedResult("unsupported"), nil
	}
	handler, exists := registry.handlers[command.CommandType]
	if !exists || strings.TrimSpace(registry.availability.ServerVersion) == "" {
		return unsupportedResult("unsupported"), nil
	}
	payload, err := handler(ctx, command.Payload)
	if err != nil {
		return CommandResult{Status: "failed", Summary: "typed adapter failed", Payload: map[string]any{"result": "failed"}}, nil
	}
	result := CommandResult{Status: "succeeded", Summary: "typed adapter completed", Payload: redactTypedPayload(payload)}
	registry.mu.Lock()
	registry.completed[command.ID] = result
	registry.mu.Unlock()
	return result, nil
}

func validateDeclaredCommandAt(command ClaimedCommand, now func() time.Time) error {
	if command.ID == "" || command.ProfileKey != ProfileKey || command.FencingToken == 0 || command.Payload == nil || command.LeaseExpiresAt.IsZero() || command.ExpiresAt.IsZero() || !now().Before(command.LeaseExpiresAt) || !now().Before(command.ExpiresAt) {
		return fmt.Errorf("invalid command")
	}
	for key, value := range command.Payload {
		if !safeCommandField(key, value) {
			return fmt.Errorf("unsafe payload")
		}
	}
	return nil
}

func safeCommandField(key string, value any) bool {
	lower := strings.ToLower(strings.TrimSpace(key))
	if lower == "" || strings.Contains(lower, "path") || strings.Contains(lower, "credential") || strings.Contains(lower, "password") || strings.Contains(lower, "sql") || strings.Contains(lower, "rcon") || strings.Contains(lower, "command") {
		return false
	}
	if text, ok := value.(string); ok {
		compact := strings.ToLower(text)
		return !strings.Contains(compact, "bearer ") && !strings.Contains(compact, "password=") && !strings.Contains(compact, "select ") && !strings.Contains(compact, "/users/")
	}
	return true
}

func unsupportedResult(code string) CommandResult {
	return CommandResult{Status: "failed", Summary: "typed operation unavailable", Payload: map[string]any{"result": code}}
}
func redactTypedPayload(payload map[string]any) map[string]any {
	result := map[string]any{}
	for key, value := range payload {
		if safeCommandField(key, value) {
			result[key] = value
		}
	}
	return result
}

type CommandGateway interface {
	ClaimCommands(context.Context, int) ([]ClaimedCommand, error)
	AckCommand(context.Context, string, uint64) (CommandAck, error)
	CompleteCommand(context.Context, string, uint64, CommandResult) (CompletedCommand, error)
}
type Dispatcher struct {
	Client    CommandGateway
	Registry  *HandlerRegistry
	PollLimit int
	Backoff   time.Duration
	Now       func() time.Time
}

func (dispatcher Dispatcher) DispatchOnce(ctx context.Context) error {
	if dispatcher.Client == nil || dispatcher.Registry == nil {
		return fmt.Errorf("dispatcher is not configured")
	}
	limit := dispatcher.PollLimit
	if limit == 0 {
		limit = 10
	}
	if limit < 1 || limit > 50 {
		return fmt.Errorf("dispatcher poll limit is invalid")
	}
	commands, err := dispatcher.Client.ClaimCommands(ctx, limit)
	if err != nil {
		return err
	}
	for _, command := range commands {
		if err := validateDeclaredCommandAt(command, dispatcher.clock()); err != nil {
			if completeErr := dispatcher.completeValidationFailure(ctx, command); completeErr != nil {
				return completeErr
			}
			continue
		}
		if _, err = dispatcher.Client.AckCommand(ctx, command.ID, command.FencingToken); err != nil {
			return err
		}
		result, executionErr := dispatcher.Registry.Execute(ctx, command)
		if executionErr != nil {
			result = CommandResult{Status: "failed", Summary: "typed adapter failed", Payload: map[string]any{"result": "failed"}}
		}
		if _, err = dispatcher.Client.CompleteCommand(ctx, command.ID, command.FencingToken, result); err != nil {
			return err
		}
	}
	return nil
}
func (dispatcher Dispatcher) completeValidationFailure(ctx context.Context, command ClaimedCommand) error {
	if command.ID == "" || command.FencingToken == 0 {
		return fmt.Errorf("claimed command is invalid")
	}
	_, err := dispatcher.Client.CompleteCommand(ctx, command.ID, command.FencingToken, unsupportedResult("validation-failed"))
	return err
}
func (dispatcher Dispatcher) clock() func() time.Time {
	if dispatcher.Now != nil {
		return dispatcher.Now
	}
	return time.Now
}
func (dispatcher Dispatcher) Run(ctx context.Context) error {
	backoff := dispatcher.Backoff
	if backoff <= 0 {
		backoff = 2 * time.Second
	}
	for {
		if err := dispatcher.DispatchOnce(ctx); err != nil {
			select {
			case <-ctx.Done():
				return ctx.Err()
			case <-time.After(backoff):
				continue
			}
		}
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-time.After(backoff):
		}
	}
}

// Runtime keeps the registered companion alive with bounded heartbeat and
// polling intervals. It owns no host connection or game credential; handlers
// are the only route to a version-bound adapter.
type RuntimeGateway interface {
	CommandGateway
	Register(context.Context) (Registration, error)
	Heartbeat(context.Context, HealthReport) (HeartbeatResult, error)
}
type Runtime struct {
	Client         RuntimeGateway
	Dispatcher     Dispatcher
	HeartbeatEvery time.Duration
	PollEvery      time.Duration
	Backoff        time.Duration
	Health         func() HealthReport
}

func (runtime Runtime) Run(ctx context.Context) error {
	if runtime.Client == nil || runtime.Dispatcher.Registry == nil {
		return fmt.Errorf("runtime is not configured")
	}
	if _, err := runtime.Client.Register(ctx); err != nil {
		return err
	}
	heartbeatEvery, pollEvery := runtime.HeartbeatEvery, runtime.PollEvery
	if heartbeatEvery < 5*time.Second {
		heartbeatEvery = 30 * time.Second
	}
	if pollEvery < time.Second {
		pollEvery = 5 * time.Second
	}
	heartbeat := time.NewTicker(heartbeatEvery)
	defer heartbeat.Stop()
	poll := time.NewTicker(pollEvery)
	defer poll.Stop()
	for {
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-heartbeat.C:
			report := HealthReport{Status: "healthy", Reason: "typed companion dispatcher ready"}
			if runtime.Health != nil {
				report = runtime.Health()
			}
			if _, err := runtime.Client.Heartbeat(ctx, report); err != nil {
				return err
			}
		case <-poll.C:
			if err := runtime.Dispatcher.DispatchOnce(ctx); err != nil {
				select {
				case <-ctx.Done():
					return ctx.Err()
				case <-time.After(runtime.backoff()):
				}
			}
		}
	}
}
func (runtime Runtime) backoff() time.Duration {
	if runtime.Backoff > 0 {
		return runtime.Backoff
	}
	return 2 * time.Second
}