amduat/tests/asl/test_asl_store_index_fs.c

#include "amduat/asl/asl_store_index_fs.h"
#include "amduat/asl/index_bloom.h"
#include "amduat/asl/ref_derive.h"
#include "amduat/asl/store.h"
#include "amduat/enc/asl1_core.h"
#include "amduat/enc/asl_core_index.h"
#include "amduat/hash/asl1.h"

#include <dirent.h>
#include <errno.h>
#include <inttypes.h>
#include <stdbool.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <unistd.h>

static bool join_path(const char *base, const char *segment, char **out_path) {
  size_t base_len;
  size_t seg_len;
  bool needs_sep;
  size_t total_len;
  char *buffer;
  size_t offset;

  if (base == NULL || segment == NULL || out_path == NULL) {
    return false;
  }
  if (base[0] == '\0' || segment[0] == '\0') {
    return false;
  }

  base_len = strlen(base);
  seg_len = strlen(segment);
  needs_sep = base[base_len - 1u] != '/';
  total_len = base_len + (needs_sep ? 1u : 0u) + seg_len + 1u;

  buffer = (char *)malloc(total_len);
  if (buffer == NULL) {
    return false;
  }

  offset = 0u;
  memcpy(buffer + offset, base, base_len);
  offset += base_len;
  if (needs_sep) {
    buffer[offset++] = '/';
  }
  memcpy(buffer + offset, segment, seg_len);
  offset += seg_len;
  buffer[offset] = '\0';

  *out_path = buffer;
  return true;
}

static bool remove_tree(const char *path) {
  struct stat st;
  DIR *dir;
  struct dirent *entry;

  if (path == NULL) {
    return false;
  }
  if (lstat(path, &st) != 0) {
    return errno == ENOENT;
  }

  if (!S_ISDIR(st.st_mode)) {
    return unlink(path) == 0;
  }

  dir = opendir(path);
  if (dir == NULL) {
    return false;
  }

  while ((entry = readdir(dir)) != NULL) {
    char *child = NULL;
    if (strcmp(entry->d_name, ".") == 0 ||
        strcmp(entry->d_name, "..") == 0) {
      continue;
    }
    if (!join_path(path, entry->d_name, &child)) {
      closedir(dir);
      return false;
    }
    if (!remove_tree(child)) {
      free(child);
      closedir(dir);
      return false;
    }
    free(child);
  }

  if (closedir(dir) != 0) {
    return false;
  }
  return rmdir(path) == 0;
}

static bool read_file(const char *path, uint8_t **out_bytes, size_t *out_len) {
  FILE *fp;
  long size;
  uint8_t *buffer;
  size_t read_len;

  if (path == NULL || out_bytes == NULL || out_len == NULL) {
    return false;
  }
  *out_bytes = NULL;
  *out_len = 0u;

  fp = fopen(path, "rb");
  if (fp == NULL) {
    return false;
  }
  if (fseek(fp, 0, SEEK_END) != 0) {
    fclose(fp);
    return false;
  }
  size = ftell(fp);
  if (size < 0) {
    fclose(fp);
    return false;
  }
  if (fseek(fp, 0, SEEK_SET) != 0) {
    fclose(fp);
    return false;
  }

  buffer = (uint8_t *)malloc((size_t)size);
  if (buffer == NULL) {
    fclose(fp);
    return false;
  }
  read_len = fread(buffer, 1u, (size_t)size, fp);
  fclose(fp);
  if (read_len != (size_t)size) {
    free(buffer);
    return false;
  }

  *out_bytes = buffer;
  *out_len = (size_t)size;
  return true;
}

static bool build_log_path(const char *root, char **out_path) {
  char *index_path = NULL;
  bool ok;

  if (root == NULL || out_path == NULL) {
    return false;
  }

  if (!join_path(root, "index", &index_path)) {
    return false;
  }
  ok = join_path(index_path, "log.asl", out_path);
  free(index_path);
  return ok;
}

