-
Notifications
You must be signed in to change notification settings - Fork 12
/
commands-reader.go
77 lines (70 loc) · 1.7 KB
/
commands-reader.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
package gw
import (
"bufio"
"bytes"
"fmt"
"io"
"strconv"
)
// CommandsReader parses a NATS connection input stream into commands
type CommandsReader struct {
io.Reader
br *bufio.Reader
}
// NewCommandsReader creates a CommandsReader
func NewCommandsReader(src io.Reader) CommandsReader {
return CommandsReader{
Reader: src,
br: bufio.NewReader(src),
}
}
// NextCommand returns the next command in the input stream
func (cr CommandsReader) NextCommand() ([]byte, error) {
return cr.nextCommand()
}
func (cr CommandsReader) nextCommand() ([]byte, error) {
var msg []byte
line, err := cr.br.ReadBytes('\n')
if err != nil {
return nil, err
}
for bytes.Equal(line, []byte("\r\n")) {
line, err = cr.br.ReadBytes('\n')
if err != nil {
return nil, err
}
}
if len(line) == 0 {
return nil, fmt.Errorf("Unexpected empty line")
}
if len(line) < 3 {
return nil, fmt.Errorf("Invalid command: %v", 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 %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 %s payload: %s", op, err)
}
size -= len(chunk)
msg = append(msg, chunk...)
}
if size != -2 {
return nil, fmt.Errorf(
"Error reading %s payload. Got %d extra bytes", op, -size-2)
}
} else {
msg = line
}
return msg, nil
}