run/runtime/protected_request.go

package runtime

import (
	"context"
	"errors"
	"fmt"
	"sync"
	"time"

	"browser.local/run/protocol"
)

const (
	ProtectedRequestStatusSucceeded = "succeeded"
	ProtectedRequestStatusFailed    = "failed"
	ProtectedRequestStatusUnknown   = "unknown"
)

// ErrProtectedRequestUnknown lets a Run-owned transport report that a request
// is syntactically safe but not one of its declared operations. It is terminal
// and affects only this request.
var ErrProtectedRequestUnknown = errors.New("protected request outcome unknown")

// ProtectedRequest contains only Platform-authorized text and logical binding.
// Handlers resolve their own private transport configuration locally; they must
// not return it in an outcome, error, log line, or result.
type ProtectedRequest struct {
	JobID            string
	ServerInstanceID string
	FencingToken     uint64
	Kind             string
	TransportKey     string
	TargetKey        string
	RequestText      string
}

type ProtectedRequestOutcome struct {
	Status string
	Stdout string
	Stderr string
}

type ProtectedRequestHandler interface {
	ExecuteProtectedRequest(context.Context, ProtectedRequest) (ProtectedRequestOutcome, error)
}

type ProtectedRequestHandlerFunc func(context.Context, ProtectedRequest) (ProtectedRequestOutcome, error)

func (fn ProtectedRequestHandlerFunc) ExecuteProtectedRequest(ctx context.Context, request ProtectedRequest) (ProtectedRequestOutcome, error) {
	return fn(ctx, request)
}

// ProtectedRequestRegistry is configured by the local Run package owner. It
// uses only logical lookup keys, so plugin and Platform payloads cannot select
// an arbitrary program, DSN, socket, or host path.
type ProtectedRequestRegistry struct {
	mu       sync.RWMutex
	handlers map[string]ProtectedRequestHandler
}

func NewProtectedRequestRegistry() *ProtectedRequestRegistry {
	return &ProtectedRequestRegistry{handlers: map[string]ProtectedRequestHandler{}}
}

func (registry *ProtectedRequestRegistry) Register(kind string, transportKey string, handler ProtectedRequestHandler) error {
	if !validProtectedRequestKind(kind) || !protocol.ValidLogicalFileKey(transportKey) || handler == nil {
		return fmt.Errorf("protected request kind, transport key, and handler are required")
	}
	registry.mu.Lock()
	defer registry.mu.Unlock()
	if registry.handlers == nil {
		registry.handlers = map[string]ProtectedRequestHandler{}
	}
	registry.handlers[protectedRequestRegistryKey(kind, transportKey)] = handler
	return nil
}

func (registry *ProtectedRequestRegistry) handler(kind string, transportKey string) (ProtectedRequestHandler, bool) {
	if registry == nil {
		return nil, false
	}
	registry.mu.RLock()
	defer registry.mu.RUnlock()
	handler, exists := registry.handlers[protectedRequestRegistryKey(kind, transportKey)]
	return handler, exists
}

func protectedRequestRegistryKey(kind string, transportKey string) string {
	return kind + "\x00" + transportKey
}

func (executor LifecycleExecutor) ExecuteProtectedRequest(ctx context.Context, assignment protocol.RunJobAssignment, input protocol.ProtectedRequestExecutionInputResponse) LifecycleExecutionResult {
	if err := protocol.ValidateRunJobAssignment(assignment); err != nil || !protectedRequestInputMatchesExecutor(input, assignment) {
		return protectedRequestFailure(assignment.Capability, ProtectedRequestStatusFailed, "protected_request_binding_invalid")
	}
	handler, exists := executor.protectedRequests.handler(input.Kind, input.TransportKey)
	if !exists {
		if input.Kind == "rcon" && assignment.ExecutionInput.SourceRCON != nil {
			return executor.executeProtectedSourceRCON(ctx, assignment, input)
		}
		return protectedRequestFailure(assignment.Capability, ProtectedRequestStatusUnknown, "protected_request_transport_unknown")
	}
	executionCtx, cancel := context.WithTimeout(ctx, time.Duration(assignment.ExecutionInput.TimeoutSeconds)*time.Second)
	defer cancel()
	outcome, err := handler.ExecuteProtectedRequest(executionCtx, ProtectedRequest{JobID: input.JobID, ServerInstanceID: input.ServerInstanceID, FencingToken: input.FencingToken, Kind: input.Kind, TransportKey: input.TransportKey, TargetKey: input.TargetKey, RequestText: input.RequestText})
	if input.Kind == "program" {
		executor.writeProtectedProgramLogs(ctx, assignment, outcome)
	}
	if errors.Is(executionCtx.Err(), context.Canceled) {
		return LifecycleExecutionResult{State: lifecycleResultStateCancelled, Progress: protocol.RunJobProgressReport{Percent: 100, Message: "protected request cancelled"}, Message: "protected request cancelled", ErrorCode: "protected_request_cancelled"}
	}
	if errors.Is(executionCtx.Err(), context.DeadlineExceeded) {
		return protectedRequestFailure(assignment.Capability, ProtectedRequestStatusFailed, "protected_request_timeout")
	}
	if errors.Is(err, ErrProtectedRequestUnknown) || outcome.Status == ProtectedRequestStatusUnknown {
		return protectedRequestFailure(assignment.Capability, ProtectedRequestStatusUnknown, "protected_request_unknown")
	}
	if err != nil || outcome.Status == ProtectedRequestStatusFailed || outcome.Status != "" && outcome.Status != ProtectedRequestStatusSucceeded {
		return protectedRequestFailure(assignment.Capability, ProtectedRequestStatusFailed, "protected_request_failed")
	}
	return LifecycleExecutionResult{State: lifecycleResultStateSucceeded, Progress: protocol.RunJobProgressReport{Percent: 100, Message: "protected request completed"}, Message: "protected request completed", ExecutionResult: protocol.RunJobExecutionResult{Kind: "protected." + input.Kind, Summary: "approved protected request executed through logical Run transport"}}
}

