#include "sleep.h"
#include "sleep_verify.h"

#include <errno.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#ifdef SLEEP_SALT
#include <sodium.h>
#endif


static uint16_t
be16(const uint8_t *p)
{
	return ((uint16_t)p[0] << 8) | p[1];
}


static uint32_t
be32(const uint8_t *p)
{
	return ((uint32_t)p[0] << 24) |
	       ((uint32_t)p[1] << 16) |
	       ((uint32_t)p[2] << 8) |
	       p[3];
}


static uint64_t
be64(const uint8_t *p)
{
	return ((uint64_t)p[0] << 56) |
	       ((uint64_t)p[1] << 48) |
	       ((uint64_t)p[2] << 40) |
	       ((uint64_t)p[3] << 32) |
	       ((uint64_t)p[4] << 24) |
	       ((uint64_t)p[5] << 16) |
	       ((uint64_t)p[6] << 8) |
	       p[7];
}


static int
load_file(const char *path,
	  uint8_t **buf,
	  size_t *len)
{
	FILE *f;
	long size;
	uint8_t *p;

	f = fopen(path, "rb");

	if (!f) {
		fprintf(stderr,
			"sleep-verify: %s: %s\n",
			path,
			strerror(errno));
		return -1;
	}

	if (fseek(f, 0, SEEK_END) != 0 ||
	    (size = ftell(f)) < 0 ||
	    fseek(f, 0, SEEK_SET) != 0) {
		fclose(f);
		return -1;
	}

	p = malloc(size ? (size_t)size : 1);

	if (!p) {
		fclose(f);
		errno = ENOMEM;
		return -1;
	}

	if (size &&
	    fread(p, 1, (size_t)size, f) != (size_t)size) {
		free(p);
		fclose(f);
		errno = EIO;
		return -1;
	}

	fclose(f);

	*buf = p;
	*len = (size_t)size;

	return 0;
}


static int
check_header(const uint8_t *buf,
	     size_t len,
	     uint32_t magic,
	     uint16_t entry_size,
	     const char *algo)
{
	size_t algo_len;

	if (len < SLEEP_HEADER_SIZE)
		return -1;

	if (be32(buf) != magic ||
	    buf[4] != 0 ||
	    be16(buf + 5) != entry_size)
		return -1;

	algo_len = buf[7];

	if (algo_len != strlen(algo))
		return -1;

	return memcmp(buf + 8, algo, algo_len) == 0 ? 0 : -1;
}


static int
verify_tree_structure(const sleep_tree_entry_t *tree,
		      size_t node_count)
{
	size_t i;

	for (i = 1; i < node_count; i += 2) {
		uint64_t left;
		uint64_t right;
		uint64_t left_len;
		uint64_t right_len;
		uint64_t length;
		uint8_t hash[32];

		if (!sleep_tree_left_child(i, &left) ||
		    !sleep_tree_right_child(i, &right))
			return -1;

		/*
		 * This parent belongs to a subtree that does not
		 * completely exist yet.
		 */
		if (right >= node_count) {
			static const uint8_t zero[40];

			if (memcmp(&tree[i],
				   zero,
				   sizeof(tree[i])) != 0) {
				fprintf(stderr,
					"  tree: nonzero incomplete node %zu\n",
					i);
				return -1;
			}

			continue;
		}

		left_len =
			be64((const uint8_t *)&tree[left].length_be);

		right_len =
			be64((const uint8_t *)&tree[right].length_be);

		if (UINT64_MAX - left_len < right_len)
			return -1;

		length = left_len + right_len;

		if (sleep_tree_hash_parent(tree[left].hash,
					   tree[right].hash,
					   length,
					   hash) != 0)
			return -1;

		if (memcmp(hash, tree[i].hash, 32) != 0 ||
		    be64((const uint8_t *)&tree[i].length_be) != length) {
			fprintf(stderr,
				"  tree: parent mismatch at node %zu\n",
				i);
			return -1;
		}
	}

	return 0;
}


