fix(push): fix peeking when push name is truncated (#3842)

* fix(push): Fix peeking when push name is truncated

* fix(push): fix peeking malformated push

* protect agains potential panic

Author: ndyakov

internal/proto/peek_push_notification_test.go

@@ -3,7 +3,10 @@ package proto
 import (
 	"bytes"
 	"fmt"
+	"io"
+	"math"
 	"math/rand"
+	"strconv"
 	"strings"
 	"testing"
 )
@@ -621,3 +624,161 @@ func TestPeekPushNotificationNameBehavior(t *testing.T) {
 		t.Log("Note: Buffer must be primed with a peek operation first")
 	})
 }
+
+// chunkedReader hands its data out in fixed-size chunks, one per Read. This
+// lets us deterministically reproduce the bufio fill boundary that landed in
+// the middle of a push frame header in issue #3839.
+type chunkedReader struct {
+	data      []byte
+	chunkSize int
+	off       int
+}
+
+func (r *chunkedReader) Read(p []byte) (int, error) {
+	if r.off >= len(r.data) {
+		return 0, io.EOF
+	}
+	n := r.chunkSize
+	if n > len(p) {
+		n = len(p)
+	}
+	if r.off+n > len(r.data) {
+		n = len(r.data) - r.off
+	}
+	copy(p, r.data[r.off:r.off+n])
+	r.off += n
+	return n, nil
+}
+
+// TestPeekPushNotificationName_TruncatedHeader covers issue #3839: when only
+// a prefix of a push frame is buffered, PeekPushNotificationName must NOT
+// return a truncated name. It should block to fetch the rest of the header
+// (so the caller can identify the notification correctly) or, when the
+// underlying read fails, return that error.
+func TestPeekPushNotificationName_TruncatedHeader(t *testing.T) {
+	const frame = ">3\r\n$7\r\nmessage\r\n$7\r\nchannel\r\n$5\r\nhello\r\n"
+
+	// Boundaries that previously caused PeekPushNotificationName to silently
+	// return a truncated name. `prefix` is the number of bytes the chunked
+	// reader hands out per Read call; with the bug, every chunk boundary that
+	// lands inside the frame header is a place where the function returned a
+	// truncated name instead of blocking for the rest.
+	cutpoints := []struct {
+		name   string
+		prefix int
+	}{
+		{"after_push_marker", 1},
+		{"inside_array_len_line", 3},
+		{"after_array_len_line", 4},
+		{"after_dollar", 5},
+		{"inside_bulk_len_line", 6},
+		{"after_bulk_len_line", 8},
+		{"inside_name", 13}, // the exact failure mode from #3839 ("messa")
+		{"before_name_crlf", 15},
+	}
+
+	for _, tc := range cutpoints {
+		t.Run(tc.name, func(t *testing.T) {
+			data := []byte(frame)
+			// Use a chunked reader that hands out tc.prefix bytes per Read
+			// call (not per fill). bufio's first fill therefore returns only
+			// tc.prefix bytes; satisfying a Peek that wants more forces
+			// additional Reads, each capped at tc.prefix bytes, until the
+			// peek window is full. That sequence reproduces the partial-fill
+			// boundary from issue #3839 where the buffered prefix landed in
+			// the middle of the push frame header.
+			rd := NewReader(&chunkedReader{data: data, chunkSize: tc.prefix})
+
+			if _, err := rd.PeekReplyType(); err != nil {
+				t.Fatalf("PeekReplyType: %v", err)
+			}
+
+			name, err := rd.PeekPushNotificationName()
+			if err != nil {
+				t.Fatalf("PeekPushNotificationName: unexpected error: %v", err)
+			}
+			if name != "message" {
+				t.Fatalf("PeekPushNotificationName: got %q, want %q", name, "message")
+			}
+
+			// And the full frame must still be readable afterwards.
+			reply, err := rd.ReadReply()
+			if err != nil {
+				t.Fatalf("ReadReply: %v", err)
+			}
+			arr, ok := reply.([]interface{})
+			if !ok || len(arr) != 3 || arr[0] != "message" || arr[1] != "channel" || arr[2] != "hello" {
+				t.Fatalf("ReadReply: got %#v, want [message channel hello]", reply)
+			}
+		})
+	}
+}
+
+// TestPeekPushNotificationName_TruncatedHeaderUnderlyingEOF covers the case
+// where the underlying connection delivers a partial frame and then closes.
+// PeekPushNotificationName must surface the read error rather than returning
+// a truncated name.
+func TestPeekPushNotificationName_TruncatedHeaderUnderlyingEOF(t *testing.T) {
+	// A push frame that ends mid-name.
+	data := []byte(">3\r\n$7\r\nmessa")
+	rd := NewReader(bytes.NewReader(data))
+
+	if _, err := rd.PeekReplyType(); err != nil {
+		t.Fatalf("PeekReplyType: %v", err)
+	}
+
+	name, err := rd.PeekPushNotificationName()
+	if err == nil {
+		t.Fatalf("PeekPushNotificationName: want error, got name=%q", name)
+	}
+	if name != "" {
+		t.Fatalf("PeekPushNotificationName: want empty name on error, got %q", name)
+	}
+}
+
+// TestPeekPushNotificationName_EmptyArrayLength asserts that ">\r\n" is
+// reported as a malformed frame rather than an incomplete prefix. RESP
+// requires at least one digit after '>' for the array length; without the
+// empty-length check in parsePushNotificationName the parser would treat
+// ">\r\n" as a valid prefix and PeekPushNotificationName would block
+// waiting for the rest of a frame that is already corrupt.
+func TestPeekPushNotificationName_EmptyArrayLength(t *testing.T) {
+	rd := NewReader(bytes.NewReader([]byte(">\r\n")))
+
+	if _, err := rd.PeekReplyType(); err != nil {
+		t.Fatalf("PeekReplyType: %v", err)
+	}
+
+	name, err := rd.PeekPushNotificationName()
+	if err == nil {
+		t.Fatalf("PeekPushNotificationName: want error for empty array length, got name=%q", name)
+	}
+	if name != "" {
+		t.Fatalf("PeekPushNotificationName: want empty name on error, got %q", name)
+	}
+	if !strings.Contains(err.Error(), "empty push notification array length") {
+		t.Fatalf("PeekPushNotificationName: error should mention empty array length, got %v", err)
+	}
+}
+
+// TestPeekPushNotificationName_OverflowingNameLength asserts that a frame
+// advertising a name length near math.MaxInt does not panic. Computing
+// next+nameLen would overflow int and wrap negative, slipping past a naive
+// "end > len(buf)" guard and panicking the backing slice; the parser must
+// instead treat it as incomplete and surface an error without crashing.
+func TestPeekPushNotificationName_OverflowingNameLength(t *testing.T) {
+	data := []byte(">1\r\n$" + strconv.Itoa(math.MaxInt) + "\r\n")
+	rd := NewReader(bytes.NewReader(data))
+
+	if _, err := rd.PeekReplyType(); err != nil {
+		t.Fatalf("PeekReplyType: %v", err)
+	}
+
+	name, err := rd.PeekPushNotificationName()
+	if err == nil {
+		t.Fatalf("PeekPushNotificationName: want error for overflowing name length, got name=%q", name)
+	}
+	if name != "" {
+		t.Fatalf("PeekPushNotificationName: want empty name on error, got %q", name)
+	}
+}