static bool build_segment_path(const char *root,
                               uint64_t segment_id,
                               char **out_path) {
  int needed;
  char *buffer;

  if (root == NULL || out_path == NULL) {
    return false;
  }

  needed = snprintf(NULL,
                    0,
                    "%s/index/segments/segment-%016" PRIx64 ".asl",
                    root,
                    segment_id);
  if (needed <= 0) {
    return false;
  }

  buffer = (char *)malloc((size_t)needed + 1u);
  if (buffer == NULL) {
    return false;
  }
  snprintf(buffer,
           (size_t)needed + 1u,
           "%s/index/segments/segment-%016" PRIx64 ".asl",
           root,
           segment_id);
  *out_path = buffer;
  return true;
}

static bool build_shard_segment_path(const char *root,
                                     uint16_t shard_id,
                                     uint64_t segment_id,
                                     char **out_path) {
  int needed;
  char *buffer;

  if (root == NULL || out_path == NULL) {
    return false;
  }

  needed = snprintf(NULL,
                    0,
                    "%s/index/shards/shard-%016" PRIx64
                    "/segments/segment-%016" PRIx64 ".asl",
                    root,
                    (uint64_t)shard_id,
                    segment_id);
  if (needed <= 0) {
    return false;
  }

  buffer = (char *)malloc((size_t)needed + 1u);
  if (buffer == NULL) {
    return false;
  }
  snprintf(buffer,
           (size_t)needed + 1u,
           "%s/index/shards/shard-%016" PRIx64
           "/segments/segment-%016" PRIx64 ".asl",
           root,
           (uint64_t)shard_id,
           segment_id);
  *out_path = buffer;
  return true;
}

static uint64_t fnv1a64(const uint8_t *data, size_t len, uint64_t seed) {
  size_t i;
  uint64_t hash = seed;

  for (i = 0; i < len; ++i) {
    hash ^= data[i];
    hash *= 1099511628211ull;
  }
  return hash;
}

static uint16_t ref_shard(amduat_reference_t ref, uint16_t shard_count) {
  uint64_t hash;
  uint8_t hash_id_bytes[2];

  if (shard_count == 0u) {
    return 0u;
  }
  hash_id_bytes[0] = (uint8_t)(ref.hash_id & 0xffu);
  hash_id_bytes[1] = (uint8_t)((ref.hash_id >> 8) & 0xffu);
  hash = fnv1a64(hash_id_bytes, sizeof(hash_id_bytes),
                 14695981039346656037ull);
  if (ref.digest.len != 0u && ref.digest.data != NULL) {
    hash = fnv1a64(ref.digest.data, ref.digest.len, hash);
  }
  return (uint16_t)(hash % shard_count);
}

static char *make_temp_root(void) {
  char *templ;
  const char template_prefix[] = "/tmp/amduat_test_asl_store_index_fs_XXXXXX";

  templ = (char *)malloc(sizeof(template_prefix));
  if (templ == NULL) {
    return NULL;
  }
  memcpy(templ, template_prefix, sizeof(template_prefix));
  if (mkdtemp(templ) == NULL) {
    free(templ);
    return NULL;
  }
  return templ;
}

