From afbdda117ab20cb0389c9498263e49dd2e6225c6 Mon Sep 17 00:00:00 2001 From: PrimeBuild-pc Date: Tue, 19 May 2026 23:16:20 +0200 Subject: [PATCH 1/2] Add safe CpuSelection affinity apply path --- Services/AffinityApplyService.cs | 341 +++++++++++++++++- Services/IProcessService.cs | 2 + Services/ProcessService.cs | 91 ++++- .../AffinityApplyServiceTests.cs | 232 ++++++++++++ 4 files changed, 654 insertions(+), 12 deletions(-) diff --git a/Services/AffinityApplyService.cs b/Services/AffinityApplyService.cs index 863233b..1ad2e3e 100644 --- a/Services/AffinityApplyService.cs +++ b/Services/AffinityApplyService.cs @@ -16,8 +16,10 @@ */ namespace ThreadPilot.Services { + using System.ComponentModel; using Microsoft.Extensions.Logging; using ThreadPilot.Models; + using ThreadPilot.Platforms.Windows; public enum AffinityApplyFailureReason { @@ -29,27 +31,347 @@ public enum AffinityApplyFailureReason ApplyFailed, } - public sealed record AffinityApplyResult( - bool Success, - long RequestedMask, - long VerifiedMask, - AffinityApplyFailureReason FailureReason, - string Message) + public static class AffinityApplyErrorCodes { + public const string None = "None"; + public const string AccessDenied = "AccessDenied"; + public const string AntiCheatOrProtectedProcessLikely = "AntiCheatOrProtectedProcessLikely"; + public const string ProcessExited = "ProcessExited"; + public const string InvalidSelection = "InvalidSelection"; + public const string InvalidTopology = "InvalidTopology"; + public const string CpuSetsUnavailable = "CpuSetsUnavailable"; + public const string LegacyFallbackUnsafe = "LegacyFallbackUnsafe"; + public const string NativeApplyFailed = "NativeApplyFailed"; + public const string UnknownError = "UnknownError"; + } + + public sealed record AffinityApplyResult + { + public bool Success { get; init; } + + public long RequestedMask { get; init; } + + public long VerifiedMask { get; init; } + + public AffinityApplyFailureReason FailureReason { get; init; } + + public string Message => string.IsNullOrWhiteSpace(this.UserMessage) ? this.TechnicalMessage : this.UserMessage; + + public string ErrorCode { get; init; } = AffinityApplyErrorCodes.None; + + public string UserMessage { get; init; } = string.Empty; + + public string TechnicalMessage { get; init; } = string.Empty; + + public bool IsAccessDenied { get; init; } + + public bool IsAntiCheatLikely { get; init; } + + public bool IsInvalidTopology { get; init; } + + public bool IsLegacyFallbackBlocked { get; init; } + + public bool UsedCpuSets { get; init; } + + public bool UsedLegacyAffinity { get; init; } + public static AffinityApplyResult Succeeded(long requestedMask, long verifiedMask) => - new(true, requestedMask, verifiedMask, AffinityApplyFailureReason.None, "Affinity applied successfully."); + new() + { + Success = true, + RequestedMask = requestedMask, + VerifiedMask = verifiedMask, + FailureReason = AffinityApplyFailureReason.None, + ErrorCode = AffinityApplyErrorCodes.None, + UserMessage = "Affinity applied successfully.", + TechnicalMessage = $"Affinity 0x{requestedMask:X} applied and verified as 0x{verifiedMask:X}.", + }; + + public static AffinityApplyResult SucceededWithCpuSets(string technicalMessage) => + new() + { + Success = true, + FailureReason = AffinityApplyFailureReason.None, + ErrorCode = AffinityApplyErrorCodes.None, + UserMessage = "Affinity applied successfully.", + TechnicalMessage = technicalMessage, + UsedCpuSets = true, + }; + + public static AffinityApplyResult SucceededWithLegacyFallback(long requestedMask, long verifiedMask) => + Succeeded(requestedMask, verifiedMask) with + { + UsedLegacyAffinity = true, + TechnicalMessage = $"CPU Sets failed; legacy affinity 0x{requestedMask:X} applied and verified as 0x{verifiedMask:X}.", + }; public static AffinityApplyResult Failed( long requestedMask, long verifiedMask, AffinityApplyFailureReason failureReason, string message) => - new(false, requestedMask, verifiedMask, failureReason, message); + new() + { + Success = false, + RequestedMask = requestedMask, + VerifiedMask = verifiedMask, + FailureReason = failureReason, + ErrorCode = MapFailureReason(failureReason), + UserMessage = message, + TechnicalMessage = message, + IsAccessDenied = failureReason == AffinityApplyFailureReason.AccessDenied, + }; + + public static AffinityApplyResult Failed( + string errorCode, + string userMessage, + string technicalMessage, + bool isAccessDenied = false, + bool isAntiCheatLikely = false, + bool isInvalidTopology = false, + bool isLegacyFallbackBlocked = false, + long requestedMask = 0, + long verifiedMask = 0, + AffinityApplyFailureReason failureReason = AffinityApplyFailureReason.ApplyFailed) => + new() + { + Success = false, + RequestedMask = requestedMask, + VerifiedMask = verifiedMask, + FailureReason = failureReason, + ErrorCode = errorCode, + UserMessage = userMessage, + TechnicalMessage = technicalMessage, + IsAccessDenied = isAccessDenied, + IsAntiCheatLikely = isAntiCheatLikely, + IsInvalidTopology = isInvalidTopology, + IsLegacyFallbackBlocked = isLegacyFallbackBlocked, + }; + + private static string MapFailureReason(AffinityApplyFailureReason failureReason) => + failureReason switch + { + AffinityApplyFailureReason.None => AffinityApplyErrorCodes.None, + AffinityApplyFailureReason.InvalidMask => AffinityApplyErrorCodes.InvalidSelection, + AffinityApplyFailureReason.ProcessTerminated => AffinityApplyErrorCodes.ProcessExited, + AffinityApplyFailureReason.AccessDenied => AffinityApplyErrorCodes.AccessDenied, + AffinityApplyFailureReason.VerificationMismatch => AffinityApplyErrorCodes.NativeApplyFailed, + AffinityApplyFailureReason.ApplyFailed => AffinityApplyErrorCodes.NativeApplyFailed, + _ => AffinityApplyErrorCodes.UnknownError, + }; } public interface IAffinityApplyService { Task ApplyAsync(ProcessModel process, long requestedMask); + + Task ApplyAsync(ProcessModel process, CpuSelection selection); + } + + internal sealed class CpuSelectionAffinityApplier + { + internal const string AccessDeniedUserMessage = + "Windows denied this change. The process may require administrator rights or may be protected."; + + internal const string AntiCheatUserMessage = + "The process appears protected by anti-cheat or process protection. ThreadPilot will not try to bypass it."; + + internal const string LegacyFallbackBlockedUserMessage = + "This CPU selection cannot be safely represented by legacy affinity APIs on this topology. CPU Sets are required for this selection."; + + internal const string InvalidSelectionUserMessage = + "This CPU selection is empty or does not match the current CPU topology."; + + private readonly Func cpuSetHandlerFactory; + private readonly Func> legacyAffinityApplier; + private readonly ILogger logger; + private readonly Action? cpuSetFailureCallback; + + public CpuSelectionAffinityApplier( + Func cpuSetHandlerFactory, + Func> legacyAffinityApplier, + ILogger logger, + Action? cpuSetFailureCallback = null) + { + this.cpuSetHandlerFactory = cpuSetHandlerFactory ?? throw new ArgumentNullException(nameof(cpuSetHandlerFactory)); + this.legacyAffinityApplier = legacyAffinityApplier ?? throw new ArgumentNullException(nameof(legacyAffinityApplier)); + this.logger = logger ?? throw new ArgumentNullException(nameof(logger)); + this.cpuSetFailureCallback = cpuSetFailureCallback; + } + + public async Task ApplyAsync(ProcessModel process, CpuSelection selection) + { + if (process == null || process.ProcessId <= 0) + { + return ProcessExited("Process is no longer running.", process); + } + + if (selection == null || (selection.CpuSetIds.Count == 0 && selection.LogicalProcessors.Count == 0)) + { + return AffinityApplyResult.Failed( + AffinityApplyErrorCodes.InvalidSelection, + InvalidSelectionUserMessage, + "CpuSelection contains neither CPU Set IDs nor logical processors.", + failureReason: AffinityApplyFailureReason.InvalidMask); + } + + var cpuSetsResult = this.TryApplyCpuSets(process, selection); + if (cpuSetsResult != null) + { + return cpuSetsResult; + } + + this.cpuSetFailureCallback?.Invoke(process); + + var legacyMask = CpuSelection.ToLegacyAffinityMaskOrNull(selection); + if (!legacyMask.HasValue || legacyMask.Value <= 0) + { + return AffinityApplyResult.Failed( + AffinityApplyErrorCodes.LegacyFallbackUnsafe, + LegacyFallbackBlockedUserMessage, + "CpuSelection cannot be represented as a non-zero single-group legacy affinity mask.", + isLegacyFallbackBlocked: true); + } + + try + { + var verifiedMask = await this.legacyAffinityApplier(process, legacyMask.Value).ConfigureAwait(false); + return AffinityApplyResult.SucceededWithLegacyFallback(legacyMask.Value, verifiedMask); + } + catch (Exception ex) when (AffinityApplyExceptionClassifier.IsAccessDenied(ex)) + { + return AccessDenied(ex, legacyMask.Value, process.ProcessorAffinity); + } + catch (Exception ex) when (AffinityApplyExceptionClassifier.IsProcessExited(ex)) + { + return ProcessExited("Process exited before legacy affinity fallback could be applied.", process, legacyMask.Value); + } + catch (Exception ex) + { + this.logger.LogWarning( + ex, + "Legacy affinity fallback failed for process {ProcessName} (PID: {ProcessId})", + process.Name, + process.ProcessId); + + return AffinityApplyResult.Failed( + AffinityApplyErrorCodes.NativeApplyFailed, + "ThreadPilot could not apply this CPU selection.", + ex.Message, + requestedMask: legacyMask.Value, + verifiedMask: process.ProcessorAffinity); + } + } + + private AffinityApplyResult? TryApplyCpuSets(ProcessModel process, CpuSelection selection) + { + try + { + var handler = this.cpuSetHandlerFactory(process); + if (!handler.IsValid) + { + this.logger.LogDebug( + "CPU Set handler is invalid for process {ProcessName} (PID: {ProcessId})", + process.Name, + process.ProcessId); + return null; + } + + if (handler.ApplyCpuSelection(selection)) + { + return AffinityApplyResult.SucceededWithCpuSets( + $"CPU Sets applied to process {process.Name} (PID: {process.ProcessId})."); + } + + return null; + } + catch (Exception ex) when (AffinityApplyExceptionClassifier.IsAccessDenied(ex)) + { + return AccessDenied(ex, 0, process.ProcessorAffinity); + } + catch (Exception ex) when (AffinityApplyExceptionClassifier.IsProcessExited(ex)) + { + return ProcessExited("Process exited before CPU Sets could be applied.", process); + } + catch (Exception ex) + { + this.logger.LogDebug( + ex, + "CPU Sets failed for process {ProcessName} (PID: {ProcessId}); evaluating legacy fallback", + process.Name, + process.ProcessId); + return null; + } + } + + private static AffinityApplyResult AccessDenied(Exception ex, long requestedMask, long verifiedMask) + { + var antiCheatLikely = AffinityApplyExceptionClassifier.IsAntiCheatLikely(ex); + return AffinityApplyResult.Failed( + antiCheatLikely + ? AffinityApplyErrorCodes.AntiCheatOrProtectedProcessLikely + : AffinityApplyErrorCodes.AccessDenied, + antiCheatLikely ? AntiCheatUserMessage : AccessDeniedUserMessage, + ex.Message, + isAccessDenied: true, + isAntiCheatLikely: antiCheatLikely, + requestedMask: requestedMask, + verifiedMask: verifiedMask, + failureReason: AffinityApplyFailureReason.AccessDenied); + } + + private static AffinityApplyResult ProcessExited(string userMessage, ProcessModel? process, long requestedMask = 0) => + AffinityApplyResult.Failed( + AffinityApplyErrorCodes.ProcessExited, + userMessage, + userMessage, + requestedMask: requestedMask, + verifiedMask: process?.ProcessorAffinity ?? 0, + failureReason: AffinityApplyFailureReason.ProcessTerminated); + } + + internal static class AffinityApplyExceptionClassifier + { + public static bool IsAccessDenied(Exception ex) => + ex is UnauthorizedAccessException || + ex is Win32Exception { NativeErrorCode: 5 } || + IsInnerAccessDenied(ex.InnerException) || + ContainsAny( + ex.Message, + "access denied", + "anti-cheat", + "anti cheat", + "protected", + "insufficient privileges"); + + public static bool IsAntiCheatLikely(Exception ex) => + ContainsAny(ex.Message, "anti-cheat", "anti cheat", "protected") || + (ex.InnerException != null && IsAntiCheatLikely(ex.InnerException)); + + public static bool IsProcessExited(Exception ex) + { + if (ex is ArgumentException) + { + return true; + } + + var message = ex.Message ?? string.Empty; + if (ex is InvalidOperationException && + ContainsAny(message, "exit", "exited", "terminated", "not running", "has no process associated")) + { + return true; + } + + return ex.InnerException != null && IsProcessExited(ex.InnerException); + } + + private static bool IsInnerAccessDenied(Exception? ex) => ex != null && IsAccessDenied(ex); + + private static bool ContainsAny(string? value, params string[] needles) + { + var source = value ?? string.Empty; + return needles.Any(needle => source.Contains(needle, StringComparison.OrdinalIgnoreCase)); + } } public sealed class AffinityApplyService : IAffinityApplyService @@ -68,6 +390,9 @@ public AffinityApplyService( this.logger = logger ?? throw new ArgumentNullException(nameof(logger)); } + public Task ApplyAsync(ProcessModel process, CpuSelection selection) => + this.processService.SetProcessorAffinity(process, selection); + public async Task ApplyAsync(ProcessModel process, long requestedMask) { ArgumentNullException.ThrowIfNull(process); diff --git a/Services/IProcessService.cs b/Services/IProcessService.cs index 22b6f26..78baf5e 100644 --- a/Services/IProcessService.cs +++ b/Services/IProcessService.cs @@ -28,6 +28,8 @@ public interface IProcessService Task SetProcessorAffinity(ProcessModel process, long affinityMask); + Task SetProcessorAffinity(ProcessModel process, CpuSelection selection); + Task SetProcessPriority(ProcessModel process, ProcessPriorityClass priority); Task SaveProcessProfile(string profileName, ProcessModel process); diff --git a/Services/ProcessService.cs b/Services/ProcessService.cs index 6b8304b..0a0f298 100644 --- a/Services/ProcessService.cs +++ b/Services/ProcessService.cs @@ -42,6 +42,7 @@ public class ProcessService : IProcessService private readonly IProcessClassifier processClassifier; private readonly IPassiveProcessErrorThrottle passiveProcessErrorThrottle; private readonly Func profilesDirectoryProvider; + private readonly CpuSelectionAffinityApplier cpuSelectionAffinityApplier; private string ProfilesDirectory => this.profilesDirectoryProvider(); @@ -65,6 +66,11 @@ public ProcessService( this.processClassifier = processClassifier ?? new ProcessClassifier(new ProcessFilterService()); this.passiveProcessErrorThrottle = passiveProcessErrorThrottle ?? new PassiveProcessErrorThrottle(); this.profilesDirectoryProvider = profilesDirectoryProvider ?? (() => StoragePaths.ProfilesDirectory); + this.cpuSelectionAffinityApplier = new CpuSelectionAffinityApplier( + this.GetOrCreateCpuSetHandler, + this.ApplyLegacyProcessorAffinityDirectAsync, + logger ?? (ILogger)Microsoft.Extensions.Logging.Abstractions.NullLogger.Instance, + process => this.cpuSetHandlers.TryRemove(process.ProcessId, out _)); StoragePaths.EnsureAppDataDirectories(); this.MigrateLegacyProfilesIfNeeded(); @@ -423,6 +429,52 @@ await Task.Run(() => }).ConfigureAwait(false); } + public async Task SetProcessorAffinity(ProcessModel process, CpuSelection selection) + { + if (process == null) + { + return AffinityApplyResult.Failed( + AffinityApplyErrorCodes.ProcessExited, + "Process is no longer running.", + "ProcessModel is null.", + failureReason: AffinityApplyFailureReason.ProcessTerminated); + } + + try + { + this.EnsureProcessOperationAllowed(process, "SetProcessAffinity"); + } + catch (Exception ex) when (AffinityApplyExceptionClassifier.IsAccessDenied(ex)) + { + this.AuditProcessOperation("SetProcessAffinity", process?.Name ?? string.Empty, success: false); + var antiCheatLikely = AffinityApplyExceptionClassifier.IsAntiCheatLikely(ex); + return AffinityApplyResult.Failed( + antiCheatLikely + ? AffinityApplyErrorCodes.AntiCheatOrProtectedProcessLikely + : AffinityApplyErrorCodes.AccessDenied, + antiCheatLikely + ? CpuSelectionAffinityApplier.AntiCheatUserMessage + : CpuSelectionAffinityApplier.AccessDeniedUserMessage, + ex.Message, + isAccessDenied: true, + isAntiCheatLikely: antiCheatLikely, + verifiedMask: process?.ProcessorAffinity ?? 0, + failureReason: AffinityApplyFailureReason.AccessDenied); + } + catch (Exception ex) when (AffinityApplyExceptionClassifier.IsProcessExited(ex)) + { + this.AuditProcessOperation("SetProcessAffinity", process?.Name ?? string.Empty, success: false); + return AffinityApplyResult.Failed( + AffinityApplyErrorCodes.ProcessExited, + "Process is no longer running.", + ex.Message, + verifiedMask: process?.ProcessorAffinity ?? 0, + failureReason: AffinityApplyFailureReason.ProcessTerminated); + } + + return await this.cpuSelectionAffinityApplier.ApplyAsync(process, selection).ConfigureAwait(false); + } + /// /// Attempts to set process affinity using CPU Sets (Windows 10+ feature). /// @@ -431,10 +483,7 @@ private bool TrySetAffinityViaCpuSets(ProcessModel process, long affinityMask) try { // Get or create CPU Set handler for this process - var handler = this.cpuSetHandlers.GetOrAdd(process.ProcessId, pid => - { - return new ProcessCpuSetHandler((uint)pid, process.Name, this.logger); - }); + var handler = this.GetOrCreateCpuSetHandler(process); // Check if handler is valid if (!handler.IsValid) @@ -470,6 +519,40 @@ private bool TrySetAffinityViaCpuSets(ProcessModel process, long affinityMask) } } + private IProcessCpuSetHandler GetOrCreateCpuSetHandler(ProcessModel process) => + this.cpuSetHandlers.GetOrAdd(process.ProcessId, pid => + { + return new ProcessCpuSetHandler((uint)pid, process.Name, this.logger); + }); + + private async Task ApplyLegacyProcessorAffinityDirectAsync(ProcessModel process, long affinityMask) + { + return await Task.Run(() => + { + try + { + using var targetProcess = Process.GetProcessById(process.ProcessId); + targetProcess.ProcessorAffinity = new IntPtr(affinityMask); + var verifiedMask = (long)targetProcess.ProcessorAffinity; + process.ProcessorAffinity = verifiedMask; + + this.AuditProcessOperation("SetProcessAffinity", process.Name, success: true); + this.logger?.LogInformation( + "Successfully applied classic ProcessorAffinity 0x{AffinityMask:X} to process {ProcessName} (PID: {ProcessId})", + affinityMask, + process.Name, + process.ProcessId); + + return verifiedMask; + } + catch + { + this.AuditProcessOperation("SetProcessAffinity", process.Name, success: false); + throw; + } + }).ConfigureAwait(false); + } + public async Task SetProcessPriority(ProcessModel process, ProcessPriorityClass priority) { this.EnsureProcessOperationAllowed(process, "SetProcessPriority"); diff --git a/Tests/ThreadPilot.Core.Tests/AffinityApplyServiceTests.cs b/Tests/ThreadPilot.Core.Tests/AffinityApplyServiceTests.cs index af592a6..10a0893 100644 --- a/Tests/ThreadPilot.Core.Tests/AffinityApplyServiceTests.cs +++ b/Tests/ThreadPilot.Core.Tests/AffinityApplyServiceTests.cs @@ -1,8 +1,10 @@ namespace ThreadPilot.Core.Tests { + using System.ComponentModel; using Microsoft.Extensions.Logging.Abstractions; using Moq; using ThreadPilot.Models; + using ThreadPilot.Platforms.Windows; using ThreadPilot.Services; public sealed class AffinityApplyServiceTests @@ -181,6 +183,160 @@ public async Task ApplyAsync_WhenApplyThrowsUnexpectedError_ReturnsApplyFailed() Assert.Equal(3, result.VerifiedMask); } + [Fact] + public async Task CpuSelectionApply_WhenCpuSetsFailAndSelectionIsSingleGroupBelow64_UsesLegacyFallback() + { + var process = new ProcessModel { ProcessId = 42, Name = "Game" }; + var cpuSets = new FakeCpuSetHandler { ApplyCpuSelectionResult = false }; + var legacy = new RecordingLegacyAffinityApplier(); + var service = CreateCpuSelectionApplier(cpuSets, legacy); + var selection = CreateSelection( + new ProcessorRef(0, 0, 0), + new ProcessorRef(0, 2, 2)); + + var result = await service.ApplyAsync(process, selection); + + Assert.True(result.Success); + Assert.Equal(AffinityApplyErrorCodes.None, result.ErrorCode); + Assert.False(result.UsedCpuSets); + Assert.True(result.UsedLegacyAffinity); + Assert.Equal(0x05, legacy.LastMask); + Assert.Equal(1, legacy.CallCount); + Assert.Equal(1, cpuSets.ApplyCpuSelectionCalls); + } + + [Fact] + public async Task CpuSelectionApply_WhenCpuSetsFailAndSelectionHasMultipleGroups_BlocksLegacyFallback() + { + var process = new ProcessModel { ProcessId = 42, Name = "Game" }; + var cpuSets = new FakeCpuSetHandler { ApplyCpuSelectionResult = false }; + var legacy = new RecordingLegacyAffinityApplier(); + var service = CreateCpuSelectionApplier(cpuSets, legacy); + var selection = CreateSelection( + new ProcessorRef(0, 0, 0), + new ProcessorRef(1, 0, 1)); + + var result = await service.ApplyAsync(process, selection); + + Assert.False(result.Success); + Assert.Equal(AffinityApplyErrorCodes.LegacyFallbackUnsafe, result.ErrorCode); + Assert.True(result.IsLegacyFallbackBlocked); + Assert.False(result.UsedLegacyAffinity); + Assert.Equal(0, legacy.CallCount); + } + + [Fact] + public async Task CpuSelectionApply_WhenCpuSetsFailAndSelectionContainsCpu64_BlocksLegacyFallback() + { + var process = new ProcessModel { ProcessId = 42, Name = "Game" }; + var cpuSets = new FakeCpuSetHandler { ApplyCpuSelectionResult = false }; + var legacy = new RecordingLegacyAffinityApplier(); + var service = CreateCpuSelectionApplier(cpuSets, legacy); + var selection = CreateSelection(new ProcessorRef(0, 64, 64)); + + var result = await service.ApplyAsync(process, selection); + + Assert.False(result.Success); + Assert.Equal(AffinityApplyErrorCodes.LegacyFallbackUnsafe, result.ErrorCode); + Assert.True(result.IsLegacyFallbackBlocked); + Assert.Equal(0, legacy.CallCount); + } + + [Fact] + public async Task CpuSelectionApply_WhenSelectionHasExplicitCpuSetIds_TriesCpuSets() + { + var process = new ProcessModel { ProcessId = 42, Name = "Game" }; + var cpuSets = new FakeCpuSetHandler { ApplyCpuSelectionResult = true }; + var legacy = new RecordingLegacyAffinityApplier(); + var service = CreateCpuSelectionApplier(cpuSets, legacy); + var selection = new CpuSelection { CpuSetIds = [101, 103] }; + + var result = await service.ApplyAsync(process, selection); + + Assert.True(result.Success); + Assert.True(result.UsedCpuSets); + Assert.False(result.UsedLegacyAffinity); + Assert.Same(selection, cpuSets.LastSelection); + Assert.Equal(0, legacy.CallCount); + } + + [Fact] + public async Task CpuSelectionApply_WhenSelectionIsEmpty_ReturnsInvalidSelection() + { + var process = new ProcessModel { ProcessId = 42, Name = "Game" }; + var cpuSets = new FakeCpuSetHandler(); + var legacy = new RecordingLegacyAffinityApplier(); + var service = CreateCpuSelectionApplier(cpuSets, legacy); + + var result = await service.ApplyAsync(process, new CpuSelection()); + + Assert.False(result.Success); + Assert.Equal(AffinityApplyErrorCodes.InvalidSelection, result.ErrorCode); + Assert.False(result.UsedCpuSets); + Assert.False(result.UsedLegacyAffinity); + Assert.Equal(0, cpuSets.ApplyCpuSelectionCalls); + Assert.Equal(0, legacy.CallCount); + } + + [Fact] + public async Task CpuSelectionApply_WhenCpuSetsThrowAccessDenied_ReturnsAccessDenied() + { + var process = new ProcessModel { ProcessId = 42, Name = "Game" }; + var cpuSets = new FakeCpuSetHandler + { + ApplyCpuSelectionException = new Win32Exception(5, "Access is denied."), + }; + var legacy = new RecordingLegacyAffinityApplier(); + var service = CreateCpuSelectionApplier(cpuSets, legacy); + var selection = CreateSelection(new ProcessorRef(0, 0, 0)); + + var result = await service.ApplyAsync(process, selection); + + Assert.False(result.Success); + Assert.Equal(AffinityApplyErrorCodes.AccessDenied, result.ErrorCode); + Assert.True(result.IsAccessDenied); + Assert.Equal(0, legacy.CallCount); + } + + [Fact] + public async Task CpuSelectionApply_WhenFallbackThrowsAccessDenied_ReturnsAccessDenied() + { + var process = new ProcessModel { ProcessId = 42, Name = "Game" }; + var cpuSets = new FakeCpuSetHandler { ApplyCpuSelectionResult = false }; + var legacy = new RecordingLegacyAffinityApplier + { + ExceptionToThrow = new UnauthorizedAccessException("Access denied."), + }; + var service = CreateCpuSelectionApplier(cpuSets, legacy); + var selection = CreateSelection(new ProcessorRef(0, 0, 0)); + + var result = await service.ApplyAsync(process, selection); + + Assert.False(result.Success); + Assert.Equal(AffinityApplyErrorCodes.AccessDenied, result.ErrorCode); + Assert.True(result.IsAccessDenied); + Assert.Equal(1, legacy.CallCount); + } + + [Fact] + public async Task CpuSelectionApply_WhenFallbackThrowsProcessExited_ReturnsProcessExited() + { + var process = new ProcessModel { ProcessId = 42, Name = "Game" }; + var cpuSets = new FakeCpuSetHandler { ApplyCpuSelectionResult = false }; + var legacy = new RecordingLegacyAffinityApplier + { + ExceptionToThrow = new ArgumentException("Process has exited."), + }; + var service = CreateCpuSelectionApplier(cpuSets, legacy); + var selection = CreateSelection(new ProcessorRef(0, 0, 0)); + + var result = await service.ApplyAsync(process, selection); + + Assert.False(result.Success); + Assert.Equal(AffinityApplyErrorCodes.ProcessExited, result.ErrorCode); + Assert.False(result.UsedLegacyAffinity); + } + private static AffinityApplyService CreateService(Mock processService) { var topologyService = new Mock(MockBehavior.Loose); @@ -199,6 +355,14 @@ private static AffinityApplyService CreateService( NullLogger.Instance); } + private static CpuSelectionAffinityApplier CreateCpuSelectionApplier( + FakeCpuSetHandler cpuSets, + RecordingLegacyAffinityApplier legacy) => + new( + _ => cpuSets, + legacy.ApplyAsync, + NullLogger.Instance); + private static Mock CreateProcessService(bool processStillRunning) { var processService = new Mock(MockBehavior.Strict); @@ -207,5 +371,73 @@ private static Mock CreateProcessService(bool processStillRunni .ReturnsAsync(processStillRunning); return processService; } + + private static CpuSelection CreateSelection(params ProcessorRef[] processors) => + new() + { + LogicalProcessors = processors.ToList(), + GlobalLogicalProcessorIndexes = processors.Select(processor => processor.GlobalIndex).ToList(), + }; + + private sealed class RecordingLegacyAffinityApplier + { + public int CallCount { get; private set; } + + public long? LastMask { get; private set; } + + public Exception? ExceptionToThrow { get; init; } + + public Task ApplyAsync(ProcessModel process, long affinityMask) + { + this.CallCount++; + this.LastMask = affinityMask; + + if (this.ExceptionToThrow != null) + { + throw this.ExceptionToThrow; + } + + process.ProcessorAffinity = affinityMask; + return Task.FromResult(affinityMask); + } + } + + private sealed class FakeCpuSetHandler : IProcessCpuSetHandler + { + public uint ProcessId => 42; + + public string ExecutableName => "Game"; + + public bool IsValid { get; init; } = true; + + public bool ApplyCpuSelectionResult { get; init; } + + public Exception? ApplyCpuSelectionException { get; init; } + + public int ApplyCpuSelectionCalls { get; private set; } + + public CpuSelection? LastSelection { get; private set; } + + public bool ApplyCpuSetMask(long affinityMask, bool clearMask = false) => false; + + public bool ApplyCpuSelection(CpuSelection? selection, bool clearSelection = false) + { + this.ApplyCpuSelectionCalls++; + this.LastSelection = selection; + + if (this.ApplyCpuSelectionException != null) + { + throw this.ApplyCpuSelectionException; + } + + return this.ApplyCpuSelectionResult; + } + + public double GetAverageCpuUsage() => 0; + + public void Dispose() + { + } + } } } From cc6b8b6cf0bc4ee99cd1dc2406608817c572a6a4 Mon Sep 17 00:00:00 2001 From: PrimeBuild-pc Date: Tue, 19 May 2026 23:29:33 +0200 Subject: [PATCH 2/2] Fix CpuSelection apply audit coverage --- Services/AffinityApplyService.cs | 30 ++++- Services/ProcessService.cs | 3 +- .../AffinityApplyServiceTests.cs | 117 +++++++++++++++++- .../ProcessServiceSecurityTests.cs | 21 ++++ 4 files changed, 166 insertions(+), 5 deletions(-) diff --git a/Services/AffinityApplyService.cs b/Services/AffinityApplyService.cs index 1ad2e3e..a01ea01 100644 --- a/Services/AffinityApplyService.cs +++ b/Services/AffinityApplyService.cs @@ -186,17 +186,20 @@ internal sealed class CpuSelectionAffinityApplier private readonly Func> legacyAffinityApplier; private readonly ILogger logger; private readonly Action? cpuSetFailureCallback; + private readonly Action? auditCallback; public CpuSelectionAffinityApplier( Func cpuSetHandlerFactory, Func> legacyAffinityApplier, ILogger logger, - Action? cpuSetFailureCallback = null) + Action? cpuSetFailureCallback = null, + Action? auditCallback = null) { this.cpuSetHandlerFactory = cpuSetHandlerFactory ?? throw new ArgumentNullException(nameof(cpuSetHandlerFactory)); this.legacyAffinityApplier = legacyAffinityApplier ?? throw new ArgumentNullException(nameof(legacyAffinityApplier)); this.logger = logger ?? throw new ArgumentNullException(nameof(logger)); this.cpuSetFailureCallback = cpuSetFailureCallback; + this.auditCallback = auditCallback; } public async Task ApplyAsync(ProcessModel process, CpuSelection selection) @@ -208,6 +211,7 @@ public async Task ApplyAsync(ProcessModel process, CpuSelec if (selection == null || (selection.CpuSetIds.Count == 0 && selection.LogicalProcessors.Count == 0)) { + this.Audit(process, success: false); return AffinityApplyResult.Failed( AffinityApplyErrorCodes.InvalidSelection, InvalidSelectionUserMessage, @@ -226,6 +230,7 @@ public async Task ApplyAsync(ProcessModel process, CpuSelec var legacyMask = CpuSelection.ToLegacyAffinityMaskOrNull(selection); if (!legacyMask.HasValue || legacyMask.Value <= 0) { + this.Audit(process, success: false); return AffinityApplyResult.Failed( AffinityApplyErrorCodes.LegacyFallbackUnsafe, LegacyFallbackBlockedUserMessage, @@ -279,18 +284,24 @@ public async Task ApplyAsync(ProcessModel process, CpuSelec if (handler.ApplyCpuSelection(selection)) { + this.Audit(process, success: true); return AffinityApplyResult.SucceededWithCpuSets( $"CPU Sets applied to process {process.Name} (PID: {process.ProcessId})."); } + // ProcessCpuSetHandler.ApplyCpuSelection currently returns only bool. A false + // can mean CPU Sets unavailable, access denied without detailed error, or + // a native apply failure. Later UX/error-model work will classify this more precisely. return null; } catch (Exception ex) when (AffinityApplyExceptionClassifier.IsAccessDenied(ex)) { + this.Audit(process, success: false); return AccessDenied(ex, 0, process.ProcessorAffinity); } catch (Exception ex) when (AffinityApplyExceptionClassifier.IsProcessExited(ex)) { + this.Audit(process, success: false); return ProcessExited("Process exited before CPU Sets could be applied.", process); } catch (Exception ex) @@ -304,6 +315,9 @@ public async Task ApplyAsync(ProcessModel process, CpuSelec } } + private void Audit(ProcessModel process, bool success) => + this.auditCallback?.Invoke(process, success); + private static AffinityApplyResult AccessDenied(Exception ex, long requestedMask, long verifiedMask) { var antiCheatLikely = AffinityApplyExceptionClassifier.IsAntiCheatLikely(ex); @@ -391,7 +405,19 @@ public AffinityApplyService( } public Task ApplyAsync(ProcessModel process, CpuSelection selection) => - this.processService.SetProcessorAffinity(process, selection); + process == null + ? Task.FromResult(AffinityApplyResult.Failed( + AffinityApplyErrorCodes.ProcessExited, + "Process is no longer running.", + "ProcessModel is null.", + failureReason: AffinityApplyFailureReason.ProcessTerminated)) + : selection == null + ? Task.FromResult(AffinityApplyResult.Failed( + AffinityApplyErrorCodes.InvalidSelection, + CpuSelectionAffinityApplier.InvalidSelectionUserMessage, + "CpuSelection is null.", + failureReason: AffinityApplyFailureReason.InvalidMask)) + : this.processService.SetProcessorAffinity(process, selection); public async Task ApplyAsync(ProcessModel process, long requestedMask) { diff --git a/Services/ProcessService.cs b/Services/ProcessService.cs index 0a0f298..0f69d70 100644 --- a/Services/ProcessService.cs +++ b/Services/ProcessService.cs @@ -70,7 +70,8 @@ public ProcessService( this.GetOrCreateCpuSetHandler, this.ApplyLegacyProcessorAffinityDirectAsync, logger ?? (ILogger)Microsoft.Extensions.Logging.Abstractions.NullLogger.Instance, - process => this.cpuSetHandlers.TryRemove(process.ProcessId, out _)); + process => this.cpuSetHandlers.TryRemove(process.ProcessId, out _), + (process, success) => this.AuditProcessOperation("SetProcessAffinity", process.Name, success)); StoragePaths.EnsureAppDataDirectories(); this.MigrateLegacyProfilesIfNeeded(); diff --git a/Tests/ThreadPilot.Core.Tests/AffinityApplyServiceTests.cs b/Tests/ThreadPilot.Core.Tests/AffinityApplyServiceTests.cs index 10a0893..32c6893 100644 --- a/Tests/ThreadPilot.Core.Tests/AffinityApplyServiceTests.cs +++ b/Tests/ThreadPilot.Core.Tests/AffinityApplyServiceTests.cs @@ -260,6 +260,25 @@ public async Task CpuSelectionApply_WhenSelectionHasExplicitCpuSetIds_TriesCpuSe Assert.Equal(0, legacy.CallCount); } + [Fact] + public async Task CpuSelectionApply_WhenCpuSetsSucceed_AuditsSuccessAndSkipsLegacyFallback() + { + var process = new ProcessModel { ProcessId = 42, Name = "Game" }; + var cpuSets = new FakeCpuSetHandler { ApplyCpuSelectionResult = true }; + var legacy = new RecordingLegacyAffinityApplier(); + var audit = new RecordingAffinityAudit(); + var service = CreateCpuSelectionApplier(cpuSets, legacy, audit); + var selection = CreateSelection(new ProcessorRef(0, 0, 0)); + + var result = await service.ApplyAsync(process, selection); + + Assert.True(result.Success); + Assert.True(result.UsedCpuSets); + Assert.False(result.UsedLegacyAffinity); + Assert.Equal(0, legacy.CallCount); + Assert.Equal([(process, true)], audit.Calls); + } + [Fact] public async Task CpuSelectionApply_WhenSelectionIsEmpty_ReturnsInvalidSelection() { @@ -278,6 +297,58 @@ public async Task CpuSelectionApply_WhenSelectionIsEmpty_ReturnsInvalidSelection Assert.Equal(0, legacy.CallCount); } + [Fact] + public async Task CpuSelectionApply_WhenSelectionIsEmpty_AuditsFailure() + { + var process = new ProcessModel { ProcessId = 42, Name = "Game" }; + var cpuSets = new FakeCpuSetHandler(); + var legacy = new RecordingLegacyAffinityApplier(); + var audit = new RecordingAffinityAudit(); + var service = CreateCpuSelectionApplier(cpuSets, legacy, audit); + + var result = await service.ApplyAsync(process, new CpuSelection()); + + Assert.False(result.Success); + Assert.Equal(AffinityApplyErrorCodes.InvalidSelection, result.ErrorCode); + Assert.Equal([(process, false)], audit.Calls); + } + + [Fact] + public async Task CpuSelectionApply_WhenLegacyFallbackIsUnsafe_AuditsFailure() + { + var process = new ProcessModel { ProcessId = 42, Name = "Game" }; + var cpuSets = new FakeCpuSetHandler { ApplyCpuSelectionResult = false }; + var legacy = new RecordingLegacyAffinityApplier(); + var audit = new RecordingAffinityAudit(); + var service = CreateCpuSelectionApplier(cpuSets, legacy, audit); + var selection = CreateSelection(new ProcessorRef(1, 64, 64)); + + var result = await service.ApplyAsync(process, selection); + + Assert.False(result.Success); + Assert.True(result.IsLegacyFallbackBlocked); + Assert.Equal(0, legacy.CallCount); + Assert.Equal([(process, false)], audit.Calls); + } + + [Fact] + public async Task CpuSelectionApply_WhenLegacyFallbackSucceeds_DoesNotAuditTwice() + { + var process = new ProcessModel { ProcessId = 42, Name = "Game" }; + var cpuSets = new FakeCpuSetHandler { ApplyCpuSelectionResult = false }; + var legacy = new RecordingLegacyAffinityApplier(); + var audit = new RecordingAffinityAudit(); + var service = CreateCpuSelectionApplier(cpuSets, legacy, audit); + var selection = CreateSelection(new ProcessorRef(0, 0, 0)); + + var result = await service.ApplyAsync(process, selection); + + Assert.True(result.Success); + Assert.True(result.UsedLegacyAffinity); + Assert.Equal(1, legacy.CallCount); + Assert.Empty(audit.Calls); + } + [Fact] public async Task CpuSelectionApply_WhenCpuSetsThrowAccessDenied_ReturnsAccessDenied() { @@ -337,6 +408,37 @@ public async Task CpuSelectionApply_WhenFallbackThrowsProcessExited_ReturnsProce Assert.False(result.UsedLegacyAffinity); } + [Fact] + public async Task ApplyCpuSelectionAsync_WhenProcessIsNull_ReturnsProcessExitedWithoutDelegating() + { + var processService = new Mock(MockBehavior.Strict); + var service = CreateService(processService); + + var result = await service.ApplyAsync(null!, CreateSelection(new ProcessorRef(0, 0, 0))); + + Assert.False(result.Success); + Assert.Equal(AffinityApplyErrorCodes.ProcessExited, result.ErrorCode); + processService.Verify( + service => service.SetProcessorAffinity(It.IsAny(), It.IsAny()), + Times.Never); + } + + [Fact] + public async Task ApplyCpuSelectionAsync_WhenSelectionIsNull_ReturnsInvalidSelectionWithoutDelegating() + { + var process = new ProcessModel { ProcessId = 42, Name = "Game" }; + var processService = new Mock(MockBehavior.Strict); + var service = CreateService(processService); + + var result = await service.ApplyAsync(process, null!); + + Assert.False(result.Success); + Assert.Equal(AffinityApplyErrorCodes.InvalidSelection, result.ErrorCode); + processService.Verify( + service => service.SetProcessorAffinity(It.IsAny(), It.IsAny()), + Times.Never); + } + private static AffinityApplyService CreateService(Mock processService) { var topologyService = new Mock(MockBehavior.Loose); @@ -357,11 +459,14 @@ private static AffinityApplyService CreateService( private static CpuSelectionAffinityApplier CreateCpuSelectionApplier( FakeCpuSetHandler cpuSets, - RecordingLegacyAffinityApplier legacy) => + RecordingLegacyAffinityApplier legacy, + RecordingAffinityAudit? audit = null) => new( _ => cpuSets, legacy.ApplyAsync, - NullLogger.Instance); + NullLogger.Instance, + null, + audit is null ? null : audit.Record); private static Mock CreateProcessService(bool processStillRunning) { @@ -402,6 +507,14 @@ public Task ApplyAsync(ProcessModel process, long affinityMask) } } + private sealed class RecordingAffinityAudit + { + public List<(ProcessModel Process, bool Success)> Calls { get; } = new(); + + public void Record(ProcessModel process, bool success) => + this.Calls.Add((process, success)); + } + private sealed class FakeCpuSetHandler : IProcessCpuSetHandler { public uint ProcessId => 42; diff --git a/Tests/ThreadPilot.Core.Tests/ProcessServiceSecurityTests.cs b/Tests/ThreadPilot.Core.Tests/ProcessServiceSecurityTests.cs index 3a95c5d..7416210 100644 --- a/Tests/ThreadPilot.Core.Tests/ProcessServiceSecurityTests.cs +++ b/Tests/ThreadPilot.Core.Tests/ProcessServiceSecurityTests.cs @@ -57,5 +57,26 @@ await Assert.ThrowsAsync(async () => security.VerifyAll(); } + + [Fact] + public async Task SetProcessorAffinity_WithInvalidCpuSelection_AuditsFailure() + { + var security = new Mock(MockBehavior.Strict); + security + .Setup(s => s.ValidateProcessOperation("Game", "SetProcessAffinity")) + .Returns(true); + security + .Setup(s => s.AuditElevatedAction("SetProcessAffinity", "Game", false)) + .Returns(Task.CompletedTask); + + var service = new ProcessService(null, security.Object); + var process = new ProcessModel { Name = "Game", ProcessId = int.MaxValue }; + + var result = await service.SetProcessorAffinity(process, new CpuSelection()); + + Assert.False(result.Success); + Assert.Equal(AffinityApplyErrorCodes.InvalidSelection, result.ErrorCode); + security.VerifyAll(); + } } }