Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions Platforms/Windows/IProcessMemoryPriorityNativeApi.cs
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,8 @@ int QueryIoPriority(SafeProcessHandle process, ref int priority, out uint return
}

int SetIoPriority(SafeProcessHandle process, ref int priority) => unchecked((int)0xC0000002);

IDisposable? AcquirePrivilegeScope(string privilegeName) => null;
}

public sealed class ProcessMemoryPriorityNativeApi : IProcessMemoryPriorityNativeApi
Expand Down Expand Up @@ -96,6 +98,9 @@ public int SetIoPriority(SafeProcessHandle process, ref int priority) =>
ProcessMemoryPriorityNativeMethods.ProcessIoPriority,
ref priority,
sizeof(int));

public IDisposable? AcquirePrivilegeScope(string privilegeName) =>
ProcessMemoryPriorityNativeMethods.AcquirePrivilegeScope(privilegeName);
}

public enum ProcessInformationClass
Expand Down
289 changes: 289 additions & 0 deletions Platforms/Windows/ProcessMemoryPriorityNativeMethods.cs
Original file line number Diff line number Diff line change
Expand Up @@ -3,12 +3,42 @@
*/
namespace ThreadPilot.Platforms.Windows
{
using System;
using System.Runtime.InteropServices;
using Microsoft.Win32.SafeHandles;

internal static partial class ProcessMemoryPriorityNativeMethods
{
internal const int ProcessIoPriority = 33;
internal const string SeIncreaseBasePriorityPrivilege = "SeIncreaseBasePriorityPrivilege";
private const uint TokenAdjustPrivileges = 0x0020;
private const uint TokenQuery = 0x0008;
private const uint SePrivilegeEnabled = 0x00000002;
private const uint TokenPrivilegesInformationClass = 3;
private static readonly object PrivilegeLock = new();
private static readonly System.Collections.Generic.Dictionary<string, PrivilegeState> ActivePrivileges = new(StringComparer.OrdinalIgnoreCase);

[StructLayout(LayoutKind.Sequential)]
internal struct Luid
{
public uint LowPart;
public int HighPart;
}

[StructLayout(LayoutKind.Sequential)]
internal struct LuidAndAttributes
{
public Luid Luid;
public uint Attributes;
}

[StructLayout(LayoutKind.Sequential)]
internal struct TokenPrivileges
{
public uint PrivilegeCount;
public Luid Luid;
public uint Attributes;
}

[LibraryImport("kernel32.dll", SetLastError = true)]
public static partial SafeProcessHandle OpenProcess(
Expand Down Expand Up @@ -46,5 +76,264 @@ public static partial int NtSetInformationProcess(
int processInformationClass,
ref int processInformation,
int processInformationLength);

[LibraryImport("kernel32.dll")]
public static partial IntPtr GetCurrentProcess();

[LibraryImport("advapi32.dll", SetLastError = true)]
[return: MarshalAs(UnmanagedType.Bool)]
public static partial bool OpenProcessToken(
IntPtr processHandle,
uint desiredAccess,
out SafeAccessTokenHandle tokenHandle);

[LibraryImport("advapi32.dll", EntryPoint = "LookupPrivilegeValueW", SetLastError = true, StringMarshalling = StringMarshalling.Utf16)]
[return: MarshalAs(UnmanagedType.Bool)]
public static partial bool LookupPrivilegeValue(
string? lpSystemName,
string lpName,
out Luid lpLuid);

[LibraryImport("advapi32.dll", SetLastError = true)]
[return: MarshalAs(UnmanagedType.Bool)]
public static partial bool AdjustTokenPrivileges(
SafeAccessTokenHandle tokenHandle,
[MarshalAs(UnmanagedType.Bool)] bool disableAllPrivileges,
ref TokenPrivileges newState,
uint bufferLength,
out TokenPrivileges previousState,
out uint returnLength);

[LibraryImport("advapi32.dll", EntryPoint = "AdjustTokenPrivileges", SetLastError = true)]
[return: MarshalAs(UnmanagedType.Bool)]
public static partial bool AdjustTokenPrivileges(
SafeAccessTokenHandle tokenHandle,
[MarshalAs(UnmanagedType.Bool)] bool disableAllPrivileges,
ref TokenPrivileges newState,
uint bufferLength,
IntPtr previousState,
IntPtr returnLength);

[LibraryImport("advapi32.dll", SetLastError = true)]
[return: MarshalAs(UnmanagedType.Bool)]
public static partial bool GetTokenInformation(
SafeAccessTokenHandle tokenHandle,
uint tokenInformationClass,
IntPtr tokenInformation,
uint tokenInformationLength,
out uint returnLength);

internal static IDisposable AcquirePrivilegeScope(string privilegeName)
{
if (string.IsNullOrWhiteSpace(privilegeName))
{
return EmptyDisposable.Instance;
}

lock (PrivilegeLock)
{
if (ActivePrivileges.TryGetValue(privilegeName, out var existingState))
{
existingState.RefCount++;
return new ScopedPrivilege(existingState);
}

if (!OpenProcessToken(GetCurrentProcess(), TokenAdjustPrivileges | TokenQuery, out var tokenHandle))
{
return EmptyDisposable.Instance;
}

using (tokenHandle)
{
if (tokenHandle.IsInvalid)
{
return EmptyDisposable.Instance;
}

if (!LookupPrivilegeValue(null, privilegeName, out var luid))
{
return EmptyDisposable.Instance;
}

var tp = new TokenPrivileges
{
PrivilegeCount = 1,
Luid = luid,
Attributes = SePrivilegeEnabled,
};

bool adjusted = AdjustTokenPrivileges(
tokenHandle,
false,
ref tp,
(uint)Marshal.SizeOf<TokenPrivileges>(),
out var previousState,
out _);

int lastError = Marshal.GetLastWin32Error();
if (!adjusted || lastError != 0)
{
return EmptyDisposable.Instance;
}

var newState = new PrivilegeState(privilegeName, previousState);
ActivePrivileges[privilegeName] = newState;
return new ScopedPrivilege(newState);
}
}
}

internal static bool TryGetPrivilegeState(string privilegeName, out bool isPresent, out bool isEnabled)
{
isPresent = false;
isEnabled = false;

if (string.IsNullOrWhiteSpace(privilegeName))
{
return false;
}

if (!OpenProcessToken(GetCurrentProcess(), TokenQuery, out var tokenHandle))
{
return false;
}

using (tokenHandle)
{
if (tokenHandle.IsInvalid)
{
return false;
}

if (!LookupPrivilegeValue(null, privilegeName, out var targetLuid))
{
return false;
}

GetTokenInformation(tokenHandle, TokenPrivilegesInformationClass, IntPtr.Zero, 0, out uint requiredLength);
if (requiredLength == 0)
{
return false;
}

IntPtr buffer = Marshal.AllocHGlobal((int)requiredLength);
try
{
if (!GetTokenInformation(tokenHandle, TokenPrivilegesInformationClass, buffer, requiredLength, out _))
{
return false;
}

int count = Marshal.ReadInt32(buffer);
int offset = 4;
for (int i = 0; i < count; i++)
{
var item = Marshal.PtrToStructure<LuidAndAttributes>(IntPtr.Add(buffer, offset));
if (item.Luid.LowPart == targetLuid.LowPart && item.Luid.HighPart == targetLuid.HighPart)
{
isPresent = true;
isEnabled = (item.Attributes & SePrivilegeEnabled) != 0;
return true;
}

offset += Marshal.SizeOf<LuidAndAttributes>();
}

return false;
}
finally
{
Marshal.FreeHGlobal(buffer);
}
}
}

internal static bool IsPrivilegeEnabled(string privilegeName)
{
return TryGetPrivilegeState(privilegeName, out var isPresent, out var isEnabled) && isPresent && isEnabled;
}

internal static bool IsPrivilegePresent(string privilegeName)
{
return TryGetPrivilegeState(privilegeName, out var isPresent, out _) && isPresent;
}

private sealed class PrivilegeState
{
public PrivilegeState(string privilegeName, TokenPrivileges previousState)
{
this.PrivilegeName = privilegeName;
this.PreviousState = previousState;
this.RefCount = 1;
}

public string PrivilegeName { get; }

public TokenPrivileges PreviousState { get; }

public int RefCount { get; set; }

public void Restore()
{
if (this.PreviousState.PrivilegeCount > 0)
{
if (OpenProcessToken(GetCurrentProcess(), TokenAdjustPrivileges | TokenQuery, out var tokenHandle))
{
using (tokenHandle)
{
if (!tokenHandle.IsInvalid)
{
var previous = this.PreviousState;
AdjustTokenPrivileges(
tokenHandle,
false,
ref previous,
0,
IntPtr.Zero,
IntPtr.Zero);
}
}
}
}
}
}

private sealed class ScopedPrivilege : IDisposable
{
private readonly PrivilegeState state;
private int disposed;

public ScopedPrivilege(PrivilegeState state)
{
this.state = state;
}

public void Dispose()
{
if (Interlocked.Exchange(ref this.disposed, 1) != 0)
{
return;
}

lock (PrivilegeLock)
{
this.state.RefCount--;
if (this.state.RefCount <= 0)
{
ActivePrivileges.Remove(this.state.PrivilegeName);
this.state.Restore();
}
}
}
}

private sealed class EmptyDisposable : IDisposable
{
public static EmptyDisposable Instance { get; } = new();

public void Dispose()
{
}
}
}
}
19 changes: 17 additions & 2 deletions Services/ProcessMemoryPriorityService.cs
Original file line number Diff line number Diff line change
Expand Up @@ -244,7 +244,22 @@ public Task<ProcessOperationResult> SetIoPriorityAsync(ProcessModel process, Pro
}

var value = (int)priority;
var status = this.nativeApi.SetIoPriority(handle, ref value);
int status;
IDisposable? privilegeScope = null;
if (priority == ProcessIoPriority.High)
{
privilegeScope = this.nativeApi.AcquirePrivilegeScope(ProcessMemoryPriorityNativeMethods.SeIncreaseBasePriorityPrivilege);
}

try
{
status = this.nativeApi.SetIoPriority(handle, ref value);
}
finally
{
privilegeScope?.Dispose();
}

if (status == 0)
{
return Task.FromResult(ProcessOperationResult.Succeeded(
Expand All @@ -254,7 +269,7 @@ public Task<ProcessOperationResult> SetIoPriorityAsync(ProcessModel process, Pro

return Task.FromResult(status switch
{
unchecked((int)0xC0000022) => ProcessOperationResult.Failed(
unchecked((int)0xC0000022) or unchecked((int)0xC0000061) => ProcessOperationResult.Failed(
AffinityApplyErrorCodes.AccessDenied,
ProcessOperationUserMessages.AccessDenied,
$"NtSetInformationProcess failed with NTSTATUS 0x{status:X8}.",
Expand Down
Loading
Loading