/*
 * 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 "block.h"

#include "blocksource.h"
#include "constants.h"
#include "iter.h"
#include "record.h"
#include "reftable-error.h"
#include "system.h"

size_t header_size(int version)
{
	switch (version) {
	case 1:
		return 24;
	case 2:
		return 28;
	}
	abort();
}

size_t footer_size(int version)
{
	switch (version) {
	case 1:
		return 68;
	case 2:
		return 72;
	}
	abort();
}

static int block_writer_register_restart(struct block_writer *w, int n,
					 int is_restart, struct reftable_buf *key)
{
	uint32_t rlen;
	int err;

	rlen = w->restart_len;
	if (rlen >= MAX_RESTARTS)
		is_restart = 0;

	if (is_restart)
		rlen++;
	if (2 + 3 * rlen + n > w->block_size - w->next)
		return REFTABLE_ENTRY_TOO_BIG_ERROR;
	if (is_restart) {
		REFTABLE_ALLOC_GROW_OR_NULL(w->restarts, w->restart_len + 1,
					    w->restart_cap);
		if (!w->restarts)
			return REFTABLE_OUT_OF_MEMORY_ERROR;
		w->restarts[w->restart_len++] = w->next;
	}

	w->next += n;

	reftable_buf_reset(&w->last_key);
	err = reftable_buf_add(&w->last_key, key->buf, key->len);
	if (err < 0)
		return err;

	w->entries++;
	return 0;
}

int block_writer_init(struct block_writer *bw, uint8_t typ, uint8_t *block,
		      uint32_t block_size, uint32_t header_off, uint32_t hash_size)
{
	bw->block = block;
	bw->hash_size = hash_size;
	bw->block_size = block_size;
	bw->header_off = header_off;
	bw->block[header_off] = typ;
	bw->next = header_off + 4;
	bw->restart_interval = 16;
	bw->entries = 0;
	bw->restart_len = 0;
	bw->last_key.len = 0;
	if (!bw->zstream) {
		REFTABLE_CALLOC_ARRAY(bw->zstream, 1);
		if (!bw->zstream)
			return REFTABLE_OUT_OF_MEMORY_ERROR;
		if (deflateInit(bw->zstream, 9) != Z_OK) {
			REFTABLE_FREE_AND_NULL(bw->zstream);
			return REFTABLE_ZLIB_ERROR;
		}
	}

	return 0;
}

uint8_t block_writer_type(struct block_writer *bw)
{
	return bw->block[bw->header_off];
}

/*
 * Adds the reftable_record to the block. Returns 0 on success and
 * appropriate error codes on failure.
 */
int block_writer_add(struct block_writer *w, struct reftable_record *rec)
{
	struct reftable_buf empty = REFTABLE_BUF_INIT;
	struct reftable_buf last =
		w->entries % w->restart_interval == 0 ? empty : w->last_key;
	struct string_view out = {
		.buf = w->block + w->next,
		.len = w->block_size - w->next,
	};
	struct string_view start = out;
	int is_restart = 0;
	int n = 0;
	int err;

	err = reftable_record_key(rec, &w->scratch);
	if (err < 0)
		goto done;

	if (!w->scratch.len) {
		err = REFTABLE_API_ERROR;
		goto done;
	}

	n = reftable_encode_key(&is_restart, out, last, w->scratch,
				reftable_record_val_type(rec));
	if (n < 0) {
		err = n;
		goto done;
	}
	string_view_consume(&out, n);

	n = reftable_record_encode(rec, out, w->hash_size);
	if (n < 0) {
		err = n;
		goto done;
	}
	string_view_consume(&out, n);

	err = block_writer_register_restart(w, start.len - out.len, is_restart,
					    &w->scratch);
done:
	return err;
}

