Skip to content

Commit 400d324

Browse files
authored
feat(encryption) [4/N] AES-GCM encryption implementation (#3963)
* tmp * nit * test(encryption): add AES-GCM known-answer regression tests The existing cipher tests all round-trip through the same code, so a symmetric change to the nonce || ciphertext || tag layout would pass them while breaking interoperability with the Java and iceberg-rust clients. Pin encryption and decryption against fixed GCM-spec vectors covering AES-128 and AES-256, with and without AAD. * fix(encryption): align AES-GCM decryption error with Java Report the same cause as the Java client when the GCM tag check fails, so the message reads consistently across implementations and makes clear the failure is bad input rather than a retryable system error. * refactor(encryption): import cryptography through try_import Import the cryptography modules with try_import rather than a spelled-out try/except, and drop not_installed now that nothing calls it. The TYPE_CHECKING block already keeps AESGCM and InvalidTag statically typed, so lazy_import is back to a single entry point.
1 parent 182169e commit 400d324

5 files changed

Lines changed: 313 additions & 1 deletion

File tree

‎mkdocs/docs/index.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,7 @@ You can mix and match optional dependencies depending on your needs:
6666
| hf | Support for Hugging Face Hub |
6767
| gcp-auth | Support for Google Cloud authentication |
6868
| entra-auth | Support for Azure Entra authentication |
69+
| encryption | Support for table encryption |
6970

7071
You either need to install `s3fs`, `adlfs`, `gcsfs`, or `pyarrow` to be able to fetch files from an object store.
7172

‎pyiceberg/encryption/ciphers.py‎

Lines changed: 117 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,117 @@
1+
# Licensed to the Apache Software Foundation (ASF) under one
2+
# or more contributor license agreements. See the NOTICE file
3+
# distributed with this work for additional information
4+
# regarding copyright ownership. The ASF licenses this file
5+
# to you under the Apache License, Version 2.0 (the
6+
# "License"); you may not use this file except in compliance
7+
# with the License. You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing,
12+
# software distributed under the License is distributed on an
13+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
# KIND, either express or implied. See the License for the
15+
# specific language governing permissions and limitations
16+
# under the License.
17+
"""AES-GCM primitives for table encryption."""
18+
19+
from __future__ import annotations
20+
21+
import os
22+
from dataclasses import dataclass, field
23+
from enum import IntEnum
24+
from typing import TYPE_CHECKING
25+
26+
from pyiceberg.utils.lazy_import import try_import
27+
28+
if TYPE_CHECKING:
29+
from cryptography.exceptions import InvalidTag
30+
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
31+
32+
33+
class AesKeySize(IntEnum):
34+
"""The supported AES key sizes, in bits."""
35+
36+
BITS_128 = 128
37+
BITS_192 = 192
38+
BITS_256 = 256
39+
40+
@property
41+
def key_length(self) -> int:
42+
"""Return the key length in bytes."""
43+
return self.value // 8
44+
45+
@classmethod
46+
def from_key_length(cls, key_length: int) -> AesKeySize:
47+
"""Return the key size for a key of `key_length` bytes."""
48+
try:
49+
return cls(key_length * 8)
50+
except ValueError as e:
51+
raise ValueError(f"Unsupported key length: {key_length} (must be 16, 24 or 32)") from e
52+
53+
54+
@dataclass(frozen=True)
55+
class SecureKey:
56+
"""An AES key of a length the spec allows, kept out of reprs and tracebacks."""
57+
58+
key: bytes = field(repr=False)
59+
60+
def __post_init__(self) -> None:
61+
"""Reject keys that are not a supported AES key length."""
62+
AesKeySize.from_key_length(len(self.key))
63+
64+
@property
65+
def key_size(self) -> AesKeySize:
66+
"""Return the size of this key."""
67+
return AesKeySize.from_key_length(len(self.key))
68+
69+
@classmethod
70+
def generate(cls, key_size: AesKeySize = AesKeySize.BITS_128) -> SecureKey:
71+
"""Generate a new key of `key_size`."""
72+
return cls(os.urandom(key_size.key_length))
73+
74+
75+
class AesGcmCipher:
76+
"""Encrypts and decrypts using AES-GCM.
77+
78+
Ciphertext is laid out as `nonce || ciphertext || tag`, matching Java and iceberg-rust.
79+
"""
80+
81+
NONCE_LENGTH = 12
82+
TAG_LENGTH = 16
83+
84+
def __init__(self, key: SecureKey) -> None:
85+
aead = try_import("cryptography.hazmat.primitives.ciphers.aead", extras_name="encryption")
86+
exceptions = try_import("cryptography.exceptions", extras_name="encryption")
87+
88+
self._aes_gcm: AESGCM = aead.AESGCM(key.key)
89+
self._invalid_tag: type[InvalidTag] = exceptions.InvalidTag
90+
91+
def encrypt(self, plaintext: bytes, aad: bytes | None = None) -> bytes:
92+
"""Encrypt `plaintext`, authenticating `aad` alongside it.
93+
94+
Args:
95+
plaintext (bytes): The data to encrypt.
96+
aad (bytes | None): Additional data to authenticate but not encrypt.
97+
"""
98+
nonce = os.urandom(self.NONCE_LENGTH)
99+
return nonce + self._aes_gcm.encrypt(nonce, plaintext, aad)
100+
101+
def decrypt(self, ciphertext: bytes, aad: bytes | None = None) -> bytes:
102+
"""Decrypt `ciphertext`, verifying `aad` alongside it.
103+
104+
Args:
105+
ciphertext (bytes): The data to decrypt, as returned by `encrypt`.
106+
aad (bytes | None): The additional data that was authenticated on encryption.
107+
"""
108+
if len(ciphertext) < self.NONCE_LENGTH + self.TAG_LENGTH:
109+
raise ValueError(
110+
f"Ciphertext too short: expected at least {self.NONCE_LENGTH + self.TAG_LENGTH} bytes, got {len(ciphertext)}"
111+
)
112+
113+
nonce, encrypted = ciphertext[: self.NONCE_LENGTH], ciphertext[self.NONCE_LENGTH :]
114+
try:
115+
return self._aes_gcm.decrypt(nonce, encrypted, aad)
116+
except self._invalid_tag as e:
117+
raise ValueError("GCM tag check failed. Possible reasons: wrong decryption key; or corrupt/tampered data") from e

‎pyproject.toml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -101,6 +101,7 @@ datafusion = ["datafusion>=53,<54"]
101101
gcp-auth = ["google-auth>=2.4.0"]
102102
entra-auth = ["azure-identity>=1.25.1"]
103103
geoarrow = ["geoarrow-pyarrow>=0.2.0"]
104+
encryption = ["cryptography>=42.0.0"]
104105

105106
[dependency-groups]
106107
dev = [

‎tests/encryption/test_ciphers.py‎

Lines changed: 189 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,189 @@
1+
# Licensed to the Apache Software Foundation (ASF) under one
2+
# or more contributor license agreements. See the NOTICE file
3+
# distributed with this work for additional information
4+
# regarding copyright ownership. The ASF licenses this file
5+
# to you under the Apache License, Version 2.0 (the
6+
# "License"); you may not use this file except in compliance
7+
# with the License. You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing,
12+
# software distributed under the License is distributed on an
13+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
# KIND, either express or implied. See the License for the
15+
# specific language governing permissions and limitations
16+
# under the License.
17+
18+
import pytest
19+
from pytest_mock import MockFixture
20+
21+
from pyiceberg.encryption.ciphers import AesGcmCipher, AesKeySize, SecureKey
22+
from pyiceberg.exceptions import NotInstalledError
23+
24+
AES128_KEY = b"0123456789012345"
25+
PLAINTEXT = b"the quick brown fox"
26+
27+
# Known-answer vectors from McGrew & Viega, "The Galois/Counter Mode of Operation
28+
# (GCM)", shared with the NIST GCM validation suite and the Java and iceberg-rust
29+
# test suites. They pin the `nonce || ciphertext || tag` layout against changes
30+
# that stay self-consistent on round trip but break cross-client interoperability.
31+
GCM_TEST_VECTORS = [
32+
pytest.param(
33+
"feffe9928665731c6d6a8f9467308308",
34+
"cafebabefacedbaddecaf888",
35+
"d9313225f88406e5a55909c5aff5269a86a7a9531534f7da2e4c303d8a318a721c3c0c95956809532fcf0e2449a6b525b16aedf5aa0de657ba637b391aafd255",
36+
"",
37+
"42831ec2217774244b7221b784d0d49ce3aa212f2c02a4e035c17e2329aca12e21d514b25466931c7d8f6a5aac84aa051ba30b396a0aac973d58e091473f5985",
38+
"4d5c2af327cd64a62cf35abd2ba6fab4",
39+
id="aes128-no-aad",
40+
),
41+
pytest.param(
42+
"feffe9928665731c6d6a8f9467308308",
43+
"cafebabefacedbaddecaf888",
44+
"d9313225f88406e5a55909c5aff5269a86a7a9531534f7da2e4c303d8a318a721c3c0c95956809532fcf0e2449a6b525b16aedf5aa0de657ba637b39",
45+
"feedfacedeadbeeffeedfacedeadbeefabaddad2",
46+
"42831ec2217774244b7221b784d0d49ce3aa212f2c02a4e035c17e2329aca12e21d514b25466931c7d8f6a5aac84aa051ba30b396a0aac973d58e091",
47+
"5bc94fbc3221a5db94fae95ae7121a47",
48+
id="aes128-with-aad",
49+
),
50+
pytest.param(
51+
"feffe9928665731c6d6a8f9467308308feffe9928665731c6d6a8f9467308308",
52+
"cafebabefacedbaddecaf888",
53+
"d9313225f88406e5a55909c5aff5269a86a7a9531534f7da2e4c303d8a318a721c3c0c95956809532fcf0e2449a6b525b16aedf5aa0de657ba637b39",
54+
"feedfacedeadbeeffeedfacedeadbeefabaddad2",
55+
"522dc1f099567d07f47f37a32a84427d643a8cdcbfe5c0c97598a2bd2555d1aa8cb08e48590dbb3da7b08b1056828838c5f61e6393ba7a0abcc9f662",
56+
"76fc6ece0f4e1768cddf8853bb2d551b",
57+
id="aes256-with-aad",
58+
),
59+
]
60+
61+
62+
@pytest.mark.parametrize(
63+
"key_length, key_size",
64+
[(16, AesKeySize.BITS_128), (24, AesKeySize.BITS_192), (32, AesKeySize.BITS_256)],
65+
)
66+
def test_key_size_from_key_length(key_length: int, key_size: AesKeySize) -> None:
67+
assert AesKeySize.from_key_length(key_length) == key_size
68+
assert key_size.key_length == key_length
69+
70+
71+
@pytest.mark.parametrize("key_length", [0, 4, 15, 20, 33])
72+
def test_key_size_rejects_invalid_key_length(key_length: int) -> None:
73+
with pytest.raises(ValueError, match=f"Unsupported key length: {key_length}"):
74+
AesKeySize.from_key_length(key_length)
75+
76+
77+
@pytest.mark.parametrize("key_length", [0, 4, 15, 20, 33])
78+
def test_secure_key_rejects_invalid_key_length(key_length: int) -> None:
79+
with pytest.raises(ValueError, match="Unsupported key length"):
80+
SecureKey(bytes(key_length))
81+
82+
83+
@pytest.mark.parametrize("key_size", list(AesKeySize))
84+
def test_secure_key_generate(key_size: AesKeySize) -> None:
85+
key = SecureKey.generate(key_size)
86+
87+
assert len(key.key) == key_size.key_length
88+
assert key.key_size == key_size
89+
assert SecureKey.generate(key_size) != key
90+
91+
92+
def test_secure_key_repr_redacts_key() -> None:
93+
key = SecureKey(AES128_KEY)
94+
95+
assert repr(key) == "SecureKey()"
96+
assert repr(AES128_KEY) not in repr(key)
97+
98+
99+
@pytest.mark.parametrize("key_size", list(AesKeySize))
100+
@pytest.mark.parametrize("aad", [None, b"", b"aad"])
101+
def test_encrypt_decrypt_round_trip(key_size: AesKeySize, aad: bytes | None) -> None:
102+
cipher = AesGcmCipher(SecureKey.generate(key_size))
103+
104+
ciphertext = cipher.encrypt(PLAINTEXT, aad)
105+
106+
assert ciphertext != PLAINTEXT
107+
assert cipher.decrypt(ciphertext, aad) == PLAINTEXT
108+
109+
110+
@pytest.mark.parametrize("key, nonce, plaintext, aad, ciphertext, tag", GCM_TEST_VECTORS)
111+
def test_decrypt_known_answer(key: str, nonce: str, plaintext: str, aad: str, ciphertext: str, tag: str) -> None:
112+
cipher = AesGcmCipher(SecureKey(bytes.fromhex(key)))
113+
stored = bytes.fromhex(nonce + ciphertext + tag)
114+
115+
assert cipher.decrypt(stored, bytes.fromhex(aad) or None) == bytes.fromhex(plaintext)
116+
117+
118+
@pytest.mark.parametrize("key, nonce, plaintext, aad, ciphertext, tag", GCM_TEST_VECTORS)
119+
def test_encrypt_known_answer(
120+
monkeypatch: pytest.MonkeyPatch, key: str, nonce: str, plaintext: str, aad: str, ciphertext: str, tag: str
121+
) -> None:
122+
monkeypatch.setattr("pyiceberg.encryption.ciphers.os.urandom", lambda _: bytes.fromhex(nonce))
123+
cipher = AesGcmCipher(SecureKey(bytes.fromhex(key)))
124+
125+
assert cipher.encrypt(bytes.fromhex(plaintext), bytes.fromhex(aad) or None) == bytes.fromhex(nonce + ciphertext + tag)
126+
127+
128+
def test_encrypt_empty_plaintext() -> None:
129+
cipher = AesGcmCipher(SecureKey(AES128_KEY))
130+
131+
assert cipher.decrypt(cipher.encrypt(b"")) == b""
132+
133+
134+
def test_ciphertext_layout() -> None:
135+
cipher = AesGcmCipher(SecureKey(AES128_KEY))
136+
137+
ciphertext = cipher.encrypt(PLAINTEXT)
138+
139+
assert len(ciphertext) == AesGcmCipher.NONCE_LENGTH + len(PLAINTEXT) + AesGcmCipher.TAG_LENGTH
140+
141+
142+
def test_nonce_is_not_reused() -> None:
143+
cipher = AesGcmCipher(SecureKey(AES128_KEY))
144+
145+
first, second = cipher.encrypt(PLAINTEXT), cipher.encrypt(PLAINTEXT)
146+
147+
assert first[: AesGcmCipher.NONCE_LENGTH] != second[: AesGcmCipher.NONCE_LENGTH]
148+
assert first != second
149+
150+
151+
def test_decrypt_with_wrong_key() -> None:
152+
ciphertext = AesGcmCipher(SecureKey(AES128_KEY)).encrypt(PLAINTEXT)
153+
154+
with pytest.raises(ValueError, match="wrong decryption key; or corrupt/tampered data"):
155+
AesGcmCipher(SecureKey(b"5432109876543210")).decrypt(ciphertext)
156+
157+
158+
def test_decrypt_with_mismatched_aad() -> None:
159+
cipher = AesGcmCipher(SecureKey(AES128_KEY))
160+
161+
ciphertext = cipher.encrypt(PLAINTEXT, b"aad")
162+
163+
with pytest.raises(ValueError, match="wrong decryption key; or corrupt/tampered data"):
164+
cipher.decrypt(ciphertext, b"other aad")
165+
166+
167+
def test_decrypt_tampered_ciphertext() -> None:
168+
cipher = AesGcmCipher(SecureKey(AES128_KEY))
169+
170+
ciphertext = bytearray(cipher.encrypt(PLAINTEXT))
171+
ciphertext[-1] ^= 0xFF
172+
173+
with pytest.raises(ValueError, match="wrong decryption key; or corrupt/tampered data"):
174+
cipher.decrypt(bytes(ciphertext))
175+
176+
177+
@pytest.mark.parametrize("length", [0, 1, 27])
178+
def test_decrypt_ciphertext_too_short(length: int) -> None:
179+
cipher = AesGcmCipher(SecureKey(AES128_KEY))
180+
181+
with pytest.raises(ValueError, match=f"Ciphertext too short: expected at least 28 bytes, got {length}"):
182+
cipher.decrypt(bytes(length))
183+
184+
185+
def test_cipher_without_cryptography_installed_raises_not_installed_error(mocker: MockFixture) -> None:
186+
mocker.patch.dict("sys.modules", {"cryptography.hazmat.primitives.ciphers.aead": None})
187+
188+
with pytest.raises(NotInstalledError, match=r"pyiceberg\[encryption\]"):
189+
AesGcmCipher(SecureKey(AES128_KEY))

‎uv.lock‎

Lines changed: 5 additions & 1 deletion
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)