fastcgi_connect/

response.rs

// Copyright 2022 jmjoy
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

//! FastCGI response types and streaming support.
//!
//! This module provides structures for handling FastCGI responses,
//! including both complete responses and streaming responses.

use std::{
    fmt::{self, Debug},
    pin::Pin,
    str,
    task::Poll,
};

use bytes::{Bytes, BytesMut};
use futures_core::stream::Stream;

#[cfg(feature = "tokio")]
use tokio::io::AsyncRead;
#[cfg(feature = "tokio")]
use tokio_util::io::ReaderStream;

#[cfg(feature = "smol")]
use smol::io::{AsyncRead, AsyncReadExt, Bytes as ReaderStream};

use tracing::debug;

use crate::{
    meta::{EndRequestRec, Header, RequestType, HEADER_LEN},
    ClientError, ClientResult,
};

/// Output of FastCGI request, contains STDOUT and STDERR data.
///
/// This structure represents a complete FastCGI response with
/// both stdout and stderr output from the FastCGI server.
#[derive(Default, Clone)]
#[non_exhaustive]
pub struct Response {
    /// The stdout output from the FastCGI server
    pub stdout: Option<Vec<u8>>,
    /// The stderr output from the FastCGI server
    pub stderr: Option<Vec<u8>>,
}

impl Debug for Response {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> {
        f.debug_struct("Response")
            .field("stdout", &self.stdout.as_deref().map(str::from_utf8))
            .field("stderr", &self.stderr.as_deref().map(str::from_utf8))
            .finish()
    }
}

/// Content type from a FastCGI response stream.
///
/// This enum represents the different types of content that can be
/// received from a FastCGI server during streaming.
pub enum Content {
    /// Standard output content from the FastCGI server
    Stdout(Bytes),
    /// Standard error content from the FastCGI server
    Stderr(Bytes),
}

/// A streaming response from a FastCGI server.
///
/// Generated by
/// [Client::execute_once_stream](crate::client::Client::execute_once_stream) or
/// [Client::execute_stream](crate::client::Client::execute_stream).
///
/// This stream yields `Content` items as they are received from the server.
pub struct ResponseStream<S: AsyncRead + Unpin> {
    #[cfg(feature = "tokio")]
    stream: ReaderStream<S>,
    #[cfg(feature = "smol")]
    stream: ReaderStream<S>,
    id: u16,
    eof: bool,
    header: Option<Header>,
    buf: BytesMut,
}

impl<S: AsyncRead + Unpin> ResponseStream<S> {
    /// Creates a new response stream.
    ///
    /// # Arguments
    ///
    /// * `stream` - The underlying stream to read from
    /// * `id` - The request ID for this response
    #[inline]
    pub(crate) fn new(stream: S, id: u16) -> Self {
        Self {
            #[cfg(feature = "tokio")]
            stream: ReaderStream::new(stream),
            #[cfg(feature = "smol")]
            stream: stream.bytes(),
            id,
            eof: false,
            header: None,
            buf: BytesMut::new(),
        }
    }

    /// Reads a FastCGI header from the buffer.
    ///
    /// Returns `None` if there isn't enough data in the buffer.
    #[inline]
    fn read_header(&mut self) -> Option<Header> {
        if self.buf.len() < HEADER_LEN {
            return None;
        }
        let buf = self.buf.split_to(HEADER_LEN);
        let header = (&buf as &[u8]).try_into().expect("failed to read header");
        Some(Header::new_from_buf(header))
    }

    /// Reads content from the buffer based on the current header.
    ///
    /// Returns `None` if there isn't enough data in the buffer.
    #[inline]
    fn read_content(&mut self) -> Option<Bytes> {
        let header = self.header.as_ref().unwrap();
        let block_length = header.content_length as usize + header.padding_length as usize;
        if self.buf.len() < block_length {
            return None;
        }
        let content = self.buf.split_to(header.content_length as usize);
        let _ = self.buf.split_to(header.padding_length as usize);
        self.header = None;
        Some(content.freeze())
    }

    /// Processes a complete FastCGI message from the buffer.
    ///
    /// Returns `Ok(Some(Content))` if a complete message was processed,
    /// `Ok(None)` if more data is needed, or an error if processing failed.
    fn process_message(&mut self) -> Result<Option<Content>, ClientError> {
        if self.buf.is_empty() {
            return Ok(None);
        }
        if self.header.is_none() {
            match self.read_header() {
                Some(header) => self.header = Some(header),
                None => return Ok(None),
            }
        }
        let header = self.header.as_ref().unwrap();
        match header.r#type.clone() {
            RequestType::Stdout => {
                if let Some(data) = self.read_content() {
                    return Ok(Some(Content::Stdout(data)));
                }
            }
            RequestType::Stderr => {
                if let Some(data) = self.read_content() {
                    return Ok(Some(Content::Stderr(data)));
                }
            }
            RequestType::EndRequest => {
                let header = header.clone();
                let Some(data) = self.read_content() else {
                    return Ok(None);
                };

                let end = EndRequestRec::new_from_buf(header, &data);
                debug!(id = self.id, ?end, "Receive from stream.");

                self.eof = true;
                end.end_request
                    .protocol_status
                    .convert_to_client_result(end.end_request.app_status)?;
                return Ok(None);
            }
            r#type => {
                self.eof = true;
                return Err(ClientError::UnknownRequestType {
                    request_type: r#type,
                });
            }
        }
        Ok(None)
    }
}

impl<S> Stream for ResponseStream<S>
where
    S: AsyncRead + Unpin,
{
    type Item = ClientResult<Content>;

    fn poll_next(
        mut self: std::pin::Pin<&mut Self>, cx: &mut std::task::Context<'_>,
    ) -> std::task::Poll<Option<Self::Item>> {
        let mut pending = false;
        loop {
            match Pin::new(&mut self.stream).poll_next(cx) {
                Poll::Ready(Some(Ok(data))) => {
                    #[cfg(feature = "smol")]
                    self.buf.extend_from_slice(&[data]);
                    #[cfg(feature = "tokio")]
                    self.buf.extend_from_slice(&data);

                    match self.process_message() {
                        Ok(Some(data)) => return Poll::Ready(Some(Ok(data))),
                        Ok(None) if self.eof => return Poll::Ready(None),
                        Ok(None) => continue,
                        Err(err) => return Poll::Ready(Some(Err(err))),
                    }
                }
                Poll::Ready(Some(Err(err))) => return Poll::Ready(Some(Err(err.into()))),
                Poll::Ready(None) => break,
                Poll::Pending => {
                    pending = true;
                    break;
                }
            }
        }
        match self.process_message() {
            Ok(Some(data)) => Poll::Ready(Some(Ok(data))),
            Ok(None) if !self.eof && pending => Poll::Pending,
            Ok(None) => Poll::Ready(None),
            Err(err) => Poll::Ready(Some(Err(err))),
        }
    }
}