fix: harden socket I/O

5 files changed, 267 insertions, 40 deletions

diff --git a/src/beaker_globals.h b/src/beaker_globals.h
index a4a98f0..96d4dcf 100644
--- a/src/beaker_globals.h
+++ b/src/beaker_globals.h
@@ -5,6 +5,22 @@
 #include <stdatomic.h>
 #include <stddef.h>
 
+#ifndef BEAKER_REQUEST_HEADER_TIMEOUT_MS
+#define BEAKER_REQUEST_HEADER_TIMEOUT_MS 5000
+#endif
+
+#ifndef BEAKER_RESPONSE_WRITE_TIMEOUT_MS
+#define BEAKER_RESPONSE_WRITE_TIMEOUT_MS 30000
+#endif
+
+typedef enum {
+  BEAKER_REQUEST_READ_OK = 0,
+  BEAKER_REQUEST_READ_CLOSED,
+  BEAKER_REQUEST_READ_TIMEOUT,
+  BEAKER_REQUEST_READ_TOO_LARGE,
+  BEAKER_REQUEST_READ_ERROR,
+} BeakerRequestReadResult;
+
 extern RouteHandler handlers[MAX_HANDLERS];
 
 extern int handler_count;
@@ -33,6 +49,16 @@ void beaker_log_request(const char *remote_addr, const char *method,
                         const char *path, int status, size_t response_size,
                         double duration_ms);
 
+BeakerRequestReadResult beaker_read_request_headers(int socket, char *buffer,
+                                                    size_t buffer_size,
+                                                    size_t *bytes_read);
+
+int beaker_configure_client_socket(int socket);
+
+int beaker_send_all(int socket, const void *buffer, size_t length);
+
+void beaker_reset_write_deadline(void);
+
 extern Locale *locales;
 
 extern int locale_count;
diff --git a/src/http.c b/src/http.c
index 4c93a27..1312d4d 100644
--- a/src/http.c
+++ b/src/http.c
@@ -75,8 +75,8 @@ void send_status(const char *status_line) {
            "\r\n",
            status_line);
 
