beaconcrypt 0.6.1

use beaconcrypt::*;
use beaconcrypt_parse::{
	CryptoFrame, InitKex, KexResponse, decode_frame, decode_init, decode_response, encode_init,
};

use std::panic::{AssertUnwindSafe, catch_unwind};

mod common;

#[derive(serde::Deserialize)]
struct SerializedStateUpdate {
	kid: u64,
	seq: u64,
	state: String,
	data: Vec<u8>,
}

fn snapshot_status(snapshot: &str) -> (u64, u64, usize) {
	let value: serde_json::Value = serde_json::from_str(snapshot).unwrap();
	(
		value["send_ctr"].as_u64().unwrap(),
		value["recv_ctr"].as_u64().unwrap(),
		value["recv_past"].as_object().unwrap().len(),
	)
}

trait TestEndpoint {
	fn encrypt_for_test(&mut self, data: &[u8], kid: u64) -> Option<Encrypted>;
	fn decrypt_for_test(&mut self, data: &[u8]) -> Option<Decrypted>;
	fn ratchet_status_for_test(&self, kid: u64) -> RatchetStatus;
}
impl TestEndpoint for Server {
	fn encrypt_for_test(&mut self, d: &[u8], k: u64) -> Option<Encrypted> {
		self.encrypt_message(d, k)
	}
	fn decrypt_for_test(&mut self, d: &[u8]) -> Option<Decrypted> {
		self.decrypt_message(d)
	}
	fn ratchet_status_for_test(&self, k: u64) -> RatchetStatus {
		self.ratchet_status(k).unwrap()
	}
}
impl TestEndpoint for Beacon {
	fn encrypt_for_test(&mut self, d: &[u8], k: u64) -> Option<Encrypted> {
		assert_eq!(k, self.server_kid());
		self.encrypt_message(d)
	}
	fn decrypt_for_test(&mut self, d: &[u8]) -> Option<Decrypted> {
		self.decrypt_message(d)
	}
	fn ratchet_status_for_test(&self, k: u64) -> RatchetStatus {
		assert_eq!(k, self.server_kid());
		self.ratchet_status().unwrap()
	}
}

const SERVER_KID: u64 = 0;
const TAG_SIZE: usize = 16;
const COMMITMENT_SIZE: usize = 64;
const CRYPTO_PAYLOAD_OVERHEAD: usize = TAG_SIZE + COMMITMENT_SIZE;
const RECEIVE_GAP_LIMIT: u64 = 50;

fn new_pair() -> (Server, Beacon) {
	let server = Server::new(SERVER_KID, None);
	let server_id = server.identity_pk().to_owned();
	let beacon = Beacon::new(SERVER_KID, server_id.as_bytes());
	assert_eq!(beacon.server_id(), &server_id);
	(server, beacon)
}

fn register_beacon(
	server: &mut Server,
	beacon: &mut Beacon,
	initial_message: Option<&[u8]>,
) -> RegResponse {
	let configured_server = server.identity_pk().clone();
	let phase_1 = beacon.get_registration_bundle().unwrap();
	let reg_out = server.get_shared_secret(&phase_1).unwrap();
	let phase_2 = server
		.build_registration_response(reg_out, initial_message)
		.unwrap();
	assert_eq!(server.pk_by_kid(phase_2.kid), Some(beacon.identity_pk()));
	let plaintext = beacon.finish_registration(&phase_2.serialized).unwrap();
	assert_eq!(beacon.server_id(), &configured_server);
	// `build_registration_response` insert a single 0xFF byte if no initial message is provided to satisfy the upstream specification's requirement
	// this allows the beacon to check that keychain derivation is correct
	assert_eq!(plaintext, initial_message.unwrap_or(&[0xFFu8; 1]));
	phase_2
}

fn corrupt_aead_ciphertext(serialized: &[u8]) -> Vec<u8> {
	let frame = decode_frame(serialized).unwrap();
	let mut ciphertext = frame.cipher_text.to_vec();
	ciphertext[0] ^= 0x01;
	serialize_crypto_frame(frame.seq, frame.key_id, &ciphertext)
}

fn serialize_crypto_frame(seq: u64, key_id: u64, ciphertext: &[u8]) -> Vec<u8> {
	common::encode_frame(CryptoFrame {
		seq,
		key_id,
		cipher_text: ciphertext,
	})
}

fn crypto_frame_seq(serialized: &[u8]) -> u64 {
	decode_frame(serialized).unwrap().seq
}

fn crypto_frame_key_id(serialized: &[u8]) -> u64 {
	decode_frame(serialized).unwrap().key_id
}

fn crypto_frame_ciphertext(serialized: &[u8]) -> Vec<u8> {
	decode_frame(serialized).unwrap().cipher_text.to_vec()
}

fn rewrite_crypto_frame_seq(serialized: &[u8], seq: u64) -> Vec<u8> {
	serialize_crypto_frame(
		seq,
		crypto_frame_key_id(serialized),
		&crypto_frame_ciphertext(serialized),
	)
}

fn corrupt_crypto_frame_commitment(serialized: &[u8]) -> Vec<u8> {
	let seq = crypto_frame_seq(serialized);
	let key_id = crypto_frame_key_id(serialized);
	let mut ciphertext = crypto_frame_ciphertext(serialized);
	let commitment_start = ciphertext.len() - COMMITMENT_SIZE;
	ciphertext[commitment_start] ^= 0x01;
	serialize_crypto_frame(seq, key_id, &ciphertext)
}

fn receive_state(crypto: &impl TestEndpoint, kid: u64) -> RatchetStatus {
	crypto.ratchet_status_for_test(kid)
}

fn cached_receive_key_count(crypto: &impl TestEndpoint, kid: u64, _start: u64, _end: u64) -> usize {
	crypto.ratchet_status_for_test(kid).receive_cache_len() as usize
}

fn assert_server_frame_tampering_is_rejected(mut tamper: impl FnMut(&mut Vec<u8>, &[u8])) {
	let (mut server, mut beacon) = new_pair();
	let response = register_beacon(&mut server, &mut beacon, None);
	let plaintext = b"message whose authenticated fields will be tampered with";
	let valid = server.encrypt_for_test(plaintext, response.kid).unwrap();
	let donor = server
		.encrypt_for_test(
			b"a second message supplies independently generated fields",
			response.kid,
		)
		.unwrap();
	let seq = crypto_frame_seq(&valid);
	let key_id = crypto_frame_key_id(&valid);
	let mut ciphertext = crypto_frame_ciphertext(&valid);
	let donor_ciphertext = crypto_frame_ciphertext(&donor);
	let receive_before = receive_state(&beacon, SERVER_KID);

	tamper(&mut ciphertext, &donor_ciphertext);
	let tampered = serialize_crypto_frame(seq, key_id, &ciphertext);

	assert!(beacon.decrypt_for_test(&tampered).is_none());
	assert_eq!(
		receive_state(&beacon, SERVER_KID),
		receive_before,
		"rejecting a tampered frame changed the receive state"
	);
	assert_eq!(
		beacon.decrypt_for_test(&valid).unwrap().plaintext,
		plaintext,
		"rejecting a tampered frame must not consume its receive key"
	);
}

fn assert_sequence_relabelling_is_rejected(
	sender: &mut impl TestEndpoint,
	receiver: &mut impl TestEndpoint,
	sender_target_kid: u64,
) {
	let first_plaintext = b"first sequence-bound message";
	let second_plaintext = b"second sequence-bound message";
	let first = sender
		.encrypt_for_test(first_plaintext, sender_target_kid)
		.unwrap();
	let second = sender
		.encrypt_for_test(second_plaintext, sender_target_kid)
		.unwrap();
	let first_seq = crypto_frame_seq(&first);
	let second_seq = crypto_frame_seq(&second);
	let receiver_remote_kid = crypto_frame_key_id(&first);
	assert_eq!(second_seq, first_seq + 1);

	let relabelled = rewrite_crypto_frame_seq(&first, second_seq);
	let receive_before = receive_state(receiver, receiver_remote_kid);
	assert!(receiver.decrypt_for_test(&relabelled).is_none());
	assert_eq!(
		receive_state(receiver, receiver_remote_kid),
		receive_before,
		"sequence relabelling changed the receive state"
	);
	assert_eq!(
		receiver.decrypt_for_test(&first).unwrap().plaintext,
		first_plaintext
	);
	assert_eq!(
		receiver.decrypt_for_test(&second).unwrap().plaintext,
		second_plaintext
	);
}

