ws/ws/message_test.go

package ws

import (
	"io"
	"bytes"
	"testing"
)


func TestSimpleMessageReading(t *testing.T) {

	cases := []struct{
		Name       string
		ReadBuffer []byte
		Expected   Message
		Err        error
	}{
		{ // FIN ; TEXT ; Unmasked -> error
			"must fail on unmasked frame",
			[]byte{0x81,0x05,0x68,0x65,0x6c,0x6c,0x6f},
			Message{},
			ErrUnmaskedFrame,
		},
		{ // FIN ; TEXT ; Unmasked -> error
			"must fail because of RSV bit 1 set",
			[]byte{0x81 | 0x40,0x10,0x00,0x00,0x00,0x00},
			Message{},
			ErrReservedBits,
		},
		{ // FIN ; TEXT ; Unmasked -> error
			"must fail because of RSV bit 2 set",
			[]byte{0x81 | 0x20,0x10,0x00,0x00,0x00,0x00},
			Message{},
			ErrReservedBits,
		},
		{ // FIN ; TEXT ; Unmasked -> error
			"must fail because of RSV bit 3 set",
			[]byte{0x81 | 0x10,0x10,0x00,0x00,0x00,0x00},
			Message{},
			ErrReservedBits,
		},
		{ // FIN ; TEXT ; hello
			"simple hello text message",
			[]byte{0x81,0x85,0x00,0x00,0x00,0x00,0x68,0x65,0x6c,0x6c,0x6f},
			Message{ true, TEXT, 5, []byte("hello") },
			nil,
		},
		{ // FIN ; BINARY ; hello
			"simple hello binary message",
			[]byte{0x82,0x85,0x00,0x00,0x00,0x00,0x68,0x65,0x6c,0x6c,0x6f},
			Message{ true, BINARY, 5, []byte("hello") },
			nil,
		},
		{ // FIN ; BINARY ; test unmasking
			"unmasking test",
			[]byte{0x82,0x88,0x01,0x02,0x03,0x04,0x10,0x20,0x30,0x40,0x50,0x60,0x70,0x80},
			Message{ true, BINARY, 8, []byte{0x11,0x22,0x33,0x44,0x51,0x62,0x73,0x84} },
			nil,
		},
		{ // FIN=0 ; TEXT ;
			"non final frame",
			[]byte{0x01,0x82,0x00,0x00,0x00,0x00,0x01,0x02},
			Message{ false, TEXT, 2, []byte{0x01,0x02} },
			nil,
		},
	}

	for _, tc := range cases{

		t.Run(tc.Name, func(t *testing.T){

			reader := bytes.NewBuffer(tc.ReadBuffer)

			got, err := readMessage(reader)

			if err != tc.Err {
				t.Errorf("Expected %v error, got %v", tc.Err, err)
			}

			// do not check message if error expected
			if tc.Err != nil {
				return
			}

			// check FIN
			if got.Final != tc.Expected.Final {
				t.Errorf("Expected FIN=%t, got %t", tc.Expected.Final, got.Final)
			}

			// check OpCode
			if got.Type != tc.Expected.Type {
				t.Errorf("Expected TYPE=%x, got %x", tc.Expected.Type, got.Type)
			}

			// check Size
			if got.Size != tc.Expected.Size {
				t.Errorf("Expected SIZE=%d, got %d", tc.Expected.Size, got.Size)
			}

			// check Data
			if string(got.Data) != string(tc.Expected.Data) {
				t.Errorf("Expected Data='%s', got '%d'", tc.Expected.Data, got.Data)
			}

		})

	}

}


