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

pyiceberg/encryption/stream.py

@@ -0,0 +1,149 @@
+# 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.
+"""Format primitives for the AGS1 stream, 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.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from pyiceberg.encryption.ciphers import AesGcmCipher
+
+_GCM_STREAM_MAGIC = b"AGS1"
+_PLAIN_BLOCK_SIZE = 1024 * 1024
+_GCM_STREAM_HEADER_LENGTH = len(_GCM_STREAM_MAGIC) + 4
+_BLOCK_OVERHEAD = AesGcmCipher.NONCE_LENGTH + AesGcmCipher.TAG_LENGTH
+_CIPHER_BLOCK_SIZE = _PLAIN_BLOCK_SIZE + _BLOCK_OVERHEAD
+_BLOCK_INDEX_LENGTH = 4
+_MAX_BLOCKS = 2 ** (8 * _BLOCK_INDEX_LENGTH) - 1
+_MIN_STREAM_LENGTH = _GCM_STREAM_HEADER_LENGTH + _BLOCK_OVERHEAD
+
+
+def _stream_block_aad(aad_prefix: bytes | None, block_index: int) -> bytes:
+    """Return the additional authenticated data for the block at `block_index`.
+
+    Args:
+        aad_prefix (bytes | None): The file's AAD prefix, from its key metadata.
+        block_index (int): The zero-based index of the block within the stream.
+    """
+    return (aad_prefix or b"") + block_index.to_bytes(_BLOCK_INDEX_LENGTH, "little")
+
+
+def _encode_stream_header() -> bytes:
+    """Encode the AGS1 header that precedes the first block."""
+    return _GCM_STREAM_MAGIC + _PLAIN_BLOCK_SIZE.to_bytes(4, "little")
+
+
+def _decode_stream_header(header: bytes) -> int:
+    """Decode an AGS1 header, returning the plaintext block size it declares.
+
+    Args:
+        header (bytes): At least `_GCM_STREAM_HEADER_LENGTH` bytes from the start of the stream.
+    """
+    if len(header) < _GCM_STREAM_HEADER_LENGTH:
+        raise ValueError(f"Invalid AGS1 header: expected {_GCM_STREAM_HEADER_LENGTH} bytes, got {len(header)}")
+
+    if (magic := header[: 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(header[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})")
+
+    return plain_block_size
+
+
+@dataclass(frozen=True)
+class _Ags1Layout:
+    """Where each block of an AGS1 stream sits, derived from the trusted encrypted file length.
+
+    Only the final block may hold less than `_PLAIN_BLOCK_SIZE` of plaintext, so the layout
+    follows from the encrypted length alone, without reading the stream.
+    """
+
+    plaintext_length: int
+    num_blocks: int
+    last_cipher_block_size: int
+
+    @classmethod
+    def from_encrypted_length(cls, encrypted_length: int) -> _Ags1Layout:
+        """Derive the layout of an AGS1 stream that occupies `encrypted_length` bytes.
+
+        Args:
+            encrypted_length (int): The stream's length, which must be the trusted `file_length` from the file's
+                `StandardKeyMetadata`, never a file system stat. The spec requires the trusted length because a
+                stat lets an attacker drop trailing blocks while every remaining block still authenticates.
+        """
+        if encrypted_length < _MIN_STREAM_LENGTH:
+            raise ValueError(f"Invalid AGS1 stream: expected at least {_MIN_STREAM_LENGTH} bytes, got {encrypted_length}")
+
+        full_blocks, cipher_bytes_in_last_block = divmod(encrypted_length - _GCM_STREAM_HEADER_LENGTH, _CIPHER_BLOCK_SIZE)
+        if cipher_bytes_in_last_block == 0:
+            num_blocks, last_cipher_block_size = full_blocks, _CIPHER_BLOCK_SIZE
+        elif cipher_bytes_in_last_block < _BLOCK_OVERHEAD:
+            raise ValueError(
+                f"Truncated AGS1 stream: last block is {cipher_bytes_in_last_block} bytes, expected at least {_BLOCK_OVERHEAD}"
+            )
+        else:
+            num_blocks, last_cipher_block_size = full_blocks + 1, cipher_bytes_in_last_block
+
+        if num_blocks > _MAX_BLOCKS:
+            raise ValueError(f"AGS1 streams hold at most {_MAX_BLOCKS} blocks, but {encrypted_length} bytes needs {num_blocks}")
+
+        return cls(
+            plaintext_length=(num_blocks - 1) * _PLAIN_BLOCK_SIZE + last_cipher_block_size - _BLOCK_OVERHEAD,
+            num_blocks=num_blocks,
+            last_cipher_block_size=last_cipher_block_size,
+        )
+
+    def _check_block_index(self, block_index: int) -> None:
+        if not 0 <= block_index < self.num_blocks:
+            raise ValueError(f"Block index out of range: {block_index} (stream holds {self.num_blocks} blocks)")
+
+    def cipher_block_size(self, block_index: int) -> int:
+        """Return the encrypted size of the block at `block_index`."""
+        self._check_block_index(block_index)
+        return self.last_cipher_block_size if block_index == self.num_blocks - 1 else _CIPHER_BLOCK_SIZE
+
+    def plain_block_size(self, block_index: int) -> int:
+        """Return the plaintext size of the block at `block_index`."""
+        return self.cipher_block_size(block_index) - _BLOCK_OVERHEAD
+
+    def encrypted_block_offset(self, block_index: int) -> int:
+        """Return the offset of the block at `block_index` within the encrypted stream."""
+        self._check_block_index(block_index)
+        return _GCM_STREAM_HEADER_LENGTH + block_index * _CIPHER_BLOCK_SIZE
+
+    def block_index_for(self, plaintext_offset: int) -> int:
+        """Return the index of the block holding `plaintext_offset`."""
+        if not 0 <= plaintext_offset < self.plaintext_length:
+            raise ValueError(f"Plaintext offset out of range: {plaintext_offset} (stream holds {self.plaintext_length} bytes)")
+        return plaintext_offset // _PLAIN_BLOCK_SIZE

tests/encryption/ags1/GenerateAgs1Fixtures.java

@@ -0,0 +1,69 @@
+/*
+ * 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 java.io.File;
+import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+import org.apache.iceberg.Files;
+import org.apache.iceberg.encryption.AesGcmOutputFile;
+import org.apache.iceberg.io.PositionOutputStream;
+
+/** Writes the AGS1 fixtures in this directory with Java's AesGcmOutputStream. See README.md. */
+public class GenerateAgs1Fixtures {
+  static final int PLAIN_BLOCK_SIZE = 1024 * 1024;
+  static final byte[] KEY = new byte[16];
+  static final byte[] AAD_PREFIX = "pyiceberg-ags1".getBytes(StandardCharsets.UTF_8);
+
+  static {
+    for (int i = 0; i < KEY.length; i++) {
+      KEY[i] = (byte) i;
+    }
+  }
+
+  /** Byte i is i % 251. The prime period shifts phase across every 1 MiB block boundary. */
+  static byte[] plaintext(int length) {
+    byte[] out = new byte[length];
+    for (int i = 0; i < length; i++) {
+      out[i] = (byte) (i % 251);
+    }
+    return out;
+  }
+
+  static void write(File dir, String name, int length, byte[] aadPrefix) throws IOException {
+    File target = new File(dir, name);
+    if (target.exists() && !target.delete()) {
+      throw new IOException("Could not delete " + target);
+    }
+
+    AesGcmOutputFile encrypted = new AesGcmOutputFile(Files.localOutput(target), KEY, aadPrefix);
+    try (PositionOutputStream stream = encrypted.create()) {
+      stream.write(plaintext(length));
+    }
+
+    System.out.printf("%-26s plaintext=%-8d encrypted=%d%n", name, length, target.length());
+  }
+
+  public static void main(String[] args) throws IOException {
+    File dir = new File(args.length > 0 ? args[0] : ".");
+    write(dir, "empty.ags1", 0, AAD_PREFIX);
+    write(dir, "partial-block.ags1", 100, AAD_PREFIX);
+    write(dir, "partial-block-no-aad.ags1", 100, null);
+    write(dir, "aligned-multi-block.ags1", 2 * PLAIN_BLOCK_SIZE, AAD_PREFIX);
+  }
+}

tests/encryption/ags1/README.md

@@ -0,0 +1,77 @@
+<!--
+  ~ 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.
+-->
+
+# AGS1 cross-client test fixtures
+
+These files were written by Java's `AesGcmOutputStream` (Apache Iceberg 1.11.0), not
+by PyIceberg. They exist so PyIceberg's AGS1 support is checked against another
+implementation's bytes rather than only against its own round trip. All four also
+decrypt with iceberg-rust 0.10.1, through its `EncryptedInputFile`.
+
+## Fixtures
+
+| File | Encrypted size | Plaintext | Pins |
+| --- | --- | --- | --- |
+| `empty.ags1` | 36 B | 0 B | Java writes an 8 byte header **plus one empty block** for an empty file, not a bare header |
+| `partial-block.ags1` | 136 B | 100 B | Header, nonce/tag layout, and a single short block |
+| `partial-block-no-aad.ags1` | 136 B | 100 B | The same stream with a null AAD prefix, so the block index alone is the AAD |
+| `aligned-multi-block.ags1` | 2097216 B | 2 MiB | Two full blocks: the little-endian block index in each block's AAD, and that a block-aligned write appends **no** trailing empty block |
+
+`empty.ags1` is worth calling out. The spec says the last block has a non-zero
+length, which makes a bare 8 byte header its encoding of an empty file. Java instead
+writes 36 bytes and its `AesGcmInputFile` rejects anything shorter, as does iceberg-rust
+since apache/iceberg-rust#3236. PyIceberg follows the implementations rather than the
+spec here: `_MIN_STREAM_LENGTH` is 36, so a bare header is rejected rather than read as
+an empty stream. apache/iceberg#18219 tracks which of the two forms writers should
+produce.
+
+Block-aligned and partial *single* block variants are deliberately not checked in.
+The 1 MiB block size is hard-coded, so each would add another 1 MiB of
+incompressible ciphertext without covering a case the four files above miss.
+
+## Parameters
+
+Every fixture uses:
+
+- **Key**: 16 bytes, `0x00` through `0x0f`
+- **AAD prefix**: ASCII `pyiceberg-ags1`, except `partial-block-no-aad.ags1`, which has none
+- **Plaintext**: byte `i` is `i % 251`. The period is prime and therefore coprime with the
+  1 MiB block size, so the pattern shifts phase at every block boundary and a
+  misordered or misindexed block is detectable from the plaintext alone
+
+## Regenerating
+
+Each block uses a fresh random nonce, so regenerating produces different bytes.
+The file lengths, and the plaintext each file decrypts to, are deterministic. Tests
+decrypt these fixtures rather than comparing them byte for byte, so a regeneration
+is safe as long as the parameters above are unchanged.
+
+From this directory, with a JDK 17 or later:
+
+<!-- markdown-link-check-disable -->
+```bash
+V=1.11.0
+for a in iceberg-core iceberg-api iceberg-bundled-guava; do
+  curl -sfLO "https://repo1.maven.org/maven2/org/apache/iceberg/$a/$V/$a-$V.jar"
+done
+java -cp "iceberg-api-$V.jar:iceberg-bundled-guava-$V.jar:iceberg-core-$V.jar" \
+  GenerateAgs1Fixtures.java .
+rm iceberg-*-$V.jar
+```
+<!-- markdown-link-check-enable -->