allow renderers to mutate Page (and frame-src CSP)

3 files changed, 16 insertions, 12 deletions

diff --git a/src/get_routes.rs b/src/get_routes.rs
index b9ad933..e07050d 100644
--- a/src/get_routes.rs
+++ b/src/get_routes.rs
@@ -81,7 +81,7 @@ fn view_blob<C: Controller>(
     match from_utf8(blob.content()) {
         Ok(text) => {
             if let Some(renderer) = get_renderer(&ctx.path) {
-                page.body.push_str(&renderer(text));
+                renderer(text, &mut page);
             } else {
                 page.body
                     .push_str(&format!("<pre>{}</pre>", html_escape(text)));
diff --git a/src/main.rs b/src/main.rs
index 4d0c1b3..a5b7ce0 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -215,6 +215,7 @@ async fn service<C: Controller>(
     let (mut parts, body) = request.into_parts();
 
     let mut script_csp = "'none'".into();
+    let mut frame_csp = "'none'".into();
 
     let mut resp = build_response(args, &*controller, &mut parts, body)
         .await
@@ -224,6 +225,9 @@ async fn service<C: Controller>(
                 if !page.script_src.is_empty() {
                     script_csp = page.script_src.join(" ");
                 }
+                if let Some(src) = page.frame_src {
+                    frame_csp = src;
+                }
                 Builder::new()
                     .content_type(mime::TEXT_HTML)
                     .body(render_page(&page, &*controller, &parts).into())
@@ -264,7 +268,8 @@ async fn service<C: Controller>(
     resp.headers_mut().insert(
         header::CONTENT_SECURITY_POLICY,
         format!(
-            "default-src 'self'; frame-src 'none'; script-src {}; style-src {}",
+            "default-src 'self'; frame-src {}; script-src {}; style-src {}",
+            frame_csp,
             script_csp,
             include_str!("static/style.css.sha"),
         )
@@ -283,6 +288,9 @@ pub struct Page {
     scripts: Vec<&'static str>,
     /// for the Content Security Policy
     script_src: Vec<&'static str>,
+
+    /// for the Content Security Policy
+    frame_src: Option<&'static str>,
 }
 
 fn render_page<C: Controller>(page: &Page, controller: &C, parts: &Parts) -> String {
@@ -518,16 +526,14 @@ impl Context {
     }
 }
 
-fn render_markdown(input: &str) -> String {
+fn render_markdown(input: &str, page: &mut Page) {
     let parser = Parser::new_ext(input, Options::all());
-    let mut out = String::new();
-    out.push_str("<div class=markdown-output>");
-    html::push_html(&mut out, parser);
-    out.push_str("</div>");
-    out
+    page.body.push_str("<div class=markdown-output>");
+    html::push_html(&mut page.body, parser);
+    page.body.push_str("</div>");
 }
 
-fn get_renderer(path: &Path) -> Option<fn(&str) -> String> {
+fn get_renderer(path: &Path) -> Option<fn(&str, &mut Page)> {
     match path.extension().map(|e| e.to_str().unwrap()) {
         Some("md") => Some(render_markdown),
         _ => None,
diff --git a/src/post_routes.rs b/src/post_routes.rs
index 8cbd3b2..392a024 100644
--- a/src/post_routes.rs
+++ b/src/post_routes.rs
@@ -363,8 +363,6 @@ async fn preview_edit<C: Controller>(
     let form: EditForm = body.into_form().await?;
     let new_text = form.text.replace("\r\n", "\n");
     let mut page = edit_text_form(&form, None, controller, &ctx, parts);
-
-    page.body
-        .push_str(&(get_renderer(&ctx.path).unwrap()(&new_text)));
+    get_renderer(&ctx.path).unwrap()(&new_text, &mut page);
     Ok(page.into())
 }