remora/crates/remora-target/src/multipart.rs

use std::fmt;
use std::time::{Duration, SystemTime};

use aws_credential_types::Credentials;
use aws_sigv4::http_request::{
    PayloadChecksumKind, SignableBody, SignableRequest, SigningSettings, UriPathNormalizationMode,
    sign,
};
use aws_sigv4::sign::v4;
use bytes::Bytes;
use http::header::{CONTENT_LENGTH, CONTENT_TYPE, HeaderName, HeaderValue};
use http::{Method, Request};
use percent_encoding::{AsciiSet, CONTROLS, utf8_percent_encode};
use quick_xml::de::from_str;
use quick_xml::se::to_string;
use serde::{Deserialize, Serialize};

use crate::config::TargetStoreConfig;
use crate::error::{TargetError, map_http_status, redact_sensitive};
use crate::identity::{CompletedPart, MultipartIdentity, REMORA_SHA256_METADATA_KEY};

/// Encode path segments for path-style S3 URLs.
const PATH_SET: &AsciiSet = &CONTROLS
    .add(b' ')
    .add(b'"')
    .add(b'#')
    .add(b'%')
    .add(b'&')
    .add(b'+')
    .add(b',')
    .add(b'<')
    .add(b'>')
    .add(b'?')
    .add(b'@')
    .add(b'[')
    .add(b'\\')
    .add(b']')
    .add(b'^')
    .add(b'`')
    .add(b'{')
    .add(b'|')
    .add(b'}');

/// One version or delete marker listed under a key/prefix.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct ObjectVersion {
    pub key: String,
    pub version_id: String,
    pub is_delete_marker: bool,
}

pub struct ObjectHead {
    pub version_id: Option<String>,
}

/// Explicit multipart control plane for destination recovery.
///
/// OpenDAL streams multipart writes but does not expose list-parts recovery.
/// This client issues path-style S3 multipart APIs with the same credentials
/// as the OpenDAL operator so exact upload IDs and re-listed parts are durable.
#[derive(Clone)]
pub struct MultipartClient {
    config: TargetStoreConfig,
    http: reqwest::Client,
}

impl fmt::Debug for MultipartClient {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        formatter
            .debug_struct("MultipartClient")
            .field("config", &self.config)
            .finish_non_exhaustive()
    }
}

impl MultipartClient {
    pub fn new(config: TargetStoreConfig) -> Result<Self, TargetError> {
        let http = reqwest::Client::builder()
            .timeout(Duration::from_secs(120))
            .build()
            .map_err(|error| TargetError::internal(redact_sensitive(error.to_string())))?;
        Ok(Self { config, http })
    }

    pub async fn list_uploads_exact(
        &self,
        key: &str,
    ) -> Result<Vec<MultipartIdentity>, TargetError> {
        let rooted_key = rooted_key(&self.config, key);
        let mut uploads = Vec::new();
        let mut key_marker = String::new();
        let mut upload_id_marker = String::new();
        loop {
            let mut uri = format!(
                "{}/{}/?uploads&prefix={}",
                self.config.endpoint,
                self.config.bucket,
                encode_query(&rooted_key)
            );
            if !key_marker.is_empty() {
                uri.push_str(&format!("&key-marker={}", encode_query(&key_marker)));
            }
            if !upload_id_marker.is_empty() {
                uri.push_str(&format!(
                    "&upload-id-marker={}",
                    encode_query(&upload_id_marker)
                ));
            }
            let request = Request::builder()
                .method(Method::GET)
                .uri(uri)
                .body(Bytes::new())
                .map_err(|error| TargetError::internal(error.to_string()))?;
            let response = self.send(request).await?;
            let status = response.status().as_u16();
            let body = response
                .text()
                .await
                .map_err(|error| TargetError::transient(redact_sensitive(error.to_string())))?;
            if status / 100 != 2 {
                return Err(map_http_status(status, &body));
            }
            let parsed: ListMultipartUploadsResult = from_str(&body)
                .map_err(|error| TargetError::internal(redact_sensitive(error.to_string())))?;
            uploads.extend(
                parsed
                    .uploads
                    .into_iter()
                    .filter(|upload| upload.key == rooted_key && !upload.upload_id.is_empty())
                    .map(|upload| MultipartIdentity {
                        target_store_id: self.config.target_store_id.clone(),
                        bucket: self.config.bucket.clone(),
                        key: key.to_owned(),
                        upload_id: upload.upload_id,
                    }),
            );
            if parsed.is_truncated {
                key_marker = parsed.next_key_marker.unwrap_or_default();
                upload_id_marker = parsed.next_upload_id_marker.unwrap_or_default();
                if key_marker.is_empty() && upload_id_marker.is_empty() {
                    return Err(TargetError::internal(
                        "truncated multipart listing omitted continuation markers",
                    ));
                }
            } else {
                break;
            }
        }
        Ok(uploads)
    }