func TestReadEOF(t *testing.T) {

	cases := []struct{
		Name          string
		ReadBuffer    []byte
		eof           bool
		unmaskedError bool
	}{
		{
			"no byte",
			[]byte{},
			true, false,
		}, {
			"only opcode",
			[]byte{0x82},
			true, false,
		}, {
			"only opcode and 0 length",
			[]byte{0x82,0x00},
			false, true,
		}, {
			"missing extended 16 bits length",
			[]byte{0x82,126},
			true, false,
		}, {
			"incomplete extended 16 bits length",
			[]byte{0x82,126, 0x00},
			true, false,
		}, {
			"complete extended 16 bits length",
			[]byte{0x82,126, 0x00, 0x00},
			false, true,
		}, {
			"missing extended 64 bits length",
			[]byte{0x82,127},
			true, false,
		}, {
			"incomplete extended 64 bits length",
			[]byte{0x82,127, 0x00},
			true, false,
		}, {
			"incomplete extended 64 bits length",
			[]byte{0x82,127, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00},
			true, false,
		}, {
			"complete extended 64 bits length",
			[]byte{0x82,127, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00},
			false, true,
		}, {
			"missing mask",
			[]byte{0x82,0x80},
			true, false,
		}, {
			"incomplete mask 1",
			[]byte{0x82,0x80, 0x00},
			true, false,
		}, {
			"incomplete mask 2",
			[]byte{0x82,0x80, 0x00, 0x00, 0x00},
			true, false,
		},{
			"complete mask",
			[]byte{0x82,0x80, 0x00, 0x00, 0x00, 0x00},
			false, false,
		},
	}

	for _, tc := range cases{

		t.Run(tc.Name, func(t *testing.T){

			reader := bytes.NewBuffer(tc.ReadBuffer)

			got, err := readMessage(reader)

			if tc.eof {

				if err != io.EOF{
					t.Errorf("Expected EOF, got %v", err)
				}

				return
			}

			if tc.unmaskedError && err != ErrUnmaskedFrame {
				t.Errorf("Expected UnmaskedFrameor, got %v", err)
			}

			if got.Size != 0x00 {
				t.Errorf("Expected a size of 0, got %d", got.Size)
			}

		})

	}

}



func TestSimpleMessageSending(t *testing.T) {

	m4b1 := make([]byte, 0x7e - 1)
	m4b2 := make([]byte, 0x7e)
	m4b3 := make([]byte, 0x7e + 1)

	m16b1 := make([]byte, 0xffff - 1)
	m16b2 := make([]byte, 0xffff)
	m16b3 := make([]byte, 0xffff + 1)

	cases := []struct{
		Name     string
		Base     Message
		Expected []byte
	}{
		{
			"simple hello text message",
			Message{ true, TEXT, 5, []byte("hello") },
			[]byte{0x81,0x05,0x68,0x65,0x6c,0x6c,0x6f},
		}, {
			"simple hello binary message",
			Message{ true, BINARY, 5, []byte("hello") },
			[]byte{0x82,0x05,0x68,0x65,0x6c,0x6c,0x6f},
		}, {
			"other simple binary message",
			Message{ true, BINARY, 8, []byte{0x10,0x20,0x30,0x40,0x50,0x60,0x70,0x80} },
			[]byte{0x82,0x08,0x10,0x20,0x30,0x40,0x50,0x60,0x70,0x80},
		}, {
			"non final frame",
			Message{ false, TEXT, 2, []byte{0x01,0x02} },
			[]byte{0x01,0x02,0x01,0x02},
		}, {
			"125 > normal length",
			Message{ true, TEXT, uint(len(m4b1)), m4b1 },
			append([]byte{0x81,0x7e-1}, m4b1...),
		}, {
			"126 > extended 16 bits length",
			Message{ true, TEXT, uint(len(m4b2)), m4b2 },
			append([]byte{0x81,126,0x00,0x7e}, m4b2...),
		}, {
			"127 > extended 16 bits length",
			Message{ true, TEXT, uint(len(m4b3)), m4b3 },
			append([]byte{0x81,126,0x00,0x7e+1}, m4b3...),
		}, {
			"fffe > extended 16 bits length",
			Message{ true, TEXT, uint(len(m16b1)), m16b1 },
			append([]byte{0x81,126, 0xff, 0xfe}, m16b1...),
		}, {
			"ffff > extended 16 bits length",
			Message{ true, TEXT, uint(len(m16b2)), m16b2 },
			append([]byte{0x81,126,0xff,0xff}, m16b2...),
		}, {
			"10000 > extended 64 bits length",
			Message{ true, TEXT, uint(len(m16b3)), m16b3 },
			append([]byte{0x81,127, 0x00,0x00,0x00,0x00,0x00,0x01,0x00,0x00,}, m16b3...),
		},
	}

	for _, tc := range cases{

		t.Run(tc.Name, func(t *testing.T){

			writer := new(bytes.Buffer)

			err := tc.Base.Send(writer)

			if err != nil {
				t.Errorf("expected no error, got %v", err)
				return
			}

			// check buffer
			if writer.String() != string(tc.Expected) {
				t.Errorf("expected '%.20x', got '%.20x'", tc.Expected, writer.String())
			}

		})

	}

}