static int test_round_trip(void) {
  amduat_asl_store_config_t config;
  amduat_asl_store_index_fs_t fs;
  amduat_asl_store_t store;
  amduat_asl_store_error_t err;
  amduat_asl_index_state_t state;
  amduat_asl_index_state_t current_state;
  amduat_artifact_t artifact;
  amduat_artifact_t loaded;
  amduat_reference_t ref;
  uint8_t payload[6];
  char *root;
  int exit_code = 1;

  root = make_temp_root();
  if (root == NULL) {
    fprintf(stderr, "temp root failed\n");
    return 1;
  }

  config.encoding_profile_id = AMDUAT_ENC_ASL1_CORE_V1;
  config.hash_id = AMDUAT_HASH_ASL1_ID_SHA256;
  if (!amduat_asl_store_index_fs_init(&fs, config, root)) {
    fprintf(stderr, "index fs init failed\n");
    goto cleanup;
  }
  amduat_asl_store_init(&store, config, amduat_asl_store_index_fs_ops(), &fs);

  memset(payload, 0x5a, sizeof(payload));
  artifact = amduat_artifact(amduat_octets(payload, sizeof(payload)));
  ref = amduat_reference(0u, amduat_octets(NULL, 0u));

  err = amduat_asl_store_put_indexed(&store, artifact, &ref, &state);
  if (err != AMDUAT_ASL_STORE_OK) {
    fprintf(stderr, "put_indexed failed: %d\n", err);
    goto cleanup;
  }

  if (!amduat_asl_index_current_state(&store, &current_state)) {
    fprintf(stderr, "current_state failed\n");
    goto cleanup;
  }
  if (current_state.log_position != state.log_position) {
    fprintf(stderr, "log position mismatch\n");
    goto cleanup;
  }

  loaded = amduat_artifact(amduat_octets(NULL, 0u));
  err = amduat_asl_store_get_indexed(&store, ref, state, &loaded);
  if (err != AMDUAT_ASL_STORE_OK) {
    fprintf(stderr, "get_indexed failed: %d\n", err);
    goto cleanup;
  }
  if (!amduat_artifact_eq(artifact, loaded)) {
    fprintf(stderr, "artifact mismatch\n");
    amduat_artifact_free(&loaded);
    goto cleanup;
  }
  amduat_artifact_free(&loaded);

  if (state.log_position > 0u) {
    amduat_asl_index_state_t cutoff = state;
    cutoff.log_position = state.log_position - 1u;
    loaded = amduat_artifact(amduat_octets(NULL, 0u));
    err = amduat_asl_store_get_indexed(&store, ref, cutoff, &loaded);
    if (err != AMDUAT_ASL_STORE_ERR_NOT_FOUND) {
      fprintf(stderr, "cutoff lookup expected not found: %d\n", err);
      amduat_artifact_free(&loaded);
      goto cleanup;
    }
  }

  {
    char *segment_path = NULL;
    uint8_t *segment_bytes = NULL;
    size_t segment_len = 0u;
    amduat_asl_core_index_segment_t segment;
    bool bloom_nonzero = false;
    size_t i;

    memset(&segment, 0, sizeof(segment));

    if (!build_segment_path(root, 1u, &segment_path)) {
      fprintf(stderr, "segment path build failed\n");
      goto cleanup;
    }
    if (!read_file(segment_path, &segment_bytes, &segment_len)) {
      fprintf(stderr, "segment read failed\n");
      free(segment_path);
      goto cleanup;
    }
    free(segment_path);
    if (!amduat_enc_asl_core_index_decode_v1(
            amduat_octets(segment_bytes, segment_len), &segment)) {
      fprintf(stderr, "segment decode failed\n");
      free(segment_bytes);
      goto cleanup;
    }
    free(segment_bytes);

    if (segment.bloom.len != AMDUAT_ASL_INDEX_BLOOM_BYTES) {
      fprintf(stderr, "segment bloom size mismatch\n");
      amduat_enc_asl_core_index_free(&segment);
      goto cleanup;
    }
    if (!amduat_asl_index_bloom_maybe_contains(segment.bloom,
                                               ref.hash_id,
                                               ref.digest)) {
      fprintf(stderr, "segment bloom missing digest\n");
      amduat_enc_asl_core_index_free(&segment);
      goto cleanup;
    }
    for (i = 0u; i < segment.bloom.len; ++i) {
      if (segment.bloom.data[i] != 0u) {
        bloom_nonzero = true;
        break;
      }
    }
    if (!bloom_nonzero) {
      fprintf(stderr, "segment bloom was empty\n");
      amduat_enc_asl_core_index_free(&segment);
      goto cleanup;
    }
    amduat_enc_asl_core_index_free(&segment);
  }

  exit_code = 0;

cleanup:
  amduat_reference_free(&ref);
  if (!remove_tree(root)) {
    fprintf(stderr, "cleanup failed\n");
    exit_code = 1;
  }
  free(root);
  return exit_code;
}

