/*

 * Copyright (C) 2012,2013 Free Software Foundation, Inc.

 * Copyright (C) 2013 Nikos Mavrogiannopoulos

 *

 * Author: Nikos Mavrogiannopoulos

 *

 * This file is part of GnuTLS.

 *

 * The GnuTLS is free software; you can redistribute it and/or

 * modify it under the terms of the GNU Lesser General Public License

 * as published by the Free Software Foundation; either version 2.1 of

 * the License, or (at your option) any later version.

 *

 * This library 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

 * Lesser 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/>

 *

 */

/* This file implements the TLS heartbeat extension.

 */

#include "errors.h"

#include "gnutls_int.h"

#include <dtls.h>

#include <record.h>

#include <ext/heartbeat.h>

#include <hello_ext.h>

#include <random.h>

#ifdef ENABLE_HEARTBEAT

/**

  * gnutls_heartbeat_enable:

  * @session: is a #gnutls_session_t type.

  * @type: one of the GNUTLS_HB_* flags

  *

  * If this function is called with the %GNUTLS_HB_PEER_ALLOWED_TO_SEND

  * @type, GnuTLS will allow heartbeat messages to be received. Moreover it also

  * request the peer to accept heartbeat messages. This function

  * must be called prior to TLS handshake.

  *

  * If the @type used is %GNUTLS_HB_LOCAL_ALLOWED_TO_SEND, then the peer

  * will be asked to accept heartbeat messages but not send ones.

  *

  * The function gnutls_heartbeat_allowed() can be used to test Whether

  * locally generated heartbeat messages can be accepted by the peer.

  *

  * Since: 3.1.2

  **/

void gnutls_heartbeat_enable(gnutls_session_t session, unsigned int type)

{

	gnutls_ext_priv_data_t epriv;

	epriv = (void*)(intptr_t)type;

	_gnutls_hello_ext_set_priv(session, GNUTLS_EXTENSION_HEARTBEAT,

				     epriv);

}

/**

  * gnutls_heartbeat_allowed:

  * @session: is a #gnutls_session_t type.

  * @type: one of %GNUTLS_HB_LOCAL_ALLOWED_TO_SEND and %GNUTLS_HB_PEER_ALLOWED_TO_SEND

  *

  * This function will check whether heartbeats are allowed

  * to be sent or received in this session. 

  *

  * Returns: Non zero if heartbeats are allowed.

  *

  * Since: 3.1.2

  **/

unsigned gnutls_heartbeat_allowed(gnutls_session_t session, unsigned int type)

{

	gnutls_ext_priv_data_t epriv;

	if (session->internals.handshake_in_progress != 0)

		return 0; /* not allowed */

	if (_gnutls_hello_ext_get_priv

	    (session, GNUTLS_EXTENSION_HEARTBEAT, &epriv) < 0)

		return 0;	/* Not enabled */

	if (type == GNUTLS_HB_LOCAL_ALLOWED_TO_SEND) {

		if (((intptr_t)epriv) & LOCAL_ALLOWED_TO_SEND)

			return 1;

	} else if (((intptr_t)epriv) & GNUTLS_HB_PEER_ALLOWED_TO_SEND)

		return 1;

	return 0;

}

#define DEFAULT_PADDING_SIZE 16

/*

 * Sends heartbeat data.

 */

static int

heartbeat_send_data(gnutls_session_t session, const void *data,

		    size_t data_size, uint8_t type)

{

	int ret, pos;

	uint8_t *response;

	response = gnutls_malloc(1 + 2 + data_size + DEFAULT_PADDING_SIZE);

	if (response == NULL)

		return gnutls_assert_val(GNUTLS_E_MEMORY_ERROR);

	pos = 0;

	response[pos++] = type;

	_gnutls_write_uint16(data_size, &response[pos]);

	pos += 2;

	memcpy(&response[pos], data, data_size);

	pos += data_size;

	ret =

	    gnutls_rnd(GNUTLS_RND_NONCE, &response[pos],

		       DEFAULT_PADDING_SIZE);

	if (ret < 0) {

		gnutls_assert();

		goto cleanup;

	}

	pos += DEFAULT_PADDING_SIZE;

	ret =

	    _gnutls_send_int(session, GNUTLS_HEARTBEAT, -1,

			     EPOCH_WRITE_CURRENT, response, pos,

			     MBUFFER_FLUSH);

      cleanup:

	gnutls_free(response);

	return ret;

}

