server: enable holding back acceptor alerts

Co-authored-by: John Howard <[email protected]>

Author: djc

src/server.rs

@@ -70,7 +70,8 @@ impl TlsAcceptor {
 pub struct LazyConfigAcceptor<IO> {
     acceptor: rustls::server::Acceptor,
     io: Option<IO>,
-    alert: Option<(rustls::Error, AcceptedAlert)>,
+    alert: Option<AlertState>,
+    send_alert: bool,
 }
 
 impl<IO> LazyConfigAcceptor<IO>
@@ -83,9 +84,33 @@ where
             acceptor,
             io: Some(io),
             alert: None,
+            send_alert: true,
         }
     }
 
+    /// Configure whether to send a TLS alert on failure.
+    pub fn send_alert(mut self, send: bool) -> Self {
+        self.send_alert = send;
+        self
+    }
+
+    /// Writes a stored alert, consuming the alert (if any) and IO.
+    pub async fn write_alert(&mut self) -> io::Result<()> {
+        let Some(alert) = self.take_alert() else {
+            return Ok(());
+        };
+
+        let Some(io) = self.take_io() else {
+            return Ok(());
+        };
+
+        WritingAlert {
+            io,
+            alert: Some(alert),
+        }
+        .await
+    }
+
     /// Takes back the client connection. Will return `None` if called more than once or if the
     /// connection has been accepted.
     ///
@@ -130,6 +155,14 @@ where
     pub fn take_io(&mut self) -> Option<IO> {
         self.io.take()
     }
+
+    pub fn take_alert(&mut self) -> Option<AcceptedAlert> {
+        match self.alert.take() {
+            Some(AlertState::Sending(_, alert)) => Some(alert),
+            Some(AlertState::Saved(alert)) => Some(alert),
+            None => None,
+        }
+    }
 }
 
 impl<IO> Future for LazyConfigAcceptor<IO>
@@ -151,17 +184,17 @@ where
                 }
             };
 
-            if let Some((err, mut alert)) = this.alert.take() {
+            if let Some(AlertState::Sending(err, mut alert)) = this.alert.take() {
                 match alert.write(&mut SyncWriteAdapter { io, cx }) {
                     Err(e) if e.kind() == io::ErrorKind::WouldBlock => {
-                        this.alert = Some((err, alert));
+                        this.alert = Some(AlertState::Sending(err, alert));
                         return Poll::Pending;
                     }
                     Ok(0) | Err(_) => {
                         return Poll::Ready(Err(io::Error::new(io::ErrorKind::InvalidData, err)))
                     }
                     Ok(_) => {
-                        this.alert = Some((err, alert));
+                        this.alert = Some(AlertState::Sending(err, alert));
                         continue;
                     }
                 };
@@ -181,9 +214,49 @@ where
                     return Poll::Ready(Ok(StartHandshake { accepted, io }));
                 }
                 Ok(None) => {}
-                Err((err, alert)) => {
-                    this.alert = Some((err, alert));
+                Err((err, alert)) => match this.send_alert {
+                    true => this.alert = Some(AlertState::Sending(err, alert)),
+                    false => {
+                        this.alert = Some(AlertState::Saved(alert));
+                        return Poll::Ready(Err(io::Error::new(io::ErrorKind::InvalidData, err)));
+                    }
+                },
+            }
+        }
+    }
+}
+
+enum AlertState {
+    Sending(rustls::Error, AcceptedAlert),
+    Saved(AcceptedAlert),
+}
+
+struct WritingAlert<IO> {
+    io: IO,
+    alert: Option<AcceptedAlert>,
+}
+
+impl<IO> Future for WritingAlert<IO>
+where
+    IO: AsyncRead + AsyncWrite + Unpin,
+{
+    type Output = Result<(), io::Error>;
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        let this = self.get_mut();
+        let io = &mut this.io;
+        loop {
+            let Some(mut alert) = this.alert.take() else {
+                return Poll::Ready(Ok(()));
+            };
+
+            match alert.write(&mut SyncWriteAdapter { io, cx }) {
+                Ok(0) => return Poll::Ready(Ok(())),
+                Ok(_) => continue,
+                Err(e) if e.kind() == io::ErrorKind::WouldBlock => {
+                    this.alert = Some(alert);
+                    return Poll::Pending;
                 }
+                Err(e) => return Poll::Ready(Err(io::Error::new(io::ErrorKind::InvalidData, e))),
             }
         }
     }