From 3e3223a05bef6c6a30de17845de6ff79faade59f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jo=C3=ABl=20Stemmer?= <stemmertech@gmail.com>
Date: Mon, 18 Jul 2022 22:38:36 +0100
Subject: [PATCH] parser/gotest: Create LimitedLineReader

---
 parser/gotest/internal/reader/reader.go      | 58 +++++++++++++++++
 parser/gotest/internal/reader/reader_test.go | 68 ++++++++++++++++++++
 2 files changed, 126 insertions(+)
 create mode 100644 parser/gotest/internal/reader/reader.go
 create mode 100644 parser/gotest/internal/reader/reader_test.go

diff --git a/parser/gotest/internal/reader/reader.go b/parser/gotest/internal/reader/reader.go
new file mode 100644
index 0000000..57b2365
--- /dev/null
+++ b/parser/gotest/internal/reader/reader.go
@@ -0,0 +1,58 @@
+package reader
+
+import (
+	"bufio"
+	"bytes"
+	"io"
+)
+
+// LimitedLineReader reads lines from an io.Reader object with a configurable
+// line size limit. Lines exceeding the limit will be truncated, but read
+// completely from the underlying io.Reader.
+type LimitedLineReader struct {
+	r     *bufio.Reader
+	limit int
+}
+
+// NewLimitedLineReader returns a LimitedLineReader to read lines from r with a
+// maximum line size of limit.
+func NewLimitedLineReader(r io.Reader, limit int) *LimitedLineReader {
+	return &LimitedLineReader{r: bufio.NewReader(r), limit: limit}
+}
+
+// ReadLine returns the next line from the underlying reader. The length of the
+// line will not exceed the configured limit. ReadLine either returns a line or
+// it returns an error, never both.
+func (r *LimitedLineReader) ReadLine() (string, error) {
+	line, isPrefix, err := r.r.ReadLine()
+	if err != nil {
+		return "", err
+	}
+
+	if !isPrefix {
+		return string(line), nil
+	}
+
+	// Line is incomplete, keep reading until we reach the end of the line.
+	var buf bytes.Buffer
+	buf.Write(line) // ignore err, always nil
+	for isPrefix {
+		line, isPrefix, err = r.r.ReadLine()
+		if err != nil {
+			return "", err
+		}
+
+		if buf.Len() >= r.limit {
+			// Stop writing to buf if we exceed the limit. We continue reading
+			// however to make sure we consume the entire line.
+			continue
+		}
+
+		buf.Write(line) // ignore err, always nil
+	}
+
+	if buf.Len() > r.limit {
+		buf.Truncate(r.limit)
+	}
+	return buf.String(), nil
+}
diff --git a/parser/gotest/internal/reader/reader_test.go b/parser/gotest/internal/reader/reader_test.go
new file mode 100644
index 0000000..8526b9d
--- /dev/null
+++ b/parser/gotest/internal/reader/reader_test.go
@@ -0,0 +1,68 @@
+package reader
+
+import (
+	"bufio"
+	"io"
+	"strings"
+	"testing"
+)
+
+const testingLimit = 4 * 1024 * 1024
+
+func TestLimitedLineReader(t *testing.T) {
+	tests := []struct {
+		desc      string
+		inputSize int
+	}{
+		{"small size", 128},
+		{"under buf size", 4095},
+		{"buf size", 4096},
+		{"multiple of buf size ", 4096 * 2},
+		{"not multiple of buf size", 10 * 1024},
+		{"bufio.MaxScanTokenSize", bufio.MaxScanTokenSize},
+		{"over bufio.MaxScanTokenSize", bufio.MaxScanTokenSize + 1},
+		{"under limit", testingLimit - 1},
+		{"at limit", testingLimit},
+		{"just over limit", testingLimit + 1},
+		{"over limit", testingLimit + 128},
+	}
+
+	for _, test := range tests {
+		t.Run(test.desc, func(t *testing.T) {
+			line1 := string(make([]byte, test.inputSize))
+			line2 := "other line"
+			input := strings.NewReader(strings.Join([]string{line1, line2}, "\n"))
+			r := NewLimitedLineReader(input, testingLimit)
+
+			got, err := r.ReadLine()
+			if err != nil {
+				t.Fatalf("ReadLine() returned error %v", err)
+			}
+
+			want := line1
+			if len(line1) > testingLimit {
+				want = want[:testingLimit]
+			}
+			if got != want {
+				t.Fatalf("ReadLine() returned incorrect line, got len %d want len %d", len(got), len(want))
+			}
+
+			got, err = r.ReadLine()
+			if err != nil {
+				t.Fatalf("ReadLine() returned error %v", err)
+			}
+			want = line2
+			if got != want {
+				t.Fatalf("ReadLine() returned incorrect line, got len %d want len %d", len(got), len(want))
+			}
+
+			got, err = r.ReadLine()
+			if err != io.EOF {
+				t.Fatalf("ReadLine() returned unexpected error, got %v want %v\n", err, io.EOF)
+			}
+			if got != "" {
+				t.Fatalf("ReadLine() returned unexpected line, got %v want nothing\n", got)
+			}
+		})
+	}
+}