remora/crates/remora-object-source/tests/compatibility.rs

use std::time::{SystemTime, UNIX_EPOCH};

use remora_domain::OpaqueId;
use remora_object_source::{
    DestinationResumeProof, ObjectResumeCheckpoint, ObjectSource, ObjectSourceConfig,
    ObjectSourceError, ObjectSourceRequest,
};
use remora_outbound::CanonicalUrl;
use remora_stream::{Sha256Accountant, SourceBody, StreamError};
use remora_target::{
    CreateObjectRequest, ObjectIdentity, TargetStore, TargetStoreConfig, attempt_object_key,
};

const PART_SIZE: usize = 5 * 1024 * 1024;

struct Prepared {
    checkpoint: ObjectResumeCheckpoint,
    source_object: ObjectIdentity,
}

fn require_env() -> bool {
    std::env::var("REMORA_OBJECT_SOURCE_REQUIRE").is_ok()
}

fn stores_from_env() -> Option<(TargetStoreConfig, TargetStore, ObjectSource)> {
    let config = match TargetStoreConfig::from_env() {
        Ok(config) => config,
        Err(error) => {
            if require_env() {
                panic!("object source env is required: {error}");
            }
            eprintln!("skipping object source compatibility test: {error}");
            return None;
        }
    };
    let target = TargetStore::connect(config.clone()).expect("target");
    let source = ObjectSource::connect(ObjectSourceConfig::from_target(&config).expect("config"))
        .expect("source");
    Some((config, target, source))
}

fn unique_id(prefix: &str) -> OpaqueId {
    let nanos = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .expect("time")
        .as_nanos();
    OpaqueId::parse(format!("{prefix}{nanos:x}")).expect("opaque id")
}

fn key(label: &str) -> String {
    attempt_object_key(
        &unique_id(&format!("att{label}")),
        &unique_id(&format!("obj{label}")),
        None,
    )
}

fn request(key: &str) -> CreateObjectRequest {
    CreateObjectRequest::parse(key).expect("request")
}

fn digest(bytes: &[u8]) -> remora_stream::IntegrityDigest {
    let mut accountant = Sha256Accountant::new();
    accountant.update(bytes);
    accountant.finalize().digest().expect("digest").clone()
}

async fn read_all(
    mut body: remora_object_source::ObjectSourceBody,
) -> Result<Vec<u8>, StreamError> {
    let mut result = Vec::new();
    let mut buffer = vec![0_u8; 64 * 1024];
    loop {
        let count = body.read(&mut buffer).await?;
        if count == 0 {
            return Ok(result);
        }
        result.extend_from_slice(&buffer[..count]);
    }
}

async fn prepare(
    target: &TargetStore,
    source: &ObjectSource,
    label: &str,
    payload: &[u8],
) -> Prepared {
    let source_key = key(&format!("src{label}"));
    let source_object = target
        .put_create_only(&request(&source_key), payload.to_vec())
        .await
        .expect("source put")
        .identity;
    let opened = source
        .open(ObjectSourceRequest::new(&source_key, "dir/cafe\u{301}.bin").expect("source request"))
        .await
        .expect("source open");
    let destination_key = key(&format!("dst{label}"));
    let session = target
        .begin_multipart(&request(&destination_key))
        .await
        .expect("multipart begin");
    target
        .upload_part(&session, 1, payload[..PART_SIZE].to_vec().into())
        .await
        .expect("part upload");
    let parts = target.list_parts(&session).await.expect("part list");
    let proof = DestinationResumeProof::new(
        CanonicalUrl::parse(&target.config().endpoint).expect("target endpoint"),
        session.identity().clone(),
        parts,
        None,
    )
    .expect("destination proof");
    drop(session);
    let checkpoint = ObjectResumeCheckpoint::new(
        opened.identity().clone(),
        opened.result_path().clone(),
        PART_SIZE as u64,
        digest(&payload[..PART_SIZE]),
        proof,
    )
    .expect("checkpoint");
    Prepared {
        checkpoint,
        source_object,
    }
}

async fn abort_checkpoint(target: &TargetStore, checkpoint: &ObjectResumeCheckpoint) {
    if let Ok((session, _)) = target
        .resume_multipart(checkpoint.destination().identity().clone(), None)
        .await
    {
        target.abort_multipart(session).await.expect("abort");
    }
}

