Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
57 changes: 41 additions & 16 deletions arvik-core/src/accept.rs
Original file line number Diff line number Diff line change
Expand Up @@ -37,26 +37,51 @@ pub fn parse_accept_encodings(value: &str) -> Vec<(String, f32)> {
///
/// An exact `;q=0` entry refuses that coding even when `*` is present; codings
/// not mentioned at all fall back to the `*` wildcard's weight.
///
/// Allocation-free: compares case-insensitively over the raw header slice.
pub fn negotiate<'a>(supported: &[&'a str], accept_encoding: &str) -> Option<&'a str> {
let entries = parse_accept_encodings(accept_encoding);
let wildcard_q = entries
.iter()
.rev()
.find(|(coding, _)| coding == "*")
.map(|(_, q)| *q)
.unwrap_or(0.0);
for candidate in supported {
let mut exact_seen = false;
let mut exact_allowed = false;
let mut wildcard_allowed = false;

supported.iter().copied().find(|candidate| {
let candidate = candidate.to_ascii_lowercase();
match entries
.iter()
.rev()
.find(|(coding, _)| *coding == candidate)
for entry in accept_encoding.split(',') {
let mut segments = entry.trim().split(';');
let coding = segments.next().map(str::trim).unwrap_or("");
if coding.eq_ignore_ascii_case(candidate) {
exact_seen = true;
exact_allowed = parse_q(segments) > 0.0;
} else if coding == "*" {
wildcard_allowed = parse_q(segments) > 0.0;
}
}

let allowed = if exact_seen {
exact_allowed
} else {
wildcard_allowed
};
if allowed {
return Some(candidate);
}
}
None
}

/// Parse the `;q=` parameter from the remaining attribute segments
/// (defaulting to 1.0).
fn parse_q<'a, I: Iterator<Item = &'a str>>(segments: I) -> f32 {
let mut quality = 1.0_f32;
for parameter in segments {
let parameter = parameter.trim();
if let Some(q) = parameter
.strip_prefix("q=")
.or_else(|| parameter.strip_prefix("Q="))
{
Some((_, q)) => *q > 0.0,
None => wildcard_q > 0.0,
quality = q.trim().parse().unwrap_or(1.0);
}
})
}
quality.clamp(0.0, 1.0)
}

#[cfg(test)]
Expand Down
23 changes: 19 additions & 4 deletions arvik-core/src/body.rs
Original file line number Diff line number Diff line change
Expand Up @@ -140,7 +140,10 @@ impl Body {
}
}

let mut buf = bytes::BytesMut::new();
// Collect frame handles first (no copying); typical requests arrive
// as a single hyper frame, which is returned zero-copy.
let mut frames: Vec<Bytes> = Vec::new();
let mut total = 0usize;
let mut body = self;
while let Some(frame) = BodyExt::frame(&mut body).await {
let data = frame?.into_data().map_err(|_| {
Expand All @@ -149,12 +152,24 @@ impl Body {
"unexpected body trailer",
)) as BoxError
})?;
if buf.len() + data.len() > limit {
total += data.len();
if total > limit {
return Err(BodyLimitError::TooLarge);
}
buf.extend_from_slice(&data);
frames.push(data);
}

match frames.len() {
0 => Ok(Bytes::new()),
1 => Ok(frames.pop().unwrap_or_default()),
_ => {
let mut buf = bytes::BytesMut::with_capacity(total);
for frame in frames {
buf.extend_from_slice(&frame);
}
Ok(buf.freeze())
}
}
Ok(buf.freeze())
}

/// Collect the entire body into a UTF-8 [`String`].
Expand Down
11 changes: 8 additions & 3 deletions arvik-core/src/handler.rs
Original file line number Diff line number Diff line change
Expand Up @@ -180,7 +180,11 @@ pub trait ErasedHandler<S>: Send + Sync {
fn clone_box(&self) -> Box<dyn ErasedHandler<S>>;

/// Call this handler, returning a boxed future.
fn call(self: Box<Self>, req: Request, state: S) -> BoxFuture<'static, Response>;
///
/// Takes `&self` — the concrete handler is cloned into the returned
/// future (it must be `Clone`), so dispatch never allocates an extra
/// boxed handler per request.
fn call(&self, req: Request, state: S) -> BoxFuture<'static, Response>;
}

