From f9dd07d3a389825df1ce3aba20d4eea72db079ee Mon Sep 17 00:00:00 2001
From: kwasniow <kwasniow@cisco.com>
Date: Fri, 17 Apr 2026 13:15:36 +0200
Subject: [PATCH] allow disabling mls context

---
 CMakeLists.txt          | 6 ++++++
 include/sframe/sframe.h | 8 ++++++++
 src/sframe.cpp          | 4 ++++
 test/sframe.cpp         | 4 ++++
 4 files changed, 22 insertions(+)

diff --git a/CMakeLists.txt b/CMakeLists.txt
index b51961a..ffad340 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -9,6 +9,7 @@ option(TESTING    "Build tests" OFF)
 option(CLANG_TIDY "Perform linting with clang-tidy" OFF)
 option(SANITIZERS "Enable sanitizers" OFF)
 option(NO_ALLOC   "Build without needing an allocator" OFF)
+option(MLS_CONTEXT "Build with MLSContext support" ON)
 option(NAMESPACE_SUFFIX "Namespace Suffix for CXX and CMake Export")
 
 if(NAMESPACE_SUFFIX)
@@ -72,6 +73,11 @@ if(NO_ALLOC)
   add_definitions(-DNO_ALLOC)
 endif()
 
+if(NOT MLS_CONTEXT)
+  message(STATUS "Disabling MLSContext")
+  add_definitions(-DSFRAME_ENABLE_MLS_CONTEXT=0)
+endif()
+
 ###
 ### Dependencies
 ###
diff --git a/include/sframe/sframe.h b/include/sframe/sframe.h
index 397d03e..094a560 100644
--- a/include/sframe/sframe.h
+++ b/include/sframe/sframe.h
@@ -30,6 +30,10 @@
 #define SFRAME_EPOCH_BITS 4
 #endif
 
+#ifndef SFRAME_ENABLE_MLS_CONTEXT
+#define SFRAME_ENABLE_MLS_CONTEXT 1
+#endif
+
 namespace SFRAME_NAMESPACE {
 
 #ifdef __cpp_exceptions
@@ -147,6 +151,8 @@ class Context
                                        input_bytes metadata);
 };
 
+#if SFRAME_ENABLE_MLS_CONTEXT
+
 // MLSContext augments Context with logic for deriving keys from MLS.  Instead
 // of adding individual keys, salts, and key IDs, the caller adds a secret for
 // an epoch, and keys / salts / key IDs are derived as needed.
@@ -219,4 +225,6 @@ class MLSContext : protected Context
   vector<std::optional<EpochKeys>, max_epochs> epoch_cache;
 };
 
+#endif // SFRAME_ENABLE_MLS_CONTEXT
+
 } // namespace SFRAME_NAMESPACE
diff --git a/src/sframe.cpp b/src/sframe.cpp
index b389cfc..542f93e 100644
--- a/src/sframe.cpp
+++ b/src/sframe.cpp
@@ -216,6 +216,8 @@ Context::unprotect_inner(const Header& header,
   return open(suite, key_and_salt.key, nonce, plaintext, aad, ciphertext);
 }
 
+#if SFRAME_ENABLE_MLS_CONTEXT
+
 ///
 /// MLSContext
 ///
@@ -417,4 +419,6 @@ MLSContext::ensure_key(KeyID key_id, KeyUsage usage)
   return Context::add_key(key_id, usage, base);
 }
 
+#endif // SFRAME_ENABLE_MLS_CONTEXT
+
 } // namespace SFRAME_NAMESPACE
diff --git a/test/sframe.cpp b/test/sframe.cpp
index 9f8adf8..a3f74dc 100644
--- a/test/sframe.cpp
+++ b/test/sframe.cpp
@@ -55,6 +55,7 @@ TEST_CASE("SFrame Round-Trip")
 
 // The MLS-based key derivation isn't covered by the RFC test vectors.  So we
 // only have round-trip tests, not known-answer tests.
+#if SFRAME_ENABLE_MLS_CONTEXT
 TEST_CASE("MLS Round-Trip")
 {
   const auto epoch_bits = 2;
@@ -231,6 +232,7 @@ TEST_CASE("MLS Failure after Purge")
   const auto dec_ab_2 = member_b.unprotect(pt_out, enc_ab_2, metadata).unwrap();
   CHECK(plaintext == to_bytes(dec_ab_2));
 }
+#endif // SFRAME_ENABLE_MLS_CONTEXT
 
 TEST_CASE("SFrame Context Remove Key")
 {
@@ -286,6 +288,7 @@ TEST_CASE("SFrame Context Remove Key - Nonexistent Key")
   CHECK_NOTHROW(ctx.remove_key(KeyID(0x99)));
 }
 
+#if SFRAME_ENABLE_MLS_CONTEXT
 TEST_CASE("MLS Remove Epoch")
 {
   const auto suite = CipherSuite::AES_GCM_128_SHA256;
@@ -344,3 +347,4 @@ TEST_CASE("MLS Remove Epoch")
   dec = to_bytes(member_b.unprotect(pt_out, enc, metadata).unwrap());
   CHECK(plaintext == dec);
 }
+#endif // SFRAME_ENABLE_MLS_CONTEXT