Refactor markup/csv: don't read all to memory (#29760)

(cherry picked from commit e79a807a8461a73bd66146d816f635b66e198c89)
This commit is contained in:
coldWater 2024-03-14 10:51:55 +08:00 committed by Earl Warren
parent 2da0628f18
commit d413a8fcac
No known key found for this signature in database
GPG key ID: 0579CB2928A78A00
2 changed files with 58 additions and 15 deletions

View file

@ -77,29 +77,62 @@ func writeField(w io.Writer, element, class, field string) error {
}
// Render implements markup.Renderer
func (Renderer) Render(ctx *markup.RenderContext, input io.Reader, output io.Writer) error {
func (r Renderer) Render(ctx *markup.RenderContext, input io.Reader, output io.Writer) error {
tmpBlock := bufio.NewWriter(output)
maxSize := setting.UI.CSV.MaxFileSize
// FIXME: don't read all to memory
rawBytes, err := io.ReadAll(input)
if maxSize == 0 {
return r.tableRender(ctx, input, tmpBlock)
}
rawBytes, err := io.ReadAll(io.LimitReader(input, maxSize+1))
if err != nil {
return err
}
if setting.UI.CSV.MaxFileSize != 0 && setting.UI.CSV.MaxFileSize < int64(len(rawBytes)) {
if _, err := tmpBlock.WriteString("<pre>"); err != nil {
return err
}
if _, err := tmpBlock.WriteString(html.EscapeString(string(rawBytes))); err != nil {
return err
}
if _, err := tmpBlock.WriteString("</pre>"); err != nil {
return err
}
return tmpBlock.Flush()
if int64(len(rawBytes)) <= maxSize {
return r.tableRender(ctx, bytes.NewReader(rawBytes), tmpBlock)
}
return r.fallbackRender(io.MultiReader(bytes.NewReader(rawBytes), input), tmpBlock)
}
func (Renderer) fallbackRender(input io.Reader, tmpBlock *bufio.Writer) error {
_, err := tmpBlock.WriteString("<pre>")
if err != nil {
return err
}
rd, err := csv.CreateReaderAndDetermineDelimiter(ctx, bytes.NewReader(rawBytes))
scan := bufio.NewScanner(input)
scan.Split(bufio.ScanRunes)
for scan.Scan() {
switch scan.Text() {
case `&`:
_, err = tmpBlock.WriteString("&amp;")
case `'`:
_, err = tmpBlock.WriteString("&#39;") // "&#39;" is shorter than "&apos;" and apos was not in HTML until HTML5.
case `<`:
_, err = tmpBlock.WriteString("&lt;")
case `>`:
_, err = tmpBlock.WriteString("&gt;")
case `"`:
_, err = tmpBlock.WriteString("&#34;") // "&#34;" is shorter than "&quot;".
default:
_, err = tmpBlock.Write(scan.Bytes())
}
if err != nil {
return err
}
}
_, err = tmpBlock.WriteString("</pre>")
if err != nil {
return err
}
return tmpBlock.Flush()
}
func (Renderer) tableRender(ctx *markup.RenderContext, input io.Reader, tmpBlock *bufio.Writer) error {
rd, err := csv.CreateReaderAndDetermineDelimiter(ctx, input)
if err != nil {
return err
}

View file

@ -4,6 +4,8 @@
package markup
import (
"bufio"
"bytes"
"strings"
"testing"
@ -29,4 +31,12 @@ func TestRenderCSV(t *testing.T) {
assert.NoError(t, err)
assert.EqualValues(t, v, buf.String())
}
t.Run("fallbackRender", func(t *testing.T) {
var buf bytes.Buffer
err := render.fallbackRender(strings.NewReader("1,<a>\n2,<b>"), bufio.NewWriter(&buf))
assert.NoError(t, err)
want := "<pre>1,&lt;a&gt;\n2,&lt;b&gt;</pre>"
assert.Equal(t, want, buf.String())
})
}