run/runtime/ue4ss_dll_extension_test.go

package runtime

import (
	"bytes"
	"context"
	"encoding/hex"
	"errors"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"testing"
	"time"

	"browser.local/run/protocol"
)

type ue4ssTestDownloader struct {
	mu        sync.Mutex
	payload   []byte
	calls     int
	lastURL   string
	started   chan struct{}
	unblock   <-chan struct{}
	startOnce sync.Once
}

func (downloader *ue4ssTestDownloader) Download(ctx context.Context, sourceURL string, destination string, _ int64) (int64, string, error) {
	downloader.mu.Lock()
	downloader.calls++
	downloader.lastURL = sourceURL
	payload := append([]byte(nil), downloader.payload...)
	started := downloader.started
	unblock := downloader.unblock
	downloader.mu.Unlock()
	if started != nil {
		downloader.startOnce.Do(func() { close(started) })
	}
	if unblock != nil {
		select {
		case <-unblock:
		case <-ctx.Done():
			return 0, "", ctx.Err()
		}
	}
	if err := os.MkdirAll(filepath.Dir(destination), 0o700); err != nil {
		return 0, "", err
	}
	if err := os.WriteFile(destination, payload, 0o600); err != nil {
		return 0, "", err
	}
	return int64(len(payload)), bytesChecksum(payload), nil
}

func (downloader *ue4ssTestDownloader) SetPayload(payload []byte) {
	downloader.mu.Lock()
	defer downloader.mu.Unlock()
	downloader.payload = append([]byte(nil), payload...)
}

func (downloader *ue4ssTestDownloader) Count() int {
	downloader.mu.Lock()
	defer downloader.mu.Unlock()
	return downloader.calls
}

func (downloader *ue4ssTestDownloader) LastURL() string {
	downloader.mu.Lock()
	defer downloader.mu.Unlock()
	return downloader.lastURL
}

type ue4ssManagedProcessSupervisor struct {
	mu      sync.Mutex
	starts  int
	command ProcessCommand
}

func (supervisor *ue4ssManagedProcessSupervisor) Start(_ context.Context, command ProcessCommand, identity ProcessIdentity, output ManagedProcessOutput) (ProcessIdentity, error) {
	supervisor.mu.Lock()
	supervisor.starts++
	supervisor.command = ProcessCommand{WorkDir: command.WorkDir, Args: append([]string(nil), command.Args...), Env: command.Env, Timeout: command.Timeout}
	supervisor.mu.Unlock()
	if output.Stdout != nil {
		_ = output.Stdout(identity, ManagedProcessLine{Text: "ue4ss managed process started", EndOffset: int64(len("ue4ss managed process started"))})
	}
	identity.State = "running"
	return identity, nil
}

func (supervisor *ue4ssManagedProcessSupervisor) Stop(_ context.Context, identity ProcessIdentity) (ProcessIdentity, error) {
	identity.State = "stopped"
	return identity, nil
}

func (supervisor *ue4ssManagedProcessSupervisor) Status(identity ProcessIdentity) ProcessIdentity {
	if identity.State == "" {
		identity.State = "stopped"
	}
	return identity
}

func (supervisor *ue4ssManagedProcessSupervisor) ResumeOutput(ManagedProcessOutput) {}

func (supervisor *ue4ssManagedProcessSupervisor) Starts() int {
	supervisor.mu.Lock()
	defer supervisor.mu.Unlock()
	return supervisor.starts
}

func (supervisor *ue4ssManagedProcessSupervisor) Command() ProcessCommand {
	supervisor.mu.Lock()
	defer supervisor.mu.Unlock()
	return ProcessCommand{WorkDir: supervisor.command.WorkDir, Args: append([]string(nil), supervisor.command.Args...), Env: supervisor.command.Env, Timeout: supervisor.command.Timeout}
}

type ue4ssExtensionFixture struct {
	assignment protocol.RunJobAssignment
	plan       protocol.RuntimeDLLExtensionPlan
	activePath string
	markerPath string
	previous   string
	stagePath  string
	configPath string
	modsPath   string
}

