hive-router 0.2.0

use std::cell::RefMut;

use http::header::CONTENT_LENGTH;
use ntex::{
    http::error::PayloadError,
    util::Extensions,
    web::{self, DefaultError, HttpRequest, WebRequest},
};
use ntex_bytes::{Bytes, BytesMut};
use ntex_http::HeaderMap;
use strum::IntoStaticStr;
use tokio_stream::StreamExt;

/// Stores the request body size in bytes.
///
/// The value comes from either:
/// - the `Content-Length` header
/// - the streamed payload, measured from bytes read.
///
/// For streamed payloads, the recorded size is the number of bytes read up to
/// the configured maximum.
///
/// Using `RequestBodySize` to store the size of the request body,
/// helps to reduce complexity in code, as otherwise,
/// we would have to return the size next within Err and Ok of `read_body_stream`.
#[derive(Debug, Clone, Copy)]
pub struct RequestBodySize(pub u64);

#[derive(Debug, thiserror::Error, IntoStaticStr)]
pub enum ReadBodyStreamError {
    #[error("Failed to read request body: {0}")]
    #[strum(serialize = "PAYLOAD_READ_ERROR")]
    // Thrown while reading the body stream with `try_next()`
    PayloadReadError(#[from] PayloadError),

    #[error("Content-Length header has invalid value")]
    #[strum(serialize = "INVALID_HEADER")]
    InvalidContentLengthHeader,

    #[error("Content-Length exceeds the maximum allowed size: {0}")]
    #[strum(serialize = "PAYLOAD_TOO_LARGE_CONTENT_LENGTH")]
    PayloadTooLargeContentLength(usize),

    #[error("Request body exceeds the maximum allowed size while reading the stream")]
    #[strum(serialize = "PAYLOAD_TOO_LARGE_BODY_STREAM")]
    PayloadTooLargeBodyStream,
}

impl ReadBodyStreamError {
    pub fn status_code(&self) -> http::StatusCode {
        match self {
            Self::PayloadReadError(_) => http::StatusCode::UNPROCESSABLE_ENTITY,
            Self::InvalidContentLengthHeader => http::StatusCode::BAD_REQUEST,
            Self::PayloadTooLargeContentLength(_) | Self::PayloadTooLargeBodyStream => {
                http::StatusCode::PAYLOAD_TOO_LARGE
            }
        }
    }

    pub fn error_code(&self) -> &'static str {
        self.into()
    }
}

#[inline]
fn write_request_body_size<R: RequestLike>(req: &R, size: u64) {
    req.extensions_mut().insert(RequestBodySize(size));
}

#[inline]
pub fn read_request_body_size(req: &HttpRequest) -> Option<u64> {
    req.extensions().get::<RequestBodySize>().map(|size| size.0)
}

/// Limit for draining a rejected request body. Keeps the connection
/// reusable for typical over-limit requests without reading the whole thing
const MAX_DRAIN_BYTES: usize = 64 * 1024;

pub async fn drain_body_stream(body_stream: &mut web::types::Payload) {
    let mut drained: usize = 0;

    while drained < MAX_DRAIN_BYTES {
        match body_stream.try_next().await {
            Ok(Some(chunk)) => {
                if chunk.is_empty() {
                    break;
                }

                drained = drained.saturating_add(chunk.len());
            }
            _ => break,
        }
    }
}

#[inline]
pub async fn read_body_stream<R: RequestLike>(
    req: &R,
    mut body_stream: web::types::Payload,
    max_size: usize,
) -> Result<Bytes, ReadBodyStreamError> {
    let content_length: Option<usize> = {
        let content_length_header = req.headers().get(CONTENT_LENGTH);
        if let Some(content_length_header) = content_length_header {
            let content_length_str = content_length_header
                .to_str()
                .map_err(|_| ReadBodyStreamError::InvalidContentLengthHeader)?;
            let content_length: usize = content_length_str
                .parse()
                .map_err(|_| ReadBodyStreamError::InvalidContentLengthHeader)?;
            if content_length > max_size {
                write_request_body_size(req, content_length as u64);

                // Drain just small amount of the request body before rejecting,
                // so the server can close the connection cleanly.
                //
                // Returning without consuming the body makes ntex reset the socket,
                // and the client might read that as a socket error instead of the router's 413
                // response.
                //
                // Note: `drain_body_stream` reads only a small amount of the body,
                // so the client will not be blocked while draining.
                drain_body_stream(&mut body_stream).await;
                return Err(ReadBodyStreamError::PayloadTooLargeContentLength(max_size));
            }
            Some(content_length)
        } else {
            None
        }
    };

    let mut body = if let Some(content_length) = content_length {
        BytesMut::with_capacity(content_length)
    } else {
        BytesMut::new()
    };

    while let Some(chunk) = body_stream.try_next().await? {
        // limit max size of in-memory payload
        if chunk.len() > max_size.saturating_sub(body.len()) {
            write_request_body_size(req, (body.len() + chunk.len()) as u64);
            return Err(ReadBodyStreamError::PayloadTooLargeBodyStream);
        }
        body.extend_from_slice(&chunk);
    }

    write_request_body_size(req, body.len() as u64);

    Ok(body.freeze())
}

pub trait RequestLike {
    fn headers(&self) -> &HeaderMap;
    fn extensions_mut(&self) -> RefMut<'_, Extensions>;
}

impl RequestLike for HttpRequest {
    fn headers(&self) -> &HeaderMap {
        self.headers()
    }

    fn extensions_mut(&self) -> RefMut<'_, Extensions> {
        self.extensions_mut()
    }
}

impl RequestLike for WebRequest<DefaultError> {
    fn headers(&self) -> &HeaderMap {
        self.headers()
    }

    fn extensions_mut(&self) -> RefMut<'_, Extensions> {
        self.extensions_mut()
    }
}