Skip to content

Commit

Permalink
SafeHandle to avoid use-after-free race condition
Browse files Browse the repository at this point in the history
  • Loading branch information
@antonfirsov
antonfirsov committed Jan 31, 2022
1 parent 8b9c79b commit a2a9355
Show file tree
Hide file tree
Showing 2 changed files with 61 additions and 37 deletions.
96 changes: 61 additions & 35 deletions src/libraries/System.Net.NameResolution/src/System/Net/NameResolutionPal.Windows.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
using System.Threading;
using System.Threading.Tasks;
using System.Diagnostics;
using Microsoft.Win32.SafeHandles;

namespace System.Net
{
Expand Down Expand Up @@ -138,17 +139,14 @@ public static unsafe string GetHostName()
{
Interop.Winsock.EnsureInitialized();

GetAddrInfoExContext* context = GetAddrInfoExContext.AllocateContext();

GetAddrInfoExState state;
GetAddrInfoExState? state = null;
try
{
state = new GetAddrInfoExState(context, hostName, justAddresses);
context->QueryStateHandle = state.CreateHandle();
state = new GetAddrInfoExState(hostName, justAddresses);
}
catch
{
GetAddrInfoExContext.FreeContext(context);
state?.Dispose();
throw;
}

Expand All @@ -158,6 +156,8 @@ public static unsafe string GetHostName()
hints.ai_flags = AddressInfoHints.AI_CANONNAME;
}

GetAddrInfoExContext* context = state.Context;

SocketError errorCode = (SocketError)Interop.Winsock.GetAddrInfoExW(
hostName, null, Interop.Winsock.NS_ALL, IntPtr.Zero, &hints, &context->Result, IntPtr.Zero, &context->Overlapped, &GetAddressInfoExCallback, &context->CancelHandle);

Expand All @@ -172,7 +172,7 @@ public static unsafe string GetHostName()
// and final result would be posted via overlapped IO.
// synchronous failure here may signal issue when GetAddrInfoExW does not work from
// impersonated context. Windows 8 and Server 2012 fail for same reason with different errorCode.
GetAddrInfoExContext.FreeContext(context);
state.Dispose();
return null;
}
else
Expand All @@ -194,10 +194,10 @@ private static unsafe void GetAddressInfoExCallback(int error, int bytes, Native

private static unsafe void ProcessResult(SocketError errorCode, GetAddrInfoExContext* context)
{
GetAddrInfoExState state = GetAddrInfoExState.FromHandleAndFree(context->QueryStateHandle);

try
{
GetAddrInfoExState state = GetAddrInfoExState.FromHandleAndFree(context->QueryStateHandle);

CancellationToken cancellationToken = state.UnregisterAndGetCancellationToken();

if (errorCode == SocketError.Success)
Expand All @@ -222,7 +222,7 @@ private static unsafe void ProcessResult(SocketError errorCode, GetAddrInfoExCon
}
finally
{
GetAddrInfoExContext.FreeContext(context);
state.Dispose();
}
}

Expand Down Expand Up @@ -360,18 +360,21 @@ private static unsafe IPAddress CreateIPv6Address(ReadOnlySpan<byte> socketAddre
return new IPAddress(address, scope);
}

private sealed unsafe class GetAddrInfoExState : IThreadPoolWorkItem
// GetAddrInfoExState is a SafeHandle that manages the lifetime of GetAddrInfoExContext*
// to make sure GetAddrInfoExCancel always takes a valid memory address regardless of the race
// between cancellation and completion callbacks.
private sealed unsafe class GetAddrInfoExState : SafeHandleZeroOrMinusOneIsInvalid, IThreadPoolWorkItem
{
private IntPtr _cancellationContextPtr;
private CancellationTokenRegistration _cancellationRegistration;

private AsyncTaskMethodBuilder<IPHostEntry> IPHostEntryBuilder;
private AsyncTaskMethodBuilder<IPAddress[]> IPAddressArrayBuilder;
private object? _result;
private volatile bool _completed;

public GetAddrInfoExState(GetAddrInfoExContext *context, string hostName, bool justAddresses)
public GetAddrInfoExState(string hostName, bool justAddresses)
: base(true)
{
_cancellationContextPtr = (IntPtr)context;
HostName = hostName;
JustAddresses = justAddresses;
if (justAddresses)
Expand All @@ -384,6 +387,10 @@ public GetAddrInfoExState(GetAddrInfoExContext *context, string hostName, bool j
IPHostEntryBuilder = AsyncTaskMethodBuilder<IPHostEntry>.Create();
_ = IPHostEntryBuilder.Task; // force initialization
}

GetAddrInfoExContext* context = GetAddrInfoExContext.AllocateContext();
context->QueryStateHandle = CreateHandle();
SetHandle((IntPtr)context);
}

public string HostName { get; }
Expand All @@ -392,46 +399,58 @@ public GetAddrInfoExState(GetAddrInfoExContext *context, string hostName, bool j

public Task Task => JustAddresses ? (Task)IPAddressArrayBuilder.Task : IPHostEntryBuilder.Task;

private GetAddrInfoExContext* CancellationContext => (GetAddrInfoExContext*)_cancellationContextPtr;
internal GetAddrInfoExContext* Context => (GetAddrInfoExContext*)handle;

public void RegisterForCancellation(CancellationToken cancellationToken)
{
if (!cancellationToken.CanBeCanceled) return;

lock (this)
if (_completed)
{
if (CancellationContext == null)
// The operation completed before registration could be done.
return;
}

_cancellationRegistration = cancellationToken.UnsafeRegister(static o =>
{
var @this = (GetAddrInfoExState)o!;
if (@this._completed)
{
// The operation completed before registration could be done.
// Escape early and avoid ObjectDisposedException in DangerousAddRef
return;
}

_cancellationRegistration = cancellationToken.UnsafeRegister(o =>
bool needRelease = false;
try
{
var @this = (GetAddrInfoExState)o!;
int cancelResult = 0;
@this.DangerousAddRef(ref needRelease);

GetAddrInfoExContext* context = @this.CancellationContext;
// If DangerousAddRef didn't throw ODE, the handle should contain a valid pointer.
GetAddrInfoExContext* context = @this.Context;

if (context != null)
// An outstanding operation will be completed with WSA_E_CANCELLED, and GetAddrInfoExCancel will return NO_ERROR.
// If this thread has lost the race between cancellation and completion, this will be a NOP
// with GetAddrInfoExCancel returning WSA_INVALID_HANDLE.
int cancelResult = Interop.Winsock.GetAddrInfoExCancel(&context->CancelHandle);
if (cancelResult != Interop.Winsock.WSA_INVALID_HANDLE && NetEventSource.Log.IsEnabled())
{
// An outstanding operation will be completed with WSA_E_CANCELLED, and GetAddrInfoExCancel will return NO_ERROR.
// If this thread has lost the race between cancellation and completion, this will be a NOP
// with GetAddrInfoExCancel returning WSA_INVALID_HANDLE.
cancelResult = Interop.Winsock.GetAddrInfoExCancel(&context->CancelHandle);
NetEventSource.Info(@this, $"GetAddrInfoExCancel returned error {cancelResult}");
}

if (cancelResult != 0 && cancelResult != Interop.Winsock.WSA_INVALID_HANDLE && NetEventSource.Log.IsEnabled())
}
finally
{
if (needRelease)
{
NetEventSource.Info(@this, $"GetAddrInfoExCancel returned error {cancelResult}");
@this.DangerousRelease();
}
}, this);
}
}

}, this);
}

public CancellationToken UnregisterAndGetCancellationToken()
{
Volatile.Write(ref _cancellationContextPtr, IntPtr.Zero);
_completed = true;

// We should not wait for pending cancellation callbacks with CTR.Dispose(),
// since we are in a completion routine and GetAddrInfoExCancel may get blocked until it's finished.
Expand Down Expand Up @@ -477,15 +496,22 @@ void IThreadPoolWorkItem.Execute()
}
}

public IntPtr CreateHandle() => GCHandle.ToIntPtr(GCHandle.Alloc(this, GCHandleType.Normal));

public static GetAddrInfoExState FromHandleAndFree(IntPtr handle)
{
GCHandle gcHandle = GCHandle.FromIntPtr(handle);
var state = (GetAddrInfoExState)gcHandle.Target!;
gcHandle.Free();
return state;
}

protected override bool ReleaseHandle()
{
GetAddrInfoExContext.FreeContext(Context);

return true;
}

private IntPtr CreateHandle() => GCHandle.ToIntPtr(GCHandle.Alloc(this, GCHandleType.Normal));
}