int block_writer_finish(struct block_writer *w)
{
	for (uint32_t i = 0; i < w->restart_len; i++) {
		reftable_put_be24(w->block + w->next, w->restarts[i]);
		w->next += 3;
	}

	reftable_put_be16(w->block + w->next, w->restart_len);
	w->next += 2;
	reftable_put_be24(w->block + 1 + w->header_off, w->next);

	/*
	 * Log records are stored zlib-compressed. Note that the compression
	 * also spans over the restart points we have just written.
	 */
	if (block_writer_type(w) == REFTABLE_BLOCK_TYPE_LOG) {
		int block_header_skip = 4 + w->header_off;
		uLongf src_len = w->next - block_header_skip, compressed_len;
		int ret;

		ret = deflateReset(w->zstream);
		if (ret != Z_OK)
			return REFTABLE_ZLIB_ERROR;

		/*
		 * Precompute the upper bound of how many bytes the compressed
		 * data may end up with. Combined with `Z_FINISH`, `deflate()`
		 * is guaranteed to return `Z_STREAM_END`.
		 */
		compressed_len = deflateBound(w->zstream, src_len);
		REFTABLE_ALLOC_GROW_OR_NULL(w->compressed, compressed_len,
					    w->compressed_cap);
		if (!w->compressed) {
			ret = REFTABLE_OUT_OF_MEMORY_ERROR;
			return ret;
		}

		w->zstream->next_out = w->compressed;
		w->zstream->avail_out = compressed_len;
		w->zstream->next_in = w->block + block_header_skip;
		w->zstream->avail_in = src_len;

		/*
		 * We want to perform all decompression in a single step, which
		 * is why we can pass Z_FINISH here. As we have precomputed the
		 * deflated buffer's size via `deflateBound()` this function is
		 * guaranteed to succeed according to the zlib documentation.
		 */
		ret = deflate(w->zstream, Z_FINISH);
		if (ret != Z_STREAM_END)
			return REFTABLE_ZLIB_ERROR;

		/*
		 * Overwrite the uncompressed data we have already written and
		 * adjust the `next` pointer to point right after the
		 * compressed data.
		 */
		memcpy(w->block + block_header_skip, w->compressed,
		       w->zstream->total_out);
		w->next = w->zstream->total_out + block_header_skip;
	}

	return w->next;
}

static int read_block(struct reftable_block_source *source,
		      struct reftable_block_data *dest, uint64_t off,
		      uint32_t sz)
{
	size_t size = block_source_size(source);
	block_source_release_data(dest);
	if (off >= size)
		return 0;
	if (off + sz > size)
		sz = size - off;
	return block_source_read_data(source, dest, off, sz);
}

