Escape all strings with character entities by default (fixes #23)

4 files changed, 107 insertions, 21 deletions

diff --git a/askama_shared/src/escaping.rs b/askama_shared/src/escaping.rs
index 54f58e4..ed4b3d7 100644
--- a/askama_shared/src/escaping.rs
+++ b/askama_shared/src/escaping.rs
@@ -1,3 +1,46 @@
+use std::fmt::{self, Display, Formatter};
+
+
+#[derive(Debug, PartialEq)]
+pub enum MarkupDisplay<T> where T: Display {
+    Safe(T),
+    Unsafe(T),
+}
+
+impl<T> MarkupDisplay<T> where T: Display {
+    pub fn mark_safe(self) -> MarkupDisplay<T> {
+        match self {
+            MarkupDisplay::Unsafe(t) => MarkupDisplay::Safe(t),
+            _ => { self },
+        }
+    }
+    pub fn unsafe_string(&self) -> String {
+        match *self {
+            MarkupDisplay::Safe(ref t) | MarkupDisplay::Unsafe(ref t) => format!("{}", t)
+        }
+    }
+}
+
+impl<T> From<T> for MarkupDisplay<T> where T: Display {
+    fn from(t: T) -> MarkupDisplay<T> {
+        MarkupDisplay::Unsafe(t)
+    }
+}
+
+impl<T> Display for MarkupDisplay<T> where T: Display {
+    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
+        match *self {
+            MarkupDisplay::Unsafe(_) => {
+                write!(f, "{}", escape(self.unsafe_string()))
+            },
+            MarkupDisplay::Safe(ref t) => {
+                t.fmt(f)
+            },
+        }
+    }
+}
+
+
 fn escapable(b: &u8) -> bool {
     *b == b'<' || *b == b'>' || *b == b'&'
 }
diff --git a/askama_shared/src/filters/mod.rs b/askama_shared/src/filters/mod.rs
index e6a5858..1aa7020 100644
--- a/askama_shared/src/filters/mod.rs
+++ b/askama_shared/src/filters/mod.rs
@@ -13,7 +13,7 @@ pub use self::json::json;
 
 use std::fmt;
 
-use escaping;
+use escaping::{self, MarkupDisplay};
 use super::Result;
 
 
@@ -21,12 +21,13 @@ use super::Result;
 // Askama or should refer to a local `filters` module. It should contain all the
 // filters shipped with Askama, even the optional ones (since optional inclusion
 // in the const vector based on features seems impossible right now).
