/*

 * Copyright 2020 Google LLC

 *

 * Use of this source code is governed by a BSD-style

 * license that can be found in the LICENSE file or at

 * https://developers.google.com/open-source/licenses/bsd

 */

#include "table.h"

#include "system.h"

#include "block.h"

#include "blocksource.h"

#include "constants.h"

#include "iter.h"

#include "record.h"

#include "reftable-error.h"

static struct reftable_table_offsets *

table_offsets_for(struct reftable_table *t, uint8_t typ)

{

  switch (typ) {

  case REFTABLE_BLOCK_TYPE_REF:

    return &t->ref_offsets;

  case REFTABLE_BLOCK_TYPE_LOG:

    return &t->log_offsets;

  case REFTABLE_BLOCK_TYPE_OBJ:

    return &t->obj_offsets;

  }

  abort();

}

enum reftable_hash reftable_table_hash_id(struct reftable_table *t)

{

  return t->hash_id;

}

const char *reftable_table_name(struct reftable_table *t)

{

  return t->name;

}

static int parse_footer(struct reftable_table *t, uint8_t *footer,

      uint8_t *header)

{

  uint8_t *f = footer;

  uint8_t first_block_typ;

  int err = 0;

  uint32_t computed_crc;

  uint32_t file_crc;

  if (memcmp(f, "REFT", 4)) {

    err = REFTABLE_FORMAT_ERROR;

    goto done;

  }

  f += 4;

  if (memcmp(footer, header, header_size(t->version))) {

    err = REFTABLE_FORMAT_ERROR;

    goto done;

  }

  f++;

  t->block_size = reftable_get_be24(f);

  f += 3;

  t->min_update_index = reftable_get_be64(f);

  f += 8;

  t->max_update_index = reftable_get_be64(f);

  f += 8;

  if (t->version == 1) {

    t->hash_id = REFTABLE_HASH_SHA1;

  } else {

    switch (reftable_get_be32(f)) {

    case REFTABLE_FORMAT_ID_SHA1:

      t->hash_id = REFTABLE_HASH_SHA1;

      break;

    case REFTABLE_FORMAT_ID_SHA256:

      t->hash_id = REFTABLE_HASH_SHA256;

      break;

    default:

      err = REFTABLE_FORMAT_ERROR;

      goto done;

    }

    f += 4;

  }

  t->ref_offsets.index_offset = reftable_get_be64(f);

  f += 8;

  t->obj_offsets.offset = reftable_get_be64(f);

  f += 8;

  t->object_id_len = t->obj_offsets.offset & ((1 << 5) - 1);

  t->obj_offsets.offset >>= 5;

  t->obj_offsets.index_offset = reftable_get_be64(f);

  f += 8;

  t->log_offsets.offset = reftable_get_be64(f);

  f += 8;

  t->log_offsets.index_offset = reftable_get_be64(f);

  f += 8;

  computed_crc = crc32(0, footer, f - footer);

  file_crc = reftable_get_be32(f);

  f += 4;

  if (computed_crc != file_crc) {

    err = REFTABLE_FORMAT_ERROR;

    goto done;

  }

  first_block_typ = header[header_size(t->version)];

  t->ref_offsets.is_present = (first_block_typ == REFTABLE_BLOCK_TYPE_REF);

  t->ref_offsets.offset = 0;

  t->log_offsets.is_present = (first_block_typ == REFTABLE_BLOCK_TYPE_LOG ||

             t->log_offsets.offset > 0);

  t->obj_offsets.is_present = t->obj_offsets.offset > 0;

  if (t->obj_offsets.is_present && !t->object_id_len) {

    err = REFTABLE_FORMAT_ERROR;

    goto done;

  }

  err = 0;

done:

  return err;

}

struct table_iter {

  struct reftable_table *table;

  uint8_t typ;

  uint64_t block_off;

  struct reftable_block block;

  struct block_iter bi;

  int is_finished;

};

static int table_iter_init(struct table_iter *ti, struct reftable_table *t)

{

  struct block_iter bi = BLOCK_ITER_INIT;

  memset(ti, 0, sizeof(*ti));

  reftable_table_incref(t);

  ti->table = t;

  ti->bi = bi;

  return 0;

}

static int table_iter_next_in_block(struct table_iter *ti,

            struct reftable_record *rec)

{

  int res = block_iter_next(&ti->bi, rec);