fn assert_invalid_future_frames_cannot_grow_receive_cache(
	sender: &mut impl TestEndpoint,
	receiver: &mut impl TestEndpoint,
	sender_target_kid: u64,
	receiver_remote_kid: u64,
) {
	let frames = (0..RECEIVE_GAP_LIMIT + 3)
		.map(|index| {
			let plaintext = format!("capacity-message-{index}").into_bytes();
			let frame = sender
				.encrypt_for_test(&plaintext, sender_target_kid)
				.unwrap();
			(plaintext, frame)
		})
		.collect::<Vec<_>>();
	let first_seq = crypto_frame_seq(&frames[0].1);
	let current_seq = first_seq - 1;
	let initial_state = receive_state(receiver, receiver_remote_kid);
	assert_eq!(initial_state.receive_sequence(), current_seq);
	assert_eq!(initial_state.receive_cache_len(), 0);
	let boundary = &frames[RECEIVE_GAP_LIMIT as usize];
	let boundary_seq = crypto_frame_seq(&boundary.1);
	assert_eq!(boundary_seq, current_seq + RECEIVE_GAP_LIMIT + 1);
	let corrupted = corrupt_crypto_frame_commitment(&frames[0].1);

	for rejected_seq in [0, boundary_seq + 1, u64::MAX] {
		let forged = rewrite_crypto_frame_seq(&corrupted, rejected_seq);
		assert!(
			receiver.decrypt_for_test(&forged).is_none(),
			"forged sequence {rejected_seq} was accepted"
		);
		assert_eq!(
			receive_state(receiver, receiver_remote_kid),
			initial_state,
			"out-of-range sequence {rejected_seq} advanced the receive state"
		);
	}

	for forged_seq in first_seq..=boundary_seq {
		let forged = rewrite_crypto_frame_seq(&corrupted, forged_seq);
		assert!(
			receiver.decrypt_for_test(&forged).is_none(),
			"invalid frame at sequence {forged_seq} was accepted"
		);
		assert_eq!(
			receive_state(receiver, receiver_remote_kid),
			initial_state,
			"invalid frame at sequence {forged_seq} changed the receive state"
		);
	}

	assert_eq!(
		receiver.decrypt_for_test(&boundary.1).unwrap().plaintext,
		boundary.0
	);
	let boundary_state = receive_state(receiver, receiver_remote_kid);
	assert_eq!(boundary_state.receive_sequence(), boundary_seq);
	assert_eq!(boundary_state.receive_cache_len(), RECEIVE_GAP_LIMIT as u8);

	let distance_two = &frames[RECEIVE_GAP_LIMIT as usize + 2];
	assert_eq!(crypto_frame_seq(&distance_two.1), boundary_seq + 2);
	assert!(receiver.decrypt_for_test(&distance_two.1).is_none());
	assert_eq!(
		receive_state(receiver, receiver_remote_kid),
		boundary_state,
		"capacity rejection changed the receive state"
	);

	let first = &frames[0];
	assert_eq!(
		receiver.decrypt_for_test(&first.1).unwrap().plaintext,
		first.0
	);
	let after_cached_consumption = receive_state(receiver, receiver_remote_kid);
	assert_eq!(
		after_cached_consumption.receive_cache_len(),
		RECEIVE_GAP_LIMIT as u8 - 1
	);

	assert_eq!(
		receiver
			.decrypt_for_test(&distance_two.1)
			.unwrap()
			.plaintext,
		distance_two.0
	);
	let recovered_state = receive_state(receiver, receiver_remote_kid);
	assert_eq!(recovered_state.receive_sequence(), boundary_seq + 2);
	assert_eq!(recovered_state.receive_cache_len(), RECEIVE_GAP_LIMIT as u8);

	let skipped = &frames[RECEIVE_GAP_LIMIT as usize + 1];
	assert_eq!(
		receiver.decrypt_for_test(&skipped.1).unwrap().plaintext,
		skipped.0
	);
}

fn assert_every_inner_payload_bit_is_authenticated(
	sender: &mut impl TestEndpoint,
	receiver: &mut impl TestEndpoint,
	sender_target_kid: u64,
) {
	let plaintext = b"every ciphertext, tag, and commitment bit is authenticated";
	let valid = sender
		.encrypt_for_test(plaintext, sender_target_kid)
		.unwrap();
	let seq = crypto_frame_seq(&valid);
	let key_id = crypto_frame_key_id(&valid);
	let payload = crypto_frame_ciphertext(&valid);
	assert!(payload.len() > CRYPTO_PAYLOAD_OVERHEAD);
	let receive_before = receive_state(receiver, key_id);

	for byte in 0..payload.len() {
		for bit in 0..u8::BITS {
			let mut mutated_payload = payload.clone();
			mutated_payload[byte] ^= 1 << bit;
			let mutated = serialize_crypto_frame(seq, key_id, &mutated_payload);
			assert!(
				receiver.decrypt_for_test(&mutated).is_none(),
				"accepted mutation at payload byte {byte}, bit {bit}"
			);
			assert_eq!(
				receive_state(receiver, key_id),
				receive_before,
				"mutation at payload byte {byte}, bit {bit} changed the receive state"
			);
		}
	}

	assert_eq!(
		receiver.decrypt_for_test(&valid).unwrap().plaintext,
		plaintext,
		"failed mutations must not consume the authentic frame's receive key"
	);
}

fn assert_valid_frame_components_cannot_be_spliced(
	sender: &mut impl TestEndpoint,
	receiver: &mut impl TestEndpoint,
	sender_target_kid: u64,
) {
	let plaintext_one = [0x11; 32];
	let plaintext_two = [0x22; 32];
	let first = sender
		.encrypt_for_test(&plaintext_one, sender_target_kid)
		.unwrap();
	let second = sender
		.encrypt_for_test(&plaintext_two, sender_target_kid)
		.unwrap();
	let first_seq = crypto_frame_seq(&first);
	let second_seq = crypto_frame_seq(&second);
	let key_id = crypto_frame_key_id(&first);
	let first_payload = crypto_frame_ciphertext(&first);
	let second_payload = crypto_frame_ciphertext(&second);
	assert_eq!(second_seq, first_seq + 1);
	assert_eq!(crypto_frame_key_id(&second), key_id);
	assert_eq!(first_payload.len(), second_payload.len());
	let receive_before = receive_state(receiver, key_id);

	let body_end = first_payload.len() - CRYPTO_PAYLOAD_OVERHEAD;
	let tag_end = first_payload.len() - COMMITMENT_SIZE;
	let mut body_splice = first_payload.clone();
	body_splice[..body_end].copy_from_slice(&second_payload[..body_end]);
	let mut tag_splice = first_payload.clone();
	tag_splice[body_end..tag_end].copy_from_slice(&second_payload[body_end..tag_end]);
	let mut commitment_splice = first_payload.clone();
	commitment_splice[tag_end..].copy_from_slice(&second_payload[tag_end..]);
	let mut authenticated_suffix_splice = first_payload.clone();
	authenticated_suffix_splice[body_end..].copy_from_slice(&second_payload[body_end..]);

	for (name, payload) in [
		("ciphertext body", body_splice),
		("AEAD tag", tag_splice),
		("commitment", commitment_splice),
		("tag and commitment", authenticated_suffix_splice),
		("complete payload under the wrong sequence", second_payload),
	] {
		let spliced = serialize_crypto_frame(first_seq, key_id, &payload);
		assert!(
			receiver.decrypt_for_test(&spliced).is_none(),
			"accepted a frame with a spliced {name}"
		);
		assert_eq!(
			receive_state(receiver, key_id),
			receive_before,
			"rejecting a frame with a spliced {name} changed the receive state"
		);
	}

	assert_eq!(
		receiver.decrypt_for_test(&first).unwrap().plaintext,
		plaintext_one
	);
	assert_eq!(
		receiver.decrypt_for_test(&second).unwrap().plaintext,
		plaintext_two
	);
}