    pub async fn head_object(&self, key: &str) -> Result<ObjectHead, TargetError> {
        let path = object_path(&self.config, key);
        let request = Request::builder()
            .method(Method::HEAD)
            .uri(format!("{}{}", self.config.endpoint, path))
            .body(Bytes::new())
            .map_err(|error| TargetError::internal(error.to_string()))?;
        let response = self.send(request).await?;
        let status = response.status().as_u16();
        if status == 404 {
            return Err(TargetError::NotFound);
        }
        if status / 100 != 2 {
            let body = response.text().await.unwrap_or_default();
            return Err(map_http_status(status, &body));
        }
        let version_id = response
            .headers()
            .get("x-amz-version-id")
            .and_then(|value| value.to_str().ok())
            .filter(|value| !value.is_empty())
            .map(str::to_owned);
        Ok(ObjectHead { version_id })
    }

    /// Enable bucket versioning when the peer supports it.
    ///
    /// Peers that return NotImplemented are reported as `false` without error.
    #[cfg(feature = "test-utils")]
    pub async fn enable_bucket_versioning(&self) -> Result<bool, TargetError> {
        let body = r#"<?xml version="1.0" encoding="UTF-8"?>
<VersioningConfiguration xmlns="http://s3.amazonaws.com/doc/2006-03-01/">
  <Status>Enabled</Status>
</VersioningConfiguration>"#;
        let uri = format!(
            "{}/{}/?versioning",
            self.config.endpoint, self.config.bucket
        );
        let request = Request::builder()
            .method(Method::PUT)
            .uri(uri)
            .header(CONTENT_TYPE, "application/xml")
            .header(CONTENT_LENGTH, body.len().to_string())
            .body(Bytes::from_static(body.as_bytes()))
            .map_err(|error| TargetError::internal(error.to_string()))?;
        let response = self.send(request).await?;
        let status = response.status().as_u16();
        if status / 100 == 2 {
            return Ok(true);
        }
        let text = response.text().await.unwrap_or_default();
        if status == 501
            || text.contains("NotImplemented")
            || text.contains("XNotImplemented")
            || text.contains("NotSupported")
        {
            return Ok(false);
        }
        Err(map_http_status(status, &text))
    }

