def test_receive_close_frame(self):
frame = ABNF(opcode=ABNF.OPCODE_CLOSE, data='time to say goodbye')
client = self.make_one(frame)
client._connect_and_receive_and_log_messages()
assert self._dummy_connection.nothing_sent is False
assert self._dummy_connection.connected is False
assert self._dummy_connection.close_reason == b'server triggered close'
def ws_send(ws, end_event):
while not end_event.is_set():
try:
try:
data = send_queue.get_nowait()
except queue.Empty:
data = log_send_queue.get(timeout=1)
for i in range(0, len(data), WS_FRAME_SIZE):
frame = data[i:i + WS_FRAME_SIZE]
last = i + WS_FRAME_SIZE >= len(data)
opcode = ABNF.OPCODE_TEXT if i == 0 else ABNF.OPCODE_CONT
ws.send_frame(ABNF.create_frame(frame, opcode, last))
except queue.Empty:
pass
except Exception:
cloudlog.exception("athenad.ws_send.exception")
end_event.set()
def _proxy(self):
channel_ports = []
channel_initialized = []
local_ports = {}
for port in self.local_ports.values():
# Setup the data channel for this port number
channel_ports.append(port)
channel_initialized.append(False)
# Setup the error channel for this port number
channel_ports.append(port)
channel_initialized.append(False)
port.python.setblocking(True)
local_ports[port.python] = port
# The data to send on the websocket socket
kubernetes_data = b''
while True:
rlist = [] # List of sockets to read from
wlist = [] # List of sockets to write to
if self.websocket.connected:
rlist.append(self.websocket)
if kubernetes_data:
wlist.append(self.websocket)
local_all_closed = True
for port in self.local_ports.values():
if port.python.fileno() != -1:
if port.error or not self.websocket.connected:
if port.data:
wlist.append(port.python)
local_all_closed = False
else:
port.python.close()
else:
rlist.append(port.python)
if port.data:
wlist.append(port.python)
local_all_closed = False
if local_all_closed and not (self.websocket.connected and kubernetes_data):
self.websocket.close()
return
r, w, _ = select.select(rlist, wlist, [])
for sock in r:
if sock == self.websocket:
opcode, frame = self.websocket.recv_data_frame(True)
if opcode == ABNF.OPCODE_BINARY:
if not frame.data:
raise RuntimeError("Unexpected frame data size")
channel = six.byte2int(frame.data)
if channel >= len(channel_ports):
raise RuntimeError("Unexpected channel number: %s" % channel)
port = channel_ports[channel]
if channel_initialized[channel]:
if channel % 2:
if port.error is None:
port.error = ''
port.error += frame.data[1:].decode()
else:
port.data += frame.data[1:]
else:
if len(frame.data) != 3:
raise RuntimeError(
"Unexpected initial channel frame data size"
)
port_number = six.byte2int(frame.data[1:2]) + (six.byte2int(frame.data[2:3]) * 256)
if port_number != port.port_number:
raise RuntimeError(
"Unexpected port number in initial channel frame: %s" % port_number
)
channel_initialized[channel] = True
elif opcode not in (ABNF.OPCODE_PING, ABNF.OPCODE_PONG, ABNF.OPCODE_CLOSE):
raise RuntimeError("Unexpected websocket opcode: %s" % opcode)
else:
port = local_ports[sock]
data = port.python.recv(1024 * 1024)
if data:
kubernetes_data += ABNF.create_frame(
port.channel + data,
ABNF.OPCODE_BINARY,
).format()
else:
port.python.close()
for sock in w:
if sock == self.websocket:
sent = self.websocket.sock.send(kubernetes_data)
kubernetes_data = kubernetes_data[sent:]
else:
port = local_ports[sock]
sent = port.python.send(port.data)
port.data = port.data[sent:]
def test_receive_text_frame(self):
frame = ABNF(opcode=ABNF.OPCODE_TEXT, data='hello dear')
client = self.make_one(frame)
client._connect_and_receive_and_log_messages()
assert self._dummy_connection.nothing_sent is True
assert self._dummy_connection.connected is True
def test_receive_ping_frame(self):
frame = ABNF(opcode=ABNF.OPCODE_PING, data='hello')
client = self.make_one(frame)
client._connect_and_receive_and_log_messages()
assert self._dummy_connection.nothing_sent is False
assert self._dummy_connection.pong_text == 'Hi!'