/*

 * 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