From 789925210420e5134da9932955d3ebbfad8aaf89 Mon Sep 17 00:00:00 2001
From: John Safranek <jasaf@mac.com>
Date: Fri, 14 Sep 2012 19:30:50 -0700
Subject: [PATCH] dtls handshake improvement

---
 cyassl/internal.h |  3 +++
 src/internal.c    | 30 ++++++++++++++++++++++++++----
 src/ssl.c         | 19 ++++++-------------
 3 files changed, 35 insertions(+), 17 deletions(-)

diff --git a/cyassl/internal.h b/cyassl/internal.h
index 3eb684c4b..b9152b125 100644
--- a/cyassl/internal.h
+++ b/cyassl/internal.h
@@ -968,6 +968,8 @@ typedef struct Keys {
     word32 dtls_peer_sequence_number;
     word32 dtls_expected_peer_sequence_number;
     word16 dtls_handshake_number;
+    word16 dtls_peer_handshake_number;
+    word16 dtls_expected_peer_handshake_number;
     word16 dtls_epoch;
     word16 dtls_peer_epoch;
     word16 dtls_expected_peer_epoch;
@@ -1492,6 +1494,7 @@ CYASSL_LOCAL void BuildTlsFinished(CYASSL* ssl, Hashes* hashes,
 #ifdef CYASSL_DTLS
     CYASSL_LOCAL int  DtlsPoolInit(CYASSL*);
     CYASSL_LOCAL void DtlsPoolSave(CYASSL*, const byte*, int);
+    CYASSL_LOCAL int  DtlsPoolTimeout(CYASSL*);
     CYASSL_LOCAL int  DtlsPoolSend(CYASSL*);
     CYASSL_LOCAL void DtlsPoolReset(CYASSL*);
 #endif /* CYASSL_DTLS */
diff --git a/src/internal.c b/src/internal.c
index 901515b42..f3ede6dd8 100644
--- a/src/internal.c
+++ b/src/internal.c
@@ -977,6 +977,7 @@ int InitSSL(CYASSL* ssl, CYASSL_CTX* ctx)
     ssl->keys.dtls_peer_sequence_number = 0;
     ssl->keys.dtls_expected_peer_sequence_number = 0;
     ssl->keys.dtls_handshake_number     = 0;
+    ssl->keys.dtls_expected_peer_handshake_number = 0;
     ssl->keys.dtls_epoch                = 0;
     ssl->keys.dtls_peer_epoch           = 0;
     ssl->keys.dtls_expected_peer_epoch  = 0;
@@ -1241,6 +1242,17 @@ void DtlsPoolReset(CYASSL* ssl)
 }
 
 
+int DtlsPoolTimeout(CYASSL* ssl)
+{
+    int result = -1;
+    if (ssl->dtls_timeout < 64) {
+        ssl->dtls_timeout *= 2;
+        result = 0;
+    }
+    return result;
+}
+
+
 int DtlsPoolSend(CYASSL* ssl)
 {
     DtlsPool *pool = ssl->dtls_pool;
@@ -1532,8 +1544,10 @@ retry:
 
 #ifdef CYASSL_DTLS
             case IO_ERR_TIMEOUT:
-                DtlsPoolSend(ssl);
-                goto retry;
+                if (DtlsPoolTimeout(ssl) == 0 && DtlsPoolSend(ssl) == 0)
+                    goto retry;
+                else
+                    return -1;
 #endif
 
             default:
@@ -1803,14 +1817,13 @@ static int GetDtlsHandShakeHeader(CYASSL* ssl, const byte* input,
 {
     word32 idx = *inOutIdx;
 
-    (void)ssl;
     *inOutIdx += HANDSHAKE_HEADER_SZ + DTLS_HANDSHAKE_EXTRA;
     
     *type = input[idx++];
     c24to32(input + idx, size);
     idx += BYTE3_LEN;
 
-    /* skip the sequence number */
+    ato16(input + idx, &ssl->keys.dtls_peer_handshake_number);
     idx += DTLS_HANDSHAKE_SEQ_SZ;
 
     c24to32(input + idx, fragOffset);
@@ -2484,6 +2497,15 @@ static int DoDtlsHandShakeMsg(CYASSL* ssl, byte* input, word32* inOutIdx,
     if (*inOutIdx + fragSz > totalSz)
         return INCOMPLETE_DATA;
 
+    if (ssl->keys.dtls_peer_handshake_number ==
+                                ssl->keys.dtls_expected_peer_handshake_number) {
+        ssl->keys.dtls_expected_peer_handshake_number++;
+    }
+    else {
+        *inOutIdx += size;
+        return 0;
+    }
+
     if (fragSz < size) {
         /* message is fragmented, knit back together */
         byte* buf = ssl->buffers.dtlsHandshake.buffer;
diff --git a/src/ssl.c b/src/ssl.c
index 6a796229b..a8d4cf808 100644
--- a/src/ssl.c
+++ b/src/ssl.c
@@ -2226,23 +2226,16 @@ int CyaSSL_dtls_get_current_timeout(CYASSL* ssl)
 
 int CyaSSL_dtls_got_timeout(CYASSL* ssl)
 {
-    int result = SSL_NOT_IMPLEMENTED;
-    (void)ssl;
-
 #ifdef CYASSL_DTLS
-    if (ssl->dtls_timeout < 64) {
-        ssl->dtls_timeout *= 2;
-        if (DtlsPoolSend(ssl) < 0)
-            result = SSL_FATAL_ERROR;
-        else
-            result = SSL_SUCCESS;
-    }
-    else {
+    int result = SSL_SUCCESS;
+    if (DtlsPoolTimeout(ssl) < 0 || DtlsPoolSend(ssl) < 0) {
         result = SSL_FATAL_ERROR;
     }
-#endif
-
     return result;
+#else
+    (void)ssl;
+    return SSL_NOT_IMPLEMENTED;
+#endif
 }