From 6bc2bf4e1bdb0d3a61c936c99e1ba671c54222b8 Mon Sep 17 00:00:00 2001 From: moon Date: Tue, 16 Sep 2025 11:22:51 +0330 Subject: compressino + fixed logic --- src/omen.c | 42 +++++++++++++++++--- src/packet_storage.c | 9 ++--- src/server.c | 108 +++++++++++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 148 insertions(+), 11 deletions(-) (limited to 'src') 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 #include #include +// dependency: zlib +#include // #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); -- cgit v1.2.3