actix-http 3.21.0

use std::{cell::Cell, convert::Infallible, pin::pin};

use actix_service::fn_service;
use tokio::io::ReadBuf;

use super::*;
use crate::{
    body::BodyStream,
    config::ServiceConfigBuilder,
    h1::{ExpectHandler, UpgradeHandler},
    test::TestBuffer,
};

/// Stops accepting writes after a shared byte budget is exhausted.
///
/// The test can restore the budget and poll the dispatcher again, as if the socket became writable.
struct LimitedWriter {
    io: TestBuffer,
    allowance: Rc<Cell<usize>>,
}

impl AsyncRead for LimitedWriter {
    fn poll_read(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &mut ReadBuf<'_>,
    ) -> Poll<io::Result<()>> {
        Pin::new(&mut self.io).poll_read(cx, buf)
    }
}

impl AsyncWrite for LimitedWriter {
    fn poll_write(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
        buf: &[u8],
    ) -> Poll<io::Result<usize>> {
        let len = buf.len().min(self.allowance.get());

        if len == 0 {
            // The test polls manually after restoring the budget, so no wakeup is needed.
            return Poll::Pending;
        }

        let written = ready!(Pin::new(&mut self.io).poll_write(cx, &buf[..len]))?;
        self.allowance.set(self.allowance.get() - written);
        Poll::Ready(Ok(written))
    }

    // as a TCP stream does, so large sized bodies are written without the copy; the default
    // `poll_write_vectored` writes one buffer per call
    fn is_write_vectored(&self) -> bool {
        true
    }

    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        Pin::new(&mut self.io).poll_flush(cx)
    }

    fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {
        Pin::new(&mut self.io).poll_shutdown(cx)
    }
}

/// Checks buffer retention and complete responses on a persistent HTTP/1 connection.
///
/// - `body_len`: Response body length in bytes, sent as one chunk for each request.
/// - `threshold`: Configured HTTP/1 write-buffer watermark in bytes. The retention limit is
///   the larger of this value and `WRITE_BUFFER_RETENTION_LIMIT`.
/// - `partial`: If true, block writes after `body_len - 1024` bytes, including headers,
///   then resume them. Requires `body_len > 1024`.
/// - `release`: If true, expect the drained buffer to retain at most `HW_BUFFER_SIZE` bytes.
///   Otherwise, expect sufficient capacity for the body and reuse of the buffer on the next request.
async fn check_buffer_retention(body_len: usize, threshold: usize, partial: bool, release: bool) {
    let mut io = TestBuffer::new("GET / HTTP/1.1\r\nHost: localhost\r\n\r\n");
    // The budget includes response headers. Leave a small unwritten tail so
    // advancing past the written bytes hides most of the original buffer capacity.
    let allowance = Rc::new(Cell::new(if partial {
        body_len - 1024
    } else {
        usize::MAX
    }));
    let services = HttpFlow::new(
        fn_service(move |_: Request| async move {
            // A single chunk of a chunked body goes into the output buffer whole, which forces it
            // to hold the large response. A sized body this large skips the buffer.
            let chunk = bytes::Bytes::from(vec![b'x'; body_len]);
            let body = BodyStream::new(futures_util::stream::once(async move {
                Ok::<_, Infallible>(chunk)
            }));
            Ok::<_, Error>(Response::ok().set_body(body))
        }),
        ExpectHandler,
        None::<UpgradeHandler>,
    );
    let dispatcher = Dispatcher::new(
        LimitedWriter {
            io: io.clone(),
            allowance: allowance.clone(),
        },
        services,
        // This sets the body-polling watermark. The retention limit is the larger
        // of this value and WRITE_BUFFER_RETENTION_LIMIT.
        ServiceConfigBuilder::new()
            .h1_write_buffer_size(threshold)
            .build(),
        None,
        OnConnectData::default(),
    );
    let mut dispatcher = pin!(dispatcher);
    let mut cx = Context::from_waker(futures_util::task::noop_waker_ref());

    // Pending is expected even after a complete response: the keep-alive
    // connection stays open and waits for another request.
    assert!(dispatcher.as_mut().poll(&mut cx).is_pending());

    if partial {
        let DispatcherState::Normal { inner } = &dispatcher.inner else {
            panic!()
        };
        assert!(!inner.write_buf.is_empty());
        assert_eq!(io.write_buf_slice().len(), body_len - 1024);

        // Let the dispatcher finish the same response after the simulated blocked write.
        allowance.set(usize::MAX);
        assert!(dispatcher.as_mut().poll(&mut cx).is_pending());
    }

    // try_reclaim does not allocate. If the old backing allocation survived,
    // reclaiming its consumed prefix exposes the large capacity to the assertion below.
    // Checking only the visible tail could otherwise let the partial-write regression pass.
    if release {
        let DispatcherStateProj::Normal { inner } = dispatcher.as_mut().project().inner.project()
        else {
            panic!()
        };
        assert!(inner.project().write_buf.try_reclaim(HW_BUFFER_SIZE));
    }

    let DispatcherState::Normal { inner } = &dispatcher.inner else {
        panic!()
    };
    assert!(inner.write_buf.is_empty());

    if release {
        assert!(
            inner.write_buf.capacity() <= HW_BUFFER_SIZE,
            "drained buffer retained {} bytes",
            inner.write_buf.capacity()
        );
    } else {
        assert!(inner.write_buf.capacity() >= body_len);
    }

    let pointer = inner.write_buf.as_ptr();
    let response = io.take_write_buf();
    // Compare the complete body separately from the generated HTTP headers.
    // A smaller retained buffer must not come at the cost of missing response bytes.
    let body_start = response
        .windows(4)
        .position(|window| window == b"\r\n\r\n")
        .unwrap()
        + 4;
    let mut body = format!("{body_len:X}\r\n").into_bytes();
    body.extend(vec![b'x'; body_len]);
    body.extend_from_slice(b"\r\n0\r\n\r\n");
    assert_eq!(&response[body_start..], body);

    // The connection must remain usable, including after a partial write.
    io.extend_read_buf("GET / HTTP/1.1\r\nHost: localhost\r\n\r\n");
    assert!(dispatcher.as_mut().poll(&mut cx).is_pending());
    let second_response = io.take_write_buf();
    let second_body_start = second_response
        .windows(4)
        .position(|window| window == b"\r\n\r\n")
        .unwrap()
        + 4;
    assert_eq!(
        &second_response[second_body_start..],
        &response[body_start..]
    );

    if !release {
        // Check reuse across requests, not just sufficient capacity after each flush.
        let DispatcherState::Normal { inner } = &dispatcher.inner else {
            panic!()
        };
        assert_eq!(
            inner.write_buf.as_ptr(),
            pointer,
            "normal buffers should be reused"
        );
    }
}

