1use crate::message::{ContentType, MessagePriority};
41use base64::{engine::general_purpose::STANDARD as BASE64, Engine as _};
42use serde::{Deserialize, Serialize};
43
44pub const WIRE_V1_MAGIC: u8 = 0xF5;
52
53pub const WIRE_VERSION_V1: u8 = 1;
58
59pub(crate) const EXT_TAG_B64_TAIL: u16 = 1;
84
85pub(crate) const EXT_TAG_REPLY_CONTEXT: u16 = 2;
99
100const B64_TAIL_MIN_LEN: usize = 64;
105
106pub(crate) fn split_b64_tail(content: &str) -> Option<(&str, Vec<u8>)> {
119 let bytes = content.as_bytes();
120 if bytes.len() < B64_TAIL_MIN_LEN {
121 return None;
122 }
123 let mut run_end = bytes.len();
125 let mut pad = 0;
126 while pad < 2 && run_end > 0 && bytes[run_end - 1] == b'=' {
127 run_end -= 1;
128 pad += 1;
129 }
130 let is_b64 = |b: u8| b.is_ascii_alphanumeric() || b == b'+' || b == b'/';
131 let mut start = run_end;
132 while start > 0 && is_b64(bytes[start - 1]) {
133 start -= 1;
134 }
135 start += (bytes.len() - start) % 4;
139 let tail = &content[start..];
140 if tail.len() < B64_TAIL_MIN_LEN {
141 return None;
142 }
143 let raw = BASE64.decode(tail).ok()?;
144 if BASE64.encode(&raw) != tail {
145 return None;
146 }
147 Some((&content[..start], raw))
148}
149
150pub(crate) fn encode_b64_tail(raw: &[u8]) -> String {
153 BASE64.encode(raw)
154}
155
156#[derive(Serialize, Deserialize)]
162pub(crate) struct WireMessageV1 {
163 pub(crate) id: [u8; 16],
164 pub(crate) sender: String,
165 pub(crate) recipient: String,
166 pub(crate) app_id: String,
167 pub(crate) priority: u8,
168 pub(crate) ttl: u8,
169 pub(crate) hop_count: u8,
170 pub(crate) timestamp: i64,
171 pub(crate) lamport_clock: u64,
172 pub(crate) content_type: u8,
173 pub(crate) content: String,
174 pub(crate) binary_content: Option<Vec<u8>>,
175 pub(crate) media_metadata_json: Option<Vec<u8>>,
179 pub(crate) metadata: Vec<(String, String)>,
181 pub(crate) requires_ack: bool,
182 pub(crate) reply_to_msg: Option<[u8; 16]>,
183 pub(crate) forwarded_from_json: Option<Vec<u8>>,
186 pub(crate) ext: Vec<(u16, Vec<u8>)>,
188}
189
190pub(crate) fn priority_to_u8(p: MessagePriority) -> u8 {
192 match p {
193 MessagePriority::Low => 0,
194 MessagePriority::Medium => 1,
195 MessagePriority::High => 2,
196 MessagePriority::Critical => 3,
197 }
198}
199
200pub(crate) fn priority_from_u8(v: u8) -> MessagePriority {
203 match v {
204 0 => MessagePriority::Low,
205 1 => MessagePriority::Medium,
206 2 => MessagePriority::High,
207 3 => MessagePriority::Critical,
208 _ => MessagePriority::Medium,
209 }
210}
211
212pub(crate) fn content_type_to_u8(c: ContentType) -> u8 {
215 match c {
216 ContentType::Text => 0,
217 ContentType::Image => 1,
218 ContentType::Video => 2,
219 ContentType::Audio => 3,
220 ContentType::VoiceNote => 4,
221 ContentType::VideoNote => 5,
222 ContentType::File => 6,
223 ContentType::FileChunk => 7,
224 ContentType::Poll => 8,
225 }
226}
227
228pub(crate) fn content_type_from_u8(v: u8) -> ContentType {
231 match v {
232 0 => ContentType::Text,
233 1 => ContentType::Image,
234 2 => ContentType::Video,
235 3 => ContentType::Audio,
236 4 => ContentType::VoiceNote,
237 5 => ContentType::VideoNote,
238 6 => ContentType::File,
239 7 => ContentType::FileChunk,
240 8 => ContentType::Poll,
241 _ => ContentType::File,
242 }
243}
244
245pub(crate) fn encode(wire: &WireMessageV1) -> crate::Result<Vec<u8>> {
247 let body = postcard::to_allocvec(wire)
248 .map_err(|e| crate::Error::SerializationError(format!("wire v1 encode: {e}")))?;
249 let mut out = Vec::with_capacity(body.len() + 1);
250 out.push(WIRE_V1_MAGIC);
251 out.extend_from_slice(&body);
252 Ok(out)
253}
254
255pub(crate) fn decode(data: &[u8]) -> crate::Result<WireMessageV1> {
260 match data.split_first() {
261 Some((&WIRE_V1_MAGIC, body)) => postcard::from_bytes::<WireMessageV1>(body)
262 .map_err(|e| crate::Error::DeserializationError(format!("wire v1 decode: {e}"))),
263 Some((other, _)) => Err(crate::Error::DeserializationError(format!(
264 "not a wire v1 frame: unexpected leading byte 0x{other:02X}"
265 ))),
266 None => Err(crate::Error::DeserializationError(
267 "not a wire v1 frame: empty input".to_string(),
268 )),
269 }
270}
271
272#[cfg(test)]
273mod tests {
274 use super::*;
275 use crate::types::LAMPORT_CLOCK_MAX;
276 use crate::{
277 AppId, ForwardInfo, LamportClock, MediaMetadata, Message, MessageId, ReplyContext,
278 Timestamp, UserId, MAX_ID_LEN, TTL,
279 };
280
281 fn sample_message() -> Message {
282 let mut m = Message::new(
283 UserId::new("alice").unwrap(),
284 UserId::new("bob").unwrap(),
285 AppId::new("com.example.chat").unwrap(),
286 "hello world",
287 );
288 m.priority = MessagePriority::High;
289 m.ttl = TTL::new(5).unwrap();
290 m.content_type = ContentType::Image;
291 m.binary_content = Some(vec![0, 1, 2, 3, 255, 254]);
292 m.lamport_clock = LamportClock::from_value(42);
293 m.metadata.insert("k1".into(), "v1".into());
294 m.metadata.insert("k2".into(), "v2".into());
295 m.requires_ack = false;
296 m.reply_to_msg = Some(MessageId::new());
297 m.media_metadata = Some(MediaMetadata {
298 mime_type: "image/png".into(),
299 file_name: "x.png".into(),
300 file_size: 1234,
301 duration_ms: None,
302 width: Some(640),
303 height: Some(480),
304 thumbnail_base64: None,
305 media_id: Some("media-1".into()),
306 download_url: Some("https://cdn.example/media-1".into()),
307 thumbnail_url: None,
308 encryption_key: Some("a2V5".into()),
309 iv: None,
310 ciphertext_hash: None,
311 sticker_provider: None,
312 sticker_remote_id: None,
313 sticker_kind: None,
314 });
315 m.forwarded_from = Some(ForwardInfo {
316 original_sender: UserId::new("carol").unwrap(),
317 original_message_id: MessageId::new(),
318 original_timestamp: Timestamp::from_millis(1_700_000_000_000),
319 forward_count: 2,
320 });
321 m
322 }
323
324 fn base_dto() -> WireMessageV1 {
326 WireMessageV1 {
327 id: [7u8; 16],
328 sender: "alice".into(),
329 recipient: "bob".into(),
330 app_id: "app".into(),
331 priority: 1,
332 ttl: 8,
333 hop_count: 0,
334 timestamp: 1,
335 lamport_clock: 0,
336 content_type: 0,
337 content: "hi".into(),
338 binary_content: None,
339 media_metadata_json: None,
340 metadata: Vec::new(),
341 requires_ack: true,
342 reply_to_msg: None,
343 forwarded_from_json: None,
344 ext: Vec::new(),
345 }
346 }
347
348 #[test]
349 fn wire_v1_round_trips_all_fields() {
350 let original = sample_message();
351 let bytes = original.to_wire_v1_bytes().unwrap();
352 assert_eq!(bytes[0], WIRE_V1_MAGIC);
353 let decoded = Message::from_wire_v1_bytes(&bytes).unwrap();
354
355 assert_eq!(decoded.id, original.id);
356 assert_eq!(decoded.sender, original.sender);
357 assert_eq!(decoded.recipient, original.recipient);
358 assert_eq!(decoded.app_id, original.app_id);
359 assert_eq!(decoded.priority, original.priority);
360 assert_eq!(decoded.ttl, original.ttl);
361 assert_eq!(decoded.hop_count, original.hop_count);
362 assert_eq!(decoded.timestamp, original.timestamp);
363 assert_eq!(decoded.lamport_clock, original.lamport_clock);
364 assert_eq!(decoded.content_type, original.content_type);
365 assert_eq!(decoded.content, original.content);
366 assert_eq!(decoded.binary_content, original.binary_content);
367 assert_eq!(decoded.metadata, original.metadata);
368 assert_eq!(decoded.requires_ack, original.requires_ack);
369 assert_eq!(decoded.reply_to_msg, original.reply_to_msg);
370 assert_eq!(decoded.forwarded_from, original.forwarded_from);
371 assert_eq!(
373 serde_json::to_vec(&decoded.media_metadata).unwrap(),
374 serde_json::to_vec(&original.media_metadata).unwrap(),
375 );
376 assert!(decoded.transport_peer_id().is_none());
378 assert_eq!(decoded.wire_codec(), crate::WireCodec::Json);
382 }
383
384 #[test]
385 fn wire_v1_round_trips_minimal_defaults() {
386 let original = Message::new(
387 UserId::new("a").unwrap(),
388 UserId::new("b").unwrap(),
389 AppId::new("c").unwrap(),
390 "hi",
391 );
392 let decoded = Message::from_wire_v1_bytes(&original.to_wire_v1_bytes().unwrap()).unwrap();
393 assert_eq!(decoded.id, original.id);
394 assert!(decoded.requires_ack);
396 assert!(decoded.binary_content.is_none());
397 assert!(decoded.media_metadata.is_none());
398 assert!(decoded.forwarded_from.is_none());
399 assert!(decoded.reply_to_msg.is_none());
400 assert!(decoded.metadata.is_empty());
401 }
402
403 #[test]
404 fn wire_v1_is_much_smaller_than_json() {
405 let m = Message::new(
406 UserId::new("a3f8c2d1-user-alice").unwrap(),
407 UserId::new("b7e4d9a2-user-bob").unwrap(),
408 AppId::new("com.example.chat").unwrap(),
409 "hi",
410 );
411 let json = m.to_bytes().unwrap();
412 let binary = m.to_wire_v1_bytes().unwrap();
413 assert!(
414 binary.len() * 2 < json.len(),
415 "binary {} should be < half of json {}",
416 binary.len(),
417 json.len()
418 );
419 }
420
421 #[test]
422 fn wire_v1_preserves_identifier_length_validation() {
423 let mut wire = base_dto();
426 wire.sender = "a".repeat(MAX_ID_LEN + 1);
427 let bytes = encode(&wire).unwrap();
428 assert!(Message::from_wire_v1_bytes(&bytes).is_err());
429 }
430
431 #[test]
432 fn wire_v1_preserves_identifier_char_validation() {
433 let mut wire = base_dto();
434 wire.recipient = "evil/path".into();
435 let bytes = encode(&wire).unwrap();
436 assert!(Message::from_wire_v1_bytes(&bytes).is_err());
437 }
438
439 #[test]
440 fn wire_v1_clamps_adversarial_lamport_clock() {
441 let mut wire = base_dto();
442 wire.lamport_clock = u64::MAX;
443 let bytes = encode(&wire).unwrap();
444 let decoded = Message::from_wire_v1_bytes(&bytes).unwrap();
445 assert_eq!(decoded.lamport_clock.value(), LAMPORT_CLOCK_MAX);
446 }
447
448 #[test]
449 fn wire_v1_tolerates_unknown_enum_discriminants() {
450 let mut wire = base_dto();
451 wire.content_type = 200;
452 wire.priority = 200;
453 let bytes = encode(&wire).unwrap();
454 let decoded = Message::from_wire_v1_bytes(&bytes).unwrap();
455 assert_eq!(decoded.content_type, ContentType::File);
456 assert_eq!(decoded.priority, MessagePriority::Medium);
457 }
458
459 #[test]
460 fn wire_v1_round_trips_poll_content_type() {
461 let mut wire = base_dto();
462 wire.content_type = content_type_to_u8(ContentType::Poll);
463 let bytes = encode(&wire).unwrap();
464 let decoded = Message::from_wire_v1_bytes(&bytes).unwrap();
465 assert_eq!(decoded.content_type, ContentType::Poll);
466 }
467
468 #[test]
469 fn wire_v1_decode_rejects_non_v1_frames() {
470 assert!(decode(&[0x7B, 0x00]).is_err()); assert!(decode(&[]).is_err()); }
473
474 #[test]
475 fn b64_tail_split_and_reconstruct_is_byte_identical() {
476 let raw: Vec<u8> = (0u8..50).collect();
478 let content = format!("__MLS_ENC__{}", BASE64.encode(&raw));
479 let (head, tail_raw) = split_b64_tail(&content).expect("tail must split");
480 assert_eq!(head, "__MLS_ENC__");
481 assert_eq!(tail_raw, raw);
482 assert_eq!(format!("{head}{}", encode_b64_tail(&tail_raw)), content);
483 }
484
485 #[test]
486 fn b64_tail_split_sheds_run_remainder_into_head_and_handles_headless() {
487 let content = format!("!{}", "A".repeat(66));
490 let (head, raw) = split_b64_tail(&content).expect("aligned suffix splits");
491 assert_eq!(head, "!AA");
492 assert_eq!(format!("{head}{}", encode_b64_tail(&raw)), content);
493
494 let content = "A".repeat(64);
496 let (head, raw) = split_b64_tail(&content).expect("headless splits");
497 assert!(head.is_empty());
498 assert_eq!(encode_b64_tail(&raw), content);
499 }
500
501 #[test]
502 fn b64_tail_split_is_utf8_boundary_safe() {
503 let content = format!("héllo → {}", "A".repeat(64));
504 let (head, raw) = split_b64_tail(&content).expect("ascii tail after unicode head");
505 assert_eq!(head, "héllo → ");
506 assert_eq!(format!("{head}{}", encode_b64_tail(&raw)), content);
507 }
508
509 #[test]
510 fn b64_tail_split_rejects_short_nonaligned_and_noncanonical() {
511 assert!(split_b64_tail("AAAA").is_none());
513 assert!(split_b64_tail(&format!("!{}", "A".repeat(63))).is_none());
515 let json_like = format!("__MLS_ENC__{{\"ciphertext\":[{}]}}", "57,".repeat(60));
517 assert!(split_b64_tail(&json_like).is_none());
518 assert!(split_b64_tail(&format!("!{}=", "B".repeat(63))).is_none());
521 assert!(split_b64_tail(&format!("{}={}", "A".repeat(64), "A".repeat(3))).is_none());
523 assert!(split_b64_tail(&"=".repeat(64)).is_none());
525 }
526
527 #[test]
528 fn wire_v1_round_trips_b64_tail_content_through_message() {
529 let mut m = Message::new(
530 UserId::new("alice").unwrap(),
531 UserId::new("bob").unwrap(),
532 AppId::new("app").unwrap(),
533 format!("__MLS_ENC__{}", BASE64.encode(vec![7u8; 90])),
534 );
535 m.requires_ack = true;
536 let bytes = m.to_wire_v1_bytes().unwrap();
537
538 let wire = decode(&bytes).unwrap();
540 assert_eq!(wire.content, "__MLS_ENC__");
541 assert_eq!(wire.ext.len(), 1);
542 assert_eq!(wire.ext[0].0, EXT_TAG_B64_TAIL);
543 assert_eq!(wire.ext[0].1, vec![7u8; 90]);
544
545 let decoded = Message::from_wire_v1_bytes(&bytes).unwrap();
547 assert_eq!(decoded.content, m.content);
548 assert!(bytes.len() < m.to_bytes().unwrap().len());
549 }
550
551 #[test]
552 fn wire_v1_plain_content_emits_no_ext_entries() {
553 let m = sample_message();
554 let wire = decode(&m.to_wire_v1_bytes().unwrap()).unwrap();
555 assert!(wire.ext.is_empty());
556 }
557
558 #[test]
559 fn wire_v1_decode_ignores_unknown_ext_tags_and_extra_tag1_entries() {
560 let mut wire = base_dto();
563 wire.content = "head:".into();
564 wire.ext = vec![
565 (99, vec![1, 2, 3]),
566 (EXT_TAG_B64_TAIL, b"AB".to_vec()),
567 (EXT_TAG_B64_TAIL, b"ZZ".to_vec()),
568 ];
569 let bytes = encode(&wire).unwrap();
570 let decoded = Message::from_wire_v1_bytes(&bytes).unwrap();
571 assert_eq!(decoded.content, format!("head:{}", BASE64.encode(b"AB")));
572 }
573
574 fn sample_reply_context() -> ReplyContext {
575 ReplyContext {
576 sender: UserId::new("carol").unwrap(),
577 text: "original text".into(),
578 timestamp: Some(Timestamp::from_millis(1_700_000_000_000)),
579 reply_media_label: Some("x.png".into()),
580 reply_content_type: Some("image".into()),
581 }
582 }
583
584 #[test]
585 fn wire_v1_round_trips_reply_context() {
586 let mut m = sample_message();
588 m.reply_context = Some(sample_reply_context());
589 let decoded = Message::from_wire_v1_bytes(&m.to_wire_v1_bytes().unwrap()).unwrap();
590 assert_eq!(decoded.reply_context, m.reply_context);
591
592 let mut m = Message::new(
594 UserId::new("a").unwrap(),
595 UserId::new("b").unwrap(),
596 AppId::new("c").unwrap(),
597 "hi",
598 );
599 m.reply_context = Some(ReplyContext {
600 sender: UserId::new("carol").unwrap(),
601 text: "quoted".into(),
602 timestamp: None,
603 reply_media_label: None,
604 reply_content_type: None,
605 });
606 let decoded = Message::from_wire_v1_bytes(&m.to_wire_v1_bytes().unwrap()).unwrap();
607 assert_eq!(decoded.reply_context, m.reply_context);
608 }
609
610 #[test]
611 fn wire_v1_reply_context_coexists_with_b64_tail() {
612 let mut m = Message::new(
613 UserId::new("alice").unwrap(),
614 UserId::new("bob").unwrap(),
615 AppId::new("app").unwrap(),
616 format!("__MLS_ENC__{}", BASE64.encode(vec![7u8; 90])),
617 );
618 m.reply_context = Some(sample_reply_context());
619 let bytes = m.to_wire_v1_bytes().unwrap();
620
621 let wire = decode(&bytes).unwrap();
622 assert_eq!(wire.ext.len(), 2);
623 assert_eq!(wire.ext[0].0, EXT_TAG_B64_TAIL);
624 assert_eq!(wire.ext[1].0, EXT_TAG_REPLY_CONTEXT);
625
626 let decoded = Message::from_wire_v1_bytes(&bytes).unwrap();
627 assert_eq!(decoded.content, m.content);
628 assert_eq!(decoded.reply_context, m.reply_context);
629 }
630
631 #[test]
632 fn wire_v1_honors_only_first_reply_context_entry() {
633 let mut wire = base_dto();
636 wire.ext = vec![
637 (
638 EXT_TAG_REPLY_CONTEXT,
639 br#"{"sender":"carol","text":"first"}"#.to_vec(),
640 ),
641 (
642 EXT_TAG_REPLY_CONTEXT,
643 br#"{"sender":"mallory","text":"second"}"#.to_vec(),
644 ),
645 ];
646 let decoded = Message::from_wire_v1_bytes(&encode(&wire).unwrap()).unwrap();
647 let rc = decoded.reply_context.expect("first entry honored");
648 assert_eq!(rc.sender.as_str(), "carol");
649 assert_eq!(rc.text, "first");
650 }
651
652 #[test]
653 fn wire_v1_rejects_malformed_reply_context_ext() {
654 let mut wire = base_dto();
657 wire.ext = vec![(EXT_TAG_REPLY_CONTEXT, b"not json".to_vec())];
658 assert!(Message::from_wire_v1_bytes(&encode(&wire).unwrap()).is_err());
659
660 let mut wire = base_dto();
662 wire.ext = vec![(
663 EXT_TAG_REPLY_CONTEXT,
664 br#"{"sender":"evil/path","text":"t"}"#.to_vec(),
665 )];
666 assert!(Message::from_wire_v1_bytes(&encode(&wire).unwrap()).is_err());
667 }
668
669 #[test]
670 fn wire_v1_golden_layout_with_reply_context_ext_is_frozen() {
671 let rc_json: &[u8] = br#"{"sender":"c","text":"t"}"#;
675 let mut wire = base_dto();
676 wire.content = String::new();
677 wire.ext = vec![(EXT_TAG_REPLY_CONTEXT, rc_json.to_vec())];
678 let mut expected_tail = vec![
679 0x01, 0x02, 0x19, ];
683 expected_tail.extend_from_slice(rc_json); let encoded = encode(&wire).unwrap();
685 assert_eq!(
686 &encoded[encoded.len() - expected_tail.len()..],
687 expected_tail
688 );
689 let decoded = Message::from_wire_v1_bytes(&encoded).unwrap();
691 let rc = decoded.reply_context.expect("reply context decoded");
692 assert_eq!(rc.sender.as_str(), "c");
693 assert_eq!(rc.text, "t");
694 }
695
696 #[test]
697 fn wire_v1_golden_layout_with_b64_tail_ext_is_frozen() {
698 let mut wire = base_dto();
702 wire.content = String::new();
703 wire.ext = vec![(EXT_TAG_B64_TAIL, vec![0xDE, 0xAD])];
704 let expected_tail: &[u8] = &[
705 0x01, 0x01, 0x02, 0xDE, 0xAD, ];
709 let encoded = encode(&wire).unwrap();
710 assert_eq!(
711 &encoded[encoded.len() - expected_tail.len()..],
712 expected_tail
713 );
714 let decoded = Message::from_wire_v1_bytes(&encoded).unwrap();
716 assert_eq!(decoded.content, BASE64.encode([0xDE, 0xAD]));
717 }
718
719 #[test]
720 fn wire_v1_golden_layout_is_frozen() {
721 let wire = WireMessageV1 {
725 id: [0u8; 16],
726 sender: "a".into(),
727 recipient: "b".into(),
728 app_id: "c".into(),
729 priority: 1,
730 ttl: 8,
731 hop_count: 0,
732 timestamp: 0,
733 lamport_clock: 0,
734 content_type: 0,
735 content: String::new(),
736 binary_content: None,
737 media_metadata_json: None,
738 metadata: Vec::new(),
739 requires_ack: true,
740 reply_to_msg: None,
741 forwarded_from_json: None,
742 ext: Vec::new(),
743 };
744 let expected: &[u8] = &[
745 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, ];
765 assert_eq!(encode(&wire).unwrap(), expected);
766 }
767}