diff --git a/src/main.cpp b/src/main.cpp index 035b32a..3460288 100644 --- a/src/main.cpp +++ b/src/main.cpp @@ -34,13 +34,14 @@ static const std::string index_page = R"( crow::response add_redirect(std::string base_url, Trie *trie, const char *url, bool secure) { - Entry *new_entry = entry_new(Redirect, url); + Entry *new_entry; + uint64_t entry_size = entry_new(&new_entry, Redirect, url); // The key already gets copied into the trie, so this pointer is safe to use // ever after unlocking the trie trie_wlock(trie); char *key; - TrieExitCode res = trie_add_random(trie, &key, new_entry, secure); + TrieExitCode res = trie_add_random(trie, &key, new_entry, entry_size, secure); trie_unlock(trie); if (res != Ok) { @@ -70,11 +71,12 @@ bool store_paste(const char *key, const char *body) { crow::response add_paste(std::string base_url, Trie *trie, const char *body, bool secure) { - Entry *new_entry = entry_new(Paste, ""); + Entry *new_entry; + uint64_t entry_size = entry_new(&new_entry, Paste, ""); trie_wlock(trie); char *key; - TrieExitCode res = trie_add_random(trie, &key, new_entry, secure); + TrieExitCode res = trie_add_random(trie, &key, new_entry, entry_size, secure); trie_unlock(trie); if (res != Ok) { @@ -169,10 +171,13 @@ int main() { [api_key, base_url, trie](const crow::request &req, std::string key) { AUTH(); - Entry *new_entry = entry_new(Redirect, req.body.c_str()); + Entry *new_entry; + uint64_t entry_size = + entry_new(&new_entry, Redirect, req.body.c_str()); trie_wlock(trie); - TrieExitCode status = trie_add(trie, key.c_str(), new_entry); + TrieExitCode status = + trie_add(trie, key.c_str(), new_entry, entry_size); trie_unlock(trie); switch (status) { @@ -209,9 +214,11 @@ int main() { [api_key, base_url, trie](const crow::request &req, std::string key) { AUTH(); - Entry *new_entry = entry_new(Paste, ""); + Entry *new_entry; + uint64_t entry_size = entry_new(&new_entry, Paste, ""); trie_wlock(trie); - TrieExitCode status = trie_add(trie, key.c_str(), new_entry); + TrieExitCode status = + trie_add(trie, key.c_str(), new_entry, entry_size); trie_unlock(trie); switch (status) { diff --git a/trie/include/trie.h b/trie/include/trie.h index 613aa68..2268418 100644 --- a/trie/include/trie.h +++ b/trie/include/trie.h @@ -42,7 +42,7 @@ typedef enum entry_type { Redirect, Paste, Unknown } EntryType; typedef struct entry { EntryType type; - char *string; + char string[]; } Entry; typedef enum trie_exit_code { @@ -52,7 +52,7 @@ typedef enum trie_exit_code { FileError } TrieExitCode; -Entry *entry_new(EntryType type, const char *string); +uint64_t entry_new(Entry **entry_ptr, EntryType type, const char *string); void entry_free(Entry *entry); @@ -89,7 +89,8 @@ TrieExitCode trie_search(Trie *trie, void **data_ptr, const char *key); * @param entry entry to add * @return 0 if added, 1 if already in trie, something else if other errors */ -TrieExitCode trie_add(Trie *trie, const char *key, void *data); +TrieExitCode trie_add(Trie *trie, const char *key, void *data, + uint64_t data_len); /** * Add an entry by generating a random string as the key. @@ -101,7 +102,7 @@ TrieExitCode trie_add(Trie *trie, const char *key, void *data); * unlocking the trie, and should be freed manually. */ TrieExitCode trie_add_random(Trie *trie, char **key_ptr, void *data, - bool secure); + uint64_t data_len, bool secure); /** * Remove an entry from this trie given its key. diff --git a/trie/src/trie.c b/trie/src/trie.c index a0e94a1..c7b0180 100644 --- a/trie/src/trie.c +++ b/trie/src/trie.c @@ -43,42 +43,45 @@ TrieExitCode trie_init(Trie **trie_ptr, const char *file_path) { return FileError; } - // We read in lines of at most 8192 characters (sounds like enough) - char buffer[8192]; - EntryType type; - Entry *entry; - char *string; - int i, j; + uint64_t key_size, data_size; + char *key; + void *data; TrieExitCode status; + size_t items_read; - while (fgets(buffer, 8192, fp)) { - i = 0; - - // Move index in buffer until we encounter first space character - while (buffer[i] != ' ') { - i++; + while (!feof(fp)) { + items_read = fread(&key_size, sizeof(uint64_t), 1, fp); + if (items_read < 1) { + break; } - // Split the buffer into two strings, the key and the payload - buffer[i] = '\0'; + key = malloc(key_size + 1); - type = entry_type_from_char(buffer[i + 1]); - - // Skip type character & its surrounding spaces - j = i + 3; - - // Now remove the newline character - while (buffer[j] != '\n') { - j++; + items_read = fread(key, 1, key_size, fp); + if (items_read < key_size) { + break; } - buffer[j] = '\0'; + key[key_size] = '\0'; - entry = entry_new(type, buffer + i + 3); - status = trie_add_no_lock(trie, buffer, entry); + items_read = fread(&data_size, sizeof(uint64_t), 1, fp); + if (items_read < 1) { + break; + } + + data = malloc(data_size); + items_read = fread(data, 1, data_size, fp); + if (items_read < data_size) { + break; + } + + status = trie_add_no_lock(trie, key, data); if (status != Ok) { trie_free(trie); + free(key); + free(data); + return status; } } @@ -275,7 +278,8 @@ TrieExitCode trie_add_no_lock(Trie *trie, const char *string, void *data) { return Ok; } -TrieExitCode trie_add(Trie *trie, const char *key, void *entry) { +TrieExitCode trie_add(Trie *trie, const char *key, void *data, + uint64_t data_len) { if (trie->file_path != NULL) { // Easiest way to make sure we don't add duplicate entries // We use an internal function that doesn't require a read lock, as we're @@ -290,23 +294,24 @@ TrieExitCode trie_add(Trie *trie, const char *key, void *entry) { return FileError; } - /* fputs(key, fp); */ - /* fputs(" ", fp); */ - /* fputc(entry_type_to_char(entry->type), fp); */ - /* fputs(" ", fp); */ - /* fputs(entry->string, fp); */ - /* fputs("\n", fp); */ + // First we write the key, then the actual data + uint64_t key_len = (uint64_t)strlen(key); + fwrite(&key_len, 1, sizeof(uint64_t), fp); + fwrite(key, 1, key_len, fp); + + fwrite(&data_len, 1, sizeof(uint64_t), fp); + fwrite(data, 1, data_len, fp); fclose(fp); } // This function *should* always return Ok. Otherwise, the function would've // exited because the string was found in the trie. - return trie_add_no_lock(trie, key, entry); + return trie_add_no_lock(trie, key, data); } TrieExitCode trie_add_random(Trie *trie, char **key_ptr, void *data, - bool secure) { + uint64_t data_len, bool secure) { // Generate random key bool ok = false; int key_length = secure ? RANDOM_KEY_LENGTH_LONG : RANDOM_KEY_LENGTH_SHORT; @@ -324,7 +329,7 @@ TrieExitCode trie_add_random(Trie *trie, char **key_ptr, void *data, ok = trie_search_node(trie, key).child == NULL; } - TrieExitCode return_value = trie_add(trie, key, data); + TrieExitCode return_value = trie_add(trie, key, data, data_len); if (return_value == Ok) { *key_ptr = key; diff --git a/trie/src/trie_entry.c b/trie/src/trie_entry.c index 1038ed6..85097c8 100644 --- a/trie/src/trie_entry.c +++ b/trie/src/trie_entry.c @@ -23,23 +23,20 @@ char entry_type_to_char(EntryType et) { } } -Entry *entry_new(EntryType type, const char *string) { - Entry *entry = malloc(sizeof(Entry)); +uint64_t entry_new(Entry **entry_ptr, EntryType type, const char *string) { + size_t str_len = strlen(string); + uint64_t entry_size = sizeof(EntryType) + str_len + 1; + Entry *entry = malloc(entry_size); entry->type = type; if (string != NULL) { - entry->string = strdup(string); + memcpy(entry->string, string, str_len + 1); } else { - entry->string = NULL; + entry->string[0] = '\0'; } - return entry; + *entry_ptr = entry; + return entry_size; } -void entry_free(Entry *entry) { - if (entry->string != NULL) { - free(entry->string); - } - - free(entry); -} +void entry_free(Entry *entry) { free(entry); }