2 files changed, 52 insertions, 0 deletions
diff --git a/app/assets/javascripts/vue_shared/components/source_viewer/plugins/index.js b/app/assets/javascripts/vue_shared/components/source_viewer/plugins/index.js
new file mode 100644
index 00000000000..c9f7e5508be
--- /dev/null
+++ b/app/assets/javascripts/vue_shared/components/source_viewer/plugins/index.js
@@ -0,0 +1,13 @@
+import { HLJS_ON_AFTER_HIGHLIGHT } from '../constants';
+import wrapComments from './wrap_comments';
+
+/**
+ * Registers our plugins for Highlight.js
+ *
+ * Plugin API: https://github.com/highlightjs/highlight.js/blob/main/docs/plugin-api.rst
+ *
+ * @param {Object} hljs - the Highlight.js instance.
+ */
+export const registerPlugins = (hljs) => {
+  hljs.addPlugin({ [HLJS_ON_AFTER_HIGHLIGHT]: wrapComments });
+};
diff --git a/app/assets/javascripts/vue_shared/components/source_viewer/plugins/wrap_comments.js b/app/assets/javascripts/vue_shared/components/source_viewer/plugins/wrap_comments.js
new file mode 100644
index 00000000000..5be92af5b55
--- /dev/null
+++ b/app/assets/javascripts/vue_shared/components/source_viewer/plugins/wrap_comments.js
@@ -0,0 +1,39 @@
+import { HLJS_COMMENT_SELECTOR } from '../constants';
+
+const createWrapper = (content) => {
+  const span = document.createElement('span');
+  span.className = HLJS_COMMENT_SELECTOR;
+  span.innerHTML = content;
+  return span.outerHTML;
+};
+
+/**
+ * Highlight.js plugin for wrapping multi-line comments in the `hljs-comment` class.
+ * This ensures that multi-line comments are rendered correctly in the GitLab UI.
+ *
+ * Plugin API: https://github.com/highlightjs/highlight.js/blob/main/docs/plugin-api.rst
+ *
+ * @param {Object} Result - an object that represents the highlighted result from Highlight.js
+ */
+export default (result) => {
+  if (!result.value.includes(HLJS_COMMENT_SELECTOR)) return;
+
+  let wrapComment = false;
+
+  // eslint-disable-next-line no-param-reassign
+  result.value = result.value // Highlight.js expects the result param to be mutated for plugins to work
+    .split('\n')
+    .map((lineContent) => {
+      const includesClosingTag = lineContent.includes('</span>');
+      if (lineContent.includes(HLJS_COMMENT_SELECTOR) && !includesClosingTag) {
+        wrapComment = true;
+        return lineContent;
+      }
+      const line = wrapComment ? createWrapper(lineContent) : lineContent;
+      if (includesClosingTag) {
+        wrapComment = false;
+      }
+      return line;
+    })
+    .join('\n');
+};