proto_test.go

package websocket

import (
	"bytes"
	"fmt"
	"testing"

	"github.com/mccutchen/websocket/internal/testing/assert"
)

func TestFrameRoundTrip(t *testing.T) {
	// Basic test to ensure that we can read back the same frame that we
	// write.
	t.Parallel()

	// write masked "client" frame to buffer
	clientFrame := NewFrame(OpcodeText, true, []byte("hello"))
	buf := &bytes.Buffer{}
	assert.NilError(t, WriteFrame(buf, NewMaskingKey(), clientFrame))

	// read "server" frame from buffer.
	serverFrame, err := ReadFrame(buf, ServerMode, len(clientFrame.Payload))
	assert.NilError(t, err)

	// ensure client and server frame match
	assert.Equal(t, serverFrame, clientFrame, "server and client frame mismatch")
}

func TestMaxFrameSize(t *testing.T) {
	// Basic test to ensure that we can read back the same frame that we
	// write.
	t.Parallel()

	// write masked "client" frame to buffer
	clientFrame := NewFrame(OpcodeText, true, []byte("hello"))
	buf := &bytes.Buffer{}
	assert.NilError(t, WriteFrame(buf, NewMaskingKey(), clientFrame))

	// read "server" frame from buffer.
	serverFrame, err := ReadFrame(buf, ServerMode, len(clientFrame.Payload)-1)
	assert.Error(t, err, ErrFrameTooLarge)
	assert.Equal(t, serverFrame, nil, "expected nil frame on error")
}

func TestRSV(t *testing.T) {
	// We don't currently support any extensions, so RSV bits are not allowed.
	// But we still need to be able properly parse and marshal them.
	marshalledFrame := func(rsvBits ...RSVBit) []byte {
		buf := &bytes.Buffer{}
		frame := NewFrame(OpcodeText, true, nil, rsvBits...)
		assert.NilError(t, WriteFrame(buf, Unmasked, frame))
		return buf.Bytes()
	}

	testCases := map[string]struct {
		rawBytes []byte
		wantRSV1 bool
		wantRSV2 bool
		wantRSV3 bool
	}{
		"no RSV bits set": {
			rawBytes: marshalledFrame(),
		},
		"RSV1 set": {
			rawBytes: marshalledFrame(RSV1),
			wantRSV1: true,
		},
		"RSV2 set": {
			rawBytes: marshalledFrame(RSV2),
			wantRSV2: true,
		},
		"RSV3 set": {
			rawBytes: marshalledFrame(RSV3),
			wantRSV3: true,
		},
		"all RSV bits set": {
			rawBytes: marshalledFrame(RSV1, RSV2, RSV3),
			wantRSV1: true,
			wantRSV2: true,
			wantRSV3: true,
		},
	}
	for name, tc := range testCases {
		t.Run(name, func(t *testing.T) {
			t.Parallel()
			buf := bytes.NewReader(tc.rawBytes)
			frame := mustReadFrame(t, buf, len(tc.rawBytes))
			assert.Equal(t, frame.RSV1(), tc.wantRSV1, "incorrect RSV1")
			assert.Equal(t, frame.RSV2(), tc.wantRSV2, "incorrect RSV2")
			assert.Equal(t, frame.RSV3(), tc.wantRSV3, "incorrect RSV3")
		})
	}
}

func TestExampleFramesFromRFC(t *testing.T) {
	// This tests every example provided in RFC 6455 section 5.7:
	// https://datatracker.ietf.org/doc/html/rfc6455#section-5.7
	testCases := map[string]struct {
		rawBytes  []byte
		wantFrame *Frame
	}{
		"single-frame unmasked text": {
			rawBytes:  []byte{0x81, 0x05, 0x48, 0x65, 0x6c, 0x6c, 0x6f},
			wantFrame: NewFrame(OpcodeText, true, []byte("Hello")),
		},
		"single-frame masked text": {
			rawBytes:  []byte{0x81, 0x85, 0x37, 0xfa, 0x21, 0x3d, 0x7f, 0x9f, 0x4d, 0x51, 0x58},
			wantFrame: NewFrame(OpcodeText, true, []byte("Hello")),
		},
		"fragmented unmasked text part 1": {
			rawBytes:  []byte{0x01, 0x03, 0x48, 0x65, 0x6c},
			wantFrame: NewFrame(OpcodeText, false, []byte("Hel")),
		},
		"fragmented unmasked text part 2": {
			rawBytes:  []byte{0x80, 0x02, 0x6c, 0x6f},
			wantFrame: NewFrame(OpcodeContinuation, true, []byte("lo")),
		},
		"unmasked ping": {
			rawBytes: []byte{
				0x89, 0x05, 0x48, 0x65, 0x6c, 0x6c, 0x6f,
			},
			wantFrame: NewFrame(OpcodePing, true, []byte("Hello")),
		},
		"masked ping response": {
			rawBytes:  []byte{0x8a, 0x85, 0x37, 0xfa, 0x21, 0x3d, 0x7f, 0x9f, 0x4d, 0x51, 0x58},
			wantFrame: NewFrame(OpcodePong, true, []byte("Hello")),
		},
		"256 bytes binary message": {
			rawBytes: append(
				[]byte{0x82, 0x7E, 0x01, 0x00},
				make([]byte, 256)...,
			),
			wantFrame: NewFrame(OpcodeBinary, true, make([]byte, 256)),
		},
		"64KiB binary message": {
			rawBytes: append(
				[]byte{0x82, 0x7F, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00},
				make([]byte, 65536)...,
			),
			wantFrame: NewFrame(OpcodeBinary, true, make([]byte, 65536)),
		},
	}

	for name, tc := range testCases {
		t.Run(name, func(t *testing.T) {
			t.Parallel()
			buf := bytes.NewReader(tc.rawBytes)
			got := mustReadFrame(t, buf, len(tc.rawBytes))
			assert.Equal(t, got, tc.wantFrame, "frames do not match")
		})
	}
}

