Improve web socket implementation to wait for message body after receiving header

This commit is contained in:
Max Sasaga 2020-11-16 23:37:36 +02:00
parent 2579b9be26
commit 14f6c25eb6

View File

@ -31,6 +31,8 @@ class Drained(WebSocketException):
class WebSocketClosed(WebSocketException): class WebSocketClosed(WebSocketException):
pass pass
class InterruptOrTimeout(WebSocketException):
pass
class WebSocketConnection(object): class WebSocketConnection(object):
connections = [] connections = []
@ -175,20 +177,25 @@ class WebSocketConnection(object):
def protected_read(num): def protected_read(num):
data = self.handler.rfile.read(num) data = self.handler.rfile.read(num)
if data is None: if data is None:
raise Drained() data = bytes()
if len(data) != num: while self.open and len(data) < num:
(read, _, _) = select.select([self.interruptPipeRecv, self.handler.rfile], [], [], 15)
if self.handler.rfile in read:
data += self.handler.rfile.read(num - len(data))
else:
if len(data) == 0:
raise InterruptOrTimeout()
else:
raise IncompleteRead()
if len(data) < num:
raise IncompleteRead() raise IncompleteRead()
return data return data
self.open = True self.open = True
while self.open: while self.open:
(read, _, _) = select.select([self.interruptPipeRecv, self.handler.rfile], [], [], 15)
if self.handler.rfile in read:
available = True
self.resetPing()
while self.open and available:
try: try:
header = protected_read(2) header = protected_read(2)
self.resetPing()
opcode = header[0] & 0x0F opcode = header[0] & 0x0F
length = header[1] & 0x7F length = header[1] & 0x7F
mask = (header[1] & 0x80) >> 7 mask = (header[1] & 0x80) >> 7
@ -197,6 +204,8 @@ class WebSocketConnection(object):
length = (header[0] << 8) + header[1] length = (header[0] << 8) + header[1]
if mask: if mask:
masking_key = protected_read(4) masking_key = protected_read(4)
data = None
if length > 0:
data = protected_read(length) data = protected_read(length)
if mask: if mask:
data = bytes([b ^ masking_key[index % 4] for (index, b) in enumerate(data)]) data = bytes([b ^ masking_key[index % 4] for (index, b) in enumerate(data)])
@ -221,8 +230,8 @@ class WebSocketConnection(object):
self.open = False self.open = False
else: else:
logger.warning("unsupported opcode: {0}".format(opcode)) logger.warning("unsupported opcode: {0}".format(opcode))
except Drained: except InterruptOrTimeout:
available = False pass
except IncompleteRead: except IncompleteRead:
logger.warning("incomplete read on websocket; closing connection") logger.warning("incomplete read on websocket; closing connection")
self.open = False self.open = False