Support custom recursion limits at build time

seanlinsley · seanlinsley · commit 7fc66894f017 · 2025-05-29T12:44:28.000-05:00
diff --git a/README.md b/README.md
@@ -467,7 +467,7 @@ configured with the required dependencies to compile the whole project.
 - `std`: Enable integration with standard library. Disable this feature for `no_std` support. This feature is enabled by default.
 - `derive`: Enable integration with `prost-derive`. Disable this feature to reduce compile times. This feature is enabled by default.
 - `prost-derive`: Deprecated. Alias for `derive` feature.
-- `no-recursion-limit`: Disable the recursion limit. The recursion limit is 100 and cannot be customized. 
+- `no-recursion-limit`: Disable the recursion limit. The recursion limit is 100, and can be changed with `prost_build::recursion_limit`.
 
 ## FAQ
 
diff --git a/prost-build/src/code_generator.rs b/prost-build/src/code_generator.rs
@@ -237,6 +237,7 @@ impl<'b> CodeGenerator<'_, 'b> {
             self.context.prost_path()
         ));
         self.append_skip_debug(&fq_message_name);
+        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));
@@ -377,6 +378,21 @@ impl<'b> CodeGenerator<'_, 'b> {
         }
     }
 
+    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)
+            .cloned()
+        {
+            push_indent(self.buf, self.depth);
+            self.buf
+                .push_str(&format!("#[prost(recursion_limit={})]", 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.context.field_attributes(fq_message_name, field_name) {
diff --git a/prost-build/src/config.rs b/prost-build/src/config.rs
@@ -50,6 +50,7 @@ pub struct Config {
     pub(crate) skip_debug: PathMap<()>,
     pub(crate) skip_protoc_run: bool,
     pub(crate) skip_source_info: bool,
+    pub(crate) recursion_limits: PathMap<u32>,
     pub(crate) include_file: Option<PathBuf>,
     pub(crate) prost_path: Option<String>,
     #[cfg(feature = "format")]
@@ -1030,6 +1031,25 @@ impl Config {
         self.compile_fds(file_descriptor_set)
     }
 
+    /// 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
+    }
+
     pub(crate) fn write_includes(
         &self,
         mut modules: Vec<&Module>,
@@ -1192,6 +1212,7 @@ impl default::Default for Config {
             skip_debug: PathMap::default(),
             skip_protoc_run: false,
             skip_source_info: false,
+            recursion_limits: PathMap::default(),
             include_file: None,
             prost_path: None,
             #[cfg(feature = "format")]
diff --git a/prost-derive/src/field/mod.rs b/prost-derive/src/field/mod.rs
@@ -225,7 +225,7 @@ impl fmt::Display for Label {
 }
 
 /// Get the items belonging to the 'prost' list attribute, e.g. `#[prost(foo, bar="baz")]`.
-fn prost_attrs(attrs: Vec<Attribute>) -> Result<Vec<Meta>, Error> {
+pub(super) fn prost_attrs(attrs: Vec<Attribute>) -> Result<Vec<Meta>, Error> {
     let mut result = Vec::new();
     for attr in attrs.iter() {
         if let Meta::List(meta_list) = &attr.meta {
diff --git a/prost-derive/src/lib.rs b/prost-derive/src/lib.rs
@@ -11,23 +11,36 @@ use proc_macro2::{Span, TokenStream};
 use quote::quote;
 use syn::{
     punctuated::Punctuated, Data, DataEnum, DataStruct, DeriveInput, Expr, Fields, FieldsNamed,
-    FieldsUnnamed, Ident, Index, Variant,
+    FieldsUnnamed, Ident, Index, Lit, Meta, Variant,
 };
 
 mod field;
 use crate::field::Field;
 
 fn try_message(input: TokenStream) -> Result<TokenStream, Error> {
     let input: DeriveInput = syn::parse2(input)?;
-
     let ident = input.ident;
 
     syn::custom_keyword!(skip_debug);
     let skip_debug = input
         .attrs
-        .into_iter()
+        .iter()
         .any(|a| a.path().is_ident("prost") && a.parse_args::<skip_debug>().is_ok());
 
+    let mut recursion_limit: u32 = 100;
+    for attr in field::prost_attrs(input.attrs.clone())? {
+        match attr {
+            Meta::NameValue(ref meta) if meta.path.is_ident("recursion_limit") => {
+                let Expr::Lit(ref lit) = meta.value else {
+                    continue;
+                };
+                let Lit::Int(ref int) = lit.lit else { continue };
+                recursion_limit = int.base10_parse()?;
+            }
+            _ => (),
+        }
+    }
+
     let variant_data = match input.data {
         Data::Struct(variant_data) => variant_data,
         Data::Enum(..) => bail!("Message can not be derived for an enum"),
@@ -203,6 +216,10 @@ fn try_message(input: TokenStream) -> Result<TokenStream, Error> {
             fn clear(&mut self) {
                 #(#clear;)*
             }
+
+            fn recursion_limit() -> u32 {
+                #recursion_limit
+            }
         }
 
         impl #impl_generics ::core::default::Default for #ident #ty_generics #where_clause {
diff --git a/prost/README.md b/prost/README.md
@@ -467,7 +467,7 @@ configured with the required dependencies to compile the whole project.
 - `std`: Enable integration with standard library. Disable this feature for `no_std` support. This feature is enabled by default.
 - `derive`: Enable integration with `prost-derive`. Disable this feature to reduce compile times. This feature is enabled by default.
 - `prost-derive`: Deprecated. Alias for `derive` feature.
-- `no-recursion-limit`: Disable the recursion limit. The recursion limit is 100 and cannot be customized. 
+- `no-recursion-limit`: Disable the recursion limit. The recursion limit is 100, and can be changed with `prost_build::recursion_limit`.
 
 ## FAQ
 
diff --git a/prost/src/encoding.rs b/prost/src/encoding.rs
@@ -37,29 +37,27 @@ 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,
+    #[doc(hidden)]
+    pub recurse_count: u32,
 }
 
-#[cfg(not(feature = "no-recursion-limit"))]
-impl Default for DecodeContext {
-    #[inline]
-    fn default() -> DecodeContext {
+impl DecodeContext {
+    #[allow(unused_variables)]
+    pub 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 {
@@ -1225,7 +1223,7 @@ mod test {
             wire_type,
             &mut roundtrip_value,
             &mut buf,
-            DecodeContext::default(),
+            DecodeContext::new(100),
         )
         .map_err(|error| TestCaseError::fail(error.to_string()))?;
 
@@ -1297,7 +1295,7 @@ mod test {
                 wire_type,
                 &mut roundtrip_value,
                 &mut buf,
-                DecodeContext::default(),
+                DecodeContext::new(100),
             )
             .map_err(|error| TestCaseError::fail(error.to_string()))?;
         }
@@ -1316,7 +1314,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());
diff --git a/prost/src/lib.rs b/prost/src/lib.rs
@@ -24,11 +24,6 @@ pub use crate::error::{DecodeError, EncodeError, UnknownEnumValue};
 pub use crate::message::Message;
 pub use crate::name::Name;
 
-// 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;
-
 // Re-export #[derive(Message, Enumeration, Oneof)].
 // Based on serde's equivalent re-export [1], but enabled by default.
 //
diff --git a/prost/src/message.rs b/prost/src/message.rs
@@ -127,7 +127,7 @@ pub trait Message: Send + Sync {
     where
         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())?;
@@ -145,12 +145,22 @@ pub trait Message: Send + Sync {
             WireType::LengthDelimited,
             self,
             &mut buf,
-            DecodeContext::default(),
+            DecodeContext::new(Self::recursion_limit()),
         )
     }
 
     /// Clears the message, resetting all fields to their default.
     fn clear(&mut self);
+
+    /// 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,
+    {
+        100
+    }
 }
 
 impl<M> Message for Box<M>
@@ -175,6 +185,9 @@ where
     fn clear(&mut self) {
         (**self).clear()
     }
+    fn recursion_limit() -> u32 {
+        M::recursion_limit()
+    }
 }
 
 #[cfg(test)]
diff --git a/tests/build.rs b/tests/build.rs
@@ -48,6 +48,7 @@ fn main() {
 
     prost_build::Config::new()
         .btree_map(["."])
+        .recursion_limit("nesting.E", 200)
         .compile_protos(&[src.join("nesting.proto")], includes)
         .unwrap();
 
diff --git a/tests/src/decode_error.rs b/tests/src/decode_error.rs
@@ -98,7 +98,7 @@ fn test_decode_error_invalid_tag() {
 fn test_decode_error_unexpected_wire_type() {
     let mut buf = [0x00].as_slice();
     let mut msg = TestAllTypesProto3::default();
-    let ctx = prost::encoding::DecodeContext::default();
+    let ctx = prost::encoding::DecodeContext::new(100);
     assert_eq!(
         msg.merge_field(1, prost::encoding::WireType::LengthDelimited, &mut buf, ctx).unwrap_err().to_string(),
         "failed to decode Protobuf message: TestAllTypesProto3.optional_int32: invalid wire type: LengthDelimited (expected Varint)"
diff --git a/tests/src/nesting.proto b/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;
+}
diff --git a/tests/src/nesting.rs b/tests/src/nesting.rs
@@ -74,3 +74,22 @@ fn test_deep_nesting_map() {
     assert!(build_and_roundtrip(50).is_ok());
     assert!(build_and_roundtrip(51).is_err());
 }
+
+#[test]
+fn test_deep_nesting_with_custom_recursion_limit() {
+    fn build_and_roundtrip(depth: usize) -> Result<(), prost::DecodeError> {
+        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());
+}

Original file line number	Diff line number	Diff line change
`@@ -225,7 +225,7 @@ impl fmt::Display for Label {`
`225`	`225`	`}`
`226`	`226`
`227`	`227`	/// Get the items belonging to the 'prost' list attribute, e.g. `#[prost(foo, bar="baz")]`.
`228`		`-fn prost_attrs(attrs: Vec<Attribute>) -> Result<Vec<Meta>, Error> {`
	`228`	`+pub(super) fn prost_attrs(attrs: Vec<Attribute>) -> Result<Vec<Meta>, Error> {`
`229`	`229`	`let mut result = Vec::new();`
`230`	`230`	`for attr in attrs.iter() {`
`231`	`231`	`if let Meta::List(meta_list) = &attr.meta {`