remora/crates/remora-http-source/src/source.rs

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

use http::{HeaderMap, HeaderValue, StatusCode, header};
use remora_domain::MemberPath;
use remora_outbound::{
    CanonicalUrl, OutboundClient, OutboundError, OutboundPolicy, OutboundRequest, Resolver,
};
use remora_stream::{IntegrityDigest, RetryTiming, capped_server_hint};
use tokio::time::timeout;

use crate::{
    HttpResumeCheckpoint, HttpSourceBody, HttpSourceError, HttpSourceIdentity, StrongEtag,
};

pub struct HttpSourceRequest {
    url: CanonicalUrl,
    result_path: MemberPath,
    credential_headers: HeaderMap,
}

impl HttpSourceRequest {
    pub fn new(url: &str, result_path: &str) -> Result<Self, HttpSourceError> {
        Ok(Self {
            url: CanonicalUrl::parse(url).map_err(|_| HttpSourceError::InvalidRequest)?,
            result_path: MemberPath::parse(result_path)
                .map_err(|_| HttpSourceError::InvalidRequest)?,
            credential_headers: HeaderMap::new(),
        })
    }

    pub fn with_credential_headers(mut self, headers: HeaderMap) -> Result<Self, HttpSourceError> {
        OutboundRequest::get(self.url.as_url().as_str())
            .map_err(|_| HttpSourceError::InvalidRequest)?
            .with_credential_headers(headers.clone())
            .map_err(|_| HttpSourceError::InvalidRequest)?;
        self.credential_headers = headers;
        Ok(self)
    }
}

impl std::fmt::Debug for HttpSourceRequest {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        formatter
            .debug_struct("HttpSourceRequest")
            .field("result_path", &self.result_path.canonical())
            .field("credential_header_count", &self.credential_headers.len())
            .finish_non_exhaustive()
    }
}

pub struct OpenedHttpSource {
    result_path: MemberPath,
    identity: Option<HttpSourceIdentity>,
    body: HttpSourceBody,
}

impl std::fmt::Debug for OpenedHttpSource {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        formatter
            .debug_struct("OpenedHttpSource")
            .field("result_path", &self.result_path.canonical())
            .field("has_resume_identity", &self.identity.is_some())
            .finish_non_exhaustive()
    }
}

impl OpenedHttpSource {
    #[must_use]
    pub fn result_path(&self) -> &MemberPath {
        &self.result_path
    }

    #[must_use]
    pub fn identity(&self) -> Option<&HttpSourceIdentity> {
        self.identity.as_ref()
    }

    #[must_use]
    pub fn into_body(self) -> HttpSourceBody {
        self.body
    }
}

pub struct HttpSource<R> {
    client: OutboundClient<R>,
    timing: RetryTiming,
}

impl<R: Resolver> HttpSource<R> {
    #[must_use]
    pub fn new(resolver: R, policy: OutboundPolicy, timing: RetryTiming) -> Self {
        let timing = timing.sanitized();
        Self {
            client: OutboundClient::builder(resolver, policy)
                .connect_timeout(timing.connect_timeout)
                .build(),
            timing,
        }
    }

    pub async fn open(
        &self,
        request: HttpSourceRequest,
    ) -> Result<OpenedHttpSource, HttpSourceError> {
        let requested_url = request.url.clone();
        let outbound = build_request(&request.url, request.credential_headers, None)?;
        let response = self.execute(outbound).await?;
        let status = response.response().status();
        if status != StatusCode::OK {
            return Err(status_error(
                status,
                response.response().headers(),
                &self.timing,
            ));
        }
        ensure_identity_encoding(response.response().headers())?;
        let expected_length = content_length(response.response().headers())?;
        let identity = resumable_identity(
            requested_url,
            response.final_url().clone(),
            response.response().headers(),
            expected_length,
        );
        Ok(OpenedHttpSource {
            result_path: request.result_path,
            identity,
            body: HttpSourceBody::new(response.into_response(), expected_length),
        })
    }

    pub async fn resume(
        &self,
        checkpoint: &HttpResumeCheckpoint,
        observed_prefix: &IntegrityDigest,
        credential_headers: HeaderMap,
    ) -> Result<OpenedHttpSource, HttpSourceError> {
        if checkpoint.prefix_integrity() != observed_prefix {
            return Err(HttpSourceError::SourceChanged);
        }
        let outbound = build_request(
            checkpoint.identity().requested_url(),
            credential_headers,
            Some(checkpoint),
        )?;
        let response = self.execute(outbound).await?;
        if response.response().status() != StatusCode::PARTIAL_CONTENT {
            return Err(match response.response().status() {
                StatusCode::UNAUTHORIZED | StatusCode::FORBIDDEN => HttpSourceError::Authorization,
                StatusCode::NOT_FOUND | StatusCode::GONE => HttpSourceError::NotFound,
                status if status.is_server_error() || status == StatusCode::TOO_MANY_REQUESTS => {
                    status_error(status, response.response().headers(), &self.timing)
                }
                _ => HttpSourceError::SourceChanged,
            });
        }
        validate_resume_response(checkpoint, &response)?;
        let remaining = checkpoint.identity().representation_length() - checkpoint.offset();
        Ok(OpenedHttpSource {
            result_path: checkpoint.result_path().clone(),
            identity: Some(checkpoint.identity().clone()),
            body: HttpSourceBody::new(response.into_response(), Some(remaining)),
        })
    }

    async fn execute(
        &self,
        request: OutboundRequest,
    ) -> Result<remora_outbound::OutboundResponse, HttpSourceError> {
        match timeout(
            self.timing.response_start_timeout,
            self.client.execute(request),
        )
        .await
        {
            Ok(Ok(response)) => Ok(response),
            Ok(Err(error)) => Err(map_outbound_error(error)),
            Err(_) => Err(HttpSourceError::Transient { retry_after: None }),
        }
    }
}

