xmtp_archive/
exporter.rs

1use super::{BACKUP_VERSION, OptionsToSave, export_stream::BatchExportStream};
2use crate::archive_options::ArchiveOptions;
3use crate::{NONCE_SIZE, util::GenericArrayExt};
4use aes_gcm::{Aes256Gcm, AesGcm, KeyInit, aead::Aead, aes::Aes256};
5use async_compression::futures::write::ZstdEncoder;
6use futures::{Stream, pin_mut, ready, task::Context};
7use futures_util::{AsyncRead, AsyncWriteExt};
8use pin_project::pin_project;
9use prost::Message;
10#[allow(deprecated)]
11use sha2::digest::{generic_array::GenericArray, typenum};
12use std::{future::Future, io, pin::Pin, sync::Arc, task::Poll};
13use xmtp_db::prelude::*;
14use xmtp_proto::xmtp::device_sync::{BackupElement, BackupMetadataSave, backup_element::Element};
15
16#[cfg(not(target_arch = "wasm32"))]
17mod file_export;
18
19#[pin_project]
20pub struct ArchiveExporter {
21    stage: Stage,
22    metadata: BackupMetadataSave,
23    #[pin]
24    stream: BatchExportStream,
25    position: usize,
26    zstd_encoder: ZstdEncoder<Vec<u8>>,
27    encoder_finished: bool,
28
29    cipher: AesGcm<Aes256, typenum::U12, typenum::U16>,
30    nonce: GenericArray<u8, typenum::U12>,
31
32    // Used to write the nonce, contains the same data as nonce.
33    nonce_buffer: Vec<u8>,
34}
35
36#[derive(Default)]
37pub(super) enum Stage {
38    #[default]
39    Nonce,
40    Metadata,
41    Elements,
42}
43
44impl ArchiveExporter {
45    #[cfg(not(target_arch = "wasm32"))]
46    pub async fn export_to_file<D>(
47        options: ArchiveOptions,
48        db: D,
49        path: impl AsRef<std::path::Path>,
50        key: &[u8],
51    ) -> Result<BackupMetadataSave, crate::ArchiveError>
52    where
53        D: DbQuery + 'static,
54    {
55        let mut exporter = Self::new(options, db, key);
56        exporter.write_to_file(path).await?;
57
58        Ok(exporter.metadata)
59    }
60
61    pub fn new<D>(options: ArchiveOptions, db: D, key: &[u8]) -> Self
62    where
63        D: DbQuery + 'static,
64    {
65        let mut nonce_buffer = BACKUP_VERSION.to_le_bytes().to_vec();
66        let nonce = xmtp_common::rand_array::<NONCE_SIZE>();
67        nonce_buffer.extend_from_slice(&nonce);
68
69        Self {
70            position: 0,
71            stage: Stage::default(),
72            stream: BatchExportStream::new(&options, Arc::new(db)),
73            metadata: BackupMetadataSave::from_options(options),
74            zstd_encoder: ZstdEncoder::new(Vec::new()),
75            encoder_finished: false,
76
77            #[allow(deprecated)]
78            cipher: Aes256Gcm::new(GenericArray::from_slice(key)),
79            #[allow(deprecated)]
80            nonce: GenericArray::clone_from_slice(&nonce),
81            nonce_buffer,
82        }
83    }
84
85    pub fn metadata(&self) -> &BackupMetadataSave {
86        &self.metadata
87    }
88}
89
90// The reason this is future_util's AsyncRead and not tokio's AsyncRead
91// is because we need this to work on WASM, and tokio's AsyncRead makes
92// some assumptions about having access to std::fs, which WASM does not have.
93//
94// To get around this, we implement AsyncRead using future_util, and use a
95// compat layer from tokio_util to be able to interact with it in tokio.
96impl AsyncRead for ArchiveExporter {
97    /// This function encrypts first, and compresses second.
98    fn poll_read(
99        self: Pin<&mut Self>,
100        cx: &mut Context<'_>,
101        buf: &mut [u8],
102    ) -> Poll<io::Result<usize>> {
103        let mut this = self.project();
104        loop {
105            // Putting this up here because we don't want to encrypt or compress the nonce.
106            if matches!(this.stage, Stage::Nonce) {
107                let amount = this.nonce_buffer.len().min(buf.len());
108                let nonce_bytes: Vec<_> = this.nonce_buffer.drain(..amount).collect();
109                buf[..amount].copy_from_slice(&nonce_bytes);
110
111                if this.nonce_buffer.is_empty() {
112                    *this.stage = Stage::Metadata;
113                }
114                return Poll::Ready(Ok(amount));
115            }
116
117            {
118                // Read from the buffer while there is data
119                let buffer_inner = this.zstd_encoder.get_ref();
120                if *this.position < buffer_inner.len() {
121                    let available = &buffer_inner[*this.position..];
122                    let amount = available.len().min(buf.len());
123                    buf[..amount].copy_from_slice(&available[..amount]);
124                    *this.position += amount;
125
126                    return Poll::Ready(Ok(amount));
127                }
128            }
129
130            // The buffer is consumed. Reset.
131            *this.position = 0;
132            this.zstd_encoder.get_mut().clear();
133
134            // Time to fill the buffer with more data 8kb at a time.
135            while this.zstd_encoder.get_ref().len() < 8_000 {
136                let element = match this.stage {
137                    Stage::Nonce => {
138                        // Should never get here due to the above logic. Error if it does.
139                        unreachable!("Nonce should not be the stage here.");
140                    }
141                    Stage::Metadata => {
142                        *this.stage = Stage::Elements;
143                        BackupElement {
144                            element: Some(Element::Metadata(this.metadata.clone())),
145                        }
146                        .encode_to_vec()
147                    }
148                    Stage::Elements => match ready!(this.stream.as_mut().poll_next(cx)) {
149                        Some(element) => element
150                            .map_err(|err| io::Error::other(err.to_string()))?
151                            .encode_to_vec(),
152                        None => {
153                            if !*this.encoder_finished {
154                                *this.encoder_finished = true;
155                                let fut = this.zstd_encoder.close();
156                                pin_mut!(fut);
157                                let _ = fut.poll(cx)?;
158                            }
159                            break;
160                        }
161                    },
162                };
163
164                let mut element = this
165                    .cipher
166                    .encrypt(this.nonce, &*element)
167                    .expect("Encryption should always work");
168                let mut bytes = (element.len() as u32).to_le_bytes().to_vec();
169                bytes.append(&mut element);
170                this.nonce.increment();
171
172                let fut = this.zstd_encoder.write(&bytes);
173                pin_mut!(fut);
174                match fut.poll(cx) {
175                    Poll::Ready(Ok(_amt)) => {}
176                    Poll::Ready(Err(err)) => return Poll::Ready(Err(err)),
177                    Poll::Pending => return Poll::Pending,
178                }
179            }
180
181            // Flush the encoder
182            if !*this.encoder_finished {
183                let fut = this.zstd_encoder.flush();
184                pin_mut!(fut);
185                let _ = fut.poll(cx)?;
186            }
187
188            if this.zstd_encoder.get_ref().is_empty() {
189                return Poll::Ready(Ok(0));
190            }
191        }
192    }
193}
xmtp_archive/exporter.rs

xmtp_archive/
exporter.rs
