[PATCH proxmox v3 4/5] rest-server: connection: avoid dropping clients during TLS check

["Max R. Carrara" <m.carrara@proxmox.com>]

While waiting for enough data to check for a TLS handshake, the
connection was considered closed as soon as the socket was reported
readable twice in a row with the same amount of data queued.

This can also happen while the client is still sending: if more data
arrives after tokio received the readiness event for the previous data
but before the peek, tokio keeps the socket marked as readable, and
the next peek sees no new data. The connection was then dropped and an
error was logged, even though the client did nothing wrong.

Only clients whose initial message is shorter than the five bytes
needed for the check are affected (both plain and TLS ones).

Note that this should occur very rarely (if at all) in practice,
because most clients send more than 5 bytes at once. I cannot actually
think of any scenario where this might happen in practice (unless you
intentionally split up your messages), but it's better to fix this
anyway.

So, only consider the connection closed if:
- the peek returns no data at all
- or if the read side of the socket was closed by the client

Both of these cases are definite signs that no further data will
arrive.

(I'm being more elaborate below than I usually am, because the changes
are not trivial at all.)

First, replace `incoming_stream.async_io()` with a plain `loop` and
await the readiness of the socket explicitly via
`incoming_stream.ready()`.

This is necessary because the closure passed to `async_io()` does not
get any readiness information, which we however need to actually check
what the client did.

Use the readiness info to check for `is_read_closed()`, which means
that no more data will arrive on the socket beyond what's already
queued. Now we know for sure that the connection is being closed.

Second, use `try_io()` for the buffer peeking + handshake check. This
is basically like `async_io()`, except that we are not awaiting the
readiness of the socket. We awaited that already.
(So, `.ready()` + `.try_io()` = `.async_io()`, kinda.)

Third, match on the result of `try_io()` to check for a `WouldBlock`
error. If this is returned, it means that the socket either wasn't
*actually* ready [1], or that the client sent < 5 bytes. Either way,
the loop continues in that case.

Fourth, await `tokio::task::consume_budget()` directly to charge each
iteration against tokio's cooperative scheduling budget (like tokio's
own IO helpers), so that the loop cannot starve other tasks or the
timeout.

To summarize: We explicitly (or "manually") check for the
readiness of the socket to see whether the client is closing the
connection or not, instead of letting `async_io()` handle this under
the hood. We also play nice with tokio's scheduling budget.

This should ensure that no connections are dropped unintentionally
anymore, even if their occurrence is (or at least should be)
incredibly rare.

[1]: https://docs.rs/tokio/latest/tokio/net/struct.TcpStream.html#method.try_io

Signed-off-by: Max R. Carrara <m.carrara@proxmox.com>
---
 proxmox-rest-server/src/connection.rs | 57 ++++++++++++++++-----------
 1 file changed, 34 insertions(+), 23 deletions(-)

diff --git a/proxmox-rest-server/src/connection.rs b/proxmox-rest-server/src/connection.rs
index 9b3c1371..f116ef47 100644
--- a/proxmox-rest-server/src/connection.rs
+++ b/proxmox-rest-server/src/connection.rs
@@ -501,64 +501,75 @@ impl AcceptBuilder {
         timeout: Duration,
     ) -> Result<TlsHandshakeCheck, Error> {
         const HANDSHAKE_BYTES_LEN: usize = 5;
 
         let future = async {
-            let mut previous_peek_len = 0;
-            incoming_stream
-                .async_io(tokio::io::Interest::READABLE, || {
+            loop {
+                // Charge each iteration against tokio's cooperative scheduling
+                // budget, like tokio's own IO helpers do, so that the loop
+                // cannot starve other tasks or the timeout.
+                tokio::task::consume_budget().await;
+
+                let ready = incoming_stream.ready(tokio::io::Interest::READABLE).await?;
+
+                let peek_res = incoming_stream.try_io(tokio::io::Interest::READABLE, || {
                     let mut buf = [0; HANDSHAKE_BYTES_LEN];
 
                     use std::os::fd::{AsRawFd, FromRawFd};
 
                     // Convert to standard lib TcpStream so we can peek without interfering
                     // with tokio's internals. Wrap the stream in ManuallyDrop in order to prevent
                     // the destructor from being called, closing the connection and messing up
                     // invariants.
                     let raw_fd = incoming_stream.as_raw_fd();
                     let std_stream =
                         unsafe { ManuallyDrop::new(std::net::TcpStream::from_raw_fd(raw_fd)) };
 
-                    let peek_res = std_stream.peek(&mut buf);
-
-                    match peek_res {
+                    match std_stream.peek(&mut buf) {
                         // If we didn't get enough bytes, raise an EAGAIN / EWOULDBLOCK which tells
-                        // tokio to await the readiness of the socket again. This should normally
-                        // only be used if the socket isn't actually ready, but is fine to do here
-                        // in our case.
+                        // tokio to clear the readiness of the socket, so that we wait for more
+                        // data. This should normally only be used if the socket isn't actually
+                        // ready, but is fine to do here in our case.
                         //
-                        // This means we will peek into the stream's queue until we got
-                        // HANDSHAKE_BYTE_LEN bytes or an error.
+                        // Only stop waiting on definite signs of a closed connection:
+                        // 1. an empty peek
+                        // 2. a closed read side
+                        //
+                        // The socket can be marked as readable again without any new data, for
+                        // example if more data arrived between tokio receiving a readiness event
+                        // and `peek()`. So, readiness *alone* is not a sign of a closed connection.
                         Ok(peek_len) if peek_len < HANDSHAKE_BYTES_LEN => {
-                            // if we detect the same peek len again but still got a readable stream,
-                            // the connection was probably closed, so abort here
-                            if peek_len == previous_peek_len {
+                            if peek_len == 0 || ready.is_read_closed() {
                                 Err(io::ErrorKind::ConnectionAborted.into())
                             } else {
-                                previous_peek_len = peek_len;
                                 Err(io::ErrorKind::WouldBlock.into())
                             }
                         }
                         // Either we got Ok(HANDSHAKE_BYTES_LEN) or some error.
                         res => res.map(|_| {
                             if contains_tls_handshake_fragment(&buf) {
                                 TlsHandshakeCheck::Initiated
                             } else {
                                 TlsHandshakeCheck::NotInitiated
                             }
                         }),
                     }
-                })
-                .await
-                .context("couldn't peek into incoming TCP stream")
+                });
+
+                match peek_res {
+                    Err(err) if err.kind() == io::ErrorKind::WouldBlock => continue,
+                    res => break res,
+                }
+            }
         };
 
-        // The elapsed error carries no information beyond the timeout itself, so do not keep it as
-        // the cause.
-        tokio::time::timeout(timeout, future).await.map_err(|_| {
-            format_err!("timed out while waiting for client to initiate TLS handshake")
-        })?
+        tokio::time::timeout(timeout, future)
+            .await
+            .map_err(|_| {
+                format_err!("timed out while waiting for client to initiate TLS handshake")
+            })?
+            .context("couldn't peek into incoming TCP stream")
     }
 }
 
 /// Checks whether an [SSL 3.0 / TLS plaintext fragment][0] being part of a
 /// SSL / TLS handshake is contained in the given buffer.
-- 
2.47.3