add more tests

Author: wxiaoguang

modules/charset/escape_stream.go

@@ -16,17 +16,21 @@ import (
 	"golang.org/x/net/html"
 )
 
+type htmlChunkReader struct {
+	in       io.Reader
+	readErr  error
+	readBuf  []byte
+	curInTag bool
+}
+
 type escapeStreamer struct {
+	htmlChunkReader
+
 	escaped         *EscapeStatus
 	locale          translation.Locale
 	ambiguousTables []*AmbiguousTable
 	allowed         map[rune]bool
 
-	in       io.Reader
-	readErr  error
-	readBuf  []byte
-	curInTag bool
-
 	out io.Writer
 }
 
@@ -35,10 +39,10 @@ func escapeStream(locale translation.Locale, in io.Reader, out io.Writer, opts .
 		escaped:         &EscapeStatus{},
 		locale:          locale,
 		ambiguousTables: AmbiguousTablesForLocale(locale),
-
-		in:      in,
-		readBuf: make([]byte, 0, 32*1024),
-
+		htmlChunkReader: htmlChunkReader{
+			in:      in,
+			readBuf: make([]byte, 0, 32*1024),
+		},
 		out: out,
 	}
 
@@ -134,7 +138,7 @@ func (e *escapeStreamer) detectAndWriteRunes(part []byte) error {
 	return e.writeDetectResults(part, results)
 }
 
-func (e *escapeStreamer) readRunes() (parts [][]byte, partInTag []bool, _ error) {
+func (e *htmlChunkReader) readRunes() (parts [][]byte, partInTag []bool, _ error) {
 	// we have read everything, eof
 	if e.readErr != nil && len(e.readBuf) == 0 {
 		return nil, nil, e.readErr
@@ -195,7 +199,7 @@ func (e *escapeStreamer) readRunes() (parts [][]byte, partInTag []bool, _ error)
 		// now we get the curPart bytes, but we can't directly use it, the last rune in it might have been cut
 		// try to decode the last rune, if it's invalid, then we cut the last byte and try again until we get a valid rune or no byte left
 		for i := curPartLen - 1; i >= 0; i-- {
-			last, lastSize := utf8.DecodeLastRune(curPart[i:])
+			last, lastSize := utf8.DecodeRune(curPart[i:])
 			if last == utf8.RuneError && lastSize == 1 {
 				curPartLen--
 			} else {
@@ -210,13 +214,23 @@ func (e *escapeStreamer) readRunes() (parts [][]byte, partInTag []bool, _ error)
 			// * at least consume 1 byte to avoid infinite loop
 			curPartLen = max(len(curPart)-utf8.UTFMax, 1)
 		}
+
+		// if curPartLen is not the same as curPart, it means we have cut some bytes,
+		// need to wait for more data if not eof
+		trailingCorrupted := curPartLen != len(curPart)
+
 		// finally, we get the real part we need
 		curPart = curPart[:curPartLen]
 		parts = append(parts, curPart)
 		partInTag = append(partInTag, e.curInTag)
 
 		pos += curPartLen
 		e.curInTag = nextInTag
+
+		if trailingCorrupted && e.readErr == nil {
+			// if the last part is corrupted, and we haven't reach eof, then we need to wait for more data to get the complete part
+			break
+		}
 	}
 
 	copy(e.readBuf, e.readBuf[pos:])

modules/charset/escape_test.go

@@ -13,6 +13,7 @@ import (
 	"code.gitea.io/gitea/modules/translation"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 type escapeControlTest struct {
@@ -165,3 +166,50 @@ func TestSettingAmbiguousUnicodeDetection(t *testing.T) {
 	_, out = EscapeControlHTML("a test", &translation.MockLocale{})
 	assert.EqualValues(t, `a test`, out)
 }
+
+func TestHTMLChunkReader(t *testing.T) {
+	type textPart struct {
+		text  string
+		isTag bool
+	}
+	testReadChunks := func(t *testing.T, chunkSize int, input string, expected []textPart) {
+		r := &htmlChunkReader{in: strings.NewReader(input), readBuf: make([]byte, 0, chunkSize)}
+		var results []textPart
+		for {
+			parts, partIsTag, err := r.readRunes()
+			if err != nil {
+				break
+			}
+			for i, part := range parts {
+				results = append(results, textPart{string(part), partIsTag[i]})
+			}
+		}
+		assert.Equal(t, expected, results, "chunk size: %d, input: %s", chunkSize, input)
+	}
+
+	testReadChunks(t, 10, "abc<def>ghi", []textPart{
+		{text: "abc", isTag: false},
+		{text: "<def>", isTag: true},
+		{text: "gh", isTag: false},
+		// -- chunk
+		{text: "i", isTag: false},
+	})
+
+	testReadChunks(t, 10, "<abc><def>ghi", []textPart{
+		{text: "<abc>", isTag: true},
+		{text: "<def>", isTag: true},
+		// -- chunk
+		{text: "ghi", isTag: false},
+	})
+
+	rune1, rune2, rune3, rune4 := "A", "é", "啊", "🌞"
+	require.Len(t, rune1, 1)
+	require.Len(t, rune2, 2)
+	require.Len(t, rune3, 3)
+	require.Len(t, rune4, 4)
+	input := "<" + rune1 + rune2 + rune3 + rune4 + ">" + rune1 + rune2 + rune3 + rune4
+	testReadChunks(t, 4, input, []textPart{{"<Aé", true}, {"啊", true}, {"🌞", true}, {">", true}, {"Aé", false}, {"啊", false}, {"🌞", false}})
+	testReadChunks(t, 5, input, []textPart{{"<Aé", true}, {"啊", true}, {"🌞>", true}, {"Aé", false}, {"啊", false}, {"🌞", false}})
+	testReadChunks(t, 6, input, []textPart{{"<Aé", true}, {"啊", true}, {"🌞>", true}, {"A", false}, {"é啊", false}, {"🌞", false}})
+	testReadChunks(t, 7, input, []textPart{{"<Aé啊", true}, {"🌞>", true}, {"A", false}, {"é啊", false}, {"🌞", false}})
+}