initial commit

1 files changed, 439 insertions, 0 deletions

diff --git a/bloom.c b/bloom.c
new file mode 100644
index 0000000..bf57819
--- /dev/null
+++ b/bloom.c
@@ -0,0 +1,439 @@
+
+#include <stdlib.h>
+#include <stdio.h>
+#include <stdint.h>
+#include <string.h>
+#include <math.h>
+#include <stdbool.h>
+#include <unistd.h>
+#include <sys/stat.h>
+#include <sys/mman.h>
+
+#include "mmh3.h"
+#include "bloom.h"
+
+static const char *const bloom_errors[] = {
+    "Success",
+    "Out of memory",
+    "Unable to open file",
+    "Unable to read file",
+    "Unable to write to file",
+    "fstat() failure",
+    "Invalid file format"
+};
+
+static const uint8_t bloom_magic[8] = { '!', 'b', 'l', 'o', 'o', 'm', 'f', '!' };
+
+static bool
+ideal_size(size_t expected, double accuracy, size_t *result)
+{
+    double bits = -((double)expected * log(accuracy)) / (log(2.0) * log(2.0));
+
+    bits = ceil(bits);
+
+    if (!isfinite(bits) || bits < 1.0 || bits > (double)SIZE_MAX) {
+        return false;
+    }
+
+    *result = (size_t)bits;
+    return true;
+}
+
+
+bloom_error_t bloom_init(bloomfilter *bf, const size_t expected, const double accuracy) {
+    if (!bf || expected == 0 ||  !isfinite(accuracy) || accuracy <= 0.0 || accuracy >= 1.0) {
+        return BF_INVALIDFILE;
+    }
+
+    float stored_accuracy = (float)accuracy;
+
+    if (!isfinite(stored_accuracy) || stored_accuracy <= 0.0f || stored_accuracy >= 1.0f) {
+        return BF_INVALIDFILE;
+    }
+
+    if (!ideal_size(expected, accuracy, &bf->size)) {
+        return BF_INVALIDFILE;
+    }
+
+    bf->bitmap_size = bf->size / 8 + ((bf->size % 8) != 0);
+
+    double hashes = ((double)bf->size / (double)expected) * log(2.0);
+
+    if (!isfinite(hashes) || hashes < 1.0 || hashes > (double)BLOOM_MAX_HASHCOUNT) {
+        return BF_INVALIDFILE;
+    }
+
+    bf->hashcount = (size_t)llround(hashes);
+
+    bf->expected = expected;
+    bf->accuracy = stored_accuracy;
+
+    snprintf(bf->name, sizeof(bf->name), "DEFAULT");
+
+    bf->bitmap = calloc(bf->bitmap_size, 1);
+
+    if (!bf->bitmap) {
+        return BF_OUTOFMEMORY;
+    }
+
+    return BF_SUCCESS;
+}
+
+void
+bloom_destroy(bloomfilter *bf) {
+    if (bf->bitmap) {
+		free(bf->bitmap);
+		bf->bitmap = NULL;
+	}
+}
+
+
+// lookup table for bloom_saturation() and bloom_saturation_count()
+static const uint8_t
+bit_count_table[256] = {
+    0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4,
+    1, 2, 2, 3, 2, 3, 3, 4, 2, 3, 3, 4, 3, 4, 4, 5,
+    1, 2, 2, 3, 2, 3, 3, 4, 2, 3, 3, 4, 3, 4, 4, 5,
+    2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6,
+    1, 2, 2, 3, 2, 3, 3, 4, 2, 3, 3, 4, 3, 4, 4, 5,
+    2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6,
+    2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6,
+    3, 4, 4, 5, 4, 5, 5, 6, 4, 5, 5, 6, 5, 6, 6, 7,
+    1, 2, 2, 3, 2, 3, 3, 4, 2, 3, 3, 4, 3, 4, 4, 5,
+    2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6,
+    2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6,
+    3, 4, 4, 5, 4, 5, 5, 6, 4, 5, 5, 6, 5, 6, 6, 7,
+    2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6,
+    3, 4, 4, 5, 4, 5, 5, 6, 4, 5, 5, 6, 5, 6, 6, 7,
+    3, 4, 4, 5, 4, 5, 5, 6, 4, 5, 5, 6, 5, 6, 6, 7,
+    4, 5, 5, 6, 5, 6, 6, 7, 5, 6, 6, 7, 6, 7, 7, 8
+};
+
+
+size_t
+bloom_saturation_count(const bloomfilter *bf) {
+	size_t count = 0;
+
+	for (size_t i = 0; i < bf->bitmap_size; i++) {
+		count += bit_count_table[bf->bitmap[i]];
+	}
+
+	return count;
+}
+
+float
+bloom_saturation(const bloomfilter *bf) {
+    size_t set_bits = bloom_saturation_count(bf);
+
+    return (float)set_bits / (float)bf->size * 100.0f;
+}
+
+
+void
+bloom_add(bloomfilter *bf, const void *element, const size_t len)
+{
+    uint64_t hash[2];
+
+    mmh3_128(element, len, 0, hash);
+
+    for (size_t i = 0; i < bf->hashcount; i++) {
+        uint64_t result = (hash[0] + i * hash[1]) % bf->size;
+
+        uint64_t byte_position = result / 8;
+        uint8_t bit_position = result % 8;
+
+        bf->bitmap[byte_position] |= (0x01 << bit_position);
+    }
+}
+
+bool
+bloom_set_name(bloomfilter *bf, const char *name)
+{
+    size_t len = strlen(name);
+
+    if (!bloom_valid_filter_name((const uint8_t *)name, len + 1))
+        return false;
+
+    memcpy(bf->name, name, len + 1);
+    return true;
+}
+
+static bloom_error_t
+bloom_write_file(const bloomfilter *bf, FILE *fp)
+{
+    bloomfilter_file bff = {0};
+
+    memcpy(bff.magic, bloom_magic, sizeof(bff.magic));
+
+    bff.version     = BLOOM_FILE_VERSION;
+    bff.size        = bf->size;
+    bff.hashcount   = bf->hashcount;
+    bff.bitmap_size = bf->bitmap_size;
+    bff.expected    = bf->expected;
+    bff.accuracy    = bf->accuracy;
+
+    memcpy(bff.name, bf->name, sizeof(bff.name));
+
+    if (fwrite(&bff, sizeof(bff), 1, fp) != 1) {
+        return BF_FWRITE;
+    }
+
+    if (fwrite(bf->bitmap, 1, bf->bitmap_size, fp) != bf->bitmap_size) {
+        return BF_FWRITE;
+    }
+
+    return BF_SUCCESS;
+}
+
+bloom_error_t
+bloom_save(const bloomfilter *bf, const char *path)
+{
+    bloom_error_t result;
+    FILE *fp = NULL;
+    int fd = -1;
+
+    size_t path_len = strlen(path);
+
+    if (path_len > SIZE_MAX - sizeof(".tmp.XXXXXX")) {
+        return BF_FOPEN;
+    }
+
+    char *tmp_path = malloc(path_len + sizeof(".tmp.XXXXXX"));
+    if (tmp_path == NULL) {
+        return BF_OUTOFMEMORY;
+    }
+
+    snprintf(tmp_path, path_len + sizeof(".tmp.XXXXXX"), "%s.tmp.XXXXXX", path);
+
+    fd = mkstemp(tmp_path);
+    if (fd < 0) {
+        free(tmp_path);
+        return BF_FOPEN;
+    }
+
+    fp = fdopen(fd, "wb");
+    if (fp == NULL) {
+        close(fd);
+        unlink(tmp_path);
+        free(tmp_path);
+        return BF_FOPEN;
+    }
+
+    result = bloom_write_file(bf, fp);
+
+    if (result != BF_SUCCESS)
+        goto fail;
+
+    if (fflush(fp) != 0) {
+        result = BF_FWRITE;
+        goto fail;
+    }
+
+    if (fsync(fd) != 0) {
+        result = BF_FWRITE;
+        goto fail;
+    }
+
+    if (fclose(fp) != 0) {
+        fp = NULL;
+        result = BF_FWRITE;
+        goto fail_closed;
+    }
+
+    fp = NULL;
+
+    if (rename(tmp_path, path) != 0) {
+        result = BF_FWRITE;
+        goto fail_closed;
+    }
+
+    free(tmp_path);
+    return BF_SUCCESS;
+
+fail:
+    fclose(fp);
+    fp = NULL;
+
+fail_closed:
+    unlink(tmp_path);
+    free(tmp_path);
+
+    return result;
+}
+
+const char *
+bloom_strerror(bloom_error_t error) {
+	if (error < 0 || error >= BF_ERRORCOUNT) {
+		return "Unknown error";
+	}
+
+	return bloom_errors[error];
+}
+
+bool
+bloom_valid_filter_name(const uint8_t *name, size_t size)
+{
+    size_t len = 0;
+
+    while (len < size && name[len] != '\0') {
+        len++;
+    }
+
+    if (len == 0 || len == size) {
+        return false;
+    }
+
+    for (size_t i = 0; i < len; i++) {
+        uint8_t c = name[i];
+
+        if (!((c >= 'a' && c <= 'z') ||
+              (c >= 'A' && c <= 'Z') ||
+              (c >= '0' && c <= '9') ||
+              c == ' ' || c == '-' || c == '_' || c == '.')) {
+            return false;
+        }
+    }
+
+    return true;
+}
+
+static bloom_error_t
+validate_mapped_header(const bloomfilter_file *header, size_t file_size)
+{
+    if (memcmp(header->magic, bloom_magic, sizeof(header->magic)) != 0) {
+        return BF_INVALIDFILE;
+    }
+
+    if (header->version != BLOOM_FILE_VERSION) return BF_INVALIDFILE;
+    if (header->size == 0)                     return BF_INVALIDFILE;
+    if (header->hashcount == 0)                return BF_INVALIDFILE;
+
+    /*
+     * Bound the amount of work for each lookup. Legitimate Bloom
+     * filters should require nowhere near this many hashes.
+     */
+    if (header->hashcount > BLOOM_MAX_HASHCOUNT) {
+        return BF_INVALIDFILE;
+    }
+
+    /*
+     * Compute ceil(size / 8) without doing size + 7,
+     * which could overflow.
+     */
+    uint64_t expected_bitmap_size = header->size / 8 + ((header->size % 8) != 0);
+    if (header->bitmap_size != expected_bitmap_size) {
+        return BF_INVALIDFILE;
+    }
+
+    /*
+     * Make sure converting bitmap_size to size_t is safe.
+     */
+    if (header->bitmap_size > SIZE_MAX) {
+        return BF_INVALIDFILE;
+    }
+
+    size_t bitmap_size = (size_t)header->bitmap_size;
+
+    /*
+     * Avoid overflow in sizeof(header) + bitmap_size
+     */
+    if (bitmap_size > SIZE_MAX - sizeof(bloomfilter_file)) {
+        return BF_INVALIDFILE;
+    }
+
+    if (sizeof(bloomfilter_file) + bitmap_size != file_size) {
+        return BF_INVALIDFILE;
+    }
+
+    /*
+     * Validate filter name. Shouldn't contain many special
+     * characters.
+     */
+    if (!bloom_valid_filter_name(header->name, sizeof(header->name))) {
+      return BF_INVALIDFILE;
+    }
+
+    if (!isfinite(header->accuracy)) {
+        return BF_INVALIDFILE;
+    }
+
+    if (header->accuracy <= 0.0f || header->accuracy >= 1.0f) {
+        return BF_INVALIDFILE;
+    }
+
+    if (header->expected == 0) {
+        return BF_INVALIDFILE;
+    }
+
+    return BF_SUCCESS;
+}
+
+bloom_error_t
+bloom_map_fd(bloomfilter_mapped *bf, int fd)
+{
+    struct stat st;
+
+    if (bf == NULL)                       return BF_INVALIDFILE;
+    if (fstat(fd, &st) < 0)               return BF_FSTAT;
+    if (st.st_size < 0)                   return BF_INVALIDFILE;
+    if ((uintmax_t)st.st_size > SIZE_MAX) return BF_INVALIDFILE;
+
+    size_t file_size = (size_t)st.st_size;
+    if (file_size < sizeof(bloomfilter_file)) return BF_INVALIDFILE;
+
+    void *mapping = mmap(NULL, file_size, PROT_READ, MAP_PRIVATE, fd, 0);
+    if (mapping == MAP_FAILED) {
+        return BF_FREAD;
+    }
+
+    const bloomfilter_file *header = (const bloomfilter_file *)mapping;
+
+    bloom_error_t err = validate_mapped_header(header, file_size);
+    if (err != BF_SUCCESS) {
+        munmap(mapping, file_size);
+        return err;
+    }
+
+    const uint8_t *bitmap = (const uint8_t *)mapping + sizeof(bloomfilter_file);
+
+    bf->header = header;
+    bf->bitmap = bitmap;
+    bf->mapping = mapping;
+    bf->mapping_size = file_size;
+
+    return BF_SUCCESS;
+}
+
+void
+bloom_unmap(bloomfilter_mapped *bf)
+{
+    if (bf->mapping) {
+        munmap(bf->mapping, bf->mapping_size);
+    }
+
+    bf->header = NULL;
+    bf->bitmap = NULL;
+    bf->mapping = NULL;
+    bf->mapping_size = 0;
+}
+
+bool
+bloom_mapped_lookup(const bloomfilter_mapped *bf, const void *element, const size_t len)
+{
+    uint64_t hash[2];
+
+    mmh3_128(element, len, 0, hash);
+
+    for (size_t i = 0; i < bf->header->hashcount; i++) {
+        uint64_t result = (hash[0] + i * hash[1]) % bf->header->size;
+
+        uint64_t byte_position = result / 8;
+        uint8_t bit_position = result % 8;
+
+        if ((bf->bitmap[byte_position] & (0x01 << bit_position)) == 0) {
+            return false;
+        }
+    }
+
+    return true;
+}