func TestUE4SSDLLExtensionSynchronizesNoOpsUpdatesAndKeepsPriorRelease(t *testing.T) {
	root := t.TempDir()
	fixture := newUE4SSExtensionFixture(t, root, true)
	payloadV1 := []byte("scum-simple-rcon DLL release one")
	downloader := &ue4ssTestDownloader{payload: payloadV1}
	supervisor := &ue4ssManagedProcessSupervisor{}
	executor := newUE4SSExtensionExecutor(root, downloader, supervisor, "windows", "amd64")

	first := executor.Execute(fixture.assignment)
	if first.State != lifecycleResultStateSucceeded {
		t.Fatalf("expected first DLL sync and start to succeed, got %+v", first)
	}
	if downloader.Count() != 1 || downloader.LastURL() != fixture.plan.ReleaseURL {
		t.Fatalf("expected one frozen release download, count=%d url=%q", downloader.Count(), downloader.LastURL())
	}
	assertUE4SSFileEquals(t, fixture.activePath, payloadV1)
	marker, found, err := loadManagedDLLExtensionMarker(fixture.markerPath)
	if err != nil || !found || !markerMatchesPlan(marker, fixture.plan) || marker.ConfigRef != "bin/ue4ss/Mods/scum_simple_rcon/config.ini" {
		t.Fatalf("expected managed release marker for first DLL, marker=%+v found=%v err=%v", marker, found, err)
	}
	config := readUE4SSFile(t, fixture.configPath)
	if !managedLoopbackRCONConfigMatches(fixture.configPath, fixture.plan.RCONPort) {
		t.Fatal("expected protected loopback RCON configuration")
	}
	password := rconPasswordFromConfig(string(config))
	if len(password) != 64 {
		t.Fatal("expected a 32-byte generated RCON password")
	}
	if _, err := hex.DecodeString(password); err != nil {
		t.Fatalf("expected hexadecimal generated RCON password: %v", err)
	}
	if info, err := os.Stat(fixture.configPath); err != nil || info.Mode().Perm() != 0o600 {
		t.Fatalf("expected protected RCON configuration permissions, info=%v err=%v", info, err)
	}
	mods := string(readUE4SSFile(t, fixture.modsPath))
	if strings.Count(mods, "scum_simple_rcon : 1\n") != 1 || strings.Contains(mods, "scum_simple_rcon : 0") {
		t.Fatalf("expected exactly one enabled managed mod entry, mods=%q", mods)
	}
	command := supervisor.Command()
	if len(command.Args) != 1 || filepath.Base(command.Args[0]) != "GameServer.exe" || strings.HasSuffix(strings.ToLower(command.Args[0]), ".dll") {
		t.Fatalf("expected normal game executable start without a DLL loader, command=%+v", command)
	}

	second := executor.Execute(fixture.assignment)
	if second.State != lifecycleResultStateSucceeded || downloader.Count() != 1 {
		t.Fatalf("expected unchanged release to skip download, result=%+v downloads=%d", second, downloader.Count())
	}
	if nextConfig := readUE4SSFile(t, fixture.configPath); !bytes.Equal(config, nextConfig) {
		t.Fatal("expected unchanged release to retain its protected RCON configuration")
	}

	payloadV2 := []byte("scum-simple-rcon DLL release two")
	planV2 := fixture.plan
	planV2.Version = "1.1.0"
	planV2.Checksum = bytesChecksum(payloadV2)
	planV2.SizeBytes = int64(len(payloadV2))
	downloader.SetPayload(payloadV2)
	updatedAssignment := fixture.assignment
	updatedAssignment.ExecutionInput.DLLExtensions = []protocol.RuntimeDLLExtensionPlan{planV2}
	updated := executor.Execute(updatedAssignment)
	if updated.State != lifecycleResultStateSucceeded || downloader.Count() != 2 {
		t.Fatalf("expected changed release to update before start, result=%+v downloads=%d", updated, downloader.Count())
	}
	assertUE4SSFileEquals(t, fixture.activePath, payloadV2)
	assertUE4SSFileEquals(t, fixture.previous, payloadV1)
	marker, found, err = loadManagedDLLExtensionMarker(fixture.markerPath)
	if err != nil || !found || !markerMatchesPlan(marker, planV2) {
		t.Fatalf("expected managed release marker for updated DLL, marker=%+v found=%v err=%v", marker, found, err)
	}
	if mods := string(readUE4SSFile(t, fixture.modsPath)); strings.Count(mods, "scum_simple_rcon : 1\n") != 1 {
		t.Fatalf("expected deterministic mod index after update, mods=%q", mods)
	}

	badPlan := planV2
	badPlan.Version = "1.2.0"
	badPlan.Checksum = bytesChecksum([]byte("expected-but-not-delivered"))
	badPlan.SizeBytes = int64(len([]byte("expected-but-not-delivered")))
	downloader.SetPayload([]byte("tampered release payload"))
	badAssignment := fixture.assignment
	badAssignment.ExecutionInput.DLLExtensions = []protocol.RuntimeDLLExtensionPlan{badPlan}
	failed := executor.Execute(badAssignment)
	if failed.State != lifecycleResultStateFailed || failed.ErrorCode != "dll_extension_verify_failed" {
		t.Fatalf("expected checksum mismatch to fail before process start, got %+v", failed)
	}
	if supervisor.Starts() != 3 {
		t.Fatalf("expected failed update not to start SCUM, starts=%d", supervisor.Starts())
	}
	assertUE4SSFileEquals(t, fixture.activePath, payloadV2)
	assertUE4SSFileEquals(t, fixture.previous, payloadV1)
	marker, found, err = loadManagedDLLExtensionMarker(fixture.markerPath)
	if err != nil || !found || !markerMatchesPlan(marker, planV2) {
		t.Fatalf("expected failed update to retain prior managed release, marker=%+v found=%v err=%v", marker, found, err)
	}
	if _, err := os.Stat(fixture.stagePath); !errors.Is(err, os.ErrNotExist) {
		t.Fatalf("expected failed staged DLL to be removed, err=%v", err)
	}
}

