refactor to use ctx manager and share logic with set_public_key_affine

2 files changed, 28 insertions, 41 deletions

diff --git a/cryptography/hazmat/backends/openssl/backend.py b/cryptography/hazmat/backends/openssl/backend.py
index 389ef0be1..3667232bb 100644
--- a/cryptography/hazmat/backends/openssl/backend.py
+++ b/cryptography/hazmat/backends/openssl/backend.py
@@ -1085,30 +1085,19 @@ class Backend(object):
         finally:
             self._lib.BN_CTX_end(bn_ctx)
 
-    def _ec_key_set_public_key_affine_coordinates(self, ctx, x, y):
+    def _ec_key_determine_group_get_set_funcs(self, ctx):
         """
-        This is a port of EC_KEY_set_public_key_affine_coordinates that was
-        added in 1.0.1.
-
-        Sets the public key point in the EC_KEY context to the affine x and y
-        values.
+        Given an EC_KEY determine the group and what methods are required to
+        get/set point coordinates.
         """
-
         assert ctx != self._ffi.NULL
 
-        bn_x = self._int_to_bn(x)
-        bn_y = self._int_to_bn(y)
-
         nid_two_field = self._lib.OBJ_sn2nid(b"characteristic-two-field")
         assert nid_two_field != self._lib.NID_undef
 
         group = self._lib.EC_KEY_get0_group(ctx)
         assert group != self._ffi.NULL
 
-        point = self._lib.EC_POINT_new(group)
-        assert point != self._ffi.NULL
-        point = self._ffi.gc(point, self._lib.EC_POINT_free)
-
         method = self._lib.EC_GROUP_method_of(group)
         assert method != self._ffi.NULL
 
@@ -1124,6 +1113,28 @@ class Backend(object):
 
         assert set_func and get_func
 
+        return set_func, get_func, group
+
+    def _ec_key_set_public_key_affine_coordinates(self, ctx, x, y):
+        """
+        This is a port of EC_KEY_set_public_key_affine_coordinates that was
+        added in 1.0.1.
+
+        Sets the public key point in the EC_KEY context to the affine x and y
+        values.
+        """
+
+        bn_x = self._int_to_bn(x)
+        bn_y = self._int_to_bn(y)
+
+        set_func, get_func, group = (
+            self._ec_key_determine_group_get_set_funcs(ctx)
+        )
+
+        point = self._lib.EC_POINT_new(group)
+        assert point != self._ffi.NULL
+        point = self._ffi.gc(point, self._lib.EC_POINT_free)
+
         with self._tmp_bn_ctx() as bn_ctx:
             check_x = self._lib.BN_CTX_get(bn_ctx)
             check_y = self._lib.BN_CTX_get(bn_ctx)
diff --git a/cryptography/hazmat/backends/openssl/ec.py b/cryptography/hazmat/backends/openssl/ec.py
index 879d674a7..369b185b9 100644
--- a/cryptography/hazmat/backends/openssl/ec.py
+++ b/cryptography/hazmat/backends/openssl/ec.py
@@ -203,35 +203,13 @@ class _EllipticCurvePublicKey(object):
                 _Reasons.UNSUPPORTED_PUBLIC_KEY_ALGORITHM)
 
     def public_numbers(self):
-        bn_ctx = self._backend._lib.BN_CTX_new()
-        assert bn_ctx != self._backend._ffi.NULL
-        bn_ctx = self._backend._ffi.gc(bn_ctx, self._backend._lib.BN_CTX_free)
-
-        group = self._backend._lib.EC_KEY_get0_group(self._ec_key)
-        assert group != self._backend._ffi.NULL
-
-        method = self._backend._lib.EC_GROUP_method_of(group)
-        assert method != self._backend._ffi.NULL
-
-        nid = self._backend._lib.EC_METHOD_get_field_type(method)
-        assert nid != self._backend._lib.NID_undef
-
-        nid_two_field = self._backend._lib.OBJ_sn2nid(
-            b"characteristic-two-field"
+        set_func, get_func, group = (
+            self._backend._ec_key_determine_group_get_set_funcs(self._ec_key)
         )
-        assert nid_two_field != self._backend._lib.NID_undef
-
-        if nid == nid_two_field and self._backend._lib.Cryptography_HAS_EC2M:
-            get_func = self._backend._lib.EC_POINT_get_affine_coordinates_GF2m
-        else:
-            get_func = self._backend._lib.EC_POINT_get_affine_coordinates_GFp
-
         point = self._backend._lib.EC_KEY_get0_public_key(self._ec_key)
         assert point != self._backend._ffi.NULL
 
-        try:
-            self._backend._lib.BN_CTX_start(bn_ctx)
-
+        with self._backend._tmp_bn_ctx() as bn_ctx:
             bn_x = self._backend._lib.BN_CTX_get(bn_ctx)
             bn_y = self._backend._lib.BN_CTX_get(bn_ctx)
 
@@ -240,8 +218,6 @@ class _EllipticCurvePublicKey(object):
 
             x = self._backend._bn_to_int(bn_x)
             y = self._backend._bn_to_int(bn_y)
-        finally:
-            self._backend._lib.BN_CTX_end(bn_ctx)
 
         return ec.EllipticCurvePublicNumbers(
             x=x,