impl<H, T, S> ErasedHandler<S> for ErasedHandlerWrapper<H, T, S>
Expand All @@ -193,8 +197,9 @@ where
Box::new(self.clone())
}

fn call(self: Box<Self>, req: Request, state: S) -> BoxFuture<'static, Response> {
let fut = self.handler.call(req, state);
fn call(&self, req: Request, state: S) -> BoxFuture<'static, Response> {
let handler = self.handler.clone();
let fut = handler.call(req, state);
Box::pin(fut)
}
}
Expand Down
12 changes: 12 additions & 0 deletions arvik-core/src/response.rs
Original file line number Diff line number Diff line change
Expand Up @@ -119,6 +119,18 @@ impl ResponseBuilder {
.body(Body::from(html.into()))
}

/// Build a plain text response from static data.
///
/// Zero-copy: both the body bytes and the Content-Type header come from
/// `'static` data, avoiding a String allocation per response.
pub fn text_static(self, text: &'static str) -> Response {
self.header(
http::header::CONTENT_TYPE,
http::header::HeaderValue::from_static("text/plain; charset=utf-8"),
)
.body(Body::from_static(text.as_bytes()))
}

/// Build a plain text response.
///
/// Sets `Content-Type: text/plain; charset=utf-8`.
Expand Down
8 changes: 6 additions & 2 deletions arvik-extract/src/form.rs
Original file line number Diff line number Diff line change
Expand Up @@ -72,8 +72,12 @@ fn form_content_type(headers: &http::HeaderMap) -> bool {
Err(_) => return false,
};

// Parse as a media type so parameter suffixes (`;charset=utf-8`) pass but
// bogus extensions (`…urlencodedfoo`) do not.
// Fast path for the exact constant.
if content_type.eq_ignore_ascii_case("application/x-www-form-urlencoded") {
return true;
}

// Parameterized variants need a real parse.
match content_type.parse::<mime::Mime>() {
Ok(mime) => mime.type_() == mime::APPLICATION && mime.subtype() == "x-www-form-urlencoded",
Err(_) => false,
Expand Down
22 changes: 14 additions & 8 deletions arvik-extract/src/json.rs
Original file line number Diff line number Diff line change
Expand Up @@ -101,15 +101,21 @@ fn json_content_type(headers: &http::HeaderMap) -> bool {
Err(_) => return false,
};

let mime: mime::Mime = match content_type.parse() {
Ok(m) => m,
Err(_) => return false,
};
// Fast path for the exact constant — avoids a full Mime parse per request.
if content_type.eq_ignore_ascii_case("application/json") {
return true;
}

// Accept application/json and application/*+json (e.g., application/vnd.api+json)
mime.type_() == mime::APPLICATION
&& (mime.subtype() == mime::JSON
|| mime.suffix().is_some_and(|suffix| suffix == mime::JSON))
// Parameterized / suffixed variants (e.g. application/vnd.api+json,
// application/json;charset=utf-8) need a real parse.
match content_type.parse::<mime::Mime>() {
Ok(mime) => {
mime.type_() == mime::APPLICATION
&& (mime.subtype() == mime::JSON
|| mime.suffix().is_some_and(|suffix| suffix == mime::JSON))
}
Err(_) => false,
}
}

#[cfg(test)]
Expand Down
97 changes: 59 additions & 38 deletions arvik-middleware/src/body_limit.rs
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,6 @@
//! ```

use std::convert::Infallible;
use std::future::Future;
use std::pin::Pin;
use std::task::{Context, Poll};