int reftable_block_init(struct reftable_block *block,
			struct reftable_block_source *source,
			uint32_t offset, uint32_t header_size,
			uint32_t table_block_size, uint32_t hash_size,
			uint8_t want_type)
{
	uint32_t guess_block_size = table_block_size ?
		table_block_size : DEFAULT_BLOCK_SIZE;
	uint32_t full_block_size = table_block_size;
	uint16_t restart_count;
	uint32_t restart_off;
	uint32_t block_size;
	uint8_t block_type;
	int err;

	err = read_block(source, &block->block_data, offset, guess_block_size);
	if (err < 0)
		goto done;

	block_type = block->block_data.data[header_size];
	if (!reftable_is_block_type(block_type)) {
		err = REFTABLE_FORMAT_ERROR;
		goto done;
	}
	if (want_type != REFTABLE_BLOCK_TYPE_ANY && block_type != want_type) {
		err = 1;
		goto done;
	}

	block_size = reftable_get_be24(block->block_data.data + header_size + 1);
	if (block_size > guess_block_size) {
		err = read_block(source, &block->block_data, offset, block_size);
		if (err < 0)
			goto done;
	}

	/*
	 * Verify that the block size covers at least the table header, block
	 * header and the 2 byte restart counter.
	 */
	if (block_size < header_size + 4 + 2) {
		err = REFTABLE_FORMAT_ERROR;
		goto done;
	}

	if (block_type == REFTABLE_BLOCK_TYPE_LOG) {
		uint32_t block_header_skip = 4 + header_size;
		uLong dst_len = block_size - block_header_skip;
		uLong src_len = block->block_data.len - block_header_skip;

		/* Log blocks specify the *uncompressed* size in their header. */
		REFTABLE_ALLOC_GROW_OR_NULL(block->uncompressed_data, block_size,
					    block->uncompressed_cap);
		if (!block->uncompressed_data) {
			err = REFTABLE_OUT_OF_MEMORY_ERROR;
			goto done;
		}

		/* Copy over the block header verbatim. It's not compressed. */
		memcpy(block->uncompressed_data, block->block_data.data, block_header_skip);

		if (!block->zstream) {
			REFTABLE_CALLOC_ARRAY(block->zstream, 1);
			if (!block->zstream) {
				err = REFTABLE_OUT_OF_MEMORY_ERROR;
				goto done;
			}

			err = inflateInit(block->zstream);
		} else {
			err = inflateReset(block->zstream);
		}
		if (err != Z_OK) {
			err = REFTABLE_ZLIB_ERROR;
			goto done;
		}

		block->zstream->next_in = block->block_data.data + block_header_skip;
		block->zstream->avail_in = src_len;
		block->zstream->next_out = block->uncompressed_data + block_header_skip;
		block->zstream->avail_out = dst_len;

		/*
		 * We know both input as well as output size, and we know that
		 * the sizes should never be bigger than `uInt_MAX` because
		 * blocks can at most be 16MB large. We can thus use `Z_FINISH`
		 * here to instruct zlib to inflate the data in one go, which
		 * is more efficient than using `Z_NO_FLUSH`.
		 */
		err = inflate(block->zstream, Z_FINISH);
		if (err != Z_STREAM_END) {
			err = REFTABLE_ZLIB_ERROR;
			goto done;
		}
		err = 0;

		if (block->zstream->total_out + block_header_skip != block_size) {
			err = REFTABLE_FORMAT_ERROR;
			goto done;
		}

		/* We're done with the input data. */
		block_source_release_data(&block->block_data);
		block->block_data.data = block->uncompressed_data;
		block->block_data.len = block_size;
		full_block_size = src_len + block_header_skip - block->zstream->avail_in;
	} else if (full_block_size == 0) {
		full_block_size = block_size;
	} else if (block_size < full_block_size && block_size < block->block_data.len &&
		   block->block_data.data[block_size] != 0) {
		/* If the block is smaller than the full block size, it is
		   padded (data followed by '\0') or the next block is
		   unaligned. */
		full_block_size = block_size;
	}

	/*
	 * Ensure that we have sufficient data available now to satisfy the
	 * claimed block size.
	 */
	if (block_size > block->block_data.len) {
		err = REFTABLE_FORMAT_ERROR;
		goto done;
	}

	restart_count = reftable_get_be16(block->block_data.data + block_size - 2);
	restart_off = block_size - 2 - 3 * restart_count;
	if (restart_off < header_size + 4 || restart_off > block_size - 2) {
		err = REFTABLE_FORMAT_ERROR;
		goto done;
	}

	block->block_type = block_type;
	block->hash_size = hash_size;
	block->restart_off = restart_off;
	block->full_block_size = full_block_size;
	block->header_off = header_size;
	block->restart_count = restart_count;

	err = 0;

done:
	if (err < 0)
		reftable_block_release(block);
	return err;
}

void reftable_block_release(struct reftable_block *block)
{
	inflateEnd(block->zstream);
	reftable_free(block->zstream);
	reftable_free(block->uncompressed_data);
	block_source_release_data(&block->block_data);
	memset(block, 0, sizeof(*block));
}

uint8_t reftable_block_type(const struct reftable_block *b)
{
	return b->block_data.data[b->header_off];
}

int reftable_block_first_key(const struct reftable_block *block, struct reftable_buf *key)
{
	int off = block->header_off + 4, n;
	struct string_view in = {
		.buf = block->block_data.data + off,
		.len = block->restart_off - off,
	};
	uint8_t extra = 0;

	reftable_buf_reset(key);

	n = reftable_decode_key(key, &extra, in);
	if (n < 0)
		return n;
	if (!key->len)
		return REFTABLE_FORMAT_ERROR;

	return 0;
}

