Support custom recursion limits at build time

prost-build/src/code_generator.rs

@@ -189,6 +189,7 @@ impl<'a> CodeGenerator<'a> {
             "#[derive(Clone, PartialEq, {}::Message)]\n",
             self.config.prost_path.as_deref().unwrap_or("::prost")
         ));
+        self.append_recursion_limit(&fq_message_name);
         self.push_indent();
         self.buf.push_str("pub struct ");
         self.buf.push_str(&to_upper_camel(&message_name));
@@ -272,6 +273,15 @@ impl<'a> CodeGenerator<'a> {
         }
     }
 
+    fn append_recursion_limit(&mut self, fq_message_name: &str) {
+        assert_eq!(b'.', fq_message_name.as_bytes()[0]);
+        if let Some(limit) = self.config.recursion_limits.get_first(fq_message_name) {
+            push_indent(self.buf, self.depth);
+            self.buf.push_str(&format!("#[RecursionLimit({})]", limit));
+            self.buf.push('\n');
+        }
+    }
+
     fn append_field_attributes(&mut self, fq_message_name: &str, field_name: &str) {
         assert_eq!(b'.', fq_message_name.as_bytes()[0]);
         for attribute in self

prost-build/src/lib.rs

@@ -245,6 +245,7 @@ pub struct Config {
     bytes_type: PathMap<BytesType>,
     type_attributes: PathMap<String>,
     field_attributes: PathMap<String>,
+    recursion_limits: PathMap<u32>,
     prost_types: bool,
     strip_enum_prefix: bool,
     out_dir: Option<PathBuf>,
@@ -468,6 +469,25 @@ impl Config {
         self
     }
 
+    /// Configure a custom recursion limit for certain messages.
+    ///
+    /// This defaults to 100, and can be disabled with the no-recursion-limit crate feature.
+    ///
+    /// # Example
+    ///
+    /// ```rust
+    /// # let mut config = prost_build::Config::new();
+    /// config.recursion_limit("my_messages.MyMessageType", 1000);
+    /// ```
+    pub fn recursion_limit<P>(&mut self, path: P, limit: u32) -> &mut Self
+    where
+        P: AsRef<str>,
+    {
+        self.recursion_limits
+            .insert(path.as_ref().to_string(), limit);
+        self
+    }
+
     /// Configures the code generator to use the provided service generator.
     pub fn service_generator(&mut self, service_generator: Box<dyn ServiceGenerator>) -> &mut Self {
         self.service_generator = Some(service_generator);
@@ -1101,6 +1121,7 @@ impl default::Default for Config {
             bytes_type: PathMap::default(),
             type_attributes: PathMap::default(),
             field_attributes: PathMap::default(),
+            recursion_limits: PathMap::default(),
             prost_types: true,
             strip_enum_prefix: true,
             out_dir: None,

prost-derive/src/lib.rs

@@ -23,6 +23,20 @@ fn try_message(input: TokenStream) -> Result<TokenStream, Error> {
 
     let ident = input.ident;
 
+    let recursion_limit: u32 = if let Some(attr) = input
+        .attrs
+        .iter()
+        .find(|attr| attr.path.is_ident("RecursionLimit"))
+    {
+        if let syn::Lit::Int(attr) = attr.parse_args().unwrap() {
+            attr.base10_parse().unwrap()
+        } else {
+            panic!("unexpected RecursionLimit type: {:?}", attr)
+        }
+    } else {
+        100
+    };
+
     let variant_data = match input.data {
         Data::Struct(variant_data) => variant_data,
         Data::Enum(..) => bail!("Message can not be derived for an enum"),
@@ -187,6 +201,10 @@ fn try_message(input: TokenStream) -> Result<TokenStream, Error> {
 
     let expanded = quote! {
         impl #impl_generics ::prost::Message for #ident #ty_generics #where_clause {
+            fn recursion_limit() -> u32 {
+                #recursion_limit
+            }
+
             #[allow(unused_variables)]
             fn encode_raw<B>(&self, buf: &mut B) where B: ::prost::bytes::BufMut {
                 #(#encode)*
@@ -238,7 +256,7 @@ fn try_message(input: TokenStream) -> Result<TokenStream, Error> {
     Ok(expanded.into())
 }
 
-#[proc_macro_derive(Message, attributes(prost))]
+#[proc_macro_derive(Message, attributes(prost, RecursionLimit))]
 pub fn message(input: TokenStream) -> TokenStream {
     try_message(input).unwrap()
 }

src/encoding.rs

@@ -195,29 +195,25 @@ pub struct DecodeContext {
     /// How many times we can recurse in the current decode stack before we hit
     /// the recursion limit.
     ///
-    /// The recursion limit is defined by `RECURSION_LIMIT` and cannot be
-    /// customized. The recursion limit can be ignored by building the Prost
-    /// crate with the `no-recursion-limit` feature.
+    /// It defaults to 100 and can be changed using `prost_build::recursion_limit`,
+    /// or it can be disabled entirely using the `no-recursion-limit` feature.
     #[cfg(not(feature = "no-recursion-limit"))]
-    recurse_count: u32,
+    pub recurse_count: u32,
 }
 
-#[cfg(not(feature = "no-recursion-limit"))]
-impl Default for DecodeContext {
-    #[inline]
-    fn default() -> DecodeContext {
+impl DecodeContext {
+    pub(crate) fn new(recursion_limit: u32) -> DecodeContext {
         DecodeContext {
-            recurse_count: crate::RECURSION_LIMIT,
+            #[cfg(not(feature = "no-recursion-limit"))]
+            recurse_count: recursion_limit,
         }
     }
-}
 
-impl DecodeContext {
     /// Call this function before recursively decoding.
     ///
     /// There is no `exit` function since this function creates a new `DecodeContext`
     /// to be used at the next level of recursion. Continue to use the old context
-    // at the previous level of recursion.
+    /// at the previous level of recursion.
     #[cfg(not(feature = "no-recursion-limit"))]
     #[inline]
     pub(crate) fn enter_recursion(&self) -> DecodeContext {
@@ -1503,7 +1499,7 @@ mod test {
             wire_type,
             &mut roundtrip_value,
             &mut buf,
-            DecodeContext::default(),
+            DecodeContext::new(100),
         )
         .map_err(|error| TestCaseError::fail(error.to_string()))?;
 
@@ -1575,7 +1571,7 @@ mod test {
                 wire_type,
                 &mut roundtrip_value,
                 &mut buf,
-                DecodeContext::default(),
+                DecodeContext::new(100),
             )
             .map_err(|error| TestCaseError::fail(error.to_string()))?;
         }
@@ -1594,7 +1590,7 @@ mod test {
             WireType::LengthDelimited,
             &mut s,
             &mut &buf[..],
-            DecodeContext::default(),
+            DecodeContext::new(100),
         );
         r.expect_err("must be an error");
         assert!(s.is_empty());

src/lib.rs

@@ -23,11 +23,6 @@ use bytes::{Buf, BufMut};
 
 use crate::encoding::{decode_varint, encode_varint, encoded_len_varint};
 
-// See `encoding::DecodeContext` for more info.
-// 100 is the default recursion limit in the C++ implementation.
-#[cfg(not(feature = "no-recursion-limit"))]
-const RECURSION_LIMIT: u32 = 100;
-
 /// Encodes a length delimiter to the buffer.
 ///
 /// See [Message.encode_length_delimited] for more info.

src/message.rs

@@ -14,6 +14,13 @@ use crate::EncodeError;
 
 /// A Protocol Buffers message.
 pub trait Message: Debug + Send + Sync {
+    /// The recursion limit for decoding protobuf messages.
+    ///
+    /// Defaults to 100. Can be customized in your build.rs or by using the no-recursion-limit crate feature.
+    fn recursion_limit() -> u32
+    where
+        Self: Sized;
+
     /// Encodes the message to a buffer.
     ///
     /// This method will panic if the buffer has insufficient capacity.
@@ -135,7 +142,7 @@ pub trait Message: Debug + Send + Sync {
         B: Buf,
         Self: Sized,
     {
-        let ctx = DecodeContext::default();
+        let ctx = DecodeContext::new(Self::recursion_limit());
         while buf.has_remaining() {
             let (tag, wire_type) = decode_key(&mut buf)?;
             self.merge_field(tag, wire_type, &mut buf, ctx.clone())?;
@@ -154,7 +161,7 @@ pub trait Message: Debug + Send + Sync {
             WireType::LengthDelimited,
             self,
             &mut buf,
-            DecodeContext::default(),
+            DecodeContext::new(Self::recursion_limit()),
         )
     }
 
@@ -166,6 +173,9 @@ impl<M> Message for Box<M>
 where
     M: Message,
 {
+    fn recursion_limit() -> u32 {
+        M::recursion_limit()
+    }
     fn encode_raw<B>(&self, buf: &mut B)
     where
         B: BufMut,

src/types.rs

@@ -20,6 +20,9 @@ use crate::{
 
 /// `google.protobuf.BoolValue`
 impl Message for bool {
+    fn recursion_limit() -> u32 {
+        0
+    }
     fn encode_raw<B>(&self, buf: &mut B)
     where
         B: BufMut,
@@ -58,6 +61,9 @@ impl Message for bool {
 
 /// `google.protobuf.UInt32Value`
 impl Message for u32 {
+    fn recursion_limit() -> u32 {
+        0
+    }
     fn encode_raw<B>(&self, buf: &mut B)
     where
         B: BufMut,
@@ -96,6 +102,9 @@ impl Message for u32 {
 
 /// `google.protobuf.UInt64Value`
 impl Message for u64 {
+    fn recursion_limit() -> u32 {
+        0
+    }
     fn encode_raw<B>(&self, buf: &mut B)
     where
         B: BufMut,
@@ -134,6 +143,10 @@ impl Message for u64 {
 
 /// `google.protobuf.Int32Value`
 impl Message for i32 {
+    fn recursion_limit() -> u32 {
+        0
+    }
+
     fn encode_raw<B>(&self, buf: &mut B)
     where
         B: BufMut,
@@ -172,6 +185,9 @@ impl Message for i32 {
 
 /// `google.protobuf.Int64Value`
 impl Message for i64 {
+    fn recursion_limit() -> u32 {
+        0
+    }
     fn encode_raw<B>(&self, buf: &mut B)
     where
         B: BufMut,
@@ -210,6 +226,9 @@ impl Message for i64 {
 
 /// `google.protobuf.FloatValue`
 impl Message for f32 {
+    fn recursion_limit() -> u32 {
+        0
+    }
     fn encode_raw<B>(&self, buf: &mut B)
     where
         B: BufMut,
@@ -248,6 +267,9 @@ impl Message for f32 {
 
 /// `google.protobuf.DoubleValue`
 impl Message for f64 {
+    fn recursion_limit() -> u32 {
+        0
+    }
     fn encode_raw<B>(&self, buf: &mut B)
     where
         B: BufMut,
@@ -286,6 +308,9 @@ impl Message for f64 {
 
 /// `google.protobuf.StringValue`
 impl Message for String {
+    fn recursion_limit() -> u32 {
+        0
+    }
     fn encode_raw<B>(&self, buf: &mut B)
     where
         B: BufMut,
@@ -324,6 +349,9 @@ impl Message for String {
 
 /// `google.protobuf.BytesValue`
 impl Message for Vec<u8> {
+    fn recursion_limit() -> u32 {
+        0
+    }
     fn encode_raw<B>(&self, buf: &mut B)
     where
         B: BufMut,
@@ -362,6 +390,9 @@ impl Message for Vec<u8> {
 
 /// `google.protobuf.BytesValue`
 impl Message for Bytes {
+    fn recursion_limit() -> u32 {
+        0
+    }
     fn encode_raw<B>(&self, buf: &mut B)
     where
         B: BufMut,
@@ -400,6 +431,9 @@ impl Message for Bytes {
 
 /// `google.protobuf.Empty`
 impl Message for () {
+    fn recursion_limit() -> u32 {
+        0
+    }
     fn encode_raw<B>(&self, _buf: &mut B)
     where
         B: BufMut,

tests/src/build.rs

@@ -54,6 +54,7 @@ fn main() {
         .compile_protos(&[src.join("ident_conversion.proto")], includes)
         .unwrap();
 
+    config.recursion_limit("nesting.E", 200);
     config
         .compile_protos(&[src.join("nesting.proto")], includes)
         .unwrap();

tests/src/lib.rs

@@ -405,6 +405,27 @@ mod tests {
         assert!(build_and_roundtrip(101).is_err());
     }
 
+    #[test]
+    fn test_deep_nesting_with_custom_recursion_limit() {
+        fn build_and_roundtrip(depth: usize) -> Result<(), prost::DecodeError> {
+            use crate::nesting::E;
+
+            let mut e = Box::new(E::default());
+            for _ in 0..depth {
+                let mut next = Box::new(E::default());
+                next.e = Some(e);
+                e = next;
+            }
+
+            let mut buf = Vec::new();
+            e.encode(&mut buf).unwrap();
+            E::decode(&*buf).map(|_| ())
+        }
+
+        assert!(build_and_roundtrip(200).is_ok());
+        assert!(build_and_roundtrip(201).is_err());
+    }
+
     #[test]
     fn test_deep_nesting_oneof() {
         fn build_and_roundtrip(depth: usize) -> Result<(), prost::DecodeError> {

tests/src/nesting.proto

@@ -23,3 +23,13 @@ message C {
 message D {
     map<string, D> m = 1;
 }
+
+message E {
+    E e = 1;
+    repeated E repeated_e = 2;
+    map<int32, E> map_e = 3;
+
+    B b = 4;
+    repeated B repeated_b = 5;
+    map<int32, B> map_b = 6;
+}