xmtp_archive/
exporter.rs

1//! Archive writing. [`export`] writes the version and the starting nonce in the
2//! clear, then one zstd stream of length-prefixed AES-GCM frames: the metadata,
3//! then every element of one database snapshot. Frames are written as rows are
4//! read, so memory does not grow with the history.
5
6use super::BACKUP_VERSION;
7use crate::archive_options::ArchiveOptions;
8use crate::{ArchiveError, NONCE_SIZE, snapshot, util::GenericArrayExt};
9use aes_gcm::aead::Aead;
10use async_compression::futures::write::ZstdEncoder;
11use futures::{FutureExt, io::AllowStdIo};
12use futures_util::AsyncWriteExt;
13use prost::Message;
14#[allow(deprecated)]
15use sha2::digest::generic_array::GenericArray;
16use std::io;
17use xmtp_db::ConnectionExt;
18use xmtp_proto::xmtp::device_sync::{BackupElement, BackupMetadataSave, backup_element::Element};
19
20/// Writes an archive of everything `options` selects, read in one snapshot
21/// measured at one export time, to `sink`, and returns its metadata. Fails,
22/// rather than omitting it, when a selected record cannot be read; on failure
23/// `sink` holds an incomplete archive that the caller must discard. `sink` is
24/// written while the snapshot's read transaction is open, so it must not use
25/// the database.
26pub fn export(
27    options: ArchiveOptions,
28    db: impl ConnectionExt,
29    key: &[u8],
30    sink: impl io::Write,
31) -> Result<BackupMetadataSave, ArchiveError> {
32    let mut writer = ElementWriter::new(key, sink)?;
33    let metadata = snapshot::read(&db, &options, |element| writer.write(element))?;
34    writer.finish()?;
35    Ok(metadata)
36}
37
38/// Streams standard archive elements to a synchronous sink.
39///
40/// The caller writes metadata first and each group before its messages. On any
41/// error, discard the sink. Only `finish` completes the compressed stream.
42#[allow(deprecated)]
43pub struct ElementWriter<W: io::Write> {
44    cipher: aes_gcm::Aes256Gcm,
45    nonce: GenericArray<u8, aes_gcm::aead::consts::U12>,
46    zstd: ZstdEncoder<AllowStdIo<W>>,
47}
48
49impl<W: io::Write> ElementWriter<W> {
50    /// Checks the key before writing the header.
51    pub fn new(key: &[u8], mut sink: W) -> Result<Self, ArchiveError> {
52        let cipher = crate::cipher(key)?;
53        let nonce = xmtp_common::rand_array::<NONCE_SIZE>();
54        sink.write_all(&BACKUP_VERSION.to_le_bytes())?;
55        sink.write_all(&nonce)?;
56        #[allow(deprecated)]
57        Ok(Self {
58            cipher,
59            nonce: GenericArray::clone_from_slice(&nonce),
60            zstd: ZstdEncoder::new(AllowStdIo::new(sink)),
61        })
62    }
63
64    /// Writes one encrypted frame and advances its nonce.
65    // implements: ARCH-001
66    pub fn write(&mut self, element: Element) -> Result<(), ArchiveError> {
67        let plaintext = BackupElement {
68            element: Some(element),
69        }
70        .encode_to_vec();
71        let ciphertext = self.cipher.encrypt(&self.nonce, &*plaintext)?;
72        self.nonce.increment();
73        let length = u32::try_from(ciphertext.len())
74            .map_err(|_| io::Error::other("archive frame exceeds u32 length"))?;
75        ready(self.zstd.write_all(&length.to_le_bytes()))?;
76        Ok(ready(self.zstd.write_all(&ciphertext))?)
77    }
78
79    /// Completes the compressed stream. A successful write alone is incomplete.
80    pub fn finish(mut self) -> Result<(), ArchiveError> {
81        Ok(ready(self.zstd.close())?)
82    }
83}
84
85/// Namespace for [`ArchiveExporter::export_to_file`], which the native
86/// bindings call.
87#[cfg(not(target_arch = "wasm32"))]
88pub struct ArchiveExporter;
89
90#[cfg(not(target_arch = "wasm32"))]
91impl ArchiveExporter {
92    /// Exports to a file at `path`, as [`export`], on tokio's blocking pool so
93    /// the snapshot never stalls an async worker. The archive is written to a
94    /// sibling temporary file and renamed over `path` only once complete, so a
95    /// failed export leaves `path` as it was. The archive is a new file:
96    /// owner-only on unix and the directory's default ACL on Windows, never
97    /// the permissions of whatever `path` held. Dropping the future cancels the
98    /// export at its next write, which counts as a failure. Must be called within a tokio runtime.
99    pub async fn export_to_file(
100        options: ArchiveOptions,
101        db: impl ConnectionExt + 'static,
102        path: impl AsRef<std::path::Path>,
103        key: &[u8],
104    ) -> Result<BackupMetadataSave, ArchiveError> {
105        let path = path.as_ref().to_owned();
106        offload(key, move |key, cancel| {
107            write_file(options, db, &path, key, cancel)
108        })
109        .await
110    }
111
112    /// Exports into a buffer, as [`export`], on tokio's blocking pool so the
113    /// snapshot never stalls an async worker. Dropping the future cancels the
114    /// export at its next write. Must be called within a tokio runtime.
115    pub async fn export_to_bytes(
116        options: ArchiveOptions,
117        db: impl ConnectionExt + 'static,
118        key: &[u8],
119    ) -> Result<Vec<u8>, ArchiveError> {
120        offload(key, move |key, cancel| {
121            let mut bytes = Vec::new();
122            export(options, db, key, Cancellable(&mut bytes, cancel))?;
123            Ok(bytes)
124        })
125        .await
126    }
127}
128
129/// Checks `key`, then runs `f` on tokio's blocking pool with a token that is
130/// cancelled when the returned future is dropped.
131#[cfg(not(target_arch = "wasm32"))]
132async fn offload<T: Send + 'static>(
133    key: &[u8],
134    f: impl FnOnce(&[u8], &tokio_util::sync::CancellationToken) -> Result<T, ArchiveError>
135    + Send
136    + 'static,
137) -> Result<T, ArchiveError> {
138    crate::check_key(key)?;
139    let key = key.to_vec();
140    let cancel = tokio_util::sync::CancellationToken::new();
141    let _cancel_on_drop = cancel.clone().drop_guard();
142    tokio::task::spawn_blocking(move || f(&key, &cancel))
143        .await
144        .map_err(io::Error::other)?
145}
146
147/// Exports to a sibling temporary file until `cancel` fires, then renames it
148/// over `path`. A failed or cancelled export removes the temporary file and
149/// leaves any archive already at `path` untouched.
150#[cfg(not(target_arch = "wasm32"))]
151pub(crate) fn write_file(
152    options: ArchiveOptions,
153    db: impl ConnectionExt,
154    path: &std::path::Path,
155    key: &[u8],
156    cancel: &tokio_util::sync::CancellationToken,
157) -> Result<BackupMetadataSave, ArchiveError> {
158    let mut partial = path.as_os_str().to_owned();
159    partial.push(format!(
160        ".{:016x}.partial",
161        u64::from_le_bytes(xmtp_common::rand_array())
162    ));
163    let partial = std::path::PathBuf::from(partial);
164    let mut create = std::fs::OpenOptions::new();
165    create.write(true).create_new(true);
166    #[cfg(unix)]
167    std::os::unix::fs::OpenOptionsExt::mode(&mut create, 0o600);
168    // Created before the cleanup below, which may then only remove a file
169    // this export owns.
170    let file = create.open(&partial)?;
171    let exported = (|| -> Result<_, ArchiveError> {
172        let mut file = io::BufWriter::new(file);
173        let metadata = export(options, db, key, Cancellable(&mut file, cancel))?;
174        file.into_inner()
175            .map_err(io::IntoInnerError::into_error)?
176            .sync_all()?;
177        live(cancel)?;
178        std::fs::rename(&partial, path)?;
179        Ok(metadata)
180    })();
181    if exported.is_err()
182        && let Err(e) = std::fs::remove_file(&partial)
183    {
184        tracing::warn!(path = %partial.display(), error = %e, "failed export left a partial archive");
185    }
186    exported
187}
188
189/// Fails once `cancel` fires, so a cancelled export takes the cleanup path.
190#[cfg(not(target_arch = "wasm32"))]
191fn live(cancel: &tokio_util::sync::CancellationToken) -> io::Result<()> {
192    if cancel.is_cancelled() {
193        return Err(io::Error::other("archive export cancelled"));
194    }
195    Ok(())
196}
197
198/// A sink that fails every write once its token is cancelled.
199#[cfg(not(target_arch = "wasm32"))]
200struct Cancellable<'a, W>(W, &'a tokio_util::sync::CancellationToken);
201
202#[cfg(not(target_arch = "wasm32"))]
203impl<W: io::Write> io::Write for Cancellable<'_, W> {
204    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
205        live(self.1)?;
206        self.0.write(buf)
207    }
208
209    fn flush(&mut self) -> io::Result<()> {
210        self.0.flush()
211    }
212}
213
214/// Resolves an encoder operation over a synchronous sink, which never pends.
215fn ready(op: impl Future<Output = io::Result<()>>) -> io::Result<()> {
216    op.now_or_never().unwrap_or_else(|| {
217        Err(io::Error::other(
218            "synchronous archive sink returned pending",
219        ))
220    })
221}
xmtp_archive/exporter.rs

xmtp_archive/
exporter.rs
