feat(encryption) [6/N] streaming block encryption

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})")

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")