From f41a6d6574b470ffaf1b8ee8b344c2fa43e5e60b Mon Sep 17 00:00:00 2001
From: Konrad Borowski <konrad@borowski.pw>
Date: Mon, 4 Nov 2019 10:13:22 +0100
Subject: [PATCH] Allow language classes

---
 Cargo.lock          |  1 +
 Cargo.toml          |  1 +
 src/models/paste.rs | 32 ++++++++++++++++++++++++++++++++
 3 files changed, 34 insertions(+)

diff --git a/Cargo.lock b/Cargo.lock
index 3924911..8355b83 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1139,6 +1139,7 @@ dependencies = [
  "env_logger 0.6.2 (registry+https://github.com/rust-lang/crates.io-index)",
  "futures 0.1.29 (registry+https://github.com/rust-lang/crates.io-index)",
  "futures-preview 0.3.0-alpha.18 (registry+https://github.com/rust-lang/crates.io-index)",
+ "itertools 0.8.1 (registry+https://github.com/rust-lang/crates.io-index)",
  "lazy_static 1.4.0 (registry+https://github.com/rust-lang/crates.io-index)",
  "log 0.4.8 (registry+https://github.com/rust-lang/crates.io-index)",
  "mime 0.3.14 (registry+https://github.com/rust-lang/crates.io-index)",
diff --git a/Cargo.toml b/Cargo.toml
index 4476a4b..5043cc9 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -15,6 +15,7 @@ diesel_migrations = "1.4.0"
 env_logger = { version = "0.6.0", default-features = false }
 futures = "0.1.28"
 futures03 = { version = "=0.3.0-alpha.18", package = "futures-preview", features = ["compat"] }
+itertools = "0.8.1"
 lazy_static = "1.3.0"
 log = "0.4.6"
 mime = "0.3.13"
diff --git a/src/models/paste.rs b/src/models/paste.rs
index f09618c..631c42c 100644
--- a/src/models/paste.rs
+++ b/src/models/paste.rs
@@ -4,10 +4,12 @@ use crate::Connection;
 use ammonia::Builder;
 use chrono::{DateTime, Utc};
 use diesel::prelude::*;
+use itertools::Itertools;
 use lazy_static::lazy_static;
 use log::info;
 use pulldown_cmark::{Options, Parser};
 use rand::seq::SliceRandom;
+use std::iter;
 use warp::Rejection;
 
 #[derive(Queryable)]
@@ -106,6 +108,20 @@ fn render_markdown(markdown: &str) -> String {
         static ref FILTER: Builder<'static> = {
             let mut builder = Builder::new();
             builder.link_rel(Some("noopener noreferrer nofollow"));
+            builder.add_generic_attributes(iter::once("class"));
+            builder.attribute_filter(|_, attribute, value| {
+                if attribute == "class" {
+                    Some(
+                        value
+                            .split_ascii_whitespace()
+                            .filter(|value| value.starts_with("language-"))
+                            .join(" ")
+                            .into(),
+                    )
+                } else {
+                    Some(value.into())
+                }
+            });
             builder
         };
     }
@@ -131,4 +147,20 @@ mod test {
     fn strikethrough_works() {
         assert_eq!(render_markdown("~~strike~~"), "<p><del>strike</del></p>\n");
     }
+
+    #[test]
+    fn code_blocks_work() {
+        assert_eq!(
+            render_markdown("```rust\nfn main() {}\n```"),
+            "<pre><code class=\"language-rust\">fn main() {}\n</code></pre>\n",
+        );
+    }
+
+    #[test]
+    fn only_language_classes_are_allowed() {
+        assert_eq!(
+            render_markdown(r#"<br class="language-a madoka language-b homura">"#),
+            "<br class=\"language-a language-b\">",
+        );
+    }
 }