    /// Start a multipart upload. `content_sha256` is optional S3 user metadata
    /// for an independently computed digest only.
    pub async fn initiate(
        &self,
        key: &str,
        content_sha256: Option<&str>,
    ) -> Result<MultipartIdentity, TargetError> {
        let path = object_path(&self.config, key);
        let mut builder = Request::builder()
            .method(Method::POST)
            .uri(format!("{}{}?uploads", self.config.endpoint, path));
        if let Some(digest) = content_sha256 {
            let header = format!("x-amz-meta-{REMORA_SHA256_METADATA_KEY}");
            builder = builder.header(
                HeaderName::from_bytes(header.as_bytes())
                    .map_err(|_| TargetError::internal("invalid metadata header name"))?,
                HeaderValue::from_str(digest)
                    .map_err(|_| TargetError::validation("invalid sha256 metadata"))?,
            );
        }
        let request = builder
            .header(CONTENT_LENGTH, "0")
            .body(Bytes::new())
            .map_err(|error| TargetError::internal(error.to_string()))?;
        let response = self.send(request).await?;
        let status = response.status().as_u16();
        let body = response
            .text()
            .await
            .map_err(|error| TargetError::transient(redact_sensitive(error.to_string())))?;
        if status / 100 != 2 {
            return Err(map_http_status(status, &body));
        }
        let parsed: InitiateMultipartUploadResult = from_str(&body)
            .map_err(|error| TargetError::internal(redact_sensitive(error.to_string())))?;
        if parsed.upload_id.is_empty() {
            return Err(TargetError::internal(
                "multipart initiate missing upload id",
            ));
        }
        Ok(MultipartIdentity {
            target_store_id: self.config.target_store_id.clone(),
            bucket: self.config.bucket.clone(),
            key: key.to_owned(),
            upload_id: parsed.upload_id,
        })
    }

    pub async fn upload_part(
        &self,
        identity: &MultipartIdentity,
        part_number: u32,
        body: bytes::Bytes,
    ) -> Result<CompletedPart, TargetError> {
        if part_number == 0 {
            return Err(TargetError::validation("part number must be >= 1"));
        }
        let path = object_path(&self.config, &identity.key);
        let uri = format!(
            "{}{}?partNumber={part_number}&uploadId={}",
            self.config.endpoint,
            path,
            encode_query(&identity.upload_id)
        );
        let len = body.len();
        let request = Request::builder()
            .method(Method::PUT)
            .uri(uri)
            .header(CONTENT_LENGTH, len.to_string())
            .body(body)
            .map_err(|error| TargetError::internal(error.to_string()))?;
        let response = self.send(request).await?;
        let status = response.status().as_u16();
        if status / 100 != 2 {
            let body = response.text().await.unwrap_or_default();
            return Err(map_http_status(status, &body));
        }
        let etag = response
            .headers()
            .get("etag")
            .and_then(|value| value.to_str().ok())
            .map(|value| value.trim_matches('"').to_owned())
            .ok_or_else(|| TargetError::internal("multipart part missing etag"))?;
        Ok(CompletedPart {
            part_number,
            etag,
            size: Some(len as u64),
        })
    }

    pub async fn list_parts(
        &self,
        identity: &MultipartIdentity,
    ) -> Result<Vec<CompletedPart>, TargetError> {
        let mut parts = Vec::new();
        let mut part_number_marker: Option<u32> = None;
        loop {
            let path = object_path(&self.config, &identity.key);
            let mut uri = format!(
                "{}{}?uploadId={}",
                self.config.endpoint,
                path,
                encode_query(&identity.upload_id)
            );
            if let Some(marker) = part_number_marker {
                uri.push_str(&format!("&part-number-marker={marker}"));
            }
            let request = Request::builder()
                .method(Method::GET)
                .uri(uri)
                .body(Bytes::new())
                .map_err(|error| TargetError::internal(error.to_string()))?;
            let response = self.send(request).await?;
            let status = response.status().as_u16();
            let body = response
                .text()
                .await
                .map_err(|error| TargetError::transient(redact_sensitive(error.to_string())))?;
            if status / 100 != 2 {
                return Err(map_http_status(status, &body));
            }
            let parsed: ListPartsResult = from_str(&body)
                .map_err(|error| TargetError::internal(redact_sensitive(error.to_string())))?;
            for part in parsed.parts {
                parts.push(CompletedPart {
                    part_number: part.part_number,
                    etag: part.etag.trim_matches('"').to_owned(),
                    size: Some(part.size),
                });
            }
            if parsed.is_truncated {
                part_number_marker = parsed
                    .next_part_number_marker
                    .or_else(|| parts.last().map(|part| part.part_number));
                if part_number_marker.is_none() {
                    break;
                }
            } else {
                break;
            }
        }
        parts.sort_by_key(|part| part.part_number);
        Ok(parts)
    }

