From d9d0cd620df04996c454892d76a92c0e11d14504 Mon Sep 17 00:00:00 2001
From: Bruno Silvestre <bruno.silvestre@gmail.com>
Date: Thu, 26 Jul 2018 11:21:54 -0300
Subject: [PATCH] Free DH parameter right after handshake

---
 src/context.c | 23 ++++-------------------
 src/ssl.c     |  5 +++++
 2 files changed, 9 insertions(+), 19 deletions(-)

diff --git a/src/context.c b/src/context.c
index d1377f1..cc9706a 100644
--- a/src/context.c
+++ b/src/context.c
@@ -196,7 +196,6 @@ static DH *dhparam_cb(SSL *ssl, int is_export, int keylength)
 {
   BIO *bio;
   lua_State *L;
-  DH *dh_tmp = NULL;
   SSL_CTX *ctx = SSL_get_SSL_CTX(ssl);
   p_context pctx = (p_context)SSL_CTX_get_app_data(ctx);
 
@@ -217,24 +216,15 @@ static DH *dhparam_cb(SSL *ssl, int is_export, int keylength)
     lua_pop(L, 2);  /* Remove values from stack */
     return NULL;
   }
-  bio = BIO_new_mem_buf((void*)lua_tostring(L, -1), 
-    lua_rawlen(L, -1));
+
+  bio = BIO_new_mem_buf((void*)lua_tostring(L, -1), lua_rawlen(L, -1));
   if (bio) {
-    dh_tmp = PEM_read_bio_DHparams(bio, NULL, NULL, NULL);
+    pctx->dh_param = PEM_read_bio_DHparams(bio, NULL, NULL, NULL);
     BIO_free(bio);
   }
 
-  /*
-   * OpenSSL exepcts the callback to maintain a reference to the DH*.  So,
-   * cache it here, and clean up the previous set of parameters.  Any remaining
-   * set is cleaned up when destroying the LuaSec context.
-   */
-  if (pctx->dh_param)
-    DH_free(pctx->dh_param);
-  pctx->dh_param = dh_tmp;
-
   lua_pop(L, 2);    /* Remove values from stack */
-  return dh_tmp;
+  return pctx->dh_param;
 }
 
 /**
@@ -668,11 +658,6 @@ static int meth_destroy(lua_State *L)
     SSL_CTX_free(ctx->context);
     ctx->context = NULL;
   }
-  if (ctx->dh_param) {
-    DH_free(ctx->dh_param);
-    ctx->dh_param = NULL;
-  }
-
   return 0;
 }
 
diff --git a/src/ssl.c b/src/ssl.c
index e5defa8..45d143d 100644
--- a/src/ssl.c
+++ b/src/ssl.c
@@ -384,6 +384,11 @@ static int meth_handshake(lua_State *L)
 {
   p_ssl ssl = (p_ssl)luaL_checkudata(L, 1, "SSL:Connection");
   int err = handshake(ssl);
+  p_context ctx = (p_context)SSL_CTX_get_app_data(SSL_get_SSL_CTX(ssl->ssl));
+  if (ctx->dh_param) {
+    DH_free(ctx->dh_param);
+    ctx->dh_param = NULL;
+  }
   if (err == IO_DONE) {
     lua_pushboolean(L, 1);
     return 1;