static int test_shard_routing(void) {
  amduat_asl_store_config_t config;
  amduat_asl_store_index_fs_t fs;
  amduat_asl_store_t store;
  char *root;
  uint8_t payload_a[2] = {0x10, 0x11};
  uint8_t payload_b[2] = {0x20, 0x21};
  amduat_reference_t ref_a;
  amduat_reference_t ref_b;
  amduat_octets_t bytes_a;
  amduat_octets_t bytes_b;
  amduat_hash_id_t hash_id;
  uint16_t shard_a;
  uint16_t shard_b;
  uint64_t local_counts[2] = {0u, 0u};
  int exit_code = 1;
  size_t i;
  bool found = false;

  root = make_temp_root();
  if (root == NULL) {
    fprintf(stderr, "temp root failed\n");
    return 1;
  }

  memset(&config, 0, sizeof(config));
  config.encoding_profile_id = AMDUAT_ENC_ASL1_CORE_V1;
  config.hash_id = AMDUAT_HASH_ASL1_ID_SHA256;
  if (!amduat_asl_store_index_fs_init(&fs, config, root)) {
    fprintf(stderr, "index fs init failed\n");
    goto cleanup;
  }
  amduat_asl_store_index_fs_set_shard_count(&fs, 2);
  amduat_asl_store_init(&store, config, amduat_asl_store_index_fs_ops(), &fs);

  hash_id = config.hash_id;
  ref_a = amduat_reference(0u, amduat_octets(NULL, 0u));
  ref_b = amduat_reference(0u, amduat_octets(NULL, 0u));
  bytes_a = amduat_octets(NULL, 0u);
  bytes_b = amduat_octets(NULL, 0u);

  for (i = 0; i < 256; ++i) {
    payload_a[0] = (uint8_t)i;
    payload_b[0] = (uint8_t)(255u - i);
    amduat_reference_free(&ref_a);
    amduat_reference_free(&ref_b);
    amduat_octets_free(&bytes_a);
    amduat_octets_free(&bytes_b);
    if (!amduat_asl_ref_derive(amduat_artifact(amduat_octets(payload_a,
                                                            sizeof(payload_a))),
                               AMDUAT_ENC_ASL1_CORE_V1,
                               hash_id,
                               &ref_a,
                               &bytes_a) ||
        !amduat_asl_ref_derive(amduat_artifact(amduat_octets(payload_b,
                                                            sizeof(payload_b))),
                               AMDUAT_ENC_ASL1_CORE_V1,
                               hash_id,
                               &ref_b,
                               &bytes_b)) {
      fprintf(stderr, "ref derive failed\n");
      goto cleanup;
    }
    shard_a = ref_shard(ref_a, 2);
    shard_b = ref_shard(ref_b, 2);
    if (shard_a != shard_b) {
      found = true;
      break;
    }
  }

  if (!found) {
    fprintf(stderr, "shard selection failed\n");
    goto cleanup;
  }

  {
    amduat_reference_t stored_ref;
    amduat_asl_index_state_t state;
    uint64_t segment_id;
    char *segment_path = NULL;
    amduat_asl_store_error_t err;

    stored_ref = amduat_reference(0u, amduat_octets(NULL, 0u));
    err = amduat_asl_store_put_indexed(&store,
                                       amduat_artifact(amduat_octets(payload_a,
                                                                    sizeof(payload_a))),
                                       &stored_ref,
                                       &state);
    if (err != AMDUAT_ASL_STORE_OK) {
      fprintf(stderr, "put shard a failed\n");
      amduat_reference_free(&stored_ref);
      goto cleanup;
    }
    amduat_reference_free(&stored_ref);
    local_counts[shard_a] += 1u;
    segment_id = ((uint64_t)shard_a << 48) | local_counts[shard_a];
    if (!build_shard_segment_path(root, shard_a, segment_id, &segment_path)) {
      fprintf(stderr, "segment path failed\n");
      goto cleanup;
    }
    if (access(segment_path, F_OK) != 0) {
      fprintf(stderr, "segment missing in shard a\n");
      free(segment_path);
      goto cleanup;
    }
    free(segment_path);
  }

  {
    amduat_reference_t stored_ref;
    amduat_asl_index_state_t state;
    uint64_t segment_id;
    char *segment_path = NULL;
    amduat_asl_store_error_t err;

    stored_ref = amduat_reference(0u, amduat_octets(NULL, 0u));
    err = amduat_asl_store_put_indexed(&store,
                                       amduat_artifact(amduat_octets(payload_b,
                                                                    sizeof(payload_b))),
                                       &stored_ref,
                                       &state);
    if (err != AMDUAT_ASL_STORE_OK) {
      fprintf(stderr, "put shard b failed\n");
      amduat_reference_free(&stored_ref);
      goto cleanup;
    }
    amduat_reference_free(&stored_ref);
    local_counts[shard_b] += 1u;
    segment_id = ((uint64_t)shard_b << 48) | local_counts[shard_b];
    if (!build_shard_segment_path(root, shard_b, segment_id, &segment_path)) {
      fprintf(stderr, "segment path failed\n");
      goto cleanup;
    }
    if (access(segment_path, F_OK) != 0) {
      fprintf(stderr, "segment missing in shard b\n");
      free(segment_path);
      goto cleanup;
    }
    free(segment_path);
  }

  exit_code = 0;

cleanup:
  amduat_reference_free(&ref_a);
  amduat_reference_free(&ref_b);
  amduat_octets_free(&bytes_a);
  amduat_octets_free(&bytes_b);
  if (!remove_tree(root)) {
    fprintf(stderr, "cleanup failed\n");
    exit_code = 1;
  }
  free(root);
  return exit_code;
}

