xmtp_archive/
importer.rs

1use super::{ArchiveError, BackupMetadata};
2use crate::{BACKUP_VERSION, NONCE_SIZE, TAG_SIZE, util::GenericArrayExt};
3use aes_gcm::{AesGcm, aead::Aead, aes::Aes256};
4use async_compression::futures::bufread::ZstdDecoder;
5use futures::{FutureExt, Stream, StreamExt, ready};
6use futures_util::{AsyncBufRead, AsyncReadExt};
7use prost::Message;
8#[allow(deprecated)]
9use sha2::digest::{generic_array::GenericArray, typenum};
10use std::{pin::Pin, task::Poll};
11use xmtp_common::{if_native, if_wasm};
12use xmtp_proto::xmtp::device_sync::{BackupElement, backup_element::Element};
13
14if_native! {
15    mod file_import;
16    type AsyncReader = Pin<Box<dyn AsyncBufRead + Send>>;
17}
18if_wasm! {
19    type AsyncReader = Pin<Box<dyn AsyncBufRead>>;
20}
21
22pub struct ArchiveImporter {
23    pub metadata: BackupMetadata,
24    decoded: Vec<u8>,
25    element_len: Option<usize>,
26    finished: bool,
27    decoder: ZstdDecoder<AsyncReader>,
28
29    cipher: AesGcm<Aes256, typenum::U12, typenum::U16>,
30    #[allow(deprecated)]
31    nonce: GenericArray<u8, typenum::U12>,
32}
33
34impl Stream for ArchiveImporter {
35    type Item = Result<BackupElement, ArchiveError>;
36
37    /// Yields elements until the stream ends cleanly or fails; after either, it yields nothing.
38    fn poll_next(
39        self: Pin<&mut Self>,
40        cx: &mut std::task::Context<'_>,
41    ) -> std::task::Poll<Option<Self::Item>> {
42        let this = self.get_mut();
43        if this.finished {
44            return Poll::Ready(None);
45        }
46        let item = ready!(this.poll_element(cx));
47        this.finished = !matches!(item, Some(Ok(_)));
48        Poll::Ready(item)
49    }
50}
51
52impl ArchiveImporter {
53    // implements: ARCH-002, ARCH-018
54    fn poll_element(
55        &mut self,
56        cx: &mut std::task::Context<'_>,
57    ) -> Poll<Option<Result<BackupElement, ArchiveError>>> {
58        let mut buffer = [0u8; 1024];
59        loop {
60            if self.element_len.is_none() && self.decoded.len() >= 4 {
61                let bytes = self.decoded.drain(..4).collect::<Vec<_>>();
62                let element_len =
63                    u32::from_le_bytes(bytes.try_into().expect("is 4 bytes")) as usize;
64                if element_len < TAG_SIZE {
65                    return Poll::Ready(Some(Err(ArchiveError::InvalidFrame(
66                        "frame shorter than its authentication tag",
67                    ))));
68                }
69                self.element_len = Some(element_len);
70            }
71
72            if let Some(element_len) = self.element_len
73                && self.decoded.len() >= element_len
74            {
75                let ciphertext = &self.decoded[..element_len];
76                let decrypted = match self.cipher.decrypt(&self.nonce, ciphertext) {
77                    Ok(decrypted) => Ok(decrypted),
78                    // Legacy archives encrypt every frame under the base nonce.
79                    Err(_) => {
80                        self.nonce.decrement();
81                        self.cipher.decrypt(&self.nonce, ciphertext)
82                    }
83                };
84                self.nonce.increment();
85                let element = decrypted
86                    .map_err(ArchiveError::from)
87                    .and_then(|decrypted| Ok(BackupElement::decode(&*decrypted)?));
88                self.decoded.drain(..element_len);
89                self.element_len = None;
90                return Poll::Ready(Some(element));
91            }
92
93            let amount = ready!(self.decoder.read(&mut buffer).poll_unpin(cx))?;
94            if amount == 0 {
95                return Poll::Ready(
96                    (self.element_len.is_some() || !self.decoded.is_empty())
97                        .then_some(Err(ArchiveError::InvalidFrame("truncated frame"))),
98                );
99            }
100            self.decoded.extend_from_slice(&buffer[..amount]);
101        }
102    }
103
104    // implements: ARCH-003
105    pub async fn load(mut reader: AsyncReader, key: &[u8]) -> Result<Self, ArchiveError> {
106        crate::check_key(key)?;
107        let mut version = [0; 2];
108        reader.read_exact(&mut version).await?;
109        let version = u16::from_le_bytes(version);
110        if version > BACKUP_VERSION {
111            return Err(ArchiveError::UnsupportedVersion(version));
112        }
113
114        let mut nonce = [0; NONCE_SIZE];
115        reader.read_exact(&mut nonce).await?;
116
117        let mut importer = Self {
118            decoder: ZstdDecoder::new(reader),
119            decoded: vec![],
120            element_len: None,
121            finished: false,
122            metadata: BackupMetadata::default(),
123
124            cipher: crate::cipher(key)?,
125            #[allow(deprecated)]
126            nonce: GenericArray::from(nonce),
127        };
128
129        let metadata = match importer.next().await {
130            Some(Ok(BackupElement {
131                element: Some(Element::Metadata(metadata)),
132            })) => metadata,
133            Some(Err(error)) => return Err(error),
134            _ => return Err(ArchiveError::MissingMetadata),
135        };
136
137        importer.metadata = BackupMetadata::from_metadata_save(metadata, version);
138        Ok(importer)
139    }
140
141    pub fn metadata(&self) -> &BackupMetadata {
142        &self.metadata
143    }
144}
145
146#[cfg(all(test, not(target_arch = "wasm32")))]
147mod tests;
148
149/// Container-level cases carried from self-hosted (#4304); `tests` holds the
150/// framing cases of this branch.
151#[cfg(all(test, not(target_arch = "wasm32")))]
152mod container_tests {
153    use super::*;
154    use crate::ENC_KEY_SIZE;
155    use async_compression::futures::write::ZstdEncoder;
156    use futures_util::{AsyncWriteExt, io::BufReader};
157    use std::sync::mpsc;
158    use xmtp_common::time::Duration;
159    use xmtp_proto::xmtp::device_sync::{BackupMetadataSave, consent_backup::ConsentSave};
160
161    const KEY: [u8; ENC_KEY_SIZE] = [3; ENC_KEY_SIZE];
162    const NONCE: [u8; NONCE_SIZE] = [0xff; NONCE_SIZE];
163
164    fn elements() -> [BackupElement; 2] {
165        [
166            BackupElement {
167                element: Some(Element::Metadata(BackupMetadataSave::default())),
168            },
169            BackupElement {
170                element: Some(Element::Consent(ConsentSave {
171                    entity: "entity".into(),
172                    ..Default::default()
173                })),
174            },
175        ]
176    }
177
178    /// Length-prefixed ciphertexts under counter nonces, or under the base nonce when `legacy`.
179    fn frames(legacy: bool, elements: &[BackupElement]) -> Vec<u8> {
180        let cipher = crate::cipher(&KEY).unwrap();
181        #[allow(deprecated)]
182        let mut nonce = GenericArray::from(NONCE);
183        elements
184            .iter()
185            .flat_map(|element| {
186                let ciphertext = cipher.encrypt(&nonce, &*element.encode_to_vec()).unwrap();
187                if !legacy {
188                    nonce.increment();
189                }
190                [(ciphertext.len() as u32).to_le_bytes().to_vec(), ciphertext].concat()
191            })
192            .collect()
193    }
194
195    /// The header for `version`, then `frames` as one complete zstd stream.
196    async fn container(version: u16, frames: &[u8]) -> Vec<u8> {
197        let mut encoder = ZstdEncoder::new(Vec::new());
198        encoder.write_all(frames).await.unwrap();
199        encoder.close().await.unwrap();
200        [&version.to_le_bytes()[..], &NONCE, &encoder.into_inner()].concat()
201    }
202
203    /// Imports `archive` on its own thread, returning the elements after metadata or the first
204    /// error. Panics if the import hangs or the stream yields anything after an error.
205    fn restore(archive: Vec<u8>) -> Result<Vec<BackupElement>, ArchiveError> {
206        let (sender, receiver) = mpsc::channel();
207        std::thread::spawn(move || {
208            let runtime = tokio::runtime::Builder::new_current_thread()
209                .build()
210                .unwrap();
211            let result = runtime.block_on(async {
212                let reader = Box::pin(BufReader::new(futures::io::Cursor::new(archive)));
213                let mut importer = ArchiveImporter::load(reader, &KEY).await?;
214                let mut restored = Vec::new();
215                while let Some(element) = importer.next().await {
216                    match element {
217                        Ok(element) => restored.push(element),
218                        Err(error) => {
219                            let after = importer.next().await;
220                            assert!(after.is_none(), "stream continued after {error}: {after:?}");
221                            return Err(error);
222                        }
223                    }
224                }
225                Ok(restored)
226            });
227            let _ = sender.send(result);
228        });
229        receiver
230            .recv_timeout(Duration::from_secs(2))
231            .expect("import hung or panicked")
232    }
233
234    /// A container this client cannot read, or one cut short anywhere, must end the import with an
235    /// error — never hang, never keep yielding, never pass for a complete backup — while valid
236    /// counter-nonce and legacy archives still restore.
237    // verifies: ARCH-003, ARCH-018
238    #[xmtp_common::test(unwrap_try = true)]
239    async fn archive_invalid_container_rejected() {
240        let [metadata, consent] = elements();
241        let counter = frames(false, &[metadata.clone(), consent.clone()]);
242        let legacy = frames(true, &[metadata.clone(), consent.clone()]);
243        let metadata_end = frames(false, &[metadata]).len();
244
245        // Both nonce forms of a version 0 archive restore every element.
246        for framed in [&counter, &legacy] {
247            assert_eq!(restore(container(0, framed).await)?, vec![consent.clone()]);
248        }
249
250        // A later version fails before any frame is read, naming the version.
251        for version in [1, 2, u16::MAX] {
252            let mut archive = container(0, &counter).await;
253            archive[..2].copy_from_slice(&version.to_le_bytes());
254            let error = restore(archive).unwrap_err();
255            assert!(
256                matches!(error, ArchiveError::UnsupportedVersion(v) if v == version),
257                "version {version} was read: {error:?}"
258            );
259            assert!(error.to_string().contains(&version.to_string()));
260        }
261
262        // Every cut of the file — header, zstd frame header, blocks, checksum — fails.
263        let archive = container(0, &counter).await;
264        for cut in 0..archive.len() {
265            assert!(
266                restore(archive[..cut].to_vec()).is_err(),
267                "archive cut at byte {cut} of {} restored",
268                archive.len()
269            );
270        }
271
272        // A complete zstd stream ending within a length prefix or ciphertext fails; only a frame
273        // boundary after the metadata frame ends cleanly.
274        for cut in 0..counter.len() {
275            let restored = restore(container(0, &counter[..cut]).await);
276            match cut {
277                0 => assert!(matches!(restored, Err(ArchiveError::MissingMetadata))),
278                _ if cut == metadata_end => assert_eq!(restored?, vec![]),
279                _ => assert!(
280                    matches!(restored, Err(ArchiveError::InvalidFrame(_))),
281                    "frames cut at byte {cut} of {}: {restored:?}",
282                    counter.len()
283                ),
284            }
285        }
286
287        // A full-length frame that authenticates under neither the counter nor the legacy nonce
288        // fails with the authentication error and ends the stream.
289        let mut forged = counter.clone();
290        *forged.last_mut()? ^= 1;
291        let restored = restore(container(0, &forged).await);
292        assert!(
293            matches!(restored, Err(ArchiveError::AesGcm(_))),
294            "forged tag: {restored:?}"
295        );
296
297        // A declared ciphertext shorter than the tag fails even when the bytes are present.
298        for len in 0..TAG_SIZE as u32 {
299            let framed = [&counter[..metadata_end], &len.to_le_bytes(), &[0; 32]].concat();
300            let restored = restore(container(0, &framed).await);
301            assert!(
302                matches!(restored, Err(ArchiveError::InvalidFrame(_))),
303                "ciphertext length {len}: {restored:?}"
304            );
305        }
306    }
307}
xmtp_archive/importer.rs

xmtp_archive/
importer.rs