func TestUE4SSDLLExtensionRollsBackWhenReleaseMarkerCannotActivate(t *testing.T) {
	root := t.TempDir()
	fixture := newUE4SSExtensionFixture(t, root, true)
	payloadV1 := []byte("scum-simple-rcon DLL release one")
	downloader := &ue4ssTestDownloader{payload: payloadV1}
	supervisor := &ue4ssManagedProcessSupervisor{}
	executor := newUE4SSExtensionExecutor(root, downloader, supervisor, "windows", "amd64")
	if result := executor.Execute(fixture.assignment); result.State != lifecycleResultStateSucceeded {
		t.Fatalf("install initial release: %+v", result)
	}

	payloadV2 := []byte("scum-simple-rcon DLL release two")
	planV2 := fixture.plan
	planV2.Version = "1.1.0"
	planV2.Checksum = bytesChecksum(payloadV2)
	planV2.SizeBytes = int64(len(payloadV2))
	downloader.SetPayload(payloadV2)
	failMarkerWrite := true
	executor.runtimeFileWriter = func(path string, body []byte, mode os.FileMode) error {
		if path == fixture.markerPath && failMarkerWrite {
			failMarkerWrite = false
			return errors.New("injected marker write failure")
		}
		return writeRuntimeAtomicFile(path, body, mode)
	}
	assignment := fixture.assignment
	assignment.ExecutionInput.DLLExtensions = []protocol.RuntimeDLLExtensionPlan{planV2}
	result := executor.Execute(assignment)
	if result.State != lifecycleResultStateFailed || result.ErrorCode != "dll_extension_activation_failed" {
		t.Fatalf("expected failed marker activation to fail closed, got %+v", result)
	}
	if supervisor.Starts() != 1 {
		t.Fatalf("expected marker activation failure not to start SCUM, starts=%d", supervisor.Starts())
	}
	assertUE4SSFileEquals(t, fixture.activePath, payloadV1)
	assertUE4SSFileEquals(t, fixture.previous, payloadV1)
	marker, found, err := loadManagedDLLExtensionMarker(fixture.markerPath)
	if err != nil || !found || !markerMatchesPlan(marker, fixture.plan) {
		t.Fatalf("expected previous release marker after rollback, marker=%+v found=%v err=%v", marker, found, err)
	}
}

func TestUE4SSDLLExtensionRejectsLinuxBeforeDownloadOrStart(t *testing.T) {
	root := t.TempDir()
	fixture := newUE4SSExtensionFixture(t, root, false)
	downloader := &ue4ssTestDownloader{payload: []byte("must not download")}
	supervisor := &ue4ssManagedProcessSupervisor{}
	executor := newUE4SSExtensionExecutor(root, downloader, supervisor, "linux", "amd64")

	result := executor.Execute(fixture.assignment)
	if result.State != lifecycleResultStateFailed || result.ErrorCode != "unsupported_extension_platform" {
		t.Fatalf("expected Linux DLL extension rejection, got %+v", result)
	}
	if downloader.Count() != 0 || supervisor.Starts() != 0 {
		t.Fatalf("expected Linux rejection before download or start, downloads=%d starts=%d", downloader.Count(), supervisor.Starts())
	}
}