  if (res == 0 && reftable_record_type(rec) == REFTABLE_BLOCK_TYPE_REF) {

    rec->u.ref.update_index += ti->table->min_update_index;

  }

  return res;

}

static void table_iter_block_done(struct table_iter *ti)

{

  reftable_block_release(&ti->block);

  block_iter_reset(&ti->bi);

}

int table_init_block(struct reftable_table *t, struct reftable_block *block,

         uint64_t next_off, uint8_t want_typ)

{

  uint32_t header_off = next_off ? 0 : header_size(t->version);

  int err;

  if (next_off >= t->size)

    return 1;

  err = reftable_block_init(block, &t->source, next_off, header_off,

          t->block_size, hash_size(t->hash_id), want_typ);

  if (err)

    reftable_block_release(block);

  return err;

}

static void table_iter_close(struct table_iter *ti)

{

  table_iter_block_done(ti);

  block_iter_close(&ti->bi);

  reftable_table_decref(ti->table);

}

static int table_iter_next_block(struct table_iter *ti)

{

  uint64_t next_block_off = ti->block_off + ti->block.full_block_size;

  int err;

  err = table_init_block(ti->table, &ti->block, next_block_off, ti->typ);

  if (err > 0)

    ti->is_finished = 1;

  if (err)

    return err;

  ti->block_off = next_block_off;

  ti->is_finished = 0;

  block_iter_init(&ti->bi, &ti->block);

  return 0;

}

static int table_iter_next(struct table_iter *ti, struct reftable_record *rec)

{

  if (reftable_record_type(rec) != ti->typ)

    return REFTABLE_API_ERROR;

  while (1) {

    int err;

    if (ti->is_finished)

      return 1;

    /*

     * Check whether the current block still has more records. If

     * so, return it. If the iterator returns positive then the

     * current block has been exhausted.

     */

    err = table_iter_next_in_block(ti, rec);

    if (err <= 0)

      return err;

    /*

     * Otherwise, we need to continue to the next block in the

     * table and retry. If there are no more blocks then the

     * iterator is drained.

     */

    err = table_iter_next_block(ti);

    if (err) {

      ti->is_finished = 1;

      return err;

    }

  }

}

static int table_iter_seek_to(struct table_iter *ti, uint64_t off, uint8_t typ)

{

  int err;

  err = table_init_block(ti->table, &ti->block, off, typ);

  if (err > 0)

    return REFTABLE_FORMAT_ERROR;

  if (err != 0)

    return err;

  ti->typ = reftable_block_type(&ti->block);

  ti->block_off = off;

  block_iter_init(&ti->bi, &ti->block);

  ti->is_finished = 0;

  return 0;

}

static int table_iter_seek_start(struct table_iter *ti, uint8_t typ, int index)

{

  struct reftable_table_offsets *offs = table_offsets_for(ti->table, typ);

  uint64_t off = offs->offset;

  if (index) {

    off = offs->index_offset;

    if (off == 0) {

      return 1;

    }

    typ = REFTABLE_BLOCK_TYPE_INDEX;

  }

  return table_iter_seek_to(ti, off, typ);

}

static int table_iter_seek_linear(struct table_iter *ti,

          struct reftable_record *want)

{

  struct reftable_buf want_key = REFTABLE_BUF_INIT;

  struct reftable_buf got_key = REFTABLE_BUF_INIT;

  struct reftable_record rec;

  int err;

  err = reftable_record_init(&rec, reftable_record_type(want));

  if (err < 0)

    goto done;

  err = reftable_record_key(want, &want_key);

  if (err < 0)

    goto done;

  /*

   * First we need to locate the block that must contain our record. To

   * do so we scan through blocks linearly until we find the first block

   * whose first key is bigger than our wanted key. Once we have found

   * that block we know that the key must be contained in the preceding

   * block.

   *

   * This algorithm is somewhat unfortunate because it means that we

   * always have to seek one block too far and then back up. But as we

   * can only decode the _first_ key of a block but not its _last_ key we

   * have no other way to do this.

   */

