use crate::message::{ContentType, MessagePriority};
use base64::{engine::general_purpose::STANDARD as BASE64, Engine as _};
use serde::{Deserialize, Serialize};
pub const WIRE_V1_MAGIC: u8 = 0xF5;
pub const WIRE_VERSION_V1: u8 = 1;
pub(crate) const EXT_TAG_B64_TAIL: u16 = 1;
pub(crate) const EXT_TAG_REPLY_CONTEXT: u16 = 2;
const B64_TAIL_MIN_LEN: usize = 64;
pub(crate) fn split_b64_tail(content: &str) -> Option<(&str, Vec<u8>)> {
let bytes = content.as_bytes();
if bytes.len() < B64_TAIL_MIN_LEN {
return None;
}
let mut run_end = bytes.len();
let mut pad = 0;
while pad < 2 && run_end > 0 && bytes[run_end - 1] == b'=' {
run_end -= 1;
pad += 1;
}
let is_b64 = |b: u8| b.is_ascii_alphanumeric() || b == b'+' || b == b'/';
let mut start = run_end;
while start > 0 && is_b64(bytes[start - 1]) {
start -= 1;
}
start += (bytes.len() - start) % 4;
let tail = &content[start..];
if tail.len() < B64_TAIL_MIN_LEN {
return None;
}
let raw = BASE64.decode(tail).ok()?;
if BASE64.encode(&raw) != tail {
return None;
}
Some((&content[..start], raw))
}
pub(crate) fn encode_b64_tail(raw: &[u8]) -> String {
BASE64.encode(raw)
}
#[derive(Serialize, Deserialize)]
pub(crate) struct WireMessageV1 {
pub(crate) id: [u8; 16],
pub(crate) sender: String,
pub(crate) recipient: String,
pub(crate) app_id: String,
pub(crate) priority: u8,
pub(crate) ttl: u8,
pub(crate) hop_count: u8,
pub(crate) timestamp: i64,
pub(crate) lamport_clock: u64,
pub(crate) content_type: u8,
pub(crate) content: String,
pub(crate) binary_content: Option<Vec<u8>>,
pub(crate) media_metadata_json: Option<Vec<u8>>,
pub(crate) metadata: Vec<(String, String)>,
pub(crate) requires_ack: bool,
pub(crate) reply_to_msg: Option<[u8; 16]>,
pub(crate) forwarded_from_json: Option<Vec<u8>>,
pub(crate) ext: Vec<(u16, Vec<u8>)>,
}
pub(crate) fn priority_to_u8(p: MessagePriority) -> u8 {
match p {
MessagePriority::Low => 0,
MessagePriority::Medium => 1,
MessagePriority::High => 2,
MessagePriority::Critical => 3,
}
}
pub(crate) fn priority_from_u8(v: u8) -> MessagePriority {
match v {
0 => MessagePriority::Low,
1 => MessagePriority::Medium,
2 => MessagePriority::High,
3 => MessagePriority::Critical,
_ => MessagePriority::Medium,
}
}
pub(crate) fn content_type_to_u8(c: ContentType) -> u8 {
match c {
ContentType::Text => 0,
ContentType::Image => 1,
ContentType::Video => 2,
ContentType::Audio => 3,
ContentType::VoiceNote => 4,
ContentType::VideoNote => 5,
ContentType::File => 6,
ContentType::FileChunk => 7,
ContentType::Poll => 8,
}
}
pub(crate) fn content_type_from_u8(v: u8) -> ContentType {
match v {
0 => ContentType::Text,
1 => ContentType::Image,
2 => ContentType::Video,
3 => ContentType::Audio,
4 => ContentType::VoiceNote,
5 => ContentType::VideoNote,
6 => ContentType::File,
7 => ContentType::FileChunk,
8 => ContentType::Poll,
_ => ContentType::File,
}
}
pub(crate) fn encode(wire: &WireMessageV1) -> crate::Result<Vec<u8>> {
let body = postcard::to_allocvec(wire)
.map_err(|e| crate::Error::SerializationError(format!("wire v1 encode: {e}")))?;
let mut out = Vec::with_capacity(body.len() + 1);
out.push(WIRE_V1_MAGIC);
out.extend_from_slice(&body);
Ok(out)
}
pub(crate) fn decode(data: &[u8]) -> crate::Result<WireMessageV1> {
match data.split_first() {
Some((&WIRE_V1_MAGIC, body)) => postcard::from_bytes::<WireMessageV1>(body)
.map_err(|e| crate::Error::DeserializationError(format!("wire v1 decode: {e}"))),
Some((other, _)) => Err(crate::Error::DeserializationError(format!(
"not a wire v1 frame: unexpected leading byte 0x{other:02X}"
))),
None => Err(crate::Error::DeserializationError(
"not a wire v1 frame: empty input".to_string(),
)),
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::types::LAMPORT_CLOCK_MAX;
use crate::{
AppId, ForwardInfo, LamportClock, MediaMetadata, Message, MessageId, ReplyContext,
Timestamp, UserId, MAX_ID_LEN, TTL,
};
fn sample_message() -> Message {
let mut m = Message::new(
UserId::new("alice").unwrap(),
UserId::new("bob").unwrap(),
AppId::new("com.example.chat").unwrap(),
"hello world",
);
m.priority = MessagePriority::High;
m.ttl = TTL::new(5).unwrap();
m.content_type = ContentType::Image;
m.binary_content = Some(vec![0, 1, 2, 3, 255, 254]);
m.lamport_clock = LamportClock::from_value(42);
m.metadata.insert("k1".into(), "v1".into());
m.metadata.insert("k2".into(), "v2".into());
m.requires_ack = false;
m.reply_to_msg = Some(MessageId::new());
m.media_metadata = Some(MediaMetadata {
mime_type: "image/png".into(),
file_name: "x.png".into(),
file_size: 1234,
duration_ms: None,
width: Some(640),
height: Some(480),
thumbnail_base64: None,
media_id: Some("media-1".into()),
download_url: Some("https://cdn.example/media-1".into()),
thumbnail_url: None,
encryption_key: Some("a2V5".into()),
iv: None,
ciphertext_hash: None,
sticker_provider: None,
sticker_remote_id: None,
sticker_kind: None,
});
m.forwarded_from = Some(ForwardInfo {
original_sender: UserId::new("carol").unwrap(),
original_message_id: MessageId::new(),
original_timestamp: Timestamp::from_millis(1_700_000_000_000),
forward_count: 2,
});
m
}
fn base_dto() -> WireMessageV1 {
WireMessageV1 {
id: [7u8; 16],
sender: "alice".into(),
recipient: "bob".into(),
app_id: "app".into(),
priority: 1,
ttl: 8,
hop_count: 0,
timestamp: 1,
lamport_clock: 0,
content_type: 0,
content: "hi".into(),
binary_content: None,
media_metadata_json: None,
metadata: Vec::new(),
requires_ack: true,
reply_to_msg: None,
forwarded_from_json: None,
ext: Vec::new(),
}
}
#[test]
fn wire_v1_round_trips_all_fields() {
let original = sample_message();
let bytes = original.to_wire_v1_bytes().unwrap();
assert_eq!(bytes[0], WIRE_V1_MAGIC);
let decoded = Message::from_wire_v1_bytes(&bytes).unwrap();
assert_eq!(decoded.id, original.id);
assert_eq!(decoded.sender, original.sender);
assert_eq!(decoded.recipient, original.recipient);
assert_eq!(decoded.app_id, original.app_id);
assert_eq!(decoded.priority, original.priority);
assert_eq!(decoded.ttl, original.ttl);
assert_eq!(decoded.hop_count, original.hop_count);
assert_eq!(decoded.timestamp, original.timestamp);
assert_eq!(decoded.lamport_clock, original.lamport_clock);
assert_eq!(decoded.content_type, original.content_type);
assert_eq!(decoded.content, original.content);
assert_eq!(decoded.binary_content, original.binary_content);
assert_eq!(decoded.metadata, original.metadata);
assert_eq!(decoded.requires_ack, original.requires_ack);
assert_eq!(decoded.reply_to_msg, original.reply_to_msg);
assert_eq!(decoded.forwarded_from, original.forwarded_from);
assert_eq!(
serde_json::to_vec(&decoded.media_metadata).unwrap(),
serde_json::to_vec(&original.media_metadata).unwrap(),
);
assert!(decoded.transport_peer_id().is_none());
assert_eq!(decoded.wire_codec(), crate::WireCodec::Json);
}
#[test]
fn wire_v1_round_trips_minimal_defaults() {
let original = Message::new(
UserId::new("a").unwrap(),
UserId::new("b").unwrap(),
AppId::new("c").unwrap(),
"hi",
);
let decoded = Message::from_wire_v1_bytes(&original.to_wire_v1_bytes().unwrap()).unwrap();
assert_eq!(decoded.id, original.id);
assert!(decoded.requires_ack);
assert!(decoded.binary_content.is_none());
assert!(decoded.media_metadata.is_none());
assert!(decoded.forwarded_from.is_none());
assert!(decoded.reply_to_msg.is_none());
assert!(decoded.metadata.is_empty());
}
#[test]
fn wire_v1_is_much_smaller_than_json() {
let m = Message::new(
UserId::new("a3f8c2d1-user-alice").unwrap(),
UserId::new("b7e4d9a2-user-bob").unwrap(),
AppId::new("com.example.chat").unwrap(),
"hi",
);
let json = m.to_bytes().unwrap();
let binary = m.to_wire_v1_bytes().unwrap();
assert!(
binary.len() * 2 < json.len(),
"binary {} should be < half of json {}",
binary.len(),
json.len()
);
}
#[test]
fn wire_v1_preserves_identifier_length_validation() {
let mut wire = base_dto();
wire.sender = "a".repeat(MAX_ID_LEN + 1);
let bytes = encode(&wire).unwrap();
assert!(Message::from_wire_v1_bytes(&bytes).is_err());
}
#[test]
fn wire_v1_preserves_identifier_char_validation() {
let mut wire = base_dto();
wire.recipient = "evil/path".into();
let bytes = encode(&wire).unwrap();
assert!(Message::from_wire_v1_bytes(&bytes).is_err());
}
#[test]
fn wire_v1_clamps_adversarial_lamport_clock() {
let mut wire = base_dto();
wire.lamport_clock = u64::MAX;
let bytes = encode(&wire).unwrap();
let decoded = Message::from_wire_v1_bytes(&bytes).unwrap();
assert_eq!(decoded.lamport_clock.value(), LAMPORT_CLOCK_MAX);
}
#[test]
fn wire_v1_tolerates_unknown_enum_discriminants() {
let mut wire = base_dto();
wire.content_type = 200;
wire.priority = 200;
let bytes = encode(&wire).unwrap();
let decoded = Message::from_wire_v1_bytes(&bytes).unwrap();
assert_eq!(decoded.content_type, ContentType::File);
assert_eq!(decoded.priority, MessagePriority::Medium);
}
#[test]
fn wire_v1_round_trips_poll_content_type() {
let mut wire = base_dto();
wire.content_type = content_type_to_u8(ContentType::Poll);
let bytes = encode(&wire).unwrap();
let decoded = Message::from_wire_v1_bytes(&bytes).unwrap();
assert_eq!(decoded.content_type, ContentType::Poll);
}
#[test]
fn wire_v1_decode_rejects_non_v1_frames() {
assert!(decode(&[0x7B, 0x00]).is_err()); assert!(decode(&[]).is_err()); }
#[test]
fn b64_tail_split_and_reconstruct_is_byte_identical() {
let raw: Vec<u8> = (0u8..50).collect();
let content = format!("__MLS_ENC__{}", BASE64.encode(&raw));
let (head, tail_raw) = split_b64_tail(&content).expect("tail must split");
assert_eq!(head, "__MLS_ENC__");
assert_eq!(tail_raw, raw);
assert_eq!(format!("{head}{}", encode_b64_tail(&tail_raw)), content);
}
#[test]
fn b64_tail_split_sheds_run_remainder_into_head_and_handles_headless() {
let content = format!("!{}", "A".repeat(66));
let (head, raw) = split_b64_tail(&content).expect("aligned suffix splits");
assert_eq!(head, "!AA");
assert_eq!(format!("{head}{}", encode_b64_tail(&raw)), content);
let content = "A".repeat(64);
let (head, raw) = split_b64_tail(&content).expect("headless splits");
assert!(head.is_empty());
assert_eq!(encode_b64_tail(&raw), content);
}
#[test]
fn b64_tail_split_is_utf8_boundary_safe() {
let content = format!("héllo → {}", "A".repeat(64));
let (head, raw) = split_b64_tail(&content).expect("ascii tail after unicode head");
assert_eq!(head, "héllo → ");
assert_eq!(format!("{head}{}", encode_b64_tail(&raw)), content);
}
#[test]
fn b64_tail_split_rejects_short_nonaligned_and_noncanonical() {
assert!(split_b64_tail("AAAA").is_none());
assert!(split_b64_tail(&format!("!{}", "A".repeat(63))).is_none());
let json_like = format!("__MLS_ENC__{{\"ciphertext\":[{}]}}", "57,".repeat(60));
assert!(split_b64_tail(&json_like).is_none());
assert!(split_b64_tail(&format!("!{}=", "B".repeat(63))).is_none());
assert!(split_b64_tail(&format!("{}={}", "A".repeat(64), "A".repeat(3))).is_none());
assert!(split_b64_tail(&"=".repeat(64)).is_none());
}
#[test]
fn wire_v1_round_trips_b64_tail_content_through_message() {
let mut m = Message::new(
UserId::new("alice").unwrap(),
UserId::new("bob").unwrap(),
AppId::new("app").unwrap(),
format!("__MLS_ENC__{}", BASE64.encode(vec![7u8; 90])),
);
m.requires_ack = true;
let bytes = m.to_wire_v1_bytes().unwrap();
let wire = decode(&bytes).unwrap();
assert_eq!(wire.content, "__MLS_ENC__");
assert_eq!(wire.ext.len(), 1);
assert_eq!(wire.ext[0].0, EXT_TAG_B64_TAIL);
assert_eq!(wire.ext[0].1, vec![7u8; 90]);
let decoded = Message::from_wire_v1_bytes(&bytes).unwrap();
assert_eq!(decoded.content, m.content);
assert!(bytes.len() < m.to_bytes().unwrap().len());
}
#[test]
fn wire_v1_plain_content_emits_no_ext_entries() {
let m = sample_message();
let wire = decode(&m.to_wire_v1_bytes().unwrap()).unwrap();
assert!(wire.ext.is_empty());
}
#[test]
fn wire_v1_decode_ignores_unknown_ext_tags_and_extra_tag1_entries() {
let mut wire = base_dto();
wire.content = "head:".into();
wire.ext = vec![
(99, vec![1, 2, 3]),
(EXT_TAG_B64_TAIL, b"AB".to_vec()),
(EXT_TAG_B64_TAIL, b"ZZ".to_vec()),
];
let bytes = encode(&wire).unwrap();
let decoded = Message::from_wire_v1_bytes(&bytes).unwrap();
assert_eq!(decoded.content, format!("head:{}", BASE64.encode(b"AB")));
}
fn sample_reply_context() -> ReplyContext {
ReplyContext {
sender: UserId::new("carol").unwrap(),
text: "original text".into(),
timestamp: Some(Timestamp::from_millis(1_700_000_000_000)),
reply_media_label: Some("x.png".into()),
reply_content_type: Some("image".into()),
}
}
#[test]
fn wire_v1_round_trips_reply_context() {
let mut m = sample_message();
m.reply_context = Some(sample_reply_context());
let decoded = Message::from_wire_v1_bytes(&m.to_wire_v1_bytes().unwrap()).unwrap();
assert_eq!(decoded.reply_context, m.reply_context);
let mut m = Message::new(
UserId::new("a").unwrap(),
UserId::new("b").unwrap(),
AppId::new("c").unwrap(),
"hi",
);
m.reply_context = Some(ReplyContext {
sender: UserId::new("carol").unwrap(),
text: "quoted".into(),
timestamp: None,
reply_media_label: None,
reply_content_type: None,
});
let decoded = Message::from_wire_v1_bytes(&m.to_wire_v1_bytes().unwrap()).unwrap();
assert_eq!(decoded.reply_context, m.reply_context);
}
#[test]
fn wire_v1_reply_context_coexists_with_b64_tail() {
let mut m = Message::new(
UserId::new("alice").unwrap(),
UserId::new("bob").unwrap(),
AppId::new("app").unwrap(),
format!("__MLS_ENC__{}", BASE64.encode(vec![7u8; 90])),
);
m.reply_context = Some(sample_reply_context());
let bytes = m.to_wire_v1_bytes().unwrap();
let wire = decode(&bytes).unwrap();
assert_eq!(wire.ext.len(), 2);
assert_eq!(wire.ext[0].0, EXT_TAG_B64_TAIL);
assert_eq!(wire.ext[1].0, EXT_TAG_REPLY_CONTEXT);
let decoded = Message::from_wire_v1_bytes(&bytes).unwrap();
assert_eq!(decoded.content, m.content);
assert_eq!(decoded.reply_context, m.reply_context);
}
#[test]
fn wire_v1_honors_only_first_reply_context_entry() {
let mut wire = base_dto();
wire.ext = vec![
(
EXT_TAG_REPLY_CONTEXT,
br#"{"sender":"carol","text":"first"}"#.to_vec(),
),
(
EXT_TAG_REPLY_CONTEXT,
br#"{"sender":"mallory","text":"second"}"#.to_vec(),
),
];
let decoded = Message::from_wire_v1_bytes(&encode(&wire).unwrap()).unwrap();
let rc = decoded.reply_context.expect("first entry honored");
assert_eq!(rc.sender.as_str(), "carol");
assert_eq!(rc.text, "first");
}
#[test]
fn wire_v1_rejects_malformed_reply_context_ext() {
let mut wire = base_dto();
wire.ext = vec![(EXT_TAG_REPLY_CONTEXT, b"not json".to_vec())];
assert!(Message::from_wire_v1_bytes(&encode(&wire).unwrap()).is_err());
let mut wire = base_dto();
wire.ext = vec![(
EXT_TAG_REPLY_CONTEXT,
br#"{"sender":"evil/path","text":"t"}"#.to_vec(),
)];
assert!(Message::from_wire_v1_bytes(&encode(&wire).unwrap()).is_err());
}
#[test]
fn wire_v1_golden_layout_with_reply_context_ext_is_frozen() {
let rc_json: &[u8] = br#"{"sender":"c","text":"t"}"#;
let mut wire = base_dto();
wire.content = String::new();
wire.ext = vec![(EXT_TAG_REPLY_CONTEXT, rc_json.to_vec())];
let mut expected_tail = vec![
0x01, 0x02, 0x19, ];
expected_tail.extend_from_slice(rc_json); let encoded = encode(&wire).unwrap();
assert_eq!(
&encoded[encoded.len() - expected_tail.len()..],
expected_tail
);
let decoded = Message::from_wire_v1_bytes(&encoded).unwrap();
let rc = decoded.reply_context.expect("reply context decoded");
assert_eq!(rc.sender.as_str(), "c");
assert_eq!(rc.text, "t");
}
#[test]
fn wire_v1_golden_layout_with_b64_tail_ext_is_frozen() {
let mut wire = base_dto();
wire.content = String::new();
wire.ext = vec![(EXT_TAG_B64_TAIL, vec![0xDE, 0xAD])];
let expected_tail: &[u8] = &[
0x01, 0x01, 0x02, 0xDE, 0xAD, ];
let encoded = encode(&wire).unwrap();
assert_eq!(
&encoded[encoded.len() - expected_tail.len()..],
expected_tail
);
let decoded = Message::from_wire_v1_bytes(&encoded).unwrap();
assert_eq!(decoded.content, BASE64.encode([0xDE, 0xAD]));
}
#[test]
fn wire_v1_golden_layout_is_frozen() {
let wire = WireMessageV1 {
id: [0u8; 16],
sender: "a".into(),
recipient: "b".into(),
app_id: "c".into(),
priority: 1,
ttl: 8,
hop_count: 0,
timestamp: 0,
lamport_clock: 0,
content_type: 0,
content: String::new(),
binary_content: None,
media_metadata_json: None,
metadata: Vec::new(),
requires_ack: true,
reply_to_msg: None,
forwarded_from_json: None,
ext: Vec::new(),
};
let expected: &[u8] = &[
0xF5, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0x01, b'a', 0x01, b'b', 0x01, b'c', 0x01, 0x08, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, ];
assert_eq!(encode(&wire).unwrap(), expected);
}
}