func TestUE4SSDLLExtensionRequiresInstalledBootstrap(t *testing.T) {
	root := t.TempDir()
	fixture := newUE4SSExtensionFixture(t, root, false)
	downloader := &ue4ssTestDownloader{payload: []byte("must not download")}
	supervisor := &ue4ssManagedProcessSupervisor{}
	executor := newUE4SSExtensionExecutor(root, downloader, supervisor, "windows", "amd64")

	result := executor.Execute(fixture.assignment)
	if result.State != lifecycleResultStateFailed || result.ErrorCode != "ue4ss_bootstrap_missing" {
		t.Fatalf("expected bootstrap precondition failure, got %+v", result)
	}
	if downloader.Count() != 0 || supervisor.Starts() != 0 {
		t.Fatalf("expected missing bootstrap to prevent download and start, downloads=%d starts=%d", downloader.Count(), supervisor.Starts())
	}
}

func TestUE4SSDLLExtensionSlowDownloadPreservesHeartbeat(t *testing.T) {
	client := newFakeWorkerClient()
	cfg := workerTestConfig(t)
	fixture := newUE4SSExtensionFixture(t, cfg.WorkspaceRoot, true)
	releaseDownload := make(chan struct{})
	downloader := &ue4ssTestDownloader{payload: []byte("scum-simple-rcon DLL release one"), started: make(chan struct{}), unblock: releaseDownload}
	supervisor := &ue4ssManagedProcessSupervisor{}
	worker, err := NewWorker(cfg, client,
		WithDependencyDownloader(downloader),
		WithManagedProcessSupervisor(supervisor),
		WithDLLExtensionRuntimeTarget("windows", "amd64"),
	)
	if err != nil {
		t.Fatalf("new worker: %v", err)
	}
	worker.state.SessionToken = "session-token"

	executionDone := make(chan LifecycleExecutionResult, 1)
	go func() { executionDone <- worker.executeAssignment(context.Background(), fixture.assignment) }()
	select {
	case <-downloader.started:
	case <-time.After(2 * time.Second):
		close(releaseDownload)
		t.Fatal("expected lifecycle job to begin its bounded DLL download")
	}

	heartbeatDone := make(chan error, 1)
	go func() { heartbeatDone <- worker.HeartbeatOnce(context.Background()) }()
	select {
	case err := <-heartbeatDone:
		if err != nil {
			close(releaseDownload)
			t.Fatalf("heartbeat during DLL download: %v", err)
		}
	case <-time.After(500 * time.Millisecond):
		close(releaseDownload)
		t.Fatal("slow DLL download blocked the control heartbeat")
	}
	if len(client.heartbeatRequests) != 1 {
		close(releaseDownload)
		t.Fatalf("expected heartbeat request during DLL download, got %d", len(client.heartbeatRequests))
	}
	close(releaseDownload)
	select {
	case result := <-executionDone:
		if result.State != lifecycleResultStateSucceeded {
			t.Fatalf("expected lifecycle job to finish after download release, got %+v", result)
		}
	case <-time.After(2 * time.Second):
		t.Fatal("lifecycle job did not finish after slow download was released")
	}
}

func newUE4SSExtensionExecutor(root string, downloader DependencyDownloader, supervisor ManagedProcessSupervisor, targetOS string, targetArch string) LifecycleExecutor {
	return NewLifecycleExecutor(
		WithLifecycleWorkspaceRoot(root),
		WithDependencyDownloader(downloader),
		WithManagedProcessSupervisor(supervisor),
		WithDLLExtensionRuntimeTarget(targetOS, targetArch),
	)
}