-  if (send(current_client_socket, http_response, strlen(http_response), 0) <
-      0) {
+  if (beaker_send_all(current_client_socket, http_response,
+                      strlen(http_response)) < 0) {
     beaker_log_errno_format("ERROR",
                             "send_status: Failed to send HTTP status.\n");
   }
@@ -91,7 +91,7 @@ void send_response(const char *html) {
   }
 
   char http_response_header[BUFFER_SIZE * 2];
-  int content_length = strlen(html);
+  size_t content_length = strlen(html);
   current_response_status = 200;
   current_response_size = content_length;
   char cookie_headers[BUFFER_SIZE];
@@ -101,20 +101,20 @@ void send_response(const char *html) {
   snprintf(http_response_header, sizeof(http_response_header),
            "HTTP/1.1 200 OK\r\n"
            "Content-Type: text/html; charset=UTF-8\r\n"
-           "Content-Length: %d\r\n"
+           "Content-Length: %zu\r\n"
            "%s"
            "Connection: close\r\n"
            "\r\n",
            content_length, cookie_headers);
 
-  if (send(current_client_socket, http_response_header,
-           strlen(http_response_header), 0) < 0) {
+  if (beaker_send_all(current_client_socket, http_response_header,
+                      strlen(http_response_header)) < 0) {
     beaker_log_errno_format("ERROR",
                             "send_response: Failed to send HTTP header.\n");
     return;
   }
 
-  if (send(current_client_socket, html, content_length, 0) < 0) {
+  if (beaker_send_all(current_client_socket, html, content_length) < 0) {
     beaker_log_errno_format("ERROR",
                             "send_response: Failed to send HTML body.\n");
     return;
@@ -146,8 +146,8 @@ void send_redirect(const char *location) {
            "\r\n",
            location, cookie_headers);
 
-  if (send(current_client_socket, http_response_header,
-           strlen(http_response_header), 0) < 0) {
+  if (beaker_send_all(current_client_socket, http_response_header,
+                      strlen(http_response_header)) < 0) {
     beaker_log_errno_format("ERROR",
                             "send_redirect: Failed to send redirect header.\n");
     return;
diff --git a/src/io.c b/src/io.c
new file mode 100644
index 0000000..81c8eac
--- /dev/null
+++ b/src/io.c
@@ -0,0 +1,185 @@
+#include "beaker_globals.h"
+#include <errno.h>
+#include <fcntl.h>
+#include <limits.h>
+#include <poll.h>
+#include <stdint.h>
+#include <string.h>
+#include <sys/socket.h>
+#include <time.h>
+
+static __thread int64_t write_deadline_ms = 0;
+
+int beaker_configure_client_socket(int socket) {
+  int flags = fcntl(socket, F_GETFL, 0);
+  if (flags < 0) {
+    return -1;
+  }
+
+  if ((flags & O_NONBLOCK) == 0 &&
+      fcntl(socket, F_SETFL, flags | O_NONBLOCK) < 0) {
+    return -1;
+  }
+
+#ifdef SO_NOSIGPIPE
+  int enabled = 1;
+  if (setsockopt(socket, SOL_SOCKET, SO_NOSIGPIPE, &enabled, sizeof(enabled)) <
+      0) {
+    return -1;
+  }
+#endif
+
+  return 0;
+}
+
+static int64_t monotonic_time_ms(void) {
+  struct timespec now;
+  if (clock_gettime(CLOCK_MONOTONIC, &now) != 0) {
+    return -1;
+  }
+  return (int64_t)now.tv_sec * 1000 + now.tv_nsec / 1000000;
+}
+
+static int wait_for_socket(int socket, short events, int64_t deadline_ms) {
+  for (;;) {
+    int64_t now_ms = monotonic_time_ms();
+    if (now_ms < 0) {
+      return -1;
+    }
+
+    int64_t remaining_ms = deadline_ms - now_ms;
+    if (remaining_ms <= 0) {
+      errno = ETIMEDOUT;
+      return -1;
+    }
+
+    int timeout_ms = remaining_ms > INT_MAX ? INT_MAX : (int)remaining_ms;
+    struct pollfd descriptor = {.fd = socket, .events = events};
+    int result = poll(&descriptor, 1, timeout_ms);
+
+    if (result > 0) {
+      if (descriptor.revents & POLLNVAL) {
+        errno = EBADF;
+        return -1;
+      }
+      if (descriptor.revents & (events | POLLERR | POLLHUP)) {
+        return 0;
+      }
+      continue;
+    }
+    if (result == 0) {
+      errno = ETIMEDOUT;
+      return -1;
+    }
+    if (errno != EINTR) {
+      return -1;
+    }
+  }
+}
+
+static bool contains_complete_headers(const char *buffer, size_t length) {
+  if (length < 4) {
+    return false;
+  }
+
+  for (size_t i = 0; i <= length - 4; i++) {
+    if (buffer[i] == '\r' && buffer[i + 1] == '\n' && buffer[i + 2] == '\r' &&
+        buffer[i + 3] == '\n') {
+      return true;
+    }
+  }
+  return false;
+}
+
+BeakerRequestReadResult beaker_read_request_headers(int socket, char *buffer,
+                                                    size_t buffer_size,
+                                                    size_t *bytes_read) {
+  if (buffer == NULL || bytes_read == NULL || buffer_size < 2) {
+    errno = EINVAL;
+    return BEAKER_REQUEST_READ_ERROR;
+  }
+
+  *bytes_read = 0;
+  buffer[0] = '\0';
+
+  int64_t now_ms = monotonic_time_ms();
+  if (now_ms < 0) {
+    return BEAKER_REQUEST_READ_ERROR;
+  }
+  int64_t deadline_ms = now_ms + BEAKER_REQUEST_HEADER_TIMEOUT_MS;
+
+  while (*bytes_read < buffer_size - 1) {
+    ssize_t result =
+        recv(socket, buffer + *bytes_read, buffer_size - 1 - *bytes_read, 0);
+    if (result > 0) {
+      *bytes_read += (size_t)result;
+      buffer[*bytes_read] = '\0';
+      if (contains_complete_headers(buffer, *bytes_read)) {
+        return BEAKER_REQUEST_READ_OK;
+      }
+      continue;
+    }
+    if (result == 0) {
+      return BEAKER_REQUEST_READ_CLOSED;
+    }
+    if (errno == EINTR) {
+      continue;
+    }
+    if (errno != EAGAIN && errno != EWOULDBLOCK) {
+      return BEAKER_REQUEST_READ_ERROR;
+    }
+    if (wait_for_socket(socket, POLLIN, deadline_ms) != 0) {
+      return errno == ETIMEDOUT ? BEAKER_REQUEST_READ_TIMEOUT
+                                : BEAKER_REQUEST_READ_ERROR;
+    }
+  }
+
+  return BEAKER_REQUEST_READ_TOO_LARGE;
+}
+
+void beaker_reset_write_deadline(void) { write_deadline_ms = 0; }
+
+int beaker_send_all(int socket, const void *buffer, size_t length) {
+  if (buffer == NULL && length != 0) {
+    errno = EINVAL;
+    return -1;
+  }
+
+  if (write_deadline_ms == 0) {
+    int64_t now_ms = monotonic_time_ms();
+    if (now_ms < 0) {
+      return -1;
+    }
+    write_deadline_ms = now_ms + BEAKER_RESPONSE_WRITE_TIMEOUT_MS;
+  }
+
+  const char *bytes = buffer;
+  size_t bytes_sent = 0;
+  while (bytes_sent < length) {
+#ifdef MSG_NOSIGNAL
+    ssize_t result =
+        send(socket, bytes + bytes_sent, length - bytes_sent, MSG_NOSIGNAL);
+#else
+    ssize_t result = send(socket, bytes + bytes_sent, length - bytes_sent, 0);
+#endif
+    if (result > 0) {
+      bytes_sent += (size_t)result;
+      continue;
+    }
+    if (result == 0) {
+      errno = EPIPE;
+      return -1;
+    }
+    if (errno == EINTR) {
+      continue;
+    }
+    if (errno != EAGAIN && errno != EWOULDBLOCK) {
+      return -1;
+    }
+    if (wait_for_socket(socket, POLLOUT, write_deadline_ms) != 0) {
+      return -1;
+    }
+  }
+
+  return 0;
+}
diff --git a/src/routing.c b/src/routing.c
index 9894768..d5f8285 100644
--- a/src/routing.c
+++ b/src/routing.c
@@ -315,7 +315,8 @@ bool serve_static_file_with_mime(const char *request_path_relative_to_static,
            "\r\n",
            mime_type, file_size);
 
-  if (send(current_client_socket, http_header, strlen(http_header), 0) < 0) {
+  if (beaker_send_all(current_client_socket, http_header, strlen(http_header)) <
+      0) {
     beaker_log_errno_format(
         "ERROR",
         "serve_static_file_with_mime: Failed to send header for '%s'.\n",
@@ -330,7 +331,7 @@ bool serve_static_file_with_mime(const char *request_path_relative_to_static,
 
   while (!send_error && !feof(fp) && !ferror(fp) &&
          (bytes_read = fread(file_buffer, 1, sizeof(file_buffer), fp)) > 0) {
-    if (send(current_client_socket, file_buffer, bytes_read, 0) < 0) {
+    if (beaker_send_all(current_client_socket, file_buffer, bytes_read) < 0) {
       beaker_log_errno_format(
           "ERROR",
           "serve_static_file_with_mime: Failed to send content for '%s'.\n",
@@ -380,20 +381,15 @@ bool serve_data(const char *data, size_t size, const char *mime_type) {
            "\r\n",
            mime_type, size);
 
-  if (send(current_client_socket, http_header, strlen(http_header), 0) < 0) {
+  if (beaker_send_all(current_client_socket, http_header, strlen(http_header)) <
+      0) {
     beaker_log_errno_format("ERROR", "serve_data: Failed to send header.\n");
     return false;
   }
 
-  size_t bytes_sent = 0;
-  while (bytes_sent < size) {
-    size_t chunk =
-        (size - bytes_sent > BUFFER_SIZE) ? BUFFER_SIZE : (size - bytes_sent);
-    if (send(current_client_socket, data + bytes_sent, chunk, 0) < 0) {
-      beaker_log_errno_format("ERROR", "serve_data: Failed to send content.\n");
-      return false;
-    }
-    bytes_sent += chunk;
+  if (beaker_send_all(current_client_socket, data, size) < 0) {
+    beaker_log_errno_format("ERROR", "serve_data: Failed to send content.\n");
+    return false;
   }
 
   return true;
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);
       }
     }