compressino + fixed logic

1 files changed, 36 insertions, 6 deletions

diff --git a/src/omen.c b/src/omen.c
index 222e702..3770d71 100644
--- a/src/omen.c
+++ b/src/omen.c
@@ -12,6 +12,8 @@
 #include <strings.h>
 #include <sys/socket.h>
 #include <unistd.h>
+// dependency: zlib
+#include <zlib.h>
 //
 #include "common.c"
 #include "ratelimit.c"
@@ -56,7 +58,7 @@ static ssize_t write_all(int fd, const void *buf, size_t len) {
     return (ssize_t) off;
 }
 
-static int read_until_headers(int fd, char **out_buf, size_t *out_len) {
+static int read_until_headers(int fd, char **out_buf, size_t *out_len, bool *enable_compression) {
     size_t total_tries = 0;
 
     size_t cap = 4096;
@@ -88,7 +90,7 @@ static int read_until_headers(int fd, char **out_buf, size_t *out_len) {
         }
         len += (size_t) r;
         if (len >= 4 && memmem(buf, len, "\r\n\r\n", 4) != NULL) {
-            if (len >= 19 && memmem(buf, len, "Transfer-Encoding: chunked", 19) != NULL) {
+            if (len >= 26 && memmem(buf, len, "Transfer-Encoding: chunked", 26) != NULL) {
                 const char *err = "HTTP/1.1 501 Not Implemented\r\n"
                                   "Content-Length: 31\r\n"
                                   "Connection: close\r\n\r\n"
@@ -97,6 +99,9 @@ static int read_until_headers(int fd, char **out_buf, size_t *out_len) {
                 free(buf);
                 return -1;
             }
+            if (len >= 17 && memmem(buf, len, "Compression: true", 17) != NULL) {
+                if (enable_compression) *enable_compression = true;
+            }
             break;
         }
         if (len > MAX_BUFFER) break;
@@ -186,9 +191,10 @@ void *omen_bound(void *arg) {
         struct timeval timeout = { 5, 0 };
         setsockopt(client_fd, SOL_SOCKET, SO_RCVTIMEO, &timeout, sizeof(timeout));
 
-        char  *hdr_buf = NULL;
-        size_t hdr_len = 0;
-        if (read_until_headers(client_fd, &hdr_buf, &hdr_len) < 0) {
+        char  *hdr_buf            = NULL;
+        size_t hdr_len            = 0;
+        bool   enable_compression = false;
+        if (read_until_headers(client_fd, &hdr_buf, &hdr_len, &enable_compression) < 0) {
             close(client_fd);
             free(hdr_buf);
             continue;
@@ -256,6 +262,30 @@ void *omen_bound(void *arg) {
             global_handler(client_ip, sizeof(client_ip), body_ptr, body_len, &out, &out_len);
         if (out_len > MAX_BODY) out_len = MAX_BODY;
 
+        // if compression is enabled, zlib compress the result before sending it
+        if (enable_compression) {
+            uLongf compressed_len    = compressBound(out_len);
+            byte  *compressed_result = malloc(compressed_len);
+            if (compress(compressed_result, &compressed_len, out, out_len) != Z_OK) {
+                fprintf(stderr, "omen_bound: compression failed\n");
+                free(compressed_result);
+                free(out);
+                out      = malloc(1);
+                out[ 0 ] = false;
+                out_len  = 1;
+                goto CONTINUE_HERE;
+            }
+
+            free(out);
+            out      = malloc(compressed_len + 1);
+            out[ 0 ] = true;
+            out_len = compressed_len + 1;
+            memcpy(out + 1, compressed_result, compressed_len);
+            free(compressed_result);
+        }
+
+    CONTINUE_HERE:
+
         send_http_response(client_fd, out, out_len);
 
         free(out);
@@ -468,7 +498,7 @@ int study(
 
     char  *hdr_buf = NULL;
     size_t hdr_len = 0;
-    if (read_until_headers(sock, &hdr_buf, &hdr_len) < 0) {
+    if (read_until_headers(sock, &hdr_buf, &hdr_len, NULL) < 0) {
         printf("failed to read headers: %s\n", strerror(errno));
         close(sock);
         free(hdr_buf);