1// SPDX-License-Identifier: AGPL-3.0-or-later
2
3package web
4
5import (
6 "context"
7 "strings"
8 "time"
9
10 "github.com/alecthomas/chroma/v2"
11 "github.com/alecthomas/chroma/v2/lexers"
12 "github.com/yuin/goldmark/ast"
13 "github.com/yuin/goldmark/renderer"
14 goldhtml "github.com/yuin/goldmark/renderer/html"
15 "github.com/yuin/goldmark/util"
16)
17
18const (
19 markdownBlockBytes = 64 << 10
20 markdownDocumentBytes = 256 << 10
21 markdownLexTime = 200 * time.Millisecond
22 markdownBudgetKey = "markdown-lex-budget"
23)
24
25// Budget belongs to the parsed document, never to the shared renderer.
26// Charge attempts, including failures, so repeated bad blocks cannot bypass it.
27type markdownLexBudget struct {
28 bytes, blocks int
29 deadline time.Time
30}
31
32type fencedRenderer struct{}
33
34// Capture goldmark's default renderer rather than duplicate its escaping rules.
35type fencedFallback struct{ render renderer.NodeRendererFunc }
36
37func (f *fencedFallback) Register(k ast.NodeKind, fn renderer.NodeRendererFunc) {
38 if k == ast.KindFencedCodeBlock {
39 f.render = fn
40 }
41}
42
43func (fencedRenderer) RegisterFuncs(r renderer.NodeRendererFuncRegisterer) {
44 var fallback fencedFallback
45 goldhtml.NewRenderer().RegisterFuncs(&fallback)
46 r.Register(ast.KindFencedCodeBlock, func(w util.BufWriter, source []byte, node ast.Node, entering bool) (ast.WalkStatus, error) {
47 n := node.(*ast.FencedCodeBlock)
48 if !entering {
49 return ast.WalkContinue, nil
50 }
51 toks, lang, ok := markdownTokens(n, source)
52 if !ok {
53 // Default rendering is plain escaped text. Omit unsafe language classes.
54 if !safeMarkdownLanguage(string(n.Language(source))) {
55 plain := ast.NewFencedCodeBlock(nil)
56 plain.SetLines(n.Lines())
57 n = plain
58 }
59 _, err := fallback.render(w, source, n, true)
60 if err != nil {
61 return ast.WalkStop, err
62 }
63 return fallback.render(w, source, n, false)
64 }
65 _, _ = w.WriteString(`<pre class="chroma"><code class="language-` + lang + `">`)
66 // EnsureNL lexers may append a newline: never emit bytes not in the block.
67 remaining := 0
68 for i := 0; i < n.Lines().Len(); i++ {
69 line := n.Lines().At(i)
70 remaining += len(line.Value(source))
71 }
72 // The tokens are not coalesced (chroma's coalescer copies strings
73 // quadratically): adjacent tokens of one class share a span here.
74 open := ""
75 for _, tk := range toks {
76 v := tk.Value
77 if len(v) > remaining {
78 v = v[:remaining]
79 }
80 remaining -= len(v)
81 if v == "" {
82 continue
83 }
84 if c := tokenClass(tk.Type); c != open {
85 if open != "" {
86 _, _ = w.WriteString(`</span>`)
87 }
88 if c != "" {
89 _, _ = w.WriteString(`<span class="` + c + `">`)
90 }
91 open = c
92 }
93 writeMarked(w, v)
94 }
95 if open != "" {
96 _, _ = w.WriteString(`</span>`)
97 }
98 _, _ = w.WriteString("</code></pre>\n")
99 return ast.WalkSkipChildren, nil
100 })
101}
102
103func safeMarkdownLanguage(s string) bool {
104 if s == "" {
105 return false
106 }
107 for _, c := range s {
108 switch {
109 case c >= 'a' && c <= 'z', c >= 'A' && c <= 'Z', c >= '0' && c <= '9', strings.ContainsRune("_+.#-", c):
110 default:
111 return false
112 }
113 }
114 return true
115}
116
117func markdownTokens(n *ast.FencedCodeBlock, source []byte) ([]chroma.Token, string, bool) {
118 var lang string
119 if n.Info != nil {
120 if words := strings.Fields(string(n.Info.Segment.Value(source))); len(words) > 0 {
121 lang = words[0]
122 }
123 }
124 if !safeMarkdownLanguage(lang) {
125 return nil, "", false
126 }
127 value, _ := n.OwnerDocument().AttributeString(markdownBudgetKey)
128 budget, ok := value.(*markdownLexBudget)
129 if !ok || budget.blocks >= 128 || time.Now().After(budget.deadline) {
130 return nil, "", false
131 }
132 size := 0
133 for i := 0; i < n.Lines().Len(); i++ {
134 line := n.Lines().At(i)
135 size += len(line.Value(source))
136 if size > markdownBlockBytes {
137 return nil, "", false
138 }
139 }
140 if size > markdownDocumentBytes-budget.bytes {
141 return nil, "", false
142 }
143 budget.blocks++
144 budget.bytes += size
145 lx := lexers.Get(lang)
146 if lx == nil {
147 return nil, "", false
148 }
149 select {
150 case lexSlots <- struct{}{}:
151 defer func() { <-lexSlots }()
152 default:
153 return nil, "", false
154 }
155 var code strings.Builder
156 code.Grow(size)
157 for i := 0; i < n.Lines().Len(); i++ {
158 line := n.Lines().At(i)
159 code.Write(line.Value(source))
160 }
161 ctx, cancel := context.WithDeadline(context.Background(), budget.deadline)
162 defer cancel()
163 toks, ok := tokenise(ctx, lx, code.String())
164 return toks, lang, ok
165}