This commit is contained in:
maxsgit 2021-05-17 15:18:02 -07:00 committed by GitHub
commit 60aded1043
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -32,6 +32,8 @@ class Drained(WebSocketException):
class WebSocketClosed(WebSocketException): class WebSocketClosed(WebSocketException):
pass pass
class InterruptOrTimeout(WebSocketException):
pass
class WebSocketConnection(object): class WebSocketConnection(object):
connections = [] connections = []
@ -176,20 +178,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
@ -198,6 +205,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)])
@ -222,8 +231,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