netget/server/grpc/
stream_codec.rs1use anyhow::{bail, ensure, Result};
3use bytes::Buf;
4use prost::Message;
5use prost_reflect::{DynamicMessage, MessageDescriptor, ReflectMessage, Value};
6use tonic::codec::{Codec, DecodeBuf, Decoder, EncodeBuf, Encoder};
7use tonic::Status;
8
9pub const MAX_MESSAGE_BYTES: usize = 4 * 1024 * 1024;
10
11pub fn check_descriptor(descriptor: &MessageDescriptor) -> Result<()> {
13 let mut pending = vec![descriptor.clone()];
14 let mut visited = std::collections::HashSet::new();
15 while let Some(message) = pending.pop() {
16 if !visited.insert(message.full_name().to_owned()) {
17 continue;
18 }
19 ensure!(visited.len() <= 100_000, "stream schema exceeds type limit");
20 ensure!(
21 message.full_name().len() <= 1024 && message.fields().len() <= 128,
22 "stream message schema exceeds name/field bound"
23 );
24 for field in message.fields() {
25 ensure!(
26 field.name().len() <= 256,
27 "stream field name exceeds 256 bytes"
28 );
29 match field.kind() {
30 prost_reflect::Kind::Bytes => {
31 bail!("bytes fields are outside the typed streaming subset")
32 }
33 prost_reflect::Kind::Message(child) => pending.push(child),
34 _ => {}
35 }
36 }
37 }
38 Ok(())
39}
40
41#[derive(Clone)]
42pub struct DynamicCodec {
43 pub encode: MessageDescriptor,
44 pub decode: MessageDescriptor,
45}
46pub struct DynamicEncoder(MessageDescriptor);
47pub struct DynamicDecoder(MessageDescriptor);
48impl Codec for DynamicCodec {
49 type Encode = DynamicMessage;
50 type Decode = DynamicMessage;
51 type Encoder = DynamicEncoder;
52 type Decoder = DynamicDecoder;
53 fn encoder(&mut self) -> Self::Encoder {
54 DynamicEncoder(self.encode.clone())
55 }
56 fn decoder(&mut self) -> Self::Decoder {
57 DynamicDecoder(self.decode.clone())
58 }
59}
60impl Encoder for DynamicEncoder {
61 type Item = DynamicMessage;
62 type Error = Status;
63 fn encode(&mut self, item: Self::Item, dst: &mut EncodeBuf<'_>) -> Result<(), Self::Error> {
64 if item.descriptor() != self.0 {
65 return Err(Status::invalid_argument(
66 "stream message type does not match method",
67 ));
68 }
69 if item.encoded_len() > MAX_MESSAGE_BYTES {
70 return Err(Status::resource_exhausted("stream message exceeds 4 MiB"));
71 }
72 item.encode(dst)
73 .map_err(|_| Status::internal("stream message encoding failed"))
74 }
75}
76impl Decoder for DynamicDecoder {
77 type Item = DynamicMessage;
78 type Error = Status;
79 fn decode(&mut self, src: &mut DecodeBuf<'_>) -> Result<Option<Self::Item>, Self::Error> {
80 if src.remaining() > MAX_MESSAGE_BYTES {
81 return Err(Status::resource_exhausted("stream message exceeds 4 MiB"));
82 }
83 DynamicMessage::decode(self.0.clone(), src)
84 .map(Some)
85 .map_err(|_| Status::invalid_argument("stream message does not decode"))
86 }
87}
88
89fn check_value(value: &Value, depth: usize, nodes: &mut usize) -> Result<()> {
93 use super::value_codec::{ValueLimitExceeded, MAX_VALUE_DEPTH, MAX_VALUE_NODES};
94 if depth > MAX_VALUE_DEPTH {
95 return Err(ValueLimitExceeded {
96 dimension: "depth",
97 limit: MAX_VALUE_DEPTH,
98 }
99 .into());
100 }
101 if *nodes >= MAX_VALUE_NODES {
102 return Err(ValueLimitExceeded {
103 dimension: "nodes",
104 limit: MAX_VALUE_NODES,
105 }
106 .into());
107 }
108 *nodes += 1;
109 match value {
110 Value::Bytes(_) => bail!("bytes fields are outside the typed streaming subset"),
111 Value::Message(message) => {
112 for (_, value) in message.fields() {
113 check_value(value, depth + 1, nodes)?;
114 }
115 }
116 Value::List(values) => {
117 for value in values {
118 check_value(value, depth + 1, nodes)?;
119 }
120 }
121 Value::Map(values) => {
122 for value in values.values() {
123 check_value(value, depth + 1, nodes)?;
124 }
125 }
126 _ => {}
127 }
128 Ok(())
129}
130fn check_message(message: &DynamicMessage) -> Result<()> {
131 let mut nodes = 1;
132 for (_, value) in message.fields() {
133 check_value(value, 1, &mut nodes)?;
134 }
135 Ok(())
136}
137pub fn to_json(message: &DynamicMessage) -> Result<serde_json::Value> {
138 check_message(message)?;
139 super::value_codec::dynamic_message_to_json(message)
140}
141pub fn from_json(
142 json: &serde_json::Value,
143 descriptor: &MessageDescriptor,
144) -> Result<DynamicMessage> {
145 check_descriptor(descriptor)?;
146 let message = super::value_codec::json_to_dynamic_message(json, descriptor)?;
147 ensure!(
148 message.encoded_len() <= MAX_MESSAGE_BYTES,
149 "stream message exceeds 4 MiB"
150 );
151 check_message(&message)?;
152 Ok(message)
153}