3 files changed, 42 insertions, 14 deletions
diff --git a/lib/banzai/filter/mermaid_filter.rb b/lib/banzai/filter/mermaid_filter.rb
new file mode 100644
index 00000000000..b545b947a2c
--- /dev/null
+++ b/lib/banzai/filter/mermaid_filter.rb
@@ -0,0 +1,20 @@
+module Banzai
+  module Filter
+    class MermaidFilter < HTML::Pipeline::Filter
+      def call
+        doc.css('pre[lang="mermaid"]').add_class('mermaid')
+        doc.css('pre[lang="mermaid"]').add_class('js-render-mermaid')
+
+        # The `<code></code>` blocks are added in the lib/banzai/filter/syntax_highlight_filter.rb
+        # We want to keep context and consistency, so we the blocks are added for all filters.
+        # Details: https://gitlab.com/gitlab-org/gitlab-ce/merge_requests/15107/diffs?diff_id=7962900#note_45495859
+        doc.css('pre[lang="mermaid"]').each do |pre|
+          document = pre.at('code')
+          document.replace(document.content)
+        end
+
+        doc
+      end
+    end
+  end
+end
diff --git a/lib/banzai/filter/syntax_highlight_filter.rb b/lib/banzai/filter/syntax_highlight_filter.rb
index 7da565043d1..a79a0154846 100644
--- a/lib/banzai/filter/syntax_highlight_filter.rb
+++ b/lib/banzai/filter/syntax_highlight_filter.rb
@@ -14,23 +14,26 @@ module Banzai
       end
 
       def highlight_node(node)
-        language = node.attr('lang')
         code = node.text
-        css_classes = "code highlight"
-        lexer = lexer_for(language)
-        lang = lexer.tag
-
-        begin
-          code = Rouge::Formatters::HTMLGitlab.format(lex(lexer, code), tag: lang)
-
-          css_classes << " js-syntax-highlight #{lang}"
-        rescue
-          lang = nil
-          # Gracefully handle syntax highlighter bugs/errors to ensure
-          # users can still access an issue/comment/etc.
+        css_classes = 'code highlight js-syntax-highlight'
+        language = node.attr('lang')
+
+        if use_rouge?(language)
+          lexer = lexer_for(language)
+          language = lexer.tag
+
+          begin
+            code = Rouge::Formatters::HTMLGitlab.format(lex(lexer, code), tag: language)
+            css_classes << " #{language}"
+          rescue
+            # Gracefully handle syntax highlighter bugs/errors to ensure
+            # users can still access an issue/comment/etc.
+
+            language = nil
+          end
         end
 
-        highlighted = %(<pre class="#{css_classes}" lang="#{lang}" v-pre="true"><code>#{code}</code></pre>)
+        highlighted = %(<pre class="#{css_classes}" lang="#{language}" v-pre="true"><code>#{code}</code></pre>)
 
         # Extracted to a method to measure it
         replace_parent_pre_element(node, highlighted)
@@ -51,6 +54,10 @@ module Banzai
         # Replace the parent `pre` element with the entire highlighted block
         node.parent.replace(highlighted)
       end
+
+      def use_rouge?(language)
+        %w(math mermaid plantuml).exclude?(language)
+      end
     end
   end
 end
diff --git a/lib/banzai/pipeline/gfm_pipeline.rb b/lib/banzai/pipeline/gfm_pipeline.rb
index 3208abfc538..55874ad50a3 100644
--- a/lib/banzai/pipeline/gfm_pipeline.rb
+++ b/lib/banzai/pipeline/gfm_pipeline.rb
@@ -14,6 +14,7 @@ module Banzai
           Filter::SyntaxHighlightFilter,
 
           Filter::MathFilter,
+          Filter::MermaidFilter,
           Filter::UploadLinkFilter,
           Filter::VideoLinkFilter,
           Filter::ImageLazyLoadFilter,