From 28edf6bafcfd8c80e2cb52c498db8930f8517fd1 Mon Sep 17 00:00:00 2001
From: Kazuki Yamaguchi <k@rhe.jp>
Date: Sat, 18 Mar 2017 22:34:19 +0900
Subject: pkey: reimplement PKey::DH#compute_key and PKey::EC#dh_compute_key

Use the new OpenSSL::PKey::PKey#derive instead of the raw
{EC,}DH_compute_key(), mainly to reduce amount of the C code.
---
 ext/openssl/ossl_pkey_dh.c | 35 -----------------------------------
 ext/openssl/ossl_pkey_ec.c | 32 --------------------------------
 lib/openssl/pkey.rb        | 33 +++++++++++++++++++++++++++++++++
 3 files changed, 33 insertions(+), 67 deletions(-)

diff --git a/ext/openssl/ossl_pkey_dh.c b/ext/openssl/ossl_pkey_dh.c
index bf4e3f93..b4f10a52 100644
--- a/ext/openssl/ossl_pkey_dh.c
+++ b/ext/openssl/ossl_pkey_dh.c
@@ -518,40 +518,6 @@ ossl_dh_generate_key(VALUE self)
     return self;
 }
 
-/*
- *  call-seq:
- *     dh.compute_key(pub_bn) -> aString
- *
- * Returns a String containing a shared secret computed from the other party's public value.
- * See DH_compute_key() for further information.
- *
- * === Parameters
- * * _pub_bn_ is a OpenSSL::BN, *not* the DH instance returned by
- *   DH#public_key as that contains the DH parameters only.
- */
-static VALUE
-ossl_dh_compute_key(VALUE self, VALUE pub)
-{
-    DH *dh;
-    const BIGNUM *pub_key, *dh_p;
-    VALUE str;
-    int len;
-
-    GetDH(self, dh);
-    DH_get0_pqg(dh, &dh_p, NULL, NULL);
-    if (!dh_p)
-	ossl_raise(eDHError, "incomplete DH");
-    pub_key = GetBNPtr(pub);
-    len = DH_size(dh);
-    str = rb_str_new(0, len);
-    if ((len = DH_compute_key((unsigned char *)RSTRING_PTR(str), pub_key, dh)) < 0) {
-	ossl_raise(eDHError, NULL);
-    }
-    rb_str_set_len(str, len);
-
-    return str;
-}
-
 /*
  * Document-method: OpenSSL::PKey::DH#set_pqg
  * call-seq:
@@ -629,7 +595,6 @@ Init_ossl_dh(void)
     rb_define_method(cDH, "public_key", ossl_dh_to_public_key, 0);
     rb_define_method(cDH, "params_ok?", ossl_dh_check_params, 0);
     rb_define_method(cDH, "generate_key!", ossl_dh_generate_key, 0);
-    rb_define_method(cDH, "compute_key", ossl_dh_compute_key, 1);
 
     DEF_OSSL_PKEY_BN(cDH, dh, p);
     DEF_OSSL_PKEY_BN(cDH, dh, q);
diff --git a/ext/openssl/ossl_pkey_ec.c b/ext/openssl/ossl_pkey_ec.c
index fc2bc6c8..0aa4e9d3 100644
--- a/ext/openssl/ossl_pkey_ec.c
+++ b/ext/openssl/ossl_pkey_ec.c
@@ -582,37 +582,6 @@ static VALUE ossl_ec_key_check_key(VALUE self)
     return Qtrue;
 }
 
-/*
- *  call-seq:
- *     key.dh_compute_key(pubkey)   => String
- *
- *  See the OpenSSL documentation for ECDH_compute_key()
- */
-static VALUE ossl_ec_key_dh_compute_key(VALUE self, VALUE pubkey)
-{
-    EC_KEY *ec;
-    EC_POINT *point;
-    int buf_len;
-    VALUE str;
-
-    GetEC(self, ec);
-    GetECPoint(pubkey, point);
-
-/* BUG: need a way to figure out the maximum string size */
-    buf_len = 1024;
-    str = rb_str_new(0, buf_len);
-/* BUG: take KDF as a block */
-    buf_len = ECDH_compute_key(RSTRING_PTR(str), buf_len, point, ec, NULL);
-    if (buf_len < 0)
-         ossl_raise(eECError, "ECDH_compute_key");
-
-    rb_str_resize(str, buf_len);
-
-    return str;
-}
-
-/* sign_setup */
-
 /*
  *  call-seq:
  *     key.dsa_sign_asn1(data)   => String
@@ -1752,7 +1721,6 @@ void Init_ossl_ec(void)
     rb_define_alias(cEC, "generate_key", "generate_key!");
     rb_define_method(cEC, "check_key", ossl_ec_key_check_key, 0);
 
-    rb_define_method(cEC, "dh_compute_key", ossl_ec_key_dh_compute_key, 1);
     rb_define_method(cEC, "dsa_sign_asn1", ossl_ec_key_dsa_sign_asn1, 1);
     rb_define_method(cEC, "dsa_verify_asn1", ossl_ec_key_dsa_verify_asn1, 2);
 /* do_sign/do_verify */
diff --git a/lib/openssl/pkey.rb b/lib/openssl/pkey.rb
index 9cc32763..be60ac2b 100644
--- a/lib/openssl/pkey.rb
+++ b/lib/openssl/pkey.rb
@@ -9,6 +9,24 @@ require_relative 'marshal'
 module OpenSSL::PKey
   class DH
     include OpenSSL::Marshal
+
+    # :call-seq:
+    #    dh.compute_key(pub_bn) -> string
+    #
+    # Returns a String containing a shared secret computed from the other
+    # party's public value.
+    #
+    # This method is provided for backwards compatibility, and calls #derive
+    # internally.
+    #
+    # === Parameters
+    # * _pub_bn_ is a OpenSSL::BN, *not* the DH instance returned by
+    #   DH#public_key as that contains the DH parameters only.
+    def compute_key(pub_bn)
+      peer = dup
+      peer.set_key(pub_bn, nil)
+      derive(peer)
+    end
   end
 
   class DSA
@@ -18,7 +36,22 @@ module OpenSSL::PKey
   if defined?(EC)
   class EC
     include OpenSSL::Marshal
+
+    # :call-seq:
+    #    ec.dh_compute_key(pubkey) -> string
+    #
+    # Derives a shared secret by ECDH. _pubkey_ must be an instance of
+    # OpenSSL::PKey::EC::Point and must belong to the same group.
+    #
+    # This method is provided for backwards compatibility, and calls #derive
+    # internally.
+    def dh_compute_key(pubkey)
+      peer = OpenSSL::PKey::EC.new(group)
+      peer.public_key = pubkey
+      derive(peer)
+    end
   end
+
   class EC::Point
     # :call-seq:
     #    point.to_bn([conversion_form]) -> OpenSSL::BN
-- 
cgit v1.2.3