From 5eaf7b37e5330c75ceab05cbd6ecaa108d78e96f Mon Sep 17 00:00:00 2001 From: Shivaram Lingamneni Date: Sat, 21 Jan 2023 19:10:17 -0500 Subject: [PATCH 1/3] reduce websocket read allocations See #2037 --- irc/ircconn.go | 44 ++++++++++++++++++++++++++++++++------------ 1 file changed, 32 insertions(+), 12 deletions(-) diff --git a/irc/ircconn.go b/irc/ircconn.go index 088909a2..f0fe947a 100644 --- a/irc/ircconn.go +++ b/irc/ircconn.go @@ -5,6 +5,7 @@ package irc import ( "bytes" + "io" "net" "unicode/utf8" @@ -93,21 +94,25 @@ func (cc *IRCStreamConn) Close() (err error) { // IRCWSConn is an IRCConn over a websocket. type IRCWSConn struct { conn *websocket.Conn + buf []byte binary bool } -func NewIRCWSConn(conn *websocket.Conn) IRCWSConn { - binary := conn.Subprotocol() == "binary.ircv3.net" - return IRCWSConn{conn: conn, binary: binary} +func NewIRCWSConn(conn *websocket.Conn) *IRCWSConn { + return &IRCWSConn{ + conn: conn, + binary: conn.Subprotocol() == "binary.ircv3.net", + buf: make([]byte, maxReadQBytes()), + } } -func (wc IRCWSConn) UnderlyingConn() *utils.WrappedConn { +func (wc *IRCWSConn) UnderlyingConn() *utils.WrappedConn { // just assume that the type is OK wConn, _ := wc.conn.UnderlyingConn().(*utils.WrappedConn) return wConn } -func (wc IRCWSConn) WriteLine(buf []byte) (err error) { +func (wc *IRCWSConn) WriteLine(buf []byte) (err error) { buf = bytes.TrimSuffix(buf, crlf) // #1483: if we have websockets at all, then we're enforcing utf8 messageType := websocket.TextMessage @@ -117,7 +122,7 @@ func (wc IRCWSConn) WriteLine(buf []byte) (err error) { return wc.conn.WriteMessage(messageType, buf) } -func (wc IRCWSConn) WriteLines(buffers [][]byte) (err error) { +func (wc *IRCWSConn) WriteLines(buffers [][]byte) (err error) { for _, buf := range buffers { err = wc.WriteLine(buf) if err != nil { @@ -127,20 +132,35 @@ func (wc IRCWSConn) WriteLines(buffers [][]byte) (err error) { return } -func (wc IRCWSConn) ReadLine() (line []byte, err error) { - messageType, line, err := wc.conn.ReadMessage() - if err == nil { +func (wc *IRCWSConn) ReadLine() (line []byte, err error) { + messageType, reader, err := wc.conn.NextReader() + switch err { + case nil: + // OK + case websocket.ErrReadLimit: + return line, ircreader.ErrReadQ + default: + return line, err + } + + n, err := io.ReadFull(reader, wc.buf) + line = wc.buf[:n] + switch err { + case io.ErrUnexpectedEOF, io.EOF: + // these are OK. io.ErrUnexpectedEOF is the good case: + // it means we read the full message and it consumed less than the full wc.buf if messageType == websocket.BinaryMessage && !utf8.Valid(line) { return line, errInvalidUtf8 } return line, nil - } else if err == websocket.ErrReadLimit { + case nil, websocket.ErrReadLimit: + // nil means we filled wc.buf without exhausting the reader: return line, ircreader.ErrReadQ - } else { + default: return line, err } } -func (wc IRCWSConn) Close() (err error) { +func (wc *IRCWSConn) Close() (err error) { return wc.conn.Close() } From 9439e9b9e14c566bda4111ce62109dacd6453ecf Mon Sep 17 00:00:00 2001 From: Shivaram Lingamneni Date: Sat, 21 Jan 2023 19:10:25 -0500 Subject: [PATCH 2/3] allow resizing the ws read buffer --- irc/ircconn.go | 18 +++++++++++++++--- 1 file changed, 15 insertions(+), 3 deletions(-) diff --git a/irc/ircconn.go b/irc/ircconn.go index f0fe947a..8deb2a81 100644 --- a/irc/ircconn.go +++ b/irc/ircconn.go @@ -102,7 +102,7 @@ func NewIRCWSConn(conn *websocket.Conn) *IRCWSConn { return &IRCWSConn{ conn: conn, binary: conn.Subprotocol() == "binary.ircv3.net", - buf: make([]byte, maxReadQBytes()), + buf: make([]byte, initialBufferSize), } } @@ -143,8 +143,7 @@ func (wc *IRCWSConn) ReadLine() (line []byte, err error) { return line, err } - n, err := io.ReadFull(reader, wc.buf) - line = wc.buf[:n] + line, err = wc.readFull(reader) switch err { case io.ErrUnexpectedEOF, io.EOF: // these are OK. io.ErrUnexpectedEOF is the good case: @@ -161,6 +160,19 @@ func (wc *IRCWSConn) ReadLine() (line []byte, err error) { } } +func (wc *IRCWSConn) readFull(reader io.Reader) (line []byte, err error) { + // XXX this is io.ReadFull with a single attempt to resize upwards + n, err := io.ReadFull(reader, wc.buf) + if err == nil && len(wc.buf) < maxReadQBytes() { + newBuf := make([]byte, maxReadQBytes()) + copy(newBuf, wc.buf[:n]) + wc.buf = newBuf + n2, err := io.ReadFull(reader, wc.buf[n:]) + return wc.buf[:n+n2], err + } + return wc.buf[:n], err +} + func (wc *IRCWSConn) Close() (err error) { return wc.conn.Close() } From abc71684f30ab6e0479549caf91b674096255a2f Mon Sep 17 00:00:00 2001 From: Shivaram Lingamneni Date: Sun, 22 Jan 2023 14:42:58 -0500 Subject: [PATCH 3/3] always validate UTF8 from websockets --- irc/ircconn.go | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/irc/ircconn.go b/irc/ircconn.go index 8deb2a81..baeeb5cf 100644 --- a/irc/ircconn.go +++ b/irc/ircconn.go @@ -133,7 +133,7 @@ func (wc *IRCWSConn) WriteLines(buffers [][]byte) (err error) { } func (wc *IRCWSConn) ReadLine() (line []byte, err error) { - messageType, reader, err := wc.conn.NextReader() + _, reader, err := wc.conn.NextReader() switch err { case nil: // OK @@ -148,7 +148,7 @@ func (wc *IRCWSConn) ReadLine() (line []byte, err error) { case io.ErrUnexpectedEOF, io.EOF: // these are OK. io.ErrUnexpectedEOF is the good case: // it means we read the full message and it consumed less than the full wc.buf - if messageType == websocket.BinaryMessage && !utf8.Valid(line) { + if !utf8.Valid(line) { return line, errInvalidUtf8 } return line, nil