  while (1) {

    struct table_iter next = *ti;

    /*

     * We must be careful to not modify underlying data of `ti`

     * because we may find that `next` does not contain our desired

     * block, but that `ti` does. In that case, we would discard

     * `next` and continue with `ti`.

     *

     * This also means that we cannot reuse allocated memory for

     * `next` here. While it would be great if we could, it should

     * in practice not be too bad given that we should only ever

     * end up doing linear seeks with at most three blocks. As soon

     * as we have more than three blocks we would have an index, so

     * we would not do a linear search there anymore.

     */

    memset(&next.block.block_data, 0, sizeof(next.block.block_data));

    next.block.zstream = NULL;

    next.block.uncompressed_data = NULL;

    next.block.uncompressed_cap = 0;

    err = table_iter_next_block(&next);

    if (err < 0)

      goto done;

    if (err > 0)

      break;

    err = reftable_block_first_key(&next.block, &got_key);

    if (err < 0)

      goto done;

    if (reftable_buf_cmp(&got_key, &want_key) > 0) {

      table_iter_block_done(&next);

      break;

    }

    table_iter_block_done(ti);

    *ti = next;

  }

  /*

   * We have located the block that must contain our record, so we seek

   * the wanted key inside of it. If the block does not contain our key

   * we know that the corresponding record does not exist.

   */

  block_iter_init(&ti->bi, &ti->block);

  err = block_iter_seek_key(&ti->bi, &want_key);

  if (err < 0)

    goto done;

  err = 0;

done:

  reftable_record_release(&rec);

  reftable_buf_release(&want_key);

  reftable_buf_release(&got_key);

  return err;

}

static int table_iter_seek_indexed(struct table_iter *ti,

           struct reftable_record *rec)

{

  struct reftable_record want_index = {

    .type = REFTABLE_BLOCK_TYPE_INDEX, .u.idx = { .last_key = REFTABLE_BUF_INIT }

  };

  struct reftable_record index_result = {

    .type = REFTABLE_BLOCK_TYPE_INDEX,

    .u.idx = { .last_key = REFTABLE_BUF_INIT },

  };

  int err;

  err = reftable_record_key(rec, &want_index.u.idx.last_key);

  if (err < 0)

    goto done;

  /*

   * The index may consist of multiple levels, where each level may have

   * multiple index blocks. We start by doing a linear search in the

   * highest layer that identifies the relevant index block as well as

   * the record inside that block that corresponds to our wanted key.

   */

  err = table_iter_seek_linear(ti, &want_index);

  if (err < 0)

    goto done;

  /*

   * Traverse down the levels until we find a non-index entry.

   */

  while (1) {

    /*

     * In case we seek a record that does not exist the index iter

     * will tell us that the iterator is over. This works because

     * the last index entry of the current level will contain the

     * last key it knows about. So in case our seeked key is larger

     * than the last indexed key we know that it won't exist.

     *

     * There is one subtlety in the layout of the index section

     * that makes this work as expected: the highest-level index is

     * at end of the section and will point backwards and thus we

     * start reading from the end of the index section, not the

     * beginning.

     *

     * If that wasn't the case and the order was reversed then the

     * linear seek would seek into the lower levels and traverse

     * all levels of the index only to find out that the key does

     * not exist.

     */

    err = table_iter_next(ti, &index_result);

    if (err != 0)

      goto done;

    err = table_iter_seek_to(ti, index_result.u.idx.offset, 0);

    if (err != 0)

      goto done;

    block_iter_init(&ti->bi, &ti->block);

    err = block_iter_seek_key(&ti->bi, &want_index.u.idx.last_key);

    if (err < 0)

      goto done;

    if (ti->typ == reftable_record_type(rec)) {

      err = 0;

      break;

    }

    if (ti->typ != REFTABLE_BLOCK_TYPE_INDEX) {

      err = REFTABLE_FORMAT_ERROR;

      goto done;

    }

  }

done:

  reftable_record_release(&want_index);

  reftable_record_release(&index_result);

  return err;

}

static int table_iter_seek(struct table_iter *ti,

         struct reftable_record *want)

{

  uint8_t typ = reftable_record_type(want);

  struct reftable_table_offsets *offs = table_offsets_for(ti->table, typ);

  int err;

  err = table_iter_seek_start(ti, reftable_record_type(want),

            !!offs->index_offset);

  if (err < 0)

    goto out;

  if (offs->index_offset)

    err = table_iter_seek_indexed(ti, want);

  else

    err = table_iter_seek_linear(ti, want);

  if (err)

    goto out;

out:

  return err;

}

static int table_iter_seek_void(void *ti, struct reftable_record *want)

{

  return table_iter_seek(ti, want);

}

static int table_iter_next_void(void *ti, struct reftable_record *rec)

