diff --git a/bindings/python/src/encryption.rs b/bindings/python/src/encryption.rs new file mode 100644 index 0000000000..4384faca4b --- /dev/null +++ b/bindings/python/src/encryption.rs @@ -0,0 +1,79 @@ +// 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. + +use iceberg::encryption::StandardKeyMetadata; +use pyo3::prelude::*; +use pyo3::types::PyBytes; + +use crate::error::to_py_err; + +/// The encryption key, AAD prefix and file length held by `StandardKeyMetadata`. +type DecodedKeyMetadata<'py> = ( + Bound<'py, PyBytes>, + Option>, + Option, +); + +/// Decode `StandardKeyMetadata` from its wire format. +#[pyfunction] +pub fn decode_standard_key_metadata<'py>( + py: Python<'py>, + data: &[u8], +) -> PyResult> { + let metadata = StandardKeyMetadata::decode(data).map_err(to_py_err)?; + + Ok(( + PyBytes::new(py, metadata.encryption_key().as_bytes()), + metadata + .aad_prefix() + .map(|aad_prefix| PyBytes::new(py, aad_prefix)), + metadata.file_length(), + )) +} + +/// Encode `StandardKeyMetadata` to its wire format. +#[pyfunction] +#[pyo3(signature = (encryption_key, aad_prefix=None, file_length=None))] +pub fn encode_standard_key_metadata<'py>( + py: Python<'py>, + encryption_key: &[u8], + aad_prefix: Option<&[u8]>, + file_length: Option, +) -> PyResult> { + let mut metadata = StandardKeyMetadata::try_new(encryption_key).map_err(to_py_err)?; + + if let Some(aad_prefix) = aad_prefix { + metadata = metadata.with_aad_prefix(aad_prefix); + } + if let Some(file_length) = file_length { + metadata = metadata.with_file_length(file_length); + } + + Ok(PyBytes::new(py, &metadata.encode().map_err(to_py_err)?)) +} + +pub fn register_module(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> { + let this = PyModule::new(py, "encryption")?; + + this.add_function(wrap_pyfunction!(decode_standard_key_metadata, &this)?)?; + this.add_function(wrap_pyfunction!(encode_standard_key_metadata, &this)?)?; + + m.add_submodule(&this)?; + py.import("sys")? + .getattr("modules")? + .set_item("pyiceberg_core.encryption", this) +} diff --git a/bindings/python/src/lib.rs b/bindings/python/src/lib.rs index baddbe93a0..c50b1ce15a 100644 --- a/bindings/python/src/lib.rs +++ b/bindings/python/src/lib.rs @@ -18,6 +18,7 @@ use pyo3::prelude::*; mod data_file; +mod encryption; mod error; mod manifest; mod transform; @@ -26,5 +27,6 @@ mod transform; fn pyiceberg_core_rust(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> { transform::register_module(py, m)?; manifest::register_module(py, m)?; + encryption::register_module(py, m)?; Ok(()) } diff --git a/bindings/python/tests/test_encryption.py b/bindings/python/tests/test_encryption.py new file mode 100644 index 0000000000..7746adf9f7 --- /dev/null +++ b/bindings/python/tests/test_encryption.py @@ -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 pytest +from pyiceberg_core import encryption + +AES128_KEY = b"0123456789012345" + + +def test_encode_decode_round_trip(): + encoded = encryption.encode_standard_key_metadata(AES128_KEY, b"ad", 1024) + + assert encryption.decode_standard_key_metadata(encoded) == (AES128_KEY, b"ad", 1024) + + +def test_encode_decode_without_optional_fields(): + encoded = encryption.encode_standard_key_metadata(AES128_KEY) + + assert encryption.decode_standard_key_metadata(encoded) == (AES128_KEY, None, None) + + +def test_encoded_wire_format(): + # A version byte, then the Avro datum. Pinned so the encoding stays compatible + # with the Java and Python implementations. + assert encryption.encode_standard_key_metadata(AES128_KEY, b"ad", 1024) == ( + b"\x01\x20" + AES128_KEY + b"\x02\x04ad\x02\x80\x10" + ) + + +@pytest.mark.parametrize("key_length", [16, 24, 32]) +def test_encode_accepts_aes_key_lengths(key_length): + encoded = encryption.encode_standard_key_metadata(bytes(key_length)) + + assert encryption.decode_standard_key_metadata(encoded) == ( + bytes(key_length), + None, + None, + ) + + +@pytest.mark.parametrize("key_length", [0, 4, 15, 20, 33]) +def test_encode_rejects_invalid_key_length(key_length): + with pytest.raises(ValueError): + encryption.encode_standard_key_metadata(bytes(key_length)) + + +@pytest.mark.parametrize("data", [b"\x02", b"\x02\x20" + AES128_KEY + b"\x00\x00"]) +def test_decode_rejects_unsupported_version(data): + with pytest.raises(ValueError, match="Unsupported key metadata version: 2"): + encryption.decode_standard_key_metadata(data) + + +def test_decode_rejects_empty_buffer(): + with pytest.raises(ValueError): + encryption.decode_standard_key_metadata(b"") diff --git a/crates/iceberg/src/encryption/key_metadata.rs b/crates/iceberg/src/encryption/key_metadata.rs index 1169698e2f..433be48d2b 100644 --- a/crates/iceberg/src/encryption/key_metadata.rs +++ b/crates/iceberg/src/encryption/key_metadata.rs @@ -206,7 +206,7 @@ mod _serde { if version != V1 { return Err(Error::new( ErrorKind::FeatureUnsupported, - format!("Cannot resolve schema for version: {version}"), + format!("Unsupported key metadata version: {version} (supported: {V1})"), )); } @@ -302,6 +302,10 @@ mod tests { assert!(result.is_err()); let err = result.unwrap_err(); assert_eq!(err.kind(), ErrorKind::FeatureUnsupported); + assert_eq!( + err.message(), + "Unsupported key metadata version: 2 (supported: 1)" + ); } #[test]