From 9ff43ba473b5158ee8028943e7bda0dcc6260fac Mon Sep 17 00:00:00 2001 From: Mounir IDRASSI Date: Fri, 16 Oct 2026 18:06:08 +0900 Subject: [PATCH] Windows: fix KDF completion lifetime Hold a requester reference until header parsing can no longer enqueue KDFs. This prevents stale completion signals from allowing cleanup while workers still use request-owned state. Add bounded tests of the production pool and caller completion protocol. --- Tests/test_windows_kdf_completion.py | 74 ++++++ Tests/windows_kdf_completion_test.c | 323 +++++++++++++++++++++++++++ src/Common/EncryptionThreadPool.c | 2 +- src/Common/EncryptionThreadPool.h | 10 +- src/Common/Volumes.c | 15 +- 5 files changed, 415 insertions(+), 9 deletions(-) create mode 100644 Tests/test_windows_kdf_completion.py create mode 100644 Tests/windows_kdf_completion_test.c diff --git a/Tests/test_windows_kdf_completion.py b/Tests/test_windows_kdf_completion.py new file mode 100644 index 00000000..c5163adf --- /dev/null +++ b/Tests/test_windows_kdf_completion.py @@ -0,0 +1,74 @@ +#!/usr/bin/env python3 +"""Check Windows KDF request lifetimes without mounting volumes or loading a driver. + +Run from an x64 MSVC developer environment: + python Tests/test_windows_kdf_completion.py [--build-dir PATH] + +The harness compiles the production thread pool with event-controlled KDF stubs. +Initialization and draining are extracted from ReadVolumeHeaderWithAbort, so the +caller and worker reference protocol are checked together. All request storage +is retained until pool shutdown joins the workers. +""" +import argparse +import os +from pathlib import Path +import shutil +import subprocess +import tempfile + +ROOT = Path(__file__).resolve().parents[1] + + +def block(source, start): + """Return a balanced source block, including its opening and closing braces.""" + first = source.index("{", start) + depth = 1 + for end in range(first + 1, len(source)): + depth += (source[end] == "{") - (source[end] == "}") + if depth == 0: + return source[first:end + 1] + raise RuntimeError("Unbalanced production source block") + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--build-dir", type=Path) + args = parser.parse_args() + if os.name != "nt" or not shutil.which("cl.exe"): + parser.error("An x64 MSVC developer environment on Windows is required") + build = (args.build_dir or Path(tempfile.mkdtemp(prefix="vc-kdf-completion-"))).resolve() + build.mkdir(parents=True, exist_ok=True) + + source = (ROOT / "src/Common/Volumes.c").read_text() + start = source.index("int ReadVolumeHeaderWithAbort (") + header = block(source, start) + condition = "if ((selected_pkcs5_prf == 0) && (encryptionThreadCount > 1))" + init_start = header.index(condition) + init = block(header, init_start) + drain_start = header.index(condition, header.index("\nret:")) + drain = block(header, drain_start) + # Keep the real reference release and wait; retain allocations for the test's + # independent final join instead of executing production cleanup here. + drain = drain[:drain.index("#if !defined(DEVICE_DRIVER)")] + "}\n" + item_end = source.index("} KeyDerivationWorkItem;") + len("} KeyDerivationWorkItem;") + item_start = source.rindex("typedef struct", 0, item_end) + for filename, contents in (("kdf_init.inc", init), ("kdf_drain.inc", drain), + ("kdf_item.inc", source[item_start:item_end])): + (build / filename).write_text(contents) + + exe = build / "windows_kdf_completion_test.exe" + command = ["cl.exe", "/nologo", "/O2", "/W3", "/WX", "/MD", + "/DUNICODE", "/D_UNICODE", f"/I{build}"] + command += [f"/I{ROOT / 'src' / directory}" + for directory in ("", "Common", "Crypto", "Crypto/Argon2/include")] + command += [str(ROOT / "Tests/windows_kdf_completion_test.c"), f"/Fe:{exe}"] + result = subprocess.run(command, cwd=build, capture_output=True, text=True) + (build / "build.log").write_text(result.stdout + result.stderr) + if result.returncode: + raise RuntimeError(result.stdout + result.stderr) + for workers in (2, 8): + subprocess.run([str(exe), str(workers)], cwd=build, check=True, timeout=90) + + +if __name__ == "__main__": + main() diff --git a/Tests/windows_kdf_completion_test.c b/Tests/windows_kdf_completion_test.c new file mode 100644 index 00000000..7cabe6cc --- /dev/null +++ b/Tests/windows_kdf_completion_test.c @@ -0,0 +1,323 @@ +/* Copyright (c) 2026 AM Crypto. Licensed under the Apache License 2.0. + * Bounded user-mode checks of the production Windows KDF reference protocol. + * Build through test_windows_kdf_completion.py, which supplies caller excerpts. + */ +#include "Common/EncryptionThreadPool.h" +#include "Common/Pkcs5.h" +#include + +static BOOL TestSetEvent (HANDLE event); +static DWORD TestWait (HANDLE event); +#undef TC_WAIT_EVENT +#define TC_WAIT_EVENT(event) TestWait (event) +#define SetEvent TestSetEvent +#include "Common/EncryptionThreadPool.c" +#undef SetEvent +#include "kdf_item.inc" + +#define MAX_REQUESTS 16 +#define PRF_COUNT (LAST_PRF_ID - FIRST_PRF_ID + 1) +typedef struct Request Request; +typedef struct +{ + Request *Owner; + KeyDerivationWorkItem *Item; + HANDLE Started, MayReturn, ReadyPublished, MayReleaseReference, WorkerDone; + int Result; + BOOL HoldAfterReady; +} Job; + +struct Request +{ + TC_EVENT *Ready, *Drained; + LONG *References; + KeyDerivationWorkItem *Items; + LONG volatile Abort; + LONG volatile Signals; + HANDLE WaitEntered, FinishDone, SignalEntered, MaySignal; + HANDLE RequesterThread; + BOOL HoldFinalSignal; + int JobCount; + Job Jobs[PRF_COUNT]; +}; + +static Request Requests[MAX_REQUESTS]; +static int RequestCount; +static __declspec(thread) Job *CurrentJob; +static __declspec(thread) Request *DrainingRequest; + +static void Require (BOOL condition, const char *message) +{ + if (!condition) + { + fprintf (stderr, "FAIL: %s\n", message); + ExitProcess (1); + } +} + +static HANDLE Event (void) +{ + HANDLE event = CreateEvent (NULL, FALSE, FALSE, NULL); + Require (event != NULL, "create test event"); + return event; +} + +static void Wait (HANDLE event) +{ + Require (WaitForSingleObject (event, 10000) == WAIT_OBJECT_0, "test barrier timed out"); +} + +static LONG References (Request *request) +{ + return InterlockedCompareExchange (request->References, 0, 0); +} + +static DWORD TestWait (HANDLE event) +{ + if (DrainingRequest && event == *DrainingRequest->Drained) + SetEvent (DrainingRequest->WaitEntered); + return WaitForSingleObject (event, INFINITE); +} + +static BOOL TestSetEvent (HANDLE event) +{ + Job *job = CurrentJob; + if (job) + { + Request *request = job->Owner; + if (event == *request->Drained) + { + Require (References (request) == 0, "completion signaled with outstanding references"); + Require (InterlockedIncrement (&request->Signals) == 1, "completion signaled twice"); + if (request->HoldFinalSignal) + { + SetEvent (request->SignalEntered); + Wait (request->MaySignal); + } + } + if (event == *request->Ready && job->HoldAfterReady) + { + BOOL result = SetEvent (event); + SetEvent (job->ReadyPublished); + Wait (job->MayReleaseReference); + return result; + } + if (event == WorkItemCompletedEvent) + { + BOOL result = SetEvent (event); + CurrentJob = NULL; + SetEvent (job->WorkerDone); + return result; + } + } + return SetEvent (event); +} + +static int InitializeRequest (Request *request) +{ + TC_EVENT *keyDerivationCompletedEvent = NULL, *noOutstandingWorkItemEvent = NULL; + LONG *outstandingWorkItemCount = NULL; + KeyDerivationWorkItem *keyDerivationWorkItems = NULL; + int keyDerivationWorkItemsSize = 0, pkcs5PrfCount = PRF_COUNT, i; +#include "kdf_init.inc" + request->Ready = keyDerivationCompletedEvent; + request->Drained = noOutstandingWorkItemEvent; + request->References = outstandingWorkItemCount; + request->Items = keyDerivationWorkItems; + return 0; +} + +static Request *NewRequest (void) +{ + Request *request; + Require (RequestCount < MAX_REQUESTS, "request fixture capacity"); + request = &Requests[RequestCount++]; + Require (InitializeRequest (request) == 0, "initialize request"); + /* Refuse to exercise an implementation without the caller reference. */ + Require (References (request) == 1, "missing requester reference"); + Require (WaitForSingleObject (*request->Drained, 0) == WAIT_TIMEOUT, "initial completion event signaled"); + request->WaitEntered = Event (); + request->FinishDone = Event (); + request->SignalEntered = Event (); + request->MaySignal = Event (); + return request; +} + +static unsigned __stdcall FinishRequest (void *context) +{ + Request *request = (Request *) context; + LONG *outstandingWorkItemCount = request->References; + TC_EVENT *noOutstandingWorkItemEvent = request->Drained; + DrainingRequest = request; +#include "kdf_drain.inc" + Require (References (request) == 0, "request finished with workers outstanding"); + DrainingRequest = NULL; + SetEvent (request->FinishDone); + return 0; +} + +static void BeginFinish (Request *request) +{ + request->RequesterThread = (HANDLE) _beginthreadex (NULL, 0, FinishRequest, request, 0, NULL); + Require (request->RequesterThread != NULL, "start requester thread"); + Wait (request->WaitEntered); + Require (WaitForSingleObject (request->FinishDone, 0) == WAIT_TIMEOUT, "request did not wait for workers"); +} + +static void EndFinish (Request *request) +{ + Wait (request->FinishDone); + Wait (request->RequesterThread); + Require (References (request) == 0 && request->Signals == 1, "last worker completion"); +} + +static Job *Submit (Request *request, int prf, BOOL holdAfterReady, int result) +{ + Job *job; + Require (request->JobCount < PRF_COUNT, "job fixture capacity"); + job = &request->Jobs[request->JobCount]; + job->Owner = request; + job->Item = &request->Items[request->JobCount++]; + job->Item->Free = FALSE; + job->Item->KeyReady = FALSE; + job->Item->DerivationResult = 0; + job->Item->Pkcs5Prf = prf; + job->Result = result; + job->HoldAfterReady = holdAfterReady; + job->Started = Event (); job->MayReturn = Event (); job->ReadyPublished = Event (); + job->MayReleaseReference = Event (); job->WorkerDone = Event (); + EncryptionThreadPoolBeginKeyDerivation (request->Ready, request->Drained, &job->Item->KeyReady, + request->References, prf, (unsigned char *) job, sizeof (*job), NULL, 1, 8, + job->Item->DerivedKey, &job->Item->DerivationResult, &request->Abort); + return job; +} + +static void JobDone (Job *job) +{ + Wait (job->WorkerDone); + Require (job->Item->KeyReady == TRUE && job->Item->DerivationResult == job->Result, + "result was not published before completion"); +} + +static void FakeKdf (const unsigned char *password, unsigned char *output, int length) +{ + Job *job = (Job *) password; + CurrentJob = job; + SetEvent (job->Started); + Wait (job->MayReturn); + memset (output, job->Owner->Abort || job->Result ? 0 : 0x5a, length); +} + +#define FAKE_KDF(name) \ +void name (const unsigned char *pwd, int pwd_len, const unsigned char *salt, int salt_len, \ + uint32 iterations, unsigned char *dk, int dklen, long volatile *abortFlag) \ +{ FakeKdf (pwd, dk, dklen); } +FAKE_KDF (derive_key_blake2s) +FAKE_KDF (derive_key_sha512) +FAKE_KDF (derive_key_whirlpool) +FAKE_KDF (derive_key_sha256) +FAKE_KDF (derive_key_streebog) +int derive_key_argon2 (const unsigned char *pwd, int pwd_len, const unsigned char *salt, int salt_len, + uint32 iterations, uint32 memory, unsigned char *dk, int dklen, long volatile *abortFlag) +{ FakeKdf (pwd, dk, dklen); return ((Job *) pwd)->Result; } + +int GetMaxPkcs5OutSize (void) { return MASTER_KEYDATA_SIZE; } +BOOL IsRamEncryptionEnabled (void) { return FALSE; } +uint64 VcGetEncryptionID (PCRYPTO_INFO ci) { return 0; } +void VcUnprotectKeys (PCRYPTO_INFO ci, uint64 id) { } +void EncryptDataUnitsCurrentThread (unsigned __int8 *data, const UINT64_STRUCT *start, TC_LARGEST_COMPILER_UINT count, PCRYPTO_INFO ci) { } +void DecryptDataUnitsCurrentThread (unsigned __int8 *data, const UINT64_STRUCT *start, TC_LARGEST_COMPILER_UINT count, PCRYPTO_INFO ci) { } + +static void RequesterLast (int count) +{ + Request *request = NewRequest (); + int i; + for (i = 0; i < count; ++i) + { + Job *job = Submit (request, FIRST_PRF_ID + i, FALSE, 0); + SetEvent (job->MayReturn); + JobDone (job); + Require (References (request) == 1 && request->Signals == 0, "idle submission phase completed"); + Require (WaitForSingleObject (*request->Drained, 0) == WAIT_TIMEOUT, "stale idle notification"); + } + FinishRequest (request); + Wait (request->FinishDone); + Require (request->Signals == 0, "requester-last path unnecessarily signaled"); +} + +static void WorkerLast (int count, BOOL cancel, int argonResult) +{ + Request *request = NewRequest (); + int i; + for (i = 0; i < count; ++i) + Submit (request, i == count - 1 ? ARGON2 : SHA512, FALSE, i == count - 1 ? argonResult : 0); + BeginFinish (request); + InterlockedExchange (&request->Abort, cancel); + for (i = 0; i < count; ++i) + SetEvent (request->Jobs[i].MayReturn); + EndFinish (request); + for (i = 0; i < count; ++i) + JobDone (&request->Jobs[i]); +} + +static void PublicationBoundaries (void) +{ + Request *request = NewRequest (); + Job *job = Submit (request, SHA256, TRUE, 0); + SetEvent (job->MayReturn); + Wait (job->ReadyPublished); + BeginFinish (request); + Require (References (request) == 1, "ready result released reference too early"); + SetEvent (job->MayReleaseReference); + EndFinish (request); + JobDone (job); + + request = NewRequest (); + request->HoldFinalSignal = TRUE; + job = Submit (request, SHA512, FALSE, 0); + BeginFinish (request); + SetEvent (job->MayReturn); + Wait (request->SignalEntered); + Require (References (request) == 0, "final notification boundary"); + Require (WaitForSingleObject (request->FinishDone, 0) == WAIT_TIMEOUT, "zero count bypassed final notification"); + SetEvent (request->MaySignal); + EndFinish (request); + JobDone (job); +} + +int main (int argc, char **argv) +{ + WORD groups; + size_t available = GetCpuCount (&groups); + size_t wanted = argc > 1 ? (size_t) atoi (argv[1]) : 2; + int i, j; + if (available < 2) { puts ("SKIP: at least two logical CPUs required"); return 0; } + if (wanted > available) wanted = available; + Require (EncryptionThreadPoolStart (available - wanted), "start production pool"); + Require (GetEncryptionThreadCount () == wanted, "pool worker count"); + RequesterLast (0); RequesterLast (1); RequesterLast (5); RequesterLast (6); + WorkerLast (2, FALSE, 0); WorkerLast (6, FALSE, 0); + WorkerLast (2, TRUE, ARGON2_OPERATION_CANCELLED); WorkerLast (6, TRUE, ARGON2_OPERATION_CANCELLED); + WorkerLast (2, FALSE, ARGON2_MEMORY_ALLOCATION_ERROR); + PublicationBoundaries (); + + /* Join every worker before releasing any request-owned allocation or event. */ + EncryptionThreadPoolStop (); + for (i = 0; i < RequestCount; ++i) + { + Request *r = &Requests[i]; + for (j = 0; j < r->JobCount; ++j) + { + Job *job = &r->Jobs[j]; + CloseHandle (job->Started); CloseHandle (job->MayReturn); CloseHandle (job->ReadyPublished); + CloseHandle (job->MayReleaseReference); CloseHandle (job->WorkerDone); + } + if (r->RequesterThread) CloseHandle (r->RequesterThread); + CloseHandle (r->WaitEntered); CloseHandle (r->FinishDone); CloseHandle (r->SignalEntered); CloseHandle (r->MaySignal); + CloseHandle (*r->Ready); CloseHandle (*r->Drained); + VirtualUnlock (r->Items, sizeof (*r->Items) * PRF_COUNT); + free (r->Ready); free (r->Drained); free (r->References); free (r->Items); + } + printf ("PASS: %d KDF lifetime scenarios with %u workers\n", RequestCount, (unsigned) wanted); + return 0; +} diff --git a/src/Common/EncryptionThreadPool.c b/src/Common/EncryptionThreadPool.c index 50e33c7f..51f34cbf 100644 --- a/src/Common/EncryptionThreadPool.c +++ b/src/Common/EncryptionThreadPool.c @@ -291,6 +291,7 @@ static TC_THREAD_PROC EncryptionThreadProc (void *threadArg) InterlockedExchange (workItem->KeyDerivation.CompletionFlag, TRUE); TC_SET_EVENT (*workItem->KeyDerivation.CompletionEvent); + // The requester retains a reference until it has finished submitting work. if (InterlockedDecrement (workItem->KeyDerivation.OutstandingWorkItemCount) == 0) TC_SET_EVENT (*workItem->KeyDerivation.NoOutstandingWorkItemEvent); @@ -581,7 +582,6 @@ void EncryptionThreadPoolBeginKeyDerivation (TC_EVENT *completionEvent, TC_EVENT workItem->KeyDerivation.pAbortKeyDerivation = pAbortKeyDerivation; InterlockedIncrement (outstandingWorkItemCount); - TC_CLEAR_EVENT (*noOutstandingWorkItemEvent); SetWorkItemState (workItem, WorkItemReady); TC_SET_EVENT (WorkItemReadyEvent); diff --git a/src/Common/EncryptionThreadPool.h b/src/Common/EncryptionThreadPool.h index d9ed44a5..50b4031b 100644 --- a/src/Common/EncryptionThreadPool.h +++ b/src/Common/EncryptionThreadPool.h @@ -32,7 +32,15 @@ typedef enum size_t GetCpuCount (WORD* pGroupCount); #endif -/* derivationResult receives 0 on success; PRFs that can fail store their native nonzero result. */ +/* + * The caller initializes outstandingWorkItemCount to 1 and noOutstandingWorkItemEvent + * to non-signaled, retaining its own reference while submissions are possible. + * After the last possible submission, it decrements the count and waits on the + * event if the result is nonzero, before releasing any caller-owned buffers. + * Each worker releases its reference after publishing completion. A worker that + * drops the count to zero signals the event, which must not be reset during submission. + * derivationResult receives 0 on success; PRFs that can fail store their native nonzero result. + */ void EncryptionThreadPoolBeginKeyDerivation (TC_EVENT *completionEvent, TC_EVENT *noOutstandingWorkItemEvent, LONG *completionFlag, LONG *outstandingWorkItemCount, int pkcs5Prf, unsigned char *password, int passwordLength, unsigned char *salt, int iterationCount, int memoryCost, unsigned char *derivedKey, LONG *derivationResult, LONG volatile *pAbortKeyDerivation); void EncryptionThreadPoolBeginReadVolumeHeaderFinalization (TC_EVENT *keyDerivationCompletedEvent, TC_EVENT *noOutstandingWorkItemEvent, LONG* outstandingWorkItemCount, void* keyInfoBuffer, int keyInfoBufferSize, void* keyDerivationWorkItems, int keyDerivationWorkItemsSize); void EncryptionThreadPoolDoWork (EncryptionThreadPoolWorkType type, uint8 *data, const UINT64_STRUCT *startUnitNo, uint32 unitCount, PCRYPTO_INFO cryptoInfo); diff --git a/src/Common/Volumes.c b/src/Common/Volumes.c index 2f4c3f79..90c52996 100644 --- a/src/Common/Volumes.c +++ b/src/Common/Volumes.c @@ -283,10 +283,11 @@ int ReadVolumeHeaderWithAbort (BOOL bBoot, unsigned char *encryptedHeader, Passw for (i = 0; i < pkcs5PrfCount; ++i) keyDerivationWorkItems[i].Free = TRUE; - *outstandingWorkItemCount = 0; + // Retain a requester reference until no further KDFs can be submitted. + *outstandingWorkItemCount = 1; #ifdef DEVICE_DRIVER KeInitializeEvent (keyDerivationCompletedEvent, SynchronizationEvent, FALSE); - KeInitializeEvent (noOutstandingWorkItemEvent, SynchronizationEvent, TRUE); + KeInitializeEvent (noOutstandingWorkItemEvent, SynchronizationEvent, FALSE); #else *keyDerivationCompletedEvent = CreateEvent (NULL, FALSE, FALSE, NULL); if (!*keyDerivationCompletedEvent) @@ -298,7 +299,7 @@ int ReadVolumeHeaderWithAbort (BOOL bBoot, unsigned char *encryptedHeader, Passw return ERR_OUTOFMEMORY; } - *noOutstandingWorkItemEvent = CreateEvent (NULL, FALSE, TRUE, NULL); + *noOutstandingWorkItemEvent = CreateEvent (NULL, FALSE, FALSE, NULL); if (!*noOutstandingWorkItemEvent) { CloseHandle (*keyDerivationCompletedEvent); @@ -736,10 +737,10 @@ ret: #if !defined(_UEFI) if ((selected_pkcs5_prf == 0) && (encryptionThreadCount > 1)) { - // Wait for all outstanding threads to finish or cancel - TC_WAIT_EVENT(*noOutstandingWorkItemEvent); - // Cleanup is now synchronous because we already waited for all threads to stop. - // The asynchronous finalization is no longer needed. + // Release the requester reference after the last possible submission. + // If workers remain, the last one signals completion before cleanup. + if (InterlockedDecrement (outstandingWorkItemCount) != 0) + TC_WAIT_EVENT (*noOutstandingWorkItemEvent); #if !defined(DEVICE_DRIVER) CloseHandle(*keyDerivationCompletedEvent); CloseHandle(*noOutstandingWorkItemEvent);