Skip to content

Commit

Permalink
Add support for CRC64NVME when the CRT is available.
Browse files Browse the repository at this point in the history
  • Loading branch information
jonathan343 committed Oct 3, 2024
1 parent cca16a2 commit e40a9f5
Show file tree
Hide file tree
Showing 2 changed files with 44 additions and 2 deletions.
18 changes: 16 additions & 2 deletions botocore/httpchecksum.py
Original file line number Diff line number Diff line change
Expand Up @@ -109,6 +109,19 @@ def digest(self):
return self._int_crc32c.to_bytes(4, byteorder="big")


class CrtCrc64NvmeChecksum(BaseChecksum):
# Note: This class is only used if the CRT is available
def __init__(self):
self._int_crc64nvme = 0

def update(self, chunk):
new_checksum = crt_checksums.crc64nvme(chunk, self._int_crc64nvme)
self._int_crc64nvme = new_checksum & 0xFFFFFFFFFFFFFFFF

def digest(self):
return self._int_crc64nvme.to_bytes(8, byteorder="big")


class Sha1Checksum(BaseChecksum):
def __init__(self):
self._checksum = sha1()
Expand Down Expand Up @@ -465,17 +478,18 @@ def _handle_bytes_response(http_response, response, algorithm):
"sha1": Sha1Checksum,
"sha256": Sha256Checksum,
}
_CRT_CHECKSUM_ALGORITHMS = ["crc32", "crc32c"]
_CRT_CHECKSUM_ALGORITHMS = ["crc32", "crc32c", "crc64nvme"]
if HAS_CRT:
# Use CRT checksum implementations if available
_CRT_CHECKSUM_CLS = {
"crc32": CrtCrc32Checksum,
"crc32c": CrtCrc32cChecksum,
"crc64nvme": CrtCrc64NvmeChecksum,
}
_CHECKSUM_CLS.update(_CRT_CHECKSUM_CLS)
# Validate this list isn't out of sync with _CRT_CHECKSUM_CLS keys
assert all(
name in _CRT_CHECKSUM_ALGORITHMS for name in _CRT_CHECKSUM_CLS.keys()
)
_SUPPORTED_CHECKSUM_ALGORITHMS = list(_CHECKSUM_CLS.keys())
_ALGORITHMS_PRIORITY_LIST = ['crc32c', 'crc32', 'sha1', 'sha256']
_ALGORITHMS_PRIORITY_LIST = ['crc32c', 'crc32', 'crc64nvme', 'sha1', 'sha256']
28 changes: 28 additions & 0 deletions tests/unit/test_httpchecksum.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@
Crc32Checksum,
CrtCrc32cChecksum,
CrtCrc32Checksum,
CrtCrc64NvmeChecksum,
Sha1Checksum,
Sha256Checksum,
StreamingChecksumBody,
Expand Down Expand Up @@ -193,6 +194,24 @@ def test_request_checksum_algorithm_model_no_crt_crc32c_unsupported(self):
str(context.exception),
)

@unittest.skipIf(HAS_CRT, "Error only expected when CRT is not available")
def test_request_checksum_algorithm_model_no_crt_crc64nvme_unsupported(
self,
):
request = self._build_request(b"")
operation_model = self._make_operation_model(
http_checksum={"requestAlgorithmMember": "Algorithm"},
)
params = {"Algorithm": "crc64nvme"}
with self.assertRaises(MissingDependencyException) as context:
resolve_request_checksum_algorithm(
request, operation_model, params
)
self.assertIn(
"Using CRC64NVME requires an additional dependency",
str(context.exception),
)

def test_request_checksum_algorithm_model_legacy_md5(self):
request = self._build_request(b"")
operation_model = self._make_operation_model(required=True)
Expand Down Expand Up @@ -691,6 +710,10 @@ def test_crt_crc32(self):
def test_crt_crc32c(self):
self.assert_base64_checksum(CrtCrc32cChecksum, "yZRlqg==")

@requires_crt()
def test_crt_crc64nvme(self):
self.assert_base64_checksum(CrtCrc64NvmeChecksum, "jSnVw/bqjr4=")


class TestCrtChecksumOverrides(unittest.TestCase):
@requires_crt()
Expand All @@ -703,6 +726,11 @@ def test_crt_crc32c_available(self):
actual_cls = _CHECKSUM_CLS.get("crc32c")
self.assertEqual(actual_cls, CrtCrc32cChecksum)

@requires_crt()
def test_crt_crc64nvme_available(self):
actual_cls = _CHECKSUM_CLS.get("crc64nvme")
self.assertEqual(actual_cls, CrtCrc64NvmeChecksum)


class TestStreamingChecksumBody(unittest.TestCase):
def setUp(self):
Expand Down

0 comments on commit e40a9f5

Please sign in to comment.