fix(proto): fix notification parser

Author: ndyakov

internal/proto/peek_push_notification_test.go

@@ -3,6 +3,7 @@ package proto
 import (
 	"bytes"
 	"fmt"
+	"math/rand"
 	"strings"
 	"testing"
 )
@@ -215,9 +216,9 @@ func TestPeekPushNotificationName(t *testing.T) {
 					// This is acceptable behavior for malformed input
 					name, err := reader.PeekPushNotificationName()
 					if err != nil {
-						t.Logf("PeekPushNotificationName errored for corrupted data %s: %v", tc.name, err)
+						t.Logf("PeekPushNotificationName errored for corrupted data %s: %v (DATA: %s)", tc.name, err, tc.data)
 					} else {
-						t.Logf("PeekPushNotificationName returned '%s' for corrupted data %s", name, tc.name)
+						t.Logf("PeekPushNotificationName returned '%s' for corrupted data NAME: %s, DATA: %s", name, tc.name, tc.data)
 					}
 				})
 			}
@@ -293,15 +294,27 @@ func TestPeekPushNotificationName(t *testing.T) {
 func createValidPushNotification(notificationName, data string) *bytes.Buffer {
 	buf := &bytes.Buffer{}
 
+	simpleOrString := rand.Intn(2) == 0
+
 	if data == "" {
+
 		// Single element notification
 		buf.WriteString(">1\r\n")
-		buf.WriteString(fmt.Sprintf("$%d\r\n%s\r\n", len(notificationName), notificationName))
+		if simpleOrString {
+			buf.WriteString(fmt.Sprintf("+%s\r\n", notificationName))
+		} else {
+			buf.WriteString(fmt.Sprintf("$%d\r\n%s\r\n", len(notificationName), notificationName))
+		}
 	} else {
 		// Two element notification
 		buf.WriteString(">2\r\n")
-		buf.WriteString(fmt.Sprintf("$%d\r\n%s\r\n", len(notificationName), notificationName))
-		buf.WriteString(fmt.Sprintf("$%d\r\n%s\r\n", len(data), data))
+		if simpleOrString {
+			buf.WriteString(fmt.Sprintf("+%s\r\n", notificationName))
+			buf.WriteString(fmt.Sprintf("+%s\r\n", data))
+		} else {
+			buf.WriteString(fmt.Sprintf("$%d\r\n%s\r\n", len(notificationName), notificationName))
+			buf.WriteString(fmt.Sprintf("$%d\r\n%s\r\n", len(notificationName), notificationName))
+		}
 	}
 
 	return buf

internal/proto/reader.go

@@ -116,33 +116,63 @@ func (r *Reader) PeekPushNotificationName() (string, error) {
 	if buf[0] != RespPush {
 		return "", fmt.Errorf("redis: can't parse push notification: %q", buf)
 	}
-	// remove push notification type and length
-	buf = buf[2:]
+
+	if len(buf) < 3 {
+		return "", fmt.Errorf("redis: can't parse push notification: %q", buf)
+	}
+
+	// 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 len(buf) < 2 {
+		return "", fmt.Errorf("redis: can't parse push notification: %q", buf)
+	}
+	// 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 {
+	if buf[0] != RespString && buf[0] != RespStatus {
 		return "", fmt.Errorf("redis: can't parse push notification name: %q", buf)
 	}
-	// skip the length of the string
-	for i := 0; i < len(buf)-1; i++ {
-		if buf[i] == '\r' && buf[i+1] == '\n' {
-			buf = buf[i+2:]
-			break
+	typeOfName := buf[0]
+	// remove the type of the push notification name
+	buf = buf[1:]
+	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)
+				}
+			}
 		}
 	}
 
+	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++ {
 		if buf[i] == '\r' && buf[i+1] == '\n' {
 			buf = buf[:i]
 			break
 		}
 	}
+
 	return util.BytesToString(buf), nil
 }