LibCore+Userland: Convert TCPServer to use the Serenity Stream API

This is intended as a real-usecase test of the Serenity Stream API, and
seemed like a good candidate due to its low amount of users.

3 files changed, 106 insertions, 42 deletions

diff --git a/Userland/Services/WebServer/Client.cpp b/Userland/Services/WebServer/Client.cpp
index cf6e663f31..dcb3e776b2 100644
--- a/Userland/Services/WebServer/Client.cpp
+++ b/Userland/Services/WebServer/Client.cpp
@@ -28,41 +28,75 @@
 
 namespace WebServer {
 
-Client::Client(NonnullRefPtr<Core::TCPSocket> socket, Core::Object* parent)
+Client::Client(Core::Stream::BufferedTCPSocket socket, Core::Object* parent)
     : Core::Object(parent)
-    , m_socket(socket)
+    , m_socket(move(socket))
 {
 }
 
 void Client::die()
 {
+    m_socket.close();
     deferred_invoke([this] { remove_from_parent(); });
 }
 
 void Client::start()
 {
-    m_socket->on_ready_to_read = [this] {
+    m_socket.on_ready_to_read = [this] {
         StringBuilder builder;
+
+        auto maybe_buffer = ByteBuffer::create_uninitialized(m_socket.buffer_size());
+        if (!maybe_buffer.has_value()) {
+            warnln("Could not create buffer for client (possibly out of memory)");
+            die();
+            return;
+        }
+
+        auto buffer = maybe_buffer.release_value();
         for (;;) {
-            auto line = m_socket->read_line();
-            if (line.is_empty())
+            auto maybe_can_read = m_socket.can_read_without_blocking();
+            if (maybe_can_read.is_error()) {
+                warnln("Failed to get the blocking status for the socket: {}", maybe_can_read.error());
+                die();
+                return;
+            }
+
+            if (!maybe_can_read.value())
+                break;
+
+            auto maybe_nread = m_socket.read_until_any_of(buffer, Array { "\r"sv, "\n"sv, "\r\n"sv });
+            if (maybe_nread.is_error()) {
+                warnln("Failed to read a line from the request: {}", maybe_nread.error());
+                die();
+                return;
+            }
+
+            if (m_socket.is_eof()) {
+                die();
                 break;
-            builder.append(line);
+            }
+
+            builder.append(StringView { buffer.data(), maybe_nread.value() });
             builder.append("\r\n");
         }
 
         auto request = builder.to_byte_buffer();
         dbgln_if(WEBSERVER_DEBUG, "Got raw request: '{}'", String::copy(request));
-        handle_request(request);
+
+        auto maybe_did_handle = handle_request(request);
+        if (maybe_did_handle.is_error()) {
+            warnln("Failed to handle the request: {}", maybe_did_handle.error());
+        }
+
         die();
     };
 }
 
-void Client::handle_request(ReadonlyBytes raw_request)
+ErrorOr<bool> Client::handle_request(ReadonlyBytes raw_request)
 {
     auto request_or_error = HTTP::HttpRequest::from_raw_request(raw_request);
     if (!request_or_error.has_value())
-        return;
+        return false;
     auto& request = request_or_error.value();
 
     if constexpr (WEBSERVER_DEBUG) {
@@ -73,16 +107,16 @@ void Client::handle_request(ReadonlyBytes raw_request)
     }
 
     if (request.method() != HTTP::HttpRequest::Method::GET) {
-        send_error_response(501, request);
-        return;
+        TRY(send_error_response(501, request));
+        return false;
     }
 
     // Check for credentials if they are required
     if (Configuration::the().credentials().has_value()) {
         bool has_authenticated = verify_credentials(request.headers());
         if (!has_authenticated) {
-            send_error_response(401, request, { "WWW-Authenticate: Basic realm=\"WebServer\", charset=\"UTF-8\"" });
-            return;
+            TRY(send_error_response(401, request, { "WWW-Authenticate: Basic realm=\"WebServer\", charset=\"UTF-8\"" }));
+            return false;
         }
     }
 
@@ -102,8 +136,8 @@ void Client::handle_request(ReadonlyBytes raw_request)
             red.append(requested_path);
             red.append("/");
 
-            send_redirect(red.to_string(), request);
-            return;
+            TRY(send_redirect(red.to_string(), request));
+            return true;
         }
 
         StringBuilder index_html_path_builder;
@@ -111,29 +145,30 @@ void Client::handle_request(ReadonlyBytes raw_request)
         index_html_path_builder.append("/index.html");
         auto index_html_path = index_html_path_builder.to_string();
         if (!Core::File::exists(index_html_path)) {
-            handle_directory_listing(requested_path, real_path, request);
-            return;
+            TRY(handle_directory_listing(requested_path, real_path, request));
+            return true;
         }
         real_path = index_html_path;
     }
 
     auto file = Core::File::construct(real_path);
     if (!file->open(Core::OpenMode::ReadOnly)) {
-        send_error_response(404, request);
-        return;
+        TRY(send_error_response(404, request));
+        return false;
     }
 
     if (file->is_device()) {
-        send_error_response(403, request);
-        return;
+        TRY(send_error_response(403, request));
+        return false;
     }
 
     Core::InputFileStream stream { file };
 
-    send_response(stream, request, Core::guess_mime_type_based_on_filename(real_path));
+    TRY(send_response(stream, request, Core::guess_mime_type_based_on_filename(real_path)));
+    return true;
 }
 
