Skip to content

Commit

Permalink
Fix unstable test
Browse files Browse the repository at this point in the history
  • Loading branch information
fantix committed Jul 18, 2018
1 parent 4b4958a commit 5b06936
Show file tree
Hide file tree
Showing 2 changed files with 40 additions and 10 deletions.
45 changes: 35 additions & 10 deletions tests/test_tcp.py
Original file line number Diff line number Diff line change
Expand Up @@ -2470,24 +2470,49 @@ def test_remote_shutdown_receives_trailing_data(self):
future = None

def server(sock):
sock.starttls(sslctx, server_side=True)
self.assertEqual(sock.recv_all(4), b'ping')
sock.send(b'pong')
incoming = ssl.MemoryBIO()
outgoing = ssl.MemoryBIO()
sslobj = sslctx.wrap_bio(incoming, outgoing, server_side=True)

while True:
try:
sslobj.do_handshake()
except ssl.SSLWantReadError:
if outgoing.pending:
sock.send(outgoing.read())
incoming.write(sock.recv(16384))
else:
if outgoing.pending:
sock.send(outgoing.read())
break

incoming.write(sock.recv(16384))
self.assertEqual(sslobj.read(4), b'ping')
sslobj.write(b'pong')
sock.send(outgoing.read())

time.sleep(0.2) # wait for the peer to fill its backlog

# send close_notify but don't wait for response
sock.setblocking(0)
with self.assertRaises(ssl.SSLWantReadError):
sock.unwrap()
sock.setblocking(1)
sslobj.unwrap()
sock.send(outgoing.read())

# should receive all data
data = sock.recv_all(CHUNK * SIZE)
self.assertEqual(len(data), CHUNK * SIZE)
data_len = 0
while True:
try:
chunk = len(sslobj.read(16384))
data_len += chunk
except ssl.SSLWantReadError:
incoming.write(sock.recv(16384))
except ssl.SSLZeroReturnError:
break

# wait for close_notify
sock.unwrap()
self.assertEqual(data_len, CHUNK * SIZE)

# verify that close_notify is received
sslobj.unwrap()

sock.close()

Expand Down
5 changes: 5 additions & 0 deletions uvloop/sslproto.pyx
Original file line number Diff line number Diff line change
Expand Up @@ -328,6 +328,11 @@ class SSLProtocol(object):
self._app_transport = None
self._wakeup_waiter(exc)

if getattr(self, '_shutdown_timeout_handle', None):
self._shutdown_timeout_handle.cancel()
if getattr(self, '_handshake_timeout_handle', None):
self._handshake_timeout_handle.cancel()

def get_buffer(self, n):
if len(self._ssl_buffer) < n:
self._ssl_buffer.extend(
Expand Down

0 comments on commit 5b06936

Please sign in to comment.