From dde4b294627eb10847ce8619dacb510011c109de Mon Sep 17 00:00:00 2001
From: toddouska <todd@wolfssl.com>
Date: Sat, 9 May 2015 11:04:47 -0700
Subject: [PATCH] add handshake done callback with ability to end connection

---
 examples/server/server.c | 20 ++++++++++++++++++++
 src/ssl.c                | 39 +++++++++++++++++++++++++++++++++++++++
 wolfssl/internal.h       |  4 ++++
 wolfssl/ssl.h            |  6 ++++++
 4 files changed, 69 insertions(+)

diff --git a/examples/server/server.c b/examples/server/server.c
index 50fb5a389..3c8ac1e65 100644
--- a/examples/server/server.c
+++ b/examples/server/server.c
@@ -60,6 +60,10 @@
     Timeval srvTo;
 #endif
 
+#ifndef NO_HANDSHAKE_DONE_CB
+    int myHsDoneCb(WOLFSSL* ssl, void* user_ctx);
+#endif
+
 
 static void NonBlockingSSL_Accept(SSL* ssl)
 {
@@ -534,6 +538,9 @@ while (1) {  /* allow resume option */
     if (ssl == NULL)
         err_sys("unable to get SSL");
 
+#ifndef NO_HANDSHAKE_DONE_CB
+    wolfSSL_SetHsDoneCb(ssl, myHsDoneCb, NULL);
+#endif
 #ifdef HAVE_CRL
     CyaSSL_EnableCRL(ssl, 0);
     CyaSSL_LoadCRL(ssl, crlPemDir, SSL_FILETYPE_PEM, CYASSL_CRL_MONITOR |
@@ -712,3 +719,16 @@ while (1) {  /* allow resume option */
 
 #endif
 
+#ifndef NO_HANDSHAKE_DONE_CB
+    int myHsDoneCb(WOLFSSL* ssl, void* user_ctx)
+    {
+        (void)user_ctx;
+        (void)ssl;
+
+        /* printf("Notified HandShake done\n"); */
+
+        /* return negative number to end TLS connection now */
+        return 0;
+    }
+#endif
+
diff --git a/src/ssl.c b/src/ssl.c
index ee0d473f0..1c027ac34 100644
--- a/src/ssl.c
+++ b/src/ssl.c
@@ -5289,6 +5289,16 @@ int wolfSSL_dtls_got_timeout(WOLFSSL* ssl)
             WOLFSSL_MSG("connect state: SECOND_REPLY_DONE");
 
         case SECOND_REPLY_DONE:
+#ifndef NO_HANDSHAKE_DONE_CB
+            if (ssl->hsDoneCb) {
+                int cbret = ssl->hsDoneCb(ssl, ssl->hsDoneCtx);
+                if (cbret < 0) {
+                    ssl->error = cbret;
+                    WOLFSSL_MSG("HandShake Done Cb don't continue error");
+                    return SSL_FATAL_ERROR;
+                }
+            }
+#endif /* NO_HANDSHAKE_DONE_CB */
             FreeHandshakeResources(ssl);
             WOLFSSL_LEAVE("SSL_connect()", SSL_SUCCESS);
             return SSL_SUCCESS;
@@ -5576,6 +5586,16 @@ int wolfSSL_dtls_got_timeout(WOLFSSL* ssl)
             WOLFSSL_MSG("accept state ACCEPT_THIRD_REPLY_DONE");
 
         case ACCEPT_THIRD_REPLY_DONE :
+#ifndef NO_HANDSHAKE_DONE_CB
+            if (ssl->hsDoneCb) {
+                int cbret = ssl->hsDoneCb(ssl, ssl->hsDoneCtx);
+                if (cbret < 0) {
+                    ssl->error = cbret;
+                    WOLFSSL_MSG("HandShake Done Cb don't continue error");
+                    return SSL_FATAL_ERROR;
+                }
+            }
+#endif /* NO_HANDSHAKE_DONE_CB */
             FreeHandshakeResources(ssl);
             WOLFSSL_LEAVE("SSL_accept()", SSL_SUCCESS);
             return SSL_SUCCESS;
@@ -5589,6 +5609,25 @@ int wolfSSL_dtls_got_timeout(WOLFSSL* ssl)
 #endif /* NO_WOLFSSL_SERVER */
 
 
+#ifndef NO_HANDSHAKE_DONE_CB
+
+int wolfSSL_SetHsDoneCb(WOLFSSL* ssl, HandShakeDoneCb cb, void* user_ctx)
+{
+    WOLFSSL_ENTER("wolfSSL_SetHsDoneCb");
+
+    if (ssl == NULL)
+        return BAD_FUNC_ARG;
+
+    ssl->hsDoneCb  = cb;
+    ssl->hsDoneCtx = user_ctx;
+
+
+    return SSL_SUCCESS;
+}
+
+#endif /* NO_HANDSHAKE_DONE_CB */
+
+
 int wolfSSL_Cleanup(void)
 {
     int ret = SSL_SUCCESS;
diff --git a/wolfssl/internal.h b/wolfssl/internal.h
index 3ead47d88..d95c2f97d 100644
--- a/wolfssl/internal.h
+++ b/wolfssl/internal.h
@@ -2111,6 +2111,10 @@ struct WOLFSSL {
     void*           verifyCbCtx;        /* cert verify callback user ctx*/
     VerifyCallback  verifyCallback;     /* cert verification callback */
     void*           heap;               /* for user overrides */
+#ifndef NO_HANDSHAKE_DONE_CB
+    HandShakeDoneCb hsDoneCb;          /*  notify user handshake done */
+    void*           hsDoneCtx;         /*  user handshake cb context  */
+#endif
     WOLFSSL_CIPHER  cipher;
     hmacfp          hmac;
     Ciphers         encrypt;
diff --git a/wolfssl/ssl.h b/wolfssl/ssl.h
index 2f72c4c65..9b9884989 100644
--- a/wolfssl/ssl.h
+++ b/wolfssl/ssl.h
@@ -1375,6 +1375,12 @@ WOLFSSL_API int wolfSSL_set_SessionTicket_cb(WOLFSSL*,
 #define WOLFSSL_CRL_MONITOR   0x01   /* monitor this dir flag */
 #define WOLFSSL_CRL_START_MON 0x02   /* start monitoring flag */
 
+
+/* notify user the hanshake is done */
+typedef int (*HandShakeDoneCb)(WOLFSSL*, void*);
+WOLFSSL_API int wolfSSL_SetHsDoneCb(WOLFSSL*, HandShakeDoneCb, void*);
+
+
 WOLFSSL_API int wolfSSL_PrintSessionStats(void);
 WOLFSSL_API int wolfSSL_get_session_stats(unsigned int* active,
                                           unsigned int* total,