static uint32_t block_restart_offset(const struct reftable_block *b, size_t idx)
{
	return reftable_get_be24(b->block_data.data + b->restart_off + 3 * idx);
}

void block_iter_init(struct block_iter *it, const struct reftable_block *block)
{
	it->block = block;
	block_iter_seek_start(it);
}

void block_iter_seek_start(struct block_iter *it)
{
	reftable_buf_reset(&it->last_key);
	it->next_off = it->block->header_off + 4;
}

struct restart_needle_less_args {
	int error;
	struct reftable_buf needle;
	const struct reftable_block *block;
};

static int restart_needle_less(size_t idx, void *_args)
{
	struct restart_needle_less_args *args = _args;
	uint32_t off = block_restart_offset(args->block, idx);
	struct string_view in = {
		.buf = args->block->block_data.data + off,
		.len = args->block->restart_off - off,
	};
	uint64_t prefix_len, suffix_len;
	uint8_t extra;
	int n;

	/*
	 * The restart offset must point to a record, which is stored before
	 * the restart table. Verify that this is the case.
	 */
	if (off >= args->block->restart_off) {
		args->error = 1;
		return -1;
	}

	/*
	 * Records at restart points are stored without prefix compression, so
	 * there is no need to fully decode the record key here. This removes
	 * the need for allocating memory.
	 */
	n = reftable_decode_keylen(in, &prefix_len, &suffix_len, &extra);
	if (n < 0 || prefix_len) {
		args->error = 1;
		return -1;
	}

	string_view_consume(&in, n);
	if (suffix_len > in.len) {
		args->error = 1;
		return -1;
	}

	n = memcmp(args->needle.buf, in.buf,
		   args->needle.len < suffix_len ? args->needle.len : suffix_len);
	if (n)
		return n < 0;
	return args->needle.len < suffix_len;
}

int block_iter_next(struct block_iter *it, struct reftable_record *rec)
{
	struct string_view in = {
		.buf = (unsigned char *) it->block->block_data.data + it->next_off,
		.len = it->block->restart_off - it->next_off,
	};
	struct string_view start = in;
	uint8_t extra = 0;
	int n = 0;

	if (it->next_off >= it->block->restart_off)
		return 1;

	n = reftable_decode_key(&it->last_key, &extra, in);
	if (n < 0)
		return -1;
	if (!it->last_key.len)
		return REFTABLE_FORMAT_ERROR;

	string_view_consume(&in, n);
	n = reftable_record_decode(rec, it->last_key, extra, in, it->block->hash_size,
				   &it->scratch);
	if (n < 0)
		return -1;
	string_view_consume(&in, n);

	it->next_off += start.len - in.len;
	return 0;
}

void block_iter_reset(struct block_iter *it)
{
	reftable_buf_reset(&it->last_key);
	it->next_off = 0;
	it->block = NULL;
}

void block_iter_close(struct block_iter *it)
{
	reftable_buf_release(&it->last_key);
	reftable_buf_release(&it->scratch);
}