fn assert_receive_window_boundary_survives_rejection_retry_and_replay(
	sender: &mut impl TestEndpoint,
	receiver: &mut impl TestEndpoint,
	sender_target_kid: u64,
	receiver_remote_kid: u64,
) {
	let frames = (0..RECEIVE_GAP_LIMIT + 1)
		.map(|index| {
			let plaintext = format!("receive-window-message-{index}").into_bytes();
			let frame = sender
				.encrypt_for_test(&plaintext, sender_target_kid)
				.unwrap();
			(plaintext, frame)
		})
		.collect::<Vec<_>>();
	let first_seq = crypto_frame_seq(&frames[0].1);
	let boundary_seq = crypto_frame_seq(&frames.last().unwrap().1);
	let sender_kid = crypto_frame_key_id(&frames[0].1);
	assert_eq!(boundary_seq - first_seq, RECEIVE_GAP_LIMIT);

	let initial_state = receive_state(receiver, receiver_remote_kid);
	for len in 0..=CRYPTO_PAYLOAD_OVERHEAD {
		let malformed = serialize_crypto_frame(boundary_seq, sender_kid, &vec![0xA5; len]);
		let result = catch_unwind(AssertUnwindSafe(|| receiver.decrypt_for_test(&malformed)));
		assert!(
			matches!(result, Ok(None)),
			"future frame with {len} payload bytes was not rejected cleanly"
		);
		assert_eq!(
			receive_state(receiver, receiver_remote_kid),
			initial_state,
			"short future frame with {len} payload bytes advanced the receive ratchet"
		);
		assert_eq!(
			cached_receive_key_count(receiver, receiver_remote_kid, first_seq, boundary_seq,),
			0,
			"short future frame with {len} payload bytes cached receive keys"
		);
	}

	let boundary = &frames.last().unwrap().1;
	let corrupted_commitment = corrupt_crypto_frame_commitment(boundary);
	assert!(receiver.decrypt_for_test(&corrupted_commitment).is_none());
	assert_eq!(
		receive_state(receiver, receiver_remote_kid),
		initial_state,
		"commitment failure at the gap boundary changed the receive state"
	);

	for attempt in 0..3 {
		assert!(
			receiver.decrypt_for_test(&corrupted_commitment).is_none(),
			"repeated invalid boundary frame was accepted on attempt {attempt}"
		);
		assert_eq!(
			receive_state(receiver, receiver_remote_kid),
			initial_state,
			"repeated invalid boundary frame changed the receive state"
		);
	}

	let corrupted_ciphertext = corrupt_aead_ciphertext(boundary);
	assert!(receiver.decrypt_for_test(&corrupted_ciphertext).is_none());
	assert_eq!(
		receive_state(receiver, receiver_remote_kid),
		initial_state,
		"AEAD failure at the gap boundary changed the receive state"
	);

	let boundary_plaintext = &frames.last().unwrap().0;
	assert_eq!(
		receiver.decrypt_for_test(boundary).unwrap().plaintext,
		boundary_plaintext.as_slice()
	);
	let boundary_state = receive_state(receiver, receiver_remote_kid);
	assert_eq!(boundary_state.receive_sequence(), boundary_seq);
	assert_eq!(boundary_state.receive_cache_len(), RECEIVE_GAP_LIMIT as u8);
	assert!(receiver.decrypt_for_test(boundary).is_none());
	assert_eq!(receive_state(receiver, receiver_remote_kid), boundary_state);

	for index in (0..frames.len() - 1).rev() {
		let (plaintext, frame) = &frames[index];
		assert_eq!(
			receiver.decrypt_for_test(frame).unwrap().plaintext,
			plaintext.as_slice(),
			"failed to decrypt cached frame at index {index}"
		);
		let after_success = receive_state(receiver, receiver_remote_kid);
		assert!(
			receiver.decrypt_for_test(frame).is_none(),
			"replayed cached frame at index {index} was accepted"
		);
		assert_eq!(
			receive_state(receiver, receiver_remote_kid),
			after_success,
			"replaying cached frame at index {index} changed the receive state"
		);
		assert_eq!(
			cached_receive_key_count(receiver, receiver_remote_kid, first_seq, boundary_seq,),
			index,
			"unexpected receive-cache size after index {index}"
		);
	}
}

fn assert_truncated_frames_reject_cleanly(
	sender: &mut impl TestEndpoint,
	receiver: &mut impl TestEndpoint,
	sender_target_kid: u64,
	receiver_remote_kid: u64,
) {
	let plaintext = b"authentic frame remains usable after every truncated prefix";
	let valid = sender
		.encrypt_for_test(plaintext, sender_target_kid)
		.unwrap();
	let corrupted = corrupt_crypto_frame_commitment(&valid);
	let initial_state = receive_state(receiver, receiver_remote_kid);

	for cut in 0..corrupted.len() {
		let result = catch_unwind(AssertUnwindSafe(|| {
			receiver.decrypt_for_test(&corrupted[..cut])
		}));
		assert!(
			matches!(result, Ok(None)),
			"serialized frame prefix of {cut} bytes was not rejected cleanly"
		);
	}

	assert_eq!(
		receive_state(receiver, receiver_remote_kid),
		initial_state,
		"a strict serialized prefix advanced the receive ratchet"
	);
	assert_eq!(
		receiver.decrypt_for_test(&valid).unwrap().plaintext,
		plaintext
	);
}

fn assert_message_size_boundaries_round_trip(
	sender: &mut impl TestEndpoint,
	receiver: &mut impl TestEndpoint,
	sender_target_kid: u64,
) {
	for len in [1, 15, 16, 17, 63, 64, 65, 255, 256, 4096] {
		let plaintext = (0..len)
			.map(|index| (index % (u8::MAX as usize + 1)) as u8)
			.collect::<Vec<_>>();
		let frame = sender
			.encrypt_for_test(&plaintext, sender_target_kid)
			.unwrap();
		assert_eq!(
			crypto_frame_ciphertext(&frame).len(),
			plaintext.len() + CRYPTO_PAYLOAD_OVERHEAD
		);
		assert_eq!(
			receiver.decrypt_for_test(&frame).unwrap().plaintext,
			plaintext,
			"failed round trip for a {len}-byte plaintext"
		);
	}
}

struct KexResponseParts {
	identity_key: Vec<u8>,
	ephemeral_key: Vec<u8>,
	kem_ciphertext: Vec<u8>,
	app_ciphertext: Vec<u8>,
	key_id: u64,
}

fn parse_kex_response(serialized: &[u8]) -> KexResponseParts {
	let response = decode_response(serialized).unwrap();
	KexResponseParts {
		identity_key: response.identity_key.to_vec(),
		ephemeral_key: response.ephemeral_key.to_vec(),
		kem_ciphertext: response.kem_cipher_text.to_vec(),
		app_ciphertext: response.app_cipher_text.to_vec(),
		key_id: response.key_id,
	}
}

fn serialize_kex_response(parts: &KexResponseParts) -> Vec<u8> {
	common::encode_response(KexResponse {
		identity_key: &parts.identity_key,
		ephemeral_key: &parts.ephemeral_key,
		kem_cipher_text: &parts.kem_ciphertext,
		app_cipher_text: &parts.app_ciphertext,
		key_id: parts.key_id,
	})
}

fn pending_registration(initial_message: &[u8]) -> (Beacon, Vec<u8>) {
	let (mut server, mut beacon) = new_pair();
	let phase_1 = beacon.get_registration_bundle().unwrap();
	let reg_out = server.get_shared_secret(&phase_1).unwrap();
	let response = server
		.build_registration_response(reg_out, Some(initial_message))
		.unwrap();
	(beacon, response.serialized)
}

