compressino + fixed logic

3 files changed, 148 insertions, 11 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);
diff --git a/src/packet_storage.c b/src/packet_storage.c
index bb8cdcd..5781661 100644
--- a/src/packet_storage.c
+++ b/src/packet_storage.c
@@ -324,10 +324,10 @@ struct data_packet *load_packet(FILE *file) {
 uint64_t set_data_packet(struct data_packet *new_packet) {
     byte temp_buffer[ sizeof(struct data_packet) ];
     memcpy(temp_buffer, new_packet, sizeof(struct data_packet));
+    uLongf compressed_len = compressBound(sizeof(struct data_packet));
     byte   compressed[ sizeof(struct data_packet) ];
-    uLongf compressed_len = sizeof(struct data_packet);
     if (compress(compressed, &compressed_len, temp_buffer, sizeof(struct data_packet)) != Z_OK) {
-        printf("set_data_packet: compression failed\n");
+        fprintf(stderr, "set_data_packet: compression failed\n");
         return invalid_index;
     }
 
@@ -343,9 +343,8 @@ uint64_t set_data_packet(struct data_packet *new_packet) {
     struct data_packet_entry *current_entry = &packet_index[ cur_index ];
 
     while (true) {
-        if (current_entry->is_full && 
-            current_entry->compressed_size == compressed_len &&
-            memcmp(current_entry->compressed_payload, compressed_payload, compressed_len) == 0) {
+        if (current_entry->is_full && current_entry->compressed_size == compressed_len
+            && memcmp(current_entry->compressed_payload, compressed_payload, compressed_len) == 0) {
             arena_free(current_entry->compressed_payload);
             current_entry->compressed_payload = compressed_payload;
             current_entry->compressed_size    = compressed_len;
diff --git a/src/server.c b/src/server.c
index 3140788..630df87 100644
--- a/src/server.c
+++ b/src/server.c
@@ -707,6 +707,114 @@ void paste_server(
             break;
         }
 
+        case query_by_id_hash: {
+            if (in_len < 1 + 4 + 32) {
+                byte* output_data = calloc(1 + 17, sizeof(byte));
+                if (!output_data) {
+                    *out     = NULL;
+                    *out_len = 0;
+                    break;
+                }
+                output_data[ 0 ] = error;
+                memcpy(output_data + 1, "input too small", 17);
+                reply(18, output_data);
+                break;
+            }
+
+            uint64_t __fetch_index = read_u64_le(in + 1);
+            hash     packet_hash;
+            bool     success;
+            s2hash(32, in + 1 + 4, &success, packet_hash);
+
+            log(_info,
+                "<%.*s> request type: get packet with hash " hash_fmt " id %li",
+                (int) ip_size,
+                ip_addr,
+                hash_ref(packet_hash),
+                __fetch_index);
+
+            if (!success) {
+                byte* output_data = calloc(1 + 12, sizeof(byte));
+                if (!output_data) {
+                    *out     = NULL;
+                    *out_len = 0;
+                    break;
+                }
+                output_data[ 0 ] = error;
+                memcpy(output_data + 1, "invalid hash", 12);
+                reply(13, output_data);
+                break;
+            }
+
+            struct data_packet_entry* entry
+                = __fetch_index >= max_storage ? NULL : &packet_index[ __fetch_index ];
+            if (entry == NULL || !entry->is_full) {
+                byte* output_data = calloc(1 + 14, sizeof(byte));
+                if (!output_data) {
+                    *out     = NULL;
+                    *out_len = 0;
+                    return;
+                }
+                output_data[ 0 ] = error;
+                memcpy(output_data + 1, "packet not found", 14);
+                reply(15, output_data);
+                break;
+            }
+
+            uLongf              uncompressed_size = sizeof(struct data_packet);
+            struct data_packet* uncompressed_data = malloc(uncompressed_size);
+
+            int ret = uncompress(
+                (byte*) uncompressed_data,
+                &uncompressed_size,
+                entry->compressed_payload,
+                entry->compressed_size);
+
+            if (ret != Z_OK) {
+                log(_error, "decompression failed with error code: %d", ret);
+                byte* output_data = calloc(1 + 20, sizeof(byte));
+                if (output_data) {
+                    output_data[ 0 ] = error;
+                    memcpy(output_data + 1, "decompression failed", 20);
+                    reply(21, output_data);
+                }
+                free(uncompressed_data);
+                break;
+            }
+
+            if (memcmp(uncompressed_data->hash, packet_hash, sizeof(hash)) != 0) {
+                log(_error, "packet hash does not match");
+                byte* output_data = calloc(1 + 17, sizeof(byte));
+                if (output_data) {
+                    output_data[ 0 ] = error;
+                    memcpy(output_data + 1, "non matching hash", 17);
+                    reply(18, output_data);
+                }
+                free(uncompressed_data);
+                break;
+            }
+
+            size_t output_len  = 1 + uncompressed_size;
+            byte*  output_data = calloc(output_len, sizeof(byte));
+            if (!output_data) {
+                free(uncompressed_data);
+                byte* err = calloc(1 + 14, sizeof(byte));
+                if (err) {
+                    err[ 0 ] = error;
+                    memcpy(err + 1, "server OOM", 10);
+                    reply(11, err);
+                }
+                break;
+            }
+
+            output_data[ 0 ] = get_packet_result;
+            memcpy(output_data + 1, uncompressed_data, uncompressed_size);
+            free(uncompressed_data);
+
+            reply(output_len, output_data);
+            break;
+        }
+
         default: {
             log(_info, "<%.*s> request type: invalid method", (int) ip_size, ip_addr);