1 files changed, 66 insertions, 42 deletions

diff --git a/ext/openssl/ossl_ssl.c b/ext/openssl/ossl_ssl.c
index b8874b84..6332121d 100644
--- a/ext/openssl/ossl_ssl.c
+++ b/ext/openssl/ossl_ssl.c
@@ -32,7 +32,7 @@ VALUE cSSLSocket;
 static VALUE eSSLErrorWaitReadable;
 static VALUE eSSLErrorWaitWritable;
 
-static ID ID_callback_state;
+static ID ID_callback_state, id_tmp_dh_callback, id_tmp_ecdh_callback;
 static VALUE sym_exception, sym_wait_readable, sym_wait_writable;
 
 static ID id_i_cert_store, id_i_ca_file, id_i_ca_path, id_i_verify_mode,
@@ -219,69 +219,90 @@ ossl_client_cert_cb(SSL *ssl, X509 **x509, EVP_PKEY **pkey)
     return 1;
 }
 
-#if !defined(OPENSSL_NO_DH)
-static VALUE
-ossl_call_tmp_dh_callback(VALUE args)
+#if !defined(OPENSSL_NO_DH) || \
+    !defined(OPENSSL_NO_EC) && defined(HAVE_SSL_CTX_SET_TMP_ECDH_CALLBACK)
+struct tmp_dh_callback_args {
+    VALUE ssl_obj;
+    ID id;
+    int type;
+    int is_export;
+    int keylength;
+};
+
+static EVP_PKEY *
+ossl_call_tmp_dh_callback(struct tmp_dh_callback_args *args)
 {
     VALUE cb, dh;
     EVP_PKEY *pkey;
 
-    cb = rb_funcall(rb_ary_entry(args, 0), rb_intern("tmp_dh_callback"), 0);
-
-    if (NIL_P(cb)) return Qfalse;
-    dh = rb_apply(cb, rb_intern("call"), args);
+    cb = rb_funcall(args->ssl_obj, args->id, 0);
+    if (NIL_P(cb))
+	return NULL;
+    dh = rb_funcall(cb, rb_intern("call"), 3,
+		    args->ssl_obj, INT2NUM(args->is_export), INT2NUM(args->keylength));
     pkey = GetPKeyPtr(dh);
-    if (EVP_PKEY_base_id(pkey) != EVP_PKEY_DH) return Qfalse;
+    if (EVP_PKEY_base_id(pkey) != args->type)
+	return NULL;
 
-    return dh;
+    return pkey;
 }
+#endif
 
-static DH*
+#if !defined(OPENSSL_NO_DH)
+static DH *
 ossl_tmp_dh_callback(SSL *ssl, int is_export, int keylength)
 {
-    VALUE args, dh, rb_ssl;
+    VALUE rb_ssl;
+    EVP_PKEY *pkey;
+    struct tmp_dh_callback_args args;
+    int state;
 
     rb_ssl = (VALUE)SSL_get_ex_data(ssl, ossl_ssl_ex_ptr_idx);
+    args.ssl_obj = rb_ssl;
+    args.id = id_tmp_dh_callback;
+    args.is_export = is_export;
+    args.keylength = keylength;
+    args.type = EVP_PKEY_DH;
+
+    pkey = (EVP_PKEY *)rb_protect((VALUE (*)(VALUE))ossl_call_tmp_dh_callback,
+				  (VALUE)&args, &state);
+    if (state) {
+	rb_ivar_set(rb_ssl, ID_callback_state, INT2NUM(state));
+	return NULL;
+    }
+    if (!pkey)
+	return NULL;
 
-    args = rb_ary_new_from_args(3, rb_ssl, INT2NUM(is_export), INT2NUM(keylength));
-
-    dh = rb_protect(ossl_call_tmp_dh_callback, args, NULL);
-    if (!RTEST(dh)) return NULL;
-
-    return EVP_PKEY_get0_DH(GetPKeyPtr(dh));
+    return EVP_PKEY_get0_DH(pkey);
 }
 #endif /* OPENSSL_NO_DH */
 
 #if !defined(OPENSSL_NO_EC) && defined(HAVE_SSL_CTX_SET_TMP_ECDH_CALLBACK)
-static VALUE
-ossl_call_tmp_ecdh_callback(VALUE args)
-{
-    VALUE cb, ecdh;
-    EVP_PKEY *pkey;
-
-    cb = rb_funcall(rb_ary_entry(args, 0), rb_intern("tmp_ecdh_callback"), 0);
-
-    if (NIL_P(cb)) return Qfalse;
-    ecdh = rb_apply(cb, rb_intern("call"), args);
-    pkey = GetPKeyPtr(ecdh);
-    if (EVP_PKEY_base_id(pkey) != EVP_PKEY_EC) return Qfalse;
-
-    return ecdh;
-}
-
-static EC_KEY*
+static EC_KEY *
 ossl_tmp_ecdh_callback(SSL *ssl, int is_export, int keylength)
 {
-    VALUE args, ecdh, rb_ssl;
+    VALUE rb_ssl;
+    EVP_PKEY *pkey;
+    struct tmp_dh_callback_args args;
+    int state;
 
     rb_ssl = (VALUE)SSL_get_ex_data(ssl, ossl_ssl_ex_ptr_idx);
+    args.ssl_obj = rb_ssl;
+    args.id = id_tmp_ecdh_callback;
+    args.is_export = is_export;
+    args.keylength = keylength;
+    args.type = EVP_PKEY_EC;
+
+    pkey = (EVP_PKEY *)rb_protect((VALUE (*)(VALUE))ossl_call_tmp_dh_callback,
+				  (VALUE)&args, &state);
+    if (state) {
+	rb_ivar_set(rb_ssl, ID_callback_state, INT2NUM(state));
+	return NULL;
+    }
+    if (!pkey)
+	return NULL;
 
-    args = rb_ary_new_from_args(3, rb_ssl, INT2NUM(is_export), INT2NUM(keylength));
-
-    ecdh = rb_protect(ossl_call_tmp_ecdh_callback, args, NULL);
-    if (!RTEST(ecdh)) return NULL;
-
-    return EVP_PKEY_get0_EC_KEY(GetPKeyPtr(ecdh));
+    return EVP_PKEY_get0_EC_KEY(pkey);
 }
 #endif
 
@@ -2688,6 +2709,9 @@ Init_ossl_ssl(void)
     sym_wait_readable = ID2SYM(rb_intern("wait_readable"));
     sym_wait_writable = ID2SYM(rb_intern("wait_writable"));
 
+    id_tmp_dh_callback = rb_intern("tmp_dh_callback");
+    id_tmp_ecdh_callback = rb_intern("tmp_ecdh_callback");
+
 #define DefIVarID(name) do \
     id_i_##name = rb_intern("@"#name); while (0)