static int
verify_data_leaves(const uint8_t *data,
		   size_t data_len,
		   const sleep_tree_entry_t *tree,
		   size_t node_count)
{
	size_t leaf_count = (node_count + 1) / 2;
	size_t offset = 0;
	size_t i;

	for (i = 0; i < leaf_count; i++) {
		size_t node = i * 2;
		uint64_t len =
			be64((const uint8_t *)&tree[node].length_be);
		uint8_t hash[32];

		if (len > data_len - offset) {
			fprintf(stderr,
				"  data: leaf %zu exceeds metadata.data\n",
				i);
			return -1;
		}

		if (sleep_tree_hash_block(data + offset,
					  (size_t)len,
					  hash) != 0)
			return -1;

		if (memcmp(hash, tree[node].hash, 32) != 0) {
			fprintf(stderr,
				"  data: leaf hash mismatch at %zu\n",
				i);
			return -1;
		}

		offset += (size_t)len;
	}

	if (offset != data_len) {
		fprintf(stderr,
			"  data: tree consumes %zu bytes, file is %zu\n",
			offset,
			data_len);
		return -1;
	}

	return 0;
}


#ifdef SLEEP_SALT

static int
verify_signatures(const uint8_t *signatures,
		  size_t signatures_len,
		  const uint8_t pubkey[32],
		  const sleep_tree_entry_t *tree,
		  size_t node_count,
		  size_t feed_length)
{
	size_t sequence;
	uint64_t roots[64];
	uint8_t message[1 + 64 * 48];
	uint8_t digest[32];

	if (signatures_len < SLEEP_HEADER_SIZE ||
	    (signatures_len - SLEEP_HEADER_SIZE) %
	    SLEEP_SIG_ENTRY_SIZE != 0)
		return -1;

	if ((signatures_len - SLEEP_HEADER_SIZE) /
	    SLEEP_SIG_ENTRY_SIZE != feed_length)
		return -1;

	for (sequence = 1;
	     sequence <= feed_length;
	     sequence++) {
		size_t root_count;
		size_t off = 0;
		size_t r;

		root_count =
			sleep_tree_full_roots(sequence,
					      roots,
					      64);

		message[off++] = SLEEP_HASH_ROOT;

		for (r = 0; r < root_count; r++) {
			uint64_t index = roots[r];
			uint64_t length;

			if (index >= node_count)
				return -1;

			memcpy(message + off,
			       tree[index].hash,
			       32);
			off += 32;

			message[off++] = (uint8_t)(index >> 56);
			message[off++] = (uint8_t)(index >> 48);
			message[off++] = (uint8_t)(index >> 40);
			message[off++] = (uint8_t)(index >> 32);
			message[off++] = (uint8_t)(index >> 24);
			message[off++] = (uint8_t)(index >> 16);
			message[off++] = (uint8_t)(index >> 8);
			message[off++] = (uint8_t)index;

			length =
				be64((const uint8_t *)
				     &tree[index].length_be);

			message[off++] = (uint8_t)(length >> 56);
			message[off++] = (uint8_t)(length >> 48);
			message[off++] = (uint8_t)(length >> 40);
			message[off++] = (uint8_t)(length >> 32);
			message[off++] = (uint8_t)(length >> 24);
			message[off++] = (uint8_t)(length >> 16);
			message[off++] = (uint8_t)(length >> 8);
			message[off++] = (uint8_t)length;
		}

		if (crypto_generichash(digest,
				       sizeof(digest),
				       message,
				       off,
				       NULL,
				       0) != 0)
			return -1;

		if (crypto_sign_verify_detached(
			signatures +
			SLEEP_HEADER_SIZE +
			(sequence - 1) * SLEEP_SIG_ENTRY_SIZE,
			digest,
			sizeof(digest),
			pubkey) != 0) {
			fprintf(stderr,
				"  signatures: failed at sequence %zu\n",
				sequence);
			return -1;
		}
	}

	return 0;
}

#endif


