add possible PONG receive during fragmented messages + add tests for 'Message.check(bool) err'

ws/message.go

@@ -251,7 +251,7 @@ func (m *Message) check(fragment bool) error{
 	}
 
 	/* (2) Waiting fragment but received standalone frame */
-	if fragment && m.Type != CONTINUATION && m.Type != CLOSE && m.Type != PING {
+	if fragment && m.Type != CONTINUATION && m.Type != CLOSE && m.Type != PING && m.Type != PONG {
 		return ErrInvalidFragment
 	}
 

ws/message_test.go

@@ -21,6 +21,24 @@ func TestSimpleMessageReading(t *testing.T) {
 			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},
@@ -271,4 +289,274 @@ func TestSimpleMessageSending(t *testing.T) {
 
 	}
 
+}
+
+
+
+
+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)
+					}
+
+				})
+
+			}
+
+		})
+
+	}
+
 }