func TestMessageCheck(t *testing.T) {

	type Case struct {
		Name            string
		Msg             Message
		WaitingFragment bool
		Expected        error
	}

	cases := []struct{
		Name            string
		Cases []Case
	}{
		{
			Name: "first fragment type",
			Cases: []Case{
				{
					"CONTINUATION must fail",
					Message{false, CONTINUATION, 0, []byte{}}, false, ErrInvalidFragment,
				}, {
					"TEXT must not fail",
					Message{false, TEXT, 0, []byte{}}, false, nil,
				}, {
					"BINARY must not fail",
					Message{false, BINARY, 0, []byte{}}, false, nil,
				}, {
					"CLOSE must fail",
					Message{false, CLOSE, 0, []byte{}}, false, ErrInvalidFragment,
				}, {
					"PING must fail",
					Message{false, PING, 0, []byte{}}, false, ErrInvalidFragment,
				}, {
					"PONG must fail",
					Message{false, PONG, 0, []byte{}}, false, ErrInvalidFragment,
				},
			},
		}, {
			Name: "frame during fragment",
			Cases: []Case{
				{
					"CONTINUATION must not fail",
					Message{true, CONTINUATION, 0, []byte{}}, true, nil,
				}, {
					"TEXT must fail",
					Message{true, TEXT, 0, []byte{}}, true, ErrInvalidFragment,
				}, {
					"BINARY must fail",
					Message{true, BINARY, 0, []byte{}}, true, ErrInvalidFragment,
				}, {
					"CLOSE must not fail",
					Message{true, CLOSE, 0, []byte{}}, true, CloseFrame,
				}, {
					"PING must not fail",
					Message{true, PING, 0, []byte{}}, true, nil,
				}, {
					"PONG must not fail",
					Message{true, PONG, 0, []byte{}}, true, nil,
				},
			},
		}, {
			Name: "125-length control frame",
			Cases: []Case{
				{
					"CLOSE must not fail",
					Message{true, CLOSE, 125, []byte{0x03,0xe8,0}}, false, CloseFrame,
				}, {
					"PING must not fail",
					Message{true, PING, 125, []byte{}}, false, nil,
				}, {
					"PONG must not fail",
					Message{true, PONG, 125, []byte{}}, false, nil,
				},
			},
		}, {
			Name: "126-length control frame",
			Cases: []Case{
				{
					"CLOSE must fail",
					Message{true, CLOSE, 126, []byte{0x03,0xe8,0}}, false, ErrTooLongControlFrame,
				}, {
					"PING must fail",
					Message{true, PING, 126, []byte{}}, false, ErrTooLongControlFrame,
				}, {
					"PONG must fail",
					Message{true, PONG, 126, []byte{}}, false, ErrTooLongControlFrame,
				},
			},
		}, {
			Name: "fragmented control frame",
			Cases: []Case{
				{
					"CLOSE must fail",
					Message{false, CLOSE, 126, []byte{0x03,0xe8,0}}, false, ErrInvalidFragment,
				}, {
					"PING must fail",
					Message{false, PING, 126, []byte{}}, false, ErrInvalidFragment,
				}, {
					"PONG must fail",
					Message{false, PONG, 126, []byte{}}, false, ErrInvalidFragment,
				},
			},
		}, {
			Name: "unexpected continuation frame",
			Cases: []Case{
				{
					"no waiting fragment final",
					Message{false, CONTINUATION, 126, nil}, false, ErrInvalidFragment,
				}, {
					"no waiting fragment non-final",
					Message{true, CONTINUATION, 126, nil}, false, ErrUnexpectedContinuation,
				},
			},
		}, {
			Name: "utf8 check",
			Cases: []Case{
				{
					"CLOSE valid reason",
					Message{true, CLOSE, 5, []byte{0x03,0xe8, 0xe2,0x82,0xa1}}, false, CloseFrame,
				}, {
					"CLOSE invalid reason byte 2",
					Message{true, CLOSE, 5, []byte{0x03,0xe8, 0xe2,0x28,0xa1}}, false, ErrInvalidPayload,
				}, {
					"CLOSE invalid reason byte 3",
					Message{true, CLOSE, 5, []byte{0x03,0xe8, 0xe2,0x82,0x28}}, false, ErrInvalidPayload,
				},
				{
					"TEXT valid reason",
					Message{true, TEXT, 3, []byte{0xe2,0x82,0xa1}}, false, nil,
				}, {
					"TEXT invalid reason byte 2",
					Message{true, TEXT, 3, []byte{0xe2,0x28,0xa1}}, false, ErrInvalidPayload,
				}, {
					"TEXT invalid reason byte 3",
					Message{true, TEXT, 3, []byte{0xe2,0x82,0x28}}, false, ErrInvalidPayload,
				},
			},
		}, {
			Name: "CLOSE status",
			Cases: []Case{
				{
					"CLOSE only 1 byte",
					Message{true, CLOSE, 1, []byte{0x03}}, false, ErrInvalidCloseStatus,
				}, {
					"valid CLOSE status 1000",
					Message{true, CLOSE, 2, []byte{0x03,0xe8}}, false, CloseFrame,
				}, {
					"invalid CLOSE status 999 under 1000",
					Message{true, CLOSE, 2, []byte{0x03,0xe7}}, false, ErrInvalidCloseStatus,
				}, {
					"valid CLOSE status 1001",
					Message{true, CLOSE, 2, []byte{0x03,0xe9}}, false, CloseFrame,
				}, {
					"valid CLOSE status 1002",
					Message{true, CLOSE, 2, []byte{0x03,0xea}}, false, CloseFrame,
				}, {
					"valid CLOSE status 1003",
					Message{true, CLOSE, 2, []byte{0x03,0xeb}}, false, CloseFrame,
				}, {
					"invalid CLOSE status 1004",
					Message{true, CLOSE, 2, []byte{0x03,0xec}}, false, ErrInvalidCloseStatus,
				}, {
					"invalid CLOSE status 1005",
					Message{true, CLOSE, 2, []byte{0x03,0xed}}, false, ErrInvalidCloseStatus,
				}, {
					"invalid CLOSE status 1006",
					Message{true, CLOSE, 2, []byte{0x03,0xee}}, false, ErrInvalidCloseStatus,
				}, {
					"valid CLOSE status 1007",
					Message{true, CLOSE, 2, []byte{0x03,0xef}}, false, CloseFrame,
				}, {
					"valid CLOSE status 1011",
					Message{true, CLOSE, 2, []byte{0x03,0xf3}}, false, CloseFrame,
				}, {
					"invalid CLOSE status 1012",
					Message{true, CLOSE, 2, []byte{0x03,0xf4}}, false, ErrInvalidCloseStatus,
				}, {
					"invalid CLOSE status 1013",
					Message{true, CLOSE, 2, []byte{0x03,0xf5}}, false, ErrInvalidCloseStatus,
				}, {
					"invalid CLOSE status 1014",
					Message{true, CLOSE, 2, []byte{0x03,0xf6}}, false, ErrInvalidCloseStatus,
				}, {
					"invalid CLOSE status 1015",
					Message{true, CLOSE, 2, []byte{0x03,0xf7}}, false, ErrInvalidCloseStatus,
				}, {
					"invalid CLOSE status 1016",
					Message{true, CLOSE, 2, []byte{0x03,0xf8}}, false, ErrInvalidCloseStatus,
				}, {
					"valid CLOSE status 1017",
					Message{true, CLOSE, 2, []byte{0x03,0xf9}}, false, CloseFrame,
				}, {
					"valid CLOSE status 1099",
					Message{true, CLOSE, 2, []byte{0x04,0x4b}}, false, CloseFrame,
				}, {
					"invalid CLOSE status 1100",
					Message{true, CLOSE, 2, []byte{0x04,0x4c}}, false, ErrInvalidCloseStatus,
				}, {
					"valid CLOSE status 1101",
					Message{true, CLOSE, 2, []byte{0x04,0x4d}}, false, CloseFrame,
				}, {
					"valid CLOSE status 1999",
					Message{true, CLOSE, 2, []byte{0x07,0xcf}}, false, CloseFrame,
				}, {
					"invalid CLOSE status 2000",
					Message{true, CLOSE, 2, []byte{0x07,0xd0}}, false, ErrInvalidCloseStatus,
				}, {
					"valid CLOSE status 2001",
					Message{true, CLOSE, 2, []byte{0x07,0xd1}}, false, CloseFrame,
				}, {
					"valid CLOSE status 2998",
					Message{true, CLOSE, 2, []byte{0x0b,0xb6}}, false, CloseFrame,
				}, {
					"invalid CLOSE status 2999",
					Message{true, CLOSE, 2, []byte{0x0b,0xb7}}, false, ErrInvalidCloseStatus,
				}, {
					"valid CLOSE status 3000",
					Message{true, CLOSE, 2, []byte{0x0b,0xb8}}, false, CloseFrame,
				},
			},
		}, {
			Name: "OpCode check",
			Cases: []Case{
				{ "0", Message{true, 0, 0, []byte{}}, false, ErrUnexpectedContinuation, },
				{ "1", Message{true, 1, 0, []byte{}}, false, nil, },
				{ "2", Message{true, 2, 0, []byte{}}, false, nil, },
				{ "3", Message{true, 3, 0, []byte{}}, false, ErrInvalidOpCode, },
				{ "4", Message{true, 4, 0, []byte{}}, false, ErrInvalidOpCode, },
				{ "5", Message{true, 5, 0, []byte{}}, false, ErrInvalidOpCode, },
				{ "6", Message{true, 6, 0, []byte{}}, false, ErrInvalidOpCode, },
				{ "7", Message{true, 7, 0, []byte{}}, false, ErrInvalidOpCode, },
				{ "8", Message{true, 8, 0, []byte{}}, false, CloseFrame, },
				{ "9", Message{true, 9, 0, []byte{}}, false, nil, },
				{ "10", Message{true, 10, 0, []byte{}}, false, nil, },
				{ "11", Message{true, 11, 0, []byte{}}, false, ErrInvalidOpCode, },
				{ "12", Message{true, 12, 0, []byte{}}, false, ErrInvalidOpCode, },
				{ "13", Message{true, 13, 0, []byte{}}, false, ErrInvalidOpCode, },
				{ "14", Message{true, 14, 0, []byte{}}, false, ErrInvalidOpCode, },
				{ "15", Message{true, 15, 0, []byte{}}, false, ErrInvalidOpCode, },
			},
		},
	}


	for _, tcc := range cases{

		t.Run(tcc.Name, func(t *testing.T){

			for _, tc := range tcc.Cases{

				t.Run(tc.Name, func(t *testing.T){

					actual := tc.Msg.check(tc.WaitingFragment)

					if actual != tc.Expected {
						t.Errorf("expected '%v', got '%v'", tc.Expected, actual)
					}

				})

			}

		})

	}

}