/**

 * gnutls_heartbeat_ping:

 * @session: is a #gnutls_session_t type.

 * @data_size: is the length of the ping payload.

 * @max_tries: if flags is %GNUTLS_HEARTBEAT_WAIT then this sets the number of retransmissions. Use zero for indefinite (until timeout).

 * @flags: if %GNUTLS_HEARTBEAT_WAIT then wait for pong or timeout instead of returning immediately.

 *

 * This function sends a ping to the peer. If the @flags is set

 * to %GNUTLS_HEARTBEAT_WAIT then it waits for a reply from the peer.

 * 

 * Note that it is highly recommended to use this function with the

 * flag %GNUTLS_HEARTBEAT_WAIT, or you need to handle retransmissions

 * and timeouts manually.

 *

 * The total TLS data transmitted as part of the ping message are given by

 * the following formula: MAX(16, @data_size)+gnutls_record_overhead_size()+3.

 *

 * Returns: %GNUTLS_E_SUCCESS on success, otherwise a negative error code.

 *

 * Since: 3.1.2

 **/

int

gnutls_heartbeat_ping(gnutls_session_t session, size_t data_size,

		      unsigned int max_tries, unsigned int flags)

{

	int ret;

	unsigned int retries = 1, diff;

	struct timespec now;

	if (data_size > MAX_HEARTBEAT_LENGTH)

		return

		    gnutls_assert_val(GNUTLS_E_UNEXPECTED_PACKET_LENGTH);

	if (gnutls_heartbeat_allowed

	    (session, GNUTLS_HB_LOCAL_ALLOWED_TO_SEND) == 0)

		return gnutls_assert_val(GNUTLS_E_INVALID_REQUEST);

	/* resume previous call if interrupted */

	if (session->internals.record_send_buffer.byte_length > 0 &&

	    session->internals.record_send_buffer.head != NULL &&

	    session->internals.record_send_buffer.head->type ==

	    GNUTLS_HEARTBEAT)

		return _gnutls_io_write_flush(session);

	switch (session->internals.hb_state) {

	case SHB_SEND1:

		if (data_size > DEFAULT_PADDING_SIZE)

			data_size -= DEFAULT_PADDING_SIZE;

		else

			data_size = 0;

		_gnutls_buffer_reset(&session->internals.hb_local_data);

		ret =

		    _gnutls_buffer_resize(&session->internals.

					  hb_local_data, data_size);

		if (ret < 0)

			return gnutls_assert_val(ret);

		ret =

		    gnutls_rnd(GNUTLS_RND_NONCE,

				session->internals.hb_local_data.data,

				data_size);

		if (ret < 0)

			return gnutls_assert_val(ret);

		gnutls_gettime(&session->internals.hb_ping_start);

		session->internals.hb_local_data.length = data_size;

		session->internals.hb_state = SHB_SEND2;

		FALLTHROUGH;

	case SHB_SEND2:

		session->internals.hb_actual_retrans_timeout_ms =

		    session->internals.hb_retrans_timeout_ms;

	      retry:

		ret =

		    heartbeat_send_data(session,

					session->internals.hb_local_data.

					data,

					session->internals.hb_local_data.

					length, HEARTBEAT_REQUEST);

		if (ret < 0)

			return gnutls_assert_val(ret);

		gnutls_gettime(&session->internals.hb_ping_sent);

		if (!(flags & GNUTLS_HEARTBEAT_WAIT)) {

			session->internals.hb_state = SHB_SEND1;

			break;

		}

		session->internals.hb_state = SHB_RECV;

		FALLTHROUGH;

	case SHB_RECV:

		ret =

		    _gnutls_recv_int(session, GNUTLS_HEARTBEAT,

				     NULL, 0, NULL,

				     session->internals.

				     hb_actual_retrans_timeout_ms);

		if (ret == GNUTLS_E_HEARTBEAT_PONG_RECEIVED) {

			session->internals.hb_state = SHB_SEND1;

			break;

		} else if (ret == GNUTLS_E_TIMEDOUT) {

			retries++;

			if (max_tries > 0 && retries > max_tries) {

				session->internals.hb_state = SHB_SEND1;

				return gnutls_assert_val(ret);

			}

			gnutls_gettime(&now;;

			diff =

			    timespec_sub_ms(&now,

					    &session->internals.

					    hb_ping_start);

			if (diff > session->internals.hb_total_timeout_ms) {

				session->internals.hb_state = SHB_SEND1;

				return

				    gnutls_assert_val(GNUTLS_E_TIMEDOUT);

			}

			session->internals.hb_actual_retrans_timeout_ms *=

			    2;

			session->internals.hb_actual_retrans_timeout_ms %=

			    MAX_DTLS_TIMEOUT;

			session->internals.hb_state = SHB_SEND2;

			goto retry;

		} else if (ret < 0) {

			session->internals.hb_state = SHB_SEND1;

			return gnutls_assert_val(ret);

		}

	}

	return 0;

}

/**

 * gnutls_heartbeat_pong:

 * @session: is a #gnutls_session_t type.

 * @flags: should be zero

 *

 * This function replies to a ping by sending a pong to the peer.

 *

 * Returns: %GNUTLS_E_SUCCESS on success, otherwise a negative error code.

 *

 * Since: 3.1.2

 **/

int gnutls_heartbeat_pong(gnutls_session_t session, unsigned int flags)

{

	int ret;

	if (session->internals.record_send_buffer.byte_length > 0 &&

	    session->internals.record_send_buffer.head != NULL &&

	    session->internals.record_send_buffer.head->type ==

	    GNUTLS_HEARTBEAT)

		return _gnutls_io_write_flush(session);

	if (session->internals.hb_remote_data.length == 0)

		return gnutls_assert_val(GNUTLS_E_INVALID_REQUEST);

	ret =

	    heartbeat_send_data(session,

				session->internals.hb_remote_data.data,

				session->internals.hb_remote_data.length,

				HEARTBEAT_RESPONSE);

	_gnutls_buffer_reset(&session->internals.hb_remote_data);

	if (ret < 0)

		return gnutls_assert_val(ret);

	return 0;

}

/*

 * Processes a heartbeat message. 

 */

int _gnutls_heartbeat_handle(gnutls_session_t session, mbuffer_st * bufel)

{

	int ret;

	unsigned type;

	unsigned pos;

	uint8_t *msg = _mbuffer_get_udata_ptr(bufel);

	size_t hb_len, len = _mbuffer_get_udata_size(bufel);

	if (gnutls_heartbeat_allowed

	    (session, GNUTLS_HB_PEER_ALLOWED_TO_SEND) == 0)

		return gnutls_assert_val(GNUTLS_E_UNEXPECTED_PACKET);

	if (len < 3 + DEFAULT_PADDING_SIZE)

		return

		    gnutls_assert_val(GNUTLS_E_UNEXPECTED_PACKET_LENGTH);

	pos = 0;

	type = msg[pos++];

	hb_len = _gnutls_read_uint16(&msg[pos]);

	if (hb_len > len - 3 - DEFAULT_PADDING_SIZE)

		return

		    gnutls_assert_val(GNUTLS_E_UNEXPECTED_PACKET_LENGTH);

	pos += 2;

	switch (type) {

	case HEARTBEAT_REQUEST:

		_gnutls_buffer_reset(&session->internals.hb_remote_data);

		ret =

		    _gnutls_buffer_resize(&session->internals.

					  hb_remote_data, hb_len);

		if (ret < 0)

			return gnutls_assert_val(ret);

		if (hb_len > 0)

			memcpy(session->internals.hb_remote_data.data,

			       &msg[pos], hb_len);

		session->internals.hb_remote_data.length = hb_len;

		return gnutls_assert_val(GNUTLS_E_HEARTBEAT_PING_RECEIVED);

	case HEARTBEAT_RESPONSE:

		if (hb_len != session->internals.hb_local_data.length)

			return

			    gnutls_assert_val(GNUTLS_E_UNEXPECTED_PACKET);

		if (hb_len > 0 &&

		    memcmp(&msg[pos],

			   session->internals.hb_local_data.data,

			   hb_len) != 0) {

			if (IS_DTLS(session))

				return gnutls_assert_val(GNUTLS_E_AGAIN);	/* ignore it */

			else

				return

				    gnutls_assert_val

				    (GNUTLS_E_UNEXPECTED_PACKET);

		}

		_gnutls_buffer_reset(&session->internals.hb_local_data);

		return gnutls_assert_val(GNUTLS_E_HEARTBEAT_PONG_RECEIVED);

	default:

		_gnutls_record_log

		    ("REC[%p]: HB: received unknown type %u\n", session,

		     type);

		return gnutls_assert_val(GNUTLS_E_UNEXPECTED_PACKET);

	}

}

/**

 * gnutls_heartbeat_get_timeout:

 * @session: is a #gnutls_session_t type.

 *

 * This function will return the milliseconds remaining

 * for a retransmission of the previously sent ping

 * message. This function is useful when ping is used in

 * non-blocking mode, to estimate when to call gnutls_heartbeat_ping()

 * if no packets have been received.

 *

 * Returns: the remaining time in milliseconds.

 *

 * Since: 3.1.2

 **/

unsigned int gnutls_heartbeat_get_timeout(gnutls_session_t session)

{

	struct timespec now;

	unsigned int diff;

	gnutls_gettime(&now;;

	diff = timespec_sub_ms(&now, &session->internals.hb_ping_sent);

	if (diff >= session->internals.hb_actual_retrans_timeout_ms)

		return 0;

	else

		return session->internals.hb_actual_retrans_timeout_ms -

		    diff;

}

/**

 * gnutls_heartbeat_set_timeouts:

 * @session: is a #gnutls_session_t type.

 * @retrans_timeout: The time at which a retransmission will occur in milliseconds

 * @total_timeout: The time at which the connection will be aborted, in milliseconds.

 *

 * This function will override the timeouts for the DTLS heartbeat

 * protocol. The retransmission timeout is the time after which a

 * message from the peer is not received, the previous request will

 * be retransmitted. The total timeout is the time after which the

 * handshake will be aborted with %GNUTLS_E_TIMEDOUT.

 *

 * Since: 3.1.2

 **/

void gnutls_heartbeat_set_timeouts(gnutls_session_t session,

				   unsigned int retrans_timeout,

				   unsigned int total_timeout)

{

	session->internals.hb_retrans_timeout_ms = retrans_timeout;

	session->internals.hb_total_timeout_ms = total_timeout;

}

static int

_gnutls_heartbeat_recv_params(gnutls_session_t session,

			      const uint8_t * data, size_t _data_size)

{

	unsigned policy;

	gnutls_ext_priv_data_t epriv;

	if (_gnutls_hello_ext_get_priv

	    (session, GNUTLS_EXTENSION_HEARTBEAT, &epriv) < 0) {

		if (session->security_parameters.entity == GNUTLS_CLIENT)

			return

			    gnutls_assert_val

			    (GNUTLS_E_RECEIVED_ILLEGAL_PARAMETER);

		return 0;	/* Not enabled */

	}

	if (_data_size == 0)

		return GNUTLS_E_UNEXPECTED_PACKET_LENGTH;

	policy = (intptr_t)epriv;

	if (data[0] == 1)

		policy |= LOCAL_ALLOWED_TO_SEND;

	else if (data[0] == 2)

		policy |= LOCAL_NOT_ALLOWED_TO_SEND;

	else

		return

		    gnutls_assert_val(GNUTLS_E_RECEIVED_ILLEGAL_PARAMETER);

	epriv = (void*)(intptr_t)policy;

	_gnutls_hello_ext_set_priv(session, GNUTLS_EXTENSION_HEARTBEAT,

				     epriv);

	return 0;

}

static int

_gnutls_heartbeat_send_params(gnutls_session_t session,

			      gnutls_buffer_st * extdata)

{

	gnutls_ext_priv_data_t epriv;

	uint8_t p;

	if (_gnutls_hello_ext_get_priv

	    (session, GNUTLS_EXTENSION_HEARTBEAT, &epriv) < 0)

		return 0;	/* nothing to send - not enabled */

	if (((intptr_t)epriv) & GNUTLS_HB_PEER_ALLOWED_TO_SEND)

		p = 1;

	else			/*if (epriv.num & GNUTLS_HB_PEER_NOT_ALLOWED_TO_SEND) */

		p = 2;

	if (_gnutls_buffer_append_data(extdata, &p, 1) < 0)

		return gnutls_assert_val(GNUTLS_E_MEMORY_ERROR);

	return 1;

}

static int

_gnutls_heartbeat_pack(gnutls_ext_priv_data_t epriv, gnutls_buffer_st * ps)

{

	int ret;

	BUFFER_APPEND_NUM(ps, (intptr_t)epriv);

	return 0;

}

static int

_gnutls_heartbeat_unpack(gnutls_buffer_st * ps,

			 gnutls_ext_priv_data_t * _priv)

{

	gnutls_ext_priv_data_t epriv;

	int ret;

	BUFFER_POP_CAST_NUM(ps, epriv);

	*_priv = epriv;

	ret = 0;

      error:

	return ret;

}

const hello_ext_entry_st ext_mod_heartbeat = {

	.name = "Heartbeat",

	.tls_id = 15,

	.gid = GNUTLS_EXTENSION_HEARTBEAT,

	.client_parse_point = GNUTLS_EXT_TLS,

	.server_parse_point = GNUTLS_EXT_TLS,

	.validity = GNUTLS_EXT_FLAG_TLS | GNUTLS_EXT_FLAG_DTLS | GNUTLS_EXT_FLAG_CLIENT_HELLO |

		    GNUTLS_EXT_FLAG_EE | GNUTLS_EXT_FLAG_TLS12_SERVER_HELLO,

	.recv_func = _gnutls_heartbeat_recv_params,

	.send_func = _gnutls_heartbeat_send_params,

	.pack_func = _gnutls_heartbeat_pack,

	.unpack_func = _gnutls_heartbeat_unpack,

	.deinit_func = NULL,

	.cannot_be_overriden = 1

};

#else

void gnutls_heartbeat_enable(gnutls_session_t session, unsigned int type)

{

}

unsigned gnutls_heartbeat_allowed(gnutls_session_t session, unsigned int type)

{

	return 0;

}

int

gnutls_heartbeat_ping(gnutls_session_t session, size_t data_size,

		      unsigned int max_tries, unsigned int flags)

{

	return gnutls_assert_val(GNUTLS_E_INVALID_REQUEST);

}

int gnutls_heartbeat_pong(gnutls_session_t session, unsigned int flags)

{

	return gnutls_assert_val(GNUTLS_E_INVALID_REQUEST);

}

unsigned int gnutls_heartbeat_get_timeout(gnutls_session_t session)

{

	return 0;

}

void gnutls_heartbeat_set_timeouts(gnutls_session_t session,

				   unsigned int retrans_timeout,

				   unsigned int total_timeout)

{

	return;

}

#endif