micro_http/codec/
response_encoder.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
//! HTTP response encoder module
//! 
//! This module provides functionality for encoding HTTP responses using a streaming approach.
//! It handles both header encoding and payload encoding through a state machine pattern.
//! 
//! # Components
//! 
//! - [`ResponseEncoder`]: Main encoder that coordinates header and payload encoding
//! - Header encoding: Uses [`HeaderEncoder`] for encoding response headers
//! - Payload handling: Uses [`PayloadEncoder`] for encoding response bodies
//! 
//! # Example
//! 
//! ```no_run
//! use micro_http::codec::ResponseEncoder;
//! use tokio_util::codec::Encoder;
//! use bytes::BytesMut;
//! 
//! let mut encoder = ResponseEncoder::new();
//! let mut buffer = BytesMut::new();
//! // ... encode response data to buffer ...
//! ```

use crate::codec::body::PayloadEncoder;
use crate::codec::header::HeaderEncoder;
use crate::protocol::{Message, PayloadSize, ResponseHead, SendError};
use bytes::{Buf, BytesMut};
use std::io;
use std::io::ErrorKind;
use tokio_util::codec::Encoder;
use tracing::error;

/// A encoder for HTTP responses that handles both headers and payload
/// 
/// The encoder operates in two phases:
/// 1. Header encoding: Encodes the response headers using [`HeaderEncoder`]
/// 2. Payload encoding: If present, encodes the response body using [`PayloadEncoder`]
pub struct ResponseEncoder {
    header_encoder: HeaderEncoder,
    payload_encoder: Option<PayloadEncoder>,
}

impl ResponseEncoder {
    /// Creates a new `ResponseEncoder` instance
    pub fn new() -> Self {
        Default::default()
    }
}

impl Default for ResponseEncoder {
    fn default() -> Self {
        Self { header_encoder: HeaderEncoder, payload_encoder: None }
    }
}

impl<D: Buf> Encoder<Message<(ResponseHead, PayloadSize), D>> for ResponseEncoder {
    type Error = SendError;

    /// Attempts to encode an HTTP response to the provided buffer
    /// 
    /// # Arguments
    /// 
    /// * `item` - The message to encode, either headers or payload
    /// * `dst` - The buffer to write the encoded data to
    /// 
    /// # Returns
    /// 
    /// - `Ok(())`: Successfully encoded the message
    /// - `Err(_)`: Encountered an encoding error
    fn encode(&mut self, item: Message<(ResponseHead, PayloadSize), D>, dst: &mut BytesMut) -> Result<(), Self::Error> {
        match item {
            Message::Header((head, payload_size)) => {
                if self.payload_encoder.is_some() {
                    error!("expect payload item but receive response head");
                    return Err(io::Error::from(ErrorKind::InvalidInput).into());
                }

                let payload_encoder = parse_payload_encoder(payload_size);
                self.payload_encoder = Some(payload_encoder);
                self.header_encoder.encode((head, payload_size), dst)
            }

            Message::Payload(payload_item) => {
                let payload_encoder = if let Some(encoder) = &mut self.payload_encoder {
                    encoder
                } else {
                    error!("expect response header but receive payload item");
                    return Err(io::Error::from(ErrorKind::InvalidInput).into());
                };

                let result = payload_encoder.encode(payload_item, dst);

                let is_eof = payload_encoder.is_finish();
                if is_eof {
                    self.payload_encoder.take();
                }

                result
            }
        }
    }
}

/// Creates a payload encoder based on the payload size
/// 
/// # Arguments
/// 
/// * `payload_size` - The size specification for the payload
/// 
/// # Returns
/// 
/// Returns a [`PayloadEncoder`] configured according to the payload size
fn parse_payload_encoder(payload_size: PayloadSize) -> PayloadEncoder {
    match payload_size {
        PayloadSize::Length(size) => PayloadEncoder::fix_length(size),
        PayloadSize::Chunked => PayloadEncoder::chunked(),
        PayloadSize::Empty => PayloadEncoder::empty(),
    }
}