Unix: complete pool work when exception capture fails

Record cloning, formatting and allocation failures without allocating, then
complete and drain the request before reporting an error. This prevents
worker loss and avoids treating incomplete cipher or KDF output as valid.

Test ordinary exceptions, capture failures and subsequent pool recovery.

Tests/encryption_thread_pool_exception_test.cpp

@@ -2,6 +2,7 @@
 #include <algorithm>
 #include <atomic>
 #include <chrono>
+#include <clocale>
 #include <condition_variable>
 #include <cstdlib>
 #include <cstring>
@@ -10,10 +11,23 @@
 #include <stdexcept>
 #include <vector>
 #include "Volume/EncryptionThreadPool.h"
+#include "Volume/Pkcs5Kdf.h"
 
 using namespace VeraCrypt;
 
 static std::atomic<int> Clones (0), LiveClones (0);
+static thread_local bool FailNextAllocation;
+
+void *operator new (size_t size)
+{
+    if (FailNextAllocation) { FailNextAllocation = false; throw std::bad_alloc(); }
+    void *memory = std::malloc (size ? size : 1);
+    if (!memory) throw std::bad_alloc();
+    return memory;
+}
+void operator delete (void *memory) noexcept { std::free (memory); }
+void *operator new[] (size_t size) { return ::operator new (size); }
+void operator delete[] (void *memory) noexcept { ::operator delete (memory); }
 
 static void Check (bool condition, const char *message)
 {
@@ -35,6 +49,17 @@ struct TrackedException : Exception
     bool Stored;
 };
 
