/*
 * Copyright (C) 2017 Red Hat, Inc.
 *
 * Author: Nikos Mavrogiannopoulos
 *
 * This file is part of GnuTLS.
 *
 * GnuTLS is free software; you can redistribute it and/or modify it
 * under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 3 of the License, or
 * (at your option) any later version.
 *
 * GnuTLS is distributed in the hope that it will be useful, but
 * WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public License
 * along with this program.  If not, see <https://www.gnu.org/licenses/>
 */

#include "utils.h"

#define TLS_EXT_SUPPORTED_VERSIONS 43
#define TLS_EXT_POST_HANDSHAKE 49

#define SKIP16(pos, _total)                                                   \
	{                                                                     \
		uint16_t _s;                                                  \
		if ((size_t)pos + 2 > (size_t)_total)                         \
			fail("error0: at %d total: %d\n", pos + 2, _total);   \
		_s = (msg->data[pos] << 8) | msg->data[pos + 1];              \
		if ((size_t)(pos + 2 + _s) > (size_t)_total)                  \
			fail("error1: at %d field: %d, total: %d\n", pos + 2, \
			     (int)_s, _total);                                \
		pos += 2 + _s;                                                \
	}

#define SKIP8(pos, _total)                                   \
	{                                                    \
		uint8_t _s;                                  \
		if ((size_t)pos + 1 > (size_t)_total)        \
			fail("error\n");                     \
		_s = msg->data[pos];                         \
		if ((size_t)(pos + 1 + _s) > (size_t)_total) \
			fail("error\n");                     \
		pos += 1 + _s;                               \
	}

typedef void (*ext_parse_func)(void *priv, gnutls_datum_t *extdata);

#define HANDSHAKE_SESSION_ID_POS 34

#if defined __clang__ || __GNUC__ > 4 || (__GNUC__ == 4 && __GNUC_MINOR__ >= 5)
#pragma GCC diagnostic push
#pragma GCC diagnostic ignored "-Wunused-function"
#endif

/* Returns 0 if the extension was not found, 1 otherwise.
 */
static unsigned find_client_extension(const gnutls_datum_t *msg, unsigned extnr,
				      void *priv, ext_parse_func cb)
{
	unsigned pos;

	if (msg->size < HANDSHAKE_SESSION_ID_POS)
		fail("invalid client hello\n");

	/* we expect the legacy version to be present */
	/* ProtocolVersion legacy_version = 0x0303 */
	if (msg->data[0] != 0x03) {
		fail("ProtocolVersion contains %d.%d\n", (int)msg->data[0],
		     (int)msg->data[1]);
	}

	pos = HANDSHAKE_SESSION_ID_POS;
	/* legacy_session_id */
	SKIP8(pos, msg->size);

	/* CipherSuites */
	SKIP16(pos, msg->size);

	/* legacy_compression_methods */
	SKIP8(pos, msg->size);

	pos += 2;

	while (pos < msg->size) {
		uint16_t type;

		if (pos + 4 > msg->size)
			fail("invalid client hello\n");

		type = (msg->data[pos] << 8) | msg->data[pos + 1];
		pos += 2;

		if (debug)
			success("Found client extension %d\n", (int)type);

		if (type != extnr) {
			SKIP16(pos, msg->size);
		} else { /* found */
			ssize_t size = (msg->data[pos] << 8) |
				       msg->data[pos + 1];
			gnutls_datum_t data;

			pos += 2;
			if (pos + size > msg->size) {
				fail("error in extension length (pos: %d, ext: %d, total: %d)\n",
				     pos, (int)size, msg->size);
			}
			data.data = &msg->data[pos];
			data.size = size;
			if (cb)
				cb(priv, &data);
			return 1;
		}
	}
	return 0;
}

static unsigned is_client_extension_last(const gnutls_datum_t *msg,
					 unsigned extnr)
{
	unsigned pos, found = 0;

	if (msg->size < HANDSHAKE_SESSION_ID_POS)
		fail("invalid client hello\n");

	/* we expect the legacy version to be present */
	/* ProtocolVersion legacy_version = 0x0303 */
	if (msg->data[0] != 0x03) {
		fail("ProtocolVersion contains %d.%d\n", (int)msg->data[0],
		     (int)msg->data[1]);
	}

	pos = HANDSHAKE_SESSION_ID_POS;
	/* legacy_session_id */
	SKIP8(pos, msg->size);

	/* CipherSuites */
	SKIP16(pos, msg->size);

	/* legacy_compression_methods */
	SKIP8(pos, msg->size);

	pos += 2;

	while (pos < msg->size) {
		uint16_t type;

		if (pos + 4 > msg->size)
			fail("invalid client hello\n");

		type = (msg->data[pos] << 8) | msg->data[pos + 1];
		pos += 2;

		if (debug)
			success("Found client extension %d\n", (int)type);

		if (type != extnr) {
			if (found) {
				success("found extension %d after %d\n", type,
					extnr);
				return 0;
			}
			SKIP16(pos, msg->size);
		} else { /* found */
			found = 1;
			SKIP16(pos, msg->size);
		}
	}

	if (found)
		return 1;
	return 0;
}

#define TLS_RANDOM_SIZE 32

static unsigned find_server_extension(const gnutls_datum_t *msg, unsigned extnr,
				      void *priv, ext_parse_func cb)
{
	unsigned pos = 0;

	success("server hello of %d bytes\n", msg->size);
	/* we expect the legacy version to be present */
	/* ProtocolVersion legacy_version = 0x0303 */
	if (msg->data[0] != 0x03 || msg->data[1] != 0x03) {
		fail("ProtocolVersion contains %d.%d\n", (int)msg->data[0],
		     (int)msg->data[1]);
	}

	if (msg->data[1] >= 0x04) {
		success("assuming TLS 1.3 or better hello format (seen %d.%d)\n",
			(int)msg->data[0], (int)msg->data[1]);
	}

	pos += 2 + TLS_RANDOM_SIZE;

	/* legacy_session_id */
	SKIP8(pos, msg->size);

	/* CipherSuite */
	pos += 2;

	/* legacy_compression_methods */
	SKIP8(pos, msg->size);

	pos += 2;

	while (pos < msg->size) {
		uint16_t type;

		if (pos + 4 > msg->size)
			fail("invalid server hello\n");

		type = (msg->data[pos] << 8) | msg->data[pos + 1];
		pos += 2;

		success("Found server extension %d\n", (int)type);

		if (type != extnr) {
			SKIP16(pos, msg->size);
		} else { /* found */
			ssize_t size = (msg->data[pos] << 8) |
				       msg->data[pos + 1];
			gnutls_datum_t data;

			pos += 2;
			if (pos + size < msg->size) {
				fail("error in server extension length (pos: %d, total: %d)\n",
				     pos, msg->size);
			}
			data.data = &msg->data[pos];
			data.size = size;
			if (cb)
				cb(priv, &data);
			return 1;
		}
	}

	return 0;
}

#if defined __clang__ || __GNUC__ > 4 || (__GNUC__ == 4 && __GNUC_MINOR__ >= 5)
#pragma GCC diagnostic pop
#endif
