unitycatalog_storage_proxy/
handler.rs

1//! The streaming HTTP handler: a blanket [`StorageProxyHandler`] over any
2//! [`StorageProxyBackend`].
3//!
4//! This is where all the wire semantics live — `Range` → `206`/`Content-Range`,
5//! `ETag`/`Accept-Ranges` relay, `If-Match` → conditional `PUT` → `412`, and
6//! **bounded-memory streaming** in both directions. A backend implementer never
7//! touches any of it; they only authorize and open a store.
8//!
9//! # Memory
10//! - `GET`/`HEAD` stream the object store's [`GetResult`] straight into the axum
11//!   response body via [`axum::body::Body::from_stream`]; the whole object is
12//!   never materialized.
13//! - `PUT` with **no** `If-Match` streams the request body through a
14//!   [`WriteMultipart`], flushing fixed-size parts — bounded regardless of object
15//!   size.
16//! - `PUT` **with** `If-Match` is an inherently single-request conditional write
17//!   on the cloud side, so the body is buffered under a hard cap
18//!   ([`MAX_BUFFERED_PUT`]); this path is for the small `_delta_log` commit files
19//!   the wasm write path conditionally writes.
20
21use axum::body::Body;
22use axum::http::{HeaderMap, StatusCode, header};
23use axum::response::{IntoResponse, Response};
24use bytes::{Bytes, BytesMut};
25use futures::StreamExt;
26use object_store::{
27    GetOptions, GetRange, ObjectStoreExt, PutMode, PutOptions, PutPayload, UpdateVersion,
28    WriteMultipart,
29};
30
31use crate::backend::{ProxyReq, StorageProxyBackend};
32use crate::error::{ProxyError, ProxyResult};
33
34/// Hard cap on the body buffered for a **conditional** (`If-Match`) `PUT`. A
35/// conditional write is a single cloud request, so the body cannot be streamed as
36/// multipart; the cap bounds proxy memory. Sized generously for Delta commit files
37/// (the only conditional writes the wasm path issues) while rejecting a runaway
38/// upload with `413`.
39pub const MAX_BUFFERED_PUT: usize = 256 * 1024 * 1024;
40
41/// The multipart part size for streaming (unconditional) `PUT`s. Matches
42/// `object_store`'s default; the proxy never buffers more than one part.
43const MULTIPART_CHUNK_SIZE: usize = 5 * 1024 * 1024;
44
45/// The HTTP-facing surface produced by the blanket impl below. A host router
46/// dispatches to these three methods; each returns a fully-formed axum
47/// [`Response`] (streaming for `GET`), or a [`ProxyError`] the router turns into a
48/// response.
49#[async_trait::async_trait]
50pub trait StorageProxyHandler<Cx>: Send + Sync + 'static {
51    /// Serve a (optionally ranged) body read.
52    async fn get(&self, req: ProxyReq, cx: Cx) -> ProxyResult<Response>;
53    /// Serve a metadata-only read.
54    async fn head(&self, req: ProxyReq, cx: Cx) -> ProxyResult<Response>;
55    /// Serve a whole-object write, consuming the request `body`.
56    async fn put(&self, req: ProxyReq, body: Body, cx: Cx) -> ProxyResult<Response>;
57}
58
59#[async_trait::async_trait]
60impl<B, Cx> StorageProxyHandler<Cx> for B
61where
62    B: StorageProxyBackend<Cx>,
63    Cx: Send + 'static,
64{
65    async fn get(&self, req: ProxyReq, cx: Cx) -> ProxyResult<Response> {
66        self.authorize(&req, &cx).await?;
67        let store = self.open(&req, &cx).await?;
68
69        let opts = GetOptions {
70            range: req.range.clone(),
71            ..Default::default()
72        };
73        let path = object_store::path::Path::from(req.key.as_str());
74        let res = store.get_opts(&path, opts).await?;
75
76        // Read metadata + effective range *before* consuming the payload stream.
77        let total = res.meta.size;
78        let etag = res.meta.e_tag.clone();
79        let effective = res.range.clone();
80        let content_len = effective.end.saturating_sub(effective.start);
81        let ranged = req.range.is_some();
82
83        let mut headers = base_read_headers(content_len, etag.as_deref());
84        let status = if ranged {
85            headers.insert(
86                header::CONTENT_RANGE,
87                header::HeaderValue::from_str(&format!(
88                    "bytes {}-{}/{}",
89                    effective.start,
90                    effective.end.saturating_sub(1),
91                    total
92                ))
93                .expect("ascii content-range"),
94            );
95            StatusCode::PARTIAL_CONTENT
96        } else {
97            StatusCode::OK
98        };
99
100        // Stream the object bytes lazily — the whole object is never buffered.
101        let stream = res
102            .into_stream()
103            .map(|chunk| chunk.map_err(std::io::Error::other));
104        let body = Body::from_stream(stream);
105        Ok((status, headers, body).into_response())
106    }
107
108    async fn head(&self, req: ProxyReq, cx: Cx) -> ProxyResult<Response> {
109        self.authorize(&req, &cx).await?;
110        let store = self.open(&req, &cx).await?;
111
112        // `head: true` fetches metadata only.
113        let opts = GetOptions {
114            head: true,
115            ..Default::default()
116        };
117        let path = object_store::path::Path::from(req.key.as_str());
118        let res = store.get_opts(&path, opts).await?;
119
120        let headers = base_read_headers(res.meta.size, res.meta.e_tag.as_deref());
121        // No body on HEAD.
122        Ok((StatusCode::OK, headers, Body::empty()).into_response())
123    }
124
125    async fn put(&self, req: ProxyReq, body: Body, cx: Cx) -> ProxyResult<Response> {
126        self.authorize(&req, &cx).await?;
127        let store = self.open(&req, &cx).await?;
128        let path = object_store::path::Path::from(req.key.as_str());
129
130        let put_result = match &req.if_match {
131            // Conditional write: an inherently single-request cloud operation.
132            // `If-Match: *` has no `object_store` equivalent — reject explicitly.
133            Some(v) if v.trim() == "*" => {
134                return Err(ProxyError::InvalidArgument(
135                    "`If-Match: *` is not supported; send a concrete ETag".into(),
136                ));
137            }
138            Some(etag) => {
139                let payload = buffer_body(body, MAX_BUFFERED_PUT).await?;
140                let opts = PutOptions {
141                    mode: PutMode::Update(UpdateVersion {
142                        e_tag: Some(etag.clone()),
143                        version: None,
144                    }),
145                    ..Default::default()
146                };
147                // A mismatch surfaces as `object_store::Error::Precondition` →
148                // `ProxyError::Storage` → 412 (see error.rs).
149                store.put_opts(&path, payload, opts).await?
150            }
151            // Unconditional write: stream through multipart, bounded memory.
152            None => {
153                let upload = store.put_multipart(&path).await?;
154                let mut writer = WriteMultipart::new_with_chunk_size(upload, MULTIPART_CHUNK_SIZE);
155                let mut stream = body.into_data_stream();
156                while let Some(chunk) = stream.next().await {
157                    let chunk =
158                        chunk.map_err(|e| ProxyError::Internal(format!("body read: {e}")))?;
159                    writer.write(&chunk);
160                }
161                writer.finish().await.map_err(ProxyError::Storage)?
162            }
163        };
164
165        let mut headers = HeaderMap::new();
166        if let Some(etag) = put_result.e_tag
167            && let Ok(v) = header::HeaderValue::from_str(&etag)
168        {
169            headers.insert(header::ETAG, v);
170        }
171        // `object_store` does not distinguish created-vs-overwritten, so return 200
172        // uniformly; the wasm `HttpStore` only reads the ETag.
173        Ok((StatusCode::OK, headers).into_response())
174    }
175}
176
177/// The headers common to `GET` (full or ranged) and `HEAD`: `Content-Length`,
178/// `Accept-Ranges: bytes`, and `ETag` when present.
179fn base_read_headers(content_length: u64, etag: Option<&str>) -> HeaderMap {
180    let mut headers = HeaderMap::new();
181    headers.insert(
182        header::CONTENT_LENGTH,
183        header::HeaderValue::from_str(&content_length.to_string()).expect("ascii content-length"),
184    );
185    headers.insert(
186        header::ACCEPT_RANGES,
187        header::HeaderValue::from_static("bytes"),
188    );
189    if let Some(etag) = etag
190        && let Ok(v) = header::HeaderValue::from_str(etag)
191    {
192        headers.insert(header::ETAG, v);
193    }
194    headers
195}
196
197/// Drain a request body into a [`PutPayload`], failing with `413` if it exceeds
198/// `cap`. Used only for conditional (`If-Match`) writes, which cannot be streamed.
199async fn buffer_body(body: Body, cap: usize) -> ProxyResult<PutPayload> {
200    let mut stream = body.into_data_stream();
201    let mut buf = BytesMut::new();
202    while let Some(chunk) = stream.next().await {
203        let chunk: Bytes = chunk.map_err(|e| ProxyError::Internal(format!("body read: {e}")))?;
204        if buf.len() + chunk.len() > cap {
205            return Err(ProxyError::PayloadTooLarge(format!(
206                "conditional PUT body exceeds the {cap}-byte cap"
207            )));
208        }
209        buf.extend_from_slice(&chunk);
210    }
211    Ok(PutPayload::from(buf.freeze()))
212}
213
214/// Parse an HTTP `Range` header value into an [`object_store::GetRange`].
215///
216/// Supports the single-range forms `bytes=a-b` ([`GetRange::Bounded`]),
217/// `bytes=a-` ([`GetRange::Offset`]), and `bytes=-n` ([`GetRange::Suffix`]).
218/// Returns `Ok(None)` for an absent header. A malformed or multi-range value is
219/// rejected with [`ProxyError::InvalidArgument`].
220pub fn parse_range(headers: &HeaderMap) -> ProxyResult<Option<GetRange>> {
221    let Some(value) = headers.get(header::RANGE) else {
222        return Ok(None);
223    };
224    let value = value
225        .to_str()
226        .map_err(|_| ProxyError::InvalidArgument("non-ASCII Range header".into()))?;
227    let spec = value
228        .strip_prefix("bytes=")
229        .ok_or_else(|| ProxyError::InvalidArgument(format!("unsupported Range unit: {value}")))?;
230    if spec.contains(',') {
231        return Err(ProxyError::InvalidArgument(
232            "multi-range requests are not supported".into(),
233        ));
234    }
235    let (start, end) = spec
236        .split_once('-')
237        .ok_or_else(|| ProxyError::InvalidArgument(format!("malformed Range: {value}")))?;
238    let parse_u64 = |s: &str| {
239        s.trim()
240            .parse::<u64>()
241            .map_err(|_| ProxyError::InvalidArgument(format!("malformed Range number in: {value}")))
242    };
243    let range = match (start.trim(), end.trim()) {
244        ("", "") => {
245            return Err(ProxyError::InvalidArgument(format!("empty Range: {value}")));
246        }
247        // bytes=-n  -> last n bytes
248        ("", n) => GetRange::Suffix(parse_u64(n)?),
249        // bytes=a-  -> from offset a
250        (a, "") => GetRange::Offset(parse_u64(a)?),
251        // bytes=a-b -> [a, b] inclusive -> [a, b+1) half-open
252        (a, b) => {
253            let a = parse_u64(a)?;
254            let b = parse_u64(b)?;
255            if b < a {
256                return Err(ProxyError::InvalidArgument(format!(
257                    "Range end before start: {value}"
258                )));
259            }
260            GetRange::Bounded(a..b + 1)
261        }
262    };
263    Ok(Some(range))
264}
265
266/// Extract the `If-Match` header value (if any) for a conditional `PUT`.
267pub fn parse_if_match(headers: &HeaderMap) -> ProxyResult<Option<String>> {
268    match headers.get(header::IF_MATCH) {
269        None => Ok(None),
270        Some(v) => Ok(Some(
271            v.to_str()
272                .map_err(|_| ProxyError::InvalidArgument("non-ASCII If-Match header".into()))?
273                .to_string(),
274        )),
275    }
276}
277
278#[cfg(test)]
279mod tests {
280    use super::*;
281
282    fn hm(name: header::HeaderName, value: &str) -> HeaderMap {
283        let mut h = HeaderMap::new();
284        h.insert(name, header::HeaderValue::from_str(value).unwrap());
285        h
286    }
287
288    #[test]
289    fn range_bounded() {
290        let r = parse_range(&hm(header::RANGE, "bytes=2-5"))
291            .unwrap()
292            .unwrap();
293        assert_eq!(r, GetRange::Bounded(2..6));
294    }
295
296    #[test]
297    fn range_offset_and_suffix() {
298        assert_eq!(
299            parse_range(&hm(header::RANGE, "bytes=100-"))
300                .unwrap()
301                .unwrap(),
302            GetRange::Offset(100)
303        );
304        assert_eq!(
305            parse_range(&hm(header::RANGE, "bytes=-500"))
306                .unwrap()
307                .unwrap(),
308            GetRange::Suffix(500)
309        );
310    }
311
312    #[test]
313    fn range_absent_is_none() {
314        assert!(parse_range(&HeaderMap::new()).unwrap().is_none());
315    }
316
317    #[test]
318    fn range_rejects_bad() {
319        assert!(parse_range(&hm(header::RANGE, "items=0-1")).is_err());
320        assert!(parse_range(&hm(header::RANGE, "bytes=0-1,4-5")).is_err());
321        assert!(parse_range(&hm(header::RANGE, "bytes=5-2")).is_err());
322        assert!(parse_range(&hm(header::RANGE, "bytes=--")).is_err());
323    }
324
325    #[test]
326    fn if_match_roundtrip() {
327        assert_eq!(
328            parse_if_match(&hm(header::IF_MATCH, "\"etag123\"")).unwrap(),
329            Some("\"etag123\"".to_string())
330        );
331        assert!(parse_if_match(&HeaderMap::new()).unwrap().is_none());
332    }
333}
unitycatalog_storage_proxy/handler.rs

unitycatalog_storage_proxy/
handler.rs
