run/runtime/ue4ss_dll_extension.go

package runtime

import (
	"context"
	"crypto/rand"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"os"
	pathpkg "path"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"browser.local/run/protocol"
)

const (
	ue4ssExtensionStateRoot           = "runtime/ue4ss-dll"
	ue4ssExtensionMarkerVersion       = 2
	maxUE4SSMetadataBytes       int64 = 16 * 1024
	maxUE4SSDLLBytes            int64 = 128 * 1024 * 1024
	maxTargetExecutableBytes    int64 = 2 * 1024 * 1024 * 1024
	managedRCONConfigMarker           = "; managed by Run UE4SS DLL extension"
)

type dllExtensionError struct {
	code    string
	message string
}

func (err dllExtensionError) Error() string { return err.message }

type managedDLLExtensionMarker struct {
	Version        int       `json:"version"`
	ReleaseVersion string    `json:"releaseVersion"`
	Checksum       string    `json:"checksum"`
	SizeBytes      int64     `json:"sizeBytes"`
	ExtensionKey   string    `json:"extensionKey"`
	ModKey         string    `json:"modKey"`
	ConfigRef      string    `json:"configRef"`
	RCONPort       int       `json:"rconPort"`
	UpdatedAt      time.Time `json:"updatedAt"`
}

func (executor LifecycleExecutor) synchronizeUE4SSDLLExtensions(ctx context.Context, assignment protocol.RunJobAssignment, template LifecycleActionTemplate, scope string) error {
	if executor.runtimeTargetOS != "windows" || executor.runtimeTargetArch != "amd64" {
		return dllExtensionError{code: "unsupported_extension_platform", message: "UE4SS DLL extensions require Windows amd64"}
	}
	if err := ctx.Err(); err != nil {
		return err
	}
	executableKey := template.TargetExecutableKey
	// Older generated plugin packages do not have targetExecutableKey yet,
	// but the existing start action already carries the same logical path
	// in SERVER_EXECUTABLE_REF. Keep those packages forward-compatible.
	if executableKey == "" && template.Environment != nil {
		executableKey = template.Environment["SERVER_EXECUTABLE_REF"]
	}
	if executableKey == "" && template.Env != nil {
		executableKey = template.Env["SERVER_EXECUTABLE_REF"]
	}
	if executableKey == "" && strings.HasSuffix(strings.ToLower(template.ExecutableKey), ".exe") {
		executableKey = template.ExecutableKey
	}
	if executableKey == "" || !strings.HasSuffix(strings.ToLower(executableKey), ".exe") {
		return dllExtensionError{code: "extension_target_executable_invalid", message: "UE4SS DLL extensions require a declared target executable"}
	}

	resolver := NewWorkspaceResolver(executor.workspaceRoot)
	targetResolver, targetScope := resolver, scope
	if deployment := assignment.ExecutionInput.Deployment; deployment != nil && deployment.ServerRoot != "" {
		root := filepath.Clean(deployment.ServerRoot)
		if root == "." || !filepath.IsAbs(root) {
			return dllExtensionError{code: "extension_target_executable_invalid", message: "declared executable root is unsafe"}
		}
		targetResolver = NewWorkspaceResolver(filepath.Dir(root))
		targetScope = root
	}
	executable, err := declaredLifecycleExecutable(targetResolver, targetScope, executableKey)
	if err != nil {
		return dllExtensionError{code: "extension_target_executable_invalid", message: "declared target executable is unavailable"}
	}
	executableChecksum, _, err := checksumRegularFile(executable, maxTargetExecutableBytes)
	if err != nil {
		return dllExtensionError{code: "extension_target_executable_invalid", message: "declared target executable cannot be verified"}
	}
	for _, plan := range assignment.ExecutionInput.DLLExtensions {
		if !strings.EqualFold(executableChecksum, plan.TargetExecutableChecksum) {
			return dllExtensionError{code: "extension_target_checksum_mismatch", message: "declared target executable does not match the extension release"}
		}
	}

	gameRootKey, err := gameRootKeyForExecutable(template.ExecutableKey)
	if err != nil {
		return dllExtensionError{code: "extension_target_executable_invalid", message: "declared target executable location is unsafe"}
	}
	if err := verifyUE4SSBootstrap(targetResolver, targetScope, gameRootKey); err != nil {
		return err
	}
	for _, plan := range assignment.ExecutionInput.DLLExtensions {
		if err := executor.synchronizeUE4SSDLLExtension(ctx, targetResolver, targetScope, gameRootKey, plan); err != nil {
			return err
		}
	}
	return nil
}

