Skip to content

Commit 3f41b59

Browse files
neon-nyanneon-nyan
committed
Fix early registeredWaitHandle access on result finalization
1 parent f4f3415 commit 3f41b59

1 file changed

Lines changed: 126 additions & 117 deletions

File tree

‎Utility/ComAsyncExtension.cs‎

Lines changed: 126 additions & 117 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,6 @@
66
using System.Threading;
77
using System.Threading.Tasks;
88

9-
// ReSharper disable AccessToModifiedClosure
109
namespace Hi3Helper.Plugin.Core.Utility;
1110

1211
/// <summary>
@@ -48,164 +47,174 @@ public static nint AsResult<T>(this Task<T> task)
4847
/// <param name="resultP">A pointer of the <see cref="ComAsyncResult"/> struct.</param>
4948
/// <returns>An awaitable managed <see cref="Task"/>.</returns>
5049
public static Task AsTask(this nint resultP)
51-
{
52-
nint waitHandleP = ComAsyncResult.GetWaitHandle(resultP);
53-
54-
RegisteredWaitHandle? registeredWaitHandle = null;
55-
TaskCompletionSource tcs = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
56-
SafeWaitHandle safeHandle = new SafeWaitHandle(waitHandleP, false);
57-
58-
WaitHandle waitHandle = new EventWaitHandle(false, EventResetMode.ManualReset)
59-
{
60-
SafeWaitHandle = safeHandle
61-
};
62-
63-
registeredWaitHandle = ThreadPool.RegisterWaitForSingleObject(waitHandle, Impl, null, -1, true);
64-
65-
return tcs.Task;
66-
67-
unsafe void Impl(object? state, bool isTimedOut)
68-
{
69-
SetResult<TaskCompletionSource, nint>(resultP.AsPointer<ComAsyncResult>(), tcs);
70-
71-
safeHandle.Dispose();
72-
waitHandle.Dispose();
73-
ComAsyncResult.FreeResult(resultP);
74-
75-
if (waitHandleP != nint.Zero)
76-
{
77-
PInvoke.CloseHandle(waitHandleP);
78-
}
79-
80-
registeredWaitHandle!.Unregister(null);
81-
}
82-
}
50+
=> new RegisteredWaitHandleContext(resultP).Task;
8351

8452
/// <summary>
8553
/// Re-marshal a pointer of <see cref="ComAsyncResult"/> struct into a managed <see cref="Task"/>.
8654
/// </summary>
8755
/// <param name="resultP">A pointer of the <see cref="ComAsyncResult"/> struct.</param>
8856
/// <returns>An awaitable managed <see cref="Task"/>.</returns>
8957
public static Task<T> AsTask<T>(this nint resultP)
90-
where T : unmanaged
58+
where T : unmanaged => (Task<T>)new RegisteredWaitHandleContext<T>(resultP).Task;
59+
60+
/// <summary>
61+
/// A wait handle context to registers and returns a result once async operation is completed.
62+
/// </summary>
63+
private class RegisteredWaitHandleContext
9164
{
92-
nint waitHandleP = ComAsyncResult.GetWaitHandle(resultP);
65+
protected nint AsyncResultP { get; }
66+
private nint UnmanagedWaitHandleP { get; }
67+
private SafeWaitHandle SafeWaitHandle { get; }
68+
private WaitHandle WaitHandle { get; }
69+
private RegisteredWaitHandle RegisteredWaitHandle { get; }
9370

94-
RegisteredWaitHandle? registeredWaitHandle = null;
95-
TaskCompletionSource<T> taskSource = new(TaskCreationOptions.RunContinuationsAsynchronously);
96-
SafeWaitHandle safeHandle = new SafeWaitHandle(waitHandleP, false);
71+
private TaskCompletionSource TaskSourceResult { get; } = new(TaskCreationOptions.RunContinuationsAsynchronously);
9772

98-
WaitHandle waitHandle = new EventWaitHandle(false, EventResetMode.ManualReset)
73+
public virtual Task Task => TaskSourceResult.Task;
74+
75+
public RegisteredWaitHandleContext(nint asyncResultP)
9976
{
100-
SafeWaitHandle = safeHandle
101-
};
77+
AsyncResultP = asyncResultP;
78+
UnmanagedWaitHandleP = ComAsyncResult.GetWaitHandle(asyncResultP);
10279

103-
registeredWaitHandle = ThreadPool.RegisterWaitForSingleObject(waitHandle, Impl, null, -1, true);
80+
SafeWaitHandle = new SafeWaitHandle(UnmanagedWaitHandleP, false);
81+
WaitHandle = new EventWaitHandle(false, EventResetMode.ManualReset)
82+
{
83+
SafeWaitHandle = SafeWaitHandle
84+
};
85+
RegisteredWaitHandle = ThreadPool.RegisterWaitForSingleObject(WaitHandle, SetResultCallback, this, -1, true);
86+
}
10487