{

  return table_iter_next(ti, rec);

}

static void table_iter_close_void(void *ti)

{

  table_iter_close(ti);

}

static struct reftable_iterator_vtable table_iter_vtable = {

  .seek = &table_iter_seek_void,

  .next = &table_iter_next_void,

  .close = &table_iter_close_void,

};

static void iterator_from_table_iter(struct reftable_iterator *it,

             struct table_iter *ti)

{

  assert(!it->ops);

  it->iter_arg = ti;

  it->ops = &table_iter_vtable;

}

int table_init_iter(struct reftable_table *t,

        struct reftable_iterator *it,

        uint8_t typ)

{

  struct reftable_table_offsets *offs = table_offsets_for(t, typ);

  if (offs->is_present) {

    struct table_iter *ti;

    REFTABLE_ALLOC_ARRAY(ti, 1);

    if (!ti)

      return REFTABLE_OUT_OF_MEMORY_ERROR;

    table_iter_init(ti, t);

    iterator_from_table_iter(it, ti);

  } else {

    iterator_set_empty(it);

  }

  return 0;

}

int reftable_table_init_ref_iterator(struct reftable_table *t,

             struct reftable_iterator *it)

{

  return table_init_iter(t, it, REFTABLE_BLOCK_TYPE_REF);

}

int reftable_table_init_log_iterator(struct reftable_table *t,

             struct reftable_iterator *it)

{

  return table_init_iter(t, it, REFTABLE_BLOCK_TYPE_LOG);

}

int reftable_table_new(struct reftable_table **out,

           struct reftable_block_source *source, char const *name)

{

  struct reftable_block_data footer = { 0 };

  struct reftable_block_data header = { 0 };

  struct reftable_table *t;

  uint64_t file_size = block_source_size(source);

  uint32_t read_size;

  ssize_t bytes_read;

  int err;

  REFTABLE_CALLOC_ARRAY(t, 1);

  if (!t) {

    err = REFTABLE_OUT_OF_MEMORY_ERROR;

    goto done;

  }

  /*

   * We need one extra byte to read the type of first block. We also

   * pretend to always be reading v2 of the format because it is larger.

   */

  read_size = header_size(2) + 1;

  if (read_size > file_size) {

    err = REFTABLE_FORMAT_ERROR;

    goto done;

  }

  bytes_read = block_source_read_data(source, &header, 0, read_size);

  if (bytes_read < 0 || (size_t)bytes_read != read_size) {

    err = REFTABLE_IO_ERROR;

    goto done;

  }

  if (memcmp(header.data, "REFT", 4)) {

    err = REFTABLE_FORMAT_ERROR;

    goto done;

  }

  t->version = header.data[4];

  if (t->version != 1 && t->version != 2) {

    err = REFTABLE_FORMAT_ERROR;

    goto done;

  }

  if (file_size < header_size(t->version) + footer_size(t->version)) {

    err = REFTABLE_FORMAT_ERROR;

    goto done;

  }

  t->size = file_size - footer_size(t->version);

  t->source = *source;

  t->name = reftable_strdup(name);

  if (!t->name) {

    err = REFTABLE_OUT_OF_MEMORY_ERROR;

    goto done;

  }

  t->hash_id = 0;

  t->refcount = 1;

  bytes_read = block_source_read_data(source, &footer, t->size,

              footer_size(t->version));

  if (bytes_read < 0 || (size_t)bytes_read != footer_size(t->version)) {

    err = REFTABLE_IO_ERROR;

    goto done;

  }

  err = parse_footer(t, footer.data, header.data);

  if (err)

    goto done;

  *out = t;

done:

  block_source_release_data(&footer);

  block_source_release_data(&header);

  if (err) {

    if (t)

      reftable_free(t->name);

    reftable_free(t);

    block_source_close(source);

  }

  return err;

}

void reftable_table_incref(struct reftable_table *t)

{

  t->refcount++;

}

void reftable_table_decref(struct reftable_table *t)

{

  if (!t)

    return;

  if (--t->refcount)

    return;

  block_source_close(&t->source);

  REFTABLE_FREE_AND_NULL(t->name);

  reftable_free(t);

}

static int reftable_table_refs_for_indexed(struct reftable_table *t,

             struct reftable_iterator *it,

             uint8_t *oid)