    pub async fn complete(
        &self,
        identity: &MultipartIdentity,
        parts: &[CompletedPart],
        if_not_exists: bool,
    ) -> Result<Option<String>, TargetError> {
        if parts.is_empty() {
            return Err(TargetError::validation(
                "multipart complete requires at least one part",
            ));
        }
        let payload = CompleteMultipartUpload {
            parts: parts
                .iter()
                .map(|part| CompletePart {
                    part_number: part.part_number,
                    etag: format!("\"{}\"", part.etag.trim_matches('"')),
                })
                .collect(),
        };
        let xml = to_string(&payload)
            .map_err(|error| TargetError::internal(redact_sensitive(error.to_string())))?;
        let path = object_path(&self.config, &identity.key);
        let uri = format!(
            "{}{}?uploadId={}",
            self.config.endpoint,
            path,
            encode_query(&identity.upload_id)
        );
        let mut builder = Request::builder()
            .method(Method::POST)
            .uri(uri)
            .header(CONTENT_TYPE, "application/xml")
            .header(CONTENT_LENGTH, xml.len().to_string());
        if if_not_exists {
            builder = builder.header("If-None-Match", "*");
        }
        let request = builder
            .body(Bytes::from(xml))
            .map_err(|error| TargetError::internal(error.to_string()))?;
        let response = self.send(request).await?;
        let status = response.status().as_u16();
        let version_id = response
            .headers()
            .get("x-amz-version-id")
            .and_then(|value| value.to_str().ok())
            .map(str::to_owned);
        let body = response
            .text()
            .await
            .map_err(|error| TargetError::transient(redact_sensitive(error.to_string())))?;
        if status / 100 != 2 {
            return Err(map_http_status(status, &body));
        }
        Ok(version_id.filter(|value| !value.is_empty()))
    }

    pub async fn abort(&self, identity: &MultipartIdentity) -> Result<(), TargetError> {
        let path = object_path(&self.config, &identity.key);
        let uri = format!(
            "{}{}?uploadId={}",
            self.config.endpoint,
            path,
            encode_query(&identity.upload_id)
        );
        let request = Request::builder()
            .method(Method::DELETE)
            .uri(uri)
            .body(Bytes::new())
            .map_err(|error| TargetError::internal(error.to_string()))?;
        let response = self.send(request).await?;
        let status = response.status().as_u16();
        if status == 404 {
            return Ok(());
        }
        if status / 100 != 2 {
            let body = response.text().await.unwrap_or_default();
            return Err(map_http_status(status, &body));
        }
        Ok(())
    }