internal/proto/reader.go

@@ -100,90 +100,159 @@ func (r *Reader) PeekReplyType() (byte, error) {
 	return b[0], nil
 }
 
+// PeekPushNotificationName returns the notification name of the next RESP3
+// push frame without consuming it. The caller is expected to have already
+// verified that the next reply is a push notification (e.g. via PeekReplyType
+// returning RespPush).
+//
+// To identify the name the method may block briefly reading more bytes from
+// the underlying connection. That is safe: once the push marker '>' has been
+// observed, the server is committed to sending the rest of the frame, so
+// fetching the next few header bytes does not race with anything the caller
+// could be waiting on. Blocking is preferred to a truncated peek, which would
+// silently misidentify the notification and cause the caller's ReadReply to
+// consume (and drop) the frame; see issue #3839.
 func (r *Reader) PeekPushNotificationName() (string, error) {
-	// "prime" the buffer by peeking at the next byte
-	c, err := r.Peek(1)
+	c, err := r.rd.Peek(1)
 	if err != nil {
 		return "", err
 	}
 	if c[0] != RespPush {
 		return "", fmt.Errorf("redis: can't peek push notification name, next reply is not a push notification")
 	}
 
-	// peek 36 bytes at most, should be enough to read the push notification name
-	toPeek := 36
-	buffered := r.Buffered()
-	if buffered == 0 {
-		return "", fmt.Errorf("redis: can't peek push notification name, no data available")
-	}
-	if buffered < toPeek {
-		toPeek = buffered
+	// Start with a peek window that covers every Redis-defined notification
+	// header (MOVING, MIGRATING, FAILED_OVER, message, pmessage, smessage,
+	// subscribe, unsubscribe, ...). If a longer name is encountered, grow
+	// the window up to maxPushHeaderPeek before giving up.
+	const initialPeek = 36
+	const maxPushHeaderPeek = 4096
+
+	peekSize := initialPeek
+	for {
+		buf, peekErr := r.rd.Peek(peekSize)
+		name, complete, parseErr := parsePushNotificationName(buf)
+		if parseErr != nil {
+			return "", parseErr
+		}
+		if complete {
+			return name, nil
+		}
+		// Parser ran out of bytes. Surface a failed underlying read before
+		// growing further; otherwise grow the peek window and retry.
+		if peekErr != nil {
+			return "", peekErr
+		}
+		if peekSize >= maxPushHeaderPeek {
+			return "", fmt.Errorf("redis: push notification header exceeds %d bytes", maxPushHeaderPeek)
+		}
+		peekSize *= 2
+		if peekSize > maxPushHeaderPeek {
+			peekSize = maxPushHeaderPeek
+		}
 	}
-	buf, err := r.rd.Peek(toPeek)
-	if err != nil {
-		return "", err
+}
+
+// parsePushNotificationName extracts the notification name from a buffered
+// RESP3 push frame prefix. The three return values are:
+//
+//   - (name, true, nil): the full name is in buf.
+//   - ("", false, nil):  buf is a valid prefix but too short to determine the
+//     name; the caller should fetch more bytes and retry.
+//   - ("", _, err):      buf is malformed.
+//
+// This split lets PeekPushNotificationName tell "incomplete header" apart
+// from "corrupt frame" without ever returning a truncated string.
+func parsePushNotificationName(buf []byte) (string, bool, error) {
+	// Need at least ">N\r" before any meaningful work.
+	if len(buf) < 3 {
+		return "", false, nil
 	}
 	if buf[0] != RespPush {
-		return "", fmt.Errorf("redis: can't parse push notification: %q", buf)
+		return "", false, fmt.Errorf("redis: can't parse push notification: %q", buf)
 	}
 
-	if len(buf) < 3 {
-		return "", fmt.Errorf("redis: can't parse push notification: %q", buf)
+	// Skip the array length line ">N\r\n".
+	const arrayLenStart = 1 // first byte after the '>' marker
+	pos, ok, err := skipDigitsThenCRLF(buf, arrayLenStart)
+	if err != nil {
+		return "", false, fmt.Errorf("redis: can't parse push notification: %w", err)
 	}
-
-	// remove push notification type
-	buf = buf[1:]
-	// remove first line - e.g. >2\r\n
-	for i := 0; i < len(buf)-1; i++ {
-		if buf[i] == '\r' && buf[i+1] == '\n' {
-			buf = buf[i+2:]
-			break
-		} else {
-			if buf[i] < '0' || buf[i] > '9' {
-				return "", fmt.Errorf("redis: can't parse push notification: %q", buf)
-			}
-		}
+	if !ok {
+		return "", false, nil
+	}
+	// Reject ">\r\n": RESP requires at least one digit for the array length.
+	// Without this check the empty length looks like a valid prefix and the
+	// caller would block fetching more bytes for a frame that is already
+	// malformed.
+	if pos-2 == arrayLenStart {
+		return "", false, fmt.Errorf("redis: empty push notification array length")
 	}
-	if len(buf) < 2 {
-		return "", fmt.Errorf("redis: can't parse push notification: %q", buf)
+
+	// First element type byte: '$' (bulk) or '+' (simple-string).
+	if pos >= len(buf) {
+		return "", false, nil
 	}
-	// next line should be $<length><string>\r\n or +<length><string>\r\n
-	// should have the type of the push notification name and it's length
-	if buf[0] != RespString && buf[0] != RespStatus {
-		return "", fmt.Errorf("redis: can't parse push notification name: %q", buf)
+	typeOfName := buf[pos]
+	if typeOfName != RespString && typeOfName != RespStatus {
+		return "", false, fmt.Errorf("redis: can't parse push notification name: %q", buf[pos:])
 	}
-	typeOfName := buf[0]
-	// remove the type of the push notification name
-	buf = buf[1:]
+	pos++
+
 	if typeOfName == RespString {
-		// remove the length of the string
-		if len(buf) < 2 {
-			return "", fmt.Errorf("redis: can't parse push notification name: %q", buf)
-		}
-		for i := 0; i < len(buf)-1; i++ {
-			if buf[i] == '\r' && buf[i+1] == '\n' {
-				buf = buf[i+2:]
-				break
-			} else {
-				if buf[i] < '0' || buf[i] > '9' {
-					return "", fmt.Errorf("redis: can't parse push notification name: %q", buf)
-				}
-			}
+		// Read "$M\r\n" then the M-byte name.
+		lenStart := pos
+		next, ok, err := skipDigitsThenCRLF(buf, pos)
+		if err != nil {
+			return "", false, fmt.Errorf("redis: can't parse push notification name length: %w", err)
+		}
+		if !ok {
+			return "", false, nil
 		}
+		if next-2 == lenStart {
+			return "", false, fmt.Errorf("redis: empty push notification name length")
+		}
+		nameLen, err := util.Atoi(buf[lenStart : next-2])
+		if err != nil {
+			return "", false, fmt.Errorf("redis: invalid push notification name length %q: %w", buf[lenStart:next-2], err)
+		}
+		if nameLen < 0 {
+			return "", false, fmt.Errorf("redis: negative push notification name length: %d", nameLen)
+		}
+		// Compare against the remaining bytes instead of computing
+		// next+nameLen: a hugely advertised length on malformed input could
+		// overflow int, wrap negative, slip past an "end > len(buf)" guard and
+		// panic the slice below. next <= len(buf) here, so the subtraction is
+		// safe.
+		if nameLen > len(buf)-next {
+			return "", false, nil
+		}
+		return util.BytesToString(buf[next : next+nameLen]), true, nil
 	}
 
-	if len(buf) < 2 {
-		return "", fmt.Errorf("redis: can't parse push notification name: %q", buf)
-	}
-	// keep only the notification name
-	for i := 0; i < len(buf)-1; i++ {
+	// RespStatus: scan for the terminating CRLF.
+	for i := pos; i < len(buf)-1; i++ {
 		if buf[i] == '\r' && buf[i+1] == '\n' {
-			buf = buf[:i]
-			break
+			return util.BytesToString(buf[pos:i]), true, nil
 		}
 	}
+	return "", false, nil
+}
 
-	return util.BytesToString(buf), nil
+// skipDigitsThenCRLF advances past zero-or-more ASCII digits and the
+// terminating "\r\n" starting at offset start in buf. It returns the position
+// after the "\r\n" and true on success; (pos, false, nil) if buf is too
+// short; or an error if a non-digit non-CR byte is encountered before the CRLF.
+func skipDigitsThenCRLF(buf []byte, start int) (int, bool, error) {
+	for pos := start; pos < len(buf)-1; pos++ {
+		if buf[pos] == '\r' && buf[pos+1] == '\n' {
+			return pos + 2, true, nil
+		}
+		if buf[pos] < '0' || buf[pos] > '9' {
+			return pos, false, fmt.Errorf("expected digit or CRLF, got %q", buf[pos])
+		}
+	}
+	return len(buf), false, nil
 }
 
 // ReadLine Return a valid reply, it will check the protocol or redis error,