Skip to content

Commit 17810cd

Browse files
LawouachLawouach
committed
masking is required, fixed a couple of minor corner cases
1 parent 362278a commit 17810cd

1 file changed

Lines changed: 23 additions & 22 deletions

File tree

‎test/test_stream.py‎

Lines changed: 23 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -10,25 +10,24 @@
1010

1111
class WSStreamTest(unittest.TestCase):
1212
def test_empty_close_message(self):
13-
f = Frame(opcode=OPCODE_CLOSE, body='', fin=1).build()
13+
f = Frame(opcode=OPCODE_CLOSE, body='', fin=1, masking_key=os.urandom(4)).build()
1414
s = Stream()
1515
self.assertEqual(s.closing, None)
1616
s.parser.send(f)
1717
self.assertEqual(type(s.closing), CloseControlMessage)
1818

1919
def test_too_large_close_message(self):
20-
f = Frame(opcode=OPCODE_CLOSE, body='*' * 330, fin=1).build()
20+
f = Frame(opcode=OPCODE_CLOSE, body='*' * 330, fin=1, masking_key=os.urandom(4)).build()
2121
s = Stream()
2222
self.assertEqual(s.closing, None)
2323
s.parser.send(f)
24-
self.assertEqual(s.closing, None)
25-
self.assertEqual(len(s.errors), 1)
26-
self.assertEqual(type(s.errors[0]), CloseControlMessage)
27-
self.assertEqual(s.errors[0].code, 1002)
24+
25+
self.assertEqual(type(s.closing), CloseControlMessage)
26+
self.assertEqual(s.closing.code, 1002)
2827

2928
def test_ping_message_received(self):
3029
msg = 'ping me'
31-
f = Frame(opcode=OPCODE_PING, body=msg, fin=1).build()
30+
f = Frame(opcode=OPCODE_PING, body=msg, fin=1, masking_key=os.urandom(4)).build()
3231
s = Stream()
3332
self.assertEqual(len(s.pings), 0)
3433
s.parser.send(f)
@@ -37,15 +36,15 @@ def test_ping_message_received(self):
3736

3837
def test_text_message_received(self):
3938
msg = 'hello there'
40-
f = Frame(opcode=OPCODE_TEXT, body=msg, fin=1).build()
39+
f = Frame(opcode=OPCODE_TEXT, body=msg, fin=1, masking_key=os.urandom(4)).build()
4140
s = Stream()
4241
self.assertEqual(len(s.messages), 0)
4342
s.parser.send(f)
4443
self.assertEqual(len(s.messages), 1)
4544

4645
def test_incremental_text_message_received(self):
4746
msg = 'hello there'
48-
f = Frame(opcode=OPCODE_TEXT, body=msg, fin=1).build()
47+
f = Frame(opcode=OPCODE_TEXT, body=msg, fin=1, masking_key=os.urandom(4)).build()
4948
s = Stream()
5049
self.assertEqual(s.has_message, False)
5150
for byte in f:
@@ -54,83 +53,85 @@ def test_incremental_text_message_received(self):
5453

5554
def test_text_message_received(self):
5655
msg = 'hello there'
57-
f = Frame(opcode=OPCODE_TEXT, body=msg, fin=1).build()
56+
f = Frame(opcode=OPCODE_TEXT, body=msg, fin=1, masking_key=os.urandom(4)).build()
5857
s = Stream()
5958
self.assertEqual(s.has_message, False)
6059
s.parser.send(f)
6160
self.assertEqual(s.message.completed, True)
6261

6362
def test_text_message_with_continuation_received(self):
6463
msg = 'hello there'
65-
f = Frame(opcode=OPCODE_TEXT, body=msg, fin=0).build()
64+
f = Frame(opcode=OPCODE_TEXT, body=msg, fin=0, masking_key=os.urandom(4)).build()
6665
s = Stream()
6766
self.assertEqual(s.has_message, False)
6867
s.parser.send(f)
6968
self.assertEqual(s.message.completed, False)
7069

7170
for i in range(3):
72-
f = Frame(opcode=OPCODE_CONTINUATION, body=msg, fin=0).build()
71+
f = Frame(opcode=OPCODE_CONTINUATION, body=msg, fin=0, masking_key=os.urandom(4)).build()
7372
s.parser.send(f)
7473
self.assertEqual(s.has_message, False)
7574
self.assertEqual(s.message.completed, False)
7675
self.assertEqual(s.message.opcode, OPCODE_TEXT)
7776

