1use bytes::{Buf, BufMut};
3use prost::Message;
4use prost_reflect::{DescriptorPool, Kind, MessageDescriptor};
5use std::{collections::HashMap, marker::PhantomData, sync::LazyLock};
6use tonic::{
7 codec::{Codec, DecodeBuf, Decoder, EncodeBuf, Encoder},
8 Status,
9};
10pub const MAX_MESSAGE_BYTES: usize = 1024 * 1024;
11pub const MAX_VALUE_BYTES: usize = 64 * 1024;
12pub const MAX_NODES: usize = 10_000;
13pub const MAX_DEPTH: usize = 32;
14static DESCRIPTORS: LazyLock<DescriptorPool> = LazyLock::new(|| {
15 DescriptorPool::decode(
16 include_bytes!(concat!(env!("OUT_DIR"), "/gnmi-descriptors.pb")).as_slice(),
17 )
18 .expect("immutable compiled OpenConfig descriptors")
19});
20pub trait WireName {
21 const NAME: &'static str;
22}
23macro_rules! names {
24 ($($name:ident),*) => { $(impl WireName for super::proto::gnmi::$name {
25 const NAME: &'static str = concat!("gnmi.", stringify!($name));
26 })* };
27}
28names!(
29 CapabilityRequest,
30 CapabilityResponse,
31 GetRequest,
32 GetResponse,
33 SetRequest,
34 SetResponse,
35 SubscribeRequest,
36 SubscribeResponse
37);
38pub struct BoundedProstCodec<T, U>(PhantomData<(T, U)>);
39impl<T, U> Default for BoundedProstCodec<T, U> {
40 fn default() -> Self {
41 Self(PhantomData)
42 }
43}
44pub struct BoundedEncoder<T>(PhantomData<T>);
45pub struct BoundedDecoder<T>(PhantomData<T>);
46impl<T, U> Codec for BoundedProstCodec<T, U>
47where
48 T: Message + WireName + Send + 'static,
49 U: Message + WireName + Default + Send + 'static,
50{
51 type Encode = T;
52 type Decode = U;
53 type Encoder = BoundedEncoder<T>;
54 type Decoder = BoundedDecoder<U>;
55 fn encoder(&mut self) -> Self::Encoder {
56 BoundedEncoder(PhantomData)
57 }
58 fn decoder(&mut self) -> Self::Decoder {
59 BoundedDecoder(PhantomData)
60 }
61}
62impl<T: Message + WireName> Encoder for BoundedEncoder<T> {
63 type Item = T;
64 type Error = Status;
65 fn encode(&mut self, item: T, output: &mut EncodeBuf<'_>) -> Result<(), Status> {
66 if item.encoded_len() > MAX_MESSAGE_BYTES {
67 return Err(Status::resource_exhausted("gNMI message exceeds 1 MiB"));
68 }
69 let bytes = item.encode_to_vec();
70 validate_wire(T::NAME, &bytes)?;
71 output.put_slice(&bytes);
72 Ok(())
73 }
74}
75impl<T: Message + Default + WireName> Decoder for BoundedDecoder<T> {
76 type Item = T;
77 type Error = Status;
78 fn decode(&mut self, input: &mut DecodeBuf<'_>) -> Result<Option<T>, Status> {
79 if input.remaining() > MAX_MESSAGE_BYTES {
80 return Err(Status::resource_exhausted("gNMI message exceeds 1 MiB"));
81 }
82 let bytes = input.copy_to_bytes(input.remaining());
83 validate_wire(T::NAME, &bytes)?;
84 T::decode(bytes)
85 .map(Some)
86 .map_err(|_| Status::invalid_argument("invalid gNMI protobuf"))
87 }
88}
89pub fn validate_wire(name: &str, bytes: &[u8]) -> Result<(), Status> {
90 if bytes.len() > MAX_MESSAGE_BYTES {
91 return Err(Status::resource_exhausted("gNMI message exceeds 1 MiB"));
92 }
93 let descriptor = DESCRIPTORS
94 .get_message_by_name(name)
95 .ok_or_else(|| Status::internal("unknown immutable gNMI descriptor"))?;
96 scan(bytes, descriptor, 0, &mut 0)
97}
98fn varint(input: &mut &[u8]) -> Result<u64, Status> {
99 let mut value = 0u64;
100 for shift in (0..70).step_by(7) {
101 let (&byte, rest) = input
102 .split_first()
103 .ok_or_else(|| Status::invalid_argument("truncated protobuf varint"))?;
104 *input = rest;
105 if shift == 63 && byte > 1 {
106 return Err(Status::invalid_argument("overflowing protobuf varint"));
107 }
108 value |= u64::from(byte & 127) << shift;
109 if byte & 128 == 0 {
110 return Ok(value);
111 }
112 }
113 Err(Status::invalid_argument("invalid protobuf varint"))
114}
115fn take<'a>(input: &mut &'a [u8], length: usize) -> Result<&'a [u8], Status> {
116 if length > input.len() {
117 return Err(Status::invalid_argument("truncated protobuf field"));
118 }
119 let (value, rest) = input.split_at(length);
120 *input = rest;
121 Ok(value)
122}
123fn repeated_limit(message: &str, field: &str) -> usize {
124 match (message, field) {
125 ("gnmi.Path", "elem" | "element") => 32,
126 ("gnmi.PathElem", "key") => 8,
127 ("gnmi.ScalarArray", "element") => 64,
128 ("gnmi.Notification", "update" | "delete") => 256,
129 ("gnmi.GetResponse", "notification") => 128,
130 ("gnmi.SetResponse", "response") => 256,
131 _ => 128,
132 }
133}
134fn field_limit(message: &str, field: &str) -> usize {
135 match (message, field) {
136 ("gnmi.Path", "origin" | "target") => 256,
137 ("gnmi.Path", "element") | ("gnmi.PathElem", "name") => 128,
138 ("gnmi.PathElem.KeyEntry", "key") => 128,
139 ("gnmi.PathElem.KeyEntry", "value") => 1024,
140 ("gnmi.ModelData", _) => 256,
141 ("gnmi.Error", "message") => 4096,
142 _ => MAX_VALUE_BYTES,
143 }
144}
145fn scan(
146 mut input: &[u8],
147 descriptor: MessageDescriptor,
148 depth: usize,
149 nodes: &mut usize,
150) -> Result<(), Status> {
151 if depth > MAX_DEPTH {
152 return Err(Status::resource_exhausted("protobuf nesting exceeds 32"));
153 }
154 let name = descriptor.full_name();
155 let mut repeated = HashMap::<u32, usize>::new();
156 while !input.is_empty() {
157 *nodes += 1;
158 if *nodes > MAX_NODES {
159 return Err(Status::resource_exhausted(
160 "protobuf/JSON nodes exceed 10000",
161 ));
162 }
163 let key = varint(&mut input)?;
164 let number = u32::try_from(key >> 3)
165 .map_err(|_| Status::invalid_argument("protobuf field number overflow"))?;
166 if number == 0 || number > 0x1fff_ffff {
167 return Err(Status::invalid_argument("invalid protobuf field number"));
168 }
169 let wire = (key & 7) as u8;
170 let field = descriptor.get_field(number);
171 if let Some(field) = &field {
172 if field.is_list() || field.is_map() {
173 let count = repeated.entry(number).or_default();
174 *count += 1;
175 if *count > repeated_limit(name, field.name()) {
176 return Err(Status::resource_exhausted(
177 "protobuf repeated field exceeds bound",
178 ));
179 }
180 }
181 if name == "gnmi.TypedValue" && matches!(number, 5 | 9 | 13)
182 || name == "gnmi.Update" && number == 2
183 || matches!(field.kind(),Kind::Message(ref child) if child.full_name()=="gnmi_ext.Extension" || child.full_name()=="google.protobuf.Any")
184 {
185 return Err(Status::unimplemented(
186 "opaque bytes, Any, legacy values and extensions are excluded",
187 ));
188 }
189 let expected = match field.kind() {
190 Kind::Message(_) | Kind::String | Kind::Bytes => 2,
191 Kind::Double | Kind::Fixed64 | Kind::Sfixed64 => 1,
192 Kind::Float | Kind::Fixed32 | Kind::Sfixed32 => 5,
193 _ => 0,
194 };
195 if wire != expected && !(field.is_list() && field.is_packed() && wire == 2) {
196 return Err(Status::invalid_argument(
197 "protobuf field has wrong wire type",
198 ));
199 }
200 }
201 match wire {
202 0 => {
203 varint(&mut input)?;
204 }
205 1 => {
206 take(&mut input, 8)?;
207 }
208 5 => {
209 take(&mut input, 4)?;
210 }
211 2 => {
212 let length = usize::try_from(varint(&mut input)?)
213 .map_err(|_| Status::resource_exhausted("protobuf field length overflow"))?;
214 let payload = take(&mut input, length)?;
215 if let Some(field) = field {
216 match field.kind() {
217 Kind::Message(child) => scan(payload, child, depth + 1, nodes)?,
218 Kind::String | Kind::Bytes => {
219 if payload.len() > field_limit(name, field.name()) {
220 return Err(Status::resource_exhausted(
221 "protobuf text/value field exceeds bound",
222 ));
223 }
224 if name == "gnmi.TypedValue" && matches!(number, 10 | 11) {
225 super::json::parse_with_budget(payload, nodes)?;
226 }
227 }
228 kind if field.is_list() && field.is_packed() => {
229 let mut packed = payload;
230 let mut count = 0usize;
231 while !packed.is_empty() {
232 match kind {
233 Kind::Double | Kind::Fixed64 | Kind::Sfixed64 => {
234 take(&mut packed, 8)?;
235 }
236 Kind::Float | Kind::Fixed32 | Kind::Sfixed32 => {
237 take(&mut packed, 4)?;
238 }
239 _ => {
240 varint(&mut packed)?;
241 }
242 }
243 count += 1;
244 if count > repeated_limit(name, field.name()) {
245 return Err(Status::resource_exhausted(
246 "packed protobuf field exceeds bound",
247 ));
248 }
249 }
250 let entries = repeated.entry(number).or_default();
251 *entries += count.saturating_sub(1);
252 *nodes += count;
253 if *entries > repeated_limit(name, field.name()) || *nodes > MAX_NODES {
254 return Err(Status::resource_exhausted(
255 "packed protobuf count/node bound",
256 ));
257 }
258 }
259 _ => {}
260 }
261 }
262 }
263 _ => {
264 return Err(Status::invalid_argument(
265 "protobuf group or reserved wire type excluded",
266 ))
267 }
268 }
269 }
270 Ok(())
271}