Prevent infinite recursion in templates

Templates are limited to 10 recursive calls total now.

CHANGELOG.md

@@ -8,6 +8,11 @@
   - Any profile values that were previously the "raw" variant (the default) that contain template syntax (e.g. `{{user_id}}`) will now be rendered as templates. In reality this is very unlikely, so this probably isn't going to break your setup
   - If you have an existing profile value tagged with `!template` it **won't** break, but it will no longer do anything
 
+### Changed
+
+- Prevent infinite recursion in templates
+  - It now triggers a helpful error instead of a panic
+
 ## [0.17.0] - 2024-04-08
 
 ### Breaking

src/cli/request.rs

@@ -197,6 +197,7 @@ impl BuildRequestCommand {
             database,
             overrides,
             prompter: Box::new(CliPrompter),
+            recursion_count: Default::default(),
         };
         let request = RequestBuilder::new(recipe, RecipeOptions::default())
             .build(&template_context)

src/template.rs

@@ -19,7 +19,10 @@ use crate::{
 use derive_more::{Deref, Display};
 use indexmap::IndexMap;
 use serde::{Deserialize, Serialize};
-use std::fmt::Debug;
+use std::{fmt::Debug, sync::atomic::AtomicU8};
+
+/// Maximum number of layers of nested templates
+const RECURSION_LIMIT: u8 = 10;
 
 /// A little container struct for all the data that the user can access via
 /// templating. Unfortunately this has to own all data so templating can be
@@ -44,6 +47,15 @@ pub struct TemplateContext {
     pub overrides: IndexMap<String, String>,
     /// A conduit to ask the user questions
     pub prompter: Box<dyn Prompter>,
+    /// A count of how many templates have *already* been rendered with this
+    /// context. This is used to prevent infinite recursion in templates. For
+    /// all external calls, you can start this at 0.
+    ///
+    /// This tracks the *total* number of recursive calls in a render tree, not
+    /// the number of *layers*. That means one template that renders 5 child
+    /// templates is the same as a template that renders a single child 5
+    /// times.
+    pub recursion_count: AtomicU8,
 }
 
 /// A immutable string that can contain templated content. The string is parsed
@@ -259,11 +271,18 @@ mod tests {
     }
 
     /// Potential error cases for a profile field
+    #[rstest]
+    #[case("{{onion_id}}", "Unknown field `onion_id`")]
+    #[case(
+        "{{nested}}",
+        "Error in nested template `{{onion_id}}`: Unknown field `onion_id`"
+    )]
+    #[case("{{recursive}}", "Template recursion limit reached")]
     #[tokio::test]
-    async fn test_field_error() {
-        let nested_template = Template::parse("{{onion_id}}".into()).unwrap();
+    async fn test_field_error(#[case] template: &str, #[case] expected: &str) {
         let profile_data = indexmap! {
-            "recursive".into() => nested_template,
+            "nested".into() => Template::parse("{{onion_id}}".into()).unwrap(),
+            "recursive".into() => Template::parse("{{recursive}}".into()).unwrap(),
         };
         let profile = create!(Profile, data: profile_data);
         let profile_id = profile.id.clone();
@@ -275,15 +294,7 @@ mod tests {
             ),
             selected_profile: Some(profile_id),
         );
-
-        assert_err!(
-            render!("{{onion_id}}", context),
-            "Unknown field `onion_id`"
-        );
-        assert_err!(
-            render!("{{recursive}}", context),
-            "Error in nested template `{{onion_id}}`: Unknown field `onion_id`"
-        );
+        assert_err!(render!(template, context), expected);
     }
 
     /// Test success cases with chained responses

src/template/error.rs

@@ -1,7 +1,7 @@
 use crate::{
     collection::{ChainId, ProfileId, RecipeId},
     http::{QueryError, RequestBuildError, RequestError},
-    template::Template,
+    template::{Template, RECURSION_LIMIT},
     util::doc_link,
 };
 use nom::error::VerboseError;
@@ -51,6 +51,13 @@ pub enum TemplateError {
         error: Box<Self>,
     },
 
+    /// Too many templates!
+    #[error(
+        "Template recursion limit reached; cannot render more than \
+        {RECURSION_LIMIT} nested templates"
+    )]
+    RecursionLimit,
+
     #[error("Error resolving chain `{chain_id}`")]
     Chain {
         chain_id: ChainId,

src/template/render.rs

@@ -6,14 +6,18 @@ use crate::{
     template::{
         error::TriggeredRequestError, parse::TemplateInputChunk, ChainError,
         Prompt, Template, TemplateChunk, TemplateContext, TemplateError,
-        TemplateKey,
+        TemplateKey, RECURSION_LIMIT,
     },
     util::ResultExt,
 };
 use async_trait::async_trait;
 use chrono::Utc;
 use futures::future::join_all;
-use std::{env, path::Path, sync::Arc};
+use std::{
+    env,
+    path::Path,
+    sync::{atomic::Ordering, Arc},
+};
 use tokio::{fs, process::Command, sync::oneshot};
 use tracing::{info, instrument, trace};
 
@@ -122,6 +126,10 @@ impl Template {
         &self,
         context: &TemplateContext,
     ) -> Result<String, TemplateError> {
+        if context.recursion_count.load(Ordering::Relaxed) >= RECURSION_LIMIT {
+            return Err(TemplateError::RecursionLimit);
+        }
+
         // Render each individual template chunk in the string
         let chunks = self.render_chunks(context).await;
 
@@ -214,8 +222,8 @@ impl<'a> TemplateSource<'a> for FieldTemplateSource<'a> {
         })?;
 
         // recursion!
-        // TODO limit recursion
         trace!(%field, %template, "Rendering recursive template");
+        context.recursion_count.fetch_add(1, Ordering::Relaxed);
         let rendered =
             template.render_stitched(context).await.map_err(|error| {
                 TemplateError::Nested {

src/test_util.rs

@@ -102,6 +102,7 @@ factori!(TemplateContext, {
         http_engine = None,
         database = CollectionDatabase::testing(),
         overrides = Default::default(),
+        recursion_count = Default::default(),
     }
 });
 

src/tui.rs

@@ -543,6 +543,7 @@ impl Tui {
             database: self.database.clone(),
             overrides: Default::default(),
             prompter,
+            recursion_count: Default::default(),
         })
     }
 }