Skip to content

Commit 0fa48a5

Browse files
committed
Reduce Reader allocation by 8 bytes
Both Reader and Writer now will only ever allocate 16 bytes for their entire usage :) This is the minimum possible while still preventing misuse of a EOFed Reader or a closed Writer.
1 parent 1e0955c commit 0fa48a5

File tree

1 file changed

+11
-12
lines changed

1 file changed

+11
-12
lines changed

websocket.go

Lines changed: 11 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -36,9 +36,6 @@ type Conn struct {
3636
closer io.Closer
3737
client bool
3838

39-
// read limit for a message in bytes.
40-
msgReadLimit int64
41-
4239
closeOnce sync.Once
4340
closeErr error
4441
closed chan struct{}
@@ -50,10 +47,13 @@ type Conn struct {
5047
writeFrameLock chan struct{}
5148
writeHeaderBuf []byte
5249
writeHeader *header
50+
// read limit for a message in bytes.
51+
msgReadLimit int64
5352

5453
// messageWriter state.
5554
writeMsgOpcode opcode
5655
writeMsgCtx context.Context
56+
readMsgLeft int64
5757

5858
// Used to ensure the previous reader is read till EOF before allowing
5959
// a new one.
@@ -371,20 +371,19 @@ func (c *Conn) reader(ctx context.Context) (MessageType, io.Reader, error) {
371371
c.readMsgHeader = h
372372
c.readFrameEOF = false
373373
c.readMaskPos = 0
374+
c.readMsgLeft = c.msgReadLimit
374375

375376
r := &messageReader{
376-
c: c,
377-
left: c.msgReadLimit,
377+
c: c,
378378
}
379379
c.previousReader = r
380380
return MessageType(h.opcode), r, nil
381381
}
382382

383383
// messageReader enables reading a data frame from the WebSocket connection.
384384
type messageReader struct {
385-
c *Conn
386-
left int64
387-
eof bool
385+
c *Conn
386+
eof bool
388387
}
389388

390389
// Read reads as many bytes as possible into p.
@@ -406,14 +405,14 @@ func (r *messageReader) read(p []byte) (int, error) {
406405
return 0, xerrors.Errorf("cannot use EOFed reader")
407406
}
408407

409-
if r.left <= 0 {
408+
if r.c.readMsgLeft <= 0 {
410409
err := xerrors.Errorf("read limited at %v bytes", r.c.msgReadLimit)
411410
r.c.Close(StatusMessageTooBig, err.Error())
412411
return 0, err
413412
}
414413

415-
if int64(len(p)) > r.left {
416-
p = p[:r.left]
414+
if int64(len(p)) > r.c.readMsgLeft {
415+
p = p[:r.c.readMsgLeft]
417416
}
418417

419418
if r.c.readFrameEOF {
@@ -441,7 +440,7 @@ func (r *messageReader) read(p []byte) (int, error) {
441440
n, err := r.c.readFramePayload(r.c.readMsgCtx, p)
442441

443442
h.payloadLength -= int64(n)
444-
r.left -= int64(n)
443+
r.c.readMsgLeft -= int64(n)
445444
if h.masked {
446445
r.c.readMaskPos = fastXOR(h.maskKey, r.c.readMaskPos, p)
447446
}

0 commit comments

Comments
 (0)