WIP: ASN.1 decoders to produce Python built-inscodecs-produce-python-builtins

Also, BitString chunked decoding optimized

4 files changed, 178 insertions, 25 deletions

diff --git a/CHANGES.rst b/CHANGES.rst
index d7b9516..e9cf42f 100644
--- a/CHANGES.rst
+++ b/CHANGES.rst
@@ -8,6 +8,8 @@ Revision 0.4.1, released XX-09-2017
 - ASN.1 decoders modified to accept the `nativeMode` boolean flag
   to produce either ASN.1 schema instance (the default) or built-in
   Python types (which is significantly faster)
+- BitString ASN.1 decoder performance improved when processing
+  chunked-encoded input
 
 Revision 0.3.5, released 16-09-2017
 -----------------------------------
diff --git a/pyasn1/codec/ber/decoder.py b/pyasn1/codec/ber/decoder.py
index 0dc7d01..0f90e26 100644
--- a/pyasn1/codec/ber/decoder.py
+++ b/pyasn1/codec/ber/decoder.py
@@ -38,12 +38,15 @@ class AbstractSimpleDecoder(AbstractDecoder):
 
     def _createComponent(self, asn1Spec, tagSet, value=noValue,
                          nativeMode=False, **options):
-        if nativeMode:
+        if nativeMode and value is not noValue:
             return value
-        if asn1Spec is None:
+
+        elif asn1Spec is None:
             return self.protoComponent.clone(value, tagSet=tagSet)
+
         elif value is noValue:
             return asn1Spec
+
         else:
             return asn1Spec.clone(value)
 
@@ -119,61 +122,103 @@ class BooleanDecoder(IntegerDecoder):
         if nativeMode:
             return value and True or False
 
-        return IntegerDecoder._createComponent(self, asn1Spec, tagSet, value and 1 or 0, **options)
+        return IntegerDecoder._createComponent(self, asn1Spec, tagSet,
+                                               value and 1 or 0, nativeMode, **options)
 
 
 class BitStringDecoder(AbstractSimpleDecoder):
     protoComponent = univ.BitString(())
     supportConstructedForm = True
 
+    def _createComponent(self, asn1Spec, tagSet, value=noValue,
+                         nativeMode=False, **options):
+        if nativeMode and value is not noValue:
+            return str(value)
+
+        return AbstractSimpleDecoder._createComponent(self, asn1Spec, tagSet, value,
+                                                      nativeMode, **options)
+
     def valueDecoder(self, substrate, asn1Spec,
                      tagSet=None, length=None, state=None,
                      decodeFun=None, substrateFun=None,
                      **options):
         head, tail = substrate[:length], substrate[length:]
+
+        if substrateFun:
+            asn1Object = self._createComponent(asn1Spec, tagSet, **options)
+            return substrateFun(asn1Object, substrate, length)
+
         if tagSet[0].tagFormat == tag.tagFormatSimple:  # XXX what tag to check?
             if not head:
                 raise error.PyAsn1Error('Empty substrate')
+
             trailingBits = oct2int(head[0])
             if trailingBits > 7:
                 raise error.PyAsn1Error(
                     'Trailing bits overflow %s' % trailingBits
                 )
+
             head = head[1:]
-            value = self.protoComponent.fromOctetString(head, trailingBits)
+            value = univ.BitString.SizedInteger.fromOctetString(head, trailingBits)
+
             return self._createComponent(asn1Spec, tagSet, value, **options), tail
 
         if not self.supportConstructedForm:
             raise error.PyAsn1Error('Constructed encoding form prohibited at %s' % self.__class__.__name__)
 
-        bitString = self._createComponent(asn1Spec, tagSet, **options)
+        # All inner fragments are of the same type, treat them as octet string
+        # TODO: nested constructed encoding won't work
+        substrateFun = self.substrateCollector
 
-        if substrateFun:
-            return substrateFun(bitString, substrate, length)
+        bitString = univ.BitString.SizedInteger()
 
         while head:
-            component, head = decodeFun(head, self.protoComponent, **options)
-            bitString += component
+            chunk, head = decodeFun(head, self.protoComponent,
+                                        substrateFun=substrateFun,
+                                        **options)
+            trailingBits = oct2int(chunk[0])
+            if trailingBits > 7:
+                raise error.PyAsn1Error(
+                    'Trailing bits overflow %s' % trailingBits
+                )
+
+            value = univ.BitString.SizedInteger.fromOctetString(chunk[1:], trailingBits)
 
-        return bitString, tail
+            bitString = (bitString << value.bitLength) | value
+
+        return self._createComponent(asn1Spec, tagSet, bitString, **options), tail
 
     def indefLenValueDecoder(self, substrate, asn1Spec,
                              tagSet=None, length=None, state=None,
                              decodeFun=None, substrateFun=None,
                              **options):
 
-        bitString = self._createComponent(asn1Spec, tagSet, **options)
-
         if substrateFun:
-            return substrateFun(bitString, substrate, length)
+            asn1Object = self._createComponent(asn1Spec, tagSet, **options)
+            return substrateFun(asn1Object, substrate, length)
+
+        # All inner fragments are of the same type, treat them as octet string
+        substrateFun = self.substrateCollector
+
+        bitString = univ.BitString.SizedInteger()
 
         while substrate:
-            component, substrate = decodeFun(substrate, self.protoComponent,
-                                             allowEoo=True, **options)
-            if component is eoo.endOfOctets:
+            chunk, substrate = decodeFun(substrate, self.protoComponent,
+                                         substrateFun=substrateFun,
+                                         allowEoo=True, **options)
+
+            if chunk is eoo.endOfOctets:
                 break
 
-            bitString += component
+            trailingBits = oct2int(chunk[0])
+            if trailingBits > 7:
+                raise error.PyAsn1Error(
+                    'Trailing bits overflow %s' % trailingBits
+                )
+
+            value = univ.BitString.SizedInteger.fromOctetString(chunk[1:], trailingBits)
+
+            bitString = (bitString << value.bitLength) | value
 
         else:
             raise error.SubstrateUnderrunError('No EOO seen before substrate ends')
@@ -192,8 +237,8 @@ class OctetStringDecoder(AbstractSimpleDecoder):
         head, tail = substrate[:length], substrate[length:]
 
         if substrateFun:
-            return substrateFun(self._createComponent(asn1Spec, tagSet, **options),
-                                substrate, length)
+            asn1Object = self._createComponent(asn1Spec, tagSet, **options)
+            return substrateFun(asn1Object, substrate, length)
 
         if tagSet[0].tagFormat == tag.tagFormatSimple:  # XXX what tag to check?
             return self._createComponent(asn1Spec, tagSet, head, **options), tail
@@ -256,7 +301,7 @@ class NullDecoder(AbstractSimpleDecoder):
 
         head, tail = substrate[:length], substrate[length:]
 
-        component = self._createComponent(asn1Spec, tagSet, **options)
+        component = self._createComponent(asn1Spec, tagSet, None, **options)
 
         if head:
             raise error.PyAsn1Error('Unexpected %d-octet substrate for Null' % length)
diff --git a/pyasn1/type/univ.py b/pyasn1/type/univ.py
index 3db6ab9..6a1e56b 100644
--- a/pyasn1/type/univ.py
+++ b/pyasn1/type/univ.py
@@ -402,6 +402,10 @@ class BitString(base.AbstractSimpleAsn1Item):
     class SizedInteger(SizedIntegerBase):
         bitLength = leadingZeroBits = None
 
+        @classmethod
+        def fromOctetString(cls, value, padding=0):
+            return cls(integer.from_bytes(value) >> padding).setBitLength(len(value) * 8 - padding)
+
         def setBitLength(self, bitLength):
             self.bitLength = bitLength
             self.leadingZeroBits = max(bitLength - integer.bitLength(self), 0)
@@ -413,6 +417,13 @@ class BitString(base.AbstractSimpleAsn1Item):
 
             return self.bitLength
 
+        def __str__(self):
+            binString = binary.bin(self)[2:]
+            return '0' * (len(self) - len(binString)) + binString
+
+        def __lshift__(self, other):
+            return self
+
     def __init__(self, value=noValue, **kwargs):
         if value is noValue or value is None:
             if kwargs:
@@ -633,8 +644,7 @@ class BitString(base.AbstractSimpleAsn1Item):
     def asBinary(self):
         """Get |ASN.1| value as a text string of bits.
         """
-        binString = binary.bin(self._value)[2:]
-        return '0' * (len(self._value) - len(binString)) + binString
+        return str(self._value)
 
     @classmethod
     def fromHexString(cls, value):
