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

@@ -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()

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 <stdio.h>
+
+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;
+}

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);

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);

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);