@@ -231,3 +231,116 @@ TEST_CASE("MLS Failure after Purge")
231231 const auto dec_ab_2 = member_b.unprotect (pt_out, enc_ab_2, metadata).unwrap ();
232232 CHECK (plaintext == to_bytes (dec_ab_2));
233233}
234+
235+ TEST_CASE (" SFrame Context Remove Key" )
236+ {
237+ const auto suite = CipherSuite::AES_GCM_128_SHA256;
238+ const auto kid = KeyID (0x07 );
239+ const auto key = from_hex (" 000102030405060708090a0b0c0d0e0f" );
240+ const auto plaintext = from_hex (" 00010203" );
241+ const auto metadata = bytes{};
242+
243+ auto pt_out = bytes (plaintext.size ());
244+ auto ct_out = bytes (plaintext.size () + Context::max_overhead);
245+
246+ auto sender = Context (suite);
247+ auto receiver = Context (suite);
248+ sender.add_key (kid, KeyUsage::protect, key).unwrap ();
249+ receiver.add_key (kid, KeyUsage::unprotect, key).unwrap ();
250+
251+ // Protect and unprotect succeed before removal
252+ auto encrypted =
253+ to_bytes (sender.protect (kid, ct_out, plaintext, metadata).unwrap ());
254+ auto decrypted =
255+ to_bytes (receiver.unprotect (pt_out, encrypted, metadata).unwrap ());
256+ CHECK (decrypted == plaintext);
257+
258+ // Remove sender key and verify protect fails
259+ sender.remove_key (kid);
260+ CHECK (sender.protect (kid, ct_out, plaintext, metadata).error ().type () ==
261+ SFrameErrorType::invalid_parameter_error);
262+
263+ // Remove receiver key and verify unprotect fails
264+ receiver.remove_key (kid);
265+ CHECK (receiver.unprotect (pt_out, encrypted, metadata).error ().type () ==
266+ SFrameErrorType::invalid_parameter_error);
267+
268+ // Re-add keys and verify round-trip works again
269+ sender.add_key (kid, KeyUsage::protect, key).unwrap ();
270+ receiver.add_key (kid, KeyUsage::unprotect, key).unwrap ();
271+
272+ encrypted =
273+ to_bytes (sender.protect (kid, ct_out, plaintext, metadata).unwrap ());
274+ decrypted =
275+ to_bytes (receiver.unprotect (pt_out, encrypted, metadata).unwrap ());
276+ CHECK (decrypted == plaintext);
277+ }
278+
279+ TEST_CASE (" SFrame Context Remove Key - Nonexistent Key" )
280+ {
281+ const auto suite = CipherSuite::AES_GCM_128_SHA256;
282+
283+ auto ctx = Context (suite);
284+
285+ // Removing a key that was never added should not throw
286+ CHECK_NOTHROW (ctx.remove_key (KeyID (0x99 )));
287+ }
288+
289+ TEST_CASE (" MLS Remove Epoch" )
290+ {
291+ const auto suite = CipherSuite::AES_GCM_128_SHA256;
292+ const auto epoch_bits = 2 ;
293+ const auto metadata = from_hex (" 00010203" );
294+ const auto plaintext = from_hex (" 04050607" );
295+ const auto sender_id = MLSContext::SenderID (0xA0A0A0A0 );
296+ const auto sframe_epoch_secret_1 = bytes (32 , 1 );
297+ const auto sframe_epoch_secret_2 = bytes (32 , 2 );
298+
299+ auto pt_out = bytes (plaintext.size ());
300+ auto ct_out = bytes (plaintext.size () + Context::max_overhead);
301+
302+ auto member_a = MLSContext (suite, epoch_bits);
303+ auto member_b = MLSContext (suite, epoch_bits);
304+
305+ // Install epoch 1 and verify round-trip
306+ const auto epoch_id_1 = MLSContext::EpochID (1 );
307+ member_a.add_epoch (epoch_id_1, sframe_epoch_secret_1);
308+ member_b.add_epoch (epoch_id_1, sframe_epoch_secret_1);
309+
310+ auto enc =
311+ member_a.protect (epoch_id_1, sender_id, ct_out, plaintext, metadata)
312+ .unwrap ();
313+ auto enc_data = to_bytes (enc);
314+ auto dec = to_bytes (member_b.unprotect (pt_out, enc_data, metadata).unwrap ());
315+ CHECK (plaintext == dec);
316+
317+ // Install epoch 2
318+ const auto epoch_id_2 = MLSContext::EpochID (2 );
319+ member_a.add_epoch (epoch_id_2, sframe_epoch_secret_2);
320+ member_b.add_epoch (epoch_id_2, sframe_epoch_secret_2);
321+
322+ // Remove only epoch 1 (not purge_before) and verify it fails
323+ member_a.remove_epoch (epoch_id_1);
324+ member_b.remove_epoch (epoch_id_1);
325+
326+ CHECK (member_a.protect (epoch_id_1, sender_id, ct_out, plaintext, metadata)
327+ .error ()
328+ .type () == SFrameErrorType::invalid_parameter_error);
329+ CHECK (member_b.unprotect (pt_out, enc_data, metadata).error ().type () ==
330+ SFrameErrorType::invalid_parameter_error);
331+
332+ // Epoch 2 should still work
333+ enc = member_a.protect (epoch_id_2, sender_id, ct_out, plaintext, metadata)
334+ .unwrap ();
335+ dec = to_bytes (member_b.unprotect (pt_out, enc, metadata).unwrap ());
336+ CHECK (plaintext == dec);
337+
338+ // Re-add epoch 1 with the same secret and verify it works again
339+ member_a.add_epoch (epoch_id_1, sframe_epoch_secret_1);
340+ member_b.add_epoch (epoch_id_1, sframe_epoch_secret_1);
341+
342+ enc = member_a.protect (epoch_id_1, sender_id, ct_out, plaintext, metadata)
343+ .unwrap ();
344+ dec = to_bytes (member_b.unprotect (pt_out, enc, metadata).unwrap ());
345+ CHECK (plaintext == dec);
346+ }
0 commit comments