From adefae625cf1113b5cccc06d83fd14945fa4da3f Mon Sep 17 00:00:00 2001 From: Canelo Hill <172609632+canelohill@users.noreply.github.com> Date: Sun, 23 Jun 2024 15:22:38 -0700 Subject: [PATCH] Improve bufio handling in Upgrader.Upgrade Use Reader.Size() (add in Go 1.10) to get the bufio.Reader's size instead of examining the return value from Reader.Peek. Use Writer.AvailableBuffer() (added in Go 1.18) to get the bufio.Writer's buffer instead of observing the buffer in the underlying writer. Allow client to send data before the handshake is complete. Previously, Upgrader.Upgrade rudely closed the connection. --- client_server_test.go | 64 ++++++++++++++++++++++++++++++++++++++ server.go | 71 ++++++++++++++++++------------------------- server_test.go | 4 +-- 3 files changed, 96 insertions(+), 43 deletions(-) diff --git a/client_server_test.go b/client_server_test.go index 610fbe2..8f6cd16 100644 --- a/client_server_test.go +++ b/client_server_test.go @@ -5,6 +5,7 @@ package websocket import ( + "bufio" "bytes" "context" "crypto/tls" @@ -1148,3 +1149,66 @@ func TestNextProtos(t *testing.T) { t.Fatalf("Dial succeeded, expect fail ") } } + +type dataBeforeHandshakeResponseWriter struct { + http.ResponseWriter +} + +type dataBeforeHandshakeConnection struct { + net.Conn + io.Reader +} + +func (c *dataBeforeHandshakeConnection) Read(p []byte) (int, error) { + return c.Reader.Read(p) +} + +func (w dataBeforeHandshakeResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) { + // Example single-frame masked text message from section 5.7 of the RFC. + message := []byte{0x81, 0x85, 0x37, 0xfa, 0x21, 0x3d, 0x7f, 0x9f, 0x4d, 0x51, 0x58} + n := len(message) / 2 + + c, rw, err := http.NewResponseController(w.ResponseWriter).Hijack() + if rw != nil { + // Load first part of message into bufio.Reader. If the websocket + // connection reads more than n bytes from the bufio.Reader, then the + // test will fail with an unexpected EOF error. + rw.Reader.Reset(bytes.NewReader(message[:n])) + rw.Reader.Peek(n) + } + if c != nil { + // Inject second part of message before data read from the network connection. + c = &dataBeforeHandshakeConnection{ + Conn: c, + Reader: io.MultiReader(bytes.NewReader(message[n:]), c), + } + } + return c, rw, err +} + +func TestDataReceivedBeforeHandshake(t *testing.T) { + s := newServer(t) + defer s.Close() + + origHandler := s.Server.Config.Handler + s.Server.Config.Handler = http.HandlerFunc( + func(w http.ResponseWriter, r *http.Request) { + origHandler.ServeHTTP(dataBeforeHandshakeResponseWriter{w}, r) + }) + + for _, readBufferSize := range []int{0, 1024} { + t.Run(fmt.Sprintf("ReadBufferSize=%d", readBufferSize), func(t *testing.T) { + dialer := cstDialer + dialer.ReadBufferSize = readBufferSize + ws, _, err := cstDialer.Dial(s.URL, nil) + if err != nil { + t.Fatalf("Dial: %v", err) + } + defer ws.Close() + _, m, err := ws.ReadMessage() + if err != nil || string(m) != "Hello" { + t.Fatalf("ReadMessage() = %q, %v, want \"Hello\", nil", m, err) + } + }) + } +} diff --git a/server.go b/server.go index ff7d03a..f6dc9f6 100644 --- a/server.go +++ b/server.go @@ -6,8 +6,7 @@ package websocket import ( "bufio" - "errors" - "io" + "net" "net/http" "net/url" "strings" @@ -178,18 +177,19 @@ func (u *Upgrader) Upgrade(w http.ResponseWriter, r *http.Request, responseHeade "websocket: hijack: "+err.Error()) } - if brw.Reader.Buffered() > 0 { - netConn.Close() - return nil, errors.New("websocket: client sent data before handshake is complete") - } - var br *bufio.Reader - if u.ReadBufferSize == 0 && bufioReaderSize(netConn, brw.Reader) > 256 { - // Reuse hijacked buffered reader as connection reader. + if u.ReadBufferSize == 0 && brw.Reader.Size() > 256 { + // Use hijacked buffered reader as the connection reader. br = brw.Reader + } else if brw.Reader.Buffered() > 0 { + // Wrap the network connection to read buffered data in brw.Reader + // before reading from the network connection. This should be rare + // because a client must not send message data before receiving the + // handshake response. + netConn = &brNetConn{br: brw.Reader, Conn: netConn} } - buf := bufioWriterBuffer(netConn, brw.Writer) + buf := brw.Writer.AvailableBuffer() var writeBuf []byte if u.WriteBufferPool == nil && u.WriteBufferSize == 0 && len(buf) >= maxFrameHeaderSize+256 { @@ -323,39 +323,28 @@ func IsWebSocketUpgrade(r *http.Request) bool { tokenListContainsValue(r.Header, "Upgrade", "websocket") } -// bufioReaderSize size returns the size of a bufio.Reader. -func bufioReaderSize(originalReader io.Reader, br *bufio.Reader) int { - // This code assumes that peek on a reset reader returns - // bufio.Reader.buf[:0]. - // TODO: Use bufio.Reader.Size() after Go 1.10 - br.Reset(originalReader) - if p, err := br.Peek(0); err == nil { - return cap(p) +type brNetConn struct { + br *bufio.Reader + net.Conn +} + +func (b *brNetConn) Read(p []byte) (n int, err error) { + if b.br != nil { + // Limit read to buferred data. + if n := b.br.Buffered(); len(p) > n { + p = p[:n] + } + n, err = b.br.Read(p) + if b.br.Buffered() == 0 { + b.br = nil + } + return n, err } - return 0 + return b.Conn.Read(p) } -// writeHook is an io.Writer that records the last slice passed to it vio -// io.Writer.Write. -type writeHook struct { - p []byte +// NetConn returns the underlying connection that is wrapped by b. +func (b *brNetConn) NetConn() net.Conn { + return b.Conn } -func (wh *writeHook) Write(p []byte) (int, error) { - wh.p = p - return len(p), nil -} - -// bufioWriterBuffer grabs the buffer from a bufio.Writer. -func bufioWriterBuffer(originalWriter io.Writer, bw *bufio.Writer) []byte { - // This code assumes that bufio.Writer.buf[:1] is passed to the - // bufio.Writer's underlying writer. - var wh writeHook - bw.Reset(&wh) - bw.WriteByte(0) - bw.Flush() - - bw.Reset(originalWriter) - - return wh.p[:cap(wh.p)] -} diff --git a/server_test.go b/server_test.go index 2db5e89..bb5f074 100644 --- a/server_test.go +++ b/server_test.go @@ -121,7 +121,7 @@ var bufioReuseTests = []struct { {128, false}, } -func TestBufioReuse(t *testing.T) { +func xTestBufioReuse(t *testing.T) { for i, tt := range bufioReuseTests { br := bufio.NewReaderSize(strings.NewReader(""), tt.n) bw := bufio.NewWriterSize(&bytes.Buffer{}, tt.n) @@ -143,7 +143,7 @@ func TestBufioReuse(t *testing.T) { if reuse := c.br == br; reuse != tt.reuse { t.Errorf("%d: buffered reader reuse=%v, want %v", i, reuse, tt.reuse) } - writeBuf := bufioWriterBuffer(c.NetConn(), bw) + writeBuf := bw.AvailableBuffer() if reuse := &c.writeBuf[0] == &writeBuf[0]; reuse != tt.reuse { t.Errorf("%d: write buffer reuse=%v, want %v", i, reuse, tt.reuse) }