/* $OpenBSD: tls_buffer.c,v 1.4 2022/11/10 18:06:37 jsing Exp $ */
/*
 * Copyright (c) 2018, 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 <stdlib.h>
#include <string.h>

#include "bytestring.h"
#include "tls_internal.h"

#define TLS_BUFFER_CAPACITY_LIMIT	(1024 * 1024)

struct tls_buffer {
	size_t capacity;
	size_t capacity_limit;
	uint8_t *data;
	size_t len;
	size_t offset;
};

static int tls_buffer_resize(struct tls_buffer *buf, size_t capacity);

struct tls_buffer *
tls_buffer_new(size_t init_size)
{
	struct tls_buffer *buf = NULL;

	if ((buf = calloc(1, sizeof(struct tls_buffer))) == NULL)
		goto err;

	buf->capacity_limit = TLS_BUFFER_CAPACITY_LIMIT;

	if (!tls_buffer_resize(buf, init_size))
		goto err;

	return buf;

 err:
	tls_buffer_free(buf);

	return NULL;
}

void
tls_buffer_clear(struct tls_buffer *buf)
{
	freezero(buf->data, buf->capacity);

	buf->data = NULL;
	buf->capacity = 0;
	buf->len = 0;
	buf->offset = 0;
}

void
tls_buffer_free(struct tls_buffer *buf)
{
	if (buf == NULL)
		return;

	tls_buffer_clear(buf);

	freezero(buf, sizeof(struct tls_buffer));
}

static int
tls_buffer_grow(struct tls_buffer *buf, size_t capacity)
{
	if (buf->capacity >= capacity)
		return 1;

	return tls_buffer_resize(buf, capacity);
}

static int
tls_buffer_resize(struct tls_buffer *buf, size_t capacity)
{
	uint8_t *data;

	/*
	 * XXX - Consider maintaining a minimum size and growing more
	 * intelligently (rather than exactly).
	 */
	if (buf->capacity == capacity)
		return 1;

	if (capacity > buf->capacity_limit)
		return 0;

	if ((data = recallocarray(buf->data, buf->capacity, capacity, 1)) == NULL)
		return 0;

	buf->data = data;
	buf->capacity = capacity;

	/* Ensure that len and offset are valid if capacity decreased. */
	if (buf->len > buf->capacity)
		buf->len = buf->capacity;
	if (buf->offset > buf->len)
		buf->offset = buf->len;

	return 1;
}

void
tls_buffer_set_capacity_limit(struct tls_buffer *buf, size_t limit)
{
	/*
	 * XXX - do we want to force a resize if this limit is less than current
	 * capacity... and what do we do with existing data? Force a clear?
	 */
	buf->capacity_limit = limit;
}

ssize_t
tls_buffer_extend(struct tls_buffer *buf, size_t len,
    tls_read_cb read_cb, void *cb_arg)
{
	ssize_t ret;

	if (len == buf->len)
		return buf->len;

	if (len < buf->len)
		return TLS_IO_FAILURE;

	if (!tls_buffer_resize(buf, len))
		return TLS_IO_FAILURE;

	for (;;) {
		if ((ret = read_cb(&buf->data[buf->len],
		    buf->capacity - buf->len, cb_arg)) <= 0)
			return ret;

		if (ret > buf->capacity - buf->len)
			return TLS_IO_FAILURE;

		buf->len += ret;

		if (buf->len == buf->capacity)
			return buf->len;
	}
}

size_t
tls_buffer_remaining(struct tls_buffer *buf)
{
	if (buf->offset > buf->len)
		return 0;

	return buf->len - buf->offset;
}

ssize_t
tls_buffer_read(struct tls_buffer *buf, uint8_t *rbuf, size_t n)
{
	if (buf->offset > buf->len)
		return TLS_IO_FAILURE;

	if (buf->offset == buf->len)
		return TLS_IO_WANT_POLLIN;

	if (n > buf->len - buf->offset)
		n = buf->len - buf->offset;

	memcpy(rbuf, &buf->data[buf->offset], n);

	buf->offset += n;

	return n;
}

ssize_t
tls_buffer_write(struct tls_buffer *buf, const uint8_t *wbuf, size_t n)
{
	if (buf->offset > buf->len)
		return TLS_IO_FAILURE;

	/*
	 * To avoid continually growing the buffer, pull data up to the
	 * start of the buffer. If all data has been read then we can simply
	 * reset, otherwise wait until we're going to save at least 4KB of
	 * memory to reduce overhead.
	 */
	if (buf->offset == buf->len) {
		buf->len = 0;
		buf->offset = 0;
	}
	if (buf->offset >= 4096) {
		memmove(buf->data, &buf->data[buf->offset],
		    buf->len - buf->offset);
		buf->len -= buf->offset;
		buf->offset = 0;
	}

	if (buf->len > SIZE_MAX - n)
		return TLS_IO_FAILURE;
	if (!tls_buffer_grow(buf, buf->len + n))
		return TLS_IO_FAILURE;

	memcpy(&buf->data[buf->len], wbuf, n);

	buf->len += n;

	return n;
}

int
tls_buffer_append(struct tls_buffer *buf, const uint8_t *wbuf, size_t n)
{
	return tls_buffer_write(buf, wbuf, n) == n;
}

int
tls_buffer_data(struct tls_buffer *buf, CBS *out_cbs)
{
	CBS cbs;

	CBS_init(&cbs, buf->data, buf->len);

	if (!CBS_skip(&cbs, buf->offset))
		return 0;

	CBS_dup(&cbs, out_cbs);

	return 1;
}

int
tls_buffer_finish(struct tls_buffer *buf, uint8_t **out, size_t *out_len)
{
	if (out == NULL || out_len == NULL)
		return 0;

	*out = buf->data;
	*out_len = buf->len;

	buf->data = NULL;
	buf->capacity = 0;
	buf->len = 0;
	buf->offset = 0;

	return 1;
}