From d4c83ba2c5268c3669dcab05fd7024386bf99f26 Mon Sep 17 00:00:00 2001
From: Matt Mastracci <matthew@mastracci.com>
Date: Tue, 2 Jan 2024 13:36:03 -0700
Subject: [PATCH] fix: don't panic if co-op budget exceeded (#22)

---
 src/stream.rs | 30 +++++++++---------------------
 1 file changed, 9 insertions(+), 21 deletions(-)

diff --git a/src/stream.rs b/src/stream.rs
index 6f0b3e7..0971a27 100644
--- a/src/stream.rs
+++ b/src/stream.rs
@@ -10,7 +10,6 @@ use crate::handshake::HandshakeResult;
 use crate::trace;
 use crate::TestOptions;
 use futures::future::poll_fn;
-use futures::task::noop_waker_ref;
 use futures::task::Context;
 use futures::task::Poll;
 use futures::task::Waker;
@@ -648,11 +647,8 @@ impl AsyncRead for TlsStream {
         } => {
           // If the handshake completed, we want to finalize it and then continue
           if handle.is_finished() {
-            let Poll::Ready(res) =
-              handle.poll_unpin(&mut Context::from_waker(noop_waker_ref()))
-            else {
-              unreachable!()
-            };
+            // This may return Pending if we've exhausted the co-op budget
+            let res = ready!(handle.poll_unpin(cx));
             self.finalize_handshake(res)?;
             continue;
           }
@@ -695,11 +691,8 @@ impl AsyncWrite for TlsStream {
         } => {
           // If the handshake completed, we want to finalize it and then continue
           if handle.is_finished() {
-            let Poll::Ready(res) =
-              handle.poll_unpin(&mut Context::from_waker(noop_waker_ref()))
-            else {
-              unreachable!()
-            };
+            // This may return Pending if we've exhausted the co-op budget
+            let res = ready!(handle.poll_unpin(cx));
             self.finalize_handshake(res)?;
             continue;
           }
@@ -752,11 +745,8 @@ impl AsyncWrite for TlsStream {
         } => {
           // If the handshake completed, we want to finalize it and then continue
           if handle.is_finished() {
-            let Poll::Ready(res) =
-              handle.poll_unpin(&mut Context::from_waker(noop_waker_ref()))
-            else {
-              unreachable!()
-            };
+            // This may return Pending if we've exhausted the co-op budget
+            let res = ready!(handle.poll_unpin(cx));
             self.finalize_handshake(res)?;
             continue;
           }
@@ -809,12 +799,10 @@ impl AsyncWrite for TlsStream {
           handle,
           ..
         } => {
+          // If the handshake completed, we want to finalize it and then continue
           if handle.is_finished() {
-            let Poll::Ready(res) =
-              handle.poll_unpin(&mut Context::from_waker(noop_waker_ref()))
-            else {
-              unreachable!()
-            };
+            // This may return Pending if we've exhausted the co-op budget
+            let res = ready!(handle.poll_unpin(cx));
             self.finalize_handshake(res)?;
             continue;
           }