xmtp_attachments/
encoding.rs

1use std::{
2    cell::Cell,
3    io::{Read, Seek, SeekFrom, Write},
4};
5
6use flate2::read::{DeflateDecoder, MultiGzDecoder, ZlibDecoder};
7use prost::Message as _;
8use xmtp_content_types::{
9    ContentCodec as _,
10    attachment::{Attachment, AttachmentCodec},
11};
12use xmtp_proto::xmtp::mls::message_contents::{Compression, EncodedContent};
13
14use crate::{AttachmentError, AttachmentFailureCause};
15
16const CONTENT_FIELD_TAG: u8 = 0x22;
17const MAX_METADATA_BYTES: usize = 65_536;
18const MAX_PARAMETER_KEY_BYTES: usize = 256;
19const MAX_DECOMPRESSED_BYTES: usize = 16_777_216;
20const COPY_CHUNK_BYTES: usize = 8192;
21
22fn invalid() -> AttachmentError {
23    AttachmentError::new(AttachmentFailureCause::NotAnAttachment)
24}
25
26fn local_storage() -> AttachmentError {
27    AttachmentError::new(AttachmentFailureCause::LocalStorage)
28}
29
30struct TrackedSource<'a, R> {
31    source: &'a mut R,
32    read_failed: &'a Cell<bool>,
33}
34
35impl<R: Read> Read for TrackedSource<'_, R> {
36    fn read(&mut self, output: &mut [u8]) -> std::io::Result<usize> {
37        self.source
38            .read(output)
39            .inspect_err(|_| self.read_failed.set(true))
40    }
41}
42
43fn decode_error(read_failed: &Cell<bool>) -> AttachmentError {
44    if read_failed.get() {
45        local_storage()
46    } else {
47        invalid()
48    }
49}
50
51/// Encode the protobuf fields before the attachment content.
52///
53/// Map entries use sorted keys, so the same inputs produce the same bytes.
54/// Task 10 must compute this prefix once per creation and reuse those bytes.
55/// The caller writes exactly `content_len` content bytes after this prefix.
56pub fn encoded_prefix(filename: Option<&str>, mime_type: &str, content_len: u64) -> Vec<u8> {
57    let mut envelope = AttachmentCodec::encode(Attachment {
58        filename: filename.map(str::to_owned),
59        mime_type: mime_type.to_owned(),
60        content: Vec::new(),
61    })
62    .expect("attachment encoding has no failure path");
63    let mut field = EncodedContent {
64        r#type: envelope.r#type.take(),
65        ..Default::default()
66    };
67    let mut prefix = field.encode_to_vec();
68    let mut parameters = envelope.parameters.drain().collect::<Vec<_>>();
69    parameters.sort_unstable_by(|left, right| left.0.cmp(&right.0));
70    field.r#type = None;
71    for (key, value) in parameters {
72        field.parameters.insert(key, value);
73        prefix.extend_from_slice(&field.encode_to_vec());
74        field.parameters.clear();
75    }
76    field.fallback = envelope.fallback.take();
77    prefix.extend_from_slice(&field.encode_to_vec());
78    if content_len != 0 {
79        prefix.push(CONTENT_FIELD_TAG);
80        encode_varint(content_len, &mut prefix);
81    }
82    prefix
83}
84
85/// Check the fields that the streaming decoder keeps before content bytes.
86/// This uses the same envelope as `encoded_prefix` and the decoder's limit.
87pub fn retained_fields_fit(filename: Option<&str>, mime_type: &str) -> bool {
88    let value_bytes = mime_type.len().saturating_add(filename.map_or(0, str::len));
89    if value_bytes > MAX_METADATA_BYTES {
90        return false;
91    }
92    let prefix = encoded_prefix(filename, mime_type, 0);
93    AttachmentDecoder::new().push(&prefix).is_ok()
94}
95
96pub fn ciphertext_len(prefix_len: usize, content_len: u64) -> u64 {
97    (prefix_len as u64)
98        .saturating_add(content_len)
99        .saturating_add(16)
100}
101
102fn encode_varint(mut value: u64, output: &mut Vec<u8>) {
103    while value >= 0x80 {
104        output.push((value as u8) | 0x80);
105        value >>= 7;
106    }
107    output.push(value as u8);
108}
109
110fn decode_varint(bytes: &[u8]) -> Result<Option<u64>, AttachmentError> {
111    if bytes.len() > 10 {
112        return Err(invalid());
113    }
114    let mut value = 0u64;
115    for (index, &byte) in bytes.iter().enumerate() {
116        if index == 9 && byte > 1 {
117            return Err(invalid());
118        }
119        value |= u64::from(byte & 0x7f) << (index * 7);
120        if byte < 0x80 {
121            return Ok(Some(value));
122        }
123    }
124    if bytes.len() == 10 {
125        Err(invalid())
126    } else {
127        Ok(None)
128    }
129}
130
131#[derive(Debug, Clone, Copy)]
132enum ParseState {
133    Tag,
134    Length { field: u64, tag_len: usize },
135    Data { remaining: usize, kind: DataKind },
136    OtherVarint { retain: bool },
137    Fixed { remaining: usize },
138}
139
140#[derive(Debug, Clone, Copy)]
141enum DataKind {
142    Content,
143    Type,
144    Parameter,
145    Skip,
146}
147
148#[derive(Debug, Clone, Copy)]
149enum EntryState {
150    Tag,
151    Length { field: u64 },
152    Data { field: u64, remaining: usize },
153    Varint,
154    Fixed { remaining: usize },
155}
156
157/// Parse one map entry with bounded key and value buffers.
158struct ParameterEntry {
159    state: EntryState,
160    header: Vec<u8>,
161    key: Vec<u8>,
162    value: Vec<u8>,
163    key_too_long: bool,
164    value_too_long: bool,
165    serialized_len: usize,
166}
167
168impl ParameterEntry {
169    fn new(prefix_len: usize) -> Self {
170        Self {
171            state: EntryState::Tag,
172            header: Vec::new(),
173            key: Vec::new(),
174            value: Vec::new(),
175            key_too_long: false,
176            value_too_long: false,
177            serialized_len: prefix_len,
178        }
179    }
180
181    fn relevant(&self) -> bool {
182        !self.key_too_long && matches!(self.key.as_slice(), b"mimeType" | b"filename")
183    }
184
185    fn push(&mut self, input: &[u8]) -> Result<(), AttachmentError> {
186        self.serialized_len = self.serialized_len.saturating_add(input.len());
187        let mut at = 0;
188        while at < input.len() {
189            match self.state {
190                EntryState::Tag | EntryState::Length { .. } | EntryState::Varint => {
191                    self.header.push(input[at]);
192                    at += 1;
193                    let Some(value) = decode_varint(&self.header)? else {
194                        continue;
195                    };
196                    match self.state {
197                        EntryState::Tag => {
198                            if value >> 3 == 0 || value > u64::from(u32::MAX) {
199                                return Err(invalid());
200                            }
201                            let field = value >> 3;
202                            let wire = value & 7;
203                            if matches!(field, 1 | 2) && wire != 2 {
204                                return Err(invalid());
205                            }
206                            self.state = match wire {
207                                0 => EntryState::Varint,
208                                1 => EntryState::Fixed { remaining: 8 },
209                                2 => EntryState::Length { field },
210                                5 => EntryState::Fixed { remaining: 4 },
211                                _ => return Err(invalid()),
212                            };
213                        }
214                        EntryState::Length { field } => {
215                            let remaining = usize::try_from(value).map_err(|_| invalid())?;
216                            match field {
217                                1 => {
218                                    self.key.clear();
219                                    self.key_too_long = remaining > MAX_PARAMETER_KEY_BYTES;
220                                }
221                                2 => {
222                                    self.value.clear();
223                                    self.value_too_long = remaining > MAX_METADATA_BYTES;
224                                }
225                                _ => {}
226                            }
227                            self.state = if remaining == 0 {
228                                EntryState::Tag
229                            } else {
230                                EntryState::Data { field, remaining }
231                            };
232                        }
233                        EntryState::Varint => self.state = EntryState::Tag,
234                        _ => unreachable!(),
235                    }
236                    self.header.clear();
237                }
238                EntryState::Data { field, remaining } => {
239                    let take = remaining.min(input.len() - at);
240                    let bytes = &input[at..at + take];
241                    match field {
242                        1 if !self.key_too_long => self.key.extend_from_slice(bytes),
243                        2 if !self.value_too_long => {
244                            self.value.extend_from_slice(bytes);
245                        }
246                        _ => {}
247                    }
248                    at += take;
249                    self.state = if take == remaining {
250                        EntryState::Tag
251                    } else {
252                        EntryState::Data {
253                            field,
254                            remaining: remaining - take,
255                        }
256                    };
257                }
258                EntryState::Fixed { remaining } => {
259                    let take = remaining.min(input.len() - at);
260                    at += take;
261                    self.state = if take == remaining {
262                        EntryState::Tag
263                    } else {
264                        EntryState::Fixed {
265                            remaining: remaining - take,
266                        }
267                    };
268                }
269            }
270        }
271        Ok(())
272    }
273
274    fn finish(self) -> Result<Option<(Vec<u8>, usize)>, AttachmentError> {
275        if !matches!(self.state, EntryState::Tag) || !self.header.is_empty() {
276            return Err(invalid());
277        }
278        if !self.relevant() {
279            return Ok(None);
280        }
281        if self.value_too_long {
282            return Err(invalid());
283        }
284        let mut entry = vec![0x0a];
285        encode_varint(self.key.len() as u64, &mut entry);
286        entry.extend_from_slice(&self.key);
287        entry.push(0x12);
288        encode_varint(self.value.len() as u64, &mut entry);
289        entry.extend_from_slice(&self.value);
290        let mut field = vec![0x12];
291        encode_varint(entry.len() as u64, &mut field);
292        field.extend_from_slice(&entry);
293        Ok(Some((field, self.serialized_len)))
294    }
295}
296
297/// Metadata of an attachment envelope.
298#[derive(Debug, Clone, PartialEq, Eq)]
299pub struct DecodedMeta {
300    pub mime_type: String,
301    pub filename: Option<String>,
302    /// True when `finish` wrote decompressed bytes to the output sink.
303    pub compressed: bool,
304}
305
306/// Content events returned by the streaming attachment decoder.
307#[derive(Debug, Clone, Copy, PartialEq, Eq)]
308pub enum ContentChunk<'a> {
309    /// Discard earlier content and start again at byte zero.
310    Reset,
311    /// Write these bytes after the last content event.
312    Bytes(&'a [u8]),
313}
314
315/// Parses a protobuf envelope across arbitrary chunk boundaries.
316///
317/// Apply every event returned by `push` to a temporary content sink, in order,
318/// before the next call. `Reset` truncates the sink to zero and rewinds it.
319/// The decoder retains only bounded metadata, not content bytes.
320pub struct AttachmentDecoder {
321    state: ParseState,
322    header: Vec<u8>,
323    metadata: Vec<u8>,
324    retained_bytes: usize,
325    parameter: Option<ParameterEntry>,
326    content_fields: usize,
327    sink_has_content: bool,
328}
329
330impl Default for AttachmentDecoder {
331    fn default() -> Self {
332        Self::new()
333    }
334}
335
336impl AttachmentDecoder {
337    pub fn new() -> Self {
338        Self {
339            state: ParseState::Tag,
340            header: Vec::new(),
341            metadata: Vec::new(),
342            retained_bytes: 0,
343            parameter: None,
344            content_fields: 0,
345            sink_has_content: false,
346        }
347    }
348
349    fn append_metadata(&mut self, bytes: &[u8]) -> Result<(), AttachmentError> {
350        self.append_retained(bytes, bytes.len())
351    }
352
353    fn append_retained(
354        &mut self,
355        bytes: &[u8],
356        serialized_len: usize,
357    ) -> Result<(), AttachmentError> {
358        if serialized_len > MAX_METADATA_BYTES.saturating_sub(self.retained_bytes) {
359            return Err(invalid());
360        }
361        self.retained_bytes += serialized_len;
362        self.metadata.extend_from_slice(bytes);
363        Ok(())
364    }
365
366    /// Return ordered content events. Byte slices borrow from this input chunk.
367    /// Repeated content fields in one input chunk keep only the last field.
368    pub fn push<'a>(&mut self, input: &'a [u8]) -> Result<Vec<ContentChunk<'a>>, AttachmentError> {
369        let mut at = 0;
370        let mut content = Vec::with_capacity(2);
371        let sink_has_content = self.sink_has_content;
372        while at < input.len() {
373            match self.state {
374                ParseState::Tag | ParseState::Length { .. } | ParseState::OtherVarint { .. } => {
375                    self.header.push(input[at]);
376                    at += 1;
377                    let Some(value) = decode_varint(&self.header)? else {
378                        continue;
379                    };
380                    match self.state {
381                        ParseState::Tag => {
382                            if value >> 3 == 0 || value > u64::from(u32::MAX) {
383                                return Err(invalid());
384                            }
385                            let field = value >> 3;
386                            let wire = value & 7;
387                            if (matches!(field, 1..=4) && wire != 2) || (field == 5 && wire != 0) {
388                                return Err(invalid());
389                            }
390                            match wire {
391                                0 => {
392                                    let retain = field == 5;
393                                    let header = std::mem::take(&mut self.header);
394                                    if retain {
395                                        self.append_metadata(&header)?;
396                                    }
397                                    self.state = ParseState::OtherVarint { retain };
398                                }
399                                1 | 5 => {
400                                    self.header.clear();
401                                    self.state = ParseState::Fixed {
402                                        remaining: if wire == 1 { 8 } else { 4 },
403                                    };
404                                }
405                                2 => {
406                                    self.state = ParseState::Length {
407                                        field,
408                                        tag_len: self.header.len(),
409                                    };
410                                    self.header.clear();
411                                }
412                                _ => return Err(invalid()),
413                            }
414                        }
415                        ParseState::Length { field, tag_len } => {
416                            let remaining = usize::try_from(value).map_err(|_| invalid())?;
417                            let kind = match field {
418                                1 => DataKind::Type,
419                                2 => DataKind::Parameter,
420                                4 => DataKind::Content,
421                                _ => DataKind::Skip,
422                            };
423                            if matches!(kind, DataKind::Content) {
424                                if self.content_fields > 0 {
425                                    content.clear();
426                                    if sink_has_content {
427                                        content.push(ContentChunk::Reset);
428                                    }
429                                }
430                                self.content_fields = self.content_fields.saturating_add(1);
431                            } else if matches!(kind, DataKind::Type) {
432                                let mut header = Vec::new();
433                                encode_varint((field << 3) | 2, &mut header);
434                                header.extend_from_slice(&self.header);
435                                // The retained bytes use a canonical tag; the limit uses the input tag.
436                                self.append_retained(&header, tag_len + self.header.len())?;
437                                if remaining
438                                    > MAX_METADATA_BYTES.saturating_sub(self.retained_bytes)
439                                {
440                                    return Err(invalid());
441                                }
442                            } else if matches!(kind, DataKind::Parameter) {
443                                self.parameter =
444                                    Some(ParameterEntry::new(tag_len + self.header.len()));
445                            }
446                            self.header.clear();
447                            if remaining == 0 && matches!(kind, DataKind::Parameter) {
448                                self.parameter.take().unwrap().finish()?;
449                            }
450                            self.state = if remaining == 0 {
451                                ParseState::Tag
452                            } else {
453                                ParseState::Data { remaining, kind }
454                            };
455                        }
456                        ParseState::OtherVarint { retain } => {
457                            let header = std::mem::take(&mut self.header);
458                            if retain {
459                                self.append_metadata(&header)?;
460                            }
461                            self.state = ParseState::Tag;
462                        }
463                        _ => unreachable!(),
464                    }
465                }
466                ParseState::Data { remaining, kind } => {
467                    let take = remaining.min(input.len() - at);
468                    let bytes = &input[at..at + take];
469                    match kind {
470                        DataKind::Content => content.push(ContentChunk::Bytes(bytes)),
471                        DataKind::Type => self.append_metadata(bytes)?,
472                        DataKind::Parameter => self.parameter.as_mut().unwrap().push(bytes)?,
473                        DataKind::Skip => {}
474                    }
475                    at += take;
476                    if take == remaining
477                        && matches!(kind, DataKind::Parameter)
478                        && let Some((field, serialized_len)) =
479                            self.parameter.take().unwrap().finish()?
480                    {
481                        self.append_retained(&field, serialized_len)?;
482                    }
483                    self.state = if take == remaining {
484                        ParseState::Tag
485                    } else {
486                        ParseState::Data {
487                            remaining: remaining - take,
488                            kind,
489                        }
490                    };
491                }
492                ParseState::Fixed { remaining } => {
493                    let take = remaining.min(input.len() - at);
494                    at += take;
495                    self.state = if take == remaining {
496                        ParseState::Tag
497                    } else {
498                        ParseState::Fixed {
499                            remaining: remaining - take,
500                        }
501                    };
502                }
503            }
504        }
505        if let Some(last) = content.last() {
506            self.sink_has_content = matches!(last, ContentChunk::Bytes(_));
507        }
508        Ok(content)
509    }
510
511    /// Validate metadata and, when needed, decompress the temporary content.
512    ///
513    /// `source` reads the content slices written by `push`. `output` receives
514    /// decompressed bytes only when compression is present. Neither sink is
515    /// published until this function and GCM authentication both succeed.
516    pub fn finish<R: Read + Seek, W: Write>(
517        self,
518        source: &mut R,
519        output: &mut W,
520    ) -> Result<DecodedMeta, AttachmentError> {
521        if !matches!(self.state, ParseState::Tag) || !self.header.is_empty() {
522            return Err(invalid());
523        }
524        let envelope = EncodedContent::decode(self.metadata.as_slice()).map_err(|_| invalid())?;
525        let ty = envelope.r#type.ok_or_else(invalid)?;
526        if ty.authority_id != "xmtp.org" || ty.type_id != "attachment" || ty.version_major != 1 {
527            return Err(invalid());
528        }
529        let compressed = match envelope.compression {
530            None => false,
531            Some(raw) if raw == Compression::Gzip as i32 => {
532                source
533                    .seek(SeekFrom::Start(0))
534                    .map_err(|_| local_storage())?;
535                let read_failed = Cell::new(false);
536                let tracked = TrackedSource {
537                    source,
538                    read_failed: &read_failed,
539                };
540                copy_bounded(&mut MultiGzDecoder::new(tracked), output, &read_failed)?;
541                true
542            }
543            Some(raw) if raw == Compression::Deflate as i32 => {
544                source
545                    .seek(SeekFrom::Start(0))
546                    .map_err(|_| local_storage())?;
547                // Validate zlib before writing. If it fails, try raw DEFLATE.
548                let read_failed = Cell::new(false);
549                let tracked = TrackedSource {
550                    source: &mut *source,
551                    read_failed: &read_failed,
552                };
553                let zlib_ok = probe_zlib(&mut ZlibDecoder::new(tracked), &read_failed)?;
554                source
555                    .seek(SeekFrom::Start(0))
556                    .map_err(|_| local_storage())?;
557                let read_failed = Cell::new(false);
558                let tracked = TrackedSource {
559                    source,
560                    read_failed: &read_failed,
561                };
562                if zlib_ok {
563                    copy_bounded(&mut ZlibDecoder::new(tracked), output, &read_failed)?;
564                } else {
565                    copy_bounded(&mut DeflateDecoder::new(tracked), output, &read_failed)?;
566                }
567                true
568            }
569            _ => return Err(invalid()),
570        };
571        Ok(DecodedMeta {
572            mime_type: envelope
573                .parameters
574                .get("mimeType")
575                .cloned()
576                .unwrap_or_default(),
577            filename: envelope.parameters.get("filename").cloned(),
578            compressed,
579        })
580    }
581}
582
583fn copy_bounded(
584    reader: &mut impl Read,
585    output: &mut impl Write,
586    read_failed: &Cell<bool>,
587) -> Result<(), AttachmentError> {
588    let mut count = 0usize;
589    let mut chunk = [0u8; COPY_CHUNK_BYTES];
590    loop {
591        let request = (MAX_DECOMPRESSED_BYTES - count + 1).min(chunk.len());
592        let read = reader
593            .read(&mut chunk[..request])
594            .map_err(|_| decode_error(read_failed))?;
595        if read == 0 {
596            return if read_failed.get() {
597                Err(local_storage())
598            } else {
599                Ok(())
600            };
601        }
602        if count + read > MAX_DECOMPRESSED_BYTES {
603            return Err(invalid());
604        }
605        output
606            .write_all(&chunk[..read])
607            .map_err(|_| AttachmentError::new(AttachmentFailureCause::LocalStorage))?;
608        count += read;
609    }
610}
611
612fn probe_zlib(reader: &mut impl Read, read_failed: &Cell<bool>) -> Result<bool, AttachmentError> {
613    let mut count = 0usize;
614    let mut chunk = [0u8; COPY_CHUNK_BYTES];
615    loop {
616        let request = (MAX_DECOMPRESSED_BYTES - count + 1).min(chunk.len());
617        let read = match reader.read(&mut chunk[..request]) {
618            Ok(read) => read,
619            Err(_) if read_failed.get() => return Err(local_storage()),
620            Err(_) => return Ok(false),
621        };
622        if read == 0 {
623            if read_failed.get() {
624                return Err(local_storage());
625            }
626            return Ok(true);
627        }
628        if count + read > MAX_DECOMPRESSED_BYTES {
629            return Err(invalid());
630        }
631        count += read;
632    }
633}
634
635#[cfg(test)]
636mod tests {
637    use std::collections::BTreeMap;
638    use std::io::{self, Cursor, Read, Seek, SeekFrom, Write as _};
639
640    use flate2::{
641        Compression as FlateCompression,
642        write::{DeflateEncoder, GzEncoder, ZlibEncoder},
643    };
644    use xmtp_content_types::{
645        ContentCodec as _,
646        attachment::{Attachment, AttachmentCodec},
647    };
648
649    use super::*;
650
651    #[derive(Clone, PartialEq, prost::Message)]
652    struct OrderedAttachmentEnvelope {
653        #[prost(message, optional, tag = "1")]
654        r#type: Option<xmtp_proto::xmtp::mls::message_contents::ContentTypeId>,
655        #[prost(btree_map = "string, string", tag = "2")]
656        parameters: BTreeMap<String, String>,
657        #[prost(string, optional, tag = "3")]
658        fallback: Option<String>,
659        #[prost(bytes = "vec", tag = "4")]
660        content: Vec<u8>,
661    }
662
663    fn envelope(content: Vec<u8>) -> EncodedContent {
664        AttachmentCodec::encode(Attachment {
665            filename: Some("report.pdf".to_owned()),
666            mime_type: "application/pdf".to_owned(),
667            content,
668        })
669        .expect("attachment encoding has no failure path")
670    }
671
672    fn decode(bytes: &[u8]) -> Result<(Vec<u8>, Vec<u8>, DecodedMeta), AttachmentError> {
673        let mut decoder = AttachmentDecoder::new();
674        let mut temporary = Cursor::new(Vec::new());
675        for byte in bytes.chunks(1) {
676            for event in decoder.push(byte)? {
677                match event {
678                    ContentChunk::Reset => {
679                        temporary.get_mut().clear();
680                        temporary.set_position(0);
681                    }
682                    ContentChunk::Bytes(slice) => {
683                        temporary.write_all(slice).map_err(|_| invalid())?;
684                    }
685                }
686            }
687        }
688        let mut decompressed = Vec::new();
689        let meta = decoder.finish(&mut temporary, &mut decompressed)?;
690        Ok((temporary.into_inner(), decompressed, meta))
691    }
692
693    fn decode_split_metadata(bytes: &[u8], split: usize) -> Result<DecodedMeta, AttachmentError> {
694        let mut decoder = AttachmentDecoder::new();
695        decoder.push(&bytes[..split])?;
696        decoder.push(&bytes[split..])?;
697        decoder.finish(&mut Cursor::new(Vec::new()), &mut Vec::new())
698    }
699
700    struct FaultingSource {
701        bytes: Cursor<Vec<u8>>,
702        seeks: usize,
703        fail_read_after_seeks: Option<usize>,
704        fail_seek: bool,
705    }
706
707    impl Read for FaultingSource {
708        fn read(&mut self, output: &mut [u8]) -> io::Result<usize> {
709            if self
710                .fail_read_after_seeks
711                .is_some_and(|count| self.seeks >= count)
712            {
713                return Err(io::Error::other("local content read failed"));
714            }
715            self.bytes.read(output)
716        }
717    }
718
719    impl Seek for FaultingSource {
720        fn seek(&mut self, position: SeekFrom) -> io::Result<u64> {
721            if self.fail_seek {
722                return Err(io::Error::other("local content seek failed"));
723            }
724            self.seeks += 1;
725            self.bytes.seek(position)
726        }
727    }
728
729    fn finish_with_faulting_source(
730        content: Vec<u8>,
731        compression: Compression,
732        fail_read_after_seeks: Option<usize>,
733        fail_seek: bool,
734    ) -> Result<DecodedMeta, AttachmentError> {
735        let mut value = envelope(content.clone());
736        value.compression = Some(compression as i32);
737        let encoded = value.encode_to_vec();
738        let mut decoder = AttachmentDecoder::new();
739        let _ = decoder.push(&encoded)?;
740        let mut source = FaultingSource {
741            bytes: Cursor::new(content),
742            seeks: 0,
743            fail_read_after_seeks,
744            fail_seek,
745        };
746        decoder.finish(&mut source, &mut Vec::new())
747    }
748
749    // verifies: ATCH-012
750    #[xmtp_common::test(unwrap_try = true)]
751    async fn encoder_matches_prost() {
752        for len in [0, 1, 127, 128, 16_384, 70_000] {
753            let content = vec![42; len];
754            // One parameter has one protobuf map entry and a stable order.
755            let expected = AttachmentCodec::encode(Attachment {
756                filename: None,
757                mime_type: "application/octet-stream".to_owned(),
758                content: content.clone(),
759            })?;
760            let mut streamed = encoded_prefix(None, "application/octet-stream", len as u64);
761            streamed.extend_from_slice(&content);
762            assert_eq!(streamed, expected.encode_to_vec(), "length {len}");
763            assert_eq!(
764                ciphertext_len(streamed.len() - len, len as u64),
765                streamed.len() as u64 + 16
766            );
767
768            let mut with_filename =
769                encoded_prefix(Some("report.pdf"), "application/pdf", len as u64);
770            assert_eq!(
771                with_filename,
772                encoded_prefix(Some("report.pdf"), "application/pdf", len as u64)
773            );
774            with_filename.extend_from_slice(&content);
775            assert_eq!(
776                EncodedContent::decode(with_filename.as_slice())?,
777                envelope(content.clone())
778            );
779            let expected = envelope(content);
780            assert_eq!(
781                with_filename,
782                OrderedAttachmentEnvelope {
783                    r#type: expected.r#type,
784                    parameters: expected.parameters.into_iter().collect(),
785                    fallback: expected.fallback,
786                    content: expected.content,
787                }
788                .encode_to_vec()
789            );
790        }
791    }
792
793    #[xmtp_common::test(unwrap_try = true)]
794    fn ciphertext_len_saturates() {
795        assert_eq!(ciphertext_len(1, u64::MAX), u64::MAX);
796        assert_eq!(ciphertext_len(usize::MAX, u64::MAX), u64::MAX);
797    }
798
799    // verifies: ATCH-012
800    #[xmtp_common::test(unwrap_try = true)]
801    async fn decoder_any_field_order() {
802        let original = envelope(b"file data".to_vec());
803        let mut fields = Vec::new();
804        let mut ty = original.clone();
805        ty.parameters.clear();
806        ty.fallback = None;
807        ty.content.clear();
808        fields.push(ty.encode_to_vec());
809        let mut params = original.clone();
810        params.r#type = None;
811        params.fallback = None;
812        params.content.clear();
813        fields.push(params.encode_to_vec());
814        let mut fallback = original.clone();
815        fallback.r#type = None;
816        fallback.parameters.clear();
817        fallback.content.clear();
818        fields.push(fallback.encode_to_vec());
819        let mut content = original.clone();
820        content.r#type = None;
821        content.parameters.clear();
822        content.fallback = None;
823        fields.push(content.encode_to_vec());
824        for order in [[0, 1, 2, 3], [3, 2, 1, 0], [1, 3, 0, 2]] {
825            let bytes = order
826                .iter()
827                .flat_map(|&index| fields[index].iter().copied())
828                .collect::<Vec<_>>();
829            let (stored, decompressed, meta) = decode(&bytes)?;
830            assert_eq!(stored, b"file data");
831            assert!(decompressed.is_empty());
832            assert_eq!(meta.mime_type, "application/pdf");
833            assert_eq!(meta.filename.as_deref(), Some("report.pdf"));
834            assert!(!meta.compressed);
835        }
836        let mut wrong = original.clone();
837        wrong.r#type.as_mut().expect("type").type_id = "text".into();
838        assert_eq!(
839            decode(&wrong.encode_to_vec()).unwrap_err().cause,
840            AttachmentFailureCause::NotAnAttachment
841        );
842        let mut large = original;
843        large
844            .parameters
845            .insert("junk".to_owned(), "a".repeat(65_537));
846        let (_, _, meta) = decode(&large.encode_to_vec())?;
847        assert_eq!(meta.mime_type, "application/pdf");
848    }
849
850    #[xmtp_common::test(unwrap_try = true)]
851    async fn decoder_skips_large_unknown_parameters() {
852        let mut value = envelope(b"file data".to_vec());
853        value.parameters.insert("junk".into(), "a".repeat(70_000));
854        value.fallback = Some("b".repeat(70_000));
855        let (stored, _, meta) = decode(&value.encode_to_vec())?;
856        assert_eq!(stored, b"file data");
857        assert_eq!(meta.mime_type, "application/pdf");
858        assert_eq!(meta.filename.as_deref(), Some("report.pdf"));
859
860        value
861            .parameters
862            .insert("filename".into(), "c".repeat(70_000));
863        assert_eq!(
864            decode(&value.encode_to_vec()).unwrap_err().cause,
865            AttachmentFailureCause::NotAnAttachment
866        );
867    }
868
869    #[xmtp_common::test(unwrap_try = true)]
870    async fn retained_fields_total_cap_is_enforced() {
871        let with_values = |mime_type: usize, filename: usize| {
872            let mut value = envelope(b"file data".to_vec());
873            value
874                .parameters
875                .insert("mimeType".into(), "m".repeat(mime_type));
876            value
877                .parameters
878                .insert("filename".into(), "f".repeat(filename));
879            value
880        };
881        // The decoder keeps the type and the mimeType and filename entries.
882        let retained_len = |value: &EncodedContent| {
883            EncodedContent {
884                r#type: value.r#type.clone(),
885                parameters: value.parameters.clone(),
886                ..Default::default()
887            }
888            .encoded_len()
889        };
890
891        // Each value is under the per-value cap. Together they are over the total.
892        let over = with_values(40_000, 30_000);
893        assert!(retained_len(&over) > MAX_METADATA_BYTES);
894        assert_eq!(
895            decode(&over.encode_to_vec()).unwrap_err().cause,
896            AttachmentFailureCause::NotAnAttachment
897        );
898
899        let filename = (20_000..30_000)
900            .find(|&len| retained_len(&with_values(40_000, len)) == MAX_METADATA_BYTES)
901            .expect("a filename length that fills the cap exactly");
902        let (stored, _, meta) = decode(&with_values(40_000, filename).encode_to_vec())?;
903        assert_eq!(stored, b"file data");
904        assert_eq!(meta.mime_type.len(), 40_000);
905        assert_eq!(meta.filename.map(|name| name.len()), Some(filename));
906
907        let one_over = with_values(40_000, filename + 1);
908        assert_eq!(retained_len(&one_over), MAX_METADATA_BYTES + 1);
909        assert_eq!(
910            decode(&one_over.encode_to_vec()).unwrap_err().cause,
911            AttachmentFailureCause::NotAnAttachment
912        );
913    }
914
915    #[xmtp_common::test(unwrap_try = true)]
916    fn overlong_outer_tags_count_toward_retained_cap() {
917        let type_field = EncodedContent {
918            r#type: envelope(Vec::new()).r#type,
919            ..Default::default()
920        }
921        .encode_to_vec();
922        assert_eq!(type_field[0], 0x0a);
923
924        let build = |value_len: usize, overlong_type: bool| {
925            let mut bytes = type_field.clone();
926            if overlong_type {
927                bytes.splice(0..1, [0x8a, 0x00]);
928            }
929            let split = if overlong_type { 1 } else { bytes.len() + 1 };
930            let mut entry = b"\x0a\x08filename\x12".to_vec();
931            encode_varint(value_len as u64, &mut entry);
932            entry.extend(std::iter::repeat_n(b'f', value_len));
933            bytes.extend_from_slice(if overlong_type {
934                &[0x12][..]
935            } else {
936                &[0x92, 0x00][..]
937            });
938            encode_varint(entry.len() as u64, &mut bytes);
939            bytes.extend_from_slice(&entry);
940            (bytes, split)
941        };
942
943        for overlong_type in [false, true] {
944            // The length varints stay three bytes at these sizes.
945            let overhead = build(64_000, overlong_type).0.len() - 64_000;
946            let at_cap = MAX_METADATA_BYTES - overhead;
947            for (size, accepted) in [(at_cap, true), (at_cap + 1, false)] {
948                let (bytes, split) = build(size, overlong_type);
949                assert_eq!(bytes.len(), MAX_METADATA_BYTES + usize::from(!accepted));
950                assert!(EncodedContent::decode(bytes.as_slice()).is_ok());
951                let result = decode_split_metadata(&bytes, split);
952                if accepted {
953                    assert!(result.is_ok(), "at cap with overlong_type={overlong_type}");
954                } else {
955                    match result {
956                        Err(error) => {
957                            assert_eq!(error.cause, AttachmentFailureCause::NotAnAttachment);
958                        }
959                        Ok(_) => panic!("one byte over cap must fail"),
960                    }
961                }
962            }
963        }
964    }
965
966    #[xmtp_common::test(unwrap_try = true)]
967    async fn decoder_repeated_content_last_wins() {
968        let mut bytes = envelope(b"first".to_vec()).encode_to_vec();
969        bytes.extend_from_slice(b"\x22\x06second");
970        assert_eq!(EncodedContent::decode(bytes.as_slice())?.content, b"second");
971        let (stored, decompressed, _) = decode(&bytes)?;
972        assert_eq!(stored, b"second");
973        assert!(decompressed.is_empty());
974    }
975
976    #[xmtp_common::test(unwrap_try = true)]
977    async fn decoder_repeated_compressed_content_last_wins() {
978        let mut zlib = ZlibEncoder::new(Vec::new(), FlateCompression::default());
979        zlib.write_all(b"second")?;
980        let zlib = zlib.finish()?;
981        let mut value = envelope(b"invalid zlib".to_vec());
982        value.compression = Some(Compression::Deflate as i32);
983        let mut bytes = value.encode_to_vec();
984        bytes.push(CONTENT_FIELD_TAG);
985        encode_varint(zlib.len() as u64, &mut bytes);
986        bytes.extend_from_slice(&zlib);
987        assert_eq!(EncodedContent::decode(bytes.as_slice())?.content, zlib);
988        let (stored, decompressed, meta) = decode(&bytes)?;
989        assert_eq!(stored, zlib);
990        assert_eq!(decompressed, b"second");
991        assert!(meta.compressed);
992    }
993
994    // verifies: ATCH-039
995    #[xmtp_common::test(unwrap_try = true)]
996    fn repeated_content_fields_coalesce_per_push() {
997        let prefix = encoded_prefix(Some("report.pdf"), "application/pdf", 0);
998        let mut repeated = Vec::with_capacity(64 * 1024);
999        for _ in 0..(64 * 1024 / 3) {
1000            repeated.extend_from_slice(b"\x22\x01x");
1001        }
1002
1003        let mut decoder = AttachmentDecoder::new();
1004        assert!(decoder.push(&prefix)?.is_empty());
1005        let mut stored = Cursor::new(Vec::new());
1006        for event in decoder.push(b"\x22\x01y")? {
1007            if let ContentChunk::Bytes(bytes) = event {
1008                stored.write_all(bytes)?;
1009            }
1010        }
1011        assert_eq!(stored.get_ref().as_slice(), b"y");
1012
1013        let events = decoder.push(&repeated)?;
1014        assert!(events.len() <= 2, "too many events: {}", events.len());
1015        assert_eq!(
1016            events
1017                .iter()
1018                .filter(|event| matches!(event, ContentChunk::Reset))
1019                .count(),
1020            1
1021        );
1022        for event in events {
1023            match event {
1024                ContentChunk::Reset => {
1025                    stored.get_mut().clear();
1026                    stored.set_position(0);
1027                }
1028                ContentChunk::Bytes(bytes) => stored.write_all(bytes)?,
1029            }
1030        }
1031        decoder.finish(&mut stored, &mut Vec::new())?;
1032        assert_eq!(stored.into_inner(), b"x");
1033
1034        let mut fresh = AttachmentDecoder::new();
1035        fresh.push(&prefix)?;
1036        let events = fresh.push(&repeated)?;
1037        assert!(events.len() <= 1, "too many events: {}", events.len());
1038        assert!(
1039            events
1040                .iter()
1041                .all(|event| !matches!(event, ContentChunk::Reset))
1042        );
1043    }
1044
1045    #[xmtp_common::test(unwrap_try = true)]
1046    async fn repeated_entry_fields_replace_not_append() {
1047        let mut bytes = envelope(b"content".to_vec()).encode_to_vec();
1048        let mut entry = Vec::new();
1049        for key in ["junk", "junk", "filename"] {
1050            entry.push(0x0a);
1051            encode_varint(key.len() as u64, &mut entry);
1052            entry.extend_from_slice(key.as_bytes());
1053        }
1054        for value in ["a".repeat(60_000), "b".repeat(60_000)] {
1055            entry.push(0x12);
1056            encode_varint(value.len() as u64, &mut entry);
1057            entry.extend_from_slice(value.as_bytes());
1058        }
1059        bytes.push(0x12);
1060        encode_varint(entry.len() as u64, &mut bytes);
1061        bytes.extend_from_slice(&entry);
1062        let expected = EncodedContent::decode(bytes.as_slice())?;
1063        assert_eq!(expected.parameters.get("filename").unwrap().len(), 60_000);
1064        assert!(
1065            expected.parameters["filename"]
1066                .bytes()
1067                .all(|byte| byte == b'b')
1068        );
1069        assert_eq!(
1070            decode(&bytes).expect_err("oversized entry must fail").cause,
1071            AttachmentFailureCause::NotAnAttachment
1072        );
1073    }
1074
1075    #[xmtp_common::test(unwrap_try = true)]
1076    async fn recognized_repeated_values_count_serialized() {
1077        let mut value = envelope(b"content".to_vec());
1078        value.parameters.remove("filename");
1079        for size in [40_000, 10_000] {
1080            let mut bytes = value.encode_to_vec();
1081            let mut entry = Vec::new();
1082            entry.extend_from_slice(b"\x0a\x08filename");
1083            for byte in *b"ab" {
1084                entry.push(0x12);
1085                encode_varint(size, &mut entry);
1086                entry.extend(std::iter::repeat_n(byte, size as usize));
1087            }
1088            bytes.push(0x12);
1089            encode_varint(entry.len() as u64, &mut bytes);
1090            bytes.extend_from_slice(&entry);
1091            if size == 40_000 {
1092                assert_eq!(
1093                    decode(&bytes).expect_err("oversized entry must fail").cause,
1094                    AttachmentFailureCause::NotAnAttachment
1095                );
1096            } else {
1097                let (_, _, meta) = decode(&bytes)?;
1098                assert_eq!(
1099                    meta.filename.as_deref(),
1100                    Some("b".repeat(size as usize).as_str())
1101                );
1102            }
1103        }
1104    }
1105
1106    // verifies: ATCH-051
1107    #[xmtp_common::test(unwrap_try = true)]
1108    async fn protobuf_keys_above_u32_are_rejected() {
1109        // Field number 2^32 with wire type 0: the key does not fit in 32 bits.
1110        let oversized_key = 1u64 << 35;
1111
1112        let mut entry = Vec::new();
1113        for (field, value) in [(0x0a, "filename"), (0x12, "report.pdf")] {
1114            entry.push(field);
1115            encode_varint(value.len() as u64, &mut entry);
1116            entry.extend_from_slice(value.as_bytes());
1117        }
1118        encode_varint(oversized_key, &mut entry);
1119        entry.push(0);
1120        let mut in_entry = envelope(b"content".to_vec()).encode_to_vec();
1121        in_entry.push(0x12);
1122        encode_varint(entry.len() as u64, &mut in_entry);
1123        in_entry.extend_from_slice(&entry);
1124        assert!(EncodedContent::decode(in_entry.as_slice()).is_err());
1125        assert_eq!(
1126            decode(&in_entry).unwrap_err().cause,
1127            AttachmentFailureCause::NotAnAttachment
1128        );
1129
1130        let mut top_level = envelope(b"content".to_vec()).encode_to_vec();
1131        encode_varint(oversized_key, &mut top_level);
1132        top_level.push(0);
1133        assert!(EncodedContent::decode(top_level.as_slice()).is_err());
1134        assert_eq!(
1135            decode(&top_level).unwrap_err().cause,
1136            AttachmentFailureCause::NotAnAttachment
1137        );
1138    }
1139
1140    #[xmtp_common::test(unwrap_try = true)]
1141    async fn entry_value_before_final_key_is_kept() {
1142        let mut envelope = envelope(b"content".to_vec());
1143        envelope.parameters.remove("filename");
1144        let mut bytes = envelope.encode_to_vec();
1145        let mut entry = Vec::new();
1146        for (field, value) in [(0x0a, "junk"), (0x12, "report.pdf"), (0x0a, "filename")] {
1147            entry.push(field);
1148            encode_varint(value.len() as u64, &mut entry);
1149            entry.extend_from_slice(value.as_bytes());
1150        }
1151        bytes.push(0x12);
1152        encode_varint(entry.len() as u64, &mut bytes);
1153        bytes.extend_from_slice(&entry);
1154
1155        let expected = EncodedContent::decode(bytes.as_slice())?;
1156        assert_eq!(expected.parameters["filename"], "report.pdf");
1157        // decode feeds the streaming parser one byte at a time.
1158        let (_, _, meta) = decode(&bytes)?;
1159        assert_eq!(meta.filename.as_deref(), Some("report.pdf"));
1160        assert_eq!(
1161            meta.filename.as_deref(),
1162            Some(expected.parameters["filename"].as_str())
1163        );
1164    }
1165
1166    // verifies: CTYPE-001, CTYPE-016
1167    #[xmtp_common::test(unwrap_try = true)]
1168    async fn decoder_accepts_later_minor_versions() {
1169        let mut value = envelope(b"file data".to_vec());
1170        let ty = value.r#type.as_mut().expect("attachment type");
1171        ty.version_minor = 3;
1172        let (stored, decompressed, meta) = decode(&value.encode_to_vec())?;
1173        assert_eq!(stored, b"file data");
1174        assert!(decompressed.is_empty());
1175        assert_eq!(meta.mime_type, "application/pdf");
1176        assert_eq!(meta.filename.as_deref(), Some("report.pdf"));
1177
1178        value
1179            .r#type
1180            .as_mut()
1181            .expect("attachment type")
1182            .version_major = 2;
1183        assert_eq!(
1184            decode(&value.encode_to_vec()).unwrap_err().cause,
1185            AttachmentFailureCause::NotAnAttachment
1186        );
1187    }
1188
1189    #[xmtp_common::test(unwrap_try = true)]
1190    async fn decoder_compressed_content() {
1191        let text = b"hello compressed XMTP";
1192        let mut zlib = ZlibEncoder::new(Vec::new(), FlateCompression::default());
1193        zlib.write_all(text)?;
1194        let zlib = zlib.finish()?;
1195        let mut raw = DeflateEncoder::new(Vec::new(), FlateCompression::default());
1196        raw.write_all(text)?;
1197        let raw = raw.finish()?;
1198        let mut first = GzEncoder::new(Vec::new(), FlateCompression::default());
1199        first.write_all(b"hello ")?;
1200        let mut second = GzEncoder::new(Vec::new(), FlateCompression::default());
1201        second.write_all(b"compressed XMTP")?;
1202        let gzip = [first.finish()?, second.finish()?].concat();
1203        for (algorithm, compressed) in [
1204            (Compression::Deflate, zlib),
1205            (Compression::Deflate, raw),
1206            (Compression::Gzip, gzip),
1207        ] {
1208            let mut value = envelope(compressed.clone());
1209            value.compression = Some(algorithm as i32);
1210            let (stored, output, meta) = decode(&value.encode_to_vec())?;
1211            assert_eq!(stored, compressed);
1212            assert_eq!(output, text);
1213            assert!(meta.compressed);
1214        }
1215        let mut oversized = ZlibEncoder::new(Vec::new(), FlateCompression::default());
1216        oversized.write_all(&vec![7u8; MAX_DECOMPRESSED_BYTES + 1])?;
1217        let mut value = envelope(oversized.finish()?);
1218        value.compression = Some(Compression::Deflate as i32);
1219        assert_eq!(
1220            decode(&value.encode_to_vec()).unwrap_err().cause,
1221            AttachmentFailureCause::NotAnAttachment
1222        );
1223        value.content = b"invalid compressed bytes".to_vec();
1224        assert_eq!(
1225            decode(&value.encode_to_vec()).unwrap_err().cause,
1226            AttachmentFailureCause::NotAnAttachment
1227        );
1228        value.compression = Some(99);
1229        assert_eq!(
1230            decode(&value.encode_to_vec()).unwrap_err().cause,
1231            AttachmentFailureCause::NotAnAttachment
1232        );
1233    }
1234
1235    #[xmtp_common::test(unwrap_try = true)]
1236    fn source_io_errors_are_local_storage() {
1237        let text = b"valid compressed content";
1238        let mut gzip = GzEncoder::new(Vec::new(), FlateCompression::default());
1239        gzip.write_all(text)?;
1240        let gzip = gzip.finish()?;
1241        let mut zlib = ZlibEncoder::new(Vec::new(), FlateCompression::default());
1242        zlib.write_all(text)?;
1243        let zlib = zlib.finish()?;
1244
1245        for (content, compression, fail_read_after_seeks, fail_seek) in [
1246            (gzip.clone(), Compression::Gzip, Some(1), false),
1247            (zlib.clone(), Compression::Deflate, Some(1), false),
1248            (zlib.clone(), Compression::Deflate, Some(2), false),
1249            (gzip, Compression::Gzip, None, true),
1250        ] {
1251            let error =
1252                finish_with_faulting_source(content, compression, fail_read_after_seeks, fail_seek)
1253                    .expect_err("source I/O failure must stop decoding");
1254            assert_eq!(error.cause, AttachmentFailureCause::LocalStorage);
1255        }
1256
1257        let mut raw = DeflateEncoder::new(Vec::new(), FlateCompression::default());
1258        raw.write_all(text)?;
1259        let mut output = Vec::new();
1260        let compressed = raw.finish()?;
1261        let mut value = envelope(compressed.clone());
1262        value.compression = Some(Compression::Deflate as i32);
1263        let mut decoder = AttachmentDecoder::new();
1264        let encoded = value.encode_to_vec();
1265        let _ = decoder.push(&encoded)?;
1266        let mut source = Cursor::new(compressed);
1267        decoder.finish(&mut source, &mut output)?;
1268        assert_eq!(output, text);
1269
1270        let mut malformed = envelope(b"not gzip".to_vec());
1271        malformed.compression = Some(Compression::Gzip as i32);
1272        assert_eq!(
1273            decode(&malformed.encode_to_vec()).unwrap_err().cause,
1274            AttachmentFailureCause::NotAnAttachment
1275        );
1276    }
1277}
xmtp_attachments/encoding.rs

xmtp_attachments/
encoding.rs
