/* $OpenBSD: buffertest.c,v 1.6 2022/07/22 19:34:55 jsing Exp $ */
/*
 * Copyright (c) 2019, 2022 Joel Sing <jsing@openbsd.org>
 *
 * Permission to use, copy, modify, and distribute this software for any
 * purpose with or without fee is hereby granted, provided that the above
 * copyright notice and this permission notice appear in all copies.
 *
 * THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
 * WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
 * MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
 * ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
 * WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
 * ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
 * OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
 */

#include <err.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include "tls_internal.h"

uint8_t testdata[] = {
	0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07,
	0x08, 0x09, 0x0a, 0x0b, 0x0c, 0x0d, 0x0e, 0x0f,
};

struct read_state {
	uint8_t *buf;
	size_t len;
	size_t offset;
};

static ssize_t
read_cb(void *buf, size_t buflen, void *cb_arg)
{
	struct read_state *rs = cb_arg;
	ssize_t n;

	if (rs->offset > rs->len)
		return TLS_IO_EOF;

	if ((size_t)(n = buflen) > (rs->len - rs->offset))
		n = rs->len - rs->offset;

	if (n == 0)
		return TLS_IO_WANT_POLLIN;

	memcpy(buf, &rs->buf[rs->offset], n);
	rs->offset += n;

	return n;
}

struct extend_test {
	size_t extend_len;
	size_t read_len;
	ssize_t want_ret;
};

const struct extend_test extend_tests[] = {
	{
		.extend_len = 4,
		.read_len = 0,
		.want_ret = TLS_IO_WANT_POLLIN,
	},
	{
		.extend_len = 4,
		.read_len = 8,
		.want_ret = 4,
	},
	{
		.extend_len = 12,
		.read_len = 8,
		.want_ret = TLS_IO_WANT_POLLIN,
	},
	{
		.extend_len = 12,
		.read_len = 10,
		.want_ret = TLS_IO_WANT_POLLIN,
	},
	{
		.extend_len = 12,
		.read_len = 12,
		.want_ret = 12,
	},
	{
		.extend_len = 16,
		.read_len = 16,
		.want_ret = 16,
	},
	{
		.extend_len = 20,
		.read_len = 1,
		.want_ret = TLS_IO_EOF,
	},
};

#define N_EXTEND_TESTS (sizeof(extend_tests) / sizeof(extend_tests[0]))

static int
tls_buffer_extend_test(void)
{
	const struct extend_test *et;
	struct tls_buffer *buf;
	struct read_state rs;
	uint8_t *data = NULL;
	size_t i, data_len;
	ssize_t ret;
	CBS cbs;
	int failed = 1;

	rs.buf = testdata;
	rs.offset = 0;

	if ((buf = tls_buffer_new(0)) == NULL)
		errx(1, "tls_buffer_new");

	for (i = 0; i < N_EXTEND_TESTS; i++) {
		et = &extend_tests[i];
		rs.len = et->read_len;

		ret = tls_buffer_extend(buf, et->extend_len, read_cb, &rs);
		if (ret != extend_tests[i].want_ret) {
			fprintf(stderr, "FAIL: Test %zd - extend returned %zd, "
			    "want %zd\n", i, ret, et->want_ret);
			goto failed;
		}

		if (!tls_buffer_data(buf, &cbs)) {
			fprintf(stderr, "FAIL: Test %zd - failed to get data\n",
			    i);
			goto failed;
		}

		if (!CBS_mem_equal(&cbs, testdata, CBS_len(&cbs))) {
			fprintf(stderr, "FAIL: Test %zd - extend buffer "
			    "mismatch", i);
			goto failed;
		}
	}

	if (!tls_buffer_finish(buf, &data, &data_len)) {
		fprintf(stderr, "FAIL: failed to finish\n");
		goto failed;
	}

	tls_buffer_free(buf);
	buf = NULL;

	if (data_len != sizeof(testdata)) {
		fprintf(stderr, "FAIL: got data length %zu, want %zu\n",
		    data_len, sizeof(testdata));
		goto failed;
	}
	if (memcmp(data, testdata, data_len) != 0) {
		fprintf(stderr, "FAIL: data mismatch\n");
		goto failed;
	}

	failed = 0;

 failed:
	tls_buffer_free(buf);
	free(data);

	return failed;
}

struct read_write_test {
	uint8_t pattern;
	size_t read;
	size_t write;
	size_t append;
	ssize_t want;
};