static int test_snapshot_truncation(void) {
  amduat_asl_store_config_t config;
  amduat_asl_store_index_fs_t fs;
  amduat_asl_store_t store;
  amduat_asl_store_error_t err;
  amduat_asl_index_state_t state;
  amduat_artifact_t artifact;
  amduat_artifact_t loaded;
  amduat_reference_t ref;
  uint8_t payload[8];
  char *root;
  char *log_path = NULL;
  uint8_t *log_before = NULL;
  uint8_t *log_after = NULL;
  size_t log_before_len = 0u;
  size_t log_after_len = 0u;
  int exit_code = 1;

  root = make_temp_root();
  if (root == NULL) {
    fprintf(stderr, "temp root failed\n");
    return 1;
  }

  memset(&config, 0, sizeof(config));
  config.encoding_profile_id = AMDUAT_ENC_ASL1_CORE_V1;
  config.hash_id = AMDUAT_HASH_ASL1_ID_SHA256;
  if (!amduat_asl_store_index_fs_init(&fs, config, root)) {
    fprintf(stderr, "index fs init failed\n");
    goto cleanup;
  }
  amduat_asl_store_init(&store, config, amduat_asl_store_index_fs_ops(), &fs);

  memset(payload, 0x33, sizeof(payload));
  artifact = amduat_artifact(amduat_octets(payload, sizeof(payload)));
  ref = amduat_reference(0u, amduat_octets(NULL, 0u));
  err = amduat_asl_store_put_indexed(&store, artifact, &ref, &state);
  if (err != AMDUAT_ASL_STORE_OK) {
    fprintf(stderr, "put_indexed failed: %d\n", err);
    goto cleanup;
  }

  if (!build_log_path(root, &log_path) ||
      !read_file(log_path, &log_before, &log_before_len)) {
    fprintf(stderr, "log read before snapshot failed\n");
    goto cleanup;
  }

  err = amduat_asl_store_index_fs_snapshot_create(&fs, 1u, NULL, NULL);
  if (err != AMDUAT_ASL_STORE_OK) {
    fprintf(stderr, "snapshot create failed: %d\n", err);
    goto cleanup;
  }

  if (!read_file(log_path, &log_after, &log_after_len)) {
    fprintf(stderr, "log read after snapshot failed\n");
    goto cleanup;
  }
  if (log_after_len >= log_before_len) {
    fprintf(stderr, "log did not shrink after snapshot\n");
    goto cleanup;
  }

  if (!amduat_asl_index_current_state(&store, &state)) {
    fprintf(stderr, "current_state failed\n");
    goto cleanup;
  }
  loaded = amduat_artifact(amduat_octets(NULL, 0u));
  err = amduat_asl_store_get_indexed(&store, ref, state, &loaded);
  if (err != AMDUAT_ASL_STORE_OK) {
    fprintf(stderr, "get_indexed after snapshot failed: %d\n", err);
    goto cleanup;
  }
  if (!amduat_artifact_eq(artifact, loaded)) {
    fprintf(stderr, "artifact mismatch after snapshot\n");
    amduat_artifact_free(&loaded);
    goto cleanup;
  }
  amduat_artifact_free(&loaded);

  exit_code = 0;

cleanup:
  free(log_path);
  free(log_before);
  free(log_after);
  amduat_reference_free(&ref);
  if (!remove_tree(root)) {
    fprintf(stderr, "cleanup failed\n");
    exit_code = 1;
  }
  free(root);
  return exit_code;
}

int main(void) {
  if (test_round_trip() != 0) {
    return 1;
  }
  if (test_snapshot_truncation() != 0) {
    return 1;
  }
  return test_shard_routing();
}