From 214a2b477e642c520c243717188cb683a9cad900 Mon Sep 17 00:00:00 2001
From: Jacob Rothstein <hi@jbr.me>
Date: Wed, 6 Mar 2024 14:48:41 -0800
Subject: [PATCH] fix: correctly poll alert.write until it finishes

---
 src/common/handshake.rs | 13 ++++++-----
 tests/test.rs           | 49 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 56 insertions(+), 6 deletions(-)

diff --git a/src/common/handshake.rs b/src/common/handshake.rs
index e9d4f641..d1272fd5 100644
--- a/src/common/handshake.rs
+++ b/src/common/handshake.rs
@@ -51,15 +51,16 @@ where
                 mut io,
                 mut alert,
                 error,
-            } => {
-                return match alert.write(&mut SyncWriteAdapter { io: &mut io, cx }) {
-                    Err(ref e) if e.kind() == io::ErrorKind::WouldBlock => {
+            } => loop {
+                match alert.write(&mut SyncWriteAdapter { io: &mut io, cx }) {
+                    Err(e) if e.kind() == io::ErrorKind::WouldBlock => {
                         *this = MidHandshake::SendAlert { io, error, alert };
-                        Poll::Pending
+                        return Poll::Pending;
                     }
-                    _ => Poll::Ready(Err((error, io))),
+                    Err(_) | Ok(0) => return Poll::Ready(Err((error, io))),
+                    Ok(_) => {}
                 };
-            }
+            },
             // Starting the handshake returned an error; fail the future immediately.
             MidHandshake::Error { io, error } => return Poll::Ready(Err((error, io))),
             _ => panic!("unexpected polling after handshake"),
diff --git a/tests/test.rs b/tests/test.rs
index 8d3921b7..bd47bebb 100644
--- a/tests/test.rs
+++ b/tests/test.rs
@@ -241,5 +241,54 @@ async fn lazy_config_acceptor_take_io() -> Result<(), rustls::Error> {
     Ok(())
 }
 
+#[tokio::test]
+async fn acceptor_alert() {
+    let (sconfig, _) = utils::make_configs();
+    // this is the client hello from https://tls12.xargs.org/#client-hello/annotated with the minor
+    // version byte changed
+    let bad_hello = [
+        0x16, 0x03, 0x01, 0x00, 0xa5, 0x01, 0x00, 0x00, 0xa1, 0x03, 0x01, 0x00, 0x01, 0x02, 0x03,
+        0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f, 0x10, 0x11, 0x12,
+        0x13, 0x14, 0x15, 0x16, 0x17, 0x18, 0x19, 0x1a, 0x1b, 0x1c, 0x1d, 0x1e, 0x1f, 0x00, 0x00,
+        0x20, 0xcc, 0xa8, 0xcc, 0xa9, 0xc0, 0x2f, 0xc0, 0x30, 0xc0, 0x2b, 0xc0, 0x2c, 0xc0, 0x13,
+        0xc0, 0x09, 0xc0, 0x14, 0xc0, 0x0a, 0x00, 0x9c, 0x00, 0x9d, 0x00, 0x2f, 0x00, 0x35, 0xc0,
+        0x12, 0x00, 0x0a, 0x01, 0x00, 0x00, 0x58, 0x00, 0x00, 0x00, 0x18, 0x00, 0x16, 0x00, 0x00,
+        0x13, 0x65, 0x78, 0x61, 0x6d, 0x70, 0x6c, 0x65, 0x2e, 0x75, 0x6c, 0x66, 0x68, 0x65, 0x69,
+        0x6d, 0x2e, 0x6e, 0x65, 0x74, 0x00, 0x05, 0x00, 0x05, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00,
+        0x0a, 0x00, 0x0a, 0x00, 0x08, 0x00, 0x1d, 0x00, 0x17, 0x00, 0x18, 0x00, 0x19, 0x00, 0x0b,
+        0x00, 0x02, 0x01, 0x00, 0x00, 0x0d, 0x00, 0x12, 0x00, 0x10, 0x04, 0x01, 0x04, 0x03, 0x05,
+        0x01, 0x05, 0x03, 0x06, 0x01, 0x06, 0x03, 0x02, 0x01, 0x02, 0x03, 0xff, 0x01, 0x00, 0x01,
+        0x00, 0x00, 0x12, 0x00, 0x00,
+    ];
+
+    // Intentionally small so that we have to call alert.write several times
+    let (mut cstream, sstream) = tokio::io::duplex(2);
+
+    let (tx, rx) = oneshot::channel();
+
+    tokio::spawn(async move {
+        cstream.write_all(&bad_hello).await.unwrap();
+        let mut buf = Vec::new();
+        cstream.read_to_end(&mut buf).await.unwrap();
+        tx.send(buf).unwrap();
+    });
+
+    let accept = LazyConfigAcceptor::new(rustls::server::Acceptor::default(), sstream);
+
+    let Ok(Ok(start_handshake)) = time::timeout(Duration::from_secs(3), accept).await else {
+        panic!("timeout");
+    };
+
+    let err = start_handshake.into_stream(sconfig).await.unwrap_err();
+
+    assert_eq!(err.to_string(), "peer is incompatible: Tls12NotOffered");
+
+    let Ok(Ok(received)) = time::timeout(Duration::from_secs(3), rx).await else {
+        panic!("failed to receive");
+    };
+
+    assert_eq!(received, [0x15, 0x03, 0x03, 0x00, 0x02, 0x02, 0x46]);
+}
+
 // Include `utils` module
 include!("utils.rs");