static int
verify_feed(const char *dir,
	    const char *name,
	    int verify_data,
	    size_t *out_length)
{
	char path[4096];

	uint8_t *tree_file = NULL;
	uint8_t *sig_file = NULL;
	uint8_t *key_file = NULL;
	uint8_t *data_file = NULL;
	uint8_t *bitfield = NULL;

	size_t tree_len = 0;
	size_t sig_len = 0;
	size_t key_len = 0;
	size_t data_len = 0;
	size_t bitfield_len = 0;

	size_t node_count;
	size_t feed_length;

	int rc = -1;

	snprintf(path,
		 sizeof(path),
		 "%s/%s.tree",
		 dir,
		 name);

	if (load_file(path, &tree_file, &tree_len) != 0)
		goto done;

	if (check_header(tree_file,
			 tree_len,
			 SLEEP_MAGIC_TREE,
			 SLEEP_TREE_ENTRY_SIZE,
			 SLEEP_ALGO_BLAKE2B) != 0) {
		fprintf(stderr, "  tree: invalid header\n");
		goto done;
	}

	if ((tree_len - SLEEP_HEADER_SIZE) %
	    sizeof(sleep_tree_entry_t) != 0)
		goto done;

	node_count =
		(tree_len - SLEEP_HEADER_SIZE) /
		sizeof(sleep_tree_entry_t);

	if (!node_count || !(node_count & 1))
		goto done;

	feed_length = (node_count + 1) / 2;

	if (verify_tree_structure(
		(const sleep_tree_entry_t *)(tree_file +
					     SLEEP_HEADER_SIZE),
		node_count) != 0)
		goto done;

	snprintf(path,
		 sizeof(path),
		 "%s/%s.key",
		 dir,
		 name);

	if (load_file(path, &key_file, &key_len) != 0 ||
	    key_len != 32)
		goto done;

	snprintf(path,
		 sizeof(path),
		 "%s/%s.signatures",
		 dir,
		 name);

	if (load_file(path, &sig_file, &sig_len) != 0)
		goto done;

	if (check_header(sig_file,
			 sig_len,
			 SLEEP_MAGIC_SIGNATURE,
			 SLEEP_SIG_ENTRY_SIZE,
			 SLEEP_ALGO_ED25519) != 0)
		goto done;

#ifdef SLEEP_SALT

	if (verify_signatures(
		sig_file,
		sig_len,
		key_file,
		(const sleep_tree_entry_t *)(tree_file +
					     SLEEP_HEADER_SIZE),
		node_count,
		feed_length) != 0)
		goto done;

#else

	fprintf(stderr,
		"  signatures: build with SLEEP_SALT\n");
	goto done;

#endif

	snprintf(path,
		 sizeof(path),
		 "%s/%s.bitfield",
		 dir,
		 name);

	if (load_file(path,
		      &bitfield,
		      &bitfield_len) != 0)
		goto done;

	if (bitfield_len < SLEEP_HEADER_SIZE ||
	    be32(bitfield) != SLEEP_MAGIC_BITFIELD ||
	    bitfield[4] != 0)
		goto done;

	/*
	 * The SLEEP V2 document says 3328 here.
	 *
	 * The supplied historical Dat fixture says 3584.
	 * We deliberately accept both until the extra 256 bytes
	 * in the historical format have been identified.
	 */
	if (be16(bitfield + 5) != SLEEP_BITFIELD_ENTRY_SIZE &&
	    be16(bitfield + 5) != 3584)
		goto done;

	if (verify_data) {
		snprintf(path,
			 sizeof(path),
			 "%s/%s.data",
			 dir,
			 name);

		if (load_file(path,
			      &data_file,
			      &data_len) != 0)
			goto done;

		if (verify_data_leaves(
			data_file,
			data_len,
			(const sleep_tree_entry_t *)(tree_file +
						     SLEEP_HEADER_SIZE),
			node_count) != 0)
			goto done;
	}

	if (out_length)
		*out_length = feed_length;

	rc = 0;

done:
	free(tree_file);
	free(sig_file);
	free(key_file);
	free(data_file);
	free(bitfield);

	return rc;
}


int
sleep_verify_dat(const char *dir)
{
	size_t metadata_length;
	size_t content_length;

	printf("sleep-verify: %s\n", dir);

	printf("metadata: ");

	if (verify_feed(dir,
			"metadata",
			1,
			&metadata_length) != 0) {
		printf("FAIL\n");
		return -1;
	}

	printf("OK (%zu entries)\n", metadata_length);

	printf("content:  ");

	if (verify_feed(dir,
			"content",
			0,
			&content_length) != 0) {
		printf("FAIL\n");
		return -1;
	}

	printf("OK (%zu entries)\n", content_length);

	printf("sleep-verify: PASS\n");

	return 0;
}


int
main(int argc, char **argv)
{
	if (argc != 2) {
		fprintf(stderr,
			"usage: %s .dat\n",
			argv[0]);
		return 2;
	}

#ifdef SLEEP_SALT

	if (sodium_init() < 0)
		return 1;

#else

	fprintf(stderr,
		"sleep-verify: built without SLEEP_SALT\n");

	return 1;

#endif

	return sleep_verify_dat(argv[1]) == 0 ? 0 : 1;
}