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