fix ecdh

README.rst

@@ -90,6 +90,15 @@ Methods
 
 * Returns: ``bytes``
 
+``ecdh(public_key, update=False)``
+
+* Parameters:
+
+    - **public_key** (``bytes``) - The other party's public key.
+    - **update** - If ``True``, will update and return ``self``.
+
+* Returns: ``coincurve.PrivateKey``
+
 ``add(scalar, update=False)``
 
 * Parameters:

coincurve/keys.py

@@ -58,6 +58,23 @@ class PrivateKey:
 
         return recoverable_to_der(signature, self.context)
 
+    def ecdh(self, public_key, update=False):
+        secret = ffi.new('unsigned char [32]')
+
+        lib.secp256k1_ecdh(
+            self.context.ctx, secret, PublicKey(public_key).public_key,
+            self.secret
+        )
+
+        secret = bytes(ffi.buffer(secret, 32))
+
+        if update:
+            self.secret = secret
+            self._update_public_key()
+            return self
+
+        return PrivateKey(secret, self.context)
+
     def add(self, scalar, update=False):
         scalar = pad_scalar(scalar)
 
@@ -240,20 +257,6 @@ class PublicKey:
         # A performance hack to avoid global bool() lookup.
         return not not verified
 
-    def ecdh(self, scalar):
-        scalar = pad_scalar(scalar)
-
-        secret = ffi.new('unsigned char [32]')
-
-        success = lib.secp256k1_ecdh(
-            self.context.ctx, secret, self.public_key, scalar
-        )
-
-        if not success:
-            raise ValueError('Scalar was invalid (zero or overflow).')
-
-        return bytes(ffi.buffer(secret, 32))
-
     def add(self, scalar, update=False):
         scalar = pad_scalar(scalar)
 

tests/test_keys.py

@@ -58,6 +58,12 @@ class TestPrivateKey:
     def test_from_der(self):
         assert PrivateKey.from_der(PRIVATE_KEY_DER).secret == PRIVATE_KEY_BYTES
 
+    def test_ecdh(self):
+        a = PrivateKey()
+        b = PrivateKey()
+
+        assert a.ecdh(b.public_key.format()) == b.ecdh(a.public_key.format())
+
     def test_add(self):
         assert PrivateKey(b'\x01').add(b'\x09').to_int() == 10
 
@@ -99,14 +105,6 @@ class TestPublicKey:
         public_key = PublicKey(PUBLIC_KEY_COMPRESSED)
         assert public_key.verify(SIGNATURE, MESSAGE)
 
-    def test_ecdh(self):
-        x = urandom(32)
-        k = urandom(32)
-        a = PublicKey.from_secret(x)
-        b = PublicKey.from_secret(k)
-
-        assert a.ecdh(k) == b.ecdh(x)
-
     def test_transform(self):
         x = urandom(32)
         k = urandom(32)