From 9cd41ac714d9bff819ece6d8cdcde064d403c671 Mon Sep 17 00:00:00 2001
From: Alex Gaynor <alex.gaynor@gmail.com>
Date: Sun, 28 Jul 2019 22:58:04 -0400
Subject: Make DER reader into a context manager (#4957)

* Make DER reader into a context manager

* Added another test case

* flake8
---
 src/cryptography/hazmat/_der.py                        | 12 +++++++++---
 src/cryptography/hazmat/primitives/asymmetric/utils.py |  9 ++++-----
 src/cryptography/x509/extensions.py                    | 18 +++++++++---------
 3 files changed, 22 insertions(+), 17 deletions(-)

(limited to 'src')

diff --git a/src/cryptography/hazmat/_der.py b/src/cryptography/hazmat/_der.py
index 3a121a85..51518d64 100644
--- a/src/cryptography/hazmat/_der.py
+++ b/src/cryptography/hazmat/_der.py
@@ -36,6 +36,13 @@ class DERReader(object):
     def __init__(self, data):
         self.data = memoryview(data)
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_value, tb):
+        if exc_value is None:
+            self.check_empty()
+
     def is_empty(self):
         return len(self.data) == 0
 
@@ -100,9 +107,8 @@ class DERReader(object):
         return body
 
     def read_single_element(self, expected_tag):
-        ret = self.read_element(expected_tag)
-        self.check_empty()
-        return ret
+        with self:
+            return self.read_element(expected_tag)
 
     def read_optional_element(self, expected_tag):
         if len(self.data) > 0 and six.indexbytes(self.data, 0) == expected_tag:
diff --git a/src/cryptography/hazmat/primitives/asymmetric/utils.py b/src/cryptography/hazmat/primitives/asymmetric/utils.py
index 43d5b9bf..14d2abee 100644
--- a/src/cryptography/hazmat/primitives/asymmetric/utils.py
+++ b/src/cryptography/hazmat/primitives/asymmetric/utils.py
@@ -12,11 +12,10 @@ from cryptography.hazmat.primitives import hashes
 
 
 def decode_dss_signature(signature):
-    seq = DERReader(signature).read_single_element(SEQUENCE)
-    r = seq.read_element(INTEGER).as_integer()
-    s = seq.read_element(INTEGER).as_integer()
-    seq.check_empty()
-    return r, s
+    with DERReader(signature).read_single_element(SEQUENCE) as seq:
+        r = seq.read_element(INTEGER).as_integer()
+        s = seq.read_element(INTEGER).as_integer()
+        return r, s
 
 
 def encode_dss_signature(r, s):
diff --git a/src/cryptography/x509/extensions.py b/src/cryptography/x509/extensions.py
index c78c76c2..5bef9945 100644
--- a/src/cryptography/x509/extensions.py
+++ b/src/cryptography/x509/extensions.py
@@ -48,17 +48,17 @@ def _key_identifier_from_public_key(public_key):
             serialization.PublicFormat.SubjectPublicKeyInfo
         )
 
-        public_key_info = DERReader(serialized).read_single_element(SEQUENCE)
-        algorithm = public_key_info.read_element(SEQUENCE)
-        public_key = public_key_info.read_element(BIT_STRING)
-        public_key_info.check_empty()
+        reader = DERReader(serialized)
+        with reader.read_single_element(SEQUENCE) as public_key_info:
+            algorithm = public_key_info.read_element(SEQUENCE)
+            public_key = public_key_info.read_element(BIT_STRING)
 
         # Double-check the algorithm structure.
-        algorithm.read_element(OBJECT_IDENTIFIER)
-        if not algorithm.is_empty():
-            # Skip the optional parameters field.
-            algorithm.read_any_element()
-        algorithm.check_empty()
+        with algorithm:
+            algorithm.read_element(OBJECT_IDENTIFIER)
+            if not algorithm.is_empty():
+                # Skip the optional parameters field.
+                algorithm.read_any_element()
 
         # BIT STRING contents begin with the number of padding bytes added. It
         # must be zero for SubjectPublicKeyInfo structures.
-- 
cgit v1.2.3