Allow language classes

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)",

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"

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\">",
+        );
+    }
 }