fn rewrite_crypto_frame_key_id(serialized: &[u8], key_id: u64) -> Vec<u8> {
	let frame = decode_frame(serialized).unwrap();
	serialize_crypto_frame(frame.seq, key_id, frame.cipher_text)
}

#[test]
fn server_uses_per_beacon_associated_data() {
	let mut server = Server::new(SERVER_KID, None);
	let server_id = server.identity_pk().to_owned();
	let mut b1 = Beacon::new(SERVER_KID, server_id.as_bytes());
	let mut b2 = Beacon::new(SERVER_KID, server_id.as_bytes());

	let b1_response = register_beacon(&mut server, &mut b1, Some(b"initial message for b1"));
	let b2_response = register_beacon(&mut server, &mut b2, Some(b"initial message for b2"));

	let to_b1 = server
		.encrypt_for_test(b"server to b1", b1_response.kid)
		.unwrap();
	let to_b2 = server
		.encrypt_for_test(b"server to b2", b2_response.kid)
		.unwrap();
	assert_eq!(
		b1.decrypt_for_test(&to_b1).unwrap().plaintext,
		b"server to b1"
	);
	assert_eq!(
		b2.decrypt_for_test(&to_b2).unwrap().plaintext,
		b"server to b2"
	);

	let from_b1 = b1.encrypt_for_test(b"b1 to server", SERVER_KID).unwrap();
	let from_b2 = b2.encrypt_for_test(b"b2 to server", SERVER_KID).unwrap();
	assert_eq!(
		server.decrypt_for_test(&from_b1).unwrap().plaintext,
		b"b1 to server"
	);
	assert_eq!(
		server.decrypt_for_test(&from_b2).unwrap().plaintext,
		b"b2 to server"
	);
}

#[test]
fn server_can_encrypt_to_beacon_a_after_registering_beacon_b() {
	let mut server = Server::new(SERVER_KID, None);
	let server_id = server.identity_pk().to_owned();
	let mut beacon_a = Beacon::new(SERVER_KID, server_id.as_bytes());
	let mut beacon_b = Beacon::new(SERVER_KID, server_id.as_bytes());

	let beacon_a_response = register_beacon(&mut server, &mut beacon_a, None);
	register_beacon(&mut server, &mut beacon_b, None);

	let message = b"server to beacon A after registering beacon B";
	let ciphertext = server
		.encrypt_for_test(message, beacon_a_response.kid)
		.unwrap();

	assert_eq!(
		beacon_a.decrypt_for_test(&ciphertext).unwrap().plaintext,
		message
	);
}

#[test]
fn server_can_decrypt_from_beacon_a_after_registering_beacon_b() {
	let mut server = Server::new(SERVER_KID, None);
	let server_id = server.identity_pk().to_owned();
	let mut beacon_a = Beacon::new(SERVER_KID, server_id.as_bytes());
	let mut beacon_b = Beacon::new(SERVER_KID, server_id.as_bytes());

	let _beacon_a_response = register_beacon(&mut server, &mut beacon_a, None);
	register_beacon(&mut server, &mut beacon_b, None);

	let message = b"beacon A to server after registering beacon B";
	let ciphertext = beacon_a.encrypt_for_test(message, SERVER_KID).unwrap();

	assert_eq!(
		server.decrypt_for_test(&ciphertext).unwrap().plaintext,
		message
	);
}

#[test]
fn server_can_decrypt_from_beacon_a_after_encrypting_to_beacon_b() {
	let mut server = Server::new(SERVER_KID, None);
	let server_id = server.identity_pk().to_owned();
	let mut beacon_a = Beacon::new(SERVER_KID, server_id.as_bytes());
	let mut beacon_b = Beacon::new(SERVER_KID, server_id.as_bytes());

	let _beacon_a_response = register_beacon(&mut server, &mut beacon_a, None);
	let beacon_b_response = register_beacon(&mut server, &mut beacon_b, None);

	let to_beacon_b = server
		.encrypt_for_test(b"server to beacon B", beacon_b_response.kid)
		.unwrap();
	assert_eq!(
		beacon_b.decrypt_for_test(&to_beacon_b).unwrap().plaintext,
		b"server to beacon B"
	);

	let from_beacon_a = beacon_a
		.encrypt_for_test(b"beacon A to server", SERVER_KID)
		.unwrap();
	assert_eq!(
		server.decrypt_for_test(&from_beacon_a).unwrap().plaintext,
		b"beacon A to server"
	);
}

#[test]
fn interleaved_registrations_use_the_correct_associated_data() {
	let mut server = Server::new(SERVER_KID, None);
	let server_id = server.identity_pk().to_owned();
	let mut b1 = Beacon::new(SERVER_KID, server_id.as_bytes());
	let mut b2 = Beacon::new(SERVER_KID, server_id.as_bytes());

	let b1_phase_1 = b1.get_registration_bundle().unwrap();
	let b2_phase_1 = b2.get_registration_bundle().unwrap();
	let b1_registration = server.get_shared_secret(&b1_phase_1).unwrap();
	let b2_registration = server.get_shared_secret(&b2_phase_1).unwrap();

	let b1_phase_2 = server
		.build_registration_response(b1_registration, Some(b"initial message for b1"))
		.unwrap();
	let b2_phase_2 = server
		.build_registration_response(b2_registration, Some(b"initial message for b2"))
		.unwrap();

	assert_eq!(
		b1.finish_registration(&b1_phase_2.serialized).unwrap(),
		b"initial message for b1"
	);
	assert_eq!(
		b2.finish_registration(&b2_phase_2.serialized).unwrap(),
		b"initial message for b2"
	);
}

#[test]
fn malformed_registration_is_rejected() {
	let (mut server, _) = new_pair();

	assert!(server.get_shared_secret(b"not a registration").is_none());
}

#[test]
fn fresh_beacon_has_no_operational_ratchet() {
	let (_, mut beacon) = new_pair();

	assert!(beacon.ratchet_status().is_none());
	assert!(beacon.encrypt_message(b"not established").is_none());
	assert!(beacon.decrypt_message(b"not established").is_none());
}

#[test]
fn registration_can_omit_initial_message() {
	let (mut server, mut beacon) = new_pair();

	let response = register_beacon(&mut server, &mut beacon, None);

	assert_eq!(response.kid, 1);
}

#[test]
fn failed_registration_response_does_not_commit_server_state() {
	let (mut server, mut rejected_beacon) = new_pair();
	let phase_1 = rejected_beacon.get_registration_bundle().unwrap();
	let registration = server.get_shared_secret(&phase_1).unwrap();
	let counter_before = server.server_kid();
	let expected_kid = counter_before + 1;

	assert!(
		server
			.build_registration_response(registration, Some(&[]))
			.is_none()
	);
	assert_eq!(server.server_kid(), counter_before);
	assert!(server.pk_by_kid(expected_kid).is_none());
	assert!(server.get_shared_secret(&phase_1).is_none());

	let server_id = server.identity_pk().to_owned();
	let mut accepted_beacon = Beacon::new(SERVER_KID, server_id.as_bytes());
	let response = register_beacon(&mut server, &mut accepted_beacon, None);
	assert_eq!(response.kid, expected_kid);
	assert_eq!(server.server_kid(), expected_kid);
}

#[test]
fn beacon_rejects_registration_response_from_wrong_server() {
	let (mut expected_server, mut beacon) = new_pair();
	let mut wrong_server = Server::new(SERVER_KID, None);

	let phase_1 = beacon.get_registration_bundle().unwrap();
	assert!(expected_server.get_shared_secret(&phase_1).is_some());
	let reg_out = wrong_server.get_shared_secret(&phase_1).unwrap();
	let phase_2 = wrong_server
		.build_registration_response(reg_out, Some(b"wrong server"))
		.unwrap();

	assert!(beacon.finish_registration(&phase_2.serialized).is_none());
}

