From 3bfc08be9130fef1f07462d543d559f5c7f533a6 Mon Sep 17 00:00:00 2001
From: Mark Chao <mchao@gitlab.com>
Date: Thu, 16 Aug 2018 23:27:20 +0800
Subject: Fix syntax highlight taking too long

Use Timeout to switch to plain rendering
---
 .../unreleased/security-2697-code-highlight-timeout.yml |  5 +++++
 lib/gitlab/highlight.rb                                 | 14 +++++++++++++-
 spec/lib/gitlab/highlight_spec.rb                       | 17 +++++++++++++++++
 3 files changed, 35 insertions(+), 1 deletion(-)
 create mode 100644 changelogs/unreleased/security-2697-code-highlight-timeout.yml

diff --git a/changelogs/unreleased/security-2697-code-highlight-timeout.yml b/changelogs/unreleased/security-2697-code-highlight-timeout.yml
new file mode 100644
index 00000000000..66ad9ff822b
--- /dev/null
+++ b/changelogs/unreleased/security-2697-code-highlight-timeout.yml
@@ -0,0 +1,5 @@
+---
+title: Set timeout for syntax highlighting
+merge_request:
+author:
+type: security
diff --git a/lib/gitlab/highlight.rb b/lib/gitlab/highlight.rb
index 5408a1a6838..0b6cc893db1 100644
--- a/lib/gitlab/highlight.rb
+++ b/lib/gitlab/highlight.rb
@@ -1,5 +1,8 @@
 module Gitlab
   class Highlight
+    TIMEOUT_BACKGROUND = 30.seconds
+    TIMEOUT_FOREGROUND = 3.seconds
+
     def self.highlight(blob_name, blob_content, repository: nil, plain: false)
       new(blob_name, blob_content, repository: repository)
         .highlight(blob_content, continue: false, plain: plain)
@@ -51,11 +54,20 @@ module Gitlab
     end
 
     def highlight_rich(text, continue: true)
-      @formatter.format(lexer.lex(text, continue: continue), tag: lexer.tag).html_safe
+      tag = lexer.tag
+      tokens = lexer.lex(text, continue: continue)
+      Timeout.timeout(timeout_time) { @formatter.format(tokens, tag: tag).html_safe }
+    rescue Timeout::Error => e
+      Gitlab::Sentry.track_exception(e)
+      highlight_plain(text)
     rescue
       highlight_plain(text)
     end
 
+    def timeout_time
+      Sidekiq.server? ? TIMEOUT_BACKGROUND : TIMEOUT_FOREGROUND
+    end
+
     def link_dependencies(text, highlighted_text)
       Gitlab::DependencyLinker.link(blob_name, text, highlighted_text)
     end
diff --git a/spec/lib/gitlab/highlight_spec.rb b/spec/lib/gitlab/highlight_spec.rb
index 29e61d15726..88f7099ff3c 100644
--- a/spec/lib/gitlab/highlight_spec.rb
+++ b/spec/lib/gitlab/highlight_spec.rb
@@ -56,5 +56,22 @@ describe Gitlab::Highlight do
 
       described_class.highlight('file.name', 'Contents')
     end
+
+    context 'timeout' do
+      subject { described_class.new('file.name', 'Contents') }
+
+      it 'utilizes timeout for web' do
+        expect(Timeout).to receive(:timeout).with(described_class::TIMEOUT_FOREGROUND).and_call_original
+
+        subject.highlight("Content")
+      end
+
+      it 'utilizes longer timeout for sidekiq' do
+        allow(Sidekiq).to receive(:server?).and_return(true)
+        expect(Timeout).to receive(:timeout).with(described_class::TIMEOUT_BACKGROUND).and_call_original
+
+        subject.highlight("Content")
+      end
+    end
   end
 end
-- 
cgit v1.2.1