structured_proxy/transcode/
codec.rs1use prost::Message;
7use prost_reflect::{Cardinality, DynamicMessage, Kind, MessageDescriptor, ReflectMessage, Value};
8use tonic::codec::{BufferSettings, Codec, DecodeBuf, Decoder, EncodeBuf, Encoder};
9use tonic::Status;
10
11#[derive(Debug, Clone)]
13pub struct DynamicEncoder;
14
15impl Encoder for DynamicEncoder {
16 type Item = DynamicMessage;
17 type Error = Status;
18
19 fn encode(&mut self, item: Self::Item, buf: &mut EncodeBuf<'_>) -> Result<(), Status> {
20 item.encode(buf)
21 .map_err(|e| Status::internal(format!("encode error: {e}")))
22 }
23
24 fn buffer_settings(&self) -> BufferSettings {
25 BufferSettings::default()
26 }
27}
28
29#[derive(Debug, Clone)]
31pub struct DynamicDecoder {
32 desc: MessageDescriptor,
33 check_required: bool,
36}
37
38impl DynamicDecoder {
39 pub fn new(desc: MessageDescriptor) -> Self {
40 let check_required = has_required_fields(&desc);
41 Self {
42 desc,
43 check_required,
44 }
45 }
46}
47
48impl Decoder for DynamicDecoder {
49 type Item = DynamicMessage;
50 type Error = Status;
51
52 fn decode(&mut self, buf: &mut DecodeBuf<'_>) -> Result<Option<Self::Item>, Status> {
59 let msg = DynamicMessage::decode(self.desc.clone(), buf)
60 .map_err(|e| Status::internal(format!("decode error: {e}")))?;
61 if self.check_required {
62 if let Some(field) = missing_required(&msg) {
63 return Err(Status::internal(format!(
64 "decode error: required field {field} is missing"
65 )));
66 }
67 }
68 Ok(Some(msg))
69 }
70
71 fn buffer_settings(&self) -> BufferSettings {
72 BufferSettings::default()
73 }
74}
75
76#[derive(Debug, Clone)]
78pub struct DynamicCodec {
79 response_desc: MessageDescriptor,
80 check_required: bool,
81}
82
83impl DynamicCodec {
84 pub fn new(response_desc: MessageDescriptor) -> Self {
85 let check_required = has_required_fields(&response_desc);
86 Self::with_required_check(response_desc, check_required)
87 }
88
89 pub(crate) fn with_required_check(
92 response_desc: MessageDescriptor,
93 check_required: bool,
94 ) -> Self {
95 Self {
96 response_desc,
97 check_required,
98 }
99 }
100}
101
102impl Codec for DynamicCodec {
103 type Encode = DynamicMessage;
104 type Decode = DynamicMessage;
105 type Encoder = DynamicEncoder;
106 type Decoder = DynamicDecoder;
107
108 fn encoder(&mut self) -> Self::Encoder {
109 DynamicEncoder
110 }
111
112 fn decoder(&mut self) -> Self::Decoder {
113 DynamicDecoder {
114 desc: self.response_desc.clone(),
115 check_required: self.check_required,
116 }
117 }
118}
119
120pub(crate) fn has_required_fields(desc: &MessageDescriptor) -> bool {
123 let mut seen = std::collections::HashSet::new();
124 let mut pending = vec![desc.clone()];
125 while let Some(desc) = pending.pop() {
126 if !seen.insert(desc.full_name().to_owned()) {
127 continue;
128 }
129 for field in desc.fields() {
130 if field.cardinality() == Cardinality::Required {
131 return true;
132 }
133 if let Kind::Message(inner) = field.kind() {
134 pending.push(inner);
135 }
136 }
137 for extension in desc.extensions() {
138 if let Kind::Message(inner) = extension.kind() {
139 pending.push(inner);
140 }
141 }
142 }
143 false
144}
145
146pub(crate) fn missing_required(msg: &DynamicMessage) -> Option<String> {
149 let desc = msg.descriptor();
150 if let Some(field) = desc
151 .fields()
152 .find(|field| field.cardinality() == Cardinality::Required && !msg.has_field(field))
153 {
154 return Some(field.full_name().to_owned());
155 }
156 let values = msg
157 .fields()
158 .map(|(_, value)| value)
159 .chain(msg.extensions().map(|(_, value)| value));
160 for value in values {
161 let missing = match value {
162 Value::Message(inner) => missing_required(inner),
163 Value::List(items) => items.iter().find_map(|item| match item {
164 Value::Message(inner) => missing_required(inner),
165 _ => None,
166 }),
167 Value::Map(entries) => entries.values().find_map(|entry| match entry {
168 Value::Message(inner) => missing_required(inner),
169 _ => None,
170 }),
171 _ => None,
172 };
173 if missing.is_some() {
174 return missing;
175 }
176 }
177 None
178}
179
180#[cfg(test)]
181mod tests;