    /// List object versions and delete markers under an optional key prefix.
    pub async fn list_object_versions(
        &self,
        prefix: &str,
    ) -> Result<Vec<ObjectVersion>, TargetError> {
        let mut versions = Vec::new();
        let mut key_marker = String::new();
        let mut version_id_marker = String::new();
        loop {
            let mut uri = format!("{}/{}/?versions", self.config.endpoint, self.config.bucket);
            let rooted_prefix = rooted_key(&self.config, prefix);
            if !rooted_prefix.is_empty() {
                uri.push_str(&format!("&prefix={}", encode_query(&rooted_prefix)));
            }
            if !key_marker.is_empty() {
                uri.push_str(&format!("&key-marker={}", encode_query(&key_marker)));
            }
            if !version_id_marker.is_empty() {
                uri.push_str(&format!(
                    "&version-id-marker={}",
                    encode_query(&version_id_marker)
                ));
            }
            let request = Request::builder()
                .method(Method::GET)
                .uri(uri)
                .body(Bytes::new())
                .map_err(|error| TargetError::internal(error.to_string()))?;
            let response = self.send(request).await?;
            let status = response.status().as_u16();
            let body = response
                .text()
                .await
                .map_err(|error| TargetError::transient(redact_sensitive(error.to_string())))?;
            if status == 501
                && (body.contains("NotImplemented")
                    || body.contains("XNotImplemented")
                    || body.contains("NotSupported"))
            {
                return Ok(Vec::new());
            }
            if status / 100 != 2 {
                return Err(map_http_status(status, &body));
            }
            let parsed: ListVersionsResult = from_str(&body)
                .map_err(|error| TargetError::internal(redact_sensitive(error.to_string())))?;
            for version in parsed.versions {
                if !version.version_id.is_empty() {
                    versions.push(ObjectVersion {
                        key: logical_key(&self.config, &version.key),
                        version_id: version.version_id,
                        is_delete_marker: false,
                    });
                }
            }
            for marker in parsed.delete_markers {
                if !marker.version_id.is_empty() {
                    versions.push(ObjectVersion {
                        key: logical_key(&self.config, &marker.key),
                        version_id: marker.version_id,
                        is_delete_marker: true,
                    });
                }
            }
            if parsed.is_truncated {
                key_marker = parsed.next_key_marker.unwrap_or_default();
                version_id_marker = parsed.next_version_id_marker.unwrap_or_default();
                if key_marker.is_empty() && version_id_marker.is_empty() {
                    break;
                }
            } else {
                break;
            }
        }
        Ok(versions)
    }

    /// Delete one exact object version or delete marker.
    #[cfg(feature = "test-utils")]
    pub async fn delete_object_version(
        &self,
        key: &str,
        version_id: &str,
    ) -> Result<(), TargetError> {
        let path = object_path(&self.config, key);
        let uri = format!(
            "{}{}?versionId={}",
            self.config.endpoint,
            path,
            encode_query(version_id)
        );
        let request = Request::builder()
            .method(Method::DELETE)
            .uri(uri)
            .body(Bytes::new())
            .map_err(|error| TargetError::internal(error.to_string()))?;
        let response = self.send(request).await?;
        let status = response.status().as_u16();
        if status == 404 || status / 100 == 2 {
            return Ok(());
        }
        let body = response.text().await.unwrap_or_default();
        Err(map_http_status(status, &body))
    }

    async fn send(&self, request: Request<Bytes>) -> Result<reqwest::Response, TargetError> {
        let (parts, body) = sign_request(&self.config, request)?.into_parts();
        self.http
            .request(parts.method, parts.uri.to_string())
            .headers(parts.headers)
            .body(body)
            .send()
            .await
            .map_err(|error| TargetError::transient(redact_sensitive(error.to_string())))
    }
}

fn object_path(config: &TargetStoreConfig, key: &str) -> String {
    let key = key.trim_start_matches('/');
    let rooted = rooted_key(config, key);
    let encoded = rooted
        .split('/')
        .map(|segment| utf8_percent_encode(segment, PATH_SET).to_string())
        .collect::<Vec<_>>()
        .join("/");
    format!("/{}/{}", config.bucket, encoded)
}

fn rooted_key(config: &TargetStoreConfig, key: &str) -> String {
    let key = key.trim_start_matches('/');
    let root = config.root.trim_matches('/');
    if root.is_empty() {
        key.to_owned()
    } else if key.is_empty() {
        root.to_owned()
    } else {
        format!("{root}/{key}")
    }
}

fn logical_key(config: &TargetStoreConfig, key: &str) -> String {
    let root = config.root.trim_matches('/');
    if root.is_empty() {
        key.to_owned()
    } else {
        key.strip_prefix(root)
            .and_then(|value| value.strip_prefix('/'))
            .unwrap_or(key)
            .to_owned()
    }
}

fn encode_query(value: &str) -> String {
    utf8_percent_encode(value, PATH_SET).to_string()
}