#[actix_rt::test]
async fn oversized_write_buffer_released_after_flush() {
    check_buffer_retention(1024 * 1024, 32 * 1024, false, true).await;
}

#[actix_rt::test]
async fn oversized_write_buffer_released_after_partial_write() {
    check_buffer_retention(1024 * 1024, 32 * 1024, true, true).await;
}

#[actix_rt::test]
async fn normal_write_buffer_reused_after_flush() {
    check_buffer_retention(1024, 32 * 1024, false, false).await;
}

#[actix_rt::test]
async fn configured_write_buffer_reused_after_flush() {
    // A custom watermark above 128 KiB also raises the retention limit.
    check_buffer_retention(256 * 1024, 1024 * 1024, false, false).await;
}

#[actix_rt::test]
async fn small_write_watermark_preserves_normal_buffer() {
    // A small watermark must not lower the 128 KiB retention limit.
    check_buffer_retention(64 * 1024, 8 * 1024, false, false).await;
}

#[actix_rt::test]
async fn buffer_above_configured_write_watermark_released() {
    // A larger watermark sets the limit directly; it is not multiplied by four.
    check_buffer_retention(512 * 1024, 256 * 1024, false, true).await;
}

#[actix_rt::test]
async fn blocked_sized_body_retains_write_buffer() {
    const BODY_LEN: usize = 64 * 1024;

    for spare in [false, true] {
        let io = TestBuffer::new("GET / HTTP/1.1\r\nHost: localhost\r\n\r\n");
        // Accept the head and part of the body, then block.
        let allowance = Rc::new(Cell::new(1024));
        let services = HttpFlow::new(
            fn_service(move |_: Request| {
                // With `spare`, the body has room to spare, as the `String` behind `.json()` does.
                let mut body = Vec::with_capacity(if spare { 2 * BODY_LEN } else { BODY_LEN });
                body.resize(BODY_LEN, b'x');
                let res = Response::ok().set_body(bytes::Bytes::from(body));
                async move { Ok::<_, Error>(res) }
            }),
            ExpectHandler,
            None::<UpgradeHandler>,
        );
        let dispatcher = Dispatcher::new(
            LimitedWriter {
                io: io.clone(),
                allowance: allowance.clone(),
            },
            services,
            ServiceConfig::default(),
            None,
            OnConnectData::default(),
        );
        let mut dispatcher = pin!(dispatcher);
        let mut cx = Context::from_waker(futures_util::task::noop_waker_ref());

        assert!(dispatcher.as_mut().poll(&mut cx).is_pending());

        let DispatcherState::Normal { inner } = &dispatcher.inner else {
            panic!()
        };
        if spare {
            // holding it would keep the spare room, so it was copied into the buffer instead
            assert!(inner.write_body.is_none(), "body with spare capacity held");
        } else {
            // The head is out, but the normal write buffer remains available for reuse.
            assert!(inner.write_body.is_some());
            assert!(
                inner.write_buf.capacity() > LW_BUFFER_SIZE
                    && inner.write_buf.capacity() <= HW_BUFFER_SIZE,
                "write buffer should retain its allocation while the body waits"
            );
            assert!(!inner.write_buf_needs_reset);
        }

        let retained_capacity = inner.write_buf.capacity();

        allowance.set(usize::MAX);
        assert!(dispatcher.as_mut().poll(&mut cx).is_pending());

        let response = io.take_write_buf();
        let body_start = response
            .windows(4)
            .position(|window| window == b"\r\n\r\n")
            .unwrap()
            + 4;
        assert_eq!(
            &response[body_start..],
            vec![b'x'; BODY_LEN],
            "spare: {spare}"
        );

        // The write buffer remains available for the next response.
        let DispatcherState::Normal { inner } = &dispatcher.inner else {
            panic!()
        };

        if spare {
            assert!(inner.write_buf.capacity() >= HW_BUFFER_SIZE);
        } else {
            assert_eq!(inner.write_buf.capacity(), retained_capacity);
        }
    }
}