Skip to main content

netget/server/gnmi/
codec.rs

1//! Check protobuf structure before prost allocates repeated fields or nested messages.
2use 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}