From 6f1de06d0dc102e2f4f74e39ef399bf05cf1effd Mon Sep 17 00:00:00 2001
From: Nick Spinale <nick@nickspinale.com>
Date: Mon, 5 Feb 2024 11:47:48 +0000
Subject: [PATCH] Make HTTP server more robust

Signed-off-by: Nick Spinale <nick@nickspinale.com>
---
 .../microkit/http-server/pds/server/core/src/lib.rs   | 11 +++++------
 .../http-server/pds/server/core/src/server.rs         |  4 +++-
 2 files changed, 8 insertions(+), 7 deletions(-)

diff --git a/crates/examples/microkit/http-server/pds/server/core/src/lib.rs b/crates/examples/microkit/http-server/pds/server/core/src/lib.rs
index f2b460077..8da6b7e5e 100644
--- a/crates/examples/microkit/http-server/pds/server/core/src/lib.rs
+++ b/crates/examples/microkit/http-server/pds/server/core/src/lib.rs
@@ -129,13 +129,12 @@ async fn use_socket_for_https<D: fat::BlockDevice + 'static, T: fat::TimeSource
     tls_config: Arc<ServerConfig>,
     mut socket: TcpSocket,
 ) -> Result<(), ClosedError<AsyncRustlsError<TcpSocketError>>> {
-    socket.accept(HTTPS_PORT).await.unwrap(); // TODO
-
-    let mut conn = ServerConnector::from(tls_config)
-        .connect(socket)
-        .unwrap()
+    socket
+        .accept(HTTPS_PORT)
         .await
-        .unwrap();
+        .map_err(AsyncRustlsError::TransitError)?;
+
+    let mut conn = ServerConnector::from(tls_config).connect(socket)?.await?;
 
     server.handle_connection(&mut conn).await?;
 
diff --git a/crates/examples/microkit/http-server/pds/server/core/src/server.rs b/crates/examples/microkit/http-server/pds/server/core/src/server.rs
index 17b9f52c9..72fc59c12 100644
--- a/crates/examples/microkit/http-server/pds/server/core/src/server.rs
+++ b/crates/examples/microkit/http-server/pds/server/core/src/server.rs
@@ -38,7 +38,9 @@ impl<D: fat::BlockDevice + 'static, T: fat::TimeSource + 'static> Server<D, T> {
             let mut i = 0;
             loop {
                 let n = conn.read(&mut buf[i..]).await?;
-                assert_ne!(n, 0);
+                if n == 0 {
+                    return Err(ClosedError::Closed);
+                }
                 i += n;
                 if is_request_complete(&buf[..i]).unwrap_or(false) {
                     break;