Skip to main content

netget/server/grpc/
stream_codec.rs

1//! Library-owned framing with the audited field-name JSON conversion for streaming RPCs.
2use 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
11/// Exclude bytes fields, including nested/map values, before JSON conversion.
12pub 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
89/// New streaming actions never accept or expose encoded protobuf bytes fields.
90/// Legacy unary conversion remains unchanged. Walk borrowed values with the same
91/// depth/node guards before converting them to model-visible JSON.
92fn 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}