{

  struct reftable_record want = {

    .type = REFTABLE_BLOCK_TYPE_OBJ,

    .u.obj = {

      .hash_prefix = oid,

      .hash_prefix_len = t->object_id_len,

    },

  };

  struct reftable_iterator oit = { NULL };

  struct reftable_record got = {

    .type = REFTABLE_BLOCK_TYPE_OBJ,

    .u.obj = { 0 },

  };

  int err = 0;

  struct indexed_table_ref_iter *itr = NULL;

  /* Look through the reverse index. */

  err = table_init_iter(t, &oit, REFTABLE_BLOCK_TYPE_OBJ);

  if (err < 0)

    goto done;

  err = iterator_seek(&oit, &want);

  if (err != 0)

    goto done;

  /* read out the reftable_obj_record */

  err = iterator_next(&oit, &got);

  if (err < 0)

    goto done;

  if (err > 0 || memcmp(want.u.obj.hash_prefix, got.u.obj.hash_prefix,

            t->object_id_len)) {

    /* didn't find it; return empty iterator */

    iterator_set_empty(it);

    err = 0;

    goto done;

  }

  err = indexed_table_ref_iter_new(&itr, t, oid, hash_size(t->hash_id),

           got.u.obj.offsets,

           got.u.obj.offset_len);

  if (err < 0)

    goto done;

  got.u.obj.offsets = NULL;

  iterator_from_indexed_table_ref_iter(it, itr);

done:

  reftable_iterator_destroy(&oit);

  reftable_record_release(&got);

  return err;

}

static int reftable_table_refs_for_unindexed(struct reftable_table *t,

               struct reftable_iterator *it,

               uint8_t *oid)

{

  struct table_iter *ti;

  struct filtering_ref_iterator *filter = NULL;

  struct filtering_ref_iterator empty = FILTERING_REF_ITERATOR_INIT;

  uint32_t oid_len = hash_size(t->hash_id);

  int err;

  REFTABLE_ALLOC_ARRAY(ti, 1);

  if (!ti) {

    err = REFTABLE_OUT_OF_MEMORY_ERROR;

    goto out;

  }

  table_iter_init(ti, t);

  err = table_iter_seek_start(ti, REFTABLE_BLOCK_TYPE_REF, 0);

  if (err < 0)

    goto out;

  filter = reftable_malloc(sizeof(*filter));

  if (!filter) {

    err = REFTABLE_OUT_OF_MEMORY_ERROR;

    goto out;

  }

  *filter = empty;

  err = reftable_buf_add(&filter->oid, oid, oid_len);

  if (err < 0)

    goto out;

  iterator_from_table_iter(&filter->it, ti);

  iterator_from_filtering_ref_iterator(it, filter);

  err = 0;

out:

  if (err < 0) {

    if (ti)

      table_iter_close(ti);

    reftable_free(ti);

    if (filter) {

      reftable_buf_release(&filter->oid);

      reftable_free(filter);

    }

  }

  return err;

}

int reftable_table_refs_for(struct reftable_table *t,

          struct reftable_iterator *it, uint8_t *oid)

{

  if (t->obj_offsets.is_present)

    return reftable_table_refs_for_indexed(t, it, oid);

  return reftable_table_refs_for_unindexed(t, it, oid);

}

uint64_t reftable_table_max_update_index(struct reftable_table *t)

{

  return t->max_update_index;

}

uint64_t reftable_table_min_update_index(struct reftable_table *t)

{

  return t->min_update_index;

}

int reftable_table_iterator_init(struct reftable_table_iterator *it,

         struct reftable_table *t)

{

  struct table_iter *ti;

  int err;

  REFTABLE_ALLOC_ARRAY(ti, 1);

  if (!ti)

    return REFTABLE_OUT_OF_MEMORY_ERROR;

  err = table_iter_init(ti, t);

  if (err < 0)

    goto out;

  it->iter_arg = ti;

  err = 0;

out:

  if (err < 0)

    reftable_free(ti);

  return err;

}

void reftable_table_iterator_release(struct reftable_table_iterator *it)

{

  if (!it->iter_arg)

    return;

  table_iter_close(it->iter_arg);

  reftable_free(it->iter_arg);

  it->iter_arg = NULL;

}

int reftable_table_iterator_next(struct reftable_table_iterator *it,

         const struct reftable_block **out)

{

  struct table_iter *ti = it->iter_arg;

  int err;

  err = table_iter_next_block(ti);

  if (err)

    return err;

  *out = &ti->block;

  return 0;

}