GitHub/ergo
3
0
Fork 0
mirror of https://github.com/ergochat/ergo.git synced 2025-02-02 07:34:06 +01:00
Code Issues Projects Releases Wiki Activity

Merge pull request #2038 from slingamn/ws_optimization.1

Browse Source 
tweaks to websocket handling
This commit is contained in:
Shivaram Lingamneni 2023-01-23 04:25:14 -08:00 committed by GitHub
commit 5ecf19d01e
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 45 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

58
irc/ircconn.go
View File

@ -5,6 +5,7 @@ package irc
import ( import (
"bytes" "bytes"
"io"
"net" "net"
"unicode/utf8" "unicode/utf8"
@ -93,21 +94,25 @@ func (cc *IRCStreamConn) Close() (err error) {
// IRCWSConn is an IRCConn over a websocket. // IRCWSConn is an IRCConn over a websocket.
type IRCWSConn struct { type IRCWSConn struct {
conn *websocket.Conn conn *websocket.Conn
buf []byte
binary bool binary bool
} }
func NewIRCWSConn(conn *websocket.Conn) IRCWSConn { func NewIRCWSConn(conn *websocket.Conn) *IRCWSConn {
binary := conn.Subprotocol() == "binary.ircv3.net" return &IRCWSConn{
return IRCWSConn{conn: conn, binary: binary} conn: conn,
binary: conn.Subprotocol() == "binary.ircv3.net",
buf: make([]byte, initialBufferSize),
}
} }
func (wc IRCWSConn) UnderlyingConn() *utils.WrappedConn { func (wc *IRCWSConn) UnderlyingConn() *utils.WrappedConn {
// just assume that the type is OK // just assume that the type is OK
wConn, _ := wc.conn.UnderlyingConn().(*utils.WrappedConn) wConn, _ := wc.conn.UnderlyingConn().(*utils.WrappedConn)
return wConn return wConn
} }
func (wc IRCWSConn) WriteLine(buf []byte) (err error) { func (wc *IRCWSConn) WriteLine(buf []byte) (err error) {
buf = bytes.TrimSuffix(buf, crlf) buf = bytes.TrimSuffix(buf, crlf)
// #1483: if we have websockets at all, then we're enforcing utf8 // #1483: if we have websockets at all, then we're enforcing utf8
messageType := websocket.TextMessage messageType := websocket.TextMessage
@ -117,7 +122,7 @@ func (wc IRCWSConn) WriteLine(buf []byte) (err error) {
return wc.conn.WriteMessage(messageType, buf) 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 { for _, buf := range buffers {
err = wc.WriteLine(buf) err = wc.WriteLine(buf)
if err != nil { if err != nil {
@ -127,20 +132,47 @@ func (wc IRCWSConn) WriteLines(buffers [][]byte) (err error) {
return return
} }
func (wc IRCWSConn) ReadLine() (line []byte, err error) { func (wc *IRCWSConn) ReadLine() (line []byte, err error) {
messageType, line, err := wc.conn.ReadMessage() _, reader, err := wc.conn.NextReader()
if err == nil { switch err {
if messageType == websocket.BinaryMessage && !utf8.Valid(line) { case nil:
// OK
case websocket.ErrReadLimit:
return line, ircreader.ErrReadQ
default:
return line, err
}
line, err = wc.readFull(reader)
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 !utf8.Valid(line) {
return line, errInvalidUtf8 return line, errInvalidUtf8
} }
return line, nil 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 return line, ircreader.ErrReadQ
} else { default:
return line, err return line, err
} }
} }
func (wc IRCWSConn) Close() (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() return wc.conn.Close()
} }