+struct UncopyableException : Exception
+{
+    Exception *CloneNew () override { throw std::bad_alloc(); }
+};
+
+struct InvalidMessageException : std::exception
+{
+    // Invalid in the C locale: exercise the real exception formatter's failure.
+    const char *what () const noexcept override { return "\xff"; }
+};
+
 class PoolView : public EncryptionThreadPool
 {
 public:
@@ -44,7 +69,7 @@ public:
 class TestMode : public EncryptionMode
 {
 public:
-    enum Failure { None, Native, Standard, Unknown };
+    enum Failure { None, Native, Standard, Unknown, CloneFailure, FormattingFailure, AllocationFailure };
     TestMode (size_t fragments, Failure failure)
         : Expected (fragments), Arrived (0), Kind (failure) { }
 
@@ -62,6 +87,7 @@ public:
 private:
     void Process (uint8 *data, uint64 count, size_t sectorSize) const
     {
+        FailNextAllocation = false;
         {
             std::unique_lock<std::mutex> lock (Gate);
             ++Arrived;
@@ -74,6 +100,9 @@ private:
         case Native: throw TrackedException ();
         case Standard: throw std::runtime_error ("fragment failure");
         case Unknown: throw 7;
+        case CloneFailure: throw UncopyableException();
+        case FormattingFailure: throw InvalidMessageException();
+        case AllocationFailure: FailNextAllocation = true; throw 7;
         case None: std::memset (data, 0xa5, count * sectorSize); break;
         }
     }
@@ -85,8 +114,68 @@ private:
     mutable std::condition_variable Ready;
 };
 
+class TestKdf : public Pkcs5Kdf
+{
+public:
+    explicit TestKdf (TestMode &mode) : Mode (mode) { }
+    int DeriveKey (const BufferPtr &key, const VolumePassword &, const ConstBufferPtr &, int) const override
+    { Mode.EncryptSectorsCurrentThread (key.Get(), 0, 1, key.Size()); return 0; }
+    int DeriveKey (const BufferPtr &key, const VolumePassword &password, const ConstBufferPtr &salt, int iterations, long volatile *) const override
+    { return DeriveKey (key, password, salt, iterations); }
+    shared_ptr<VeraCrypt::Hash> GetHash () const override { return shared_ptr<VeraCrypt::Hash>(); }
+    int GetIterationCount (int) const override { return 1; }
+    std::wstring GetName () const override { return L"exception-test"; }
+    Pkcs5Kdf *Clone () const override { return new TestKdf (Mode); }
+private:
+    TestMode &Mode;
+};
+
+// The runner extracts the actual drain and exception checks from VolumeHeader.cpp.
+#include "kdf_exception_completion.inc"
+
+static int KdfExceptions (size_t width)
+{
+    for (int kind = TestMode::None; kind <= TestMode::AllocationFailure; ++kind)
+    {
+        TestMode mode (width, static_cast<TestMode::Failure> (kind));
+        shared_ptr<Pkcs5Kdf> kdf (new TestKdf (mode));
+        VolumePassword password;
+        uint8 salt[64] = {};
+        SyncEvent completed, noOutstandingWorkItemEvent;
+        SharedVal<size_t> outstanding (0);
+        long volatile abortKeyDerivation = 0;
+        bool workItemsDrained = false;
+        std::vector<shared_ptr<EncryptionThreadPool::KeyDerivationWorkItem>> items;
+        for (size_t i = 0; i < width; ++i)
+        {
+            shared_ptr<EncryptionThreadPool::KeyDerivationWorkItem> item (new EncryptionThreadPool::KeyDerivationWorkItem (kdf, 64));
+            items.push_back (item);
+            EncryptionThreadPool::BeginKeyDerivation (*item, password, 0, ConstBufferPtr (salt, sizeof salt), completed, noOutstandingWorkItemEvent, outstanding, &abortKeyDerivation);
+        }
+        while (!std::all_of (items.begin(), items.end(), [] (const shared_ptr<EncryptionThreadPool::KeyDerivationWorkItem> &item) { return item->Completed.Get(); }))
+            completed.Wait();
+        for (const auto &item : items)
+        {
+            int observed = TestMode::None;
+            try { HandleKdfFailure (item, noOutstandingWorkItemEvent, width, workItemsDrained, abortKeyDerivation); }
+            catch (const TrackedException &) { observed = TestMode::Native; }
+            catch (const ExternalException &) { observed = TestMode::Standard; }
+            catch (const UnknownException &) { observed = TestMode::Unknown; }
+            Check (observed == std::min (kind, static_cast<int> (TestMode::Unknown)), "KDF failure was lost or changed");
+            Check (item->ExceptionCaptureFailed == (kind >= TestMode::CloneFailure), "KDF capture failure was not recorded");
+            if (kind == TestMode::None)
+                for (size_t i = 0; i < item->DerivedKey.Size(); ++i) Check (item->DerivedKey[i] == 0xa5, "KDF recovery failed");
+        }
+        DrainKeyDerivationWorkItems (noOutstandingWorkItemEvent, width, workItemsDrained);
+        Check (outstanding.Get() == 0, "KDF drain returned early");
+        Check (kind == TestMode::None || abortKeyDerivation == 1, "KDF failure did not abort the header read");
+    }
+    return TestMode::AllocationFailure + 1;
+}
+
 int main ()
 {
+    std::setlocale (LC_CTYPE, "C");
     EncryptionThreadPool::Start ();
     if (!EncryptionThreadPool::IsRunning ())
     {
@@ -101,7 +190,7 @@ int main ()
         for (int direction = 0; direction < 2; ++direction)
         {
             auto type = direction ? EncryptionThreadPool::WorkType::DecryptDataUnits : EncryptionThreadPool::WorkType::EncryptDataUnits;
-            for (int kind = TestMode::Native; kind <= TestMode::Unknown; ++kind)
+            for (int kind = TestMode::Native; kind <= TestMode::AllocationFailure; ++kind)
             {
                 TestMode mode (width, static_cast<TestMode::Failure> (kind));
                 int observed = TestMode::None;
@@ -110,7 +199,7 @@ int main ()
                 catch (const TrackedException &) { observed = TestMode::Native; }
                 catch (const ExternalException &) { observed = TestMode::Standard; }
                 catch (const UnknownException &) { observed = TestMode::Unknown; }
-                Check (observed == kind, "requester received the wrong exception");
+                Check (observed == std::min (kind, static_cast<int> (TestMode::Unknown)), "requester received the wrong exception");
                 Check (LiveClones == 0, "stored exception leaked");
                 if (kind == TestMode::Native) Check (Clones == 1, "more than one fragment exception retained");
                 ++checks;
@@ -121,7 +210,9 @@ int main ()
             for (uint8 byte : data) Check (byte == 0xa5, "pool failed after exception propagation");
             ++checks;
         }
+        checks += KdfExceptions (width);
     }
     EncryptionThreadPool::Stop ();
-    std::cout << "PASS: " << checks << " fragment exception and recovery checks" << std::endl;
+    Check (LiveClones == 0, "stored KDF exception leaked");
+    std::cout << "PASS: " << checks << " fragment/KDF exception, capture-failure and recovery checks" << std::endl;
 }

Tests/test_encryption_thread_pool_exceptions.py

@@ -1,5 +1,5 @@
 #!/usr/bin/env python3
-"""Check concurrent fragment exceptions in the production Unix encryption pool.
+"""Check fragment/KDF exceptions and exception-capture failures in the Unix pool.
 
 Uses real pool, encryption-mode dispatch and platform synchronization code with
 an event-controlled test cipher. Requires Linux, Python 3 and a C++11 compiler.
@@ -16,6 +16,16 @@ import tempfile
 ROOT = Path(__file__).resolve().parents[1]
 
 
+def source_block(source, start):
+    first = source.index("{", start)
+    depth = 1
+    for end in range(first + 1, len(source)):
+        depth += (source[end] == "{") - (source[end] == "}")
+        if depth == 0:
+            return source[start:end + 1]
+    raise RuntimeError("Unbalanced production source block")
+
+
 def build_test(build, sanitize, test_source, extra_sources=(), extra_flags=()):
     build.mkdir(parents=True, exist_ok=True)
     sources = ["Volume/EncryptionThreadPool.cpp", "Volume/EncryptionMode.cpp"]
@@ -51,7 +61,19 @@ def main():
     if not sys.platform.startswith("linux"):
         parser.error("This build runner requires Linux")
     build = (args.build_dir or Path(tempfile.mkdtemp(prefix="vc-pool-exceptions-"))).resolve()
-    exe = build_test(build, args.sanitize, "encryption_thread_pool_exception_test.cpp")
+    build.mkdir(parents=True, exist_ok=True)
+    header = (ROOT / "src/Volume/VolumeHeader.cpp").read_text()
+    drain = source_block(header, header.index("static void DrainKeyDerivationWorkItems ("))
+    failure = source_block(header, header.index("if (keyDerivationWorkItem->ItemException.get()"))
+    (build / "kdf_exception_completion.inc").write_text(drain + "\n" + """
+static void HandleKdfFailure (
+    shared_ptr<EncryptionThreadPool::KeyDerivationWorkItem> keyDerivationWorkItem,
+    SyncEvent &noOutstandingWorkItemEvent, size_t enqueuedWorkItemCount,
+    bool &workItemsDrained, long volatile &abortKeyDerivation)
+{
+""" + failure + "\n}\n")
+    exe = build_test(build, args.sanitize, "encryption_thread_pool_exception_test.cpp",
+                     ("Volume/Pkcs5Kdf.cpp", "Volume/VolumePassword.cpp"), ["-I" + str(build)])
     result = subprocess.run([str(exe)], timeout=60)
     if result.returncode not in (0, 77):
         raise SystemExit(result.returncode)

src/Volume/EncryptionThreadPool.cpp

@@ -27,6 +27,21 @@
 
 namespace VeraCrypt
 {
+	// Called from a catch handler. Failure to copy/format an exception must not
+	// bypass completion or turn a failed operation into a successful one.
+	static bool CaptureWorkItemException (unique_ptr <Exception> &itemException) noexcept
+	{
+		try
+		{
+			try { throw; }
+			catch (Exception &e) { itemException.reset (e.CloneNew()); }
+			catch (exception &e) { itemException.reset (new ExternalException (SRC_POS, StringConverter::ToExceptionString (e))); }
+			catch (...) { itemException.reset (new UnknownException (SRC_POS)); }
+			return itemException.get() != nullptr;
+		}
+		catch (...) { return false; }
+	}
+
 	bool EncryptionThreadPool::AcquirePoolReference ()
 	{
 		uint32 references = PoolReferences.load();
@@ -45,7 +60,7 @@ namespace VeraCrypt
 	}
 
 	EncryptionThreadPool::KeyDerivationWorkItem::KeyDerivationWorkItem (shared_ptr <Pkcs5Kdf> kdf, size_t derivedKeySize)
-		: Completed (false), DerivedKey (derivedKeySize), Kdf (kdf), Processed (false), Result (0)
+		: Completed (false), DerivedKey (derivedKeySize), ExceptionCaptureFailed (false), Kdf (kdf), Processed (false), Result (0)
 	{
 	}
 
@@ -60,6 +75,7 @@ namespace VeraCrypt
 		{
 			keyDerivationWorkItem.Completed.Set (false);
 			keyDerivationWorkItem.ItemException.reset();
+			keyDerivationWorkItem.ExceptionCaptureFailed = false;
 			keyDerivationWorkItem.Processed = false;
 			keyDerivationWorkItem.Result = abortFlag && *abortFlag ? ERR_USER_ABORT
 				: keyDerivationWorkItem.Kdf->DeriveKey (keyDerivationWorkItem.DerivedKey, password, pim, salt, abortFlag);
@@ -86,6 +102,7 @@ namespace VeraCrypt
 
 		keyDerivationWorkItem.Completed.Set (false);
 		keyDerivationWorkItem.ItemException.reset();
+		keyDerivationWorkItem.ExceptionCaptureFailed = false;
 		keyDerivationWorkItem.Processed = false;
 		keyDerivationWorkItem.Result = 0;
 
@@ -177,6 +194,7 @@ namespace VeraCrypt
 
 			firstFragmentWorkItem->OutstandingFragmentCount.Set (fragmentCount);
 			firstFragmentWorkItem->ItemException.reset();
+			firstFragmentWorkItem->ExceptionCaptureFailed = false;
 
 			while (fragmentCount-- > 0)
 			{
@@ -215,12 +233,15 @@ namespace VeraCrypt
 		unique_ptr <Exception> itemException;
 		if (firstFragmentWorkItem->ItemException.get())
 			itemException = move_ptr(firstFragmentWorkItem->ItemException);
+		bool exceptionCaptureFailed = firstFragmentWorkItem->ExceptionCaptureFailed.load();
 
 		firstFragmentWorkItem->State.Set (WorkItem::State::Free);
 		WorkItemCompletedEvent.Signal();
 
 		if (itemException.get())
 			itemException->Throw();
+		if (exceptionCaptureFailed)
+			throw UnknownException (SRC_POS);
 	}
 
 	void EncryptionThreadPool::Start ()
@@ -395,37 +416,23 @@ namespace VeraCrypt
 						throw ParameterIncorrect (SRC_POS);
 					}
 				}
-				catch (Exception &e)
-				{
-					if (workItem->Type == WorkType::DeriveKey)
-						workItem->KeyDerivation.WorkItem->ItemException.reset (e.CloneNew());
-					else
-					{
-						ScopeLock exceptionLock (FragmentExceptionMutex);
-						if (!workItem->FirstFragment->ItemException.get())
-							workItem->FirstFragment->ItemException.reset (e.CloneNew());
-					}
-				}
-				catch (exception &e)
-				{
-					if (workItem->Type == WorkType::DeriveKey)
-						workItem->KeyDerivation.WorkItem->ItemException.reset (new ExternalException (SRC_POS, StringConverter::ToExceptionString (e)));
-					else
-					{
-						ScopeLock exceptionLock (FragmentExceptionMutex);
-						if (!workItem->FirstFragment->ItemException.get())
-							workItem->FirstFragment->ItemException.reset (new ExternalException (SRC_POS, StringConverter::ToExceptionString (e)));
-					}
-				}
 				catch (...)
 				{
 					if (workItem->Type == WorkType::DeriveKey)
-						workItem->KeyDerivation.WorkItem->ItemException.reset (new UnknownException (SRC_POS));
+					{
+						KeyDerivationWorkItem &item = *workItem->KeyDerivation.WorkItem;
+						item.ExceptionCaptureFailed = !CaptureWorkItemException (item.ItemException);
+					}
 					else
 					{
-						ScopeLock exceptionLock (FragmentExceptionMutex);
-						if (!workItem->FirstFragment->ItemException.get())
-							workItem->FirstFragment->ItemException.reset (new UnknownException (SRC_POS));
+						WorkItem &firstFragment = *workItem->FirstFragment;
+						try
+						{
+							ScopeLock exceptionLock (FragmentExceptionMutex);
+							if (!firstFragment.ItemException.get() && !firstFragment.ExceptionCaptureFailed)
+								firstFragment.ExceptionCaptureFailed = !CaptureWorkItemException (firstFragment.ItemException);
+						}
+						catch (...) { firstFragment.ExceptionCaptureFailed = true; }
 					}
 				}
 

src/Volume/EncryptionThreadPool.h

@@ -51,6 +51,7 @@ namespace VeraCrypt
 
 			struct WorkItem *FirstFragment;
 			unique_ptr <Exception> ItemException;
+			std::atomic <bool> ExceptionCaptureFailed;
 			SyncEvent ItemCompletedEvent;
 			SharedVal <size_t> OutstandingFragmentCount;
 			SharedVal <State::Enum> State;
@@ -90,6 +91,7 @@ namespace VeraCrypt
 			SharedVal <bool> Completed;
 			SecureBuffer DerivedKey;
 			unique_ptr <Exception> ItemException;
+			bool ExceptionCaptureFailed;
 			shared_ptr <Pkcs5Kdf> Kdf;
 			bool Processed;
 			int Result;

src/Volume/VolumeHeader.cpp

@@ -144,12 +144,14 @@ namespace VeraCrypt
 							++processedWorkItemCount;
 							processed = true;
 
-							if (keyDerivationWorkItem->ItemException.get())
+							if (keyDerivationWorkItem->ItemException.get() || keyDerivationWorkItem->ExceptionCaptureFailed)
 							{
 								// KDF exceptions are fatal setup/runtime errors; candidate failures are reported via Result.
 								abortKeyDerivation = 1;
 								DrainKeyDerivationWorkItems (noOutstandingWorkItemEvent, enqueuedWorkItemCount, workItemsDrained);
-								keyDerivationWorkItem->ItemException->Throw();
+								if (keyDerivationWorkItem->ItemException.get())
+									keyDerivationWorkItem->ItemException->Throw();
+								throw UnknownException (SRC_POS);
 							}
 
 							if (keyDerivationWorkItem->Result != 0)