#[test]
fn server_rejects_pending_registration_from_a_different_server() {
	let (mut accepting_server, mut beacon) = new_pair();
	let mut other_server = Server::new(SERVER_KID, None);
	let phase_1 = beacon.get_registration_bundle().unwrap();
	let pending = accepting_server.get_shared_secret(&phase_1).unwrap();

	assert!(
		other_server
			.build_registration_response(pending, Some(b"wrong server"))
			.is_none()
	);
	assert_eq!(other_server.server_kid(), SERVER_KID);
	assert!(other_server.pk_by_kid(SERVER_KID + 1).is_none());
	assert!(accepting_server.get_shared_secret(&phase_1).is_none());

	let accepting_identity = accepting_server.identity_pk().to_owned();
	let mut fresh_beacon = Beacon::new(SERVER_KID, accepting_identity.as_bytes());
	let response = register_beacon(
		&mut accepting_server,
		&mut fresh_beacon,
		Some(b"right server"),
	);
	assert_eq!(response.kid, SERVER_KID + 1);
}

#[test]
fn beacon_can_encrypt_to_server() {
	let (mut server, mut beacon) = new_pair();
	let response = register_beacon(&mut server, &mut beacon, Some(&[0xFF; 32]));
	let message = b"beacon to server";

	let ciphertext = beacon
		.encrypt_for_test(message.as_slice(), SERVER_KID)
		.unwrap();
	let decrypted = server.decrypt_for_test(&ciphertext).unwrap();

	assert_eq!(decrypted.key_id, response.kid);
	assert_eq!(decrypted.plaintext, message);
}

#[test]
fn beacon_frame_identifies_its_sender() {
	let (mut server, mut beacon) = new_pair();
	let response = register_beacon(&mut server, &mut beacon, Some(&[0xFF; 32]));
	let message = b"authenticated beacon to server";

	let ciphertext = beacon
		.encrypt_for_test(message.as_slice(), SERVER_KID)
		.unwrap();
	let decrypted = server.decrypt_for_test(&ciphertext).unwrap();

	assert_eq!(decrypted.key_id, response.kid);
	assert_eq!(decrypted.plaintext, message);
}

#[test]
fn server_frame_identifies_its_sender() {
	let (mut server, mut beacon) = new_pair();
	let response = register_beacon(&mut server, &mut beacon, None);
	let message = b"authenticated server to beacon";

	let ciphertext = server.encrypt_for_test(message, response.kid).unwrap();
	let decrypted = beacon.decrypt_for_test(&ciphertext).unwrap();

	assert_eq!(decrypted.key_id, SERVER_KID);
	assert_eq!(decrypted.plaintext, message);
}

#[test]
fn authenticated_beacon_frame_rejects_tampering() {
	let (mut server, mut beacon) = new_pair();
	let response = register_beacon(&mut server, &mut beacon, Some(&[0xFF; 32]));

	let mut ciphertext = beacon
		.encrypt_for_test(b"beacon to server", SERVER_KID)
		.unwrap();
	let last = ciphertext.len() - 1;
	ciphertext[last] ^= 0x01;

	let receive_before = receive_state(&server, response.kid);
	assert!(server.decrypt_for_test(&ciphertext).is_none());
	assert_eq!(receive_state(&server, response.kid), receive_before);
}

#[test]
fn authenticated_server_frame_rejects_tampering() {
	let (mut server, mut beacon) = new_pair();
	let response = register_beacon(&mut server, &mut beacon, Some(&[0xFF; 32]));

	let mut ciphertext = server
		.encrypt_for_test(b"server to beacon", response.kid)
		.unwrap();
	let last = ciphertext.len() - 1;
	ciphertext[last] ^= 0x01;

	let receive_before = receive_state(&beacon, SERVER_KID);
	assert!(beacon.decrypt_for_test(&ciphertext).is_none());
	assert_eq!(receive_state(&beacon, SERVER_KID), receive_before);
}

#[test]
fn decrypt_rejects_wrong_direction() {
	let (mut server, mut beacon) = new_pair();
	let response = register_beacon(&mut server, &mut beacon, Some(&[0xFF; 32]));

	let server_to_beacon = server
		.encrypt_for_test(b"server to beacon", response.kid)
		.unwrap();
	let server_receive_before = receive_state(&server, response.kid);
	assert!(server.decrypt_for_test(&server_to_beacon).is_none());
	assert_eq!(receive_state(&server, response.kid), server_receive_before);

	let beacon_to_server = beacon
		.encrypt_for_test(b"beacon to server", SERVER_KID)
		.unwrap();
	let beacon_receive_before = receive_state(&beacon, SERVER_KID);
	assert!(beacon.decrypt_for_test(&beacon_to_server).is_none());
	assert_eq!(receive_state(&beacon, SERVER_KID), beacon_receive_before);
}

#[test]
fn beacon_cannot_decrypt_message_for_different_beacon() {
	let mut server = Server::new(SERVER_KID, None);
	let server_id = server.identity_pk().to_owned();
	let mut b1 = Beacon::new(SERVER_KID, server_id.as_bytes());
	let mut b2 = Beacon::new(SERVER_KID, server_id.as_bytes());
	let b1_response = register_beacon(&mut server, &mut b1, Some(&[0xFF; 32]));
	let b2_response = register_beacon(&mut server, &mut b2, Some(&[0xFF; 32]));

	let for_b1 = server
		.encrypt_for_test(b"for b1 only", b1_response.kid)
		.unwrap();
	let for_b2 = server
		.encrypt_for_test(b"for b2 only", b2_response.kid)
		.unwrap();

	let b2_receive_before = receive_state(&b2, SERVER_KID);
	assert!(b2.decrypt_for_test(&for_b1).is_none());
	assert_eq!(receive_state(&b2, SERVER_KID), b2_receive_before);
	assert_eq!(
		b2.decrypt_for_test(&for_b2).unwrap().plaintext,
		b"for b2 only"
	);
	assert_eq!(
		b1.decrypt_for_test(&for_b1).unwrap().plaintext,
		b"for b1 only"
	);
}

#[test]
fn ciphertext_cannot_be_replayed() {
	let (mut server, mut beacon) = new_pair();
	let response = register_beacon(&mut server, &mut beacon, Some(&[0xFF; 32]));
	let message = b"one shot";

	let ciphertext = server.encrypt_for_test(message, response.kid).unwrap();
	let first = beacon.decrypt_for_test(&ciphertext).unwrap();

	assert_eq!(first.plaintext, message);
	let receive_after_success = receive_state(&beacon, SERVER_KID);
	assert!(beacon.decrypt_for_test(&ciphertext).is_none());
	assert_eq!(receive_state(&beacon, SERVER_KID), receive_after_success);
}

#[test]
fn beacon_can_retry_decryption_after_corrupted_aead_message() {
	let (mut server, mut beacon) = new_pair();
	let response = register_beacon(&mut server, &mut beacon, Some(&[0xFF; 32]));
	let message = b"server to beacon";

	let ciphertext = server.encrypt_for_test(message, response.kid).unwrap();
	let corrupted = corrupt_aead_ciphertext(&ciphertext);

	let receive_before = receive_state(&beacon, SERVER_KID);
	assert!(beacon.decrypt_for_test(&corrupted).is_none());
	assert_eq!(receive_state(&beacon, SERVER_KID), receive_before);
	assert_eq!(
		beacon.decrypt_for_test(&ciphertext).unwrap().plaintext,
		message
	);
}

#[test]
fn server_can_retry_decryption_after_corrupted_aead_message() {
	let (mut server, mut beacon) = new_pair();
	let response = register_beacon(&mut server, &mut beacon, Some(&[0xFF; 32]));
	let message = b"beacon to server";

	let ciphertext = beacon.encrypt_for_test(message, SERVER_KID).unwrap();
	let corrupted = corrupt_aead_ciphertext(&ciphertext);

	let receive_before = receive_state(&server, response.kid);
	assert!(server.decrypt_for_test(&corrupted).is_none());
	assert_eq!(receive_state(&server, response.kid), receive_before);
	assert_eq!(
		server.decrypt_for_test(&ciphertext).unwrap().plaintext,
		message
	);
}