func newUE4SSExtensionFixture(t *testing.T, root string, bootstrap bool) ue4ssExtensionFixture {
	t.Helper()
	assignment := lifecycleAssignment(protocol.RunCapabilityProcessStart)
	assignment.TargetKey = "actions/start.json"
	assignment.ExecutionInput.WorkspaceScope = "run-local"
	scope, err := NewWorkspaceResolver(root).Scope(assignment.ServerInstanceID, assignment.ExecutionInput.WorkspaceScope)
	if err != nil {
		t.Fatalf("create workspace scope: %v", err)
	}
	writeUE4SSFixtureFile(t, filepath.Join(scope, "actions", "start.json"), []byte(`{"version":1,"action":"start","mode":"supervised","executableKey":"bin/GameServer.exe"}`), 0o600)
	executable := []byte("game server executable fixture")
	writeUE4SSFixtureFile(t, filepath.Join(scope, "bin", "GameServer.exe"), executable, 0o700)
	if bootstrap {
		writeUE4SSFixtureFile(t, filepath.Join(scope, "bin", "dwmapi.dll"), []byte("UE4SS proxy fixture"), 0o600)
		writeUE4SSFixtureFile(t, filepath.Join(scope, "bin", "UE4SS.dll"), []byte("UE4SS loader fixture"), 0o600)
		writeUE4SSFixtureFile(t, filepath.Join(scope, "bin", "ue4ss", "Mods", "mods.txt"), []byte("OtherMod : 1\nscum_simple_rcon : 0\nscum_simple_rcon : 1\n"), 0o600)
	}
	plan := protocol.RuntimeDLLExtensionPlan{
		Key:                      "scum-simple-rcon",
		Version:                  "1.0.0",
		ReleaseURL:               "https://cdn.npc0.com/scum_simple_rcon_ue4s.dll",
		Checksum:                 bytesChecksum([]byte("scum-simple-rcon DLL release one")),
		SizeBytes:                int64(len([]byte("scum-simple-rcon DLL release one"))),
		TargetKey:                "ue4ss/scum-simple-rcon",
		ModKey:                   "scum_simple_rcon",
		DLLRef:                   "ue4ss/Mods/scum_simple_rcon/dlls/main.dll",
		TargetExecutableChecksum: bytesChecksum(executable),
		UE4SSABI:                 "ue4ss-3.0",
		RCONPort:                 27015,
	}
	assignment.ExecutionInput.DLLExtensions = []protocol.RuntimeDLLExtensionPlan{plan}
	return ue4ssExtensionFixture{
		assignment: assignment,
		plan:       plan,
		activePath: filepath.Join(scope, "bin", "ue4ss", "Mods", plan.ModKey, "dlls", "main.dll"),
		markerPath: filepath.Join(scope, "runtime", "ue4ss-dll", "ue4ss", "scum-simple-rcon", "release.json"),
		previous:   filepath.Join(scope, "runtime", "ue4ss-dll", "ue4ss", "scum-simple-rcon", "previous.dll"),
		stagePath:  filepath.Join(scope, "runtime", "ue4ss-dll", "ue4ss", "scum-simple-rcon", "download.staged"),
		configPath: filepath.Join(scope, "bin", "ue4ss", "Mods", plan.ModKey, "config.ini"),
		modsPath:   filepath.Join(scope, "bin", "ue4ss", "Mods", "mods.txt"),
	}
}

func writeUE4SSFixtureFile(t *testing.T, path string, body []byte, mode os.FileMode) {
	t.Helper()
	if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil {
		t.Fatalf("create fixture directory: %v", err)
	}
	if err := os.WriteFile(path, body, mode); err != nil {
		t.Fatalf("write fixture file: %v", err)
	}
}

func assertUE4SSFileEquals(t *testing.T, path string, expected []byte) {
	t.Helper()
	if actual := readUE4SSFile(t, path); !bytes.Equal(actual, expected) {
		t.Fatalf("unexpected managed file contents at %s", filepath.Base(path))
	}
}

func readUE4SSFile(t *testing.T, path string) []byte {
	t.Helper()
	body, err := os.ReadFile(path)
	if err != nil {
		t.Fatalf("read fixture file %s: %v", filepath.Base(path), err)
	}
	return body
}

func rconPasswordFromConfig(content string) string {
	inRCON := false
	for _, rawLine := range strings.Split(strings.ReplaceAll(content, "\r\n", "\n"), "\n") {
		line := strings.TrimSpace(rawLine)
		if line == "[rcon]" {
			inRCON = true
			continue
		}
		if strings.HasPrefix(line, "[") {
			inRCON = false
			continue
		}
		if inRCON && strings.HasPrefix(line, "password=") {
			return strings.TrimPrefix(line, "password=")
		}
	}
	return ""
}