@@ -675,7 +685,7 @@ class BitString(base.AbstractSimpleAsn1Item):
         value: :class:`str` (Py2) or :class:`bytes` (Py3)
             Text string like '\\\\x01\\\\xff' (Py2) or b'\\\\x01\\\\xff' (Py3)
         """
-        return cls(cls.SizedInteger(integer.from_bytes(value) >> padding).setBitLength(len(value) * 8 - padding))
+        return cls(cls.SizedInteger.fromOctetString(value, padding))
 
     def prettyIn(self, value):
         if octets.isStringType(value):
@@ -734,7 +744,7 @@ class BitString(base.AbstractSimpleAsn1Item):
             )
 
     def prettyOut(self, value):
-        return '\'%s\'' % str(self)
+        return "'%s'" % self
 
 
 try:
diff --git a/tests/codec/ber/test_decoder.py b/tests/codec/ber/test_decoder.py
index 635ec19..d0eddb9 100644
--- a/tests/codec/ber/test_decoder.py
+++ b/tests/codec/ber/test_decoder.py
@@ -14,7 +14,7 @@ from tests.base import BaseTestCase
 
 from pyasn1.type import tag, namedtype, opentype, univ, char
 from pyasn1.codec.ber import decoder, eoo
-from pyasn1.compat.octets import ints2octs, str2octs, null
+from pyasn1.compat.octets import ints2octs, str2octs, null, isOctetsType
 from pyasn1.error import PyAsn1Error
 
 
@@ -132,6 +132,24 @@ class BooleanDecoderTestCase(BaseTestCase):
             assert 0, 'wrong tagFormat worked out'
 
 
+class BooleanDecoderNatoveModeTestCase(BaseTestCase):
+    def testTrue(self):
+        value, rest = decoder.decode(
+            ints2octs((1, 1, 1)), nativeMode=True
+        )
+        assert not rest
+        assert value is True
+        assert isinstance(value, bool)
+
+    def testFalse(self):
+        value, rest = decoder.decode(
+            ints2octs((1, 1, 0)), nativeMode=True
+        )
+        assert not rest
+        assert value is False
+        assert isinstance(value, bool)
+
+
 class BitStringDecoderTestCase(BaseTestCase):
     def testDefMode(self):
         assert decoder.decode(
@@ -174,6 +192,40 @@ class BitStringDecoderTestCase(BaseTestCase):
             assert 0, 'accepted mis-encoded bit-string constructed out of an integer'
 
 
+class BitStringDecoderNativeModeTestCase(BaseTestCase):
+    def testDefMode(self):
+        value, rest = decoder.decode(
+            ints2octs((3, 3, 1, 169, 138)), nativeMode=True
+        )
+        assert not rest
+        assert value == '101010011000101'
+        assert isinstance(value, str)
+
+    def testIndefMode(self):
+        value, rest = decoder.decode(
+            ints2octs((3, 3, 1, 169, 138)), nativeMode=True
+        )
+        assert not rest
+        assert value == '101010011000101'
+        assert isinstance(value, str)
+
+    def testDefModeChunked(self):
+        value, rest = decoder.decode(
+            ints2octs((35, 8, 3, 2, 0, 169, 3, 2, 1, 138)), nativeMode=True
+        )
+        assert not rest
+        assert value == '101010011000101'
+        assert isinstance(value, str)
+
+    def testIndefModeChunked(self):
+        value, rest = decoder.decode(
+            ints2octs((35, 128, 3, 2, 0, 169, 3, 2, 1, 138, 0, 0)), nativeMode=True
+        )
+        assert not rest
+        assert value == '101010011000101'
+        assert isinstance(value, str)
+
+
 class OctetStringDecoderTestCase(BaseTestCase):
     def testDefMode(self):
         assert decoder.decode(
@@ -213,6 +265,41 @@ class OctetStringDecoderTestCase(BaseTestCase):
             (4, 4, 81, 117, 105, 99, 4, 4, 107, 32, 98, 114, 4, 4, 111, 119, 110, 32, 4, 3, 102, 111, 120, 0, 0)), str2octs(''))
 
 
+class OctetStringDecoderNativeModeTestCase(BaseTestCase):
+    def testDefMode(self):
+        value, rest = decoder.decode(
+            ints2octs((4, 15, 81, 117, 105, 99, 107, 32, 98, 114, 111, 119, 110, 32, 102, 111, 120)), nativeMode=True
+        )
+        assert not rest
+        assert value == str2octs('Quick brown fox')
+        assert isOctetsType(value)
+
+    def testIndefMode(self):
+        value, rest = decoder.decode(
+            ints2octs((36, 128, 4, 15, 81, 117, 105, 99, 107, 32, 98, 114, 111, 119, 110, 32, 102, 111, 120, 0, 0)), nativeMode=True
+        )
+        assert not rest
+        assert value == str2octs('Quick brown fox')
+        assert isOctetsType(value)
+
+    def testDefModeChunked(self):
+        value, rest = decoder.decode(
+            ints2octs(
+                (36, 23, 4, 4, 81, 117, 105, 99, 4, 4, 107, 32, 98, 114, 4, 4, 111, 119, 110, 32, 4, 3, 102, 111, 120)), nativeMode=True
+        )
+        assert not rest
+        assert value == str2octs('Quick brown fox')
+        assert isOctetsType(value)
+
+    def testIndefModeChunked(self):
+        value, rest = decoder.decode(
+            ints2octs((36, 128, 4, 4, 81, 117, 105, 99, 4, 4, 107, 32, 98, 114, 4, 4, 111, 119, 110, 32, 4, 3, 102, 111, 120, 0, 0)), nativeMode=True
+        )
+        assert not rest
+        assert value == str2octs('Quick brown fox')
+        assert isOctetsType(value)
+
+
 class ExpTaggedOctetStringDecoderTestCase(BaseTestCase):
     def setUp(self):
         BaseTestCase.setUp(self)
@@ -287,6 +374,15 @@ class NullDecoderTestCase(BaseTestCase):
             assert 0, 'wrong tagFormat worked out'
 
 
+class NullDecoderNativeModeTestCase(BaseTestCase):
+    def testNull(self):
+        value, rest = decoder.decode(
+            ints2octs((5, 0)), nativeMode=True
+        )
+        assert not rest
+        assert value is None
+
+
 # Useful analysis of OID encoding issues could be found here:
 # http://www.viathinksoft.de/~daniel-marschall/asn.1/oid_facts.html
 class ObjectIdentifierDecoderTestCase(BaseTestCase):