Prevent panic in Org mode HighlightCodeBlock (#17140)

When rendering source in org mode there is a mistake in the highlight code that
causes a panic.

This PR fixes this.

Fix #17139

Signed-off-by: Andrew Thornton <art27@cantab.net>

3 files changed, 46 insertions, 12 deletions

diff --git a/modules/highlight/highlight.go b/modules/highlight/highlight.go
index 079f7a44bd..6684fbe842 100644
--- a/modules/highlight/highlight.go
+++ b/modules/highlight/highlight.go
@@ -66,17 +66,6 @@ func Code(fileName, code string) string {
 	if len(code) > sizeLimit {
 		return code
 	}
-	formatter := html.New(html.WithClasses(true),
-		html.WithLineNumbers(false),
-		html.PreventSurroundingPre(true),
-	)
-	if formatter == nil {
-		log.Error("Couldn't create chroma formatter")
-		return code
-	}
-
-	htmlbuf := bytes.Buffer{}
-	htmlw := bufio.NewWriter(&htmlbuf)
 
 	var lexer chroma.Lexer
 	if val, ok := highlightMapping[filepath.Ext(fileName)]; ok {
@@ -97,6 +86,18 @@ func Code(fileName, code string) string {
 		}
 		cache.Add(fileName, lexer)
 	}
+	return CodeFromLexer(lexer, code)
+}
+
+// CodeFromLexer returns a HTML version of code string with chroma syntax highlighting classes
+func CodeFromLexer(lexer chroma.Lexer, code string) string {
+	formatter := html.New(html.WithClasses(true),
+		html.WithLineNumbers(false),
+		html.PreventSurroundingPre(true),
+	)
+
+	htmlbuf := bytes.Buffer{}
+	htmlw := bufio.NewWriter(&htmlbuf)
 
 	iterator, err := lexer.Tokenise(nil, string(code))
 	if err != nil {
diff --git a/modules/markup/orgmode/orgmode.go b/modules/markup/orgmode/orgmode.go
index 7e9f1f45c5..b035e04a1f 100644
--- a/modules/markup/orgmode/orgmode.go
+++ b/modules/markup/orgmode/orgmode.go
@@ -12,6 +12,7 @@ import (
 	"strings"
 
 	"code.gitea.io/gitea/modules/highlight"
+	"code.gitea.io/gitea/modules/log"
 	"code.gitea.io/gitea/modules/markup"
 	"code.gitea.io/gitea/modules/setting"
 	"code.gitea.io/gitea/modules/util"
@@ -51,6 +52,12 @@ func (Renderer) SanitizerRules() []setting.MarkupSanitizerRule {
 func Render(ctx *markup.RenderContext, input io.Reader, output io.Writer) error {
 	htmlWriter := org.NewHTMLWriter()
 	htmlWriter.HighlightCodeBlock = func(source, lang string, inline bool) string {
+		defer func() {
+			if err := recover(); err != nil {
+				log.Error("Panic in HighlightCodeBlock: %v\n%s", err, log.Stack(2))
+				panic(err)
+			}
+		}()
 		var w strings.Builder
 		if _, err := w.WriteString(`<pre>`); err != nil {
 			return ""
@@ -80,7 +87,7 @@ func Render(ctx *markup.RenderContext, input io.Reader, output io.Writer) error
 			}
 			lexer = chroma.Coalesce(lexer)
 
-			if _, err := w.WriteString(highlight.Code(lexer.Config().Filenames[0], source)); err != nil {
+			if _, err := w.WriteString(highlight.CodeFromLexer(lexer, source)); err != nil {
 				return ""
 			}
 		}
diff --git a/modules/markup/orgmode/orgmode_test.go b/modules/markup/orgmode/orgmode_test.go
index da89326e9e..81d0d66a76 100644
--- a/modules/markup/orgmode/orgmode_test.go
+++ b/modules/markup/orgmode/orgmode_test.go
@@ -57,3 +57,29 @@ func TestRender_Images(t *testing.T) {
 	test("[[file:"+url+"]]",
 		"<p><img src=\""+result+"\" alt=\""+result+"\" title=\""+result+"\" /></p>")
 }
+
+func TestRender_Source(t *testing.T) {
+	setting.AppURL = AppURL
+	setting.AppSubURL = AppSubURL
+
+	test := func(input, expected string) {
+		buffer, err := RenderString(&markup.RenderContext{
+			URLPrefix: setting.AppSubURL,
+		}, input)
+		assert.NoError(t, err)
+		assert.Equal(t, strings.TrimSpace(expected), strings.TrimSpace(buffer))
+	}
+
+	test(`#+begin_src go
+// HelloWorld prints "Hello World"
+func HelloWorld() {
+	fmt.Println("Hello World")
+}
+#+end_src
+`, `<div class="src src-go">
+<pre><code class="chroma language-go"><span class="c1">// HelloWorld prints &#34;Hello World&#34;
+</span><span class="c1"></span><span class="kd">func</span> <span class="nf">HelloWorld</span><span class="p">()</span> <span class="p">{</span>
+	<span class="nx">fmt</span><span class="p">.</span><span class="nf">Println</span><span class="p">(</span><span class="s">&#34;Hello World&#34;</span><span class="p">)</span>
+<span class="p">}</span></code></pre>
+</div>`)
+}