#[test]
fn encrypt_and_update_returns_the_advanced_send_state() {
	let (mut server, mut beacon) = new_pair();
	let response = register_beacon(&mut server, &mut beacon, None);
	let message = b"server to beacon with structured updated state";
	let status_before = server.ratchet_status(response.kid).unwrap();

	let update: SendState = server.encrypt_and_update(message, response.kid).unwrap();
	let status = server.ratchet_status(response.kid).unwrap();

	assert_eq!(update.kid, response.kid);
	assert_eq!(update.seq, crypto_frame_seq(&update.data));
	assert_eq!(
		snapshot_status(update.state.as_str()),
		(
			status.send_sequence(),
			status.receive_sequence(),
			status.receive_cache_len() as usize,
		)
	);
	assert_eq!(status.send_sequence(), status_before.send_sequence() + 1);
	assert_eq!(status.receive_sequence(), status_before.receive_sequence());
	let decrypted = beacon.decrypt_for_test(&update.data).unwrap();
	assert_eq!(decrypted.key_id, SERVER_KID);
	assert_eq!(decrypted.plaintext, message);
}

#[test]
fn decrypt_and_update_returns_the_advanced_receive_state() {
	let (mut server, mut beacon) = new_pair();
	let response = register_beacon(&mut server, &mut beacon, None);
	let message = b"beacon to server with structured updated state";
	let ciphertext = beacon.encrypt_for_test(message, SERVER_KID).unwrap();
	let status_before = server.ratchet_status(response.kid).unwrap();

	let update: RecvState = server.decrypt_and_update(&ciphertext).unwrap();
	let status = server.ratchet_status(response.kid).unwrap();

	assert_eq!(update.kid, response.kid);
	assert_eq!(update.seq, ciphertext.seq);
	assert_eq!(update.data, message);
	assert_eq!(
		snapshot_status(update.state.as_str()),
		(
			status.send_sequence(),
			status.receive_sequence(),
			status.receive_cache_len() as usize,
		)
	);
	assert_eq!(status.send_sequence(), status_before.send_sequence());
	assert_eq!(
		status.receive_sequence(),
		status_before.receive_sequence() + 1
	);
}

#[test]
fn structured_update_failures_return_none() {
	let (mut server, _) = new_pair();

	let encrypted: Option<SendState> = server.encrypt_and_update(b"message", u64::MAX);
	assert!(encrypted.is_none());

	let decrypted: Option<RecvState> = server.decrypt_and_update(b"not a frame");
	assert!(decrypted.is_none());
}

#[test]
fn encrypt_and_update_json_returns_the_advanced_send_state() {
	let (mut server, mut beacon) = new_pair();
	let response = register_beacon(&mut server, &mut beacon, None);
	let message = b"server to beacon with updated state";
	let status_before = server.ratchet_status(response.kid).unwrap();

	let serialized = server
		.encrypt_and_update_json(message, response.kid)
		.unwrap();
	let update: SerializedStateUpdate =
		serde_json::from_str(&serialized).expect("send update should be valid JSON");
	let status = server.ratchet_status(response.kid).unwrap();

	assert_eq!(update.kid, response.kid);
	assert_eq!(update.seq, crypto_frame_seq(&update.data));
	assert_eq!(
		snapshot_status(&update.state),
		(
			status.send_sequence(),
			status.receive_sequence(),
			status.receive_cache_len() as usize,
		)
	);
	assert_eq!(status.send_sequence(), status_before.send_sequence() + 1);
	assert_eq!(status.receive_sequence(), status_before.receive_sequence());
	let decrypted = beacon.decrypt_for_test(&update.data).unwrap();
	assert_eq!(decrypted.key_id, SERVER_KID);
	assert_eq!(decrypted.plaintext, message);
}

#[test]
fn decrypt_and_update_json_returns_the_advanced_receive_state() {
	let (mut server, mut beacon) = new_pair();
	let response = register_beacon(&mut server, &mut beacon, None);
	let message = b"beacon to server with updated state";
	let ciphertext = beacon.encrypt_for_test(message, SERVER_KID).unwrap();
	let status_before = server.ratchet_status(response.kid).unwrap();

	let serialized = server.decrypt_and_update_json(&ciphertext).unwrap();
	let update: SerializedStateUpdate =
		serde_json::from_str(&serialized).expect("receive update should be valid JSON");
	let status = server.ratchet_status(response.kid).unwrap();

	assert_eq!(update.kid, response.kid);
	assert_eq!(update.seq, ciphertext.seq);
	assert_eq!(update.data, message);
	assert_eq!(
		snapshot_status(&update.state),
		(
			status.send_sequence(),
			status.receive_sequence(),
			status.receive_cache_len() as usize,
		)
	);
	assert_eq!(status.send_sequence(), status_before.send_sequence());
	assert_eq!(
		status.receive_sequence(),
		status_before.receive_sequence() + 1
	);
}

#[test]
fn update_json_failures_return_none() {
	let (mut server, _) = new_pair();

	assert!(
		server
			.encrypt_and_update_json(b"message", u64::MAX)
			.is_none()
	);
	assert!(server.decrypt_and_update_json(b"not a frame").is_none());
}

#[test]
fn malformed_crypto_frame_ciphertext_lengths_are_rejected_without_panicking() {
	let (mut server, mut beacon) = new_pair();
	let response = register_beacon(&mut server, &mut beacon, None);
	let valid = server
		.encrypt_for_test(b"valid after malformed frames", response.kid)
		.unwrap();
	let seq = crypto_frame_seq(&valid);
	let receive_before = receive_state(&beacon, SERVER_KID);

	for len in [0, 1, 15, 16, 63, 64, 79, 80] {
		let malformed = serialize_crypto_frame(seq, SERVER_KID, &vec![0xA5; len]);
		let result = catch_unwind(AssertUnwindSafe(|| beacon.decrypt_for_test(&malformed)));
		assert!(
			matches!(result, Ok(None)),
			"ciphertext length {len} was not rejected cleanly"
		);
		assert_eq!(
			receive_state(&beacon, SERVER_KID),
			receive_before,
			"ciphertext length {len} changed the receive state"
		);
	}

	assert_eq!(
		beacon.decrypt_for_test(&valid).unwrap().plaintext,
		b"valid after malformed frames"
	);
}

#[test]
fn empty_plaintext_is_rejected_without_advancing_send_ratchets() {
	let (mut server, mut beacon) = new_pair();
	let response = register_beacon(&mut server, &mut beacon, None);
	let server_sequence = server.ratchet_status(response.kid).unwrap().send_sequence();
	let beacon_sequence = beacon.ratchet_status().unwrap().send_sequence();

	assert!(server.encrypt_for_test(b"", response.kid).is_none());
	assert!(beacon.encrypt_for_test(b"", SERVER_KID).is_none());
	assert_eq!(
		server.ratchet_status(response.kid).unwrap().send_sequence(),
		server_sequence
	);
	assert_eq!(
		beacon.ratchet_status().unwrap().send_sequence(),
		beacon_sequence
	);

	let to_beacon = server.encrypt_for_test(b"non-empty", response.kid).unwrap();
	let to_server = beacon.encrypt_for_test(b"non-empty", SERVER_KID).unwrap();
	assert_eq!(
		beacon.decrypt_for_test(&to_beacon).unwrap().plaintext,
		b"non-empty"
	);
	assert_eq!(
		server.decrypt_for_test(&to_server).unwrap().plaintext,
		b"non-empty"
	);
}

#[test]
fn crypto_frame_sequence_is_bound_in_both_directions() {
	let (mut server, mut beacon) = new_pair();
	let response = register_beacon(&mut server, &mut beacon, None);
	assert_sequence_relabelling_is_rejected(&mut server, &mut beacon, response.kid);

	let (mut server, mut beacon) = new_pair();
	let response = register_beacon(&mut server, &mut beacon, None);
	assert_sequence_relabelling_is_rejected(&mut beacon, &mut server, SERVER_KID);
	assert_eq!(
		server
			.ratchet_status(response.kid)
			.unwrap()
			.receive_cache_len(),
		0
	);
}

