From 1f5bfeb7d7060e126205d54de968feb07183f8a5 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ren=C3=A9=20Kijewski?= <rene.kijewski@fu-berlin.de>
Date: Sun, 7 Jul 2024 10:09:14 +0200
Subject: [PATCH] Add `Safe` and `MaybeSafe` wrappers for custom filters

Also a documentation that there is no `Unsafe` wrapper, because that's
the default anyway.
---
 rinja/src/filters/escape.rs | 102 ++++++++++++++++++++++++++++++++++++
 rinja/src/filters/mod.rs    |   5 +-
 2 files changed, 106 insertions(+), 1 deletion(-)

diff --git a/rinja/src/filters/escape.rs b/rinja/src/filters/escape.rs
index 9fcca7a1..faa88401 100644
--- a/rinja/src/filters/escape.rs
+++ b/rinja/src/filters/escape.rs
@@ -188,6 +188,108 @@ impl<'a, T: HtmlSafeMarker + ?Sized> AutoEscape for &AutoEscaper<'a, T, Html> {
     }
 }
 
+/// Mark the output of a filter as "maybe safe"
+///
+/// This struct can be used as a transparent return type of custom filters that want to mark
+/// their output as "safe" depending on some circumstances, i.e. that their output maybe does not
+/// need to be escaped.
+///
+/// If the filter is not used as the last element in the filter chain, then any assumption is void.
+/// Let the next filter decide if the output is safe or not.
+pub struct MaybeSafe<T: fmt::Display> {
+    pub text: T,
+    pub safe: bool,
+}
+
+const _: () = {
+    // This is the fallback. The filter is not the last element of the filter chain.
+    impl<T: fmt::Display> fmt::Display for MaybeSafe<T> {
+        #[inline]
+        fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+            write!(f, "{}", self.text)
+        }
+    }
+
+    impl<'a, T: fmt::Display, E: Escaper> AutoEscape for AutoEscaper<'a, MaybeSafe<T>, E> {
+        type Escaped = Wrapped<'a, T, E>;
+        type Error = Infallible;
+
+        #[inline]
+        fn rinja_auto_escape(&self) -> Result<Self::Escaped, Self::Error> {
+            match self.text.safe {
+                true => Ok(Wrapped::Safe(&self.text.text)),
+                false => Ok(Wrapped::Unsafe(&self.text.text, self.escaper)),
+            }
+        }
+    }
+
+    pub enum Wrapped<'a, T: fmt::Display + ?Sized, E: Escaper> {
+        Unsafe(&'a T, E),
+        Safe(&'a T),
+    }
+
+    impl<T: fmt::Display + ?Sized, E: Escaper> fmt::Display for Wrapped<'_, T, E> {
+        fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+            match *self {
+                Self::Unsafe(text, escaper) => EscapeDisplay(text, escaper).fmt(f),
+                Self::Safe(text) => write!(f, "{text}"),
+            }
+        }
+    }
+};
+
+/// Mark the output of a filter as "safe"
+///
+/// This struct can be used as a transparent return type of custom filters that want to mark their
+/// output as "safe" no matter what, i.e. that their output does not need to be escaped.
+///
+/// If the filter is not used as the last element in the filter chain, then any assumption is void.
+/// Let the next filter decide if the output is safe or not.
+pub struct Safe<T: fmt::Display>(pub T);
+
+const _: () = {
+    // This is the fallback. The filter is not the last element of the filter chain.
+    impl<T: fmt::Display> fmt::Display for Safe<T> {
+        #[inline]
+        fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+            write!(f, "{}", self.0)
+        }
+    }
+
+    impl<'a, T: fmt::Display, E: Escaper> AutoEscape for AutoEscaper<'a, Safe<T>, E> {
+        type Escaped = &'a T;
+        type Error = Infallible;
+
+        #[inline]
+        fn rinja_auto_escape(&self) -> Result<Self::Escaped, Self::Error> {
+            Ok(&self.text.0)
+        }
+    }
+};
+
+/// There is not need to mark the output of a custom filter as "unsafe"; this is simply the default
+pub struct Unsafe<T: fmt::Display>(pub T);
+
+const _: () = {
+    // This is the fallback. The filter is not the last element of the filter chain.
+    impl<T: fmt::Display> fmt::Display for Unsafe<T> {
+        #[inline]
+        fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
+            write!(f, "{}", self.0)
+        }
+    }
+
+    impl<'a, T: fmt::Display, E: Escaper> AutoEscape for AutoEscaper<'a, Unsafe<T>, E> {
+        type Escaped = EscapeDisplay<&'a T, E>;
+        type Error = Infallible;
+
+        #[inline]
+        fn rinja_auto_escape(&self) -> Result<Self::Escaped, Self::Error> {
+            Ok(EscapeDisplay(&self.text.0, self.escaper))
+        }
+    }
+};
+
 macro_rules! mark_html_safe {
     ($($ty:ty),* $(,)?) => {$(
         impl HtmlSafeMarker for $ty {}
diff --git a/rinja/src/filters/mod.rs b/rinja/src/filters/mod.rs
index 3504d103..c25de5c5 100644
--- a/rinja/src/filters/mod.rs
+++ b/rinja/src/filters/mod.rs
@@ -11,7 +11,10 @@ use std::cell::Cell;
 use std::convert::Infallible;
 use std::fmt::{self, Write};
 
-pub use escape::{e, escape, safe, AutoEscape, AutoEscaper, Escaper, Html, HtmlSafeMarker, Text};
+pub use escape::{
+    e, escape, safe, AutoEscape, AutoEscaper, Escaper, Html, HtmlSafeMarker, MaybeSafe, Safe, Text,
+    Unsafe,
+};
 #[cfg(feature = "humansize")]
 use humansize::{ISizeFormatter, ToF64, DECIMAL};
 #[cfg(feature = "serde_json")]