# vim: tabstop=4 shiftwidth=4 softtabstop=4 # Copyright(c)2013 NTT corp. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); you may # not use this file except in compliance with the License. You may obtain # a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, WITHOUT # WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the # License for the specific language governing permissions and limitations # under the License. """ Unit tests for websocket """ import unittest from websockify import websocket class FakeSocket: def __init__(self): self.data = b'' def send(self, buf): self.data += buf return len(buf) class AcceptTestCase(unittest.TestCase): def test_success(self): ws = websocket.WebSocket() sock = FakeSocket() ws.accept(sock, {'upgrade': 'websocket', 'Sec-WebSocket-Version': '13', 'Sec-WebSocket-Key': 'DKURYVK9cRFul1vOZVA56Q=='}) self.assertEqual(sock.data[:13], b'HTTP/1.1 101 ') self.assertTrue(b'\r\nUpgrade: websocket\r\n' in sock.data) self.assertTrue(b'\r\nConnection: Upgrade\r\n' in sock.data) self.assertTrue(b'\r\nSec-WebSocket-Accept: pczpYSQsvE1vBpTQYjFQPcuoj6M=\r\n' in sock.data) def test_bad_version(self): ws = websocket.WebSocket() sock = FakeSocket() self.assertRaises(Exception, ws.accept, sock, {'upgrade': 'websocket', 'Sec-WebSocket-Key': 'DKURYVK9cRFul1vOZVA56Q=='}) self.assertRaises(Exception, ws.accept, sock, {'upgrade': 'websocket', 'Sec-WebSocket-Version': '5', 'Sec-WebSocket-Key': 'DKURYVK9cRFul1vOZVA56Q=='}) self.assertRaises(Exception, ws.accept, sock, {'upgrade': 'websocket', 'Sec-WebSocket-Version': '20', 'Sec-WebSocket-Key': 'DKURYVK9cRFul1vOZVA56Q=='}) def test_bad_upgrade(self): ws = websocket.WebSocket() sock = FakeSocket() self.assertRaises(Exception, ws.accept, sock, {'Sec-WebSocket-Version': '13', 'Sec-WebSocket-Key': 'DKURYVK9cRFul1vOZVA56Q=='}) self.assertRaises(Exception, ws.accept, sock, {'upgrade': 'websocket2', 'Sec-WebSocket-Version': '13', 'Sec-WebSocket-Key': 'DKURYVK9cRFul1vOZVA56Q=='}) def test_missing_key(self): ws = websocket.WebSocket() sock = FakeSocket() self.assertRaises(Exception, ws.accept, sock, {'upgrade': 'websocket', 'Sec-WebSocket-Version': '13'}) def test_protocol(self): class ProtoSocket(websocket.WebSocket): def select_subprotocol(self, protocol): return 'gazonk' ws = ProtoSocket() sock = FakeSocket() ws.accept(sock, {'upgrade': 'websocket', 'Sec-WebSocket-Version': '13', 'Sec-WebSocket-Key': 'DKURYVK9cRFul1vOZVA56Q==', 'Sec-WebSocket-Protocol': 'foobar,gazonk'}) self.assertEqual(sock.data[:13], b'HTTP/1.1 101 ') self.assertTrue(b'\r\nSec-WebSocket-Protocol: gazonk\r\n' in sock.data) def test_no_protocol(self): ws = websocket.WebSocket() sock = FakeSocket() ws.accept(sock, {'upgrade': 'websocket', 'Sec-WebSocket-Version': '13', 'Sec-WebSocket-Key': 'DKURYVK9cRFul1vOZVA56Q=='}) self.assertEqual(sock.data[:13], b'HTTP/1.1 101 ') self.assertFalse(b'\r\nSec-WebSocket-Protocol:' in sock.data) def test_missing_protocol(self): ws = websocket.WebSocket() sock = FakeSocket() self.assertRaises(Exception, ws.accept, sock, {'upgrade': 'websocket', 'Sec-WebSocket-Version': '13', 'Sec-WebSocket-Key': 'DKURYVK9cRFul1vOZVA56Q==', 'Sec-WebSocket-Protocol': 'foobar,gazonk'}) def test_unsupported_protocol(self): class ProtoSocket(websocket.WebSocket): def select_subprotocol(self, protocol): return 'oddball' ws = ProtoSocket() sock = FakeSocket() self.assertRaises(Exception, ws.accept, sock, {'upgrade': 'websocket', 'Sec-WebSocket-Version': '13', 'Sec-WebSocket-Key': 'DKURYVK9cRFul1vOZVA56Q==', 'Sec-WebSocket-Protocol': 'foobar,gazonk'}) class PingPongTest(unittest.TestCase): def setUp(self): self.ws = websocket.WebSocket() self.sock = FakeSocket() self.ws.accept(self.sock, {'upgrade': 'websocket', 'Sec-WebSocket-Version': '13', 'Sec-WebSocket-Key': 'DKURYVK9cRFul1vOZVA56Q=='}) self.assertEqual(self.sock.data[:13], b'HTTP/1.1 101 ') self.sock.data = b'' def test_ping(self): self.ws.ping() self.assertEqual(self.sock.data, b'\x89\x00') def test_pong(self): self.ws.pong() self.assertEqual(self.sock.data, b'\x8a\x00') def test_ping_data(self): self.ws.ping(b'foo') self.assertEqual(self.sock.data, b'\x89\x03foo') def test_pong_data(self): self.ws.pong(b'foo') self.assertEqual(self.sock.data, b'\x8a\x03foo') class HyBiEncodeDecodeTestCase(unittest.TestCase): def test_decode_hybi_text(self): buf = b'\x81\x85\x37\xfa\x21\x3d\x7f\x9f\x4d\x51\x58' ws = websocket.WebSocket() res = ws._decode_hybi(buf) self.assertEqual(res['fin'], 1) self.assertEqual(res['opcode'], 0x1) self.assertEqual(res['masked'], True) self.assertEqual(res['length'], len(buf)) self.assertEqual(res['payload'], b'Hello') def test_decode_hybi_binary(self): buf = b'\x82\x04\x01\x02\x03\x04' ws = websocket.WebSocket() res = ws._decode_hybi(buf) self.assertEqual(res['fin'], 1) self.assertEqual(res['opcode'], 0x2) self.assertEqual(res['length'], len(buf)) self.assertEqual(res['payload'], b'\x01\x02\x03\x04') def test_decode_hybi_extended_16bit_binary(self): data = (b'\x01\x02\x03\x04' * 65) # len > 126 -- len == 260 buf = b'\x82\x7e\x01\x04' + data ws = websocket.WebSocket() res = ws._decode_hybi(buf) self.assertEqual(res['fin'], 1) self.assertEqual(res['opcode'], 0x2) self.assertEqual(res['length'], len(buf)) self.assertEqual(res['payload'], data) def test_decode_hybi_extended_64bit_binary(self): data = (b'\x01\x02\x03\x04' * 65) # len > 126 -- len == 260 buf = b'\x82\x7f\x00\x00\x00\x00\x00\x00\x01\x04' + data ws = websocket.WebSocket() res = ws._decode_hybi(buf) self.assertEqual(res['fin'], 1) self.assertEqual(res['opcode'], 0x2) self.assertEqual(res['length'], len(buf)) self.assertEqual(res['payload'], data) def test_decode_hybi_multi(self): buf1 = b'\x01\x03\x48\x65\x6c' buf2 = b'\x80\x02\x6c\x6f' ws = websocket.WebSocket() res1 = ws._decode_hybi(buf1) self.assertEqual(res1['fin'], 0) self.assertEqual(res1['opcode'], 0x1) self.assertEqual(res1['length'], len(buf1)) self.assertEqual(res1['payload'], b'Hel') res2 = ws._decode_hybi(buf2) self.assertEqual(res2['fin'], 1) self.assertEqual(res2['opcode'], 0x0) self.assertEqual(res2['length'], len(buf2)) self.assertEqual(res2['payload'], b'lo') def test_encode_hybi_basic(self): ws = websocket.WebSocket() res = ws._encode_hybi(0x1, b'Hello') expected = b'\x81\x05\x48\x65\x6c\x6c\x6f' self.assertEqual(res, expected)