105-
return taskSource.Task;
88+
private static void SetResultCallback(object? state, bool isTimedOut)
89+
{
90+
RegisteredWaitHandleContext ctx = (RegisteredWaitHandleContext)state!;
91+
ctx.SetResultCallbackCore();
92+
}
10693

107-
unsafe void Impl(object? state, bool isTimedOut)
94+
protected virtual unsafe void SetResultCallbackCore()
10895
{
109-
SetResult<TaskCompletionSource<T>, T>(resultP.AsPointer<ComAsyncResult>(), taskSource);
96+
SetResult<TaskCompletionSource, nint>(AsyncResultP.AsPointer<ComAsyncResult>(), TaskSourceResult);
97+
CloseHandles();
98+
}
11099

111-
safeHandle.Dispose();
112-
waitHandle.Dispose();
113-
ComAsyncResult.FreeResult(resultP);
100+
/// <summary>
101+
/// Closes wait handles and free <see cref="ComAsyncResult"/> struct.
102+
/// </summary>
103+
protected void CloseHandles()
104+
{
105+
SafeWaitHandle.Dispose();
106+
WaitHandle.Dispose();
107+
ComAsyncResult.FreeResult(AsyncResultP);
114108

115-
if (waitHandleP != nint.Zero)
109+
if (UnmanagedWaitHandleP != nint.Zero)
116110
{
117-
PInvoke.CloseHandle(waitHandleP);
111+
PInvoke.CloseHandle(UnmanagedWaitHandleP);
118112
}
119113

120-
registeredWaitHandle!.Unregister(null);
114+
RegisteredWaitHandle.Unregister(null);
121115
}
122-
}
123116