fn build_request(
    url: &CanonicalUrl,
    credentials: HeaderMap,
    checkpoint: Option<&HttpResumeCheckpoint>,
) -> Result<OutboundRequest, HttpSourceError> {
    let mut public = HeaderMap::new();
    public.insert(
        header::ACCEPT_ENCODING,
        HeaderValue::from_static("identity"),
    );
    if let Some(checkpoint) = checkpoint {
        public.insert(
            header::RANGE,
            HeaderValue::from_str(&format!("bytes={}-", checkpoint.offset()))
                .map_err(|_| HttpSourceError::InvalidRequest)?,
        );
        public.insert(
            header::IF_MATCH,
            HeaderValue::from_str(checkpoint.identity().etag().as_str())
                .map_err(|_| HttpSourceError::InvalidRequest)?,
        );
    }
    OutboundRequest::get(url.as_url().as_str())
        .and_then(|request| request.with_public_headers(public))
        .and_then(|request| request.with_credential_headers(credentials))
        .map_err(|_| HttpSourceError::InvalidRequest)
}

fn ensure_identity_encoding(headers: &HeaderMap) -> Result<(), HttpSourceError> {
    match headers.get(header::CONTENT_ENCODING) {
        None => Ok(()),
        Some(value) if value.as_bytes().eq_ignore_ascii_case(b"identity") => Ok(()),
        Some(_) => Err(HttpSourceError::InvalidResponse),
    }
}

fn content_length(headers: &HeaderMap) -> Result<Option<u64>, HttpSourceError> {
    headers
        .get(header::CONTENT_LENGTH)
        .map(|value| {
            value
                .to_str()
                .map_err(|_| HttpSourceError::InvalidResponse)?
                .parse::<u64>()
                .map_err(|_| HttpSourceError::InvalidResponse)
        })
        .transpose()
}

fn resumable_identity(
    requested_url: CanonicalUrl,
    final_url: CanonicalUrl,
    headers: &HeaderMap,
    length: Option<u64>,
) -> Option<HttpSourceIdentity> {
    let etag = headers
        .get(header::ETAG)?
        .to_str()
        .ok()
        .and_then(|value| StrongEtag::parse(value).ok())?;
    Some(HttpSourceIdentity::from_parts(
        requested_url,
        final_url,
        etag,
        length?,
    ))
}

fn validate_resume_response(
    checkpoint: &HttpResumeCheckpoint,
    response: &remora_outbound::OutboundResponse,
) -> Result<(), HttpSourceError> {
    ensure_identity_encoding(response.response().headers())
        .map_err(|_| HttpSourceError::SourceChanged)?;
    if response.final_url() != checkpoint.identity().final_url() {
        return Err(HttpSourceError::SourceChanged);
    }
    let etag = response
        .response()
        .headers()
        .get(header::ETAG)
        .and_then(|value| value.to_str().ok())
        .and_then(|value| StrongEtag::parse(value).ok())
        .ok_or(HttpSourceError::SourceChanged)?;
    if &etag != checkpoint.identity().etag() {
        return Err(HttpSourceError::SourceChanged);
    }
    let total = checkpoint.identity().representation_length();
    let remaining = total - checkpoint.offset();
    if content_length(response.response().headers())? != Some(remaining) {
        return Err(HttpSourceError::SourceChanged);
    }
    let range = response
        .response()
        .headers()
        .get(header::CONTENT_RANGE)
        .and_then(|value| value.to_str().ok())
        .ok_or(HttpSourceError::SourceChanged)?;
    let expected = format!("bytes {}-{}/{total}", checkpoint.offset(), total - 1);
    if range != expected {
        return Err(HttpSourceError::SourceChanged);
    }
    Ok(())
}

fn status_error(status: StatusCode, headers: &HeaderMap, timing: &RetryTiming) -> HttpSourceError {
    match status {
        StatusCode::UNAUTHORIZED | StatusCode::FORBIDDEN => HttpSourceError::Authorization,
        StatusCode::NOT_FOUND | StatusCode::GONE => HttpSourceError::NotFound,
        status if status == StatusCode::TOO_MANY_REQUESTS || status.is_server_error() => {
            HttpSourceError::Transient {
                retry_after: parse_retry_after(headers, timing, SystemTime::now()),
            }
        }
        _ => HttpSourceError::InvalidResponse,
    }
}

fn parse_retry_after(
    headers: &HeaderMap,
    timing: &RetryTiming,
    now: SystemTime,
) -> Option<Duration> {
    let value = headers.get(header::RETRY_AFTER)?.to_str().ok()?;
    let hint = value
        .parse::<u64>()
        .ok()
        .map(Duration::from_secs)
        .or_else(|| {
            httpdate::parse_http_date(value)
                .ok()?
                .duration_since(now)
                .ok()
        });
    capped_server_hint(hint, timing)
}

fn map_outbound_error(error: OutboundError) -> HttpSourceError {
    match error {
        OutboundError::ResolutionFailed | OutboundError::Transport => {
            HttpSourceError::Transient { retry_after: None }
        }
        OutboundError::InvalidUrl | OutboundError::InvalidHeaders => {
            HttpSourceError::InvalidRequest
        }
        OutboundError::EmptyResolution
        | OutboundError::AddressDenied
        | OutboundError::SchemePortDenied
        | OutboundError::HttpsDowngrade => HttpSourceError::Authorization,
        OutboundError::InvalidRedirect | OutboundError::RedirectLimit => {
            HttpSourceError::InvalidResponse
        }
    }
}