#[test]
fn invalid_future_frames_cannot_grow_receive_cache_beyond_gap() {
	let (mut server, mut beacon) = new_pair();
	let response = register_beacon(&mut server, &mut beacon, None);
	assert_invalid_future_frames_cannot_grow_receive_cache(
		&mut server,
		&mut beacon,
		response.kid,
		SERVER_KID,
	);

	let (mut server, mut beacon) = new_pair();
	let response = register_beacon(&mut server, &mut beacon, None);
	assert_invalid_future_frames_cannot_grow_receive_cache(
		&mut beacon,
		&mut server,
		SERVER_KID,
		response.kid,
	);
}

#[test]
fn every_inner_payload_bit_is_authenticated_in_both_directions() {
	let (mut server, mut beacon) = new_pair();
	let response = register_beacon(&mut server, &mut beacon, None);
	assert_every_inner_payload_bit_is_authenticated(&mut server, &mut beacon, response.kid);

	let (mut server, mut beacon) = new_pair();
	register_beacon(&mut server, &mut beacon, None);
	assert_every_inner_payload_bit_is_authenticated(&mut beacon, &mut server, SERVER_KID);
}

#[test]
fn valid_frame_components_cannot_be_spliced_in_both_directions() {
	let (mut server, mut beacon) = new_pair();
	let response = register_beacon(&mut server, &mut beacon, None);
	assert_valid_frame_components_cannot_be_spliced(&mut server, &mut beacon, response.kid);

	let (mut server, mut beacon) = new_pair();
	register_beacon(&mut server, &mut beacon, None);
	assert_valid_frame_components_cannot_be_spliced(&mut beacon, &mut server, SERVER_KID);
}

#[test]
fn receive_window_boundary_survives_rejection_retry_and_replay_in_both_directions() {
	let (mut server, mut beacon) = new_pair();
	let response = register_beacon(&mut server, &mut beacon, None);
	assert_receive_window_boundary_survives_rejection_retry_and_replay(
		&mut server,
		&mut beacon,
		response.kid,
		SERVER_KID,
	);

	let (mut server, mut beacon) = new_pair();
	let response = register_beacon(&mut server, &mut beacon, None);
	assert_receive_window_boundary_survives_rejection_retry_and_replay(
		&mut beacon,
		&mut server,
		SERVER_KID,
		response.kid,
	);
}

#[test]
fn every_truncated_frame_prefix_rejects_cleanly_in_both_directions() {
	let (mut server, mut beacon) = new_pair();
	let response = register_beacon(&mut server, &mut beacon, None);
	assert_truncated_frames_reject_cleanly(&mut server, &mut beacon, response.kid, SERVER_KID);

	let (mut server, mut beacon) = new_pair();
	let response = register_beacon(&mut server, &mut beacon, None);
	assert_truncated_frames_reject_cleanly(&mut beacon, &mut server, SERVER_KID, response.kid);
}

#[test]
fn message_size_boundaries_round_trip_in_both_directions() {
	let (mut server, mut beacon) = new_pair();
	let response = register_beacon(&mut server, &mut beacon, None);
	assert_message_size_boundaries_round_trip(&mut server, &mut beacon, response.kid);

	let (mut server, mut beacon) = new_pair();
	register_beacon(&mut server, &mut beacon, None);
	assert_message_size_boundaries_round_trip(&mut beacon, &mut server, SERVER_KID);
}

#[test]
fn commitment_corruption_is_rejected() {
	assert_server_frame_tampering_is_rejected(|ciphertext, _| {
		let commitment_start = ciphertext.len() - COMMITMENT_SIZE;
		ciphertext[commitment_start] ^= 0x01;
	});
}

#[test]
fn missing_tag_is_rejected() {
	assert_server_frame_tampering_is_rejected(|ciphertext, _| {
		let tag_start = ciphertext.len() - COMMITMENT_SIZE - TAG_SIZE;
		ciphertext.drain(tag_start..tag_start + TAG_SIZE);
	});
}

#[test]
fn missing_commitment_is_rejected() {
	assert_server_frame_tampering_is_rejected(|ciphertext, _| {
		ciphertext.truncate(ciphertext.len() - COMMITMENT_SIZE);
	});
}

#[test]
fn tag_corruption_is_rejected() {
	assert_server_frame_tampering_is_rejected(|ciphertext, _| {
		let tag_start = ciphertext.len() - COMMITMENT_SIZE - TAG_SIZE;
		ciphertext[tag_start] ^= 0x01;
	});
}

#[test]
fn swapped_tag_and_commitment_are_rejected() {
	assert_server_frame_tampering_is_rejected(|ciphertext, _| {
		let tag_start = ciphertext.len() - COMMITMENT_SIZE - TAG_SIZE;
		let commitment_start = ciphertext.len() - COMMITMENT_SIZE;
		let mut swapped = ciphertext[..tag_start].to_vec();
		swapped.extend_from_slice(&ciphertext[commitment_start..]);
		swapped.extend_from_slice(&ciphertext[tag_start..commitment_start]);
		*ciphertext = swapped;
	});
}

#[test]
fn beacon_rejects_tampered_registration_ephemeral_key() {
	let (mut beacon, response) = pending_registration(b"initial message");
	let mut parts = parse_kex_response(&response);
	parts.ephemeral_key[0] ^= 0x01;

	assert!(
		beacon
			.finish_registration(&serialize_kex_response(&parts))
			.is_none()
	);
}

#[test]
fn beacon_rejects_tampered_registration_kem_ciphertext() {
	let (mut beacon, response) = pending_registration(b"initial message");
	let mut parts = parse_kex_response(&response);
	parts.kem_ciphertext[0] ^= 0x01;

	assert!(
		beacon
			.finish_registration(&serialize_kex_response(&parts))
			.is_none()
	);
}

#[test]
fn failed_initial_ciphertext_is_terminal_and_clears_registration_state() {
	let (mut beacon, response) = pending_registration(b"initial message");
	let mut parts = parse_kex_response(&response);
	parts.app_ciphertext[0] ^= 0x01;

	assert!(
		beacon
			.finish_registration(&serialize_kex_response(&parts))
			.is_none()
	);
	assert!(beacon.finish_registration(&response).is_none());
	assert!(beacon.get_registration_bundle().is_none());
	assert!(beacon.get_prekey_pk().is_none());
	assert!(beacon.get_prekey_sk().is_none());
	assert!(beacon.get_onetime_pk().is_none());
	assert!(beacon.get_onetime_sk().is_none());
	assert!(beacon.pq_pk().is_none());
	assert!(beacon.pq_sk().is_none());
	assert!(beacon.associated_data().is_none());

	assert!(beacon.ratchet_status().is_none());
	assert!(beacon.encrypt_message(b"aborted").is_none());
}

#[test]
fn malformed_initial_ciphertext_is_rejected_without_panicking() {
	for malformed_app_ciphertext in [
		Vec::new(),
		vec![0xA5; 79],
		serialize_crypto_frame(1, SERVER_KID, &[0xA5; 79]),
	] {
		let (mut beacon, response) = pending_registration(b"initial message");
		let mut parts = parse_kex_response(&response);
		parts.app_ciphertext = malformed_app_ciphertext;

		let result = catch_unwind(AssertUnwindSafe(|| {
			beacon.finish_registration(&serialize_kex_response(&parts))
		}));
		assert!(matches!(result, Ok(None)));
		assert!(beacon.get_onetime_sk().is_none());
		assert!(beacon.pq_sk().is_none());
		assert!(beacon.associated_data().is_none());
	}
}

#[test]
fn beacon_rejects_tampered_registration_key_id() {
	let (mut beacon, response) = pending_registration(b"initial message");
	let mut parts = parse_kex_response(&response);
	parts.key_id = parts.key_id.wrapping_add(1);

	assert!(
		beacon
			.finish_registration(&serialize_kex_response(&parts))
			.is_none()
	);
}