fn sign_request(
    config: &TargetStoreConfig,
    mut request: Request<Bytes>,
) -> Result<Request<Bytes>, TargetError> {
    let identity = Credentials::new(
        &config.access_key_id,
        &config.secret_access_key,
        None,
        None,
        "remora-target",
    )
    .into();
    let mut settings = SigningSettings::default();
    // S3-compatible stores (including Garage) require x-amz-content-sha256 and
    // reject URI path normalization in the canonical request.
    settings.payload_checksum_kind = PayloadChecksumKind::XAmzSha256;
    settings.uri_path_normalization_mode = UriPathNormalizationMode::Disabled;
    let signing_params = v4::SigningParams::builder()
        .identity(&identity)
        .region(&config.region)
        .name("s3")
        .time(SystemTime::now())
        .settings(settings)
        .build()
        .map_err(|error| TargetError::internal(redact_sensitive(error.to_string())))?
        .into();
    let signable = SignableRequest::new(
        request.method().as_str(),
        request.uri().to_string(),
        request
            .headers()
            .iter()
            .filter_map(|(name, value)| value.to_str().ok().map(|value| (name.as_str(), value))),
        SignableBody::Bytes(request.body().as_ref()),
    )
    .map_err(|error| TargetError::internal(redact_sensitive(error.to_string())))?;
    let (instructions, _signature) = sign(signable, &signing_params)
        .map_err(|error| TargetError::internal(redact_sensitive(error.to_string())))?
        .into_parts();
    instructions.apply_to_request_http1x(&mut request);
    Ok(request)
}

#[derive(Debug, Deserialize)]
struct InitiateMultipartUploadResult {
    #[serde(rename = "UploadId", default)]
    upload_id: String,
}

#[derive(Debug, Deserialize)]
struct ListMultipartUploadsResult {
    #[serde(rename = "Upload", default)]
    uploads: Vec<ListedMultipartUpload>,
    #[serde(rename = "IsTruncated", default)]
    is_truncated: bool,
    #[serde(rename = "NextKeyMarker", default)]
    next_key_marker: Option<String>,
    #[serde(rename = "NextUploadIdMarker", default)]
    next_upload_id_marker: Option<String>,
}

#[derive(Debug, Deserialize)]
struct ListedMultipartUpload {
    #[serde(rename = "Key", default)]
    key: String,
    #[serde(rename = "UploadId", default)]
    upload_id: String,
}

#[derive(Debug, Deserialize)]
struct ListPartsResult {
    #[serde(rename = "Part", default)]
    parts: Vec<ListPart>,
    #[serde(rename = "IsTruncated", default)]
    is_truncated: bool,
    #[serde(rename = "NextPartNumberMarker", default)]
    next_part_number_marker: Option<u32>,
}

#[derive(Debug, Deserialize)]
struct ListPart {
    #[serde(rename = "PartNumber")]
    part_number: u32,
    #[serde(rename = "ETag")]
    etag: String,
    #[serde(rename = "Size", default)]
    size: u64,
}

#[derive(Debug, Serialize)]
#[serde(rename = "CompleteMultipartUpload")]
struct CompleteMultipartUpload {
    #[serde(rename = "Part")]
    parts: Vec<CompletePart>,
}

#[derive(Debug, Serialize)]
struct CompletePart {
    #[serde(rename = "PartNumber")]
    part_number: u32,
    #[serde(rename = "ETag")]
    etag: String,
}

#[derive(Debug, Deserialize)]
struct ListVersionsResult {
    #[serde(rename = "Version", default)]
    versions: Vec<ListedVersion>,
    #[serde(rename = "DeleteMarker", default)]
    delete_markers: Vec<ListedVersion>,
    #[serde(rename = "IsTruncated", default)]
    is_truncated: bool,
    #[serde(rename = "NextKeyMarker", default)]
    next_key_marker: Option<String>,
    #[serde(rename = "NextVersionIdMarker", default)]
    next_version_id_marker: Option<String>,
}

#[derive(Debug, Deserialize)]
struct ListedVersion {
    #[serde(rename = "Key", default)]
    key: String,
    #[serde(rename = "VersionId", default)]
    version_id: String,
}