#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn object_source_and_destination_resume_matrix() {
    let Some((config, target, source)) = stores_from_env() else {
        return;
    };
    let payload = [vec![b'a'; PART_SIZE], vec![b'b'; 8192]].concat();
    let replacement = [vec![b'x'; PART_SIZE], vec![b'y'; 8192]].concat();

    let full_key = key("full");
    let full = target
        .put_create_only(&request(&full_key), payload.clone())
        .await
        .expect("full put");
    let opened = source
        .open(ObjectSourceRequest::new(&full_key, "member.bin").expect("request"))
        .await
        .expect("open");
    assert_eq!(
        opened.identity().endpoint(),
        &CanonicalUrl::parse(&config.endpoint).expect("canonical fixture endpoint")
    );
    assert_eq!(opened.identity().bucket(), config.bucket);
    assert_eq!(opened.identity().key(), full_key);
    assert_eq!(opened.identity().size(), payload.len() as u64);
    assert!(opened.identity().is_resumable());
    assert_eq!(read_all(opened.into_body()).await.expect("body"), payload);
    target
        .delete_exact(&full.identity)
        .await
        .expect("full delete");

    let success = prepare(&target, &source, "success", &payload).await;
    drop(source);
    let restarted_source =
        ObjectSource::connect(ObjectSourceConfig::from_target(&config).expect("restart config"))
            .expect("restart source");
    let restarted_target = TargetStore::connect(config.clone()).expect("restart target");
    let mut substituted_config = config.clone();
    substituted_config.endpoint = "http://substituted.invalid:3900".to_owned();
    let substituted_target = TargetStore::connect(substituted_config).expect("substituted target");
    assert_eq!(
        restarted_source
            .resume(
                &success.checkpoint,
                &digest(&payload[..PART_SIZE]),
                &substituted_target,
            )
            .await
            .expect_err("destination endpoint changed"),
        ObjectSourceError::DestinationChanged
    );
    let resumed = restarted_source
        .resume(
            &success.checkpoint,
            &digest(&payload[..PART_SIZE]),
            &restarted_target,
        )
        .await
        .expect("resume after restart");
    assert_eq!(resumed.source().result_path().canonical(), "dir/café.bin");
    let (resumed_source, destination) = resumed.into_parts();
    assert_eq!(
        read_all(resumed_source.into_body()).await.expect("suffix"),
        payload[PART_SIZE..]
    );
    restarted_target
        .abort_multipart(destination)
        .await
        .expect("success abort");
    restarted_target
        .delete_exact(&success.source_object)
        .await
        .expect("success source delete");

    let changed_source = prepare(
        &restarted_target,
        &restarted_source,
        "sourcechange",
        &payload,
    )
    .await;
    restarted_target
        .put_overwrite(
            &request(&changed_source.source_object.key),
            replacement.clone(),
        )
        .await
        .expect("source overwrite");
    assert_eq!(
        restarted_source
            .resume(
                &changed_source.checkpoint,
                &digest(&payload[..PART_SIZE]),
                &restarted_target,
            )
            .await
            .expect_err("source changed"),
        ObjectSourceError::SourceChanged
    );
    abort_checkpoint(&restarted_target, &changed_source.checkpoint).await;
    restarted_target
        .delete_exact(&changed_source.source_object)
        .await
        .expect("changed source delete");

    let changed_part = prepare(&restarted_target, &restarted_source, "partchange", &payload).await;
    let (session, _) = restarted_target
        .resume_multipart(
            changed_part.checkpoint.destination().identity().clone(),
            None,
        )
        .await
        .expect("part session");
    restarted_target
        .upload_part(&session, 1, vec![b'z'; PART_SIZE].into())
        .await
        .expect("part overwrite");
    drop(session);
    assert_eq!(
        restarted_source
            .resume(
                &changed_part.checkpoint,
                &digest(&payload[..PART_SIZE]),
                &restarted_target,
            )
            .await
            .expect_err("destination changed"),
        ObjectSourceError::DestinationChanged
    );
    abort_checkpoint(&restarted_target, &changed_part.checkpoint).await;
    restarted_target
        .delete_exact(&changed_part.source_object)
        .await
        .expect("part source delete");

    let aborted = prepare(&restarted_target, &restarted_source, "aborted", &payload).await;
    let (session, _) = restarted_target
        .resume_multipart(aborted.checkpoint.destination().identity().clone(), None)
        .await
        .expect("aborted session");
    restarted_target
        .abort_multipart(session)
        .await
        .expect("abort");
    assert_eq!(
        restarted_source
            .resume(
                &aborted.checkpoint,
                &digest(&payload[..PART_SIZE]),
                &restarted_target,
            )
            .await
            .expect_err("aborted upload"),
        ObjectSourceError::DestinationChanged
    );
    restarted_target
        .delete_exact(&aborted.source_object)
        .await
        .expect("aborted source delete");

    let completed = prepare(&restarted_target, &restarted_source, "completed", &payload).await;
    let (session, parts) = restarted_target
        .resume_multipart(completed.checkpoint.destination().identity().clone(), None)
        .await
        .expect("completion session");
    let completed_destination = restarted_target
        .complete_multipart(session, &parts)
        .await
        .expect("ambiguous completion");
    assert_eq!(
        restarted_source
            .resume(
                &completed.checkpoint,
                &digest(&payload[..PART_SIZE]),
                &restarted_target,
            )
            .await
            .expect_err("completed upload"),
        ObjectSourceError::DestinationChanged
    );
    restarted_target
        .delete_exact(&completed_destination.identity)
        .await
        .expect("completed destination delete");
    restarted_target
        .delete_exact(&completed.source_object)
        .await
        .expect("completed source delete");

    if restarted_target
        .enable_bucket_versioning()
        .await
        .expect("enable versioning")
    {
        let versioned = prepare(&restarted_target, &restarted_source, "versioned", &payload).await;
        assert!(versioned.checkpoint.source().version_id().is_some());
        restarted_target
            .put_overwrite(&request(&versioned.source_object.key), replacement)
            .await
            .expect("versioned overwrite");
        let resumed = restarted_source
            .resume(
                &versioned.checkpoint,
                &digest(&payload[..PART_SIZE]),
                &restarted_target,
            )
            .await
            .expect("version-pinned resume");
        let (source, destination) = resumed.into_parts();
        assert_eq!(
            read_all(source.into_body()).await.expect("old version"),
            payload[PART_SIZE..]
        );
        restarted_target
            .abort_multipart(destination)
            .await
            .expect("versioned abort");
        restarted_target
            .purge_versions_with_prefix("_remora")
            .await
            .expect("version purge");
    } else if std::env::var("REMORA_TARGET_PROVIDER").as_deref() == Ok("versitygw") {
        panic!("VersityGW must support version-pinned WP-12 qualification");
    }
}