From 5536bca34e218a1e17a55e5ba36c22fab5784f44 Mon Sep 17 00:00:00 2001
From: Raphael 'kena' Poss <knz@thaumogen.net>
Date: Wed, 18 Oct 2023 12:59:26 +0200
Subject: [PATCH] fix(key): support very long buffered input (#570)

---
 key.go      | 36 ++++++++++++++++++++++++++++++++++--
 key_test.go | 21 ++++++++++++++++++++-
 2 files changed, 54 insertions(+), 3 deletions(-)

diff --git a/key.go b/key.go
index b712228..fe438a6 100644
--- a/key.go
+++ b/key.go
@@ -543,6 +543,8 @@ var spaceRunes = []rune{' '}
 func readInputs(ctx context.Context, msgs chan<- Msg, input io.Reader) error {
 	var buf [256]byte
 
+	var leftOverFromPrevIteration []byte
+loop:
 	for {
 		// Read and block.
 		numBytes, err := input.Read(buf[:])
@@ -550,11 +552,31 @@ func readInputs(ctx context.Context, msgs chan<- Msg, input io.Reader) error {
 			return fmt.Errorf("error reading input: %w", err)
 		}
 		b := buf[:numBytes]
+		if leftOverFromPrevIteration != nil {
+			b = append(leftOverFromPrevIteration, b...)
+		}
+
+		// If we had a short read (numBytes < len(buf)), we're sure that
+		// the end of this read is an event boundary, so there is no doubt
+		// if we are encountering the end of the buffer while parsing a message.
+		// However, if we've succeeded in filling up the buffer, there may
+		// be more data in the OS buffer ready to be read in, to complete
+		// the last message in the input. In that case, we will retry with
+		// the left over data in the next iteration.
+		canHaveMoreData := numBytes == len(buf)
 
 		var i, w int
 		for i, w = 0, 0; i < len(b); i += w {
 			var msg Msg
-			w, msg = detectOneMsg(b[i:])
+			w, msg = detectOneMsg(b[i:], canHaveMoreData)
+			if w == 0 {
+				// Expecting more bytes beyond the current buffer. Try waiting
+				// for more input.
+				leftOverFromPrevIteration = make([]byte, 0, len(b[i:])+len(buf))
+				leftOverFromPrevIteration = append(leftOverFromPrevIteration, b[i:]...)
+				continue loop
+			}
+
 			select {
 			case msgs <- msg:
 			case <-ctx.Done():
@@ -565,12 +587,13 @@ func readInputs(ctx context.Context, msgs chan<- Msg, input io.Reader) error {
 				return err
 			}
 		}
+		leftOverFromPrevIteration = nil
 	}
 }
 
 var unknownCSIRe = regexp.MustCompile(`^\x1b\[[\x30-\x3f]*[\x20-\x2f]*[\x40-\x7e]`)
 
-func detectOneMsg(b []byte) (w int, msg Msg) {
+func detectOneMsg(b []byte, canHaveMoreData bool) (w int, msg Msg) {
 	// Detect mouse events.
 	const mouseEventLen = 6
 	if len(b) >= mouseEventLen && b[0] == '\x1b' && b[1] == '[' && b[2] == 'M' {
@@ -618,6 +641,15 @@ func detectOneMsg(b []byte) (w int, msg Msg) {
 			break
 		}
 	}
+	if i >= len(b) && canHaveMoreData {
+		// We have encountered the end of the input buffer. Alas, we can't
+		// be sure whether the data in the remainder of the buffer is
+		// complete (maybe there was a short read). Instead of sending anything
+		// dumb to the message channel, do a short read. The outer loop will
+		// handle this case by extending the buffer as necessary.
+		return 0, nil
+	}
+
 	// If we found at least one rune, we report the bunch of them as
 	// a single KeyRunes or KeySpace event.
 	if len(runes) > 0 {
diff --git a/key_test.go b/key_test.go
index d466b05..ae8643b 100644
--- a/key_test.go
+++ b/key_test.go
@@ -200,7 +200,7 @@ func TestDetectOneMsg(t *testing.T) {
 
 	for _, tc := range td {
 		t.Run(fmt.Sprintf("%q", string(tc.seq)), func(t *testing.T) {
-			width, msg := detectOneMsg(tc.seq)
+			width, msg := detectOneMsg(tc.seq, false /* canHaveMoreData */)
 			if width != len(tc.seq) {
 				t.Errorf("parser did not consume the entire input: got %d, expected %d", width, len(tc.seq))
 			}
@@ -211,6 +211,25 @@ func TestDetectOneMsg(t *testing.T) {
 	}
 }
 
+func TestReadLongInput(t *testing.T) {
+	input := strings.Repeat("a", 1000)
+	msgs := testReadInputs(t, bytes.NewReader([]byte(input)))
+	if len(msgs) != 1 {
+		t.Errorf("expected 1 messages, got %d", len(msgs))
+	}
+	km := msgs[0]
+	k := Key(km.(KeyMsg))
+	if k.Type != KeyRunes {
+		t.Errorf("expected key runes, got %d", k.Type)
+	}
+	if len(k.Runes) != 1000 || !reflect.DeepEqual(k.Runes, []rune(input)) {
+		t.Errorf("unexpected runes: %+v", k)
+	}
+	if k.Alt {
+		t.Errorf("unexpected alt")
+	}
+}
+
 func TestReadInput(t *testing.T) {
 	type test struct {
 		keyname string