[StructLayout(LayoutKind.Sequential)]
Expand Down
2 changes: 0 additions & 2 deletions src/libraries/System.Net.NameResolution/tests/FunctionalTests/GetHostAddressesTest.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -177,7 +177,6 @@ public async Task DnsGetHostAddresses_PreCancelledToken_Throws()
[Collection(nameof(DisableParallelization))]
public class GetHostAddressesTest_Cancellation
{
[OuterLoop]
[Fact]
[ActiveIssue("https://github.com/dotnet/runtime/issues/33378", TestPlatforms.AnyUnix)] // Cancellation of an outstanding getaddrinfo is not supported on *nix.
public async Task DnsGetHostAddresses_PostCancelledToken_Throws()
Expand All @@ -195,7 +194,6 @@ public async Task DnsGetHostAddresses_PostCancelledToken_Throws()
}

// This is a regression test for https://github.com/dotnet/runtime/issues/63552
[OuterLoop]
[Fact]
[ActiveIssue("https://github.com/dotnet/runtime/issues/33378", TestPlatforms.AnyUnix)] // Cancellation of an outstanding getaddrinfo is not supported on *nix.
public async Task DnsGetHostAddresses_ResolveParallelCancelOnFailure_AllCallsReturn()
Expand Down

0 comments on commit a2a9355

Please sign in to comment.