gproxy-transform 3.0.0-alpha.0

GPROXY v3 pairwise protocol transforms
Documentation 
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
120
121
122
123
124
125
126
127

use bytes::Bytes;
use gproxy_protocol::StreamFraming;

use super::super::json_array::{JsonArrayDecoder, JsonArrayEncoder};
use super::super::{SseDecoder, SseFrame};
use crate::TransformError;

pub(super) enum FrameDecoder {
    Sse(SseDecoder),
    JsonArray(JsonArrayDecoder),
}

pub(super) enum FrameEncoder {
    Sse,
    JsonArray {
        decoder: SseDecoder,
        encoder: JsonArrayEncoder,
        done: bool,
    },
}

impl FrameDecoder {
    pub(super) fn new(framing: StreamFraming) -> Result<Self, TransformError> {
        match framing {
            StreamFraming::Sse => Ok(Self::Sse(SseDecoder::default())),
            StreamFraming::JsonArray => Ok(Self::JsonArray(JsonArrayDecoder::default())),
            StreamFraming::WebSocket => Err(TransformError::shape(
                "stream framing",
                "websocket framing requires a websocket transform",
            )),
        }
    }

    pub(super) fn push(&mut self, chunk: &[u8]) -> Result<Vec<SseFrame>, TransformError> {
        match self {
            Self::Sse(decoder) => decoder.push(chunk),
            Self::JsonArray(decoder) => decoder.push(chunk),
        }
    }

    pub(super) fn finish(&mut self) -> Result<Vec<SseFrame>, TransformError> {
        match self {
            Self::Sse(decoder) => Ok(decoder.finish()?.into_iter().collect()),
            Self::JsonArray(decoder) => decoder.finish(),
        }
    }
}

impl FrameEncoder {
    pub(super) fn new(framing: StreamFraming) -> Result<Self, TransformError> {
        match framing {
            StreamFraming::Sse => Ok(Self::Sse),
            StreamFraming::JsonArray => Ok(Self::JsonArray {
                decoder: SseDecoder::default(),
                encoder: JsonArrayEncoder::default(),
                done: false,
            }),
            StreamFraming::WebSocket => Err(TransformError::shape(
                "stream framing",
                "websocket framing requires a websocket transform",
            )),
        }
    }

    pub(super) fn push(&mut self, chunks: Vec<Bytes>) -> Result<Vec<Bytes>, TransformError> {
        match self {
            Self::Sse => Ok(chunks),
            Self::JsonArray {
                decoder,
                encoder,
                done,
            } => {
                let mut output = Vec::new();
                for chunk in chunks {
                    for frame in decoder.push(&chunk)? {
                        encode_array_frame(encoder, done, frame, &mut output)?;
                    }
                }
                Ok(output)
            }
        }
    }

    pub(super) fn finish(&mut self) -> Result<Vec<Bytes>, TransformError> {
        match self {
            Self::Sse => Ok(Vec::new()),
            Self::JsonArray {
                decoder,
                encoder,
                done,
            } => {
                let mut output = Vec::new();
                if let Some(frame) = decoder.finish()? {
                    encode_array_frame(encoder, done, frame, &mut output)?;
                }
                output.push(encoder.finish()?);
                Ok(output)
            }
        }
    }
}

fn encode_array_frame(
    encoder: &mut JsonArrayEncoder,
    done: &mut bool,
    frame: SseFrame,
    output: &mut Vec<Bytes>,
) -> Result<(), TransformError> {
    if frame.data == "[DONE]" {
        if *done {
            return Err(TransformError::shape(
                "stream framing",
                "duplicate [DONE] sentinel",
            ));
        }
        *done = true;
        return Ok(());
    }
    if *done {
        return Err(TransformError::shape(
            "stream framing",
            "data followed [DONE] sentinel",
        ));
    }
    output.push(encoder.push(&frame.data)?);
    Ok(())
}