diff --git a/ws/message.go b/ws/message.go index 6381f34..a0fe513 100644 --- a/ws/message.go +++ b/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 } diff --git a/ws/message_test.go b/ws/message_test.go index bf1c2c6..f248631 100644 --- a/ws/message_test.go +++ b/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) + } + + }) + + } + + }) + + } + } \ No newline at end of file