-pub const BUILT_IN_FILTERS: [&str; 9] = [
+pub const BUILT_IN_FILTERS: [&str; 10] = [
     "e",
     "escape",
     "format",
     "lower",
     "lowercase",
+    "safe",
     "trim",
     "upper",
     "uppercase",
@@ -34,15 +35,32 @@ pub const BUILT_IN_FILTERS: [&str; 9] = [
 ];
 
 
+pub fn safe<D, I>(v: I) -> Result<MarkupDisplay<D>>
+where
+    D: fmt::Display,
+    MarkupDisplay<D>: From<I>
+{
+    let res: MarkupDisplay<D> = v.into();
+    Ok(res.mark_safe())
+}
+
 /// Escapes `&`, `<` and `>` in strings
-pub fn escape(s: &fmt::Display) -> Result<String> {
-    let s = format!("{}", s);
-    Ok(escaping::escape(s))
+pub fn escape<D, I>(i: I) -> Result<MarkupDisplay<String>>
+where
+    D: fmt::Display,
+    MarkupDisplay<D>: From<I>
+{
+    let md: MarkupDisplay<D> = i.into();
+    Ok(MarkupDisplay::Safe(escaping::escape(md.unsafe_string())))
 }
 
 /// Alias for the `escape()` filter
-pub fn e(s: &fmt::Display) -> Result<String> {
-    escape(s)
+pub fn e<D, I>(i: I) -> Result<MarkupDisplay<String>>
+where
+    D: fmt::Display,
+    MarkupDisplay<D>: From<I>
+{
+    escape(i)
 }
 
 /// Formats arguments according to the specified format
diff --git a/askama_shared/src/generator.rs b/askama_shared/src/generator.rs
index 5c4b7f8..a7fba18 100644
--- a/askama_shared/src/generator.rs
+++ b/askama_shared/src/generator.rs
@@ -190,31 +190,34 @@ impl<'a> Generator<'a> {
 
     /* Visitor methods for expression types */
 
-    fn visit_num_lit(&mut self, s: &str) {
+    fn visit_num_lit(&mut self, s: &str) -> DisplayWrap {
         self.write(s);
+        DisplayWrap::Unwrapped
     }
 
-    fn visit_str_lit(&mut self, s: &str) {
+    fn visit_str_lit(&mut self, s: &str) -> DisplayWrap {
         self.write(&format!("\"{}\"", s));
+        DisplayWrap::Unwrapped
     }
 
-    fn visit_var(&mut self, s: &str) {
+    fn visit_var(&mut self, s: &str) -> DisplayWrap {
         if self.locals.contains(s) {
             self.write(s);
         } else {
             self.write(&format!("self.{}", s));
         }
+        DisplayWrap::Unwrapped
     }
 
-    fn visit_attr(&mut self, obj: &Expr, attr: &str) {
+    fn visit_attr(&mut self, obj: &Expr, attr: &str) -> DisplayWrap {
         if let Expr::Var(name) = *obj {
             if name == "loop" {
                 self.write("_loop_index");
                 if attr == "index" {
                     self.write(" + 1");
-                    return;
+                    return DisplayWrap::Unwrapped;
                 } else if attr == "index0" {
-                    return;
+                    return DisplayWrap::Unwrapped;
                 } else {
                     panic!("unknown loop variable");
                 }
@@ -222,6 +225,7 @@ impl<'a> Generator<'a> {
         }
         self.visit_expr(obj);
         self.write(&format!(".{}", attr));
+        DisplayWrap::Unwrapped
     }
 
     fn _visit_filter_args(&mut self, args: &[Expr]) {
@@ -254,13 +258,13 @@ impl<'a> Generator<'a> {
         self.write(")?");
     }
 
-    fn visit_filter(&mut self, name: &str, args: &[Expr]) {
+    fn visit_filter(&mut self, name: &str, args: &[Expr]) -> DisplayWrap {
         if name == "format" {
             self._visit_format_filter(args);
-            return;
+            return DisplayWrap::Unwrapped;
         } else if name == "join" {
             self._visit_join_filter(args);
-            return;
+            return DisplayWrap::Unwrapped;
         }
 
         if filters::BUILT_IN_FILTERS.contains(&name) {
@@ -271,21 +275,28 @@ impl<'a> Generator<'a> {
 
         self._visit_filter_args(args);
         self.write(")?");
+        if name == "safe" || name == "escape" || name == "e" {
+            DisplayWrap::Wrapped
+        } else {
+            DisplayWrap::Unwrapped
+        }
     }
 
-    fn visit_binop(&mut self, op: &str, left: &Expr, right: &Expr) {
+    fn visit_binop(&mut self, op: &str, left: &Expr, right: &Expr) -> DisplayWrap {
         self.visit_expr(left);
         self.write(&format!(" {} ", op));
         self.visit_expr(right);
+        DisplayWrap::Unwrapped
     }
 
-    fn visit_group(&mut self, inner: &Expr) {
+    fn visit_group(&mut self, inner: &Expr) -> DisplayWrap {
         self.write("(");
         self.visit_expr(inner);
         self.write(")");
+        DisplayWrap::Unwrapped
     }
 
-    fn visit_method_call(&mut self, obj: &Expr, method: &str, args: &[Expr]) {
+    fn visit_method_call(&mut self, obj: &Expr, method: &str, args: &[Expr]) -> DisplayWrap {
         self.visit_expr(obj);
         self.write(&format!(".{}(", method));
         for (i, arg) in args.iter().enumerate() {
@@ -295,9 +306,10 @@ impl<'a> Generator<'a> {
             self.visit_expr(arg);
         }
         self.write(")");
+        DisplayWrap::Unwrapped
     }
 
-    fn visit_expr(&mut self, expr: &Expr) {
+    fn visit_expr(&mut self, expr: &Expr) -> DisplayWrap {
         match *expr {
             Expr::NumLit(s) => self.visit_num_lit(s),
             Expr::StrLit(s) => self.visit_str_lit(s),
@@ -347,8 +359,15 @@ impl<'a> Generator<'a> {
 
     fn write_expr(&mut self, ws: &WS, s: &Expr) {
         self.handle_ws(ws);
+        self.write("let askama_expr = &");
+        let wrapped = self.visit_expr(s);
+        self.writeln(";");
+
         self.write("writer.write_fmt(format_args!(\"{}\", ");
-        self.visit_expr(s);
+        self.write(match wrapped {
+            DisplayWrap::Wrapped => "askama_expr",
+            DisplayWrap::Unwrapped => "&::askama::MarkupDisplay::from(askama_expr)",
+        });
         self.writeln("))?;");
     }
 
@@ -729,4 +748,9 @@ impl<'a, T: 'a> SetChain<'a, T> where T: cmp::Eq + hash::Hash {
     }
 }
 
+enum DisplayWrap {
+    Wrapped,
+    Unwrapped,
+}
+
 type MacroMap<'a> = HashMap<&'a str, (WS, &'a str, Vec<&'a str>, Vec<Node<'a>>, WS)>;
diff --git a/askama_shared/src/lib.rs b/askama_shared/src/lib.rs
index 35e779c..b5bc9af 100644
--- a/askama_shared/src/lib.rs
+++ b/askama_shared/src/lib.rs
@@ -10,6 +10,7 @@ extern crate serde;
 #[cfg(feature = "serde-json")]
 extern crate serde_json;
 
+pub use escaping::MarkupDisplay;
 pub use errors::{Error, Result};
 pub mod filters;
 pub mod path;