func TestIncompleteFrames(t *testing.T) {
	testCases := map[string]struct {
		rawBytes []byte
		wantErr  string
	}{
		"2-byte extended payload can't be read": {
			rawBytes: []byte{0x82, 0x7E},
			wantErr:  "error reading 2-byte extended payload length: EOF",
		},
		"8-byte extended payload can't be read": {
			rawBytes: []byte{0x82, 0x7F},
			wantErr:  "error reading 8-byte extended payload length: EOF",
		},
		"mask can't be read": {
			rawBytes: []byte{0x81, 0x85},
			wantErr:  "error reading mask key: EOF",
		},
		"payload can't be read": {
			rawBytes: []byte{0x81, 0x05},
			wantErr:  "error reading 5 byte payload: EOF",
		},
	}

	for name, tc := range testCases {
		t.Run(name, func(t *testing.T) {
			t.Parallel()
			buf := bytes.NewReader(tc.rawBytes)
			_, err := ReadFrame(buf, ClientMode, 70000)
			assert.Error(t, err, tc.wantErr)
		})
	}
}

// ============================================================================
// Fuzzers
// ============================================================================
func FuzzReadFrame(f *testing.F) {
	// Set up seed corpus
	var testCases [][]byte
	// Example frames from RFC 6455 section 5.7
	// https://datatracker.ietf.org/doc/html/rfc6455#section-5.7
	testCases = append(testCases, [][]byte{
		// single-frame unmasked text
		{0x81, 0x05, 0x48, 0x65, 0x6c, 0x6c, 0x6f},
		// single-frame masked text
		{0x81, 0x85, 0x37, 0xfa, 0x21, 0x3d, 0x7f, 0x9f, 0x4d, 0x51, 0x58},
		// fragmented unmasked text part 1
		{0x01, 0x03, 0x48, 0x65, 0x6c},
		// fragmented unmasked text part 2
		{0x80, 0x02, 0x6c, 0x6f},
		// unmasked ping
		{0x89, 0x05, 0x48, 0x65, 0x6c, 0x6c, 0x6f},
		// masked ping response
		{0x8a, 0x85, 0x37, 0xfa, 0x21, 0x3d, 0x7f, 0x9f, 0x4d, 0x51, 0x58},
		// 256 bytes binary message
		append([]byte{0x82, 0x7E, 0x01, 0x00}, make([]byte, 256)...),
		// 64KiB binary message
		append([]byte{0x82, 0x7F, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00}, make([]byte, 65536)...),
	}...)

	for _, tc := range testCases {
		f.Add(tc)
	}

	f.Fuzz(func(t *testing.T, input []byte) {
		modes := []Mode{ClientMode, ServerMode}
		for _, mode := range modes {
			frame, err := ReadFrame(bytes.NewReader(input), mode, 1<<20)
			if err != nil {
				t.Skipf("skipping eror: %s", err)
				return
			}
			t.Logf("frame: %s", frame)
		}
	})
}

// ============================================================================
// Benchmarks
// ============================================================================
var benchMarkFrameSizes = []int{
	// test odd sizes to cover edge cases (e.g. in applyMask's optimization)
	(1 << 10) + 1, // ~1 KiB
	(1 << 20) + 1, // ~1 MiB
}

func BenchmarkReadFrame(b *testing.B) {
	for _, size := range benchMarkFrameSizes {
		frame := makeFrame(OpcodeText, true, size)
		buf := &bytes.Buffer{}
		assert.NilError(b, WriteFrame(buf, NewMaskingKey(), frame))

		// Run sub-benchmarks for each payload size
		b.Run(formatSize(size), func(b *testing.B) {
			src := bytes.NewReader(buf.Bytes())
			b.SetBytes(int64(buf.Len()))
			b.ResetTimer()
			for b.Loop() {
				_, _ = src.Seek(0, 0)
				frame2, err := ReadFrame(src, ServerMode, size)
				if err != nil {
					b.Fatalf("unexpected error: %v", err)
				}
				assert.Equal(b, len(frame2.Payload), len(frame.Payload), "payload length")
			}
		})
	}
}

func BenchmarkWriteFrame(b *testing.B) {
	for _, size := range benchMarkFrameSizes {
		b.Run(formatSize(size), func(b *testing.B) {
			frame := makeFrame(OpcodeText, true, size)
			mask := NewMaskingKey()
			buf := &bytes.Buffer{}

			// Write the frame to the buffer once to get the size.
			assert.NilError(b, WriteFrame(buf, mask, frame))
			expectedSize := len(buf.Bytes())
			b.SetBytes(int64(expectedSize))
			b.ResetTimer()

			for b.Loop() {
				buf.Reset()
				assert.NilError(b, WriteFrame(buf, mask, frame))
				assert.Equal(b, buf.Len(), expectedSize, "payload length")
			}
		})
	}
}

func makeFrame(opcode Opcode, fin bool, payloadLen int) *Frame {
	payload := make([]byte, payloadLen)
	for i := range payload {
		payload[i] = 0x20 + byte(i%95) // Map to range 0x20 (space) to 0x7E (~)
	}
	return NewFrame(opcode, fin, payload)
}

func formatSize(b int) string {
	const unit = 1024
	if b < unit {
		return fmt.Sprintf("%dB", b)
	}
	div, exp := int64(unit), 0
	for n := b / unit; n >= unit; n /= unit {
		div *= unit
		exp++
	}
	return fmt.Sprintf("%.0f%ciB",
		float64(b)/float64(div), "KMGTPE"[exp])
}