eventstream.py 20 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616
  1. # Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License"). You
  4. # may not use this file except in compliance with the License. A copy of
  5. # the License is located at
  6. #
  7. # http://aws.amazon.com/apache2.0/
  8. #
  9. # or in the "license" file accompanying this file. This file is
  10. # distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
  11. # ANY KIND, either express or implied. See the License for the specific
  12. # language governing permissions and limitations under the License.
  13. """Binary Event Stream Decoding """
  14. from binascii import crc32
  15. from struct import unpack
  16. from botocore.exceptions import EventStreamError
  17. # byte length of the prelude (total_length + header_length + prelude_crc)
  18. _PRELUDE_LENGTH = 12
  19. _MAX_HEADERS_LENGTH = 128 * 1024 # 128 Kb
  20. _MAX_PAYLOAD_LENGTH = 16 * 1024 ** 2 # 16 Mb
  21. class ParserError(Exception):
  22. """Base binary flow encoding parsing exception. """
  23. pass
  24. class DuplicateHeader(ParserError):
  25. """Duplicate header found in the event. """
  26. def __init__(self, header):
  27. message = 'Duplicate header present: "%s"' % header
  28. super(DuplicateHeader, self).__init__(message)
  29. class InvalidHeadersLength(ParserError):
  30. """Headers length is longer than the maximum. """
  31. def __init__(self, length):
  32. message = 'Header length of %s exceeded the maximum of %s' % (
  33. length, _MAX_HEADERS_LENGTH
  34. )
  35. super(InvalidHeadersLength, self).__init__(message)
  36. class InvalidPayloadLength(ParserError):
  37. """Payload length is longer than the maximum. """
  38. def __init__(self, length):
  39. message = 'Payload length of %s exceeded the maximum of %s' % (
  40. length, _MAX_PAYLOAD_LENGTH
  41. )
  42. super(InvalidPayloadLength, self).__init__(message)
  43. class ChecksumMismatch(ParserError):
  44. """Calculated checksum did not match the expected checksum. """
  45. def __init__(self, expected, calculated):
  46. message = 'Checksum mismatch: expected 0x%08x, calculated 0x%08x' % (
  47. expected, calculated
  48. )
  49. super(ChecksumMismatch, self).__init__(message)
  50. class NoInitialResponseError(ParserError):
  51. """An event of type initial-response was not received.
  52. This exception is raised when the event stream produced no events or
  53. the first event in the stream was not of the initial-response type.
  54. """
  55. def __init__(self):
  56. message = 'First event was not of the initial-response type'
  57. super(NoInitialResponseError, self).__init__(message)
  58. class DecodeUtils(object):
  59. """Unpacking utility functions used in the decoder.
  60. All methods on this class take raw bytes and return a tuple containing
  61. the value parsed from the bytes and the number of bytes consumed to parse
  62. that value.
  63. """
  64. UINT8_BYTE_FORMAT = '!B'
  65. UINT16_BYTE_FORMAT = '!H'
  66. UINT32_BYTE_FORMAT = '!I'
  67. INT8_BYTE_FORMAT = '!b'
  68. INT16_BYTE_FORMAT = '!h'
  69. INT32_BYTE_FORMAT = '!i'
  70. INT64_BYTE_FORMAT = '!q'
  71. PRELUDE_BYTE_FORMAT = '!III'
  72. # uint byte size to unpack format
  73. UINT_BYTE_FORMAT = {
  74. 1: UINT8_BYTE_FORMAT,
  75. 2: UINT16_BYTE_FORMAT,
  76. 4: UINT32_BYTE_FORMAT,
  77. }
  78. @staticmethod
  79. def unpack_true(data):
  80. """This method consumes none of the provided bytes and returns True.
  81. :type data: bytes
  82. :param data: The bytes to parse from. This is ignored in this method.
  83. :rtype: tuple
  84. :rtype: (bool, int)
  85. :returns: The tuple (True, 0)
  86. """
  87. return True, 0
  88. @staticmethod
  89. def unpack_false(data):
  90. """This method consumes none of the provided bytes and returns False.
  91. :type data: bytes
  92. :param data: The bytes to parse from. This is ignored in this method.
  93. :rtype: tuple
  94. :rtype: (bool, int)
  95. :returns: The tuple (False, 0)
  96. """
  97. return False, 0
  98. @staticmethod
  99. def unpack_uint8(data):
  100. """Parse an unsigned 8-bit integer from the bytes.
  101. :type data: bytes
  102. :param data: The bytes to parse from.
  103. :rtype: (int, int)
  104. :returns: A tuple containing the (parsed integer value, bytes consumed)
  105. """
  106. value = unpack(DecodeUtils.UINT8_BYTE_FORMAT, data[:1])[0]
  107. return value, 1
  108. @staticmethod
  109. def unpack_uint32(data):
  110. """Parse an unsigned 32-bit integer from the bytes.
  111. :type data: bytes
  112. :param data: The bytes to parse from.
  113. :rtype: (int, int)
  114. :returns: A tuple containing the (parsed integer value, bytes consumed)
  115. """
  116. value = unpack(DecodeUtils.UINT32_BYTE_FORMAT, data[:4])[0]
  117. return value, 4
  118. @staticmethod
  119. def unpack_int8(data):
  120. """Parse a signed 8-bit integer from the bytes.
  121. :type data: bytes
  122. :param data: The bytes to parse from.
  123. :rtype: (int, int)
  124. :returns: A tuple containing the (parsed integer value, bytes consumed)
  125. """
  126. value = unpack(DecodeUtils.INT8_BYTE_FORMAT, data[:1])[0]
  127. return value, 1
  128. @staticmethod
  129. def unpack_int16(data):
  130. """Parse a signed 16-bit integer from the bytes.
  131. :type data: bytes
  132. :param data: The bytes to parse from.
  133. :rtype: tuple
  134. :rtype: (int, int)
  135. :returns: A tuple containing the (parsed integer value, bytes consumed)
  136. """
  137. value = unpack(DecodeUtils.INT16_BYTE_FORMAT, data[:2])[0]
  138. return value, 2
  139. @staticmethod
  140. def unpack_int32(data):
  141. """Parse a signed 32-bit integer from the bytes.
  142. :type data: bytes
  143. :param data: The bytes to parse from.
  144. :rtype: tuple
  145. :rtype: (int, int)
  146. :returns: A tuple containing the (parsed integer value, bytes consumed)
  147. """
  148. value = unpack(DecodeUtils.INT32_BYTE_FORMAT, data[:4])[0]
  149. return value, 4
  150. @staticmethod
  151. def unpack_int64(data):
  152. """Parse a signed 64-bit integer from the bytes.
  153. :type data: bytes
  154. :param data: The bytes to parse from.
  155. :rtype: tuple
  156. :rtype: (int, int)
  157. :returns: A tuple containing the (parsed integer value, bytes consumed)
  158. """
  159. value = unpack(DecodeUtils.INT64_BYTE_FORMAT, data[:8])[0]
  160. return value, 8
  161. @staticmethod
  162. def unpack_byte_array(data, length_byte_size=2):
  163. """Parse a variable length byte array from the bytes.
  164. The bytes are expected to be in the following format:
  165. [ length ][0 ... length bytes]
  166. where length is an unsigned integer represented in the smallest number
  167. of bytes to hold the maximum length of the array.
  168. :type data: bytes
  169. :param data: The bytes to parse from.
  170. :type length_byte_size: int
  171. :param length_byte_size: The byte size of the preceeding integer that
  172. represents the length of the array. Supported values are 1, 2, and 4.
  173. :rtype: (bytes, int)
  174. :returns: A tuple containing the (parsed byte array, bytes consumed).
  175. """
  176. uint_byte_format = DecodeUtils.UINT_BYTE_FORMAT[length_byte_size]
  177. length = unpack(uint_byte_format, data[:length_byte_size])[0]
  178. bytes_end = length + length_byte_size
  179. array_bytes = data[length_byte_size:bytes_end]
  180. return array_bytes, bytes_end
  181. @staticmethod
  182. def unpack_utf8_string(data, length_byte_size=2):
  183. """Parse a variable length utf-8 string from the bytes.
  184. The bytes are expected to be in the following format:
  185. [ length ][0 ... length bytes]
  186. where length is an unsigned integer represented in the smallest number
  187. of bytes to hold the maximum length of the array and the following
  188. bytes are a valid utf-8 string.
  189. :type data: bytes
  190. :param bytes: The bytes to parse from.
  191. :type length_byte_size: int
  192. :param length_byte_size: The byte size of the preceeding integer that
  193. represents the length of the array. Supported values are 1, 2, and 4.
  194. :rtype: (str, int)
  195. :returns: A tuple containing the (utf-8 string, bytes consumed).
  196. """
  197. array_bytes, consumed = DecodeUtils.unpack_byte_array(
  198. data, length_byte_size)
  199. return array_bytes.decode('utf-8'), consumed
  200. @staticmethod
  201. def unpack_uuid(data):
  202. """Parse a 16-byte uuid from the bytes.
  203. :type data: bytes
  204. :param data: The bytes to parse from.
  205. :rtype: (bytes, int)
  206. :returns: A tuple containing the (uuid bytes, bytes consumed).
  207. """
  208. return data[:16], 16
  209. @staticmethod
  210. def unpack_prelude(data):
  211. """Parse the prelude for an event stream message from the bytes.
  212. The prelude for an event stream message has the following format:
  213. [total_length][header_length][prelude_crc]
  214. where each field is an unsigned 32-bit integer.
  215. :rtype: ((int, int, int), int)
  216. :returns: A tuple of ((total_length, headers_length, prelude_crc),
  217. consumed)
  218. """
  219. return (unpack(DecodeUtils.PRELUDE_BYTE_FORMAT, data), _PRELUDE_LENGTH)
  220. def _validate_checksum(data, checksum, crc=0):
  221. # To generate the same numeric value across all Python versions and
  222. # platforms use crc32(data) & 0xffffffff.
  223. computed_checksum = crc32(data, crc) & 0xFFFFFFFF
  224. if checksum != computed_checksum:
  225. raise ChecksumMismatch(checksum, computed_checksum)
  226. class MessagePrelude(object):
  227. """Represents the prelude of an event stream message. """
  228. def __init__(self, total_length, headers_length, crc):
  229. self.total_length = total_length
  230. self.headers_length = headers_length
  231. self.crc = crc
  232. @property
  233. def payload_length(self):
  234. """Calculates the total payload length.
  235. The extra minus 4 bytes is for the message CRC.
  236. :rtype: int
  237. :returns: The total payload length.
  238. """
  239. return self.total_length - self.headers_length - _PRELUDE_LENGTH - 4
  240. @property
  241. def payload_end(self):
  242. """Calculates the byte offset for the end of the message payload.
  243. The extra minus 4 bytes is for the message CRC.
  244. :rtype: int
  245. :returns: The byte offset from the beginning of the event stream
  246. message to the end of the payload.
  247. """
  248. return self.total_length - 4
  249. @property
  250. def headers_end(self):
  251. """Calculates the byte offset for the end of the message headers.
  252. :rtype: int
  253. :returns: The byte offset from the beginning of the event stream
  254. message to the end of the headers.
  255. """
  256. return _PRELUDE_LENGTH + self.headers_length
  257. class EventStreamMessage(object):
  258. """Represents an event stream message. """
  259. def __init__(self, prelude, headers, payload, crc):
  260. self.prelude = prelude
  261. self.headers = headers
  262. self.payload = payload
  263. self.crc = crc
  264. def to_response_dict(self, status_code=200):
  265. message_type = self.headers.get(':message-type')
  266. if message_type == 'error' or message_type == 'exception':
  267. status_code = 400
  268. return {
  269. 'status_code': status_code,
  270. 'headers': self.headers,
  271. 'body': self.payload
  272. }
  273. class EventStreamHeaderParser(object):
  274. """ Parses the event headers from an event stream message.
  275. Expects all of the header data upfront and creates a dictionary of headers
  276. to return. This object can be reused multiple times to parse the headers
  277. from multiple event stream messages.
  278. """
  279. # Maps header type to appropriate unpacking function
  280. # These unpacking functions return the value and the amount unpacked
  281. _HEADER_TYPE_MAP = {
  282. # boolean_true
  283. 0: DecodeUtils.unpack_true,
  284. # boolean_false
  285. 1: DecodeUtils.unpack_false,
  286. # byte
  287. 2: DecodeUtils.unpack_int8,
  288. # short
  289. 3: DecodeUtils.unpack_int16,
  290. # integer
  291. 4: DecodeUtils.unpack_int32,
  292. # long
  293. 5: DecodeUtils.unpack_int64,
  294. # byte_array
  295. 6: DecodeUtils.unpack_byte_array,
  296. # string
  297. 7: DecodeUtils.unpack_utf8_string,
  298. # timestamp
  299. 8: DecodeUtils.unpack_int64,
  300. # uuid
  301. 9: DecodeUtils.unpack_uuid,
  302. }
  303. def __init__(self):
  304. self._data = None
  305. def parse(self, data):
  306. """Parses the event stream headers from an event stream message.
  307. :type data: bytes
  308. :param data: The bytes that correspond to the headers section of an
  309. event stream message.
  310. :rtype: dict
  311. :returns: A dicionary of header key, value pairs.
  312. """
  313. self._data = data
  314. return self._parse_headers()
  315. def _parse_headers(self):
  316. headers = {}
  317. while self._data:
  318. name, value = self._parse_header()
  319. if name in headers:
  320. raise DuplicateHeader(name)
  321. headers[name] = value
  322. return headers
  323. def _parse_header(self):
  324. name = self._parse_name()
  325. value = self._parse_value()
  326. return name, value
  327. def _parse_name(self):
  328. name, consumed = DecodeUtils.unpack_utf8_string(self._data, 1)
  329. self._advance_data(consumed)
  330. return name
  331. def _parse_type(self):
  332. type, consumed = DecodeUtils.unpack_uint8(self._data)
  333. self._advance_data(consumed)
  334. return type
  335. def _parse_value(self):
  336. header_type = self._parse_type()
  337. value_unpacker = self._HEADER_TYPE_MAP[header_type]
  338. value, consumed = value_unpacker(self._data)
  339. self._advance_data(consumed)
  340. return value
  341. def _advance_data(self, consumed):
  342. self._data = self._data[consumed:]
  343. class EventStreamBuffer(object):
  344. """Streaming based event stream buffer
  345. A buffer class that wraps bytes from an event stream providing parsed
  346. messages as they become available via an iterable interface.
  347. """
  348. def __init__(self):
  349. self._data = b''
  350. self._prelude = None
  351. self._header_parser = EventStreamHeaderParser()
  352. def add_data(self, data):
  353. """Add data to the buffer.
  354. :type data: bytes
  355. :param data: The bytes to add to the buffer to be used when parsing
  356. """
  357. self._data += data
  358. def _validate_prelude(self, prelude):
  359. if prelude.headers_length > _MAX_HEADERS_LENGTH:
  360. raise InvalidHeadersLength(prelude.headers_length)
  361. if prelude.payload_length > _MAX_PAYLOAD_LENGTH:
  362. raise InvalidPayloadLength(prelude.payload_length)
  363. def _parse_prelude(self):
  364. prelude_bytes = self._data[:_PRELUDE_LENGTH]
  365. raw_prelude, _ = DecodeUtils.unpack_prelude(prelude_bytes)
  366. prelude = MessagePrelude(*raw_prelude)
  367. self._validate_prelude(prelude)
  368. # The minus 4 removes the prelude crc from the bytes to be checked
  369. _validate_checksum(prelude_bytes[:_PRELUDE_LENGTH-4], prelude.crc)
  370. return prelude
  371. def _parse_headers(self):
  372. header_bytes = self._data[_PRELUDE_LENGTH:self._prelude.headers_end]
  373. return self._header_parser.parse(header_bytes)
  374. def _parse_payload(self):
  375. prelude = self._prelude
  376. payload_bytes = self._data[prelude.headers_end:prelude.payload_end]
  377. return payload_bytes
  378. def _parse_message_crc(self):
  379. prelude = self._prelude
  380. crc_bytes = self._data[prelude.payload_end:prelude.total_length]
  381. message_crc, _ = DecodeUtils.unpack_uint32(crc_bytes)
  382. return message_crc
  383. def _parse_message_bytes(self):
  384. # The minus 4 includes the prelude crc to the bytes to be checked
  385. message_bytes = self._data[_PRELUDE_LENGTH-4:self._prelude.payload_end]
  386. return message_bytes
  387. def _validate_message_crc(self):
  388. message_crc = self._parse_message_crc()
  389. message_bytes = self._parse_message_bytes()
  390. _validate_checksum(message_bytes, message_crc, crc=self._prelude.crc)
  391. return message_crc
  392. def _parse_message(self):
  393. crc = self._validate_message_crc()
  394. headers = self._parse_headers()
  395. payload = self._parse_payload()
  396. message = EventStreamMessage(self._prelude, headers, payload, crc)
  397. self._prepare_for_next_message()
  398. return message
  399. def _prepare_for_next_message(self):
  400. # Advance the data and reset the current prelude
  401. self._data = self._data[self._prelude.total_length:]
  402. self._prelude = None
  403. def next(self):
  404. """Provides the next available message parsed from the stream
  405. :rtype: EventStreamMessage
  406. :returns: The next event stream message
  407. """
  408. if len(self._data) < _PRELUDE_LENGTH:
  409. raise StopIteration()
  410. if self._prelude is None:
  411. self._prelude = self._parse_prelude()
  412. if len(self._data) < self._prelude.total_length:
  413. raise StopIteration()
  414. return self._parse_message()
  415. def __next__(self):
  416. return self.next()
  417. def __iter__(self):
  418. return self
  419. class EventStream(object):
  420. """Wrapper class for an event stream body.
  421. This wraps the underlying streaming body, parsing it for individual events
  422. and yielding them as they come available through the iterator interface.
  423. The following example uses the S3 select API to get structured data out of
  424. an object stored in S3 using an event stream.
  425. **Example:**
  426. ::
  427. from botocore.session import Session
  428. s3 = Session().create_client('s3')
  429. response = s3.select_object_content(
  430. Bucket='bucketname',
  431. Key='keyname',
  432. ExpressionType='SQL',
  433. RequestProgress={'Enabled': True},
  434. Expression="SELECT * FROM S3Object s",
  435. InputSerialization={'CSV': {}},
  436. OutputSerialization={'CSV': {}},
  437. )
  438. # This is the event stream in the response
  439. event_stream = response['Payload']
  440. end_event_received = False
  441. with open('output', 'wb') as f:
  442. # Iterate over events in the event stream as they come
  443. for event in event_stream:
  444. # If we received a records event, write the data to a file
  445. if 'Records' in event:
  446. data = event['Records']['Payload']
  447. f.write(data)
  448. # If we received a progress event, print the details
  449. elif 'Progress' in event:
  450. print(event['Progress']['Details'])
  451. # End event indicates that the request finished successfully
  452. elif 'End' in event:
  453. print('Result is complete')
  454. end_event_received = True
  455. if not end_event_received:
  456. raise Exception("End event not received, request incomplete.")
  457. """
  458. def __init__(self, raw_stream, output_shape, parser, operation_name):
  459. self._raw_stream = raw_stream
  460. self._output_shape = output_shape
  461. self._operation_name = operation_name
  462. self._parser = parser
  463. self._event_generator = self._create_raw_event_generator()
  464. def __iter__(self):
  465. for event in self._event_generator:
  466. parsed_event = self._parse_event(event)
  467. if parsed_event:
  468. yield parsed_event
  469. def _create_raw_event_generator(self):
  470. event_stream_buffer = EventStreamBuffer()
  471. for chunk in self._raw_stream.stream():
  472. event_stream_buffer.add_data(chunk)
  473. for event in event_stream_buffer:
  474. yield event
  475. def _parse_event(self, event):
  476. response_dict = event.to_response_dict()
  477. parsed_response = self._parser.parse(response_dict, self._output_shape)
  478. if response_dict['status_code'] == 200:
  479. return parsed_response
  480. else:
  481. raise EventStreamError(parsed_response, self._operation_name)
  482. def get_initial_response(self):
  483. try:
  484. initial_event = next(self._event_generator)
  485. event_type = initial_event.headers.get(':event-type')
  486. if event_type == 'initial-response':
  487. return initial_event
  488. except StopIteration:
  489. pass
  490. raise NoInitialResponseError()
  491. def close(self):
  492. """Closes the underlying streaming body. """
  493. self._raw_stream.close()