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

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;

pub struct ResponseEncoder {
    header_encoder: HeaderEncoder,
    payload_encoder: Option<PayloadEncoder>,
}

impl ResponseEncoder {
    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;

    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
            }
        }
    }
}

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(),
    }
}