RPCv2 Serializer Support (#3396)

botocore/serialize.py

@@ -42,7 +42,9 @@
 import calendar
 import datetime
 import json
+import math
 import re
+import struct
 from xml.etree import ElementTree
 
 from botocore import validate
@@ -431,6 +433,246 @@ def _serialize_type_blob(self, serialized, value, shape, key):
         serialized[key] = self._get_base64(value)
 
 
+class CBORSerializer(Serializer):
+    UNSIGNED_INT_MAJOR_TYPE = 0
+    NEGATIVE_INT_MAJOR_TYPE = 1
+    BLOB_MAJOR_TYPE = 2
+    STRING_MAJOR_TYPE = 3
+    LIST_MAJOR_TYPE = 4
+    MAP_MAJOR_TYPE = 5
+    TAG_MAJOR_TYPE = 6
+    FLOAT_AND_SIMPLE_MAJOR_TYPE = 7
+
+    def _serialize_data_item(self, serialized, value, shape, key=None):
+        method = getattr(self, f'_serialize_type_{shape.type_name}')
+        if method is None:
+            raise ValueError(
+                f"Unrecognized C2J type: {shape.type_name}, unable to "
+                f"serialize request"
+            )
+        method(serialized, value, shape, key)
+
+    def _serialize_type_integer(self, serialized, value, shape, key):
+        if value >= 0:
+            major_type = self.UNSIGNED_INT_MAJOR_TYPE
+        else:
+            major_type = self.NEGATIVE_INT_MAJOR_TYPE
+            # The only differences in serializing negative and positive integers is
+            # that for negative, we set the major type to 1 and set the value to -1
+            # minus the value
+            value = -1 - value
+        additional_info, num_bytes = self._get_additional_info_and_num_bytes(
+            value
+        )
+        initial_byte = self._get_initial_byte(major_type, additional_info)
+        if num_bytes == 0:
+            serialized.extend(initial_byte)
+        else:
+            serialized.extend(initial_byte + value.to_bytes(num_bytes, "big"))
+
+    def _serialize_type_long(self, serialized, value, shape, key):
+        self._serialize_type_integer(serialized, value, shape, key)
+
+    def _serialize_type_blob(self, serialized, value, shape, key):
+        if isinstance(value, str):
+            value = value.encode('utf-8')
+        elif not isinstance(value, (bytes, bytearray)):
+            # We support file-like objects for blobs; these already have been
+            # validated to ensure they have a read method
+            value = value.read()
+        length = len(value)
+        additional_info, num_bytes = self._get_additional_info_and_num_bytes(
+            length
+        )
+        initial_byte = self._get_initial_byte(
+            self.BLOB_MAJOR_TYPE, additional_info
+        )
+        if num_bytes == 0:
+            serialized.extend(initial_byte)
+        else:
+            serialized.extend(initial_byte + length.to_bytes(num_bytes, "big"))
+        serialized.extend(value)
+
+    def _serialize_type_string(self, serialized, value, shape, key):
+        encoded = value.encode('utf-8')
+        length = len(encoded)
+        additional_info, num_bytes = self._get_additional_info_and_num_bytes(
+            length
+        )
+        initial_byte = self._get_initial_byte(self.STRING_MAJOR_TYPE, length)
+        if num_bytes == 0:
+            serialized.extend(initial_byte + encoded)
+        else:
+            serialized.extend(
+                initial_byte + length.to_bytes(num_bytes, "big") + encoded
+            )
+
+    def _serialize_type_list(self, serialized, value, shape, key):
+        length = len(value)
+        additional_info, num_bytes = self._get_additional_info_and_num_bytes(
+            length
+        )
+        initial_byte = self._get_initial_byte(
+            self.LIST_MAJOR_TYPE, additional_info
+        )
+        if num_bytes == 0:
+            serialized.extend(initial_byte)
+        else:
+            serialized.extend(initial_byte + length.to_bytes(num_bytes, "big"))
+        for item in value:
+            self._serialize_data_item(serialized, item, shape.member)
+
+    def _serialize_type_map(self, serialized, value, shape, key):
+        length = len(value)
+        additional_info, num_bytes = self._get_additional_info_and_num_bytes(
+            length
+        )
+        initial_byte = self._get_initial_byte(
+            self.MAP_MAJOR_TYPE, additional_info
+        )
+        if num_bytes == 0:
+            serialized.extend(initial_byte)
+        else:
+            serialized.extend(initial_byte + length.to_bytes(num_bytes, "big"))
+        for key_item, item in value.items():
+            self._serialize_data_item(serialized, key_item, shape.key)
+            self._serialize_data_item(serialized, item, shape.value)
+
+    def _serialize_type_structure(self, serialized, value, shape, key):
+        if key is not None:
+            # For nested structures, we need to serialize the key first
+            self._serialize_data_item(serialized, key, shape.key_shape)
+
+        # Remove `None` values from the dictionary
+        value = {k: v for k, v in value.items() if v is not None}
+
+        map_length = len(value)
+        additional_info, num_bytes = self._get_additional_info_and_num_bytes(
+            map_length
+        )
+        initial_byte = self._get_initial_byte(
+            self.MAP_MAJOR_TYPE, additional_info
+        )
+        if num_bytes == 0:
+            serialized.extend(initial_byte)
+        else:
+            serialized.extend(
+                initial_byte + map_length.to_bytes(num_bytes, "big")
+            )
+
+        members = shape.members
+        for member_key, member_value in value.items():
+            member_shape = members[member_key]
+            if 'name' in member_shape.serialization:
+                member_key = member_shape.serialization['name']
+            if member_value is not None:
+                self._serialize_type_string(serialized, member_key, None, None)
+                self._serialize_data_item(
+                    serialized, member_value, member_shape
+                )
+
+    def _serialize_type_timestamp(self, serialized, value, shape, key):
+        timestamp = self._convert_timestamp_to_str(value)
+        tag = 1  # Use tag 1 for unix timestamp
+        initial_byte = self._get_initial_byte(self.TAG_MAJOR_TYPE, tag)
+        serialized.extend(initial_byte)  # Tagging the timestamp
+        additional_info, num_bytes = self._get_additional_info_and_num_bytes(
+            timestamp
+        )
+
+        if num_bytes == 0:
+            initial_byte = self._get_initial_byte(
+                self.UNSIGNED_INT_MAJOR_TYPE, timestamp
+            )
+            serialized.extend(initial_byte)
+        else:
+            initial_byte = self._get_initial_byte(
+                self.UNSIGNED_INT_MAJOR_TYPE, additional_info
+            )
+            serialized.extend(
+                initial_byte + timestamp.to_bytes(num_bytes, "big")
+            )
+
+    def _serialize_type_float(self, serialized, value, shape, key):
+        if self._is_special_number(value):
+            serialized.extend(
+                self._get_bytes_for_special_numbers(value)
+            )  # Handle special values like NaN or Infinity
+        else:
+            initial_byte = self._get_initial_byte(
+                self.FLOAT_AND_SIMPLE_MAJOR_TYPE, 26
+            )
+            serialized.extend(initial_byte + struct.pack(">f", value))
+
+    def _serialize_type_double(self, serialized, value, shape, key):
+        if self._is_special_number(value):
+            serialized.extend(
+                self._get_bytes_for_special_numbers(value)
+            )  # Handle special values like NaN or Infinity
+        else:
+            initial_byte = self._get_initial_byte(
+                self.FLOAT_AND_SIMPLE_MAJOR_TYPE, 27
+            )
+            serialized.extend(initial_byte + struct.pack(">d", value))
+
+    def _serialize_type_boolean(self, serialized, value, shape, key):
+        additional_info = 21 if value else 20
+        serialized.extend(
+            self._get_initial_byte(
+                self.FLOAT_AND_SIMPLE_MAJOR_TYPE, additional_info
+            )
+        )
+
+    def _get_additional_info_and_num_bytes(self, value):
+        # Values under 24 can be stored in the initial byte and don't need further
+        # encoding
+        if value < 24:
+            return value, 0
+        # Values between 24 and 255 (inclusive) can be stored in 1 byte and
+        # correspond to additional info 24
+        elif value < 256:
+            return 24, 1
+        # Values up to 65535 can be stored in two bytes and correspond to additional
+        # info 25
+        elif value < 65536:
+            return 25, 2
+        # Values up to 4294967296 can be stored in four bytes and correspond to
+        # additional info 26
+        elif value < 4294967296:
+            return 26, 4
+        # The maximum number of bytes in a definite length data items is 8 which
+        # to additional info 27
+        else:
+            return 27, 8
+
+    def _get_initial_byte(self, major_type, additional_info):
+        # The highest order three bits are the major type, so we need to bitshift the
+        # major type by 5
+        major_type_bytes = major_type << 5
+        return (major_type_bytes | additional_info).to_bytes(1, "big")
+
+    def _is_special_number(self, value):
+        return any(
+            [
+                value == float('inf'),
+                value == float('-inf'),
+                math.isnan(value),
+            ]
+        )
+
+    def _get_bytes_for_special_numbers(self, value):
+        additional_info = 25
+        initial_byte = self._get_initial_byte(
+            self.FLOAT_AND_SIMPLE_MAJOR_TYPE, additional_info
+        )
+        if value == float('inf'):
+            return initial_byte + struct.pack(">H", 0x7C00)
+        elif value == float('-inf'):
+            return initial_byte + struct.pack(">H", 0xFC00)
+        elif math.isnan(value):
+            return initial_byte + struct.pack(">H", 0x7E00)
+
+
 class BaseRestSerializer(Serializer):
     """Base class for rest protocols.
 
@@ -669,6 +911,45 @@ def _convert_header_value(self, shape, value):
             return value
 
 
+class BaseRpcV2Serializer(Serializer):
+    """Base class for RPCv2 protocols.
+
+    The only variance between the various RPCv2 protocols is the
+    way that the body is serialized.  All other aspects (headers, uri, etc.)
+    are the same and logic for serializing those aspects lives here.
+
+    Subclasses must implement the ``_serialize_body_params``  and
+    ``_serialize_headers`` methods.
+
+    """
+
+    def serialize_to_request(self, parameters, operation_model):
+        serialized = self._create_default_request()
+        service_name = operation_model.service_model.metadata['targetPrefix']
+        operation_name = operation_model.name
+        serialized['url_path'] = (
+            f'/service/{service_name}/operation/{operation_name}'
+        )
+
+        input_shape = operation_model.input_shape
+        if input_shape is not None:
+            self._serialize_payload(parameters, serialized, input_shape)
+
+        self._serialize_headers(serialized, operation_model)
+
+        return serialized
+
+    def _serialize_payload(self, parameters, serialized, shape):
+        body_payload = self._serialize_body_params(parameters, shape)
+        serialized['body'] = body_payload
+
+    def _serialize_headers(self, serialized, operation_model):
+        raise NotImplementedError("_serialize_headers")
+
+    def _serialize_body_params(self, parameters, shape):
+        raise NotImplementedError("_serialize_body_params")
+
+
 class RestJSONSerializer(BaseRestSerializer, JSONSerializer):
     def _serialize_empty_body(self):
         return b'{}'
@@ -802,10 +1083,35 @@ def _default_serialize(self, xmlnode, params, shape, name):
         node.text = str(params)
 
 
+class RpcV2CBORSerializer(BaseRpcV2Serializer, CBORSerializer):
+    TIMESTAMP_FORMAT = 'unixtimestamp'
+
+    def _serialize_body_params(self, parameters, input_shape):
+        body = bytearray()
+        self._serialize_data_item(body, parameters, input_shape)
+        return bytes(body)
+
+    def _serialize_headers(self, serialized, operation_model):
+        serialized['headers']['smithy-protocol'] = 'rpc-v2-cbor'
+
+        if operation_model.has_event_stream_output:
+            header_val = 'application/vnd.amazon.eventstream'
+        else:
+            header_val = 'application/cbor'
+
+        has_body = serialized['body'] != b''
+        has_content_type = has_header('Content-Type', serialized['headers'])
+
+        serialized['headers']['Accept'] = header_val
+        if not has_content_type and has_body:
+            serialized['headers']['Content-Type'] = header_val
+
+
 SERIALIZERS = {
     'ec2': EC2Serializer,
     'query': QuerySerializer,
     'json': JSONSerializer,
     'rest-json': RestJSONSerializer,
     'rest-xml': RestXMLSerializer,
+    'smithy-rpc-v2-cbor': RpcV2CBORSerializer,
 }