async def readexactly(self, n: int) -> bytes:
if self._exception is not None:
raise self._exception
blocks = [] # type: List[bytes]
while n > 0:
block = await self.read(n)
if not block:
partial = b"".join(blocks)
raise asyncio.IncompleteReadError(partial, len(partial) + n)
blocks.append(block)
n -= len(block)
return b"".join(blocks)
async def read(self, n, datatype, exact):
"""Read data from the channel"""
recv_buf = self._recv_buf[datatype]
buf = '' if self._encoding else b''
data = []
async with self._read_locks[datatype]:
while True:
while recv_buf and n != 0:
if isinstance(recv_buf[0], Exception):
if data:
break
else:
exc = recv_buf.pop(0)
if isinstance(exc, SoftEOFReceived):
n = 0
break
else:
raise exc
l = len(recv_buf[0])
if l > n > 0:
data.append(recv_buf[0][:n])
recv_buf[0] = recv_buf[0][n:]
self._recv_buf_len -= n
n = 0
break
data.append(recv_buf.pop(0))
self._recv_buf_len -= l
n -= l
if self._maybe_resume_reading():
continue
if n == 0 or (n > 0 and data and not exact) or \
(n < 0 and recv_buf) or self._eof_received:
break
await self._block_read(datatype)
buf = buf.join(data)
if n > 0 and exact:
raise asyncio.IncompleteReadError(buf, len(buf) + n)
return buf
async def readexactly(self, n):
"""Read response data from Pageant agent"""
if self._writing:
cwnd = _find_agent_window()
if not cwnd.SendMessage(win32con.WM_COPYDATA, None, self._cds):
raise OSError(errno.EIO, 'Unable to send agent request')
self._writing = False
self._mapfile.seek(0)
result = self._mapfile.read(n)
if len(result) != n:
raise asyncio.IncompleteReadError(result, n)
return result
def readuntil(self, separator):
seplen = len(separator)
offset = 0
while True:
buflen = len(self._buffer)
if buflen - offset >= seplen:
isep = self._buffer.find(separator, offset)
if isep != -1:
break
offset = buflen + 1 - seplen
if self._eof:
chunk = bytes(self._buffer)
self._buffer.clear()
raise asyncio.IncompleteReadError(chunk, None)
yield from self._wait_for_data('readuntil')
chunk = self._buffer[:isep + seplen]
del self._buffer[:isep + seplen]
self._maybe_resume_transport()
return bytes(chunk)
async def read_redis(self) -> Union[bytes, List[bytes]]:
"""
Return a parsed Redis object or an exception when something wrong happened.
"""
if self._exception is not None:
raise self._exception
while True:
obj = self.hi_reader.gets()
if obj is not False:
self._maybe_resume_transport()
return obj
if self._eof:
self.hi_reader = hiredis.Reader()
raise asyncio.IncompleteReadError(b'<redis>', None)
await self._wait_for_data('read_redis')
def readexactly(self, n):
'''
Implementation of StreamReader.readexactly which has an O(n) instead of
O(2n) memory footprint (as is the case for asyncio in python 3.5.2). Also
it is 25% to 35% faster in terms of throughput. See also
https://github.com/python/asyncio/issues/394.
'''
if self._exception is not None:
raise self._exception
if not n:
return b''
if len(self._buffer) >= n and not self._eof:
data = self._buffer[:n]
del self._buffer[:n]
self._maybe_resume_transport()
return data
data = bytearray(n)
pos = 0
while n:
if not self._buffer:
yield from self._wait_for_data('readexactly')
if self._eof or not self._buffer:
raise asyncio.IncompleteReadError(data[:pos], pos + n)
available = len(self._buffer)
if available <= n:
data[pos:pos + available] = self._buffer
self._buffer.clear()
n -= available
pos += available
else:
data[pos:pos + n] = self._buffer[:n]
del self._buffer[:n]
n = 0
self._maybe_resume_transport()
return data
def read(self, n, datatype, exact):
"""Read data from the channel"""
recv_buf = self._recv_buf[datatype]
buf = '' if self._chan.get_encoding() else b''
data = []
while True:
while recv_buf and n != 0:
if isinstance(recv_buf[0], Exception):
if data:
break
else:
raise recv_buf.pop(0)
l = len(recv_buf[0])
if n > 0 and l > n:
data.append(recv_buf[0][:n])
recv_buf[0] = recv_buf[0][n:]
self._recv_buf_len -= n
n = 0
break
data.append(recv_buf.pop(0))
self._recv_buf_len -= l
n -= l
if self._read_paused and self._recv_buf_len < self._limit:
self._read_paused = False
self._chan.resume_reading()
continue
if n == 0 or (n > 0 and data and not exact) or self._eof_received:
break
yield from self._block_read(datatype)
buf = buf.join(data)
if n > 0 and exact:
raise asyncio.IncompleteReadError(buf, len(buf) + n)
return buf
class HttpPacketCodec(PacketCodec):
tag = None
obfuscate_tag = None
def encode_packet(self, data):
"""
Encode packet to binary.
Args:
self: (todo): write your description
data: (todo): write your description
"""
return ('POST /api HTTP/1.1\r\n'
'Host: {}:{}\r\n'
'Content-Type: application/x-www-form-urlencoded\r\n'
'Connection: keep-alive\r\n'
'Keep-Alive: timeout=100000, max=10000000\r\n'
'Content-Length: {}\r\n\r\n'
.format(self._conn._ip, self._conn._port, len(data))
.encode('ascii') + data)
async def read_packet(self, reader):
"""
Read a packet from the socket.
Args:
self: (todo): write your description
reader: (todo): write your description
"""
while True:
line = await reader.readline()
if not line or line[-1] != b'\n':
raise asyncio.IncompleteReadError(line, None)
if line.lower().startswith(b'content-length: '):
await reader.readexactly(2)
length = int(line[16:-2])
return await reader.readexactly(length)
def test_that_when_authentication_is_successful_connection_lasts_until_its_end(
self, event_loop):
with override_settings(
CONCENT_PRIVATE_KEY=CONCENT_PRIVATE_KEY,
CONCENT_PUBLIC_KEY=CONCENT_PUBLIC_KEY,
SIGNING_SERVICE_PUBLIC_KEY=SIGNING_SERVICE_PUBLIC_KEY,
):
self.middleman._is_signing_service_connection_active = False
some_bytes = b"some bytes"
mocked_reader = prepare_mocked_reader(
some_bytes,
side_effect=asyncio.IncompleteReadError(some_bytes, 123))
mocked_writer = prepare_mocked_writer()
with mock.patch.object(
self.middleman,
"_authenticate_signing_service",
new=async_stream_actor_mock(return_value=True)):
event_loop.run_until_complete(
self.middleman._handle_service_connection(
mocked_reader, mocked_writer))
self.crash_logger_mock.assert_not_called()
def _handle_responses(self, d, line_handler, fetch_handler, incomplete_line=b''):
if not d:
return
data = incomplete_line + d
match_fetch_message = fetch_message_with_literal_data_re.match(data)
if match_fetch_message:
head, crlf, tail = data.partition(CRLF)
msg_size = match_fetch_message.group('size')
# we want to cut -----------------------
# ...here |
# so 4+1 v
# b'* 3 FETCH (UID 3 RFC822 {4}\r\nmail)\r\n...
end_message_index_with_parenthesis = int(msg_size) + 1
fetch_handler(head + crlf + tail[0:end_message_index_with_parenthesis], end_message_index_with_parenthesis)
after_fetch = tail[end_message_index_with_parenthesis:]
self._handle_responses(after_fetch, line_handler, fetch_handler)
else:
line, separator, tail = data.partition(CRLF)
if not separator:
raise asyncio.IncompleteReadError(data, b'line should end with CRLF')
else:
line_handler(line.decode())
self._handle_responses(tail, line_handler, fetch_handler)
async def readexactly(self, n: int) -> bytes:
raise asyncio.IncompleteReadError(b"", n)
def readexactly(self, size):
data = self.read(size)
if len(data) < size:
raise asyncio.IncompleteReadError(data, size)
else:
return data
async def readuntil(self, separators=(b'\n', b'\r')):
"""Read data from the stream until ``separator`` is found.
On success, the data and separator will be removed from the
internal buffer (consumed). Returned data will include the
separator at the end.
Configured stream limit is used to check result. Limit sets the
maximal length of data that can be returned, not counting the
separator.
If an EOF occurs and the complete separator is still not found,
an IncompleteReadError exception will be raised, and the internal
buffer will be reset. The IncompleteReadError.partial attribute
may contain the separator partially.
If the data cannot be read because of over limit, a
LimitOverrunError exception will be raised, and the data
will be left in the internal buffer, so it can be read again.
"""
assert separators
# seplen = len(separator)
# 这里强制设为1无关紧要，只要separators的元素都不为空就行
seplen = 1
if seplen == 0:
raise ValueError('Separator should be at least one-byte string')
if self._exception is not None:
raise self._exception
# Consume whole buffer except last bytes, which length is
# one less than seplen. Let's check corner cases with
# separator='SEPARATOR':
# * we have received almost complete separator (without last
# byte). i.e buffer='some textSEPARATO'. In this case we
# can safely consume len(separator) - 1 bytes.
# * last byte of buffer is first byte of separator, i.e.
# buffer='abcdefghijklmnopqrS'. We may safely consume
# everything except that last byte, but this require to
# analyze bytes of buffer that match partial separator.
# This is slow and/or require FSM. For this case our
# implementation is not optimal, since require rescanning
# of data that is known to not belong to separator. In
# real world, separator will not be so long to notice
# performance problems. Even when reading MIME-encoded
# messages :)
# `offset` is the number of bytes from the beginning of the buffer
# where there is no occurrence of `separator`.
offset = 0
# Loop until we find `separator` in the buffer, exceed the buffer size,
# or an EOF has happened.
while True:
buflen = len(self._buffer)
# Check if we now have enough data in the buffer for `separator` to
# fit.
if buflen - offset >= seplen:
# 多种处理方式，并且有顺序的，先检测 \n，然后是\r
# 并且保证\r不在最后一字节，以优先匹配\r\n
for separator in separators:
isep = self._buffer.find(separator, offset)
try:
if separator == b'\r' and isep != -1 and self._buffer[
isep + 1] == b'\n':
continue
except IndexError:
pass
if isep != -1:
# `separator` is in the buffer. `isep` will be used later
# to retrieve the data.
break
else:
# see upper comment for explanation.
offset = buflen + 1 - seplen
if offset > self._limit:
raise asyncio.LimitOverrunError(
'Separator is not found, and chunk exceed the limit',
offset)
break
# Complete message (with full separator) may be present in buffer
# even when EOF flag is set. This may happen when the last chunk
# adds data which makes separator be found. That's why we check for
# EOF *ater* inspecting the buffer.
if self._eof:
chunk = bytes(self._buffer)
self._buffer.clear()
raise asyncio.IncompleteReadError(chunk, None)
# _wait_for_data() will resume reading if stream was paused.
await self._wait_for_data('readuntil')
if isep > self._limit:
raise asyncio.LimitOverrunError(
'Separator is found, but chunk is longer than limit', isep)
chunk = self._buffer[:isep + seplen]
del self._buffer[:isep + seplen]
self._maybe_resume_transport()
return bytes(chunk)
async def readexactly(self, n: int) -> bytes:
data = await self.read(n)
if len(data) < n:
raise asyncio.IncompleteReadError(data, n)
return data
async def __connect_request_response(self):
data = ConnectRequest()
data.request_id = UtilInternal.get_request_id()
data.sdk_version = SDK_VERSION
data.other_supported_protocol_versions = self.__OLD_SUPPORTED_PROTOCOL_VERSIONS
data.protocol_version = VersionInfo.PROTOCOL_VERSION.value
if self.auth_token is not None:
data.auth_token = self.auth_token
# Write the connect version
self.__writer.write(
UtilInternal.int_to_bytes(self.__CONNECT_VERSION, 1))
# Write request to socket
frame = MessageFrame(operation=Operation.Connect,
payload=cbor2.dumps(data.as_dict()))
for b in UtilInternal.encode_frame(frame):
self.__writer.write(b)
await self.__writer.drain()
# Read connect version
connect_response_version_byte = await self.__reader.read(n=1)
if len(connect_response_version_byte) == 0:
raise asyncio.IncompleteReadError(connect_response_version_byte, 1)
connect_response_version = UtilInternal.int_from_bytes(
connect_response_version_byte)
if connect_response_version != self.__CONNECT_VERSION:
self.logger.error(
"Unexpected response from the server, Connect version: %s.",
connect_response_version)
raise ConnectFailedException(
"Failed to establish connection with the server")
# Read connect response
response = await self.__read_message_frame() # type: MessageFrame
if response.operation == Operation.ConnectResponse:
payload = cbor2.loads(response.payload)
response = ConnectResponse.from_dict(
payload) # type: ConnectResponse
self.logger.debug("Received ConnectResponse from server: %s",
response)
else:
self.logger.error("Received data with unexpected operation %s.",
response.operation)
raise ConnectFailedException(
"Failed to establish connection with the server")
if response.status != ResponseStatusCode.Success:
self.logger.error(
"Received ConnectResponse with unexpected status %s.",
response.status)
raise ConnectFailedException(
"Failed to establish connection with the server")
if data.protocol_version != response.protocol_version:
self.logger.warn(
"SDK with version %s using Protocol version %s is not fully compatible with Server with version %s. "
"Client has connected in a compatibility mode using protocol version %s. "
"Some features will not work as expected",
SDK_VERSION,
data.protocol_version,
response.server_version,
response.protocol_version,
)
async def readuntil(self, separator, datatype):
"""Read data from the channel until a separator is seen"""
if separator is _NEWLINE:
separator = '\n' if self._encoding else b'\n'
elif not separator:
raise ValueError('Separator cannot be empty')
if isinstance(separator, (str, bytes)):
separators = [separator]
else:
separators = list(separator)
seplen = max(len(sep) for sep in separators)
bar = '|' if self._encoding else b'|'
pat = re.compile(bar.join(map(re.escape, separators)))
recv_buf = self._recv_buf[datatype]
buf = '' if self._encoding else b''
curbuf = 0
buflen = 0
async with self._read_locks[datatype]:
while True:
while curbuf < len(recv_buf):
if isinstance(recv_buf[curbuf], Exception):
if buf:
recv_buf[:curbuf] = []
self._recv_buf_len -= buflen
raise asyncio.IncompleteReadError(buf, None)
else:
exc = recv_buf.pop(0)
if isinstance(exc, SoftEOFReceived):
return buf
else:
raise exc
buf += recv_buf[curbuf]
start = max(buflen + 1 - seplen, 0)
match = pat.search(buf, start)
if match:
idx = match.end()
recv_buf[:curbuf] = []
recv_buf[0] = buf[idx:]
buf = buf[:idx]
self._recv_buf_len -= idx
if not recv_buf[0]:
recv_buf.pop(0)
self._maybe_resume_reading()
return buf
buflen += len(recv_buf[curbuf])
curbuf += 1
if self._read_paused or self._eof_received:
recv_buf[:curbuf] = []
self._recv_buf_len -= buflen
self._maybe_resume_reading()
raise asyncio.IncompleteReadError(buf, None)
await self._block_read(datatype)
async def readexactly(self, n):
data = self.io.read(n)
if len(data) != n:
raise asyncio.IncompleteReadError(data, n)
return data