From dd0ba315067c69f95e54da06725e4b40ec4d9dcb Mon Sep 17 00:00:00 2001
From: Chewing_Bever <roosensjef@gmail.com>
Date: Mon, 29 May 2023 15:34:46 +0200
Subject: [PATCH] feat: support adding up to 4 headers

---
 include/http.h                | 14 +++++++++++
 include/http_loop.h           |  3 +++
 src/http/http_res_names.c     |  5 ++++
 src/http_loop/http_loop_ctx.c |  8 +++++++
 src/http_loop/http_loop_res.c | 44 ++++++++++++++++++++++++++++++-----
 src/main.c                    |  1 +
 6 files changed, 69 insertions(+), 6 deletions(-)

diff --git a/include/http.h b/include/http.h
index dda8164..a5a82a5 100644
--- a/include/http.h
+++ b/include/http.h
@@ -9,6 +9,7 @@
 #define HTTP_MAX_ALLOWED_HEADERS 16
 
 extern const char *http_response_type_names[5][32];
+extern const char *http_header_names[];
 
 typedef enum http_request_method {
   http_get = 0,
@@ -111,6 +112,17 @@ typedef enum http_response_type {
   http_network_authentication_required = 511
 } http_response_type;
 
+typedef enum http_header {
+  http_header_connection = 0,
+  http_header_redirect = 1
+} http_header;
+
+typedef struct http_response_header {
+  http_header type;
+  const char *value;
+  bool owned;
+} http_response_header;
+
 typedef struct http_response {
   http_response_type type;
   const char *head;
@@ -121,6 +133,8 @@ typedef struct http_response {
   size_t body_written;
   // If false, the body won't be freed
   bool owns_body;
+  http_response_header headers[4];
+  size_t header_count;
 } http_response;
 
 #endif
diff --git a/include/http_loop.h b/include/http_loop.h
index 81ac4b5..f3344bb 100644
--- a/include/http_loop.h
+++ b/include/http_loop.h
@@ -97,6 +97,9 @@ void http_loop_process_request(event_loop_conn *conn);
 void http_loop_res_set_body(http_loop_ctx *ctx, const char *body,
                             size_t body_len, bool owned);
 
+void http_loop_res_add_header(http_loop_ctx *ctx, http_header type,
+                              const char *value, bool owned);
+
 /**
  * Initialize a new http loop
  */
diff --git a/src/http/http_res_names.c b/src/http/http_res_names.c
index 5160d8d..d565853 100644
--- a/src/http/http_res_names.c
+++ b/src/http/http_res_names.c
@@ -86,4 +86,9 @@ const char *http_response_type_names[5][32] = {
   },
 };
 
+const char *http_header_names[] = {
+  "Connection",
+  "Redirect"
+};
+
 // clang-format on
diff --git a/src/http_loop/http_loop_ctx.c b/src/http_loop/http_loop_ctx.c
index 432281f..39fdc16 100644
--- a/src/http_loop/http_loop_ctx.c
+++ b/src/http_loop/http_loop_ctx.c
@@ -33,6 +33,14 @@ void http_loop_ctx_reset(http_loop_ctx *ctx) {
     free((void *)ctx->res.body);
   }
 
+  for (size_t i = 0; i < ctx->res.header_count; i++) {
+    if (ctx->res.headers[i].owned) {
+      free((void *)ctx->res.headers[i].value);
+    }
+  }
+
+  ctx->res.header_count = 0;
+
   ctx->res.body = NULL;
 
   ctx->res.type = 0;
diff --git a/src/http_loop/http_loop_res.c b/src/http_loop/http_loop_res.c
index 581829d..104272a 100644
--- a/src/http_loop/http_loop_res.c
+++ b/src/http_loop/http_loop_res.c
@@ -7,8 +7,14 @@
 
 static const char *http_response_format = "HTTP/1.1 %i %s\n"
                                           "Server: lander/" LANDER_VERSION "\n"
-                                          "Content-Length: %lu\n\n";
+                                          "Content-Length: %lu\n";
 
+/*
+ * This function precalculates the size of the total buffer required using
+ * snprintf. When this function is called with a buf size of 0, it never tries
+ * to write any data, but it does return the amount of bytes that would be
+ * written.
+ */
 void http_loop_init_header(http_response *res) {
   if (res->type == 0) {
     res->type = http_ok;
@@ -17,16 +23,33 @@ void http_loop_init_header(http_response *res) {
   const char *response_type_name =
       http_response_type_names[res->type / 100 - 1][res->type % 100];
 
-  // Running snprintf with size 0 prevents it from writing any bytes, while
-  // still letting it calculate how many bytes it would have written
+  // First we calculate the size of the start of the header
   int buf_size = snprintf(NULL, 0, http_response_format, res->type,
                           response_type_name, res->body_len);
+
+  // We add each header's required size
+  for (size_t i = 0; i < res->header_count; i++) {
+    buf_size +=
+        snprintf(NULL, 0, "%s: %s\n", http_header_names[res->headers[i].type],
+                 res->headers[i].value);
+  }
+
+  // The + 1 is required to store the final null byte, but we will replace it
+  // with the required final newline
   char *buf = malloc(buf_size + 1);
-  sprintf(buf, http_response_format, res->type, response_type_name,
-          res->body_len);
+  buf_size = sprintf(buf, http_response_format, res->type, response_type_name,
+                     res->body_len);
+
+  for (size_t i = 0; i < res->header_count; i++) {
+    buf_size +=
+        sprintf(&buf[buf_size], "%s: %s\n",
+                http_header_names[res->headers[i].type], res->headers[i].value);
+  }
+
+  buf[buf_size] = '\n';
 
   res->head = buf;
-  res->head_len = buf_size;
+  res->head_len = buf_size + 1;
 }
 
 void http_loop_write_response(event_loop_conn *conn) {
@@ -73,3 +96,12 @@ void http_loop_res_set_body(http_loop_ctx *ctx, const char *body,
   ctx->res.body_len = body_len;
   ctx->res.owns_body = owned;
 }
+
+void http_loop_res_add_header(http_loop_ctx *ctx, http_header type,
+                              const char *value, bool owned) {
+  ctx->res.headers[ctx->res.header_count].type = type;
+  ctx->res.headers[ctx->res.header_count].value = value;
+  ctx->res.headers[ctx->res.header_count].owned = owned;
+
+  ctx->res.header_count++;
+}
diff --git a/src/main.c b/src/main.c
index 20e8c85..c490d08 100644
--- a/src/main.c
+++ b/src/main.c
@@ -18,6 +18,7 @@ bool lander_get_index(event_loop_conn *conn) {
   http_loop_ctx *ctx = conn->ctx;
 
   http_loop_res_set_body(ctx, index_page, sizeof(index_page) - 1, false);
+  http_loop_res_add_header(ctx, http_header_connection, "close", false);
 
   conn->state = event_loop_conn_state_res;