const struct read_write_test read_write_tests[] = {
	{
		.read = 2048,
		.want = TLS_IO_WANT_POLLIN,
	},
	{
		.pattern = 0xdb,
		.write = 2048,
		.want = 2048,
	},
	{
		.pattern = 0xbd,
		.append = 2048,
		.want = 1,
	},
	{
		.pattern = 0xdb,
		.read = 2048,
		.want = 2048,
	},
	{
		.pattern = 0xfe,
		.append = 1024,
		.want = 1,
	},
	{
		.pattern = 0xbd,
		.read = 1000,
		.want = 1000,
	},
	{
		.pattern = 0xbd,
		.read = 1048,
		.want = 1048,
	},
	{
		.pattern = 0xdb,
		.write = 2048,
		.want = 2048,
	},
	{
		.pattern = 0xbd,
		.append = 1024,
		.want = 1,
	},
	{
		.pattern = 0xee,
		.append = 4096,
		.want = 1,
	},
	{
		.pattern = 0xfe,
		.append = 1,
		.want = 0,
	},
	{
		.pattern = 0xfe,
		.write = 1,
		.want = TLS_IO_FAILURE,
	},
	{
		.pattern = 0xfe,
		.read = 1024,
		.want = 1024,
	},
	{
		.pattern = 0xdb,
		.read = 2048,
		.want = 2048,
	},
	{
		.pattern = 0xbd,
		.read = 1024,
		.want = 1024,
	},
	{
		.pattern = 0xee,
		.read = 1024,
		.want = 1024,
	},
	{
		.pattern = 0xee,
		.read = 4096,
		.want = 3072,
	},
	{
		.read = 2048,
		.want = TLS_IO_WANT_POLLIN,
	},
};

#define N_READ_WRITE_TESTS (sizeof(read_write_tests) / sizeof(read_write_tests[0]))

static int
tls_buffer_read_write_test(void)
{
	const struct read_write_test *rwt;
	struct tls_buffer *buf = NULL;
	uint8_t *rbuf = NULL, *wbuf = NULL;
	ssize_t n;
	size_t i;
	int ret;
	int failed = 1;

	if ((buf = tls_buffer_new(0)) == NULL)
		errx(1, "tls_buffer_new");

	tls_buffer_set_capacity_limit(buf, 8192);

	for (i = 0; i < N_READ_WRITE_TESTS; i++) {
		rwt = &read_write_tests[i];

		if (rwt->append > 0) {
			free(wbuf);
			if ((wbuf = malloc(rwt->append)) == NULL)
				errx(1, "malloc");
			memset(wbuf, rwt->pattern, rwt->append);
			if ((ret = tls_buffer_append(buf, wbuf, rwt->append)) !=
			    rwt->want) {
				fprintf(stderr, "FAIL: test %zu - "
				    "tls_buffer_append() = %d, want %zu\n",
				    i, ret, rwt->want);
				goto failed;
			}
		}

		if (rwt->write > 0) {
			free(wbuf);
			if ((wbuf = malloc(rwt->write)) == NULL)
				errx(1, "malloc");
			memset(wbuf, rwt->pattern, rwt->write);
			if ((n = tls_buffer_write(buf, wbuf, rwt->write)) !=
			    rwt->want) {
				fprintf(stderr, "FAIL: test %zu - "
				    "tls_buffer_write() = %zi, want %zu\n",
				    i, n, rwt->want);
				goto failed;
			}
		}

		if (rwt->read > 0) {
			free(rbuf);
			if ((rbuf = calloc(1, rwt->read)) == NULL)
				errx(1, "malloc");
			if ((n = tls_buffer_read(buf, rbuf, rwt->read)) !=
			    rwt->want) {
				fprintf(stderr, "FAIL: test %zu - "
				    "tls_buffer_read() = %zi, want %zu\n",
				    i, n, rwt->want);
				goto failed;
			}
			if (rwt->want > 0) {
				free(wbuf);
				if ((wbuf = malloc(rwt->want)) == NULL)
					errx(1, "malloc");
				memset(wbuf, rwt->pattern, rwt->want);
				if (memcmp(rbuf, wbuf, rwt->want) != 0) {
					fprintf(stderr, "FAIL: test %zu - "
					    "read byte mismatch\n", i);
					goto failed;
				}
			}
		}
	}

	failed = 0;

 failed:
	tls_buffer_free(buf);
	free(rbuf);
	free(wbuf);

	return failed;
}

int
main(int argc, char **argv)
{
	int failed = 0;

	failed |= tls_buffer_extend_test();
	failed |= tls_buffer_read_write_test();

	return failed;
}