func declaredLifecycleExecutable(resolver WorkspaceResolver, scope string, executableKey string) (string, error) {
	if !protocol.ValidLogicalFileKey(filepath.ToSlash(executableKey)) {
		return "", fmt.Errorf("declared executable key is unsafe")
	}
	return resolver.ExistingTarget(scope, executableKey)
}

func (executor LifecycleExecutor) synchronizeUE4SSDLLExtension(ctx context.Context, resolver WorkspaceResolver, scope string, gameRootKey string, plan protocol.RuntimeDLLExtensionPlan) error {
	activeKey := gameRelativeKey(gameRootKey, plan.DLLRef)
	configRef := managedRCONConfigRef(gameRootKey, plan.ModKey)
	activePath, _, err := resolver.WritableTarget(scope, activeKey)
	if err != nil {
		return dllExtensionError{code: "dll_extension_workspace_failed", message: "extension deployment workspace is unavailable"}
	}
	markerPath, stagePath, previousPath, err := extensionStatePaths(resolver, scope, plan)
	if err != nil {
		return dllExtensionError{code: "dll_extension_workspace_failed", message: "extension state workspace is unavailable"}
	}
	marker, markerFound, err := loadManagedDLLExtensionMarker(markerPath)
	if err != nil {
		return dllExtensionError{code: "dll_extension_state_failed", message: "extension release state cannot be read"}
	}
	unchanged := markerFound && markerMatchesDeployment(marker, plan, configRef) && managedDLLMatchesPlan(activePath, plan)
	if !unchanged {
		_ = os.Remove(stagePath)
		defer os.Remove(stagePath)
		downloadedSize, downloadedChecksum, downloadErr := executor.dependencyDownloader.Download(ctx, plan.ReleaseURL, stagePath, plan.SizeBytes)
		if downloadErr != nil {
			if ctx.Err() != nil {
				return ctx.Err()
			}
			return dllExtensionError{code: "dll_extension_download_failed", message: "declared DLL download failed"}
		}
		verifiedChecksum, verifiedSize, verifyErr := checksumRegularFile(stagePath, plan.SizeBytes)
		if verifyErr != nil || downloadedSize != plan.SizeBytes || verifiedSize != plan.SizeBytes || !strings.EqualFold(downloadedChecksum, plan.Checksum) || !strings.EqualFold(verifiedChecksum, plan.Checksum) {
			return dllExtensionError{code: "dll_extension_verify_failed", message: "declared DLL did not match its fixed release checksum"}
		}
	}
	if err := executor.ensureLoopbackRCONConfig(resolver, scope, gameRootKey, plan); err != nil {
		return err
	}
	if err := executor.ensureUE4SSModsIndex(resolver, scope, gameRootKey, plan.ModKey); err != nil {
		return err
	}
	if unchanged {
		return nil
	}
	if err := executor.activateManagedDLLExtension(activePath, stagePath, previousPath, markerPath, plan, configRef); err != nil {
		return err
	}
	return nil
}

func verifyUE4SSBootstrap(resolver WorkspaceResolver, scope string, gameRootKey string) error {
	for _, filename := range []string{"dwmapi.dll", "UE4SS.dll"} {
		if _, err := resolver.ExistingTarget(scope, gameRelativeKey(gameRootKey, filename)); err != nil {
			return dllExtensionError{code: "ue4ss_bootstrap_missing", message: "required UE4SS bootstrap files are not installed"}
		}
	}
	if err := existingRuntimeDirectory(scope, gameRelativeKey(gameRootKey, "ue4ss")); err != nil {
		return dllExtensionError{code: "ue4ss_bootstrap_missing", message: "required UE4SS bootstrap files are not installed"}
	}
	return nil
}

