Skip to content

Commit

Permalink
Remove common wsutil.Reader allocations (#189)
Browse files Browse the repository at this point in the history
![image](https://github.com/gobwas/ws/assets/5663952/b7f5a544-f6f0-44cf-b38c-e89fd3de0bd4)

* `ws.ReadHeader` allocates a 12 byte temporary buffer to read the header into. Since `io.ReadFull` is used, it escapes to the heap.

* `wsutil.NewCipherReader` is used for many calls.

Since the `NextFrame` shouldn't have concurrent calls, both of these should be possible to have internally.

I am forced to copy "readHeader" from `ws`, since there is not way to eliminate the alloc otherwise as far as I can tell.
  • Loading branch information
klauspost authored Oct 30, 2023
1 parent a6c9785 commit 516805a
Showing 1 changed file with 90 additions and 6 deletions.
96 changes: 90 additions & 6 deletions wsutil/reader.go
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
package wsutil

import (
"encoding/binary"
"errors"
"io"
"io/ioutil"
Expand Down Expand Up @@ -56,10 +57,12 @@ type Reader struct {
OnContinuation FrameHandlerFunc
OnIntermediate FrameHandlerFunc

opCode ws.OpCode // Used to store message op code on fragmentation.
frame io.Reader // Used to as frame reader.
raw io.LimitedReader // Used to discard frames without cipher.
utf8 UTF8Reader // Used to check UTF8 sequences if CheckUTF8 is true.
opCode ws.OpCode // Used to store message op code on fragmentation.
frame io.Reader // Used to as frame reader.
raw io.LimitedReader // Used to discard frames without cipher.
utf8 UTF8Reader // Used to check UTF8 sequences if CheckUTF8 is true.
tmp [ws.MaxHeaderSize - 2]byte // Used for reading headers.
cr *CipherReader // Used by NextFrame() to unmask frame payload.
}

// NewReader creates new frame reader that reads from r keeping given state to
Expand Down Expand Up @@ -165,7 +168,7 @@ func (r *Reader) Discard() (err error) {
// Note that next NextFrame() call must be done after receiving or discarding
// all current message bytes.
func (r *Reader) NextFrame() (hdr ws.Header, err error) {
hdr, err = ws.ReadHeader(r.Source)
hdr, err = r.readHeader(r.Source)
if err == io.EOF && r.fragmented() {
// If we are in fragmented state EOF means that is was totally
// unexpected.
Expand Down Expand Up @@ -196,7 +199,12 @@ func (r *Reader) NextFrame() (hdr ws.Header, err error) {

frame := io.Reader(&r.raw)
if hdr.Masked {
frame = NewCipherReader(frame, hdr.Mask)
if r.cr == nil {
r.cr = NewCipherReader(frame, hdr.Mask)
} else {
r.cr.Reset(frame, hdr.Mask)
}
frame = r.cr
}

for _, x := range r.Extensions {
Expand Down Expand Up @@ -261,6 +269,82 @@ func (r *Reader) reset() {
r.opCode = 0
}

// readHeader reads a frame header from in.
func (r *Reader) readHeader(in io.Reader) (h ws.Header, err error) {
// Make slice of bytes with capacity 12 that could hold any header.
//
// The maximum header size is 14, but due to the 2 hop reads,
// after first hop that reads first 2 constant bytes, we could reuse 2 bytes.
// So 14 - 2 = 12.
bts := r.tmp[:2]

// Prepare to hold first 2 bytes to choose size of next read.
_, err = io.ReadFull(in, bts)
if err != nil {
return h, err
}
const bit0 = 0x80

h.Fin = bts[0]&bit0 != 0
h.Rsv = (bts[0] & 0x70) >> 4
h.OpCode = ws.OpCode(bts[0] & 0x0f)

var extra int

if bts[1]&bit0 != 0 {
h.Masked = true
extra += 4
}

length := bts[1] & 0x7f
switch {
case length < 126:
h.Length = int64(length)

case length == 126:
extra += 2

case length == 127:
extra += 8

default:
err = ws.ErrHeaderLengthUnexpected
return h, err
}

if extra == 0 {
return h, err
}

// Increase len of bts to extra bytes need to read.
// Overwrite first 2 bytes that was read before.
bts = bts[:extra]
_, err = io.ReadFull(in, bts)
if err != nil {
return h, err
}

switch {
case length == 126:
h.Length = int64(binary.BigEndian.Uint16(bts[:2]))
bts = bts[2:]

case length == 127:
if bts[0]&0x80 != 0 {
err = ws.ErrHeaderLengthMSB
return h, err
}
h.Length = int64(binary.BigEndian.Uint64(bts[:8]))
bts = bts[8:]
}

if h.Masked {
copy(h.Mask[:], bts)
}

return h, nil
}

// NextReader prepares next message read from r. It returns header that
// describes the message and io.Reader to read message's payload. It returns
// non-nil error when it is not possible to read message's initial frame.
Expand Down

0 comments on commit 516805a

Please sign in to comment.