124-
/// <summary>
125-
/// Sets result of the async operation based on its return type into the <see cref="TaskCompletionSource"/>.
126-
/// </summary>
127-
/// <typeparam name="TSource">A type of <see cref="TaskCompletionSource"/> or <see cref="TaskCompletionSource{TResult}"/></typeparam>
128-
/// <typeparam name="T">A type of the return value from <see cref="ComAsyncResult"/>.</typeparam>
129-
/// <param name="asyncResult">A pointer to the <see cref="ComAsyncResult"/> struct.</param>
130-
/// <param name="tcs">An instance of <see cref="TaskCompletionSource"/> or <see cref="TaskCompletionSource{TResult}"/> to set the exception or result to.</param>
131-
/// <exception cref="InvalidCastException">Whether the <typeparamref name="TSource"/> is not either <see cref="TaskCompletionSource"/> or <see cref="TaskCompletionSource{TResult}"/></exception>
132-
private static unsafe void SetResult<TSource, T>(ComAsyncResult* asyncResult, TSource tcs)
133-
where T : unmanaged
134-
{
135-
// Try check the TCS kind.
136-
switch (tcs)
117+
/// <summary>
118+
/// Sets result of the async operation based on its return type into the <see cref="TaskCompletionSource"/>.
119+
/// </summary>
120+
/// <typeparam name="TSource">A type of <see cref="TaskCompletionSource"/> or <see cref="TaskCompletionSource{TResult}"/></typeparam>
121+
/// <typeparam name="T">A type of the return value from <see cref="ComAsyncResult"/>.</typeparam>
122+
/// <param name="asyncResult">A pointer to the <see cref="ComAsyncResult"/> struct.</param>
123+
/// <param name="tcs">An instance of <see cref="TaskCompletionSource"/> or <see cref="TaskCompletionSource{TResult}"/> to set the exception or result to.</param>
124+
/// <exception cref="InvalidCastException">Whether the <typeparamref name="TSource"/> is not either <see cref="TaskCompletionSource"/> or <see cref="TaskCompletionSource{TResult}"/></exception>
125+
protected static unsafe void SetResult<TSource, T>(ComAsyncResult* asyncResult, TSource tcs)
126+
where T : unmanaged
137127
{
138-
case TaskCompletionSource<T> asTcsRes:
139-
SetResultInner(asyncResult, () => asTcsRes.SetResult(*(T*)asyncResult->_resultP), asTcsRes.SetException, asTcsRes.SetCanceled);
140-
return;
141-
case TaskCompletionSource asTcs:
142-
SetResultInner(asyncResult, asTcs.SetResult, asTcs.SetException, asTcs.SetCanceled);
143-
return;
144-
}
128+
// Try check the TCS kind.
129+
switch (tcs)
130+
{
131+
case TaskCompletionSource<T> asTcsRes:
132+
SetResultInner(asyncResult, () => asTcsRes.SetResult(*(T*)asyncResult->_resultP), asTcsRes.SetException, asTcsRes.SetCanceled);
133+
return;
134+
case TaskCompletionSource asTcs:
135+
SetResultInner(asyncResult, asTcs.SetResult, asTcs.SetException, asTcs.SetCanceled);
136+
return;
137+
}
145138

146-
// Bada bing, bada bong. Throw the cast exception.
147-
throw new InvalidCastException($"Cannot cast {nameof(TSource)} to either {nameof(TaskCompletionSource<>)} or {nameof(TaskCompletionSource)}");
148-
}
139+
// Bada bing, bada bong. Throw the cast exception.
140+
throw new InvalidCastException($"Cannot cast {nameof(TSource)} to either {nameof(TaskCompletionSource<>)} or {nameof(TaskCompletionSource)}");
141+
}
149142

150-
private static unsafe void SetResultInner(ComAsyncResult* asyncResult, Action actionSetRes, SetExceptionDelegate actionSetExc, Action actionSetCancel)
151-
{
152-
// IsCancelled is set without IsFaulty being set. So, try check if the exception is not empty. If yes, then set the default exception.
153-
if (asyncResult->IsCancelled)
143+
private static unsafe void SetResultInner(ComAsyncResult* asyncResult, Action actionSetRes, SetExceptionDelegate actionSetExc, Action actionSetCancel)
154144
{
155-
if (asyncResult->ExceptionMemory.IsEmpty)
145+
// IsCancelled is set without IsFaulty being set. So, try check if the exception is not empty. If yes, then set the default exception.
146+
if (asyncResult->IsCancelled)
147+
{
148+
if (asyncResult->ExceptionMemory.IsEmpty)
149+
{
150+
actionSetExc(new TaskCanceledException());
151+
return;
152+
}
153+
154+
// Otherwise, set the exception from ExceptionMemory
155+
SetException(asyncResult, actionSetExc, actionSetCancel);
156+
return;
157+
}
158+
159+
// If it isn't faulty, invoke result.
160+
if (!asyncResult->IsFaulty)
156161
{
157-
actionSetExc(new TaskCanceledException());
162+
actionSetRes();
158163
return;
159164
}
160165

161-
// Otherwise, set the exception from ExceptionMemory
166+
// Otherwise, invoke the exception.
162167
SetException(asyncResult, actionSetExc, actionSetCancel);
163-
return;
164168
}
165169

166-
// If it isn't faulty, invoke result.
167-
if (!asyncResult->IsFaulty)
170+
/// <summary>
171+
/// Set the exception provided by the <see cref="ComAsyncResult.ExceptionMemory"/>.
172+
/// </summary>
173+
/// <param name="asyncResult">A pointer to the <see cref="ComAsyncResult"/> struct.</param>
174+
/// <param name="actionSetExc">A callback to set the exception into the <see cref="TaskCompletionSource"/></param>
175+
/// <param name="actionSetCancel">A callback to set the canceled status into the <see cref="TaskCompletionSource"/></param>
176+
/// <remarks>
177+
/// If the buffer provided by the <see cref="ComAsyncResult.ExceptionMemory"/> is empty, then a generic <see cref="COMException"/> will be set.
178+
/// </remarks>
179+
private static unsafe void SetException(ComAsyncResult* asyncResult, SetExceptionDelegate actionSetExc, Action actionSetCancel)
168180
{
169-
actionSetRes();
170-
return;
171-
}
181+
StackTrace currentStackTrace = new(true);
182+
Exception? exception = ComAsyncException.GetExceptionFromHandle(asyncResult->ExceptionMemory);
172183

173-
// Otherwise, invoke the exception.
174-
SetException(asyncResult, actionSetExc, actionSetCancel);
175-
}
184+
#if DEBUG
185+
Debug.Assert(exception != null, $"Exception shouldn't be null! Result address is: 0x{(nint)asyncResult:x8}");
186+
#endif
176187

