RustProxy/src/cache/range.rs

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ByteRange {
    start: Option<u64>,
    end: Option<u64>,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RangeSelection {
    None,
    Single(ByteRange),
    Unsupported,
}

pub fn request_range(method: &Method, headers: &HeaderMap) -> RangeSelection {
    let mut values = headers.get_all("range").iter();
    let Some(value) = values.next() else {
        return RangeSelection::None;
    };
    // Multiple field lines are equivalent to a combined range-set; RustProxy
    // intentionally delegates those to the origin with multipart ranges.
    if values.next().is_some() || headers.get_all("if-range").iter().count() > 1 {
        return RangeSelection::Unsupported;
    }
    // RustProxy serves one byte range itself. HEAD, unknown units, malformed
    // syntax and multipart requests remain transparent origin pass-through.
    if method != Method::GET {
        return RangeSelection::Unsupported;
    }
    let Ok(value) = value.to_str() else {
        return RangeSelection::Unsupported;
    };
    let Some(spec) = value.trim().strip_prefix("bytes=") else {
        return RangeSelection::Unsupported;
    };
    if spec.is_empty() || spec.contains(',') {
        return RangeSelection::Unsupported;
    }
    let Some((left, right)) = spec.split_once('-') else {
        return RangeSelection::Unsupported;
    };
    let start = if left.trim().is_empty() {
        None
    } else {
        match left.trim().parse::<u64>() {
            Ok(value) => Some(value),
            Err(_) => return RangeSelection::Unsupported,
        }
    };
    let end = if right.trim().is_empty() {
        None
    } else {
        match right.trim().parse::<u64>() {
            Ok(value) => Some(value),
            Err(_) => return RangeSelection::Unsupported,
        }
    };
    if start.is_none() && end.is_none() {
        return RangeSelection::Unsupported;
    }
    if let (Some(start), Some(end)) = (start, end) {
        if start > end {
            return RangeSelection::Unsupported;
        }
    }
    RangeSelection::Single(ByteRange { start, end })
}

fn if_range_matches(headers: &HeaderMap, if_range: Option<&HeaderValue>) -> bool {
    let Some(if_range) = if_range else {
        return true;
    };
    let Ok(candidate) = if_range.to_str() else {
        return false;
    };
    let candidate = candidate.trim();
    if candidate.starts_with('"') || candidate.starts_with("W/\"") {
        if candidate.starts_with("W/") {
            return false;
        }
        return headers
            .get("etag")
            .and_then(|value| value.to_str().ok())
            .is_some_and(|etag| !etag.trim().starts_with("W/") && etag.trim() == candidate);
    }
    let Some(candidate_date) = httpdate::parse_http_date(candidate).ok() else {
        return false;
    };
    headers
        .get("last-modified")
        .and_then(|value| value.to_str().ok())
        .and_then(|value| httpdate::parse_http_date(value).ok())
        .is_some_and(|modified| modified <= candidate_date)
}

fn resolve_byte_range(range: ByteRange, len: usize) -> Option<(usize, usize)> {
    if len == 0 {
        return None;
    }
    let len_u64 = len as u64;
    match (range.start, range.end) {
        (Some(start), Some(end)) => {
            if start >= len_u64 {
                return None;
            }
            let end = end.min(len_u64.saturating_sub(1));
            Some((start as usize, end as usize))
        }
        (Some(start), None) => {
            if start >= len_u64 {
                None
            } else {
                Some((start as usize, len - 1))
            }
        }
        (None, Some(suffix)) => {
            if suffix == 0 {
                return None;
            }
            let take = suffix.min(len_u64) as usize;
            Some((len - take, len - 1))
        }
        (None, None) => None,
    }
}

impl Hit {
    pub fn apply_range(mut self, range: ByteRange, if_range: Option<&HeaderValue>) -> Self {
        if !if_range_matches(self.response.headers(), if_range) {
            return self;
        }
        let full_len = self.body.len();
        self.response
            .headers_mut()
            .insert("accept-ranges", HeaderValue::from_static("bytes"));
        self.response.extensions_mut().insert(crate::compression::Skip);
        let Some((start, end)) = resolve_byte_range(range, full_len) else {
            *self.response.status_mut() = StatusCode::RANGE_NOT_SATISFIABLE;
            *self.response.body_mut() = Body::empty();
            self.response.headers_mut().insert(
                "content-range",
                HeaderValue::from_str(&format!("bytes */{full_len}")).expect("generated range"),
            );
            self.response
                .headers_mut()
                .insert("content-length", HeaderValue::from_static("0"));
            return self;
        };
        let slice = self.body.slice(start..=end);
        *self.response.status_mut() = StatusCode::PARTIAL_CONTENT;
        *self.response.body_mut() = Body::from(slice.clone());
        self.response.headers_mut().insert(
            "content-range",
            HeaderValue::from_str(&format!("bytes {start}-{end}/{full_len}"))
                .expect("generated range"),
        );
        self.response.headers_mut().insert(
            "content-length",
            HeaderValue::from_str(&slice.len().to_string()).expect("generated length"),
        );
        self.body = slice;
        self
    }
}

impl Revalidation {
    pub fn response_with_range(
        &self,
        headers: HeaderMap,
        range: ByteRange,
        if_range: Option<&HeaderValue>,
    ) -> Response {
        let hit = Hit {
            response: self.response(headers),
            body: self.body.clone(),
            hits: self.hits,
            transformed: self.transformed,
            remaining_ttl: self.ttl.saturating_sub(self.age),
            stale: false,
            background_revalidate: false,
            resident_age_seconds: self.resident_age_at_lookup.as_secs(),
        };
        hit.apply_range(range, if_range).response
    }
}