func gameRootKeyForExecutable(executableKey string) (string, error) {
	normalized := filepath.ToSlash(executableKey)
	if !protocol.ValidLogicalFileKey(normalized) || strings.HasPrefix(normalized, "/") || strings.Contains(normalized, `\`) {
		return "", fmt.Errorf("executable key is unsafe")
	}
	parent := pathpkg.Dir(normalized)
	if parent == "." {
		return "", nil
	}
	return parent, nil
}

func gameRelativeKey(gameRootKey string, relativeKey string) string {
	if gameRootKey == "" {
		return relativeKey
	}
	return gameRootKey + "/" + relativeKey
}

func extensionStatePaths(resolver WorkspaceResolver, scope string, plan protocol.RuntimeDLLExtensionPlan) (string, string, string, error) {
	baseKey := ue4ssExtensionStateRoot + "/" + plan.TargetKey
	markerPath, _, err := resolver.WritableTarget(scope, baseKey+"/release.json")
	if err != nil {
		return "", "", "", err
	}
	stagePath, _, err := resolver.WritableTarget(scope, baseKey+"/download.staged")
	if err != nil {
		return "", "", "", err
	}
	previousPath, _, err := resolver.WritableTarget(scope, baseKey+"/previous.dll")
	if err != nil {
		return "", "", "", err
	}
	return markerPath, stagePath, previousPath, nil
}

func loadManagedDLLExtensionMarker(path string) (managedDLLExtensionMarker, bool, error) {
	body, found, err := readBoundedRegularFile(path, maxUE4SSMetadataBytes)
	if err != nil || !found {
		return managedDLLExtensionMarker{}, found, err
	}
	var marker managedDLLExtensionMarker
	if err := json.Unmarshal(body, &marker); err != nil {
		return managedDLLExtensionMarker{}, false, nil
	}
	if marker.Version != ue4ssExtensionMarkerVersion || !protocol.ValidLogicalFileKey(marker.ExtensionKey) || !protocol.ValidLogicalFileKey(marker.ModKey) || !managedRCONConfigRefForMod(marker.ConfigRef, marker.ModKey) || !protocolValidSHA256(marker.Checksum) || marker.SizeBytes < 1 || marker.ReleaseVersion == "" || marker.RCONPort < 1024 || marker.RCONPort > 65535 {
		return managedDLLExtensionMarker{}, false, nil
	}
	return marker, true, nil
}

func markerMatchesPlan(marker managedDLLExtensionMarker, plan protocol.RuntimeDLLExtensionPlan) bool {
	return marker.Version == ue4ssExtensionMarkerVersion && marker.ReleaseVersion == plan.Version && strings.EqualFold(marker.Checksum, plan.Checksum) && marker.SizeBytes == plan.SizeBytes && marker.ExtensionKey == plan.Key && marker.ModKey == plan.ModKey && marker.RCONPort == plan.RCONPort
}

func markerMatchesDeployment(marker managedDLLExtensionMarker, plan protocol.RuntimeDLLExtensionPlan, configRef string) bool {
	return markerMatchesPlan(marker, plan) && marker.ConfigRef == configRef
}

func managedDLLMatchesPlan(path string, plan protocol.RuntimeDLLExtensionPlan) bool {
	checksum, size, err := checksumRegularFile(path, plan.SizeBytes)
	return err == nil && size == plan.SizeBytes && strings.EqualFold(checksum, plan.Checksum)
}

func (executor LifecycleExecutor) activateManagedDLLExtension(activePath string, stagePath string, previousPath string, markerPath string, plan protocol.RuntimeDLLExtensionPlan, configRef string) error {
	if _, _, err := checksumRegularFile(stagePath, plan.SizeBytes); err != nil {
		return dllExtensionError{code: "dll_extension_verify_failed", message: "staged DLL cannot be verified"}
	}
	previousMarker, previousMarkerFound, markerErr := readBoundedRegularFile(markerPath, maxUE4SSMetadataBytes)
	if markerErr != nil {
		return dllExtensionError{code: "dll_extension_state_failed", message: "extension release state cannot be read"}
	}
	activeExists := false
	if _, _, err := checksumRegularFile(activePath, maxUE4SSDLLBytes); err == nil {
		activeExists = true
		if err := copyRegularFileAtomic(activePath, previousPath, maxUE4SSDLLBytes, 0o600); err != nil {
			return dllExtensionError{code: "dll_extension_activation_failed", message: "previous DLL could not be retained"}
		}
	} else if !errors.Is(err, os.ErrNotExist) {
		return dllExtensionError{code: "dll_extension_activation_failed", message: "current DLL cannot be safely replaced"}
	}
	if err := os.Rename(stagePath, activePath); err != nil {
		return dllExtensionError{code: "dll_extension_activation_failed", message: "verified DLL could not be activated"}
	}
	if err := os.Chmod(activePath, 0o600); err != nil {
		rollbackManagedDLLExtension(activePath, previousPath, activeExists)
		return dllExtensionError{code: "dll_extension_activation_failed", message: "activated DLL permissions could not be secured"}
	}
	marker := managedDLLExtensionMarker{Version: ue4ssExtensionMarkerVersion, ReleaseVersion: plan.Version, Checksum: strings.ToLower(plan.Checksum), SizeBytes: plan.SizeBytes, ExtensionKey: plan.Key, ModKey: plan.ModKey, ConfigRef: configRef, RCONPort: plan.RCONPort, UpdatedAt: time.Now().UTC()}
	body, err := json.Marshal(marker)
	if err != nil || executor.writeRuntimeFile(markerPath, body, 0o600) != nil {
		rollbackManagedDLLExtension(activePath, previousPath, activeExists)
		if previousMarkerFound {
			_ = executor.writeRuntimeFile(markerPath, previousMarker, 0o600)
		} else {
			_ = os.Remove(markerPath)
		}
		return dllExtensionError{code: "dll_extension_activation_failed", message: "extension release state could not be activated"}
	}
	return nil
}

func rollbackManagedDLLExtension(activePath string, previousPath string, activeExists bool) {
	if activeExists {
		_ = copyRegularFileAtomic(previousPath, activePath, maxUE4SSDLLBytes, 0o600)
		return
	}
	_ = os.Remove(activePath)
}

func (executor LifecycleExecutor) ensureLoopbackRCONConfig(resolver WorkspaceResolver, scope string, gameRootKey string, plan protocol.RuntimeDLLExtensionPlan) error {
	configKey := managedRCONConfigRef(gameRootKey, plan.ModKey)
	configPath, _, err := resolver.WritableTarget(scope, configKey)
	if err != nil {
		return dllExtensionError{code: "dll_extension_config_failed", message: "loopback RCON configuration cannot be prepared"}
	}
	if managedLoopbackRCONConfigMatches(configPath, plan.RCONPort) {
		return nil
	}
	password, err := randomRCONPassword()
	if err != nil {
		return dllExtensionError{code: "dll_extension_config_failed", message: "loopback RCON configuration cannot be secured"}
	}
	body := fmt.Sprintf("%s\n[rcon]\nbind_address=127.0.0.1\nport=%d\npassword=%s\n", managedRCONConfigMarker, plan.RCONPort, password)
	if err := executor.writeRuntimeFile(configPath, []byte(body), 0o600); err != nil {
		return dllExtensionError{code: "dll_extension_config_failed", message: "loopback RCON configuration cannot be written"}
	}
	return nil
}

func managedRCONConfigRef(gameRootKey string, modKey string) string {
	return gameRelativeKey(gameRootKey, "ue4ss/Mods/"+modKey+"/config.ini")
}

func managedRCONConfigRefForMod(configRef string, modKey string) bool {
	baseRef := "ue4ss/Mods/" + modKey + "/config.ini"
	return protocol.ValidLogicalFileKey(configRef) && (configRef == baseRef || strings.HasSuffix(configRef, "/"+baseRef))
}

func managedLoopbackRCONConfigMatches(path string, port int) bool {
	body, found, err := readBoundedRegularFile(path, maxUE4SSMetadataBytes)
	if err != nil || !found {
		return false
	}
	content := strings.ReplaceAll(string(body), "\r\n", "\n")
	if !strings.Contains(content, managedRCONConfigMarker) {
		return false
	}
	values := map[string]string{}
	inRCON := false
	for _, rawLine := range strings.Split(content, "\n") {
		line := strings.TrimSpace(rawLine)
		if line == "[rcon]" {
			inRCON = true
			continue
		}
		if strings.HasPrefix(line, "[") {
			inRCON = false
			continue
		}
		if !inRCON || line == "" || strings.HasPrefix(line, ";") || strings.HasPrefix(line, "#") {
			continue
		}
		key, value, ok := strings.Cut(line, "=")
		if !ok {
			continue
		}
		key = strings.TrimSpace(key)
		if key != "bind_address" && key != "port" && key != "password" {
			continue
		}
		if _, duplicate := values[key]; duplicate {
			return false
		}
		values[key] = strings.TrimSpace(value)
	}
	configuredPort, err := strconv.Atoi(values["port"])
	if err != nil || values["bind_address"] != "127.0.0.1" || configuredPort != port || len(values["password"]) != 64 {
		return false
	}
	_, err = hex.DecodeString(values["password"])
	return err == nil
}

func randomRCONPassword() (string, error) {
	bytes := make([]byte, 32)
	if _, err := rand.Read(bytes); err != nil {
		return "", err
	}
	return hex.EncodeToString(bytes), nil
}

func (executor LifecycleExecutor) ensureUE4SSModsIndex(resolver WorkspaceResolver, scope string, gameRootKey string, modKey string) error {
	modsKey := gameRelativeKey(gameRootKey, "ue4ss/Mods/mods.txt")
	modsPath, _, err := resolver.WritableTarget(scope, modsKey)
	if err != nil {
		return dllExtensionError{code: "dll_extension_mods_failed", message: "UE4SS mods index cannot be prepared"}
	}
	body, found, err := readBoundedRegularFile(modsPath, maxUE4SSMetadataBytes)
	if err != nil {
		return dllExtensionError{code: "dll_extension_mods_failed", message: "UE4SS mods index cannot be read"}
	}
	content := ""
	if found {
		content = strings.ReplaceAll(string(body), "\r\n", "\n")
	}
	lines := strings.Split(content, "\n")
	if content == "" {
		lines = nil
	}
	updated := make([]string, 0, len(lines)+1)
	declared := false
	for _, line := range lines {
		if modsIndexLineKey(line) == modKey {
			if !declared {
				updated = append(updated, modKey+" : 1")
				declared = true
			}
			continue
		}
		updated = append(updated, line)
	}
	if !declared {
		updated = append(updated, modKey+" : 1")
	}
	next := strings.Join(updated, "\n")
	if !strings.HasSuffix(next, "\n") {
		next += "\n"
	}
	if content == next {
		return nil
	}
	if err := executor.writeRuntimeFile(modsPath, []byte(next), 0o600); err != nil {
		return dllExtensionError{code: "dll_extension_mods_failed", message: "UE4SS mods index cannot be updated"}
	}
	return nil
}

func modsIndexLineKey(line string) string {
	withoutComment := strings.SplitN(line, "#", 2)[0]
	parts := strings.SplitN(strings.TrimSpace(withoutComment), ":", 2)
	if len(parts) != 2 {
		return ""
	}
	return strings.TrimSpace(parts[0])
}

func existingRuntimeDirectory(scope string, key string) error {
	if !protocol.ValidLogicalFileKey(key) || filepath.IsAbs(key) || strings.Contains(key, `\`) {
		return fmt.Errorf("directory key is unsafe")
	}
	cleanScope, err := filepath.Abs(scope)
	if err != nil {
		return err
	}
	current := cleanScope
	for _, part := range strings.Split(filepath.ToSlash(key), "/") {
		current = filepath.Join(current, part)
		info, err := os.Lstat(current)
		if err != nil || info.Mode()&os.ModeSymlink != 0 || !info.IsDir() {
			return fmt.Errorf("required directory is unavailable")
		}
	}
	return nil
}

func checksumRegularFile(path string, maxBytes int64) (string, int64, error) {
	info, err := os.Lstat(path)
	if err != nil {
		return "", 0, err
	}
	if info.Mode()&os.ModeSymlink != 0 || !info.Mode().IsRegular() {
		return "", 0, fmt.Errorf("file is not regular")
	}
	file, err := os.Open(path)
	if err != nil {
		return "", 0, err
	}
	defer file.Close()
	hash := sha256.New()
	size, err := io.Copy(hash, io.LimitReader(file, maxBytes+1))
	if err != nil {
		return "", 0, err
	}
	if size > maxBytes {
		return "", size, fmt.Errorf("file exceeds maximum size")
	}
	return "sha256:" + hex.EncodeToString(hash.Sum(nil)), size, nil
}

func copyRegularFileAtomic(source string, destination string, maxBytes int64, mode os.FileMode) error {
	checksum, size, err := checksumRegularFile(source, maxBytes)
	if err != nil || checksum == "" || size < 1 {
		if err != nil {
			return err
		}
		return fmt.Errorf("source file is empty")
	}
	input, err := os.Open(source)
	if err != nil {
		return err
	}
	defer input.Close()
	temporary := destination + ".copying"
	output, err := os.OpenFile(temporary, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, mode)
	if err != nil {
		return err
	}
	remove := true
	defer func() {
		_ = output.Close()
		if remove {
			_ = os.Remove(temporary)
		}
	}()
	written, err := io.Copy(output, io.LimitReader(input, maxBytes+1))
	if err != nil || written != size || written > maxBytes {
		if err != nil {
			return err
		}
		return fmt.Errorf("source file changed during copy")
	}
	if err := output.Sync(); err != nil {
		return err
	}
	if err := output.Close(); err != nil {
		return err
	}
	if err := os.Rename(temporary, destination); err != nil {
		return err
	}
	remove = false
	return os.Chmod(destination, mode)
}

func readBoundedRegularFile(path string, maxBytes int64) ([]byte, bool, error) {
	info, err := os.Lstat(path)
	if errors.Is(err, os.ErrNotExist) {
		return nil, false, nil
	}
	if err != nil {
		return nil, false, err
	}
	if info.Mode()&os.ModeSymlink != 0 || !info.Mode().IsRegular() || info.Size() > maxBytes {
		return nil, false, fmt.Errorf("file is not a bounded regular file")
	}
	file, err := os.Open(path)
	if err != nil {
		return nil, false, err
	}
	defer file.Close()
	body, err := io.ReadAll(io.LimitReader(file, maxBytes+1))
	if err != nil || int64(len(body)) > maxBytes {
		if err != nil {
			return nil, false, err
		}
		return nil, false, fmt.Errorf("file exceeds maximum size")
	}
	return body, true, nil
}

func protocolValidSHA256(value string) bool {
	if len(value) != len("sha256:")+64 || !strings.HasPrefix(value, "sha256:") {
		return false
	}
	_, err := hex.DecodeString(strings.TrimPrefix(value, "sha256:"))
	return err == nil
}

func dllExtensionLifecycleFailure(err error) LifecycleExecutionResult {
	if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
		return LifecycleExecutionResult{State: lifecycleResultStateCancelled, Progress: protocol.RunJobProgressReport{Percent: 100, Message: "DLL extension synchronization cancelled"}, Message: "DLL extension synchronization cancelled", ErrorCode: "dll_extension_cancelled"}
	}
	var extensionErr dllExtensionError
	if errors.As(err, &extensionErr) {
		return lifecycleFailure(extensionErr.code, extensionErr.message)
	}
	return lifecycleFailure("dll_extension_sync_failed", "DLL extension synchronization failed")
}