From ca760c570e1104f76d378d4c0102308ba1e59788 Mon Sep 17 00:00:00 2001 From: BAder82t <41265463+BAder82t@users.noreply.github.com> Date: Mon, 20 Apr 2026 14:03:32 +0200 Subject: [PATCH 1/4] feat: expose ClearBootstrapPrecom to release cached bootstrap data (#533) FHECKKSRNS::m_bootPrecomMap is populated by EvalBootstrapSetup and EvalBootstrapPrecompute but never released by the existing ClearEval*Keys or ReleaseAllContexts APIs when a user still holds a CryptoContext reference, so the documented "clear" paths leave that memory alive for the lifetime of the process. Add a virtual FHEBase::ClearBootstrapPrecom() (default no-op), override it in FHECKKSRNS to clear m_bootPrecomMap, and expose matching forwarders on SchemeBase and CryptoContextImpl so callers can free the precomputations without destroying the context. --- src/pke/include/cryptocontext.h | 13 +++++++++++++ src/pke/include/scheme/ckksrns/ckksrns-fhe.h | 4 ++++ src/pke/include/schemebase/base-fhe.h | 9 +++++++++ src/pke/include/schemebase/base-scheme.h | 5 +++++ 4 files changed, 31 insertions(+) diff --git a/src/pke/include/cryptocontext.h b/src/pke/include/cryptocontext.h index 28dfd641c..65e517b50 100644 --- a/src/pke/include/cryptocontext.h +++ b/src/pke/include/cryptocontext.h @@ -620,6 +620,19 @@ class CryptoContextImpl : public Serializable { */ static void ClearStaticMapsAndVectors(); + /** + * @brief Release the scheme-level bootstrap precomputations held by this + * CryptoContext (e.g. FHECKKSRNS::m_bootPrecomMap). Lets callers free + * this memory without destroying the context itself, since the + * existing ClearEval*Keys / ReleaseAllContexts APIs leave scheme-level + * caches alive while any user reference keeps the context alive + * (issue #533). No-op for schemes without bootstrap caches. + */ + void ClearBootstrapPrecom() { + if (m_scheme) + m_scheme->ClearBootstrapPrecom(); + } + /** * @brief Serializes either all EvalMult keys (if keyTag is empty) or the EvalMult keys for keyTag * diff --git a/src/pke/include/scheme/ckksrns/ckksrns-fhe.h b/src/pke/include/scheme/ckksrns/ckksrns-fhe.h index b610e3073..de337bf7c 100644 --- a/src/pke/include/scheme/ckksrns/ckksrns-fhe.h +++ b/src/pke/include/scheme/ckksrns/ckksrns-fhe.h @@ -139,6 +139,10 @@ class FHECKKSRNS : public FHERNS { public: virtual ~FHECKKSRNS() = default; + void ClearBootstrapPrecom() override { + m_bootPrecomMap.clear(); + } + //------------------------------------------------------------------------------ // Bootstrap Wrapper //------------------------------------------------------------------------------ diff --git a/src/pke/include/schemebase/base-fhe.h b/src/pke/include/schemebase/base-fhe.h index 0ac85dd8c..bde1b0e62 100644 --- a/src/pke/include/schemebase/base-fhe.h +++ b/src/pke/include/schemebase/base-fhe.h @@ -67,6 +67,15 @@ class FHEBase { public: virtual ~FHEBase() = default; + /** + * Release any scheme-level cached data produced by EvalBootstrapSetup / + * EvalBootstrapPrecompute (e.g. the FHECKKSRNS m_bootPrecomMap). The + * default is a no-op for schemes that do not cache bootstrap + * precomputations. Exposed so callers can free the memory held by a + * CryptoContext without destroying the context itself (issue #533). + */ + virtual void ClearBootstrapPrecom() {} + /** * Bootstrap functionality: * There are three methods that have to be called in this specific order: diff --git a/src/pke/include/schemebase/base-scheme.h b/src/pke/include/schemebase/base-scheme.h index c5725dc95..18d5c31ef 100644 --- a/src/pke/include/schemebase/base-scheme.h +++ b/src/pke/include/schemebase/base-scheme.h @@ -1147,6 +1147,11 @@ class SchemeBase { m_FHE->EvalBootstrapSetup(cc, levelBudget, dim1, slots, correctionFactor, precompute, BTSlotsEncoding); } + void ClearBootstrapPrecom() { + if (m_FHE) + m_FHE->ClearBootstrapPrecom(); + } + std::shared_ptr>> EvalBootstrapKeyGen(const PrivateKey privateKey, uint32_t slots) { VerifyFHEEnabled(__func__); From 8360daa6caa4c912dda44544f41cf741952070a0 Mon Sep 17 00:00:00 2001 From: BAder82t <41265463+BAder82t@users.noreply.github.com> Date: Wed, 22 Apr 2026 17:35:51 +0200 Subject: [PATCH 2/4] feat: release CKKS scheme caches and trim allocator (#533) Extends the initial ClearBootstrapPrecom work so callers can actually drop every piece of scheme-level memory a CKKS CryptoContext retains, without destroying the context itself. - Add per-slot ClearBootstrapPrecom(slots) and ClearBootstrapPrecomExcept(keep) on FHEBase / SchemeBase / CryptoContextImpl so callers can evict a single transient precomputation or keep only a hot entry. - Add HasBootstrapPrecom() inspectors (whole-map and per-slot) for tests and memory-diagnostic code. - Introduce FHEBase::ClearSchemeSwitchPrecom() overridden by SWITCHCKKSRNS to drop the intermediate CKKS/BinFHE contexts, CKKS<->FHEW switching keys, auxiliary ciphertext, and the precomputed decoding-matrix plaintexts that EvalCKKStoFHEWSetup / EvalSchemeSwitchingSetup populate and which the existing Clear*Keys paths never released. - Add CryptoContextFactory::ReleaseAllContextsAndTrim() which extends ReleaseAllContexts() by also invoking the scheme-level clears on every tracked context and asking the allocator to return pages to the OS via the new TrimAllocator() helper (malloc_trim on glibc, no-op elsewhere). - Mark the cache-clear APIs noexcept so they are safe to call from destructors and shutdown paths. Adds UnitTestCKKSCacheClear (15 tests) covering per-slot / keep-except / full-release semantics, scheme-switch precom release, inspector agreement, and the ReleaseAllContextsAndTrim path. Adds ckks-release-memory example demonstrating the three common usage patterns. Closes #533 --- src/core/include/utils/memory.h | 21 + src/core/lib/utils/memory.cpp | 35 ++ src/pke/examples/ckks-release-memory.cpp | 121 +++++ src/pke/include/cryptocontext.h | 65 ++- src/pke/include/cryptocontextfactory.h | 60 +++ src/pke/include/scheme/ckksrns/ckksrns-fhe.h | 23 +- .../scheme/ckksrns/ckksrns-schemeswitching.h | 28 + src/pke/include/schemebase/base-fhe.h | 41 +- src/pke/include/schemebase/base-scheme.h | 29 +- src/pke/lib/cryptocontextfactory.cpp | 6 +- .../utckksrns/UnitTestCKKSCacheClear.cpp | 482 ++++++++++++++++++ 11 files changed, 906 insertions(+), 5 deletions(-) create mode 100644 src/pke/examples/ckks-release-memory.cpp create mode 100644 src/pke/unittest/utckksrns/UnitTestCKKSCacheClear.cpp diff --git a/src/core/include/utils/memory.h b/src/core/include/utils/memory.h index a4d637c8e..986fd7a25 100644 --- a/src/core/include/utils/memory.h +++ b/src/core/include/utils/memory.h @@ -63,6 +63,27 @@ void MoveAppend(std::vector& dst, std::vector& src) { */ void secure_memset(volatile void* mem, uint8_t c, size_t len); +/** + * @brief Ask the C allocator to return as many free pages to the OS as it can. + * + * Calling the cache-clear APIs (ClearBootstrapPrecom, ClearSchemeSwitchPrecom, + * ClearAllCKKSCaches, ReleaseAllContexts) frees memory back to the allocator, + * but whether that translates into RSS shrinking is allocator-dependent: + * glibc malloc keeps freed pages in its arena, mimalloc and jemalloc have + * their own purge policies. This helper calls the allocator's trim/purge + * primitive where available so the drop is visible at the OS level. + * + * Implementations: + * - glibc: malloc_trim(0) + * - Apple: malloc_zone_pressure_relief on every default zone + * - MSVC: _heapmin() + * - others: no-op + * + * Safe to call at any time. Returns true if the platform has a trim primitive, + * false if the call was a no-op on this build (issue #533). + */ +bool TrimAllocator(); + } // namespace lbcrypto #endif // LBCRYPTO_UTILS_MEMORY_H diff --git a/src/core/lib/utils/memory.cpp b/src/core/lib/utils/memory.cpp index 89b0083fd..57c663840 100644 --- a/src/core/lib/utils/memory.cpp +++ b/src/core/lib/utils/memory.cpp @@ -30,6 +30,15 @@ //================================================================================== #include "utils/memory.h" +#if defined(__APPLE__) + #include + #include +#elif defined(__GLIBC__) + #include +#elif defined(_MSC_VER) + #include +#endif + namespace lbcrypto { void secure_memset(volatile void* mem, uint8_t c, size_t len) { @@ -38,4 +47,30 @@ void secure_memset(volatile void* mem, uint8_t c, size_t len) { *(ptr + i) = c; } +bool TrimAllocator() { +#if defined(__GLIBC__) + // Returns non-zero if any memory was released; absence of trim is not an error. + malloc_trim(0); + return true; +#elif defined(__APPLE__) + // Walk every registered zone and pressure-relieve it. goal=0 asks libmalloc + // to return as much as it can without capping. + vm_address_t* zones = nullptr; + unsigned count = 0; + if (malloc_get_all_zones(mach_task_self(), nullptr, &zones, &count) == KERN_SUCCESS && zones) { + for (unsigned i = 0; i < count; ++i) { + auto* z = reinterpret_cast(zones[i]); + if (z) + malloc_zone_pressure_relief(z, 0); + } + } + return true; +#elif defined(_MSC_VER) + _heapmin(); + return true; +#else + return false; +#endif +} + } // namespace lbcrypto diff --git a/src/pke/examples/ckks-release-memory.cpp b/src/pke/examples/ckks-release-memory.cpp new file mode 100644 index 000000000..d40ecdf00 --- /dev/null +++ b/src/pke/examples/ckks-release-memory.cpp @@ -0,0 +1,121 @@ +//================================================================================== +// BSD 2-Clause License +// +// Copyright (c) 2014-2026, NJIT, Duality Technologies Inc. and other contributors +// +// All rights reserved. +// +// Author TPOC: contact@openfhe.org +// +// Redistribution and use in source and binary forms, with or without +// modification, are permitted provided that the following conditions are met: +// +// 1. Redistributions of source code must retain the above copyright notice, this +// list of conditions and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright notice, +// this list of conditions and the following disclaimer in the documentation +// and/or other materials provided with the distribution. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +// FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +// DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +// CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +// OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +//================================================================================== + +/* + Releasing CKKS scheme-level memory without destroying the CryptoContext + (issue #533). + + CKKS bootstrap caches (FHECKKSRNS::m_bootPrecomMap) and scheme-switch caches + (SWITCHCKKSRNS: ccLWE, ccKS, switching keys, decoding-matrix plaintexts) are + owned by the scheme, which lives as long as the CryptoContext. Before these + APIs, the only way to free them was to drop every CryptoContext handle plus + call CryptoContextFactory::ReleaseAllContexts() — if any caller kept a + context ref alive, the caches survived. + + This example demonstrates the three common patterns: + + 1. Clear one slot's bootstrap precom and reuse the context + 2. Clear all bootstrap precom except a hot slot count + 3. Full release + allocator trim at shutdown +*/ + +#include "openfhe.h" +#include "utils/memory.h" + +#include + +using namespace lbcrypto; + +static CryptoContext BuildBootstrapContext() { + CCParams parameters; + SecretKeyDist skDist = UNIFORM_TERNARY; + parameters.SetSecretKeyDist(skDist); + parameters.SetSecurityLevel(HEStd_NotSet); + parameters.SetRingDim(1 << 12); + +#if NATIVEINT == 128 + parameters.SetScalingModSize(78); + parameters.SetFirstModSize(89); + parameters.SetScalingTechnique(FIXEDAUTO); +#else + parameters.SetScalingModSize(59); + parameters.SetFirstModSize(60); + parameters.SetScalingTechnique(FLEXIBLEAUTO); +#endif + + std::vector levelBudget = {4, 4}; + uint32_t depth = 10 + FHECKKSRNS::GetBootstrapDepth(levelBudget, skDist); + parameters.SetMultiplicativeDepth(depth); + + auto cc = GenCryptoContext(parameters); + cc->Enable(PKE); + cc->Enable(KEYSWITCH); + cc->Enable(LEVELEDSHE); + cc->Enable(ADVANCEDSHE); + cc->Enable(FHE); + return cc; +} + +int main() { + auto cc = BuildBootstrapContext(); + uint32_t ringDim = cc->GetRingDimension(); + uint32_t hotSlots = ringDim / 2; + uint32_t tmpSlots = ringDim / 4; + + std::cout << "ring dim = " << ringDim << "\n"; + + // Pattern 1: transient slot count, then clear just that entry. + cc->EvalBootstrapSetup({4, 4}, {0, 0}, tmpSlots); + std::cout << "after transient setup: " + << "bootstrap cache present? " << std::boolalpha << cc->HasBootstrapPrecom(tmpSlots) << "\n"; + cc->ClearBootstrapPrecom(tmpSlots); + std::cout << "after per-slot clear: " << cc->HasBootstrapPrecom(tmpSlots) << "\n"; + + // Pattern 2: hot-slot workload with occasional transient bootstraps. + cc->EvalBootstrapSetup({4, 4}, {0, 0}, hotSlots); + cc->EvalBootstrapSetup({4, 4}, {0, 0}, tmpSlots); + std::cout << "with two entries, keep-except hot slots:\n"; + cc->ClearBootstrapPrecomExcept(hotSlots); + std::cout << " hot still present? " << cc->HasBootstrapPrecom(hotSlots) << "\n"; + std::cout << " tmp gone? " << !cc->HasBootstrapPrecom(tmpSlots) << "\n"; + + // Pattern 3: shutdown — release everything and ask the allocator to + // return pages to the OS. Useful right before a memory-sensitive phase + // of the surrounding application (e.g. forking a child process). + CryptoContextFactory::ReleaseAllContextsAndTrim(); + // After this call, static eval-key maps are empty, every context's + // scheme-level caches are released, and the allocator has been asked to + // trim. `cc` is still a valid handle but its caches are gone. + std::cout << "after ReleaseAllContextsAndTrim:\n"; + std::cout << " cc bootstrap caches? " << cc->HasBootstrapPrecom() << "\n"; + + return 0; +} diff --git a/src/pke/include/cryptocontext.h b/src/pke/include/cryptocontext.h index 65e517b50..a37b21374 100644 --- a/src/pke/include/cryptocontext.h +++ b/src/pke/include/cryptocontext.h @@ -628,11 +628,74 @@ class CryptoContextImpl : public Serializable { * caches alive while any user reference keeps the context alive * (issue #533). No-op for schemes without bootstrap caches. */ - void ClearBootstrapPrecom() { + void ClearBootstrapPrecom() noexcept { if (m_scheme) m_scheme->ClearBootstrapPrecom(); } + /** + * @brief Release the bootstrap precomputation entry associated with a + * specific number of slots. Useful for callers who re-use a CryptoContext + * with multiple slot counts and want to free only the stale ones + * (issue #533). No-op if no entry exists for the given slot count or the + * scheme does not cache per-slot bootstrap data. + */ + void ClearBootstrapPrecom(uint32_t slots) noexcept { + if (m_scheme) + m_scheme->ClearBootstrapPrecom(slots); + } + + /** + * @brief Release every bootstrap precomputation entry except the one keyed + * by @p keepSlots. Handy for workloads that repeatedly bootstrap a fixed + * slot count but occasionally bootstrap transient slot counts — reclaims + * the transient caches without invalidating the hot path (issue #533). + */ + void ClearBootstrapPrecomExcept(uint32_t keepSlots) noexcept { + if (m_scheme) + m_scheme->ClearBootstrapPrecomExcept(keepSlots); + } + + /** + * @brief Release scheme-switch precomputations (intermediate CKKS/FHEW + * contexts, switching keys, and linear-transform plaintexts) held by this + * CryptoContext (issue #533). No-op for schemes that do not support + * CKKS<->FHEW switching. + */ + void ClearSchemeSwitchPrecom() noexcept { + if (m_scheme) + m_scheme->ClearSchemeSwitchPrecom(); + } + + /** + * @brief Convenience: release all scheme-level caches held by this + * CryptoContext (bootstrap + scheme-switch). This does NOT touch the + * static eval-key maps — call ClearEvalMultKeys / ClearEvalAutomorphismKeys + * for those (issue #533). + */ + void ClearAllCKKSCaches() noexcept { + ClearBootstrapPrecom(); + ClearSchemeSwitchPrecom(); + } + + /** + * @brief Inspectors for scheme-level caches. Return false for schemes that + * do not maintain the corresponding cache, and whenever the cache is empty + * (issue #533). Intended for callers that need to reason about memory + * without reaching into internal types. + */ + bool HasBootstrapPrecom() const noexcept { + return m_scheme && m_scheme->HasBootstrapPrecom(); + } + + bool HasBootstrapPrecom(uint32_t slots) const noexcept { + return m_scheme && m_scheme->HasBootstrapPrecom(slots); + } + + bool HasSchemeSwitchPrecom() const noexcept { + return m_scheme && m_scheme->HasSchemeSwitchPrecom(); + } + /** * @brief Serializes either all EvalMult keys (if keyTag is empty) or the EvalMult keys for keyTag * diff --git a/src/pke/include/cryptocontextfactory.h b/src/pke/include/cryptocontextfactory.h index 11fbda613..57f25372f 100644 --- a/src/pke/include/cryptocontextfactory.h +++ b/src/pke/include/cryptocontextfactory.h @@ -35,6 +35,7 @@ #include "cryptocontext-fwd.h" #include "lattice/lat-hal.h" #include "scheme/scheme-id.h" +#include "utils/memory.h" #include #include @@ -55,6 +56,12 @@ class CryptoParametersBase; template class CryptoContextFactory { static std::vector> AllContexts; + // When AutoReleaseMode is true, newly built contexts are NOT tracked in + // AllContexts. Dropping the caller's last handle then runs the context + // destructor immediately, which releases every scheme-level cache and + // static eval-key ref without needing a manual ReleaseAllContexts call + // (issue #533). Default is false to preserve pre-fix semantics. + static bool s_autoReleaseMode; protected: static CryptoContext FindContext(std::shared_ptr> params, @@ -62,12 +69,62 @@ class CryptoContextFactory { static void AddContext(CryptoContext); public: + /** + * Enable or disable AutoRelease mode. When enabled, the factory no longer + * holds a strong reference to newly created contexts; the caller becomes + * the sole owner and cache memory is reclaimed automatically when they + * drop the last reference. Existing tracked contexts are released when + * the switch is flipped so the behavior is consistent (issue #533). + * + * Trade-off: FindContext / GetContext deduplication only inspects the + * tracked list, so two identical parameter sets created in AutoRelease + * mode will produce two distinct contexts instead of aliasing. Enable + * this only if you don't rely on factory-level deduplication. + */ + static void SetAutoReleaseMode(bool enabled) { + if (enabled && !AllContexts.empty()) { + // Flush tracked contexts so the semantics are consistent: from + // now on, no context is tracked by the factory. + ReleaseAllContexts(); + } + s_autoReleaseMode = enabled; + } + + static bool IsAutoReleaseMode() noexcept { + return s_autoReleaseMode; + } + + static void ReleaseAllContexts() { + // Drop scheme-level caches (bootstrap + scheme-switch) on every live + // context before clearing the static maps. This ensures memory held + // by FHECKKSRNS::m_bootPrecomMap and SWITCHCKKSRNS state is freed + // even when users still hold CryptoContext shared pointers after + // this call returns (issue #533). + for (auto& cc : AllContexts) { + if (cc) + cc->ClearAllCKKSCaches(); + } if (AllContexts.size() > 0) AllContexts[0]->ClearStaticMapsAndVectors(); AllContexts.clear(); } + /** + * Release every live context's scheme-level caches, clear the static key + * maps, and ask the C allocator to return freed pages to the OS. Behaves + * like ReleaseAllContexts() followed by TrimAllocator() (issue #533). + * + * Use this when heap usage drops after clearing but RSS doesn't — the C + * allocator is holding freed pages in its arena. On platforms without a + * trim primitive (non-glibc Linux, FreeBSD, WASM), this degrades into the + * plain ReleaseAllContexts() path. + */ + static void ReleaseAllContextsAndTrim() { + ReleaseAllContexts(); + TrimAllocator(); + } + static int GetContextCount() { return AllContexts.size(); } @@ -89,6 +146,9 @@ class CryptoContextFactory { template <> std::vector> CryptoContextFactory::AllContexts; +template <> +bool CryptoContextFactory::s_autoReleaseMode; + } // namespace lbcrypto #endif diff --git a/src/pke/include/scheme/ckksrns/ckksrns-fhe.h b/src/pke/include/scheme/ckksrns/ckksrns-fhe.h index de337bf7c..6126e8634 100644 --- a/src/pke/include/scheme/ckksrns/ckksrns-fhe.h +++ b/src/pke/include/scheme/ckksrns/ckksrns-fhe.h @@ -139,10 +139,31 @@ class FHECKKSRNS : public FHERNS { public: virtual ~FHECKKSRNS() = default; - void ClearBootstrapPrecom() override { + void ClearBootstrapPrecom() noexcept override { m_bootPrecomMap.clear(); } + void ClearBootstrapPrecom(uint32_t slots) noexcept override { + m_bootPrecomMap.erase(slots); + } + + void ClearBootstrapPrecomExcept(uint32_t keepSlots) noexcept override { + for (auto it = m_bootPrecomMap.begin(); it != m_bootPrecomMap.end();) { + if (it->first != keepSlots) + it = m_bootPrecomMap.erase(it); + else + ++it; + } + } + + bool HasBootstrapPrecom() const noexcept override { + return !m_bootPrecomMap.empty(); + } + + bool HasBootstrapPrecom(uint32_t slots) const noexcept override { + return m_bootPrecomMap.find(slots) != m_bootPrecomMap.end(); + } + //------------------------------------------------------------------------------ // Bootstrap Wrapper //------------------------------------------------------------------------------ diff --git a/src/pke/include/scheme/ckksrns/ckksrns-schemeswitching.h b/src/pke/include/scheme/ckksrns/ckksrns-schemeswitching.h index 46db09a18..d8ba23838 100644 --- a/src/pke/include/scheme/ckksrns/ckksrns-schemeswitching.h +++ b/src/pke/include/scheme/ckksrns/ckksrns-schemeswitching.h @@ -57,6 +57,34 @@ class SWITCHCKKSRNS : public FHERNS { public: virtual ~SWITCHCKKSRNS() = default; + /** + * Release cached scheme-switch state so it can be freed without destroying + * the CryptoContext (issue #533). Resets the intermediate CKKS and BinFHE + * contexts, the CKKS<->FHEW switching keys, the auxiliary ciphertext, and + * the precomputed decoding matrix. + * + * Only drops our own references. If the caller also retained a + * BinFHEContext handle via GetBinCCForSchemeSwitch(), their copy and its + * bootstrap keys are left untouched — memory is reclaimed only when the + * last owner releases it. + */ + void ClearSchemeSwitchPrecom() noexcept override { + // Swap with a temporary to guarantee capacity is released. shrink_to_fit + // is non-binding on some standard library implementations. + std::vector().swap(m_U0Pre); + m_ccLWE.reset(); + m_ccKS.reset(); + m_CKKStoFHEWswk.reset(); + m_FHEWtoCKKSswk.reset(); + m_ctxtKS.reset(); + } + + bool HasSchemeSwitchPrecom() const noexcept override { + return !m_U0Pre.empty() || static_cast(m_ccLWE) || static_cast(m_ccKS) || + static_cast(m_CKKStoFHEWswk) || static_cast(m_FHEWtoCKKSswk) || + static_cast(m_ctxtKS); + } + //------------------------------------------------------------------------------ // Scheme Switching Wrappers //------------------------------------------------------------------------------ diff --git a/src/pke/include/schemebase/base-fhe.h b/src/pke/include/schemebase/base-fhe.h index bde1b0e62..2c45a5479 100644 --- a/src/pke/include/schemebase/base-fhe.h +++ b/src/pke/include/schemebase/base-fhe.h @@ -74,7 +74,46 @@ class FHEBase { * precomputations. Exposed so callers can free the memory held by a * CryptoContext without destroying the context itself (issue #533). */ - virtual void ClearBootstrapPrecom() {} + virtual void ClearBootstrapPrecom() noexcept {} + + /** + * Release only the bootstrap precomputation entry associated with the given + * number of slots. No-op if no entry exists or the scheme does not cache + * per-slot bootstrap data (issue #533). + */ + virtual void ClearBootstrapPrecom(uint32_t slots) noexcept {} + + /** + * Release all bootstrap precomputation entries except the one keyed by + * @p keepSlots. Handy for callers that oscillate between a fixed slot count + * and transient ones — reclaims every transient cache without throwing away + * the hot entry (issue #533). + */ + virtual void ClearBootstrapPrecomExcept(uint32_t /*keepSlots*/) noexcept {} + + /** + * Release any scheme-switch cached data produced by the CKKS<->FHEW scheme + * switching setup (switching keys, intermediate CKKS/LWE contexts, + * precomputed linear-transform plaintexts). Default is a no-op; overridden + * by SWITCHCKKSRNS. Exposed so callers can free this memory without + * destroying the context (issue #533). + */ + virtual void ClearSchemeSwitchPrecom() noexcept {} + + /** + * Inspectors used by tests and memory-diagnostic code to check whether the + * scheme currently holds any cached data. Default: false. Overridden by + * FHECKKSRNS (bootstrap) and SWITCHCKKSRNS (scheme-switch). + */ + virtual bool HasBootstrapPrecom() const noexcept { + return false; + } + virtual bool HasBootstrapPrecom(uint32_t /*slots*/) const noexcept { + return false; + } + virtual bool HasSchemeSwitchPrecom() const noexcept { + return false; + } /** * Bootstrap functionality: diff --git a/src/pke/include/schemebase/base-scheme.h b/src/pke/include/schemebase/base-scheme.h index 18d5c31ef..b71fea1b6 100644 --- a/src/pke/include/schemebase/base-scheme.h +++ b/src/pke/include/schemebase/base-scheme.h @@ -1147,11 +1147,38 @@ class SchemeBase { m_FHE->EvalBootstrapSetup(cc, levelBudget, dim1, slots, correctionFactor, precompute, BTSlotsEncoding); } - void ClearBootstrapPrecom() { + void ClearBootstrapPrecom() noexcept { if (m_FHE) m_FHE->ClearBootstrapPrecom(); } + void ClearBootstrapPrecom(uint32_t slots) noexcept { + if (m_FHE) + m_FHE->ClearBootstrapPrecom(slots); + } + + void ClearBootstrapPrecomExcept(uint32_t keepSlots) noexcept { + if (m_FHE) + m_FHE->ClearBootstrapPrecomExcept(keepSlots); + } + + void ClearSchemeSwitchPrecom() noexcept { + if (m_SchemeSwitch) + m_SchemeSwitch->ClearSchemeSwitchPrecom(); + } + + bool HasBootstrapPrecom() const noexcept { + return m_FHE && m_FHE->HasBootstrapPrecom(); + } + + bool HasBootstrapPrecom(uint32_t slots) const noexcept { + return m_FHE && m_FHE->HasBootstrapPrecom(slots); + } + + bool HasSchemeSwitchPrecom() const noexcept { + return m_SchemeSwitch && m_SchemeSwitch->HasSchemeSwitchPrecom(); + } + std::shared_ptr>> EvalBootstrapKeyGen(const PrivateKey privateKey, uint32_t slots) { VerifyFHEEnabled(__func__); diff --git a/src/pke/lib/cryptocontextfactory.cpp b/src/pke/lib/cryptocontextfactory.cpp index 5e1cdb59a..bd3a4c7ca 100644 --- a/src/pke/lib/cryptocontextfactory.cpp +++ b/src/pke/lib/cryptocontextfactory.cpp @@ -38,6 +38,9 @@ namespace lbcrypto { template <> std::vector> CryptoContextFactory::AllContexts = {}; +template <> +bool CryptoContextFactory::s_autoReleaseMode = false; + template CryptoContext CryptoContextFactory::FindContext(std::shared_ptr> params, std::shared_ptr> scheme) { @@ -55,7 +58,8 @@ CryptoContext CryptoContextFactory::FindContext(std::shared_pt template void CryptoContextFactory::AddContext(CryptoContext cc) { - CryptoContextFactory::AllContexts.push_back(cc); + if (!CryptoContextFactory::s_autoReleaseMode) + CryptoContextFactory::AllContexts.push_back(cc); if (cc->GetEncodingParams()->GetPlaintextRootOfUnity() != 0) { PackedEncoding::SetParams(cc->GetCyclotomicOrder(), cc->GetEncodingParams()); diff --git a/src/pke/unittest/utckksrns/UnitTestCKKSCacheClear.cpp b/src/pke/unittest/utckksrns/UnitTestCKKSCacheClear.cpp new file mode 100644 index 000000000..98c6e6cf8 --- /dev/null +++ b/src/pke/unittest/utckksrns/UnitTestCKKSCacheClear.cpp @@ -0,0 +1,482 @@ +//================================================================================== +// BSD 2-Clause License +// +// Copyright (c) 2014-2026, NJIT, Duality Technologies Inc. and other contributors +// +// All rights reserved. +// +// Author TPOC: contact@openfhe.org +// +// Redistribution and use in source and binary forms, with or without +// modification, are permitted provided that the following conditions are met: +// +// 1. Redistributions of source code must retain the above copyright notice, this +// list of conditions and the following disclaimer. +// +// 2. Redistributions in binary form must reproduce the above copyright notice, +// this list of conditions and the following disclaimer in the documentation +// and/or other materials provided with the distribution. +// +// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +// FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +// DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +// CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +// OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +//================================================================================== + +// Coverage for the cache-clear APIs added for issue #533: +// CryptoContextImpl::ClearBootstrapPrecom() -- full clear +// CryptoContextImpl::ClearBootstrapPrecom(uint32_t) -- per-slot clear +// CryptoContextImpl::ClearSchemeSwitchPrecom() +// CryptoContextImpl::ClearAllCKKSCaches() +// CryptoContextFactory::ReleaseAllContexts() -- now also drops +// scheme-level caches + +#include "binfhecontext.h" +#include "openfhe.h" +#include "scheme/ckksrns/ckksrns-fhe.h" +#include "scheme/ckksrns/ckksrns-schemeswitching.h" +#include "scheme/scheme-swch-params.h" +#include "utils/memory.h" + +#include "gtest/gtest.h" + +#include + +#if defined(__GLIBC__) + #include +#elif defined(__APPLE__) + #include +#endif + +using namespace lbcrypto; + +namespace { + +CryptoContext MakeBootstrapCC(uint32_t ringDim = 1 << 8) { + CCParams params; + SecretKeyDist skDist = UNIFORM_TERNARY; + params.SetSecretKeyDist(skDist); + params.SetSecurityLevel(HEStd_NotSet); + params.SetRingDim(ringDim); +#if NATIVEINT == 128 + params.SetScalingModSize(78); + params.SetFirstModSize(89); + params.SetScalingTechnique(FIXEDAUTO); +#else + params.SetScalingModSize(59); + params.SetFirstModSize(60); + params.SetScalingTechnique(FLEXIBLEAUTO); +#endif + std::vector levelBudget = {1, 1}; + uint32_t depth = 2 + FHECKKSRNS::GetBootstrapDepth(levelBudget, skDist); + params.SetMultiplicativeDepth(depth); + + auto cc = GenCryptoContext(params); + cc->Enable(PKE); + cc->Enable(KEYSWITCH); + cc->Enable(LEVELEDSHE); + cc->Enable(ADVANCEDSHE); + cc->Enable(FHE); + return cc; +} + +// Returns bytes currently in use by the heap allocator, or 0 if no portable +// probe is available. Only used by MemoryDropsAfterBootstrapClear, which also +// guards itself by platform. The absolute value is noisy (other threads, +// background allocations, allocator bookkeeping); the test asserts a relative +// drop, not an exact number. +size_t HeapInUseBytes() { +#if defined(__GLIBC__) + auto info = mallinfo2(); + return info.uordblks; +#elif defined(__APPLE__) + malloc_statistics_t s{}; + malloc_zone_statistics(nullptr, &s); + return s.size_in_use; +#else + return 0; +#endif +} + +CryptoContext MakeBFVContext() { + CCParams params; + params.SetPlaintextModulus(65537); + params.SetMultiplicativeDepth(2); + auto cc = GenCryptoContext(params); + cc->Enable(PKE); + cc->Enable(KEYSWITCH); + cc->Enable(LEVELEDSHE); + return cc; +} + +CryptoContext MakeBGVContext() { + CCParams params; + params.SetPlaintextModulus(65537); + params.SetMultiplicativeDepth(2); + auto cc = GenCryptoContext(params); + cc->Enable(PKE); + cc->Enable(KEYSWITCH); + cc->Enable(LEVELEDSHE); + return cc; +} + +CryptoContext MakeSchemeSwitchCC() { + CCParams params; + params.SetMultiplicativeDepth(3); + params.SetFirstModSize(60); + params.SetScalingModSize(50); + params.SetScalingTechnique(FLEXIBLEAUTOEXT); + params.SetSecurityLevel(HEStd_NotSet); + params.SetRingDim(1 << 12); + params.SetBatchSize(16); + auto cc = GenCryptoContext(params); + cc->Enable(PKE); + cc->Enable(KEYSWITCH); + cc->Enable(LEVELEDSHE); + cc->Enable(SCHEMESWITCH); + return cc; +} + +} // namespace + +class UTCKKSCacheClear : public ::testing::Test { +protected: + void TearDown() override { + CryptoContextFactory::ReleaseAllContexts(); + } +}; + +// Baseline: a context with no EvalBootstrapSetup / scheme-switch setup reports +// empty caches and the clear APIs are safe no-ops. +TEST_F(UTCKKSCacheClear, DefaultContextReportsNoCaches) { + auto cc = MakeBootstrapCC(); + + EXPECT_FALSE(cc->HasBootstrapPrecom()); + EXPECT_FALSE(cc->HasBootstrapPrecom(8)); + EXPECT_FALSE(cc->HasSchemeSwitchPrecom()); + + EXPECT_NO_THROW(cc->ClearBootstrapPrecom()); + EXPECT_NO_THROW(cc->ClearBootstrapPrecom(8)); + EXPECT_NO_THROW(cc->ClearSchemeSwitchPrecom()); + EXPECT_NO_THROW(cc->ClearAllCKKSCaches()); +} + +// Full clear drops every slot-keyed entry in FHECKKSRNS::m_bootPrecomMap. +TEST_F(UTCKKSCacheClear, FullBootstrapClear) { + auto cc = MakeBootstrapCC(); + uint32_t numSlots = cc->GetRingDimension() / 2; + + cc->EvalBootstrapSetup({1, 1}, {0, 0}, numSlots); + EXPECT_TRUE(cc->HasBootstrapPrecom()); + EXPECT_TRUE(cc->HasBootstrapPrecom(numSlots)); + + cc->ClearBootstrapPrecom(); + EXPECT_FALSE(cc->HasBootstrapPrecom()); + EXPECT_FALSE(cc->HasBootstrapPrecom(numSlots)); +} + +// Per-slot clear removes only the targeted slot entry and leaves others alone. +TEST_F(UTCKKSCacheClear, PerSlotBootstrapClear) { + auto cc = MakeBootstrapCC(); + uint32_t ringDim = cc->GetRingDimension(); + uint32_t slotsLarge = ringDim / 2; + uint32_t slotsSmall = ringDim / 4; + + cc->EvalBootstrapSetup({1, 1}, {0, 0}, slotsLarge); + cc->EvalBootstrapSetup({1, 1}, {0, 0}, slotsSmall); + ASSERT_TRUE(cc->HasBootstrapPrecom(slotsLarge)); + ASSERT_TRUE(cc->HasBootstrapPrecom(slotsSmall)); + + cc->ClearBootstrapPrecom(slotsLarge); + EXPECT_FALSE(cc->HasBootstrapPrecom(slotsLarge)); + EXPECT_TRUE(cc->HasBootstrapPrecom(slotsSmall)); + EXPECT_TRUE(cc->HasBootstrapPrecom()); // map still non-empty + + // Clearing a non-existent slot is a safe no-op. + EXPECT_NO_THROW(cc->ClearBootstrapPrecom(slotsLarge)); + + cc->ClearBootstrapPrecom(slotsSmall); + EXPECT_FALSE(cc->HasBootstrapPrecom()); +} + +// The CKKS<->FHEW scheme-switch precomputation holds an entire BinFHEContext, +// intermediate CKKS context, two switching keys, a ciphertext and a precomputed +// linear-transform matrix. ClearSchemeSwitchPrecom() must release all of them. +TEST_F(UTCKKSCacheClear, SchemeSwitchPrecomClear) { + auto cc = MakeSchemeSwitchCC(); + auto kp = cc->KeyGen(); + + SchSwchParams p; + p.SetSecurityLevelCKKS(HEStd_NotSet); + p.SetSecurityLevelFHEW(TOY); + p.SetCtxtModSizeFHEWLargePrec(25); + p.SetNumSlotsCKKS(16); + auto lweSk = cc->EvalCKKStoFHEWSetup(p); + cc->EvalCKKStoFHEWKeyGen(kp, lweSk); + EXPECT_TRUE(cc->HasSchemeSwitchPrecom()); + + cc->ClearSchemeSwitchPrecom(); + EXPECT_FALSE(cc->HasSchemeSwitchPrecom()); + EXPECT_FALSE(static_cast(cc->GetScheme()->GetBinCCForSchemeSwitch())); +} + +// ClearAllCKKSCaches() combines both clears. Covers the convenience path +// ReleaseAllContexts() now relies on. +TEST_F(UTCKKSCacheClear, ClearAllCKKSCachesClearsBoth) { + auto cc = MakeSchemeSwitchCC(); + auto kp = cc->KeyGen(); + + // Scheme-switch state. + SchSwchParams p; + p.SetSecurityLevelCKKS(HEStd_NotSet); + p.SetSecurityLevelFHEW(TOY); + p.SetCtxtModSizeFHEWLargePrec(25); + p.SetNumSlotsCKKS(16); + auto lweSk = cc->EvalCKKStoFHEWSetup(p); + cc->EvalCKKStoFHEWKeyGen(kp, lweSk); + + // Bootstrap-like state cannot be produced on this smaller scheme-switch + // context without heavy parameters, so only exercise the scheme-switch arm + // here; the bootstrap arm is covered by the tests above. + ASSERT_TRUE(cc->HasSchemeSwitchPrecom()); + + cc->ClearAllCKKSCaches(); + EXPECT_FALSE(cc->HasSchemeSwitchPrecom()); + EXPECT_FALSE(cc->HasBootstrapPrecom()); +} + +// ReleaseAllContexts() must drop scheme-level caches even when the caller +// keeps its own shared_ptr to the CryptoContext. This is the original #533 +// regression: previously the static maps were cleared but the scheme-owned +// m_bootPrecomMap (and scheme-switch state) survived. +TEST_F(UTCKKSCacheClear, ReleaseAllContextsDropsSchemeCaches) { + auto cc = MakeBootstrapCC(); + uint32_t numSlots = cc->GetRingDimension() / 2; + + cc->EvalBootstrapSetup({1, 1}, {0, 0}, numSlots); + ASSERT_TRUE(cc->HasBootstrapPrecom(numSlots)); + + CryptoContextFactory::ReleaseAllContexts(); + + // The factory no longer owns the context, but this test still does. + EXPECT_FALSE(cc->HasBootstrapPrecom()); + EXPECT_FALSE(cc->HasBootstrapPrecom(numSlots)); +} + +TEST_F(UTCKKSCacheClear, ReleaseAllContextsDropsSchemeSwitchCaches) { + auto cc = MakeSchemeSwitchCC(); + auto kp = cc->KeyGen(); + + SchSwchParams p; + p.SetSecurityLevelCKKS(HEStd_NotSet); + p.SetSecurityLevelFHEW(TOY); + p.SetCtxtModSizeFHEWLargePrec(25); + p.SetNumSlotsCKKS(16); + auto lweSk = cc->EvalCKKStoFHEWSetup(p); + cc->EvalCKKStoFHEWKeyGen(kp, lweSk); + ASSERT_TRUE(cc->HasSchemeSwitchPrecom()); + + CryptoContextFactory::ReleaseAllContexts(); + + EXPECT_FALSE(cc->HasSchemeSwitchPrecom()); +} + +// BFV and BGV do not maintain scheme-level caches; the inspectors must report +// false and the clear APIs must be safe no-ops regardless of whether FHE is +// even enabled. This pins the default FHEBase path so a future BFV/BGV cache +// cannot silently go unreleased. +TEST_F(UTCKKSCacheClear, NonCKKSSchemesReportNoCachesAndClearsAreNoOps) { + auto bfv = MakeBFVContext(); + EXPECT_FALSE(bfv->HasBootstrapPrecom()); + EXPECT_FALSE(bfv->HasBootstrapPrecom(8)); + EXPECT_FALSE(bfv->HasSchemeSwitchPrecom()); + EXPECT_NO_THROW(bfv->ClearBootstrapPrecom()); + EXPECT_NO_THROW(bfv->ClearBootstrapPrecom(8)); + EXPECT_NO_THROW(bfv->ClearSchemeSwitchPrecom()); + EXPECT_NO_THROW(bfv->ClearAllCKKSCaches()); + + auto bgv = MakeBGVContext(); + EXPECT_FALSE(bgv->HasBootstrapPrecom()); + EXPECT_FALSE(bgv->HasBootstrapPrecom(8)); + EXPECT_FALSE(bgv->HasSchemeSwitchPrecom()); + EXPECT_NO_THROW(bgv->ClearBootstrapPrecom()); + EXPECT_NO_THROW(bgv->ClearBootstrapPrecom(8)); + EXPECT_NO_THROW(bgv->ClearSchemeSwitchPrecom()); + EXPECT_NO_THROW(bgv->ClearAllCKKSCaches()); + + // ReleaseAllContexts must tolerate mixed-scheme context sets. + EXPECT_NO_THROW(CryptoContextFactory::ReleaseAllContexts()); +} + +// Regression for the behavior-risk noted above: when a caller keeps their own +// BinFHEContext handle via GetBinCCForSchemeSwitch(), ClearSchemeSwitchPrecom +// must not clear the bootstrap keys on that external handle. It only drops +// our owning reference. +TEST_F(UTCKKSCacheClear, ClearSchemeSwitchDoesNotMutateExternalBinFHE) { + auto cc = MakeSchemeSwitchCC(); + auto kp = cc->KeyGen(); + + SchSwchParams p; + p.SetSecurityLevelCKKS(HEStd_NotSet); + p.SetSecurityLevelFHEW(TOY); + p.SetCtxtModSizeFHEWLargePrec(25); + p.SetNumSlotsCKKS(16); + auto lweSk = cc->EvalCKKStoFHEWSetup(p); + cc->EvalCKKStoFHEWKeyGen(kp, lweSk); + + auto externalCcLWE = cc->GetBinCCForSchemeSwitch(); + ASSERT_TRUE(static_cast(externalCcLWE)); + // Generate FHEW bootstrap keys on the external handle so we can verify + // they survive our cache clear. + externalCcLWE->BTKeyGen(lweSk); + + cc->ClearSchemeSwitchPrecom(); + + // Our reference is gone, but the caller's is still alive and usable. + EXPECT_FALSE(cc->HasSchemeSwitchPrecom()); + ASSERT_TRUE(static_cast(externalCcLWE)); + EXPECT_NO_THROW(externalCcLWE->GetParams()->GetLWEParams()->Getn()); +} + +// ClearBootstrapPrecomExcept keeps the specified slot entry and drops every +// other entry. Covers the "oscillating between two slot counts" workload: a +// long-running bootstrap at N slots plus transient smaller setups that should +// not accumulate. +TEST_F(UTCKKSCacheClear, ClearBootstrapPrecomExceptKeepsOneEntry) { + auto cc = MakeBootstrapCC(); + uint32_t ringDim = cc->GetRingDimension(); + uint32_t slotsKeep = ringDim / 2; + uint32_t slotsDropA = ringDim / 4; + uint32_t slotsDropB = ringDim / 8; + + cc->EvalBootstrapSetup({1, 1}, {0, 0}, slotsKeep); + cc->EvalBootstrapSetup({1, 1}, {0, 0}, slotsDropA); + cc->EvalBootstrapSetup({1, 1}, {0, 0}, slotsDropB); + ASSERT_TRUE(cc->HasBootstrapPrecom(slotsKeep)); + ASSERT_TRUE(cc->HasBootstrapPrecom(slotsDropA)); + ASSERT_TRUE(cc->HasBootstrapPrecom(slotsDropB)); + + cc->ClearBootstrapPrecomExcept(slotsKeep); + + EXPECT_TRUE(cc->HasBootstrapPrecom(slotsKeep)); + EXPECT_FALSE(cc->HasBootstrapPrecom(slotsDropA)); + EXPECT_FALSE(cc->HasBootstrapPrecom(slotsDropB)); + + // Keeping a slot that doesn't exist clears the whole map (there is + // nothing to keep). + cc->ClearBootstrapPrecomExcept(12345); + EXPECT_FALSE(cc->HasBootstrapPrecom()); +} + +// Regression probe for the core #533 contract: after a clear, the allocator's +// in-use bytes must actually shrink. The previous broken behavior (before the +// fix that started this branch) left m_bootPrecomMap alive, so the in-use +// size after clear equalled the peak. +// +// The delta is noisy — parallel threads, allocator bookkeeping, lazy frees — +// so we only assert direction, not magnitude, and only on platforms where we +// can probe the allocator portably. If a future refactor reintroduces a +// self-reference that keeps the precom alive, this test catches it. +#if defined(__GLIBC__) || defined(__APPLE__) +TEST_F(UTCKKSCacheClear, MemoryDropsAfterBootstrapClear) { + // Use a slightly larger ring dim than other tests so the per-slot + // precomputations are large enough to be visible above allocator noise. + auto cc = MakeBootstrapCC(1 << 10); + uint32_t numSlots = cc->GetRingDimension() / 2; + + size_t before = HeapInUseBytes(); + cc->EvalBootstrapSetup({1, 1}, {0, 0}, numSlots); + size_t peak = HeapInUseBytes(); + cc->ClearBootstrapPrecom(); + lbcrypto::TrimAllocator(); + size_t after = HeapInUseBytes(); + + ASSERT_GT(before, 0u) << "heap probe returned 0; platform guard misconfigured"; + EXPECT_GT(peak, before) << "EvalBootstrapSetup did not grow heap as expected"; + EXPECT_LT(after, peak) << "ClearBootstrapPrecom did not release memory"; +} +#endif + +// TrimAllocator must be safe to call at any time and return a platform +// capability flag. It's hard to assert RSS drops portably, so this just pins +// the contract: never throws, returns true on platforms with a trim +// primitive (glibc, Apple, MSVC), false elsewhere. +TEST_F(UTCKKSCacheClear, TrimAllocatorIsSafe) { + EXPECT_NO_THROW(lbcrypto::TrimAllocator()); + EXPECT_NO_THROW(lbcrypto::TrimAllocator()); // idempotent +#if defined(__GLIBC__) || defined(__APPLE__) || defined(_MSC_VER) + EXPECT_TRUE(lbcrypto::TrimAllocator()); +#endif +} + +// ReleaseAllContextsAndTrim is a convenience: does everything ReleaseAllContexts +// does, plus the allocator trim. The cache-clear behavior must match. +TEST_F(UTCKKSCacheClear, ReleaseAllContextsAndTrimClearsSchemeCaches) { + auto cc = MakeBootstrapCC(); + uint32_t numSlots = cc->GetRingDimension() / 2; + + cc->EvalBootstrapSetup({1, 1}, {0, 0}, numSlots); + ASSERT_TRUE(cc->HasBootstrapPrecom(numSlots)); + + CryptoContextFactory::ReleaseAllContextsAndTrim(); + + EXPECT_FALSE(cc->HasBootstrapPrecom()); + EXPECT_FALSE(cc->HasBootstrapPrecom(numSlots)); +} + +// Opt-in AutoRelease mode: the factory does not track newly created contexts, +// so the caller's last handle drop destructs the context immediately and +// every scheme-level cache is released with no manual ReleaseAllContexts +// call. Verifies the flag toggles, existing contexts are flushed on enable, +// and GetContextCount() reflects the untracked state (issue #533). +TEST_F(UTCKKSCacheClear, AutoReleaseModeFactoryDoesNotTrack) { + ASSERT_FALSE(CryptoContextFactory::IsAutoReleaseMode()); + auto tracked = MakeBootstrapCC(); + EXPECT_GE(CryptoContextFactory::GetContextCount(), 1); + + CryptoContextFactory::SetAutoReleaseMode(true); + EXPECT_TRUE(CryptoContextFactory::IsAutoReleaseMode()); + // Flipping the flag should have flushed the tracked list. + EXPECT_EQ(CryptoContextFactory::GetContextCount(), 0); + + // Building a new context under AutoRelease must not register it. + int before = CryptoContextFactory::GetContextCount(); + { + auto ephemeral = MakeBootstrapCC(); + uint32_t numSlots = ephemeral->GetRingDimension() / 2; + ephemeral->EvalBootstrapSetup({1, 1}, {0, 0}, numSlots); + ASSERT_TRUE(ephemeral->HasBootstrapPrecom(numSlots)); + EXPECT_EQ(CryptoContextFactory::GetContextCount(), before); + } + // ephemeral is out of scope; its caches should have been auto-released + // by ~CryptoContextImpl(). We can't inspect the ex-ephemeral directly, + // but GetContextCount() must still be 0 — nothing leaked into the + // factory list. + EXPECT_EQ(CryptoContextFactory::GetContextCount(), 0); + + // Reset for subsequent tests. + CryptoContextFactory::SetAutoReleaseMode(false); +} + +// The scheme's internal FHE object matches the behavior exposed on the +// CryptoContext facade. This pins the dispatch path used by callers who work +// directly with GetScheme(). +TEST_F(UTCKKSCacheClear, FHEInspectorsAgreeWithFacade) { + auto cc = MakeBootstrapCC(); + uint32_t numSlots = cc->GetRingDimension() / 2; + + cc->EvalBootstrapSetup({1, 1}, {0, 0}, numSlots); + auto scheme = cc->GetScheme(); + ASSERT_TRUE(scheme->HasBootstrapPrecom()); + ASSERT_TRUE(scheme->HasBootstrapPrecom(numSlots)); + + scheme->ClearBootstrapPrecom(); + EXPECT_FALSE(cc->HasBootstrapPrecom()); + EXPECT_FALSE(scheme->HasBootstrapPrecom(numSlots)); +} From dab669e10cc51e97e22e5d779f05d98cb1237a61 Mon Sep 17 00:00:00 2001 From: BAder82t <41265463+BAder82t@users.noreply.github.com> Date: Thu, 30 Apr 2026 17:44:30 +0300 Subject: [PATCH 3/4] test: use 128-bit-compatible params in scheme-switch helper The test helper used modulus sizes and a scaling style that only work on 32/64-bit builds, so it threw on the 128-bit CI job. Pick larger moduli and FIXEDAUTO scaling when NATIVEINT == 128, matching what UnitTestSchemeSwitch.cpp already does. Production code is unchanged. --- src/pke/unittest/utckksrns/UnitTestCKKSCacheClear.cpp | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/src/pke/unittest/utckksrns/UnitTestCKKSCacheClear.cpp b/src/pke/unittest/utckksrns/UnitTestCKKSCacheClear.cpp index 98c6e6cf8..231139227 100644 --- a/src/pke/unittest/utckksrns/UnitTestCKKSCacheClear.cpp +++ b/src/pke/unittest/utckksrns/UnitTestCKKSCacheClear.cpp @@ -129,9 +129,17 @@ CryptoContext MakeBGVContext() { CryptoContext MakeSchemeSwitchCC() { CCParams params; params.SetMultiplicativeDepth(3); +#if NATIVEINT == 128 + // FLEXIBLE* scaling is unsupported on the 128-bit backend; use the + // FIXED* path with the larger moduli that UnitTestSchemeSwitch uses. + params.SetFirstModSize(89); + params.SetScalingModSize(78); + params.SetScalingTechnique(FIXEDAUTO); +#else params.SetFirstModSize(60); params.SetScalingModSize(50); params.SetScalingTechnique(FLEXIBLEAUTOEXT); +#endif params.SetSecurityLevel(HEStd_NotSet); params.SetRingDim(1 << 12); params.SetBatchSize(16); From 61df8bc6795704c7b6130f19e4f42de37345d28c Mon Sep 17 00:00:00 2001 From: BAder82t <41265463+BAder82t@users.noreply.github.com> Date: Fri, 1 May 2026 11:22:19 +0300 Subject: [PATCH 4/4] test: skip MemoryDropsAfterBootstrapClear when heap probe is unavailable When the build links an alternative allocator like TCMalloc (WITH_TCM=ON), mallinfo2() and malloc_zone_statistics() return 0 because they only see the system malloc that TCMalloc replaced. The relative-drop check is unmeasurable in that case, so call GTEST_SKIP() instead of failing the assertion. The TrimAllocator contract test (TrimAllocatorIsSafe) still runs and pins the no-throw + capability-flag behavior on every backend. --- src/pke/unittest/utckksrns/UnitTestCKKSCacheClear.cpp | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/src/pke/unittest/utckksrns/UnitTestCKKSCacheClear.cpp b/src/pke/unittest/utckksrns/UnitTestCKKSCacheClear.cpp index 231139227..9a30e26ac 100644 --- a/src/pke/unittest/utckksrns/UnitTestCKKSCacheClear.cpp +++ b/src/pke/unittest/utckksrns/UnitTestCKKSCacheClear.cpp @@ -399,13 +399,19 @@ TEST_F(UTCKKSCacheClear, MemoryDropsAfterBootstrapClear) { uint32_t numSlots = cc->GetRingDimension() / 2; size_t before = HeapInUseBytes(); + if (before == 0) { + // mallinfo2() / malloc_zone_statistics return 0 when an alternative + // allocator (e.g. TCMalloc via WITH_TCM=ON) replaces the system malloc. + // The heap probe cannot see through it, so the relative-drop assertion + // is unmeasurable; skip rather than fail. + GTEST_SKIP() << "heap probe unavailable on this allocator"; + } cc->EvalBootstrapSetup({1, 1}, {0, 0}, numSlots); size_t peak = HeapInUseBytes(); cc->ClearBootstrapPrecom(); lbcrypto::TrimAllocator(); size_t after = HeapInUseBytes(); - ASSERT_GT(before, 0u) << "heap probe returned 0; platform guard misconfigured"; EXPECT_GT(peak, before) << "EvalBootstrapSetup did not grow heap as expected"; EXPECT_LT(after, peak) << "ClearBootstrapPrecom did not release memory"; }