kvarn-fastcgi-client 0.9.0

// 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.

use crate::{
    meta::{EndRequestRec, Header, RequestType},
    ClientError, ClientResult,
};
use std::{cmp::min, fmt, fmt::Debug, str};
use tokio::io::{AsyncRead, AsyncReadExt};
use tracing::debug;

/// Output of fastcgi request, contains STDOUT and STDERR.
#[derive(Default, Clone)]
#[non_exhaustive]
pub struct Response {
    pub stdout: Option<Vec<u8>>,
    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()
    }
}

pub enum Content<'a> {
    Stdout(&'a [u8]),
    Stderr(&'a [u8]),
}

#[derive(PartialEq)]
enum ReadStep {
    Content,
    Padding,
}

/// Generated by
/// [Client::execute_once_stream](crate::client::Client::execute_once_stream) or
/// [Client::execute_stream](crate::client::Client::execute_stream).
///
/// The [ResponseStream] does not implement `futures::Stream`, because
/// `futures::Stream` does not yet support GAT, so manually provide the
/// [next](ResponseStream::next) method, which support the `while let` syntax.
pub struct ResponseStream<S: AsyncRead + Unpin> {
    stream: S,
    id: u16,

    ended: bool,

    header: Option<Header>,

    content_buf: Vec<u8>,
    content_read: usize,

    read_step: ReadStep,
}

impl<S: AsyncRead + Unpin + Send> ResponseStream<S> {
    #[inline]
    pub(crate) fn new(stream: S, id: u16) -> Self {
        Self {
            stream,
            id,
            ended: false,
            header: None,
            content_buf: vec![0; 4096],
            content_read: 0,
            read_step: ReadStep::Content,
        }
    }

    pub async fn next(&mut self) -> Option<ClientResult<Content<'_>>> {
        if self.ended {
            return None;
        }

        loop {
            if self.header.is_none() {
                match Header::new_from_stream(&mut self.stream).await {
                    Ok(header) => {
                        self.header = Some(header);
                    }
                    Err(err) => {
                        self.ended = true;
                        return Some(Err(err.into()));
                    }
                };
            }

            let header = self.header.as_ref().unwrap();

            match header.r#type.clone() {
                RequestType::Stdout => match self.read_step {
                    ReadStep::Content => {
                        return self
                            .read_to_content(
                                header.content_length as usize,
                                Content::Stdout,
                                Self::prepare_for_read_padding,
                            )
                            .await;
                    }
                    ReadStep::Padding => {
                        self.read_to_content(
                            header.padding_length as usize,
                            Content::Stdout,
                            Self::prepare_for_read_header,
                        )
                        .await;
                        continue;
                    }
                },
                RequestType::Stderr => match self.read_step {
                    ReadStep::Content => {
                        return self
                            .read_to_content(
                                header.content_length as usize,
                                Content::Stderr,
                                Self::prepare_for_read_padding,
                            )
                            .await;
                    }
                    ReadStep::Padding => {
                        self.read_to_content(
                            header.padding_length as usize,
                            Content::Stderr,
                            Self::prepare_for_read_header,
                        )
                        .await;
                        continue;
                    }
                },
                RequestType::EndRequest => {
                    let end_request_rec =
                        match EndRequestRec::from_header(header, &mut self.stream).await {
                            Ok(rec) => rec,
                            Err(err) => {
                                self.ended = true;
                                return Some(Err(err.into()));
                            }
                        };
                    debug!(id = self.id, ?end_request_rec, "Receive from stream.");

                    self.ended = true;

                    return match end_request_rec
                        .end_request
                        .protocol_status
                        .convert_to_client_result(end_request_rec.end_request.app_status)
                    {
                        Ok(_) => None,
                        Err(err) => Some(Err(err)),
                    };
                }
                r#type => {
                    self.ended = true;
                    return Some(Err(ClientError::UnknownRequestType {
                        request_type: r#type,
                    }));
                }
            }
        }
    }

    async fn read_to_content<'a, T: 'a>(
        &'a mut self, length: usize, content_fn: impl FnOnce(&'a [u8]) -> T,
        prepare_for_next_fn: impl FnOnce(&mut Self),
    ) -> Option<ClientResult<T>> {
        let content_len = self.content_buf.len();
        let read = match self
            .stream
            .read(&mut self.content_buf[..min(content_len, length - self.content_read)])
            .await
        {
            Ok(read) => read,
            Err(err) => {
                self.ended = true;
                return Some(Err(err.into()));
            }
        };

        self.content_read += read;
        if self.content_read >= length {
            self.content_read = 0;
            prepare_for_next_fn(self);
        }

        Some(Ok(content_fn(&self.content_buf[..read])))
    }

    fn prepare_for_read_padding(&mut self) {
        self.read_step = ReadStep::Padding;
    }

    fn prepare_for_read_header(&mut self) {
        self.header = None;
        self.read_step = ReadStep::Content;
    }
}