def setUp(self):
self.channel = DummyChannel()
request = Request(self.channel, False)
transport = WebSocketTransport(request)
handler = TestHandler(transport)
transport._attachHandler(handler)
self.decoder = WebSocketFrameDecoder(request, handler)
# taken straight from the IETF spec, masking added where appropriate
self.hello = "\x81\x85\x37\xfa\x21\x3d\x7f\x9f\x4d\x51\x58"
self.frag_hello = ("\x01\x83\x12\x21\x65\x23\x5a\x44\x09",
"\x80\x82\x63\x34\xf1\x00\x0f\x5b")
self.binary_orig = "\x3f" * 256
self.binary = ("\x82\xfe\x01\x00\x12\x6d\xa6\x23" +
"\x2d\x52\x99\x1c" * 64)
self.long_binary = ("\x82\xff\x00\x00\x00\x00\x00\x01\x00\x00" +
"\x12\x6d\xa6\x23" +
"\x2d\x52\x99\x1c" * 16384)
self.long_binary_orig = "\x3f" * 65536
self.ping = "\x89\x85\x56\x23\x88\x23\x1e\x46\xe4\x4f\x39"
self.pong = "\x8a\x85\xde\x41\x0f\x34\x96\x24\x63\x58\xb1"
self.pong_unmasked = "\x8a\x05\x48\x65\x6c\x6c\x6f"
# code 1000, message "Normal Closure"
self.close = ("\x88\x90\x34\x23\x87\xde\x37\xcb\xc9\xb1\x46"
"\x4e\xe6\xb2\x14\x60\xeb\xb1\x47\x56\xf5\xbb")
## close message can be empty or with normal close code (1000)
self.empty_unmasked_close_list = ("\x88\x00", "\x88\x02\x03\xe8")
self.empty_text = "\x81\x80\x00\x01\x02\x03"
self.cont_empty_text = "\x00\x80\x00\x01\x02\x03"
class WebSocketFrameDecoderXor(TestCase):
"""
Test the vitamined xor function of the decoder (_xxor)
"""
def setUp(self):
self.channel = DummyChannel()
request = Request(self.channel, False)
transport = WebSocketTransport(request)
handler = WebSocketHandler(transport)
self.decoder = WebSocketFrameDecoder(request, handler)
self.decoder._mask = "Key!"
self.decoder._rkeys = {}
self.decoder._maskIndex = 0
def xor(self, data):
"""
Reference version
"""
return ''.join(chr(ord(x) ^ ord(y)) for (x,y) in izip(data, cycle(self.decoder._mask)))
def test_xxorShort(self):
"""
Short data (less than 16 bytes)
"""
data = "Hello"
self.assertEquals(self.decoder._xxor(data), self.xor(data))
def test_xxorLong(self):
"""
Long data (more than 16 bytes). Make sure the data is not a multiple of 8
Test the long xor capability
"""
data = "Hello" * 500 + "trail"
self.assertEquals(self.decoder._xxor(data), self.xor(data))
def test_xxorSplit(self):
"""
Long data (4005 bytes), processed in chunks of 37 bytes
Test the rolling key mechanism
"""
data = "Hello" * 500 + "trail"
chunk_size = 37
self.assertEquals("".join(self.decoder._xxor(data[i:i+chunk_size])
for i in range(0, len(data), chunk_size)),
self.xor(data))
def setUp(self):
self.channel = DummyChannel()
request = Request(self.channel, False)
transport = WebSocketTransport(request)
handler = WebSocketHandler(transport)
self.decoder = WebSocketFrameDecoder(request, handler)
self.decoder._mask = "Key!"
self.decoder._rkeys = {}
self.decoder._maskIndex = 0
def connectionMade(self):
self.data = ""
self.extensions = []
self.isSecure = (self.factory.scheme == "wss")
self.transport.connectionLost = self.connectionLost
self.sendHeader()
self.decoder = WebSocketFrameDecoder(self.transport,
self,
isClient=True,
assembleFragments=self.assembleFragments)
self._firstFragment = True
class WebSocketServerFrameDecoderTestCase(TestCase):
"""
Test for C{WebSocketFrameDecoder}.
"""
def setUp(self):
self.channel = DummyChannel()
request = Request(self.channel, False)
transport = WebSocketTransport(request)
handler = TestHandler(transport)
transport._attachHandler(handler)
self.decoder = WebSocketFrameDecoder(request, handler)
# taken straight from the IETF spec, masking added where appropriate
self.hello = "\x81\x85\x37\xfa\x21\x3d\x7f\x9f\x4d\x51\x58"
self.frag_hello = ("\x01\x83\x12\x21\x65\x23\x5a\x44\x09",
"\x80\x82\x63\x34\xf1\x00\x0f\x5b")
self.binary_orig = "\x3f" * 256
self.binary = ("\x82\xfe\x01\x00\x12\x6d\xa6\x23" +
"\x2d\x52\x99\x1c" * 64)
self.long_binary = ("\x82\xff\x00\x00\x00\x00\x00\x01\x00\x00" +
"\x12\x6d\xa6\x23" +
"\x2d\x52\x99\x1c" * 16384)
self.long_binary_orig = "\x3f" * 65536
self.ping = "\x89\x85\x56\x23\x88\x23\x1e\x46\xe4\x4f\x39"
self.pong = "\x8a\x85\xde\x41\x0f\x34\x96\x24\x63\x58\xb1"
self.pong_unmasked = "\x8a\x05\x48\x65\x6c\x6c\x6f"
# code 1000, message "Normal Closure"
self.close = ("\x88\x90\x34\x23\x87\xde\x37\xcb\xc9\xb1\x46"
"\x4e\xe6\xb2\x14\x60\xeb\xb1\x47\x56\xf5\xbb")
## close message can be empty or with normal close code (1000)
self.empty_unmasked_close_list = ("\x88\x00", "\x88\x02\x03\xe8")
self.empty_text = "\x81\x80\x00\x01\x02\x03"
self.cont_empty_text = "\x00\x80\x00\x01\x02\x03"
def assertOneDecodingError(self):
"""
Assert that exactly one L{DecodingError} has been logged and return
that error.
"""
errors = self.flushLoggedErrors(DecodingError)
self.assertEquals(len(errors), 1)
return errors[0]
def test_oneTextFrame(self):
"""
We can send one frame handled with one C{dataReceived} call.
"""
self.decoder.dataReceived(self.hello)
self.assertEquals(self.decoder.handler.frames, ["Hello"])
def test_chunkedTextFrame(self):
"""
We can send one text frame handled with multiple C{dataReceived} calls.
"""
# taken straight from the IETF spec
for part in (self.hello[:1], self.hello[1:3],
self.hello[3:7], self.hello[7:]):
self.decoder.dataReceived(part)
self.assertEquals(self.decoder.handler.frames, ["Hello"])
def test_fragmentedTextFrame(self):
"""
We can send a fragmented frame handled with one C{dataReceived} call.
"""
self.decoder.dataReceived("".join(self.frag_hello))
self.assertEquals(self.decoder.handler.frames, ["Hello"])
def test_chunkedfragmentedTextFrame(self):
"""
We can send a fragmented text frame handled with multiple
C{dataReceived} calls.
"""
# taken straight from the IETF spec
for part in (self.frag_hello[0][:3], self.frag_hello[0][3:]):
self.decoder.dataReceived(part)
for part in (self.frag_hello[1][:1], self.frag_hello[1][1:]):
self.decoder.dataReceived(part)
self.assertEquals(self.decoder.handler.frames, ["Hello"])
def test_twoFrames(self):
"""
We can send two frames together and they will be correctly parsed.
"""
self.decoder.dataReceived("".join(self.frag_hello) + self.hello)
self.assertEquals(self.decoder.handler.frames, ["Hello"] * 2)
def test_controlInterleaved(self):
"""
A control message (in this case a pong) can appear between the
fragmented frames.
"""
data = self.frag_hello[0] + self.pong + self.frag_hello[1]
for part in data[:2], data[2:7], data[7:8], data[8:14], data[14:]:
self.decoder.dataReceived(part)
self.assertEquals(self.decoder.handler.frames, ["Hello"])
self.assertEquals(self.decoder.handler.pongs, ["Hello"])
def test_binaryFrameAllAtOnce(self):
"""
We can send a binary frame that uses a longer length field.
"""
self.decoder.dataReceived(self.binary)
self.assertEquals(self.decoder.handler.binaryFrames,
[self.binary_orig])
def test_binaryFrame(self):
"""
We can send a binary frame that uses a longer length field.
Split data in some interesting locations.
"""
data = self.binary + self.hello
for part in data[:3], data[3:4], data[4:]:
self.decoder.dataReceived(part)
self.assertEquals(self.decoder.handler.binaryFrames,
[self.binary_orig])
self.assertEquals(self.decoder.handler.frames, ["Hello"])
def test_binaryFrameMultiSplit(self):
"""
We can send a binary frame that uses a longer length field,
plus one other frame.
Split in all header locations.
"""
data = self.binary + self.hello
for i in range(10):
self.decoder.handler.binaryFrames = []
self.decoder.handler.frames = []
for part in data[:i], data[i:]:
self.decoder.dataReceived(part)
self.assertEquals(self.decoder.handler.binaryFrames,
[self.binary_orig],
"Failed at i=%d" % i)
self.assertEquals(self.decoder.handler.frames, ["Hello"])
for i in range(10):
self.decoder.handler.binaryFrames = []
self.decoder.handler.frames = []
for part in data[:len(self.binary)+i], data[len(self.binary)+i:]:
self.decoder.dataReceived(part)
self.assertEquals(self.decoder.handler.binaryFrames,
[self.binary_orig],
"Failed (last) at i=%d" % i)
self.assertEquals(self.decoder.handler.frames, ["Hello"])
def test_binaryFrameOneByOne(self):
"""
We can send a binary frame that uses a longer length field.
Bytes are sent one by one.
"""
data = self.binary + self.hello
for part in data:
self.decoder.dataReceived(part)
self.assertEquals(self.decoder.handler.binaryFrames,
[self.binary_orig])
self.assertEquals(self.decoder.handler.frames, ["Hello"])
def test_longBinaryFrame(self):
"""
We can send a binary frame that uses a very long length field.
"""
data = self.long_binary + self.hello
for part in data[:3], data[3:4], data[4:]:
self.decoder.dataReceived(part)
self.assertEquals(self.decoder.handler.binaryFrames,
[self.long_binary_orig])
self.assertEquals(self.decoder.handler.frames, ["Hello"])
def test_pingInterleaved(self):
"""
We can get a ping frame in the middle of a fragmented frame and we'll
correctly send a pong resonse.
"""
data = self.frag_hello[0] + self.ping + self.frag_hello[1]
for part in data[:12], data[12:16], data[16:]:
self.decoder.dataReceived(part)
self.assertEquals(self.decoder.handler.frames, ["Hello"])
result = self.channel.transport.written.getvalue()
headers, response = result.split('\r\n\r\n')
self.assertEquals(response, self.pong_unmasked)
def test_pingInterleavedOneByOne(self):
"""
We can get a ping frame in the middle of a fragmented frame and we'll
correctly send a pong resonse.
Send bytes one by one.
"""
data = self.frag_hello[0] + self.ping + self.frag_hello[1]
for part in data:
self.decoder.dataReceived(part)
self.assertEquals(self.decoder.handler.frames, ["Hello"])
result = self.channel.transport.written.getvalue()
headers, response = result.split('\r\n\r\n')
self.assertEquals(response, self.pong_unmasked)
def test_close(self):
"""
A close frame causes the remaining data to be discarded and the
connection to be closed.
"""
self.decoder.dataReceived(self.hello + self.close + "crap" * 20)
self.assertEquals(self.decoder.handler.frames, ["Hello"])
self.assertEquals(self.decoder.handler.closes,
[(1000, "Normal Closure")])
result = self.channel.transport.written.getvalue()
headers, response = result.split('\r\n\r\n')
self.assertIn(response, self.empty_unmasked_close_list)
self.assertTrue(self.channel.transport.disconnected)
def test_emptyFrame(self):
"""
An empty text frame is correctly parsed.
"""
self.decoder.dataReceived(self.empty_text)
self.assertEquals(self.decoder.handler.frames, [""])
def test_emptyFrameInterleaved(self):
"""
An empty fragmented frame and a interleaved pong message are received
and parsed.
"""
data = (self.frag_hello[0] + self.cont_empty_text +
self.pong + self.frag_hello[1])
for part in data[:1], data[1:8], data[8:17], data[17:]:
self.decoder.dataReceived(part)
self.assertEquals(self.decoder.handler.frames, ["Hello"])
self.assertEquals(self.decoder.handler.pongs, ["Hello"])
class WebSocketClient(Protocol, TimeoutMixin, _PauseableMixin):
"""
Protocol for websocket clients.
@ivar connectionEstablished: a callable which will be invoqued when the
connection has been established.
@ivar handshakeError: A one-argument callable which will be invoked when
the handshake failed, the argument being a string describing
the reason for the failure.
@ivar frameReceived: A one-argument callable which will be invoked when
the terminal chunk is received. It will be invoked with all bytes
which were delivered to this protocol which came after the terminal
chunk.
@ivar fragmentReceived: A one-argument callable which will be invoked each
time a fragment of data is received for streaming purpose.
The argument is the plain data as received on the wire.
@ivar dataError: A one-argument callable which will be invoked when
the handshake failed, the argument being a string describing
the reason for the failure.
"""
extensions = []
assembleFragments = False
def connectionMade(self):
self.data = ""
self.extensions = []
self.isSecure = (self.factory.scheme == "wss")
self.transport.connectionLost = self.connectionLost
self.sendHeader()
self.decoder = WebSocketFrameDecoder(self.transport,
self,
isClient=True,
assembleFragments=self.assembleFragments)
self._firstFragment = True
def sendHeader(self):
self.in_handshake = True
self.fields = [
"GET %s HTTP/1.1\r\n" % self.factory.path,
"Host: %s:%s\r\n" % (self.factory.host, self.factory.port),
## XXX: Can be added later: Proxy-authorization
]
if self.isSecure:
if not has_ssl:
self.handshakeError("The PyOpenSSL library is not installed."
" Cannot use the wss scheme")
self.transport.loseConnection()
return
self.transport.startTLS(ClientTLSContext(), self.factory)
else:
self.sendHeaderEnd()
def sendHeaderEnd(self):
fields = self.fields
self.nonce = b64encode(generate_key())
## These fields must be shuffled
field_list = [
"Upgrade: WebSocket\r\n",
"Connection: Upgrade\r\n",
"Sec-WebSocket-Key: %s\r\n" % self.nonce,
"Origin: http://%s\r\n" % self.factory.origin,
"Sec-WebSocket-Version: 13\r\n"
]
if self.factory.subprotocolsAvailable:
field_list.append("Sec-WebSocket-Protocol: %s\r\n"
% ",".join(self.factory.subprotocolsAvailable))
if self.factory.extensionsAvailable:
field_list.append("Sec-WebSocket-Extensions: %s\r\n"
% ",".join(ext.requestHeader
for ext in self.factory.extensionsAvailable))
if self.factory.extra_headers != None:
field_list.extend(self.factory.extra_headers)
field_list.append("\r\n")
all_fields = ''.join(fields + field_list)
self.transport.write(all_fields)
def writeBinary(self, data, fragmented=False):
"""
Treat the given frame as a binary frame and send it to the client.
@param frame: a binary C{str} to send to the client.
@type frame: C{str}
"""
self.sendFrame(OPCODE_BINARY, data, fragmented)
def write(self, data, fragmented=False):
"""
Treat the given utf-8 data as a text frame and send it to the client.
@param frame: a I{UTF-8} encoded C{str} to send to the client.
@type frame: C{str}
"""
self.sendFrame(OPCODE_TEXT, data.encode('utf-8'), fragmented=fragmented)
def writeFragment(self, data, final=False):
"""
Send a fragment of a text frame to the client.
@param fragment: a I{UTF-8} encoded C{str} to send to the client.
@type fragment: C{str}
"""
self.sendFrame(OPCODE_TEXT, data.encode('utf-8'), fragmented=True, final=final)
write.writeFragment = writeFragment
def writeBinaryFragment(self, data, final=False):
"""
Send a fragment of a binary frame to the client.
@param fragment: a C{str} to send to the client.
@type fragment: C{str}
"""
self.sendFrame(OPCODE_BINARY, data, fragmented=True, final=final)
writeBinary.writeFragment = writeBinaryFragment
def sendPing(self, data=""):
"""
Send a simple ping packet.
The factory should implement pongReceived to check the pong response
"""
self.sendFrame(OPCODE_PING, data)
def writeSequence(self, frames):
"""
Send a sequence of text frames to the connected client.
"""
for frame in frames:
self.sendFrame(OPCODE_TEXT, frame)
def writeBinarySequence(self, frames):
"""
Send a sequence of binary frames to the connected client.
"""
for frame in frames:
self.sendFrame(OPCODE_BINARY, frame)
def sendFrame(self, opcode, payload, fragmented=False, final=False, **kwargs):
"""
Send a frame with the given opcode and payload to the client. If the
L{fragmented} parameter is set, the message frame will contain a flag
saying it's part of a fragmented payload, by default data is sent as a
self-contained frame. Note that if you use fragmentation support, it is
up to you to correctly set the first frame's opcode and then use
L{OPCODE_CONT} on the following continuation frames.
Payloads sent using this method are never masked.
@param opcode: the opcode as defined in rfc6455
@type opcode: C{int}
@param payload: the frame's payload
@type payload: C{str}
@param fragmented: should the frame be marked as part of a fragmented payload
@type fragmented: C{bool}
"""
if opcode not in ALL_OPCODES:
raise ValueError, "Invalid opcode 0x%X" % opcode
## Process extensions
if self.extensions:
if opcode == OPCODE_TEXT:
for ext in self.extensions:
payload = ext.processOutgoingFrame(payload, **kwargs)
elif opcode == OPCODE_BINARY:
for ext in self.extensions:
payload = ext.processOutgoingBinaryFrame(payload, **kwargs)
length = len(payload)
if fragmented:
if final:
if self._firstFragment:
## Special case when final fragment is also
## the first one: same as unfragmented
data = [0x80 | opcode]
else:
## Final fragment, opcode set to 0
data = [0x80]
self._firstFragment = True
elif self._firstFragment:
## First fragment, final set to 0
data = [opcode]
self._firstFragment = False
else:
## Continuation
data = [OPCODE_CONT]
else:
data = [0x80 | opcode]
## Create the frame header, depending on the length
if length < 126:
spec = "!BBI"
data.append(0x80|length)
elif length > 65535:
# same for even longer frames
spec = "!BBQI"
data.append(0xff)
data.append(length)
else:
## add a 16-bit int to the spec and append 126 value, which means
## "interpret the next two bytes"
spec = "!BBHI"
data.append(0xfe)
data.append(length)
## Add the mask
mask = getrandbits(32)
data.append(mask)
header = pack(spec, *data)
payload = header + xxor(mask, payload)
self.transport.write(payload)
def dataReceived(self, data):
if self.in_handshake:
self.data += data
try:
remainingData = self.checkHandshake()
except HandshakeResponseError, error:
self.handshakeError(error)
self.transport.loseConnection()
return
## Header is over, witch to data handling mode
self.in_handshake = False
self.data = ""
## Just got response from server, might send something back
self.connectionEstablished()
self.decoder.dataReceived(remainingData)
else:
