several updates in the heartbeat handling code.

14 files changed, 249 insertions, 226 deletions

diff --git a/lib/ext/heartbeat.c b/lib/ext/heartbeat.c
index 677a91dc77..d27b94f1bd 100644
--- a/lib/ext/heartbeat.c
+++ b/lib/ext/heartbeat.c
@@ -65,10 +65,10 @@ _gnutls_heartbeat (unsigned policy)
 /**
   * gnutls_heartbeat_allowed:
   * @session: is a #gnutls_session_t structure.
-  * @type: non zero is for sending, and zero for receiving
+  * @type: one of %GNUTLS_HB_LOCAL_ALLOWED_TO_SEND and %GNUTLS_HB_PEER_ALLOWED_TO_SEND
   *
   * This function will check whether heartbeats are allowed
-  * in this session.
+  * to be sent or received in this session. 
   *
   * Returns: Non zero if heartbeats are allowed.
   *
@@ -82,7 +82,7 @@ gnutls_heartbeat_allowed (gnutls_session_t session, unsigned int type)
       (session, GNUTLS_EXTENSION_HEARTBEAT, &epriv) < 0)
     return 0;                   /* Not enabled */
 
-  if (type != 0)
+  if (type == GNUTLS_HB_LOCAL_ALLOWED_TO_SEND)
     {
       if (epriv.num & LOCAL_ALLOWED_TO_SEND)
         return 1;
@@ -105,12 +105,12 @@ heartbeat_allow_send (gnutls_session_t session)
   return gnutls_heartbeat_allowed(session, 1);
 }
 
-/**
- * _gnutls_heartbeat_send_data:
+/*-
+ * heartbeat_send_data:
  * @session: is a #gnutls_session_t structure.
  * @data: contains the data to send.
  * @data_size: is the length of the data.
- * @request: true if Request message is about to be send.
+ * @start_timer: true if the heartbeat timer is to be started.
  *
  * This function has the similar semantics with gnutls_record_send() The only
  * difference is that it uses GNUTLS_HEARTBEAT content type.
@@ -123,16 +123,14 @@ heartbeat_allow_send (gnutls_session_t session)
  *   number of bytes sent might be less than @data_size.  The maximum
  *   number of bytes this function can send in a single call depends
  *   on the negotiated maximum record size.
- **/
-ssize_t
-_gnutls_heartbeat_send_data (gnutls_session_t session, const void *data,
-                             size_t data_size, int request)
+ -*/
+static int
+heartbeat_send_data (gnutls_session_t session, const void *data,
+                     size_t data_size, uint8_t type)
 {
   int ret;
-  char pr[128];
   gnutls_buffer_st response;
   uint8_t payload[16];
-  uint8_t type = request ? HEARTBEAT_REQUEST : HEARTBEAT_RESPONSE;
   _gnutls_buffer_init (&response);
 
   ret = gnutls_rnd (GNUTLS_RND_RANDOM, payload, 16);
@@ -146,23 +144,7 @@ _gnutls_heartbeat_send_data (gnutls_session_t session, const void *data,
   ret = _gnutls_send_int (session, GNUTLS_HEARTBEAT, -1,
                           EPOCH_WRITE_CURRENT, response.data,
                           response.length, MBUFFER_FLUSH);
-  if (request)
-    {
-      if (ret >= 0)
-        {
-          _gnutls_record_log
-              ("REC[%p]: HB %zu bytes sent OK [%d in packet], saved for response verification:%s\n",
-               session, data_size, ret, _gnutls_bin2hex ((uint8_t *) data,
-                                                         data_size, pr,
-                                                         sizeof (pr),
-                                                         NULL));
-          session->internals.dtls.heartbeat_timeout = HEARTBEAT_TIMEOUT;
-          gettime (&session->internals.dtls.heartbeat_sent);
-          _gnutls_buffer_reset (&session->internals.heartbeat_payload);
-          BUFFER_APPEND (&session->internals.heartbeat_payload, data,
-                         data_size);
-        }
-    }
+
   _gnutls_record_log ("REC[%p]: HB sent: %d\n", session, ret);
 
   _gnutls_buffer_clear (&response);
@@ -172,27 +154,25 @@ _gnutls_heartbeat_send_data (gnutls_session_t session, const void *data,
 /**
  * gnutls_heartbeat_ping:
  * @session: is a #gnutls_session_t structure.
- * @data_size: is the length of the random data.
- * @wait_for_it: wait for pong or timeout instead of returning
- * immediately.
- *
- * This function has the similar semantics with gnutls_record_send().
- * The only difference is that it uses GNUTLS_HEARTBEAT content type
- * and auto-generate data to send.
+ * @data_size: is the length of the ping payload.
+ * @max_tries: if flags is %GNUTLS_HEARTBEAT_WAIT then this sets the number of retransmissions.
+ * @flags: if %GNUTLS_HEARTBEAT_WAIT then wait for pong or timeout instead of returning immediately.
  *
- * This function send HeartBeat request message with proper padding.
+ * 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.
  *
- *
- * Returns: The number of bytes sent, or a negative error code. The
- *   number of bytes sent might be less than @data_size.  The maximum
- *   number of bytes this function can send in a single call depends
- *   on the negotiated maximum record size. GNUTLS_E_HEARTBEAT_FLIGHT
- *   is returned if HB Request is alredy in flight.
+ * Returns: %GNUTLS_E_SUCCESS on success, otherwise a negative error code.
  **/
-ssize_t
-gnutls_heartbeat_ping (gnutls_session_t session, size_t data_size)
+int
+gnutls_heartbeat_ping (gnutls_session_t session, size_t data_size, 
+                       unsigned int max_tries, unsigned int flags)
 {
-  int ret = GNUTLS_E_HEARTBEAT_FLIGHT;
+  int ret;
+  unsigned int retries = 0;
 
   if (data_size > MAX_HEARTBEAT_LENGTH)
     return gnutls_assert_val (GNUTLS_E_UNEXPECTED_PACKET_LENGTH);
@@ -200,42 +180,101 @@ gnutls_heartbeat_ping (gnutls_session_t session, size_t data_size)
   if (!heartbeat_allow_send (session))
     return GNUTLS_E_INVALID_REQUEST;
 
-  _gnutls_record_log
-      ("REC[%p]: sending HB_REQUEST with length: %zu to peer\n",
-       session, data_size);
-
-  if (gnutls_heartbeat_timeout (session, 1) == GNUTLS_E_ILLEGAL_PARAMETER)
+  switch(session->internals.hb_state)
     {
-      uint8_t data[data_size];
-      ret = _gnutls_rnd (GNUTLS_RND_NONCE, data, data_size);
-      if (ret >= 0)
-        ret = _gnutls_heartbeat_send_data (session, data, data_size, 1);
+      case SHB_SEND1:
+        _gnutls_record_log
+         ("REC[%p]: sending HB_REQUEST with length: %zu to peer\n", session, data_size);
+
+        _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);
+        session->internals.hb_local_data.length = data_size;
+
+        session->internals.hb_state = SHB_SEND2;
+      case SHB_SEND2:
+        session->internals.hb_timeout = HEARTBEAT_TIMEOUT;
+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);
+
+        gettime (&session->internals.hb_ping_sent);
+            
+        if (!(flags & GNUTLS_HEARTBEAT_WAIT))
+          {
+            session->internals.hb_state = SHB_SEND1;
+            break;
+          }
+        
+        session->internals.hb_state = SHB_RECV;
+
+      case SHB_RECV:
+        ret = _gnutls_recv_int(session, GNUTLS_HEARTBEAT, -1, NULL, 0, NULL, session->internals.hb_timeout);
+        if (ret == GNUTLS_E_HEARTBEAT_PONG_RECEIVED)
+          {
+            session->internals.hb_state = SHB_SEND1;
+            break;
+          }
+        else if (ret == GNUTLS_E_TIMEDOUT)
+          {
+            retries++;
+            if (retries > max_tries)
+              {
+                session->internals.hb_state = SHB_SEND1;
+                return gnutls_assert_val(ret);
+              }
+
+            session->internals.hb_timeout *= 2;
+            session->internals.hb_timeout %= 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);
+          }
     }
-  else
-    _gnutls_record_log
-        ("REC[%p]: HB_REQUEST with length %zu already in-flight: %d\n",
-         session, data_size, gnutls_heartbeat_timeout (session, 1));
 
-  return ret;
+  return 0;
 }
 
 /**
- * simple wrapper for ping with random length
+ * gnutls_heartbeat_pong:
+ * @session: is a #gnutls_session_t structure.
+ * @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.
  **/
-ssize_t
-gnutls_heartbeat_ping_rnd (gnutls_session_t session)
+int
+gnutls_heartbeat_pong (gnutls_session_t session, unsigned int flags)
 {
-  uint8_t rnd;
-  int ret;
-  
-  ret = gnutls_rnd (GNUTLS_RND_NONCE, &rnd, 1);
-  if (ret < 0)
-    return gnutls_assert_val(ret);
-  
-  return gnutls_heartbeat_ping (session, rnd + 1);
+int ret;
+
+  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);
+  return ret;
 }
 
-/**
+/*-
  * Process HB message in buffer.
  * @bufel: the (suspected) HeartBeat message (TLV+padding)
  *
@@ -244,11 +283,12 @@ gnutls_heartbeat_ping_rnd (gnutls_session_t session)
  * GNUTLS_E_AGAIN if processed OK
  * GNUTLS_E_HEARTBEAT_PONG_FAILED on response send failure for request message
  * GNUTLS_E_RECEIVED_ILLEGAL_PARAMETER on payload mismatch for response message
- **/
+ -*/
 int
 _gnutls_heartbeat_handle (gnutls_session_t session, mbuffer_st * bufel)
 {
   char pr[128];
+  int ret;
   uint8_t *msg = _mbuffer_get_udata_ptr (bufel);
   size_t hb_len, len = _mbuffer_get_udata_size (bufel);
 
@@ -261,108 +301,49 @@ _gnutls_heartbeat_handle (gnutls_session_t session, mbuffer_st * bufel)
   hb_len = _gnutls_read_uint16 (msg + 1);
   if (hb_len > len - 3)
     return gnutls_assert_val (GNUTLS_E_UNEXPECTED_PACKET);
-
+  
   switch (msg[0])
     {
     case HEARTBEAT_REQUEST:
       _gnutls_record_log
-          ("REC[%p]: received HEARTBEAT_REQUEST, responding...\n",
-           session);
-      if (_gnutls_heartbeat_send_data (session, msg + 3, hb_len, 0) >= 0)
-        return GNUTLS_E_AGAIN;  /* HB responded, no APP_DATA so needs to be called again */
-      else
-        {                       /* immediate response failed, TBD: save received data somewhere and let upper layers handle it, loosing single ping is non-critical for HB */
-          return GNUTLS_E_HEARTBEAT_PONG_FAILED;
-        }
-      break;
+          ("REC[%p]: received HEARTBEAT_REQUEST\n", session);
+
+      _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);
+      
+      memcpy(session->internals.hb_remote_data.data, msg+3, hb_len);
+      session->internals.hb_remote_data.length = hb_len;
+
+      return GNUTLS_E_HEARTBEAT_PING_RECEIVED;
 
     case HEARTBEAT_RESPONSE:
 
-      if (session->internals.heartbeat_payload.length != hb_len)
-        return gnutls_assert_val (GNUTLS_E_UNEXPECTED_PACKET_LENGTH);
+      if (hb_len != session->internals.hb_local_data.length)
+        return gnutls_assert_val (GNUTLS_E_UNEXPECTED_PACKET);
 
-      if (memcmp (msg + 3, session->internals.heartbeat_payload.data,
+      if (memcmp (msg + 3, session->internals.hb_local_data.data,
                   hb_len) != 0)
         {
           _gnutls_record_log ("REC[%p]: HB: %s - received\n", session,
                               _gnutls_bin2hex (msg + 3, hb_len, pr,
                                                sizeof (pr), NULL));
-          _gnutls_record_log ("REC[%p]: HB: %s - expected\n", session,
-                              _gnutls_bin2hex (session->internals.
-                                               heartbeat_payload.data,
-                                               hb_len, pr, sizeof (pr),
-                                               NULL));
-          return gnutls_assert_val (GNUTLS_E_RECEIVED_ILLEGAL_PARAMETER);
+          return gnutls_assert_val (GNUTLS_E_UNEXPECTED_PACKET);
         }
-      _gnutls_record_log
-          ("REC[%p]: HB: %zu response bytes received OK (%d msec. left before timeout)\n",
-           session, hb_len, gnutls_heartbeat_timeout (session, 1));
 
-      session->internals.dtls.heartbeat_timeout = HEARTBEAT_TIMEOUT;
-      _gnutls_buffer_reset (&session->internals.heartbeat_payload);
-      return GNUTLS_E_AGAIN;
+      _gnutls_buffer_reset (&session->internals.hb_local_data);
 
+      return GNUTLS_E_HEARTBEAT_PONG_RECEIVED;
     default:
       _gnutls_record_log
           ("REC[%p]: HB: received unknown type %u\n",
            session, msg[0]);
       return gnutls_assert_val (GNUTLS_E_UNEXPECTED_PACKET);
     }
-
-  return gnutls_assert_val (GNUTLS_E_INTERNAL_ERROR);
 }
 
-/**
- * Update HB timeouts: should be called on retransmit to set new timeout.
- * @check_only: guarantees lack of side-effects (no variables are written)
- *
- * Returns:
- * number of milliseconds left before timeout expiration OR
- * GNUTLS_E_TIMEDOUT if timeout expired
- * GNUTLS_E_SUCCESS if timeout updated
- * GNUTLS_E_ILLEGAL_PARAMETER if no HB message is in-flight
- **/
-int
-gnutls_heartbeat_timeout (gnutls_session_t session, int check_only)
-{
-  struct timespec now;
-  unsigned int ms;
-  if (session->internals.heartbeat_payload.length)
-    {
-      _gnutls_debug_log
-          ("HB: %zu bytes are in-flight already: %u msec timeout.\n",
-           session->internals.heartbeat_payload.length,
-           (unsigned int) session->internals.dtls.heartbeat_timeout);
-      gettime (&now);
-      ms = _dtls_timespec_sub_ms (&now,
-                                  &session->internals.dtls.heartbeat_sent);
-      if (ms > session->internals.dtls.heartbeat_timeout)
-        {
-          if (check_only)
-            return GNUTLS_E_TIMEDOUT;
-          _gnutls_buffer_reset (&session->internals.heartbeat_payload);
-
-          if (session->internals.dtls.heartbeat_timeout * 2 >
-              MAX_HEARTBEAT_TIMEOUT)
-            {                   /* update impossible */
-              session->internals.dtls.heartbeat_timeout =
-                  HEARTBEAT_TIMEOUT;
-              return GNUTLS_E_TIMEDOUT;
-            }
-          else
-            {
-              session->internals.dtls.heartbeat_timeout *= 2;
-              gettime (&session->internals.dtls.heartbeat_sent);
-              return GNUTLS_E_SUCCESS;
-            }
-        }
-      else
-        {
-          return ms;
-        }
-    }
-  return GNUTLS_E_ILLEGAL_PARAMETER;
-}
 
 static int
 _gnutls_heartbeat_recv_params (gnutls_session_t session,
diff --git a/lib/ext/heartbeat.h b/lib/ext/heartbeat.h
index 7c4b2991d7..085d14c2a9 100644
--- a/lib/ext/heartbeat.h
+++ b/lib/ext/heartbeat.h
@@ -28,7 +28,7 @@
 #define HEARTBEAT_REQUEST 1
 #define HEARTBEAT_RESPONSE 2
 
-#define MAX_HEARTBEAT_LENGTH 16384
+#define MAX_HEARTBEAT_LENGTH DEFAULT_MAX_RECORD_SIZE
 #define HEARTBEAT_TIMEOUT 1000
 #define MAX_HEARTBEAT_TIMEOUT 60000
 
@@ -42,11 +42,6 @@ extern extension_entry_st ext_mod_heartbeat;
 typedef uint8_t heartbeat_policy_t;
 
 int _gnutls_heartbeat_handle (gnutls_session_t session, mbuffer_st * bufel);
-ssize_t _gnutls_heartbeat_send_data (gnutls_session_t session,
-                                     const void *data, size_t data_size,
-                                     int request);
-ssize_t gnutls_heartbeat_ping (gnutls_session_t session, size_t data_size);
-ssize_t gnutls_heartbeat_ping_rnd (gnutls_session_t session);
 int _gnutls_heartbeat_enabled (gnutls_session_t session, int local);
 int gnutls_heartbeat_timeout (gnutls_session_t session, int check_only);
 #endif
diff --git a/lib/gnutls_buffers.c b/lib/gnutls_buffers.c
index 3dc10d2a2c..30db0336ed 100644
--- a/lib/gnutls_buffers.c
+++ b/lib/gnutls_buffers.c
@@ -49,6 +49,7 @@
 #include <gnutls_dtls.h>
 #include <system.h>
 #include <gnutls_constate.h> /* gnutls_epoch_get */
+#include <gnutls_handshake.h> /* remaining_time() */
 #include <errno.h>
 #include <system.h>
 #include "debug.h"
@@ -174,24 +175,27 @@ int errno_to_gerr(int err)
 
 static ssize_t
 _gnutls_dgram_read (gnutls_session_t session, mbuffer_st **bufel,
-		    gnutls_pull_func pull_func, unsigned int ms)
+		    gnutls_pull_func pull_func, unsigned int *ms)
 {
   ssize_t i, ret;
   uint8_t *ptr;
+  struct timespec t1, t2;
   size_t max_size = _gnutls_get_max_decrypted_data(session);
   size_t recv_size = MAX_RECV_SIZE(session);
   gnutls_transport_ptr_t fd = session->internals.transport_recv_ptr;
+  unsigned int diff;
 
   if (recv_size > max_size)
     recv_size = max_size;
 
   session->internals.direction = 0;
 
-  if (ms)
+  if (ms && *ms > 0)
     {
-      ret = _gnutls_io_check_recv(session, ms);
+      ret = _gnutls_io_check_recv(session, *ms);
       if (ret < 0)
         return gnutls_assert_val(ret);
+      gettime(&t1);
     }
 
   *bufel = _mbuffer_alloc (0, max_size);
@@ -228,6 +232,16 @@ _gnutls_dgram_read (gnutls_session_t session, mbuffer_st **bufel,
       _mbuffer_set_udata_size (*bufel, i);
     }
 
+  if (ms && *ms > 0)
+    {
+      gettime(&t2);
+      diff = _dtls_timespec_sub_ms(&t2, &t1);
+      if (diff < *ms)
+        *ms -= diff;
+      else
+        return gnutls_assert_val(GNUTLS_E_TIMEDOUT);
+    }
+
   _gnutls_read_log ("READ: read %d bytes from %p\n", (int) i, fd);
   
   return i;
@@ -239,7 +253,7 @@ cleanup:
 
 static ssize_t
 _gnutls_stream_read (gnutls_session_t session, mbuffer_st **bufel,
-		     size_t size, gnutls_pull_func pull_func, unsigned int ms)
+		     size_t size, gnutls_pull_func pull_func, unsigned int *ms)
 {
   size_t left;
   ssize_t i = 0;
@@ -263,9 +277,9 @@ _gnutls_stream_read (gnutls_session_t session, mbuffer_st **bufel,
   left = size;
   while (left > 0)
     {
-      if (ms)
+      if (ms && *ms > 0)
         {
-          ret = _gnutls_io_check_recv(session, ms);
+          ret = _gnutls_io_check_recv(session, *ms);
           if (ret < 0)
             return gnutls_assert_val(ret);
             
@@ -314,12 +328,12 @@ _gnutls_stream_read (gnutls_session_t session, mbuffer_st **bufel,
       left -= i;
       (*bufel)->msg.size += i;
 
-      if (ms && left)
+      if (ms && *ms > 0)
         {
           gettime(&t2);
           diff = _dtls_timespec_sub_ms(&t2, &t1);
-          if (diff < ms)
-            ms -= diff;
+          if (diff < *ms)
+            *ms -= diff;
           else
             return gnutls_assert_val(GNUTLS_E_TIMEDOUT);
         }
@@ -341,7 +355,7 @@ finish:
  */
 static ssize_t
 _gnutls_read (gnutls_session_t session, mbuffer_st **bufel,
-	      size_t size, gnutls_pull_func pull_func, unsigned int ms)
+	      size_t size, gnutls_pull_func pull_func, unsigned int *ms)
 {
   if (IS_DTLS (session))
     /* Size is not passed, since a whole datagram will be read. */
@@ -402,7 +416,10 @@ _gnutls_writev (gnutls_session_t session, const giovec_t * giovec,
   return i;
 }
 
-/* This function is like recv(with MSG_PEEK). But it does not return -1 on error.
+/* 
+ * @ms: a pointer to the number of milliseconds to wait for data. Use zero or NULL for indefinite.
+ *
+ * This function is like recv(with MSG_PEEK). But it does not return -1 on error.
  * It does return gnutls_errno instead.
  * This function reads data from the socket and keeps them in a buffer, of up to
  * MAX_RECV_SIZE. 
@@ -416,7 +433,7 @@ _gnutls_writev (gnutls_session_t session, const giovec_t * giovec,
  */
 ssize_t
 _gnutls_io_read_buffered (gnutls_session_t session, size_t total,
-                          content_type_t recv_type, unsigned int ms)
+                          content_type_t recv_type, unsigned int *ms)
 {
   ssize_t ret = 0;
   size_t min;
@@ -1170,6 +1187,7 @@ _gnutls_handshake_io_recv_int (gnutls_session_t session,
                                handshake_buffer_st * hsk, unsigned int optional)
 {
   int ret;
+  unsigned int tleft = 0;
 
   ret = get_last_packet(session, htype, hsk, optional);
   if (ret != GNUTLS_E_AGAIN && ret != GNUTLS_E_INTERRUPTED && ret != GNUTLS_E_REQUESTED_DATA_NOT_AVAILABLE)
@@ -1195,9 +1213,17 @@ _gnutls_handshake_io_recv_int (gnutls_session_t session,
         return gnutls_assert_val(ret);
     }
 
+  if (htype != (unsigned)-1)
+    {
+      ret = handshake_remaining_time(session);
+      if (ret < 0)
+        return gnutls_assert_val(ret);
+      tleft = ret;
+    }
+
   /* if we don't have a complete message waiting for us, try 
    * receiving more */
-  ret = _gnutls_recv_in_buffers(session, GNUTLS_HANDSHAKE, htype);
+  ret = _gnutls_recv_in_buffers(session, GNUTLS_HANDSHAKE, htype, tleft);
   if (ret < 0)
     return gnutls_assert_val_fatal(ret);
 
diff --git a/lib/gnutls_buffers.h b/lib/gnutls_buffers.h
index 52106acf41..193e97f21f 100644
--- a/lib/gnutls_buffers.h
+++ b/lib/gnutls_buffers.h
@@ -53,7 +53,7 @@ record_check_unprocessed (gnutls_session_t session)
 int _gnutls_record_buffer_get (content_type_t type,
                                gnutls_session_t session, uint8_t * data,
                                size_t length, uint8_t seq[8]);
-ssize_t _gnutls_io_read_buffered (gnutls_session_t, size_t n, content_type_t, unsigned int ms);
+ssize_t _gnutls_io_read_buffered (gnutls_session_t, size_t n, content_type_t, unsigned int *ms);
 int _gnutls_io_clear_peeked_data (gnutls_session_t session);
 
 ssize_t _gnutls_io_write_buffered (gnutls_session_t session,
@@ -109,7 +109,7 @@ _gnutls_parse_record_buffered_msgs (gnutls_session_t session);
 
 ssize_t
 _gnutls_recv_in_buffers (gnutls_session_t session, content_type_t type,
-                         gnutls_handshake_description_t htype);
+                         gnutls_handshake_description_t htype, unsigned int ms);
 
 #define _gnutls_handshake_io_buffer_clear( session) \
         _mbuffer_head_clear( &session->internals.handshake_send_buffer); \
diff --git a/lib/gnutls_dtls.c b/lib/gnutls_dtls.c
index ae6faf9aeb..54c56815e8 100644
--- a/lib/gnutls_dtls.c
+++ b/lib/gnutls_dtls.c
@@ -179,7 +179,7 @@ static int is_next_hpacket_expected(gnutls_session_t session)
 int ret;
 
   /* htype is arbitrary */
-  ret = _gnutls_recv_in_buffers(session, GNUTLS_HANDSHAKE, GNUTLS_HANDSHAKE_FINISHED);
+  ret = _gnutls_recv_in_buffers(session, GNUTLS_HANDSHAKE, GNUTLS_HANDSHAKE_FINISHED, 0);
   if (ret < 0)
     return gnutls_assert_val(ret);
   
diff --git a/lib/gnutls_errors.c b/lib/gnutls_errors.c
index b43771f11c..2a14f74a13 100644
--- a/lib/gnutls_errors.c
+++ b/lib/gnutls_errors.c
@@ -93,10 +93,10 @@ static const gnutls_error_entry error_algorithms[] = {
                GNUTLS_E_NO_CERTIFICATE_FOUND, 1),
   ERROR_ENTRY (N_("The given DSA key is incompatible with the selected TLS protocol."),
                GNUTLS_E_INCOMPAT_DSA_KEY_WITH_TLS_PROTOCOL, 1),
-  ERROR_ENTRY (N_("HeartBeat Response to Request has failed for some reason: pong dropped."),
-               GNUTLS_E_HEARTBEAT_PONG_FAILED, 1),
-  ERROR_ENTRY (N_("HeartBeat message is already in-flight."),
-               GNUTLS_E_HEARTBEAT_FLIGHT, 1),
+  ERROR_ENTRY (N_("A heartbeat pong message was received."),
+               GNUTLS_E_HEARTBEAT_PONG_RECEIVED, 0),
+  ERROR_ENTRY (N_("A heartbeat ping message was received."),
+               GNUTLS_E_HEARTBEAT_PING_RECEIVED, 0),
   ERROR_ENTRY (N_("There is already a crypto algorithm with lower priority."),
                GNUTLS_E_CRYPTO_ALREADY_REGISTERED, 1),
 
diff --git a/lib/gnutls_handshake.c b/lib/gnutls_handshake.c
index edb736b4c5..1eb5d6b454 100644
--- a/lib/gnutls_handshake.c
+++ b/lib/gnutls_handshake.c
@@ -2725,6 +2725,12 @@ _gnutls_recv_handshake_final (gnutls_session_t session, int init)
   int ret = 0;
   uint8_t ch;
   unsigned int ccs_len = 1;
+  unsigned int tleft;
+  
+  ret = handshake_remaining_time(session);
+  if (ret < 0)
+    return gnutls_assert_val(ret);
+  tleft = ret;
 
   switch (STATE)
     {
@@ -2747,7 +2753,8 @@ _gnutls_recv_handshake_final (gnutls_session_t session, int init)
       if (gnutls_protocol_get_version (session) == GNUTLS_DTLS0_9)
         ccs_len = 3;
 
-      ret = _gnutls_recv_int (session, GNUTLS_CHANGE_CIPHER_SPEC, -1, &ch, ccs_len, NULL);
+      ret = _gnutls_recv_int (session, GNUTLS_CHANGE_CIPHER_SPEC, -1, &ch, ccs_len, NULL, 
+                              tleft);
       if (ret <= 0)
         {
           ERR ("recv ChangeCipherSpec", ret);
diff --git a/lib/gnutls_handshake.h b/lib/gnutls_handshake.h
index 22a55f7dec..4fa14f3d9c 100644
--- a/lib/gnutls_handshake.h
+++ b/lib/gnutls_handshake.h
@@ -23,6 +23,8 @@
 #ifndef HANDSHAKE_H
 #define HANDSHAKE_H
 
+#include <gnutls_errors.h>
+
 int _gnutls_send_handshake (gnutls_session_t session, mbuffer_st * bufel,
                             gnutls_handshake_description_t type);
 int _gnutls_recv_hello_request (gnutls_session_t session, void *data,
@@ -55,4 +57,17 @@ void _gnutls_handshake_hash_buffers_clear (gnutls_session_t session);
 #define AGAIN(target) (STATE==target?1:0)
 #define AGAIN2(state, target) (state==target?1:0)
 
+inline static int handshake_remaining_time(gnutls_session_t session)
+{
+  if (session->internals.handshake_endtime)
+    {
+      time_t now = gnutls_time(0);
+      if (now < session->internals.handshake_endtime)
+        return (session->internals.handshake_endtime - now) * 1000;
+      else 
+        return gnutls_assert_val(GNUTLS_E_TIMEDOUT);
+    }
+  return 0;
+}
+
 #endif
diff --git a/lib/gnutls_int.h b/lib/gnutls_int.h
index d80c2a7255..5462789096 100644
--- a/lib/gnutls_int.h
+++ b/lib/gnutls_int.h
@@ -230,6 +230,13 @@ typedef enum handshake_state_t
   STATE60 = 60, STATE61, STATE62, STATE70, STATE71
 } handshake_state_t;
 
+typedef enum heartbeat_state_t
+{  
+  SHB_SEND1 = 0, 
+  SHB_SEND2, 
+  SHB_RECV, 
+} heartbeat_state_t;
+
 #include <gnutls_str.h>
 
 /* This is the maximum number of algorithms (ciphers or macs etc).
@@ -660,9 +667,6 @@ typedef struct
    */
   time_t async_term;
 
-  struct timespec heartbeat_sent; /* timestamp: when last HeartBeat Request was sent*/
-  time_t heartbeat_timeout; /* current timeout, in milliseconds*/
-
   /* last retransmission triggered by record layer */
   struct timespec last_retransmit;
   unsigned int packets_dropped;
@@ -890,14 +894,19 @@ typedef struct
   /* if set it means that the master key was set using
    * gnutls_session_set_master() rather than being negotiated. */
   unsigned int premaster_set:1;
-
+  
   unsigned int cb_tls_unique_len;
   unsigned char cb_tls_unique[MAX_VERIFY_DATA_SIZE];
   
   unsigned int handshake_endtime; /* end time in seconds */
   unsigned int handshake_timeout_ms; /* timeout in milliseconds */
 
-  gnutls_buffer_st heartbeat_payload; /* store in-flight payload for heartbeat extension*/
+  gnutls_buffer_st hb_local_data;
+  gnutls_buffer_st hb_remote_data;
+  struct timespec hb_ping_sent; /* timestamp: when last HeartBeat ping was sent*/
+  unsigned int hb_timeout; /* current timeout, in milliseconds*/
+
+  heartbeat_state_t hb_state; /* for ping */
   
   /* If you add anything here, check _gnutls_handshake_internal_state_clear().
    */
diff --git a/lib/gnutls_record.c b/lib/gnutls_record.c
index 2a616338e0..9f91f3d15c 100644
--- a/lib/gnutls_record.c
+++ b/lib/gnutls_record.c
@@ -220,7 +220,7 @@ gnutls_bye (gnutls_session_t session, gnutls_close_request_t how)
         {
           do
             {
-              ret = _gnutls_recv_int (session, GNUTLS_ALERT, -1, NULL, 0, NULL);
+              ret = _gnutls_recv_int (session, GNUTLS_ALERT, -1, NULL, 0, NULL, 0);
             }
           while (ret == GNUTLS_E_GOT_APPLICATION_DATA);
 
@@ -874,7 +874,7 @@ record_read_headers (gnutls_session_t session,
 static int recv_headers( gnutls_session_t session, content_type_t type, 
                          gnutls_handshake_description_t htype, 
                          struct tls_record_st* record,
-                         unsigned int ms)
+                         unsigned int *ms)
 {
 int ret;
 gnutls_datum_t raw; /* raw headers */
@@ -952,7 +952,9 @@ gnutls_datum_t raw; /* raw headers */
 
 #define MAX_EMPTY_PACKETS_SEQUENCE 4
 
-/* This will receive record layer packets and add them to 
+/* @ms: is the number of milliseconds to wait for data. Use zero for indefinite.
+ *
+ * This will receive record layer packets and add them to 
  * application_data_buffer and handshake_data_buffer.
  *
  * If the htype is not -1 then handshake timeouts
@@ -960,7 +962,7 @@ gnutls_datum_t raw; /* raw headers */
  */
 ssize_t
 _gnutls_recv_in_buffers (gnutls_session_t session, content_type_t type,
-                         gnutls_handshake_description_t htype)
+                         gnutls_handshake_description_t htype, unsigned int ms)
 {
   uint64 *packet_sequence;
   uint8_t *ciphertext;
@@ -970,7 +972,6 @@ _gnutls_recv_in_buffers (gnutls_session_t session, content_type_t type,
   record_parameters_st *record_params;
   record_state_st *record_state;
   struct tls_record_st record;
-  time_t now, tleft = 0;
 
 begin:
 
@@ -1003,17 +1004,8 @@ begin:
 
   record_state = &record_params->read;
 
-  if (htype != (unsigned)-1 && session->internals.handshake_endtime > 0)
-    {
-      now = gnutls_time(0);
-      if (now < session->internals.handshake_endtime)
-        tleft = (session->internals.handshake_endtime - now) * 1000;
-      else
-        return gnutls_assert_val(GNUTLS_E_TIMEDOUT);
-    }
-
   /* receive headers */
-  ret = recv_headers(session, type, htype, &record, tleft);
+  ret = recv_headers(session, type, htype, &record, &ms);
   if (ret < 0)
     {
       ret = gnutls_assert_val_fatal(ret);
@@ -1025,20 +1017,11 @@ begin:
   else
     packet_sequence = &record_state->sequence_number;
 
-  if (htype != (unsigned)-1 && session->internals.handshake_endtime > 0)
-    {
-      now = gnutls_time(0);
-      if (now < session->internals.handshake_endtime)
-        tleft = (session->internals.handshake_endtime - now) * 1000;
-      else
-        return gnutls_assert_val(GNUTLS_E_TIMEDOUT);
-    }
-
   /* Read the packet data and insert it to record_recv_buffer.
    */
   ret =
        _gnutls_io_read_buffered (session, record.packet_size,
-                                 record.type, tleft);
+                                 record.type, &ms);
   if (ret != record.packet_size)
     {
       gnutls_assert();
@@ -1210,7 +1193,8 @@ recv_error:
 ssize_t
 _gnutls_recv_int (gnutls_session_t session, content_type_t type,
                   gnutls_handshake_description_t htype,
-                  uint8_t * data, size_t data_size, void* seq)
+                  uint8_t * data, size_t data_size, void* seq,
+                  unsigned int ms)
 {
   int ret;
 
@@ -1241,7 +1225,7 @@ _gnutls_recv_int (gnutls_session_t session, content_type_t type,
   if (ret != 0)
     return ret;
 
-  ret = _gnutls_recv_in_buffers(session, type, htype);
+  ret = _gnutls_recv_in_buffers(session, type, htype, ms);
   if (ret < 0 && ret != GNUTLS_E_SESSION_EOF)
     return gnutls_assert_val(ret);
 
@@ -1312,8 +1296,9 @@ gnutls_record_send (gnutls_session_t session, const void *data,
  * The number of bytes received might be less than the requested @data_size.
  **/
 ssize_t
-gnutls_record_recv (gnutls_session_t session, void *data, size_t data_size) {
-    return _gnutls_recv_int (session, GNUTLS_APPLICATION_DATA, -1, data, data_size, NULL);
+gnutls_record_recv (gnutls_session_t session, void *data, size_t data_size) 
+{
+    return _gnutls_recv_int (session, GNUTLS_APPLICATION_DATA, -1, data, data_size, NULL, 0);
 }
 
 /**
@@ -1341,5 +1326,5 @@ gnutls_record_recv_seq (gnutls_session_t session, void *data, size_t data_size,
   unsigned char *seq)
 {
   return _gnutls_recv_int (session, GNUTLS_APPLICATION_DATA, -1, data,
-                           data_size, seq);
+                           data_size, seq, 0);
 }
diff --git a/lib/gnutls_record.h b/lib/gnutls_record.h
index 4151673f1d..12cc31a592 100644
--- a/lib/gnutls_record.h
+++ b/lib/gnutls_record.h
@@ -32,7 +32,7 @@ ssize_t _gnutls_send_int (gnutls_session_t session, content_type_t type,
                           size_t sizeofdata, unsigned int mflags);
 ssize_t _gnutls_recv_int (gnutls_session_t session, content_type_t type,
                           gnutls_handshake_description_t, uint8_t * data,
-                          size_t sizeofdata, void* seq);
+                          size_t sizeofdata, void* seq, unsigned int ms);
 int _gnutls_get_max_decrypted_data(gnutls_session_t session);
 
 #endif
diff --git a/lib/gnutls_state.c b/lib/gnutls_state.c
index a2feb5b53e..6adfd417fa 100644
--- a/lib/gnutls_state.c
+++ b/lib/gnutls_state.c
@@ -326,7 +326,8 @@ gnutls_init (gnutls_session_t * session, unsigned int flags)
 
   /* Initialize buffers */
   _gnutls_buffer_init (&(*session)->internals.handshake_hash_buffer);
-  _gnutls_buffer_init (&(*session)->internals.heartbeat_payload);
+  _gnutls_buffer_init (&(*session)->internals.hb_remote_data);
+  _gnutls_buffer_init (&(*session)->internals.hb_local_data);
 
   _mbuffer_head_init (&(*session)->internals.record_buffer);
   _mbuffer_head_init (&(*session)->internals.record_send_buffer);
@@ -441,7 +442,8 @@ gnutls_deinit (gnutls_session_t session)
       }
 
   _gnutls_buffer_clear (&session->internals.handshake_hash_buffer);
-  _gnutls_buffer_clear (&session->internals.heartbeat_payload);
+  _gnutls_buffer_clear (&session->internals.hb_remote_data);
+  _gnutls_buffer_clear (&session->internals.hb_local_data);
 
   _mbuffer_head_clear (&session->internals.record_buffer);
   _mbuffer_head_clear (&session->internals.record_recv_buffer);
diff --git a/lib/includes/gnutls/gnutls.h.in b/lib/includes/gnutls/gnutls.h.in
index b1685cd80a..ce77a8a300 100644
--- a/lib/includes/gnutls/gnutls.h.in
+++ b/lib/includes/gnutls/gnutls.h.in
@@ -857,8 +857,11 @@ gnutls_ecc_curve_t gnutls_ecc_curve_get(gnutls_session_t session);
 
 /* Record layer functions.
  */
-  ssize_t gnutls_heartbeat_ping (gnutls_session_t session, size_t data_size);
-  ssize_t gnutls_heartbeat_ping_rnd (gnutls_session_t session);
+#define GNUTLS_HEARTBEAT_WAIT 1
+  int gnutls_heartbeat_ping (gnutls_session_t session, size_t data_size, 
+                       unsigned int max_tries, unsigned int flags);
+  int gnutls_heartbeat_pong (gnutls_session_t session, unsigned int flags);
+
   ssize_t gnutls_record_send (gnutls_session_t session, const void *data,
                               size_t data_size);
   ssize_t gnutls_record_recv (gnutls_session_t session, void *data,
@@ -912,10 +915,11 @@ gnutls_ecc_curve_t gnutls_ecc_curve_get(gnutls_session_t session);
 #define GNUTLS_HB_PEER_ALLOWED_TO_SEND (1)
 #define GNUTLS_HB_PEER_NOT_ALLOWED_TO_SEND (1<<1)
 
-    /* Heartbeat */
-    void gnutls_heartbeat_enable (gnutls_session_t session, unsigned int type);
-    int gnutls_heartbeat_allowed (gnutls_session_t session, unsigned int type);
-    int gnutls_heartbeat_timeout (gnutls_session_t session, int check_only);
+  /* Heartbeat */
+  void gnutls_heartbeat_enable (gnutls_session_t session, unsigned int type);
+
+#define GNUTLS_HB_LOCAL_ALLOWED_TO_SEND (1<<2)
+  int gnutls_heartbeat_allowed (gnutls_session_t session, unsigned int type);
 
   /* Safe renegotiation */
   int gnutls_safe_renegotiation_status (gnutls_session_t session);
@@ -1977,8 +1981,8 @@ typedef int (*gnutls_pin_callback_t) (void *userdata, int attempt,
 #define GNUTLS_E_OPENPGP_PREFERRED_KEY_ERROR -215
 #define GNUTLS_E_INCOMPAT_DSA_KEY_WITH_TLS_PROTOCOL -216
 
-#define GNUTLS_E_HEARTBEAT_FLIGHT -298
-#define GNUTLS_E_HEARTBEAT_PONG_FAILED -299
+#define GNUTLS_E_HEARTBEAT_PONG_RECEIVED -292
+#define GNUTLS_E_HEARTBEAT_PING_RECEIVED -293
 
 /* PKCS11 related */
 #define GNUTLS_E_PKCS11_ERROR -300
diff --git a/lib/libgnutls.map b/lib/libgnutls.map
index 032ee644cb..110209cff9 100644
--- a/lib/libgnutls.map
+++ b/lib/libgnutls.map
@@ -302,7 +302,7 @@ GNUTLS_1_4
     gnutls_record_send;
     gnutls_record_set_max_size;
     gnutls_heartbeat_ping;
-    gnutls_heartbeat_ping_rnd;
+    gnutls_heartbeat_pong;
     gnutls_rehandshake;
     gnutls_rsa_export_get_modulus_bits;
     gnutls_rsa_export_get_pubkey;
@@ -337,7 +337,6 @@ GNUTLS_1_4
     gnutls_sign_list;
     gnutls_heartbeat_enable;
     gnutls_heartbeat_allowed;
-    gnutls_heartbeat_timeout;
     gnutls_srp_1024_group_generator;
     gnutls_srp_1024_group_prime;
     gnutls_srp_1536_group_generator;