From eee260f15430ef212d588b0990d5eedc203714c4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Ken-H=C3=A5vard=20Lieng?= Date: Fri, 11 Jan 2019 04:53:50 +0100 Subject: [PATCH] Read lines with a bufio.Scanner, reuse input buffer, ignore empty messages, handle multiple spaces between tags and prefix --- pkg/irc/client.go | 4 ++- pkg/irc/conn.go | 14 +++++--- pkg/irc/conn_test.go | 6 ++-- pkg/irc/message.go | 28 ++++++++------- pkg/irc/message_test.go | 78 +++++++++++++++++++++++++---------------- 5 files changed, 80 insertions(+), 50 deletions(-) diff --git a/pkg/irc/client.go b/pkg/irc/client.go index 107795ab..fe7327af 100644 --- a/pkg/irc/client.go +++ b/pkg/irc/client.go @@ -30,7 +30,8 @@ type Client struct { connected bool registered bool dialer *net.Dialer - reader *bufio.Reader + recvBuf []byte + scan *bufio.Scanner backoff *backoff.Backoff out chan string @@ -51,6 +52,7 @@ func NewClient(nick, username string) *Client { out: make(chan string, 32), quit: make(chan struct{}), reconnect: make(chan struct{}), + recvBuf: make([]byte, 0, 4096), backoff: &backoff.Backoff{ Jitter: true, }, diff --git a/pkg/irc/conn.go b/pkg/irc/conn.go index 3257fed7..10023e6f 100644 --- a/pkg/irc/conn.go +++ b/pkg/irc/conn.go @@ -2,6 +2,7 @@ package irc import ( "bufio" + "bytes" "crypto/tls" "crypto/x509" "errors" @@ -151,7 +152,8 @@ func (c *Client) connect() error { c.connected = true c.connChange(true, nil) - c.reader = bufio.NewReader(c.conn) + c.scan = bufio.NewScanner(c.conn) + c.scan.Buffer(c.recvBuf, cap(c.recvBuf)) c.register() @@ -185,8 +187,7 @@ func (c *Client) recv() { defer c.sendRecv.Done() for { - line, err := c.reader.ReadString('\n') - if err != nil { + if !c.scan.Scan() { select { case <-c.quit: return @@ -203,7 +204,12 @@ func (c *Client) recv() { } } - msg := parseMessage(line) + b := bytes.Trim(c.scan.Bytes(), " ") + if len(b) == 0 { + continue + } + + msg := ParseMessage(string(b)) if msg == nil { close(c.quit) c.connChange(false, ErrBadProtocol) diff --git a/pkg/irc/conn_test.go b/pkg/irc/conn_test.go index cee00680..45456b09 100644 --- a/pkg/irc/conn_test.go +++ b/pkg/irc/conn_test.go @@ -129,19 +129,21 @@ func TestRecv(t *testing.T) { buf.WriteString("CMD\r\n") buf.WriteString("PING :test\r\n") buf.WriteString("001 foo\r\n") - c.reader = bufio.NewReader(buf) + c.scan = bufio.NewScanner(buf) c.sendRecv.Add(1) go c.recv() assert.Equal(t, "PONG :test\r\n", <-conn.hook) assert.Equal(t, &Message{Command: "CMD"}, <-c.Messages) + assert.Equal(t, &Message{Command: Ping, Params: []string{"test"}}, <-c.Messages) + assert.Equal(t, &Message{Command: ReplyWelcome, Params: []string{"foo"}}, <-c.Messages) } func TestRecvTriggersReconnect(t *testing.T) { c := testClient() c.conn = &mockConn{} - c.reader = bufio.NewReader(bytes.NewBufferString("001 bob\r\n")) + c.scan = bufio.NewScanner(bytes.NewBufferString("001 bob\r\n")) done := make(chan struct{}) ok := false go func() { diff --git a/pkg/irc/message.go b/pkg/irc/message.go index 08cf045e..da109d41 100644 --- a/pkg/irc/message.go +++ b/pkg/irc/message.go @@ -22,8 +22,7 @@ func (m *Message) LastParam() string { return "" } -func parseMessage(line string) *Message { - line = strings.Trim(line, "\r\n ") +func ParseMessage(line string) *Message { msg := Message{} if strings.HasPrefix(line, "@") { @@ -35,21 +34,24 @@ func parseMessage(line string) *Message { if len(tags) > 0 { msg.Tags = map[string]string{} - } - for _, tag := range tags { - key, val := splitParam(tag) - if key == "" { - continue - } + for _, tag := range tags { + key, val := splitParam(tag) + if key == "" { + continue + } - if val != "" { - msg.Tags[key] = unescapeTag(val) - } else { - msg.Tags[key] = "" + if val != "" { + msg.Tags[key] = unescapeTag(val) + } else { + msg.Tags[key] = "" + } } } + for line[next+1] == ' ' { + next++ + } line = line[next+1:] } @@ -73,7 +75,7 @@ func parseMessage(line string) *Message { cmdEnd := len(line) trailing := "" - if i := strings.Index(line, " :"); i > 0 { + if i := strings.Index(line, " :"); i >= 0 { cmdEnd = i trailing = line[i+2:] } diff --git a/pkg/irc/message_test.go b/pkg/irc/message_test.go index cb923964..8c965daf 100644 --- a/pkg/irc/message_test.go +++ b/pkg/irc/message_test.go @@ -12,7 +12,7 @@ func TestParseMessage(t *testing.T) { expected *Message }{ { - ":user CMD #chan :some message\r\n", + ":user CMD #chan :some message", &Message{ Prefix: "user", Nick: "user", @@ -20,7 +20,7 @@ func TestParseMessage(t *testing.T) { Params: []string{"#chan", "some message"}, }, }, { - ":nick!user@host.com CMD a b\r\n", + ":nick!user@host.com CMD a b", &Message{ Prefix: "nick!user@host.com", Nick: "nick", @@ -28,80 +28,80 @@ func TestParseMessage(t *testing.T) { Params: []string{"a", "b"}, }, }, { - "CMD a b :\r\n", + "CMD a b :", &Message{ Command: "CMD", Params: []string{"a", "b", ""}, }, }, { - "CMD a b\r\n", + "CMD a b", &Message{ Command: "CMD", Params: []string{"a", "b"}, }, }, { - "CMD\r\n", + "CMD", &Message{ Command: "CMD", }, }, { - "CMD :tests and stuff\r\n", + "CMD :tests and stuff", &Message{ Command: "CMD", Params: []string{"tests and stuff"}, }, }, { - ":nick@host.com CMD\r\n", + ":nick@host.com CMD", &Message{ Prefix: "nick@host.com", Nick: "nick", Command: "CMD", }, }, { - ":ni@ck!user!name@host!.com CMD\r\n", + ":ni@ck!user!name@host!.com CMD", &Message{ Prefix: "ni@ck!user!name@host!.com", Nick: "ni@ck", Command: "CMD", }, }, { - "CMD #cake pie \r\n", + "CMD #cake pie ", &Message{ Command: "CMD", Params: []string{"#cake", "pie"}, }, }, { - " CMD #cake pie\r\n", + " CMD #cake pie", &Message{ Command: "CMD", Params: []string{"#cake", "pie"}, }, }, { - "CMD #cake ::pie\r\n", + "CMD #cake ::pie", &Message{ Command: "CMD", Params: []string{"#cake", ":pie"}, }, }, { - "CMD #cake : pie\r\n", + "CMD #cake : pie", &Message{ Command: "CMD", Params: []string{"#cake", " pie"}, }, }, { - "CMD #cake :pie :P <3\r\n", + "CMD #cake :pie :P <3", &Message{ Command: "CMD", Params: []string{"#cake", "pie :P <3"}, }, }, { - "CMD #cake :pie!\r\n", + "CMD #cake :pie!", &Message{ Command: "CMD", Params: []string{"#cake", "pie!"}, }, }, { - "@x=y CMD\r\n", + "@x=y CMD", &Message{ Tags: map[string]string{ "x": "y", @@ -109,7 +109,7 @@ func TestParseMessage(t *testing.T) { Command: "CMD", }, }, { - "@x=y :nick!user@host.com CMD\r\n", + "@x=y :nick!user@host.com CMD", &Message{ Tags: map[string]string{ "x": "y", @@ -119,7 +119,7 @@ func TestParseMessage(t *testing.T) { Command: "CMD", }, }, { - "@x=y :nick!user@host.com CMD :pie and cake\r\n", + "@x=y :nick!user@host.com CMD :pie and cake", &Message{ Tags: map[string]string{ "x": "y", @@ -130,7 +130,19 @@ func TestParseMessage(t *testing.T) { Params: []string{"pie and cake"}, }, }, { - "@x=y;a=b CMD\r\n", + "@x=y :nick!user@host.com CMD beans rainbows :pie and cake", + &Message{ + Tags: map[string]string{ + "x": "y", + }, + Prefix: "nick!user@host.com", + Nick: "nick", + Command: "CMD", + Params: []string{"beans", "rainbows", "pie and cake"}, + }, + }, + { + "@x=y;a=b CMD", &Message{ Tags: map[string]string{ "x": "y", @@ -139,7 +151,7 @@ func TestParseMessage(t *testing.T) { Command: "CMD", }, }, { - "@x=y;a=\\\\\\:\\s\\r\\n CMD\r\n", + "@x=y;a=\\\\\\:\\s\\r\\n CMD", &Message{ Tags: map[string]string{ "x": "y", @@ -151,23 +163,29 @@ func TestParseMessage(t *testing.T) { } for _, tc := range cases { - assert.Equal(t, tc.expected, parseMessage(tc.input)) + assert.Equal(t, tc.expected, ParseMessage(tc.input)) + } +} + +func BenchmarkParseMessage(b *testing.B) { + for i := 0; i < b.N; i++ { + ParseMessage("@x=y :nick!user@host.com CMD beans rainbows :pie and cake") } } func TestLastParam(t *testing.T) { - assert.Equal(t, "some message", parseMessage(":user CMD #chan :some message\r\n").LastParam()) - assert.Equal(t, "", parseMessage("NO_PARAMS").LastParam()) + assert.Equal(t, "some message", ParseMessage(":user CMD #chan :some message").LastParam()) + assert.Equal(t, "", ParseMessage("NO_PARAMS").LastParam()) } -func TestBadMessagePanic(t *testing.T) { - parseMessage("@\r\n") - parseMessage("@ :\r\n") - parseMessage("@ :\r\n") - parseMessage(":user\r\n") - parseMessage(":\r\n") - parseMessage(":") - parseMessage("") +func TestBadMessage(t *testing.T) { + assert.Nil(t, ParseMessage("@")) + assert.Nil(t, ParseMessage("@ :")) + assert.Nil(t, ParseMessage("@ :")) + assert.Nil(t, ParseMessage("@ :")) + assert.Nil(t, ParseMessage(":user")) + assert.Nil(t, ParseMessage(":")) + assert.Nil(t, ParseMessage("")) } func TestParseISupport(t *testing.T) {