int block_iter_seek_key(struct block_iter *it, struct reftable_buf *want)
{
	struct restart_needle_less_args args = {
		.needle = *want,
		.block = it->block,
	};
	struct reftable_record rec;
	int err = 0;
	size_t i;

	err = reftable_record_init(&rec, reftable_block_type(it->block));
	if (err < 0)
		goto done;

	/*
	 * Perform a binary search over the block's restart points, which
	 * avoids doing a linear scan over the whole block. Like this, we
	 * identify the section of the block that should contain our key.
	 *
	 * Note that we explicitly search for the first restart point _greater_
	 * than the sought-after record, not _greater or equal_ to it. In case
	 * the sought-after record is located directly at the restart point we
	 * would otherwise start doing the linear search at the preceding
	 * restart point. While that works alright, we would end up scanning
	 * too many record.
	 */
	i = binsearch(it->block->restart_count, &restart_needle_less, &args);
	if (args.error) {
		err = REFTABLE_FORMAT_ERROR;
		goto done;
	}

	/*
	 * Now there are multiple cases:
	 *
	 *   - `i == 0`: The wanted record is smaller than the record found at
	 *     the first restart point. As the first restart point is the first
	 *     record in the block, our wanted record cannot be located in this
	 *     block at all. We still need to position the iterator so that the
	 *     next call to `block_iter_next()` will yield an end-of-iterator
	 *     signal.
	 *
	 *   - `i == restart_count`: The wanted record was not found at any of
	 *     the restart points. As there is no restart point at the end of
	 *     the section the record may thus be contained in the last block.
	 *
	 *   - `i > 0`: The wanted record must be contained in the section
	 *     before the found restart point. We thus do a linear search
	 *     starting from the preceding restart point.
	 */
	if (i > 0)
		it->next_off = block_restart_offset(it->block, i - 1);
	else
		it->next_off = it->block->header_off + 4;

	/*
	 * We're looking for the last entry less than the wanted key so that
	 * the next call to `block_reader_next()` would yield the wanted
	 * record. We thus don't want to position our iterator at the sought
	 * after record, but one before. To do so, we have to go one entry too
	 * far and then back up.
	 */
	while (1) {
		size_t prev_off = it->next_off;

		err = block_iter_next(it, &rec);
		if (err < 0)
			goto done;
		if (err > 0) {
			it->next_off = prev_off;
			err = 0;
			goto done;
		}

		err = reftable_record_key(&rec, &it->last_key);
		if (err < 0)
			goto done;

		/*
		 * Check whether the current key is greater or equal to the
		 * sought-after key. In case it is greater we know that the
		 * record does not exist in the block and can thus abort early.
		 * In case it is equal to the sought-after key we have found
		 * the desired record.
		 *
		 * Note that we store the next record's key record directly in
		 * `last_key` without restoring the key of the preceding record
		 * in case we need to go one record back. This is safe to do as
		 * `block_iter_next()` would return the ref whose key is equal
		 * to `last_key` now, and naturally all keys share a prefix
		 * with themselves.
		 */
		if (reftable_buf_cmp(&it->last_key, want) >= 0) {
			it->next_off = prev_off;
			goto done;
		}
	}

done:
	reftable_record_release(&rec);
	return err;
}

static int block_iter_seek_void(void *it, struct reftable_record *want)
{
	struct reftable_buf buf = REFTABLE_BUF_INIT;
	struct block_iter *bi = it;
	int err;

	if (bi->block->block_type != want->type)
		return REFTABLE_API_ERROR;

	err = reftable_record_key(want, &buf);
	if (err < 0)
		goto out;

	err = block_iter_seek_key(it, &buf);
	if (err < 0)
		goto out;

	err = 0;

out:
	reftable_buf_release(&buf);
	return err;
}

static int block_iter_next_void(void *it, struct reftable_record *rec)
{
	return block_iter_next(it, rec);
}

static void block_iter_close_void(void *it)
{
	block_iter_close(it);
}

static struct reftable_iterator_vtable block_iter_vtable = {
	.seek = &block_iter_seek_void,
	.next = &block_iter_next_void,
	.close = &block_iter_close_void,
};

int reftable_block_init_iterator(const struct reftable_block *b,
				 struct reftable_iterator *it)
{
	struct block_iter *bi;

	REFTABLE_CALLOC_ARRAY(bi, 1);
	block_iter_init(bi, b);

	assert(!it->ops);
	it->iter_arg = bi;
	it->ops = &block_iter_vtable;

	return 0;
}

void block_writer_release(struct block_writer *bw)
{
	deflateEnd(bw->zstream);
	REFTABLE_FREE_AND_NULL(bw->zstream);
	REFTABLE_FREE_AND_NULL(bw->restarts);
	REFTABLE_FREE_AND_NULL(bw->compressed);
	reftable_buf_release(&bw->scratch);
	reftable_buf_release(&bw->last_key);
	/* the block is not owned. */
}