use std::ops::Bound;
use futures_io::AsyncWrite;
use futures_lite::io::AsyncWriteExt as _;
use futures_util::StreamExt as _;
use super::format::{SnapshotEncoder, SnapshotEntry};
use crate::LixError;
use crate::storage_adapter::{
MAX_SCAN_PAGE_ROWS, Storage, StorageAdapter, StorageAdapterRead as _, StorageBeginScanOptions,
StorageCoreProjection, StorageKeyRange, StorageProjectedValue,
StorageReadDurability as ReadDurability, StorageReadOptions, StorageSession,
};
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct SnapshotExportReport {
pub entry_count: u64,
pub payload_bytes: u64,
pub digest: [u8; 32],
}
#[expect(missing_debug_implementations)]
pub struct SnapshotExportBuilder<StorageImpl>
where
StorageImpl: Storage + Clone + Send + Sync + 'static,
{
storage: Option<StorageAdapter<StorageSession<StorageImpl>>>,
durability: ReadDurability,
preflight_error: Option<LixError>,
remote: Option<(crate::ServerOptions, String)>,
connected_remote: Option<RemoteSnapshotExport>,
local_partial_replica: bool,
}
#[derive(Clone)]
struct RemoteSnapshotExport {
http: crate::sync::AuthorityHttp,
url: String,
session_id: Option<String>,
}
impl<StorageImpl> SnapshotExportBuilder<StorageImpl>
where
StorageImpl: Storage + Clone + Send + Sync + 'static,
{
pub(crate) fn new(storage: StorageAdapter<StorageSession<StorageImpl>>) -> Self {
Self {
storage: Some(storage),
durability: ReadDurability::Visible,
preflight_error: None,
remote: None,
connected_remote: None,
local_partial_replica: false,
}
}
pub(crate) fn from_local_partial_replica(mut self) -> Self {
self.local_partial_replica = true;
self
}
pub(crate) fn from_sync_server(
mut self,
server: crate::ServerOptions,
expected_account_id: String,
) -> Self {
self.remote = Some((server, expected_account_id));
self
}
pub(crate) fn reject_connected_replica(mut self) -> Self {
self.preflight_error = Some(LixError::new(
LixError::CODE_INVALID_PARAM,
"a connected replica is a sparse cache and cannot export a canonical repository snapshot; export from the authority",
));
self
}
pub fn durability(mut self, durability: ReadDurability) -> Self {
self.durability = durability;
self
}
pub async fn write_to<W>(self, writer: &mut W) -> Result<SnapshotExportReport, LixError>
where
W: AsyncWrite + Unpin + Send + ?Sized,
{
if let Some(error) = self.preflight_error {
return Err(error);
}
if let Some(remote) = self.connected_remote {
return remote.write_to(writer, self.durability).await;
}
if let Some((server, expected_account_id)) = self.remote {
let http = crate::sync::authority_http(&server.headers)?;
let client =
crate::authority_client::open_protocol_client(http, server.url, None).await?;
let mut lease = SnapshotAuthorityLease(Some(client));
let client = lease.0.as_ref().expect("new export lease owns its client");
let result = async {
if client.active_account_id().await? != expected_account_id {
return Err(LixError::new(
LixError::CODE_INVALID_PARAM,
"snapshot authority changed the authenticated account",
));
}
RemoteSnapshotExport {
http: client.http().clone(),
url: client.join_path("snapshot")?,
session_id: client.session_id(),
}
.write_to(writer, self.durability)
.await
}
.await;
let close = lease
.0
.take()
.expect("export owns its client")
.close()
.await;
let report = result?;
close?;
return Ok(report);
}
let storage = self.storage.ok_or_else(|| {
LixError::new(
LixError::CODE_INTERNAL_ERROR,
"snapshot export has no source",
)
})?;
let read = storage
.begin_read(StorageReadOptions {
durability: self.durability,
..StorageReadOptions::default()
})
.await?;
if !self.local_partial_replica && crate::sync::has_any_sync_replica_state(&read).await? {
return Err(LixError::new(
LixError::CODE_INVALID_PARAM,
"a persisted replica is a sparse cache and cannot export a canonical repository snapshot; export from the authority",
));
}
if self.local_partial_replica {
if !crate::init::is_partial_repository_protocol(&read).await?
|| crate::sync::load_partial_replica_state(&read)
.await?
.is_none()
{
return Err(LixError::new(
LixError::CODE_INVALID_PARAM,
"local partial export requires a valid partial replica receipt",
));
}
}
let mut encoder = SnapshotEncoder::new_with_partial(
writer,
crate::init::CURRENT_FORMAT_VERSION,
self.local_partial_replica,
)
.await?;
for space in super::snapshot_spaces() {
let mut cursor = read
.begin_scan(
space,
StorageKeyRange {
lower: Bound::Unbounded,
upper: Bound::Unbounded,
},
StorageBeginScanOptions {
projection: StorageCoreProjection::FullValue,
..StorageBeginScanOptions::default()
},
)
.await?;
loop {
let (entries, has_more) = cursor.next_page(MAX_SCAN_PAGE_ROWS).await?.into_parts();
for entry in entries {
if space == crate::sync::SYNC_AUTHORITY_STATE_SPACE
&& entry.key == crate::sync::authority_state_key()
{
continue;
}
let StorageProjectedValue::FullValue(value) = entry.value else {
return Err(LixError::new(
LixError::CODE_INTERNAL_ERROR,
"full-value snapshot scan returned a key-only entry",
));
};
encoder
.write_entry(&SnapshotEntry {
space_id: space.id.0,
key: entry.key.0,
value,
})
.await?;
}
if !has_more {
break;
}
}
}
let trailer = encoder.finish().await?;
Ok(SnapshotExportReport {
entry_count: trailer.entry_count,
payload_bytes: trailer.payload_bytes,
digest: trailer.digest,
})
}
}
struct SnapshotAuthorityLease(
Option<crate::authority_client::ProtocolClient<crate::sync::AuthorityHttp>>,
);
impl Drop for SnapshotAuthorityLease {
fn drop(&mut self) {
let Some(client) = self.0.take() else {
return;
};
let http = client.http().clone();
crate::authority_client::ProtocolHttp::spawn(
&http,
Box::pin(async move {
let _ = client.close().await;
}),
);
}
}
impl SnapshotExportBuilder<crate::Memory> {
pub(crate) fn remote(
http: crate::sync::AuthorityHttp,
url: Result<String, LixError>,
session_id: Option<String>,
) -> Self {
let (connected_remote, preflight_error) = match url {
Ok(url) => (
Some(RemoteSnapshotExport {
http,
url,
session_id,
}),
None,
),
Err(error) => (None, Some(error)),
};
Self {
storage: None,
durability: ReadDurability::Visible,
preflight_error,
remote: None,
connected_remote,
local_partial_replica: false,
}
}
}
impl RemoteSnapshotExport {
async fn write_to<W>(
self,
writer: &mut W,
durability: ReadDurability,
) -> Result<SnapshotExportReport, LixError>
where
W: AsyncWrite + Unpin + Send + ?Sized,
{
use crate::authority_client::{ProtocolHttp as _, ProtocolHttpRequest};
let mut headers = vec![
(
"accept".to_owned(),
"application/vnd.lix.snapshot".to_owned(),
),
(
"lix-snapshot-durability".to_owned(),
authority_snapshot_durability(durability).to_owned(),
),
];
if let Some(session_id) = self.session_id {
headers.push(("lix-session-id".to_owned(), session_id));
}
let response = self
.http
.request_stream(ProtocolHttpRequest {
method: "GET".to_owned(),
url: self.url,
headers,
body: None,
})
.await?;
if !(200..300).contains(&response.status) {
return Err(LixError::new(
"LIX_REMOTE_REQUEST_FAILED",
format!(
"authority snapshot request failed with HTTP {}",
response.status
),
));
}
write_verified_snapshot_stream(response.body, writer).await
}
}
fn authority_snapshot_durability(durability: ReadDurability) -> &'static str {
match durability {
ReadDurability::Visible => "visible",
ReadDurability::Durable => "durable",
}
}
async fn write_verified_snapshot_stream<S, W>(
mut stream: S,
writer: &mut W,
) -> Result<SnapshotExportReport, LixError>
where
S: futures_core::Stream<Item = Result<bytes::Bytes, LixError>> + Unpin,
W: AsyncWrite + Unpin + Send + ?Sized,
{
let mut total_bytes = 0_u64;
let mut header = [0_u8; super::format::HEADER_BYTES];
let mut header_bytes = 0_usize;
let mut tail = Vec::with_capacity(super::format::TRAILER_BYTES);
let mut hasher = blake3::Hasher::new();
while let Some(chunk) = stream.next().await {
let chunk = chunk?;
total_bytes = total_bytes
.checked_add(u64::try_from(chunk.len()).map_err(|_| {
LixError::new(
LixError::CODE_INVALID_PARAM,
"authority snapshot is too large",
)
})?)
.ok_or_else(|| {
LixError::new(
LixError::CODE_INVALID_PARAM,
"authority snapshot is too large",
)
})?;
let mut remaining = chunk.as_ref();
if header_bytes < header.len() {
let take = remaining.len().min(header.len() - header_bytes);
header[header_bytes..header_bytes + take].copy_from_slice(&remaining[..take]);
header_bytes += take;
remaining = &remaining[take..];
if header_bytes == header.len() {
if super::format::decode_streamed_snapshot_header(&header)?.partial_replica {
return Err(super::format::invalid_snapshot(
"authority returned a partial replica snapshot",
));
}
writer
.write_all(&header)
.await
.map_err(authority_snapshot_io_error)?;
}
}
if header_bytes == header.len() && !remaining.is_empty() {
writer
.write_all(remaining)
.await
.map_err(authority_snapshot_io_error)?;
retain_trailer_and_hash_payload(&mut hasher, &mut tail, remaining);
}
}
if header_bytes != header.len() {
return Err(super::format::invalid_snapshot(
"authority snapshot is truncated in its header",
));
}
let trailer = super::format::decode_streamed_snapshot_trailer(
total_bytes,
*hasher.finalize().as_bytes(),
&tail,
)?;
writer.flush().await.map_err(authority_snapshot_io_error)?;
Ok(SnapshotExportReport {
entry_count: trailer.entry_count,
payload_bytes: trailer.payload_bytes,
digest: trailer.digest,
})
}
fn retain_trailer_and_hash_payload(hasher: &mut blake3::Hasher, tail: &mut Vec<u8>, bytes: &[u8]) {
let finalized = tail
.len()
.saturating_add(bytes.len())
.saturating_sub(super::format::TRAILER_BYTES);
let from_tail = finalized.min(tail.len());
hasher.update(&tail[..from_tail]);
tail.drain(..from_tail);
let from_bytes = finalized - from_tail;
hasher.update(&bytes[..from_bytes]);
tail.extend_from_slice(&bytes[from_bytes..]);
debug_assert!(tail.len() <= super::format::TRAILER_BYTES);
}
fn authority_snapshot_io_error(error: std::io::Error) -> LixError {
LixError::new(
LixError::CODE_UNKNOWN,
format!("write authority snapshot: {error}"),
)
}
#[cfg(test)]
mod tests {
use bytes::Bytes;
use futures_lite::stream;
use super::*;
#[tokio::test]
async fn verified_remote_snapshot_stream_preserves_chunked_bytes_and_report() {
let mut valid = Vec::new();
let trailer = SnapshotEncoder::new(&mut valid, 76)
.await
.expect("encode header")
.finish()
.await
.expect("encode trailer");
let chunks = valid
.chunks(3)
.map(Bytes::copy_from_slice)
.map(Ok)
.collect::<Vec<Result<Bytes, LixError>>>();
let mut forwarded = Vec::new();
let report = write_verified_snapshot_stream(stream::iter(chunks), &mut forwarded)
.await
.expect("valid authority snapshot");
assert_eq!(forwarded, valid);
assert_eq!(report.entry_count, trailer.entry_count);
assert_eq!(report.payload_bytes, trailer.payload_bytes);
assert_eq!(report.digest, trailer.digest);
}
#[tokio::test]
async fn corrupt_remote_headers_are_rejected_before_any_bytes_are_forwarded() {
let mut valid = Vec::new();
SnapshotEncoder::new(&mut valid, 76)
.await
.expect("encode header")
.finish()
.await
.expect("encode trailer");
for (label, offset, value) in [
("magic", 0, b'X'),
(
"container version",
crate::snapshot::format::MAGIC.len() + 1,
2,
),
("checksum", crate::snapshot::format::MAGIC.len() + 2, 0xff),
(
"reserved flags",
crate::snapshot::format::MAGIC.len() + 3,
1,
),
] {
let mut corrupt = valid.clone();
corrupt[offset] = value;
let chunks = stream::iter([Ok(Bytes::from(corrupt))]);
let mut forwarded = Vec::new();
let error = write_verified_snapshot_stream(chunks, &mut forwarded)
.await
.expect_err("corrupt authority header must fail");
assert_eq!(error.code, LixError::CODE_INVALID_SNAPSHOT, "{label}");
assert!(forwarded.is_empty(), "{label}");
}
}
#[test]
fn remote_snapshot_durability_is_preserved_on_the_authority_request() {
assert_eq!(
authority_snapshot_durability(ReadDurability::Visible),
"visible"
);
assert_eq!(
authority_snapshot_durability(ReadDurability::Durable),
"durable"
);
}
}