func (executor LifecycleExecutor) executeProtectedSourceRCON(ctx context.Context, assignment protocol.RunJobAssignment, input protocol.ProtectedRequestExecutionInputResponse) LifecycleExecutionResult {
	result := executor.ExecuteSourceRCON(ctx, assignment, input.RequestText)
	if result.State != lifecycleResultStateSucceeded {
		return result
	}
	return LifecycleExecutionResult{State: lifecycleResultStateSucceeded, Progress: protocol.RunJobProgressReport{Percent: 100, Message: "protected RCON request delivered"}, Message: "protected RCON request delivered", ExecutionResult: protocol.RunJobExecutionResult{Kind: "protected.rcon", Summary: "approved protected RCON request delivered through Source RCON"}}
}

func protectedRequestInputMatchesExecutor(input protocol.ProtectedRequestExecutionInputResponse, assignment protocol.RunJobAssignment) bool {
	return protocol.ValidProtectedRequestExecutionInput(input) && input.JobID == assignment.JobID && input.ServerInstanceID == assignment.ServerInstanceID && input.FencingToken == assignment.FencingToken && input.TargetKey == assignment.TargetKey && input.TransportKey == assignment.ExecutionInput.RemoteAdapterKey && input.Kind == protectedRequestKindForCapability(assignment.Capability)
}

func protectedRequestKindForCapability(capability string) string {
	switch capability {
	case protocol.RunCapabilityRemoteRunProtectedSQL:
		return "sql"
	case protocol.RunCapabilityRemoteRunProtectedRCON:
		return "rcon"
	case protocol.RunCapabilityRemoteRunProgram:
		return "program"
	default:
		return ""
	}
}

func validProtectedRequestKind(kind string) bool {
	return kind == "sql" || kind == "rcon" || kind == "program"
}

func (executor LifecycleExecutor) writeProtectedProgramLogs(ctx context.Context, assignment protocol.RunJobAssignment, outcome ProtectedRequestOutcome) {
	for _, item := range []struct{ stream, body string }{{"management-program.stdout", outcome.Stdout}, {"management-program.stderr", outcome.Stderr}} {
		for _, line := range splitProtectedProgramLines(item.body) {
			_ = executor.logSink.Append(ctx, assignment, item.stream, line)
		}
	}
}

func splitProtectedProgramLines(value string) []string {
	return splitRawLogLines(value)
}

func protectedRequestFailure(capability string, status string, code string) LifecycleExecutionResult {
	kind := protectedRequestKindForCapability(capability)
	if kind == "" {
		kind = "unknown"
	}
	if status == ProtectedRequestStatusUnknown {
		return LifecycleExecutionResult{State: lifecycleResultStateFailed, Progress: protocol.RunJobProgressReport{Percent: 100, Message: "protected request outcome is unknown"}, Message: "protected request outcome is unknown", ErrorCode: "protected_request_unknown", ExecutionResult: protocol.RunJobExecutionResult{Kind: "protected." + kind + ".unknown", Summary: "protected request outcome is unknown"}}
	}
	return LifecycleExecutionResult{State: lifecycleResultStateFailed, Progress: protocol.RunJobProgressReport{Percent: 100, Message: "protected request failed"}, Message: "protected request failed", ErrorCode: code, ExecutionResult: protocol.RunJobExecutionResult{Kind: "protected." + kind + ".failed", Summary: "protected request failed safely"}}
}