-void Client::send_response(InputStream& response, HTTP::HttpRequest const& request, String const& content_type)
+ErrorOr<void> Client::send_response(InputStream& response, HTTP::HttpRequest const& request, String const& content_type)
 {
     StringBuilder builder;
     builder.append("HTTP/1.0 200 OK\r\n");
@@ -146,7 +181,8 @@ void Client::send_response(InputStream& response, HTTP::HttpRequest const& reque
     builder.append("\r\n");
     builder.append("\r\n");
 
-    m_socket->write(builder.to_string());
+    auto builder_contents = builder.to_byte_buffer();
+    TRY(m_socket.write(builder_contents));
     log_response(200, request);
 
     char buffer[PAGE_SIZE];
@@ -155,11 +191,22 @@ void Client::send_response(InputStream& response, HTTP::HttpRequest const& reque
         if (response.unreliable_eof() && size == 0)
             break;
 
-        m_socket->write({ buffer, size });
+        ReadonlyBytes write_buffer { buffer, size };
+        while (!write_buffer.is_empty()) {
+            auto nwritten = TRY(m_socket.write(write_buffer));
+
+            if (nwritten == 0) {
+                dbgln("EEEEEE got 0 bytes written!");
+            }
+
+            write_buffer = write_buffer.slice(nwritten);
+        }
     } while (true);
+
+    return {};
 }
 
-void Client::send_redirect(StringView redirect_path, HTTP::HttpRequest const& request)
+ErrorOr<void> Client::send_redirect(StringView redirect_path, HTTP::HttpRequest const& request)
 {
     StringBuilder builder;
     builder.append("HTTP/1.0 301 Moved Permanently\r\n");
@@ -168,9 +215,11 @@ void Client::send_redirect(StringView redirect_path, HTTP::HttpRequest const& re
     builder.append("\r\n");
     builder.append("\r\n");
 
-    m_socket->write(builder.to_string());
+    auto builder_contents = builder.to_byte_buffer();
+    TRY(m_socket.write(builder_contents));
 
     log_response(301, request);
+    return {};
 }
 
 static String folder_image_data()
@@ -195,7 +244,7 @@ static String file_image_data()
     return cache;
 }
 
-void Client::handle_directory_listing(String const& requested_path, String const& real_path, HTTP::HttpRequest const& request)
+ErrorOr<void> Client::handle_directory_listing(String const& requested_path, String const& real_path, HTTP::HttpRequest const& request)
 {
     StringBuilder builder;
 
@@ -270,10 +319,10 @@ void Client::handle_directory_listing(String const& requested_path, String const
 
     auto response = builder.to_string();
     InputMemoryStream stream { response.bytes() };
-    send_response(stream, request, "text/html");
+    return send_response(stream, request, "text/html");
 }
 
-void Client::send_error_response(unsigned code, HTTP::HttpRequest const& request, Vector<String> const& headers)
+ErrorOr<void> Client::send_error_response(unsigned code, HTTP::HttpRequest const& request, Vector<String> const& headers)
 {
     auto reason_phrase = HTTP::HttpResponse::reason_phrase_for_code(code);
     StringBuilder builder;
@@ -292,9 +341,12 @@ void Client::send_error_response(unsigned code, HTTP::HttpRequest const& request
     builder.appendff("{} ", code);
     builder.append(reason_phrase);
     builder.append("</h1></body></html>");
-    m_socket->write(builder.to_string());
+
+    auto builder_contents = builder.to_byte_buffer();
+    TRY(m_socket.write(builder_contents));
 
     log_response(code, request);
+    return {};
 }
 
 void Client::log_response(unsigned code, HTTP::HttpRequest const& request)
diff --git a/Userland/Services/WebServer/Client.h b/Userland/Services/WebServer/Client.h
index 5318776b46..a8a306f3ed 100644
--- a/Userland/Services/WebServer/Client.h
+++ b/Userland/Services/WebServer/Client.h
@@ -7,8 +7,9 @@
 #pragma once
 
 #include <LibCore/Object.h>
-#include <LibCore/TCPSocket.h>
+#include <LibCore/Stream.h>
 #include <LibHTTP/Forward.h>
+#include <LibHTTP/HttpRequest.h>
 
 namespace WebServer {
 
@@ -19,18 +20,18 @@ public:
     void start();
 
 private:
-    Client(NonnullRefPtr<Core::TCPSocket>, Core::Object* parent);
+    Client(Core::Stream::BufferedTCPSocket, Core::Object* parent);
 
-    void handle_request(ReadonlyBytes);
-    void send_response(InputStream&, HTTP::HttpRequest const&, String const& content_type);
-    void send_redirect(StringView redirect, HTTP::HttpRequest const&);
-    void send_error_response(unsigned code, HTTP::HttpRequest const&, Vector<String> const& headers = {});
+    ErrorOr<bool> handle_request(ReadonlyBytes);
+    ErrorOr<void> send_response(InputStream&, HTTP::HttpRequest const&, String const& content_type);
+    ErrorOr<void> send_redirect(StringView redirect, HTTP::HttpRequest const&);
+    ErrorOr<void> send_error_response(unsigned code, HTTP::HttpRequest const&, Vector<String> const& headers = {});
     void die();
     void log_response(unsigned code, HTTP::HttpRequest const&);
-    void handle_directory_listing(String const& requested_path, String const& real_path, HTTP::HttpRequest const&);
+    ErrorOr<void> handle_directory_listing(String const& requested_path, String const& real_path, HTTP::HttpRequest const&);
     bool verify_credentials(Vector<HTTP::HttpRequest::Header> const&);
 
-    NonnullRefPtr<Core::TCPSocket> m_socket;
+    Core::Stream::BufferedTCPSocket m_socket;
 };
 
 }
diff --git a/Userland/Services/WebServer/main.cpp b/Userland/Services/WebServer/main.cpp
index eb032eec3b..d78aef0fbb 100644
--- a/Userland/Services/WebServer/main.cpp
+++ b/Userland/Services/WebServer/main.cpp
@@ -72,9 +72,20 @@ ErrorOr<int> serenity_main(Main::Arguments arguments)
     auto server = TRY(Core::TCPServer::try_create());
 
     server->on_ready_to_accept = [&] {
-        auto client_socket = server->accept();
-        VERIFY(client_socket);
-        auto client = WebServer::Client::construct(client_socket.release_nonnull(), server);
+        auto maybe_client_socket = server->accept();
+        if (maybe_client_socket.is_error()) {
+            warnln("Failed to accept the client: {}", maybe_client_socket.error());
+            return;
+        }
+
+        auto maybe_buffered_socket = Core::Stream::BufferedTCPSocket::create(maybe_client_socket.release_value());
+        if (maybe_buffered_socket.is_error()) {
+            warnln("Could not obtain a buffered socket for the client: {}", maybe_buffered_socket.error());
+            return;
+        }
+
+        VERIFY(!maybe_buffered_socket.value().set_blocking(true).is_error());
+        auto client = WebServer::Client::construct(maybe_buffered_socket.release_value(), server);
         client->start();
     };