fix: harden socket I/O

1 files changed, 39 insertions, 19 deletions

diff --git a/src/server.c b/src/server.c
index b95b8c1..5ac20a0 100644
--- a/src/server.c
+++ b/src/server.c
@@ -190,17 +190,6 @@ static int initialize_server_socket(const char *ip, int port,
   return 0;
 }
 
-static int set_socket_blocking(int fd) {
-  int flags = fcntl(fd, F_GETFL, 0);
-  if (flags < 0)
-    return -1;
-
-  if ((flags & O_NONBLOCK) == 0)
-    return 0;
-
-  return fcntl(fd, F_SETFL, flags & ~O_NONBLOCK);
-}
-
 static void finish_client_connection(int socket, char *requested_path,
                                      const char *method, const char *path,
                                      const struct timespec *started_at) {
@@ -227,6 +216,7 @@ void handle_client_connection(int new_socket) {
   clock_gettime(CLOCK_MONOTONIC, &started_at);
   current_response_status = 0;
   current_response_size = 0;
+  beaker_reset_write_deadline();
   memset(&current_request_info, 0, sizeof(RequestInfo));
 
   struct sockaddr_in client_addr;
@@ -239,8 +229,33 @@ void handle_client_connection(int new_socket) {
     strcpy(current_request_info.remote_addr, "-");
   }
 
-  ssize_t bytes_read = read(new_socket, buffer, BUFFER_SIZE - 1);
-  if (bytes_read < 0) {
+  size_t bytes_read = 0;
+  BeakerRequestReadResult read_result = beaker_read_request_headers(
+      new_socket, buffer, sizeof(buffer), &bytes_read);
+  if (read_result == BEAKER_REQUEST_READ_TIMEOUT) {
+    beaker_log("WARN", "handle_client_connection: Request header timed out.\n");
+    send_status("408 Request Timeout");
+    finish_client_connection(new_socket, requested_path, method, log_path,
+                             &started_at);
+    return;
+  }
+  if (read_result == BEAKER_REQUEST_READ_TOO_LARGE) {
+    beaker_log("WARN",
+               "handle_client_connection: Request headers too large.\n");
+    send_status("431 Request Header Fields Too Large");
+    finish_client_connection(new_socket, requested_path, method, log_path,
+                             &started_at);
+    return;
+  }
+  if (read_result == BEAKER_REQUEST_READ_CLOSED) {
+    if (bytes_read > 0) {
+      send_status("400 Bad Request");
+    }
+    finish_client_connection(new_socket, requested_path, method, log_path,
+                             &started_at);
+    return;
+  }
+  if (read_result == BEAKER_REQUEST_READ_ERROR) {
     beaker_log_errno_format(
         "ERROR",
         "handle_client_connection: Failed to read from client socket.\n");
@@ -248,8 +263,6 @@ void handle_client_connection(int new_socket) {
                              &started_at);
     return;
   }
-  buffer[bytes_read] = '\0';
-
   strncpy(current_request_buffer, buffer, BUFFER_SIZE - 1);
   current_request_buffer[BUFFER_SIZE - 1] = '\0';
 
@@ -337,7 +350,10 @@ void handle_client_connection(int new_socket) {
              "Connection: close\r\n"
              "\r\n%s",
              strlen(not_found_html), not_found_html);
-    send(new_socket, not_found_response, strlen(not_found_response), 0);
+    if (beaker_send_all(new_socket, not_found_response,
+                        strlen(not_found_response)) < 0) {
+      beaker_log_errno("Failed to send 404 response");
+    }
     current_response_status = 404;
     current_response_size = strlen(not_found_html);
   }
@@ -398,8 +414,8 @@ void beaker_run_with_threads(const char *ip, int port, int num_workers) {
     int new_socket;
     while ((new_socket = accept(server_fd, (struct sockaddr *)&address,
                                 (socklen_t *)&addrlen)) >= 0) {
-      if (set_socket_blocking(new_socket) < 0) {
-        beaker_log_errno("fcntl clear O_NONBLOCK failed");
+      if (beaker_configure_client_socket(new_socket) < 0) {
+        beaker_log_errno("failed to configure client socket");
         close(new_socket);
         continue;
       }
@@ -407,7 +423,11 @@ void beaker_run_with_threads(const char *ip, int port, int num_workers) {
         beaker_log("WARN", "work queue full; rejecting connection");
         const char *busy_response =
             "HTTP/1.1 503 Service Unavailable\r\nContent-Length: 0\r\n\r\n";
-        send(new_socket, busy_response, strlen(busy_response), 0);
+        beaker_reset_write_deadline();
+        if (beaker_send_all(new_socket, busy_response, strlen(busy_response)) <
+            0) {
+          beaker_log_errno("failed to send busy response");
+        }
         close(new_socket);
       }
     }