Skip to content

Commit 286b5c9

Browse files
committed
Use bufio.Reader returned from hijack in upgrade
Use the bufio.Reader returned from hijack if the reader's buffer size is equal to the buffer size specified in Upgrader.ReadBufferSize.
1 parent 3f3e394 commit 286b5c9

File tree

3 files changed

+38
-7
lines changed

3 files changed

+38
-7
lines changed

conn.go

Lines changed: 20 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -265,6 +265,10 @@ type Conn struct {
265265
}
266266

267267
func newConn(conn net.Conn, isServer bool, readBufferSize, writeBufferSize int) *Conn {
268+
return newConnBRW(conn, isServer, readBufferSize, writeBufferSize, nil)
269+
}
270+
271+
func newConnBRW(conn net.Conn, isServer bool, readBufferSize, writeBufferSize int, brw *bufio.ReadWriter) *Conn {
268272
mu := make(chan bool, 1)
269273
mu <- true
270274

@@ -274,13 +278,28 @@ func newConn(conn net.Conn, isServer bool, readBufferSize, writeBufferSize int)
274278
if readBufferSize < maxControlFramePayloadSize {
275279
readBufferSize = maxControlFramePayloadSize
276280
}
281+
282+
// Reuse the supplied brw.Reader if brw.Reader's buf is the requested size.
283+
var br *bufio.Reader
284+
if brw != nil && brw.Reader != nil {
285+
// This code assumes that peek on a reset reader returns
286+
// bufio.Reader.buf[:0].
287+
brw.Reader.Reset(conn)
288+
if p, err := brw.Reader.Peek(0); err == nil && cap(p) == readBufferSize {
289+
br = brw.Reader
290+
}
291+
}
292+
if br == nil {
293+
br = bufio.NewReaderSize(conn, readBufferSize)
294+
}
295+
277296
if writeBufferSize == 0 {
278297
writeBufferSize = defaultWriteBufferSize
279298
}
280299

281300
c := &Conn{
282301
isServer: isServer,
283-
br: bufio.NewReaderSize(conn, readBufferSize),
302+
br: br,
284303
conn: conn,
285304
mu: mu,
286305
readFinal: true,

conn_test.go

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -463,3 +463,17 @@ func TestFailedConnectionReadPanic(t *testing.T) {
463463
}
464464
t.Fatal("should not get here")
465465
}
466+
467+
func TestBufioReaderReuse(t *testing.T) {
468+
brw := bufio.NewReadWriter(bufio.NewReader(nil), nil)
469+
c := newConnBRW(nil, false, 0, 0, brw)
470+
if c.br != brw.Reader {
471+
t.Error("connection did not reuse bufio.Reader")
472+
}
473+
474+
brw = bufio.NewReadWriter(bufio.NewReaderSize(nil, 1234), nil) // size must not equal bufio.defaultBufSize
475+
c = newConnBRW(nil, false, 0, 0, brw)
476+
if c.br == brw.Reader {
477+
t.Error("connection reuse bufio.Reader with wrong size")
478+
}
479+
}

server.go

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -152,27 +152,25 @@ func (u *Upgrader) Upgrade(w http.ResponseWriter, r *http.Request, responseHeade
152152

153153
var (
154154
netConn net.Conn
155-
br *bufio.Reader
156155
err error
157156
)
158157

159158
h, ok := w.(http.Hijacker)
160159
if !ok {
161160
return u.returnError(w, r, http.StatusInternalServerError, "websocket: response does not implement http.Hijacker")
162161
}
163-
var rw *bufio.ReadWriter
164-
netConn, rw, err = h.Hijack()
162+
var brw *bufio.ReadWriter
163+
netConn, brw, err = h.Hijack()
165164
if err != nil {
166165
return u.returnError(w, r, http.StatusInternalServerError, err.Error())
167166
}
168-
br = rw.Reader
169167

170-
if br.Buffered() > 0 {
168+
if brw.Reader.Buffered() > 0 {
171169
netConn.Close()
172170
return nil, errors.New("websocket: client sent data before handshake is complete")
173171
}
174172

175-
c := newConn(netConn, true, u.ReadBufferSize, u.WriteBufferSize)
173+
c := newConnBRW(netConn, true, u.ReadBufferSize, u.WriteBufferSize, brw)
176174
c.subprotocol = subprotocol
177175

178176
if compress {

0 commit comments

Comments
 (0)