177-
/// <summary>
178-
/// Set the exception provided by the <see cref="ComAsyncResult.ExceptionMemory"/>.
179-
/// </summary>
180-
/// <param name="asyncResult">A pointer to the <see cref="ComAsyncResult"/> struct.</param>
181-
/// <param name="actionSetExc">A callback to set the exception into the <see cref="TaskCompletionSource"/></param>
182-
/// <param name="actionSetCancel">A callback to set the canceled status into the <see cref="TaskCompletionSource"/></param>
183-
/// <remarks>
184-
/// If the buffer provided by the <see cref="ComAsyncResult.ExceptionMemory"/> is empty, then a generic <see cref="COMException"/> will be set.
185-
/// </remarks>
186-
private static unsafe void SetException(ComAsyncResult* asyncResult, SetExceptionDelegate actionSetExc, Action actionSetCancel)
187-
{
188-
StackTrace currentStackTrace = new(true);
189-
Exception? exception = ComAsyncException.GetExceptionFromHandle(asyncResult->ExceptionMemory);
188+
// ReSharper disable once ConditionIsAlwaysTrueOrFalseAccordingToNullableAPIContract
189+
if (exception == null)
190+
{
191+
actionSetExc(new COMException());
192+
return;
193+
}
190194

191-
#if DEBUG
192-
Debug.Assert(exception != null, $"Exception shouldn't be null! Result address is: 0x{(nint)asyncResult:x8}");
193-
#endif
195+
exception.SetExceptionRemoteStackTrace() += Environment.NewLine;
196+
exception.SetExceptionStackTrace() = currentStackTrace.ToString();
197+
actionSetExc(exception);
194198

195-
// ReSharper disable once ConditionIsAlwaysTrueOrFalseAccordingToNullableAPIContract
196-
if (exception == null)
197-
{
198-
actionSetExc(new COMException());
199-
return;
199+
if (exception is OperationCanceledException)
200+
{
201+
actionSetCancel();
202+
}
200203
}
204+
}
205+
206+
/// <inheritdoc/>
207+
private class RegisteredWaitHandleContext<T>(nint asyncResultP) : RegisteredWaitHandleContext(asyncResultP)
208+
where T : unmanaged
209+
{
210+
private TaskCompletionSource<T> TaskSourceResult { get; } = new(TaskCreationOptions.RunContinuationsAsynchronously);
201211

202-
exception.SetExceptionRemoteStackTrace() += Environment.NewLine;
203-
exception.SetExceptionStackTrace() = currentStackTrace.ToString();
204-
actionSetExc(exception);
212+
public override Task Task => TaskSourceResult.Task;
205213

206-
if (exception is OperationCanceledException)
214+
protected override unsafe void SetResultCallbackCore()
207215
{
208-
actionSetCancel();
216+
SetResult<TaskCompletionSource<T>, T>(AsyncResultP.AsPointer<ComAsyncResult>(), TaskSourceResult);
217+
CloseHandles();
209218
}
210219
}
211220
}

0 commit comments

Comments
 (0)