diff --git a/pyiceberg/encryption/stream.py b/pyiceberg/encryption/stream.py new file mode 100644 index 0000000000..39a88bb098 --- /dev/null +++ b/pyiceberg/encryption/stream.py @@ -0,0 +1,118 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +"""AES GCM Stream (AGS1) encryption, used to encrypt manifests and manifest lists. + +An AGS1 stream is an 8 byte header followed by a sequence of AES-GCM blocks:: + + "AGS1" || plain_block_size (4 bytes, little endian) + nonce || ciphertext || tag (block 0, up to plain_block_size of plaintext) + nonce || ciphertext || tag (block 1..n, the last of which may be shorter) + +Each block authenticates `aad_prefix || block_index` as additional data, so blocks cannot +be reordered or moved between files. + +The spec gives the last block a non-zero length, which makes a bare header its encoding of an +empty file. Java, iceberg-rust and PyIceberg all require at least one block instead, so an +empty file is a header followed by a single empty block. apache/iceberg#18219 tracks which of +the two forms writers should produce. + +`_encrypt_ags1` and `_decrypt_ags1` work on whole files, matching how PyIceberg reads and writes +Avro, so their implementation can move to pyiceberg-core without changing callers. They stay +private until the manifest reader and writer use them. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from pyiceberg.encryption.ciphers import AesGcmCipher, SecureKey + +if TYPE_CHECKING: + from pyiceberg.encryption.key_metadata import StandardKeyMetadata + +_GCM_STREAM_MAGIC = b"AGS1" +_PLAIN_BLOCK_SIZE = 1024 * 1024 +_GCM_STREAM_HEADER = _GCM_STREAM_MAGIC + _PLAIN_BLOCK_SIZE.to_bytes(4, "little") +_GCM_STREAM_HEADER_LENGTH = len(_GCM_STREAM_HEADER) +_BLOCK_OVERHEAD = AesGcmCipher.NONCE_LENGTH + AesGcmCipher.TAG_LENGTH +_CIPHER_BLOCK_SIZE = _PLAIN_BLOCK_SIZE + _BLOCK_OVERHEAD +_MIN_STREAM_LENGTH = _GCM_STREAM_HEADER_LENGTH + _BLOCK_OVERHEAD + + +def _encrypt_ags1(plaintext: bytes, key_metadata: StandardKeyMetadata) -> bytes: + """Encrypt `plaintext` into an AGS1 stream. + + An empty plaintext is written as the header plus one empty block, as Java and iceberg-rust do. + + Args: + plaintext (bytes): The file contents to encrypt. + key_metadata (StandardKeyMetadata): The file's key and AAD prefix. Its `file_length` is ignored, and + callers must record the length of the returned stream as the file's `file_length`. + """ + cipher = AesGcmCipher(SecureKey(key_metadata.encryption_key)) + starts = range(0, len(plaintext), _PLAIN_BLOCK_SIZE) or range(1) + return _GCM_STREAM_HEADER + b"".join( + cipher.encrypt(plaintext[start : start + _PLAIN_BLOCK_SIZE], _stream_block_aad(key_metadata.aad_prefix, index)) + for index, start in enumerate(starts) + ) + + +def _decrypt_ags1(data: bytes, key_metadata: StandardKeyMetadata) -> bytes: + """Decrypt an AGS1 stream, authenticating every block. + + Args: + data (bytes): The encrypted file contents. + key_metadata (StandardKeyMetadata): The file's key, AAD prefix and `file_length`. The length is required and + must match `data`: dropping trailing blocks leaves every remaining block valid, so only the trusted length + from key metadata exposes a truncated file. + """ + if key_metadata.file_length is None: + raise ValueError("Cannot decrypt AGS1 stream: key metadata has no file_length") + + if len(data) != key_metadata.file_length: + raise ValueError(f"Invalid AGS1 stream: key metadata expects {key_metadata.file_length} bytes, got {len(data)}") + + if len(data) < _MIN_STREAM_LENGTH: + raise ValueError(f"Invalid AGS1 stream: expected at least {_MIN_STREAM_LENGTH} bytes, got {len(data)}") + + _check_stream_header(data) + + last_block_length = (len(data) - _GCM_STREAM_HEADER_LENGTH) % _CIPHER_BLOCK_SIZE + if 0 < last_block_length < _BLOCK_OVERHEAD: + raise ValueError(f"Truncated AGS1 stream: last block is {last_block_length} bytes, expected at least {_BLOCK_OVERHEAD}") + + cipher = AesGcmCipher(SecureKey(key_metadata.encryption_key)) + starts = range(_GCM_STREAM_HEADER_LENGTH, len(data), _CIPHER_BLOCK_SIZE) + return b"".join( + cipher.decrypt(data[start : start + _CIPHER_BLOCK_SIZE], _stream_block_aad(key_metadata.aad_prefix, index)) + for index, start in enumerate(starts) + ) + + +def _stream_block_aad(aad_prefix: bytes | None, block_index: int) -> bytes: + """Return the additional authenticated data for the block at `block_index`.""" + return (aad_prefix or b"") + block_index.to_bytes(4, "little") + + +def _check_stream_header(data: bytes) -> None: + """Check that `data` starts with an AGS1 header declaring the supported block size.""" + if (magic := data[: len(_GCM_STREAM_MAGIC)]) != _GCM_STREAM_MAGIC: + raise ValueError(f"Invalid AGS1 header: magic {magic!r} does not match {_GCM_STREAM_MAGIC!r}") + + plain_block_size = int.from_bytes(data[len(_GCM_STREAM_MAGIC) : _GCM_STREAM_HEADER_LENGTH], "little") + if plain_block_size != _PLAIN_BLOCK_SIZE: + raise ValueError(f"Unsupported AGS1 block size: {plain_block_size} (expected {_PLAIN_BLOCK_SIZE})") diff --git a/tests/encryption/test_stream.py b/tests/encryption/test_stream.py new file mode 100644 index 0000000000..5195dfe52e --- /dev/null +++ b/tests/encryption/test_stream.py @@ -0,0 +1,167 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +import pytest + +from pyiceberg.encryption.key_metadata import StandardKeyMetadata +from pyiceberg.encryption.stream import ( + _BLOCK_OVERHEAD, + _CIPHER_BLOCK_SIZE, + _GCM_STREAM_HEADER, + _GCM_STREAM_HEADER_LENGTH, + _MIN_STREAM_LENGTH, + _PLAIN_BLOCK_SIZE, + _decrypt_ags1, + _encrypt_ags1, + _stream_block_aad, +) + +KEY = b"0123456789012345" +AAD_PREFIX = b"0123456789abcdef" +KEY_METADATA = StandardKeyMetadata(encryption_key=KEY, aad_prefix=AAD_PREFIX) + +# The header a Java `AesGcmOutputStream` writes: "AGS1" then 1 MiB as a little-endian int32. +JAVA_HEADER = b"AGS1\x00\x00\x10\x00" + + +def decrypt(stream: bytes, key: bytes = KEY, aad_prefix: bytes | None = AAD_PREFIX) -> bytes: + """Decrypt `stream`, trusting its own length as the key metadata's `file_length`.""" + return _decrypt_ags1(stream, StandardKeyMetadata(encryption_key=key, aad_prefix=aad_prefix, file_length=len(stream))) + + +def patterned_plaintext(length: int) -> bytes: + """Return a plaintext whose byte `i` is `i % 251`, so a misordered block changes the bytes.""" + return (bytes(range(251)) * (length // 251 + 1))[:length] + + +def test_format_constants() -> None: + assert _GCM_STREAM_HEADER == JAVA_HEADER + assert _PLAIN_BLOCK_SIZE == 1024 * 1024 + assert _GCM_STREAM_HEADER_LENGTH == 8 + assert _BLOCK_OVERHEAD == 28 + assert _CIPHER_BLOCK_SIZE == _PLAIN_BLOCK_SIZE + _BLOCK_OVERHEAD + assert _MIN_STREAM_LENGTH == 36 + + +@pytest.mark.parametrize( + "block_index, expected", + [(0, b"\x00\x00\x00\x00"), (1, b"\x01\x00\x00\x00"), (258, b"\x02\x01\x00\x00"), (2**32 - 1, b"\xff\xff\xff\xff")], +) +def test_stream_block_aad_encodes_the_index_little_endian(block_index: int, expected: bytes) -> None: + assert _stream_block_aad(None, block_index) == expected + assert _stream_block_aad(b"", block_index) == expected + assert _stream_block_aad(AAD_PREFIX, block_index) == AAD_PREFIX + expected + + +@pytest.mark.parametrize("aad_prefix", [AAD_PREFIX, None]) +@pytest.mark.parametrize("plaintext_length", [0, 1, 100, _PLAIN_BLOCK_SIZE, _PLAIN_BLOCK_SIZE + 7, 2 * _PLAIN_BLOCK_SIZE]) +def test_round_trip(plaintext_length: int, aad_prefix: bytes | None) -> None: + plaintext = patterned_plaintext(plaintext_length) + + stream = _encrypt_ags1(plaintext, StandardKeyMetadata(encryption_key=KEY, aad_prefix=aad_prefix)) + + assert stream.startswith(JAVA_HEADER) + assert decrypt(stream, aad_prefix=aad_prefix) == plaintext + + +@pytest.mark.parametrize( + "plaintext_length, encrypted_length", + [ + (0, _MIN_STREAM_LENGTH), + (100, _GCM_STREAM_HEADER_LENGTH + _BLOCK_OVERHEAD + 100), + (_PLAIN_BLOCK_SIZE, _GCM_STREAM_HEADER_LENGTH + _CIPHER_BLOCK_SIZE), + (_PLAIN_BLOCK_SIZE + 7, _GCM_STREAM_HEADER_LENGTH + _CIPHER_BLOCK_SIZE + _BLOCK_OVERHEAD + 7), + (2 * _PLAIN_BLOCK_SIZE, _GCM_STREAM_HEADER_LENGTH + 2 * _CIPHER_BLOCK_SIZE), + ], +) +def test_encrypted_length(plaintext_length: int, encrypted_length: int) -> None: + """An empty file gets one empty block, and a block-aligned file gets no trailing empty block, as in Java.""" + assert len(_encrypt_ags1(bytes(plaintext_length), KEY_METADATA)) == encrypted_length + + +def test_decrypt_requires_file_length() -> None: + stream = _encrypt_ags1(b"data", KEY_METADATA) + + with pytest.raises(ValueError, match="key metadata has no file_length"): + _decrypt_ags1(stream, KEY_METADATA) + + +def test_a_truncated_stream_still_authenticates() -> None: + """Dropping a trailing block leaves every remaining block valid, so the length must come from key metadata.""" + stream = _encrypt_ags1(patterned_plaintext(2 * _PLAIN_BLOCK_SIZE), KEY_METADATA) + + truncated = stream[: _GCM_STREAM_HEADER_LENGTH + _CIPHER_BLOCK_SIZE] + + assert decrypt(truncated) == patterned_plaintext(_PLAIN_BLOCK_SIZE) + + +@pytest.mark.parametrize("change", ["truncate", "extend"]) +def test_decrypt_rejects_a_length_other_than_the_trusted_length(change: str) -> None: + stream = _encrypt_ags1(patterned_plaintext(2 * _PLAIN_BLOCK_SIZE), KEY_METADATA) + key_metadata = StandardKeyMetadata(encryption_key=KEY, aad_prefix=AAD_PREFIX, file_length=len(stream)) + changed = stream[: _GCM_STREAM_HEADER_LENGTH + _CIPHER_BLOCK_SIZE] if change == "truncate" else stream + b"\x00" + + with pytest.raises(ValueError, match=f"key metadata expects {len(stream)} bytes, got {len(changed)}"): + _decrypt_ags1(changed, key_metadata) + + +@pytest.mark.parametrize("length", [0, 7, _GCM_STREAM_HEADER_LENGTH, _MIN_STREAM_LENGTH - 1]) +def test_decrypt_rejects_a_stream_shorter_than_one_block(length: int) -> None: + """A header alone is not a stream, matching Java's `MIN_STREAM_LENGTH` and iceberg-rust.""" + with pytest.raises(ValueError, match=f"expected at least {_MIN_STREAM_LENGTH} bytes, got {length}"): + decrypt((JAVA_HEADER + bytes(_BLOCK_OVERHEAD))[:length]) + + +def test_decrypt_rejects_the_wrong_magic() -> None: + with pytest.raises(ValueError, match="magic b'AGS2' does not match b'AGS1'"): + decrypt(b"AGS2" + JAVA_HEADER[4:] + bytes(_BLOCK_OVERHEAD)) + + +def test_decrypt_rejects_an_unsupported_block_size() -> None: + with pytest.raises(ValueError, match=f"Unsupported AGS1 block size: 512 \\(expected {_PLAIN_BLOCK_SIZE}\\)"): + decrypt(b"AGS1" + (512).to_bytes(4, "little") + bytes(_BLOCK_OVERHEAD)) + + +@pytest.mark.parametrize("last_block_length", [1, _BLOCK_OVERHEAD - 1]) +def test_decrypt_rejects_a_truncated_last_block(last_block_length: int) -> None: + with pytest.raises(ValueError, match=f"last block is {last_block_length} bytes, expected at least {_BLOCK_OVERHEAD}"): + decrypt(JAVA_HEADER + bytes(_CIPHER_BLOCK_SIZE + last_block_length)) + + +def test_decrypt_rejects_reordered_blocks() -> None: + stream = _encrypt_ags1(bytes(2 * _PLAIN_BLOCK_SIZE), KEY_METADATA) + first, second = ( + stream[_GCM_STREAM_HEADER_LENGTH : _GCM_STREAM_HEADER_LENGTH + _CIPHER_BLOCK_SIZE], + stream[_GCM_STREAM_HEADER_LENGTH + _CIPHER_BLOCK_SIZE :], + ) + + with pytest.raises(ValueError, match="wrong decryption key; or corrupt/tampered data"): + decrypt(JAVA_HEADER + second + first) + + +def test_decrypt_rejects_another_files_aad_prefix() -> None: + stream = _encrypt_ags1(bytes(100), KEY_METADATA) + + with pytest.raises(ValueError, match="wrong decryption key; or corrupt/tampered data"): + decrypt(stream, aad_prefix=b"another file's prefix") + + +def test_decrypt_rejects_the_wrong_key() -> None: + stream = _encrypt_ags1(bytes(100), KEY_METADATA) + + with pytest.raises(ValueError, match="wrong decryption key; or corrupt/tampered data"): + decrypt(stream, key=b"1123456789012345")