#[test]
fn beacon_generates_only_one_registration_bundle() {
	let (_, mut beacon) = new_pair();
	assert!(beacon.get_registration_bundle().is_some());
	assert!(beacon.get_registration_bundle().is_none());
}

#[test]
fn server_rejects_replayed_registration_bundle() {
	let (mut server, mut beacon) = new_pair();
	let phase_1 = beacon.get_registration_bundle().unwrap();

	assert!(server.get_shared_secret(&phase_1).is_some());
	assert!(server.get_shared_secret(&phase_1).is_none());
}

#[test]
fn registration_request_field_lengths_reject_before_consuming_replay_material() {
	let (mut server, mut beacon) = new_pair();
	let original = beacon.get_registration_bundle().unwrap();
	let init = decode_init(&original).unwrap();
	let fields = [
		init.identity_key,
		init.pre_key,
		init.one_time_key,
		init.pq_key,
	];
	assert_eq!(fields.map(<[u8]>::len), [33, 98, 98, 1_249]);
	assert_eq!(common::encode_unchecked_init(init), original);
	for field in 0..fields.len() {
		for length in [fields[field].len() - 1, fields[field].len() + 1] {
			let mut changed = fields.map(<[u8]>::to_vec);
			changed[field].resize(length, 0);
			let candidate = common::encode_unchecked_init(InitKex {
				identity_key: &changed[0],
				pre_key: &changed[1],
				one_time_key: &changed[2],
				pq_key: &changed[3],
			});
			assert!(
				decode_init(&candidate).is_none(),
				"field {field}, length {length}"
			);
			assert!(server.get_shared_secret(&candidate).is_none());
			assert_eq!(server.server_kid(), SERVER_KID);
			assert!(server.pk_by_kid(1).is_none());
		}
	}
	assert!(server.get_shared_secret(&original).is_some());
	assert!(server.get_shared_secret(&original).is_none());
}

#[test]
fn registration_response_field_length_rejection_is_terminal() {
	for (field, expected_length) in [(0, 32), (1, 32), (2, 1_088)] {
		for length in [expected_length - 1, expected_length + 1] {
			let (mut beacon, original) = pending_registration(b"field length checks");
			assert_eq!(
				common::encode_unchecked_response(decode_response(&original).unwrap()),
				original
			);
			let mut parts = parse_kex_response(&original);
			let changed = match field {
				0 => &mut parts.identity_key,
				1 => &mut parts.ephemeral_key,
				_ => &mut parts.kem_ciphertext,
			};
			assert_eq!(changed.len(), expected_length);
			changed.resize(length, 0);
			let candidate = common::encode_unchecked_response(KexResponse {
				identity_key: &parts.identity_key,
				ephemeral_key: &parts.ephemeral_key,
				kem_cipher_text: &parts.kem_ciphertext,
				app_cipher_text: &parts.app_ciphertext,
				key_id: parts.key_id,
			});
			assert!(
				decode_response(&candidate).is_none(),
				"field {field}, length {length}"
			);
			assert!(beacon.finish_registration(&candidate).is_none());
			assert!(beacon.ratchet_status().is_none());
			assert!(beacon.associated_data().is_none());
			assert!(beacon.get_prekey_sk().is_none());
			assert!(beacon.get_onetime_sk().is_none());
			assert!(beacon.pq_sk().is_none());
			assert!(beacon.finish_registration(&original).is_none());
		}
	}
}

#[test]
fn replay_detection_uses_identity_and_one_time_key_across_distinct_signed_prekeys() {
	use beaconcrypt_core::pqxdh::{KEM_TYPE_X25519, KEY_ROLE_PREKEY};
	use libsodium_rs::{crypto_kx, crypto_sign};

	let seed = [0x4a; 32];
	let mut server = Server::new(SERVER_KID, Some(&seed));
	let mut beacon = Beacon::new(SERVER_KID, server.identity_pk().as_bytes());
	let original = beacon.get_registration_bundle().unwrap();
	let init = decode_init(&original).unwrap();
	let replacement = crypto_kx::KeyPair::generate().unwrap();
	let mut typed_prekey = vec![KEM_TYPE_X25519, KEY_ROLE_PREKEY];
	typed_prekey.extend_from_slice(replacement.public_key.as_bytes());
	let signed_prekey = crypto_sign::sign(&typed_prekey, beacon.identity_sk()).unwrap();
	assert_ne!(signed_prekey, init.pre_key);
	let mut changed = vec![0; original.len() + 64];
	let written = encode_init(
		InitKex {
			identity_key: init.identity_key,
			pre_key: &signed_prekey,
			one_time_key: init.one_time_key,
			pq_key: init.pq_key,
		},
		&mut changed,
	);
	assert_ne!(written, 0);
	changed.truncate(written);
	assert_ne!(original, changed);
	let changed_init = decode_init(&changed).unwrap();
	assert_eq!(changed_init.identity_key, init.identity_key);
	assert_eq!(changed_init.one_time_key, init.one_time_key);

	// Both canonical requests are valid before their shared registration material is consumed.
	for request in [&original, &changed] {
		let mut fresh = Server::new(SERVER_KID, Some(&seed));
		assert!(fresh.get_shared_secret(request).is_some());
	}
	assert!(server.get_shared_secret(&original).is_some());
	assert!(server.get_shared_secret(&changed).is_none());
	assert!(server.get_shared_secret(&original).is_none());
}

#[test]
fn malformed_registration_framing_does_not_consume_replay_material() {
	let (mut server, mut beacon) = new_pair();
	let phase_1 = beacon.get_registration_bundle().unwrap();
	let before = server.server_kid();
	for cut in 0..phase_1.len() {
		assert!(server.get_shared_secret(&phase_1[..cut]).is_none());
		assert_eq!(server.server_kid(), before);
	}
	for suffix in [&[0u8][..], &[0u8; 8][..], phase_1.as_slice()] {
		let mut malformed = phase_1.clone();
		malformed.extend_from_slice(suffix);
		assert!(server.get_shared_secret(&malformed).is_none());
		assert_eq!(server.server_kid(), before);
	}
	let pending = server.get_shared_secret(&phase_1).unwrap();
	assert!(server.get_shared_secret(&phase_1).is_none());
	let response = server.build_registration_response(pending, None).unwrap();
	assert_eq!(
		beacon.finish_registration(&response.serialized),
		Some(vec![0xFF])
	);
}

#[test]
fn server_key_id_exhaustion_does_not_wrap_or_commit_a_peer() {
	let server_identity_kid = u64::MAX - 1;
	let mut server = Server::new(server_identity_kid, None);
	let server_id = server.identity_pk().to_owned();
	let mut last_beacon = Beacon::new(server_identity_kid, server_id.as_bytes());

	let last = register_beacon(&mut server, &mut last_beacon, None);
	assert_eq!(last.kid, u64::MAX);
	assert_eq!(server.server_kid(), u64::MAX);

	let mut rejected_beacon = Beacon::new(server_identity_kid, server_id.as_bytes());
	let phase_1 = rejected_beacon.get_registration_bundle().unwrap();
	let pending = server.get_shared_secret(&phase_1).unwrap();

	assert!(server.build_registration_response(pending, None).is_none());
	assert_eq!(server.server_kid(), u64::MAX);
	assert!(server.pk_by_kid(0).is_none());
}

#[test]
fn encrypted_message_cannot_be_relabelled_to_an_unknown_key_id() {
	let (mut server, mut beacon) = new_pair();
	let registration = register_beacon(&mut server, &mut beacon, None);
	let ciphertext = beacon
		.encrypt_for_test(b"authenticated beacon message", SERVER_KID)
		.unwrap();
	let relabelled = rewrite_crypto_frame_key_id(&ciphertext, registration.kid + 1);

	let receive_before = receive_state(&server, registration.kid);
	assert!(server.decrypt_for_test(&relabelled).is_none());
	assert_eq!(receive_state(&server, registration.kid), receive_before);
	assert_eq!(
		server.decrypt_for_test(&ciphertext).unwrap().plaintext,
		b"authenticated beacon message"
	);
}