trillium-http 1.3.3

use crate::{
    BufWriter, Buffer, Conn, ConnectionStatus, Error, Headers, HttpContext, KnownHeaderName,
    Method, ProtocolSession, ReceivedBody, Result, Status, TypeSet, Version, after_send::AfterSend,
    conn::ReceivedBodyState, util::encoding,
};
use futures_lite::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
use memchr::memmem::Finder;
use std::{
    borrow::Cow,
    io::Write,
    sync::Arc,
    time::{Instant, SystemTime},
};

impl<Transport> Conn<Transport>
where
    Transport: AsyncRead + AsyncWrite + Unpin + Send + Sync + 'static,
{
    /// HTTP/1.x response-header finalization. Parallel to
    /// `finalize_response_headers_h2` and `finalize_response_headers_h3` — keep the
    /// three in sync when changing universal policy (e.g. how Date is set).
    /// Differences are version-intrinsic: chunked Transfer-Encoding and
    /// `Connection: close` are h1-only; the h2/h3 paths strip `H1_ONLY_HEADERS`
    /// and let the framing layer signal end-of-stream.
    pub(super) fn finalize_response_headers_1x(&mut self) {
        if self.status == Some(Status::SwitchingProtocols) {
            return;
        }

        self.response_headers
            .try_insert_with(KnownHeaderName::Date, || {
                httpdate::fmt_http_date(SystemTime::now())
            });

        if !matches!(self.status, Some(Status::NotModified | Status::NoContent)) {
            // Upgrade path: don't default to `Content-Length: 0` — body bytes will be
            // written post-handoff. A prelude response body is sent as an open chunked
            // stream (see `Body::keep_open`), so force chunked and drop any Content-Length.
            // With no prelude body, honor an explicit Content-Length; otherwise chunked.
            let has_content_length = if self.upgrade {
                if self.response_body.as_ref().is_some_and(|b| !b.is_empty()) {
                    self.response_headers.remove(KnownHeaderName::ContentLength);
                    false
                } else {
                    self.response_headers
                        .has_header(KnownHeaderName::ContentLength)
                }
            } else if let Some(len) = self.body_len() {
                self.response_headers
                    .try_insert(KnownHeaderName::ContentLength, len);
                true
            } else {
                self.response_headers
                    .has_header(KnownHeaderName::ContentLength)
            };

            if self.version == Version::Http1_1 && !has_content_length {
                self.response_headers
                    .insert(KnownHeaderName::TransferEncoding, "chunked");
            } else {
                self.response_headers
                    .remove(KnownHeaderName::TransferEncoding);
            }
        }

        if self.context.swansong.state().is_shutting_down() {
            self.response_headers
                .insert(KnownHeaderName::Connection, "close");
        }
    }

    pub(crate) async fn send(mut self) -> Result<ConnectionStatus<Transport>> {
        let mut output_buffer = Vec::with_capacity(self.context.config.response_buffer_len);
        self.write_headers(&mut output_buffer)?;

        // Read before the bufwriter borrows `self.transport`, so the body block can consult it.
        let upgrading = self.should_upgrade();

        let max_buf = self.context.config.response_buffer_max_len;
        let mut bufwriter = BufWriter::new_with_buffer(output_buffer, &mut self.transport, max_buf);

        if self.method != Method::Head
            && !matches!(self.status, Some(Status::NotModified | Status::NoContent))
            && let Some(mut body) = self.response_body.take()
        {
            let chunked = body.len().is_none();

            body.ensure_chunked_framing();
            if upgrading {
                // Leave the chunked stream unterminated for the following upgrade to close.
                body.set_keep_open();
            }

            let loops_per_yield = self.context.config.copy_loops_per_yield;

            bufwriter.copy_from(&mut body, loops_per_yield).await?;

            // When an upgrade follows, the upgrade owns the terminator; the body's trailers
            // (if any) ride onto the `Upgrade` and merge with whatever the upgrade handler
            // emits. Skip the terminator stitch here.
            if !upgrading {
                // Chunked-trailer-section stitch. `Body::poll_read` emitted the last-chunk
                // marker `0\r\n` at EOF and stopped there; we own the rest of the framing
                // because trailers are structured `Headers` (not bytes) and the terminating
                // CRLF closes the trailer-section. See `Body::poll_read`'s `len: None`
                // branch for the full rationale.
                if let Some(trailers) = body.trailers() {
                    log::trace!("sending trailers:\n{trailers}");
                    write_headers_or_trailers(bufwriter.buffer_mut(), &trailers, &self.context)?;
                    // we don't store the trailers anywhere because the conn is about to be dropped
                }

                if chunked {
                    write!(bufwriter.buffer_mut(), "\r\n")?;
                }
            }
        }

        bufwriter.flush().await?;

        self.after_send.call(true.into());
        self.finish().await
    }

    pub(super) fn needs_100_continue(&self) -> bool {
        self.request_body_state.is_unread()
            && self.version == Version::Http1_1
            && self
                .request_headers
                .eq_ignore_ascii_case(KnownHeaderName::Expect, "100-continue")
    }

    #[allow(clippy::needless_borrow, clippy::needless_borrows_for_generic_args)]
    pub(super) fn build_request_body(&mut self) -> ReceivedBody<'_, Transport> {
        ReceivedBody::new_with_config(
            self.request_content_length().ok().flatten(),
            &mut self.buffer,
            &mut self.transport,
            &mut self.request_body_state,
            None,
            encoding(&self.request_headers),
            &self.context.config,
        )
        .with_trailers(&mut self.request_trailers)
        .with_protocol_session(self.protocol_session.clone())
    }

    /// Resolve the initial [`ReceivedBodyState`] for the incoming h1 request body from
    /// the parsed headers. h1 requests without explicit framing default to an empty
    /// body — read-to-close on inbound has no sender-side end-of-request signal.
    fn initial_request_body_state(request_headers: &Headers) -> ReceivedBodyState {
        let chunked = request_headers.has_header(KnownHeaderName::TransferEncoding);
        let content_length = if chunked {
            None
        } else {
            request_headers
                .get_str(KnownHeaderName::ContentLength)
                .and_then(|s| s.parse().ok())
                .or(Some(0))
        };
        ReceivedBodyState::new_h1(content_length, chunked)
    }

    fn validate_headers(request_headers: &Headers) -> Result<()> {
        if request_headers
            .get_values(KnownHeaderName::ContentLength)
            .is_some_and(|v| v.len() > 1)
        {
            return Err(Error::UnexpectedHeader(
                KnownHeaderName::ContentLength.into(),
            ));
        }

        if let Some(te) = request_headers.get_values(KnownHeaderName::TransferEncoding)
            && te
                .as_str()
                .is_none_or(|te_str| !te_str.eq_ignore_ascii_case("chunked"))
        {
            return Err(Error::UnexpectedHeader(
                KnownHeaderName::TransferEncoding.into(),
            ));
        }

        if request_headers.has_header(KnownHeaderName::ContentLength)
            && request_headers.has_header(KnownHeaderName::TransferEncoding)
        {
            return Err(Error::UnexpectedHeader(
                KnownHeaderName::ContentLength.into(),
            ));
        }

        Ok(())
    }

    #[cfg(not(feature = "parse"))]
    pub(crate) async fn new_internal(
        context: Arc<HttpContext>,
        mut transport: Transport,
        mut buffer: Buffer,
    ) -> Result<Self> {
        use crate::{HeaderName, HeaderValue};
        use httparse::{EMPTY_HEADER, Request};

        let (head_size, start_time) = Self::head(&mut transport, &mut buffer, &context).await?;

        let mut headers = vec![EMPTY_HEADER; context.config.max_headers];
        let mut httparse_req = Request::new(&mut headers);

        let status = httparse_req.parse(&buffer[..]).map_err(|e| match e {
            httparse::Error::HeaderName => Error::InvalidHeaderName,
            httparse::Error::HeaderValue => Error::InvalidHeaderValue("unknown".into()),
            httparse::Error::Status => Error::InvalidStatus,
            httparse::Error::TooManyHeaders => Error::HeadersTooLong,
            httparse::Error::Version => Error::InvalidVersion,
            _ => Error::InvalidHead,
        })?;

        if status.is_partial() {
            return Err(Error::InvalidHead);
        }

        let method = match httparse_req.method {
            Some(method) => match method.parse() {
                Ok(method) => method,
                Err(_) => return Err(Error::UnrecognizedMethod(method.to_string())),
            },
            None => return Err(Error::MissingMethod),
        };

        let version = match httparse_req.version {
            Some(0) => Version::Http1_0,
            Some(1) => Version::Http1_1,
            _ => return Err(Error::InvalidVersion),
        };

        let mut request_headers = Headers::new();
        for header in httparse_req.headers {
            use std::str::FromStr;

            let header_name = HeaderName::from_str(header.name)?;
            let header_value = HeaderValue::from(header.value.to_owned());
            request_headers.append(header_name, header_value);
        }

        Self::validate_headers(&request_headers)?;

        let mut path = Cow::Owned(
            httparse_req
                .path
                .ok_or(Error::RequestPathMissing)?
                .to_owned(),
        );

        let mut authority = None;

        if method == Method::Connect {
            authority = Some(path);
            path = Cow::Borrowed("/");
        }

        log::trace!("received:\n{method} {path} {version}\n{request_headers}");

        let response_headers = context
            .shared_state()
            .get::<Headers>()
            .cloned()
            .unwrap_or_default();

        buffer.ignore_front(head_size);

        let request_body_state = Self::initial_request_body_state(&request_headers);

        Ok(Self {
            transport,
            request_headers,
            method,
            version,
            path,
            buffer,
            response_headers,
            status: None,
            state: TypeSet::new(),
            response_body: None,
            request_body_state,
            secure: false,
            after_send: AfterSend::default(),
            start_time,
            peer_ip: None,
            context,
            authority,
            scheme: None,
            protocol: None,
            protocol_session: ProtocolSession::Http1,
            request_trailers: None,
            upgrade: false,
        })
    }

    #[cfg(feature = "parse")]
    pub(crate) async fn new_internal(
        context: Arc<HttpContext>,
        mut transport: Transport,
        mut buffer: Buffer,
    ) -> Result<Self> {
        let (head_size, start_time) = Self::head(&mut transport, &mut buffer, &context).await?;

        let first_line_index = Finder::new(b"\r\n")
            .find(&buffer)
            .ok_or(Error::InvalidHead)?;

        let mut spaces = memchr::memchr_iter(b' ', &buffer[..first_line_index]);
        let first_space = spaces.next().ok_or(Error::MissingMethod)?;
        let method = Method::parse(&buffer[0..first_space])?;
        let second_space = spaces.next().ok_or(Error::RequestPathMissing)?;
        let mut path: Cow<'static, str> = Cow::Owned(
            std::str::from_utf8(&buffer[first_space + 1..second_space])
                .map_err(|_| Error::RequestPathMissing)?
                .to_string(),
        );

        if path.is_empty() {
            return Err(Error::InvalidHead);
        }

        let version = Version::parse(&buffer[second_space + 1..first_line_index])?;
        if !matches!(version, Version::Http1_1 | Version::Http1_0) {
            return Err(Error::UnsupportedVersion(version));
        }

        let request_headers = Headers::parse(&buffer[first_line_index + 2..head_size])?;

        Self::validate_headers(&request_headers)?;

        let mut authority = None;

        if method == Method::Connect {
            authority = Some(path);
            path = Cow::Borrowed("/");
        }

        let response_headers = context
            .shared_state()
            .get::<Headers>()
            .cloned()
            .unwrap_or_default();

        buffer.ignore_front(head_size);

        let request_body_state = Self::initial_request_body_state(&request_headers);

        Ok(Self {
            context,
            transport,
            request_headers,
            method,
            version,
            path,
            buffer,
            response_headers,
            status: None,
            state: TypeSet::new(),
            response_body: None,
            request_body_state,
            secure: false,
            after_send: AfterSend::default(),
            start_time,
            peer_ip: None,
            authority,
            scheme: None,
            protocol: None,
            protocol_session: ProtocolSession::Http1,
            request_trailers: None,
            upgrade: false,
        })
    }

    async fn head(
        transport: &mut Transport,
        buf: &mut Buffer,
        context: &HttpContext,
    ) -> Result<(usize, Instant)> {
        let mut len = 0;
        let mut start_with_read = buf.is_empty();
        let mut instant = None;
        let finder = Finder::new(b"\r\n\r\n");
        loop {
            if len >= context.config.head_max_len {
                return Err(Error::HeadersTooLong);
            }

            let bytes = if start_with_read {
                buf.expand();
                if len == 0 {
                    context
                        .swansong
                        .interrupt(transport.read(buf))
                        .await
                        .ok_or(Error::Closed)??
                } else {
                    transport.read(&mut buf[len..]).await?
                }
            } else {
                start_with_read = true;
                buf.len()
            };

            if instant.is_none() {
                instant = Some(Instant::now());
            }

            let search_start = len.max(3) - 3;
            let search = finder.find(&buf[search_start..]);

            if let Some(index) = search {
                buf.truncate(len + bytes);
                return Ok((search_start + index + 4, instant.unwrap()));
            }

            len += bytes;

            if bytes == 0 {
                return if len == 0 {
                    Err(Error::Closed)
                } else {
                    Err(Error::InvalidHead)
                };
            }
        }
    }

    async fn next(mut self) -> Result<Self> {
        // Drain unless we set up 100-continue and the client never started sending: in
        // that case no body bytes are coming and draining would block.
        if !self.needs_100_continue() {
            self.build_request_body().drain().await?;
        }
        Conn::new_internal(self.context, self.transport, self.buffer).await
    }

    fn should_close(&self) -> bool {
        let has_token = |headers: &Headers, token: &str| {
            headers
                .get_str(KnownHeaderName::Connection)
                .is_some_and(|v| v.split(',').any(|t| t.trim().eq_ignore_ascii_case(token)))
        };

        if has_token(&self.request_headers, "close") || has_token(&self.response_headers, "close") {
            true
        } else if has_token(&self.request_headers, "keep-alive")
            && has_token(&self.response_headers, "keep-alive")
        {
            false
        } else {
            self.version == Version::Http1_0
        }
    }

    async fn finish(self) -> Result<ConnectionStatus<Transport>> {
        if self.should_close() {
            Ok(ConnectionStatus::Close)
        } else if self.should_upgrade() {
            Ok(ConnectionStatus::Upgrade(self.into()))
        } else {
            match self.next().await {
                Err(Error::Closed) => {
                    log::trace!("connection closed by client");
                    Ok(ConnectionStatus::Close)
                }
                Err(e) => Err(e),
                Ok(conn) => Ok(ConnectionStatus::Conn(conn)),
            }
        }
    }

    fn request_content_length(&self) -> Result<Option<u64>> {
        if self
            .request_headers
            .has_header(KnownHeaderName::TransferEncoding)
        {
            Ok(None)
        } else if let Some(cl) = self.request_headers.get_str(KnownHeaderName::ContentLength) {
            cl.parse()
                .map(Some)
                .map_err(|_| Error::InvalidHeaderValue(KnownHeaderName::ContentLength.into()))
        } else if matches!(self.version, Version::Http2 | Version::Http3) {
            // h2 and h3 frame the body via stream-level END_STREAM; there's no equivalent of
            // h1's implicit "no content-length means empty body" default.
            Ok(None)
        } else {
            Ok(Some(0))
        }
    }

    pub(super) fn body_len(&self) -> Option<u64> {
        match self.response_body {
            Some(ref body) => body.len(),
            None => Some(0),
        }
    }

    fn write_headers(&mut self, output_buffer: &mut Vec<u8>) -> Result<()> {
        let status = self.status().unwrap_or(Status::NotFound);

        write!(
            output_buffer,
            "{} {} {}\r\n",
            self.version,
            status as u16,
            status.canonical_reason()
        )?;

        self.finalize_headers();

        log::trace!(
            "sending:\n{} {}\n{}",
            self.version,
            status,
            self.response_headers
        );

        write_headers_or_trailers(output_buffer, &self.response_headers, &self.context)?;

        write!(output_buffer, "\r\n")?;

        Ok(())
    }
}

/// Writes the HTTP/1.1 chunked header or trailer section + terminating CRLF to `writer`.
pub(crate) fn write_headers_or_trailers(
    output_buffer: &mut Vec<u8>,
    headers: &Headers,
    context: &HttpContext,
) -> Result<()> {
    let panic_on_invalid = context.config.panic_on_invalid_response_headers;

    for (name, values) in headers {
        if name.is_valid() {
            for value in values {
                if value.is_valid() {
                    write!(output_buffer, "{name}: ")?;
                    output_buffer.extend_from_slice(value.as_ref());
                    write!(output_buffer, "\r\n")?;
                } else if panic_on_invalid {
                    panic!("invalid response header value {value:?} for header {name}");
                } else {
                    log::error!("skipping invalid header value {value:?} for header {name}");
                }
            }
        } else if panic_on_invalid {
            panic!("invalid response header name {name:?}");
        } else {
            log::error!("skipping invalid header with name {name:?}");
        }
    }
    Ok(())
}