Fix PUB command parsing + add tests

commands-reader.go

@@ -29,27 +29,34 @@ func (cr CommandsReader) nextCommand() ([]byte, error) {
 		return nil, err
 	}
 
-	if bytes.Equal(line[0:3], []byte("MSG")) {
+	if len(line) == 0 {
+		return nil, fmt.Errorf("Unexpected empty line")
+	}
+	if len(line) < 3 {
+		return nil, fmt.Errorf("Invalid command: '%s'", line)
+	}
+	op := line[0:3]
+	if bytes.Equal(op, []byte("MSG")) || bytes.Equal(op, []byte("PUB")) {
 		msg = line[:]
 		splitted := bytes.Split(line, []byte(" "))
 		sizeStr := splitted[len(splitted)-1]
 		sizeStr = sizeStr[:len(sizeStr)-2]
 		size, err := strconv.Atoi(string(sizeStr))
 		if err != nil {
-			return nil, fmt.Errorf("Error reading MSG size: %s", err)
+			return nil, fmt.Errorf("Error reading %s size: %s", op, err)
 		}
 		// the '-2' is to account for the trailing \r\n which is after the payload
 		for size > -2 {
 			chunk, err := cr.br.ReadBytes('\n')
 			if err != nil {
-				return nil, fmt.Errorf("Error reading MSG payload: %s", err)
+				return nil, fmt.Errorf("Error reading %s payload: %s", op, err)
 			}
 			size -= len(chunk)
 			msg = append(msg, chunk...)
 		}
 		if size != -2 {
 			return nil, fmt.Errorf(
-				"Error reading MSG payload. Got %d extra bytes", -size-2)
+				"Error reading %s payload. Got %d extra bytes", op, -size-2)
 		}
 	} else {
 		msg = line

commands-reader_test.go

@@ -0,0 +1,50 @@
+package gw
+
+import (
+	"bytes"
+	"testing"
+
+	"gotest.tools/assert"
+)
+
+func TestCommandsReader(t *testing.T) {
+	for _, tt := range []struct {
+		name     string
+		commands []string
+		err      string
+	}{
+		{
+			name: "base",
+			commands: []string{
+				"INFO {}\r\n",
+				"MSG test 1 3\r\n123\r\n",
+				"MSG test 1 3\r\n1\r\n\r\n",
+				"PUB test 3\r\n1\r\n\r\n",
+			},
+		},
+	} {
+		t.Run(tt.name, func(t *testing.T) {
+			var buf bytes.Buffer
+			for _, s := range tt.commands {
+				buf.WriteString(s)
+			}
+			reader := NewCommandsReader(&buf)
+			for _, expected := range tt.commands {
+				next, err := reader.nextCommand()
+				if err != nil {
+					t.Fatal(err)
+				}
+				assert.Equal(t, expected, string(next))
+			}
+			_, err := reader.nextCommand()
+			if err == nil {
+				t.Fatal("Expected an error")
+			}
+			if tt.err != "" {
+				assert.Equal(t, tt.err, err.Error())
+			} else {
+				assert.Equal(t, "EOF", err.Error())
+			}
+		})
+	}
+}