Expand Down Expand Up @@ -66,11 +65,11 @@ pub struct RequestBodyLimitService<S> {
impl<S> Service<Request> for RequestBodyLimitService<S>
where
S: Service<Request, Response = Response, Error = Infallible> + Clone + Send + 'static,
S::Future: Send + 'static,
S::Future: Send + Unpin + 'static,
{
type Response = Response;
type Error = Infallible;
type Future = Pin<Box<dyn Future<Output = Result<Response, Infallible>> + Send + 'static>>;
type Future = RequestBodyLimitFuture<S>;

fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
self.inner.poll_ready(cx)
Expand All @@ -79,42 +78,67 @@ where
fn call(&mut self, req: Request) -> Self::Future {
let limit = self.limit;

// Fast-path: reject based on Content-Length immediately. Only this
// rare rejection path pays an allocation — the accepted path returns
// the inner future directly with no extra boxing.
if let Some(content_length) = req
.headers()
.get(http::header::CONTENT_LENGTH)
.and_then(|v| v.to_str().ok())
.and_then(|s| s.parse::<usize>().ok())
{
if content_length > limit {
tracing::warn!(
content_length = content_length,
limit = limit,
"Request body exceeds size limit (Content-Length check)"
);
return RequestBodyLimitFuture::Reject(std::future::ready(Ok(payload_too_large(
limit,
))));
}
}

// For streaming bodies: pass the body through unbuffered with a
// counting cap — readers get an error once the limit is crossed.
let (parts, body) = req.into_request_parts();
let limited = Body::new(LimitedBody {
inner: body,
limit: limit as u64,
read: 0,
});
let req = Request::from_request_parts(parts, limited);

let cloned = self.inner.clone();
let mut inner = std::mem::replace(&mut self.inner, cloned);
RequestBodyLimitFuture::Accept(inner.call(req))
}
}

Box::pin(async move {
// Fast-path: reject based on Content-Length header immediately.
if let Some(content_length) = req
.headers()
.get(http::header::CONTENT_LENGTH)
.and_then(|v| v.to_str().ok())
.and_then(|s| s.parse::<usize>().ok())
{
if content_length > limit {
tracing::warn!(
content_length = content_length,
limit = limit,
"Request body exceeds size limit (Content-Length check)"
);
return Ok::<Response, Infallible>(payload_too_large(limit));
}
}
/// Happy path delegates to the inner service's own (unboxed) future; only
/// the rare over-limit rejection uses a tiny ready future.
pub enum RequestBodyLimitFuture<S>
where
S: Service<Request, Response = Response, Error = Infallible>,
{
Accept(S::Future),
Reject(std::future::Ready<Result<Response, Infallible>>),
}

impl<S> std::future::Future for crate::body_limit::RequestBodyLimitFuture<S>
where
S: Service<Request, Response = Response, Error = Infallible>,
S::Future: Send + Unpin + 'static,
{
type Output = Result<Response, Infallible>;

// For streaming bodies: pass the body through unbuffered with a
// counting cap — readers get an error once the limit is crossed.
let (parts, body) = req.into_request_parts();
let limited = Body::new(LimitedBody {
inner: body,
limit: limit as u64,
read: 0,
});
let req = Request::from_request_parts(parts, limited);
inner.call(req).await
})
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
match &mut *self {
Self::Accept(fut) => std::pin::Pin::new(fut).poll(cx),
Self::Reject(fut) => std::pin::Pin::new(fut).poll(cx),
}
}
}

/// Streaming frame filter that errors once more than `limit` bytes have been
/// read. Frames themselves are passed through untouched — nothing is buffered.
struct LimitedBody {
inner: Body,
Expand Down Expand Up @@ -235,18 +259,15 @@ where
{
type Response = Response;
type Error = Infallible;
type Future = Pin<Box<dyn Future<Output = Result<Response, Infallible>> + Send + 'static>>;
type Future = S::Future;

fn poll_ready(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
self.inner.poll_ready(cx)
}

fn call(&mut self, mut req: Request) -> Self::Future {
let cloned = self.inner.clone();
let mut inner = std::mem::replace(&mut self.inner, cloned);

req.extensions_mut().insert(DefaultBodyLimit(self.limit));
Box::pin(inner.call(req))
self.inner.call(req)
}
}

Expand Down
Loading
Loading