78-
f = Frame(opcode=OPCODE_CONTINUATION, body=msg, fin=1).build()
77+
f = Frame(opcode=OPCODE_CONTINUATION, body=msg, fin=1, masking_key=os.urandom(4)).build()
7978
s.parser.send(f)
8079
self.assertEqual(s.has_message, True)
8180
self.assertEqual(s.message.completed, True)
8281
self.assertEqual(s.message.opcode, OPCODE_TEXT)
8382

8483
def test_text_message_with_continuation_and_ping_in_between(self):
8584
msg = 'hello there'
86-
f = Frame(opcode=OPCODE_TEXT, body=msg, fin=0).build()
85+
key = os.urandom(4)
86+
f = Frame(opcode=OPCODE_TEXT, body=msg, fin=0, masking_key=os.urandom(4)).build()
8787
s = Stream()
8888
self.assertEqual(s.has_message, False)
8989
s.parser.send(f)
9090
self.assertEqual(s.message.completed, False)
9191

9292
for i in range(3):
93-
f = Frame(opcode=OPCODE_CONTINUATION, body=msg, fin=0).build()
93+
f = Frame(opcode=OPCODE_CONTINUATION, body=msg, fin=0, masking_key=os.urandom(4)).build()
9494
s.parser.send(f)
9595
self.assertEqual(s.has_message, False)
9696
self.assertEqual(s.message.completed, False)
9797
self.assertEqual(s.message.opcode, OPCODE_TEXT)
9898

99-
f = Frame(opcode=OPCODE_PING, body='ping me', fin=1).build()
99+
f = Frame(opcode=OPCODE_PING, body='ping me', fin=1, masking_key=os.urandom(4)).build()
100100
self.assertEqual(len(s.pings), i)
101101
s.parser.send(f)
102102
self.assertEqual(len(s.pings), i+1)
103103

104-
f = Frame(opcode=OPCODE_CONTINUATION, body=msg, fin=1).build()
104+
f = Frame(opcode=OPCODE_CONTINUATION, body=msg, fin=1, masking_key=os.urandom(4)).build()
105105
s.parser.send(f)
106106
self.assertEqual(s.has_message, True)
107107
self.assertEqual(s.message.opcode, OPCODE_TEXT)
108108
self.assertEqual(s.message.completed, True)
109109

110110
def test_binary_message_received(self):
111111
msg = os.urandom(16)
112-
f = Frame(opcode=OPCODE_BINARY, body=msg, fin=1).build()
112+
f = Frame(opcode=OPCODE_BINARY, body=msg, fin=1, masking_key=os.urandom(4)).build()
113113
s = Stream()
114114
self.assertEqual(s.has_message, False)
115115
s.parser.send(f)
116116
self.assertEqual(s.message.completed, True)
117117

118118
def test_binary_message_with_continuation_received(self):
119119
msg = os.urandom(16)
120-
f = Frame(opcode=OPCODE_BINARY, body=msg, fin=0).build()
120+
key = os.urandom(4)
121+
f = Frame(opcode=OPCODE_BINARY, body=msg, fin=0, masking_key=key).build()
121122
s = Stream()
122123
self.assertEqual(s.has_message, False)
123124
s.parser.send(f)
124-
self.assertEqual(s.message.completed, False)
125+
self.assertEqual(s.has_message, False)
125126

126127
for i in range(3):
127-
f = Frame(opcode=OPCODE_CONTINUATION, body=msg, fin=0).build()
128+
f = Frame(opcode=OPCODE_CONTINUATION, body=msg, fin=0, masking_key=key).build()
128129
s.parser.send(f)
129130
self.assertEqual(s.has_message, False)
130131
self.assertEqual(s.message.completed, False)
131132
self.assertEqual(s.message.opcode, OPCODE_BINARY)
132133

133-
f = Frame(opcode=OPCODE_CONTINUATION, body=msg, fin=1).build()
134+
f = Frame(opcode=OPCODE_CONTINUATION, body=msg, fin=1, masking_key=key).build()
134135
s.parser.send(f)
135136
self.assertEqual(s.has_message, True)
136137
self.assertEqual(s.message.completed, True)

0 commit comments

Comments
 (0)