iceberg/encryption/
key_metadata.rs

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//! Avro-serialized key metadata format compatible with Java's
19//! `org.apache.iceberg.encryption.StandardKeyMetadata`.
20
21use std::fmt;
22
23use aes_gcm::aead::OsRng;
24use aes_gcm::aead::rand_core::RngCore;
25
26use super::{AesKeySize, SecureKey};
27use crate::{Error, ErrorKind, Result};
28
29/// Standard key metadata for Iceberg table encryption.
30///
31/// Contains the Data Encryption Key (DEK), AAD prefix, and optional file
32/// length. Byte-compatible with Java's `StandardKeyMetadata` via Avro
33/// serialization.
34///
35/// Wire format: `[version byte (0x01)] [Avro binary datum]`
36#[derive(Clone, PartialEq, Eq)]
37pub struct StandardKeyMetadata {
38    encryption_key: SecureKey,
39    aad_prefix: Option<Box<[u8]>>,
40    file_length: Option<u64>,
41}
42
43impl fmt::Debug for StandardKeyMetadata {
44    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
45        f.debug_struct("StandardKeyMetadata")
46            .field("encryption_key", &self.encryption_key)
47            .field(
48                "aad_prefix",
49                &self
50                    .aad_prefix
51                    .as_ref()
52                    .map(|b| format!("[{} bytes]", b.len())),
53            )
54            .field("file_length", &self.file_length)
55            .finish()
56    }
57}
58
59impl StandardKeyMetadata {
60    /// Creates a new `StandardKeyMetadata` from raw key bytes.
61    pub fn try_new(encryption_key: &[u8]) -> Result<Self> {
62        Ok(Self::from(SecureKey::new(encryption_key)?))
63    }
64
65    /// Generates a `StandardKeyMetadata` carrying a fresh random DEK of
66    /// `key_size` together with a fresh random AAD prefix.
67    pub(crate) fn generate(key_size: AesKeySize) -> Self {
68        Self::from(SecureKey::generate(key_size)).with_aad_prefix(&generate_aad_prefix())
69    }
70
71    /// Adds an AAD prefix.
72    pub fn with_aad_prefix(mut self, aad_prefix: &[u8]) -> Self {
73        self.aad_prefix = Some(aad_prefix.into());
74        self
75    }
76
77    /// Adds a file length.
78    pub fn with_file_length(mut self, length: u64) -> Self {
79        self.file_length = Some(length);
80        self
81    }
82
83    /// Returns the plaintext Data Encryption Key.
84    pub fn encryption_key(&self) -> &SecureKey {
85        &self.encryption_key
86    }
87
88    /// Returns the AAD prefix.
89    pub fn aad_prefix(&self) -> Option<&[u8]> {
90        self.aad_prefix.as_deref()
91    }
92
93    /// Returns the optional file length.
94    pub fn file_length(&self) -> Option<u64> {
95        self.file_length
96    }
97
98    /// Encodes to Java-compatible format: `[0x01] [Avro binary datum]`
99    pub fn encode(&self) -> Result<Box<[u8]>> {
100        _serde::StandardKeyMetadataV1::from(self).encode()
101    }
102
103    /// Decodes from Java-compatible format.
104    pub fn decode(bytes: &[u8]) -> Result<Self> {
105        _serde::StandardKeyMetadataV1::decode(bytes).and_then(Self::try_from)
106    }
107}
108
109impl From<SecureKey> for StandardKeyMetadata {
110    /// Creates a `StandardKeyMetadata` from an already-validated key.
111    fn from(encryption_key: SecureKey) -> Self {
112        Self {
113            encryption_key,
114            aad_prefix: None,
115            file_length: None,
116        }
117    }
118}
119
120/// AAD prefix length in bytes.
121const AAD_PREFIX_LENGTH: usize = 16;
122
123fn generate_aad_prefix() -> Box<[u8]> {
124    let mut prefix = vec![0u8; AAD_PREFIX_LENGTH];
125    OsRng.fill_bytes(&mut prefix);
126    prefix.into_boxed_slice()
127}
128
129mod _serde {
130    use std::io::Cursor;
131    use std::sync::{Arc, LazyLock};
132
133    use apache_avro::{Schema as AvroSchema, from_avro_datum, from_value, to_avro_datum, to_value};
134    use serde::{Deserialize, Serialize};
135
136    use super::*;
137    use crate::avro::schema_to_avro_schema;
138    use crate::spec::{NestedField, PrimitiveType, Schema, Type};
139
140    pub(super) const V1: u8 = 1;
141
142    /// Avro schema for StandardKeyMetadata V1, derived from Iceberg schema.
143    pub(super) static AVRO_SCHEMA_V1: LazyLock<AvroSchema> = LazyLock::new(|| {
144        let schema = Schema::builder()
145            .with_fields(vec![
146                Arc::new(NestedField::required(
147                    0,
148                    "encryption_key",
149                    Type::Primitive(PrimitiveType::Binary),
150                )),
151                Arc::new(NestedField::optional(
152                    1,
153                    "aad_prefix",
154                    Type::Primitive(PrimitiveType::Binary),
155                )),
156                Arc::new(NestedField::optional(
157                    2,
158                    "file_length",
159                    Type::Primitive(PrimitiveType::Long),
160                )),
161            ])
162            .build()
163            .expect("Failed to build StandardKeyMetadata Iceberg schema");
164
165        schema_to_avro_schema("StandardKeyMetadata", &schema)
166            .expect("Failed to convert StandardKeyMetadata to Avro schema")
167    });
168
169    /// Serde struct for Avro serialization of [`StandardKeyMetadata`] V1.
170    /// Field names must match [`AVRO_SCHEMA_V1`] exactly.
171    #[derive(Serialize, Deserialize)]
172    pub(super) struct StandardKeyMetadataV1 {
173        pub encryption_key: serde_bytes::ByteBuf,
174        pub aad_prefix: Option<serde_bytes::ByteBuf>,
175        pub file_length: Option<u64>,
176    }
177
178    impl StandardKeyMetadataV1 {
179        pub(super) fn encode(&self) -> Result<Box<[u8]>> {
180            let value = to_value(self)
181                .and_then(|v| v.resolve(&AVRO_SCHEMA_V1))
182                .map_err(|e| {
183                    Error::new(ErrorKind::Unexpected, "Failed to encode key metadata")
184                        .with_source(e)
185                })?;
186
187            let datum = to_avro_datum(&AVRO_SCHEMA_V1, value).map_err(|e| {
188                Error::new(ErrorKind::Unexpected, "Failed to encode key metadata").with_source(e)
189            })?;
190
191            let mut result = Vec::with_capacity(1 + datum.len());
192            result.push(V1);
193            result.extend_from_slice(&datum);
194            Ok(result.into_boxed_slice())
195        }
196
197        pub(super) fn decode(bytes: &[u8]) -> Result<Self> {
198            if bytes.is_empty() {
199                return Err(Error::new(
200                    ErrorKind::DataInvalid,
201                    "Empty key metadata buffer",
202                ));
203            }
204
205            let version = bytes[0];
206            if version != V1 {
207                return Err(Error::new(
208                    ErrorKind::FeatureUnsupported,
209                    format!("Cannot resolve schema for version: {version}"),
210                ));
211            }
212
213            let mut reader = Cursor::new(&bytes[1..]);
214            let value = from_avro_datum(&AVRO_SCHEMA_V1, &mut reader, None).map_err(|e| {
215                Error::new(ErrorKind::DataInvalid, "Failed to decode key metadata").with_source(e)
216            })?;
217
218            from_value(&value).map_err(|e| {
219                Error::new(
220                    ErrorKind::DataInvalid,
221                    "Failed to decode key metadata fields",
222                )
223                .with_source(e)
224            })
225        }
226    }
227
228    impl From<&StandardKeyMetadata> for StandardKeyMetadataV1 {
229        fn from(metadata: &StandardKeyMetadata) -> Self {
230            Self {
231                encryption_key: serde_bytes::ByteBuf::from(metadata.encryption_key.as_bytes()),
232                aad_prefix: metadata
233                    .aad_prefix
234                    .as_ref()
235                    .map(|b| serde_bytes::ByteBuf::from(b.as_ref())),
236                file_length: metadata.file_length,
237            }
238        }
239    }
240
241    impl TryFrom<StandardKeyMetadataV1> for StandardKeyMetadata {
242        type Error = Error;
243
244        fn try_from(v1: StandardKeyMetadataV1) -> Result<Self> {
245            let encryption_key = SecureKey::new(&v1.encryption_key).map_err(|e| {
246                Error::new(
247                    ErrorKind::DataInvalid,
248                    "Invalid encryption key in key metadata",
249                )
250                .with_source(e)
251            })?;
252            Ok(Self {
253                encryption_key,
254                aad_prefix: v1.aad_prefix.map(|b| b.into_vec().into_boxed_slice()),
255                file_length: v1.file_length,
256            })
257        }
258    }
259}
260
261#[cfg(test)]
262mod tests {
263    use super::*;
264
265    #[test]
266    fn test_roundtrip() {
267        let key = b"0123456789012345";
268        let aad = b"1234567890123456";
269
270        let metadata = StandardKeyMetadata::try_new(key)
271            .unwrap()
272            .with_aad_prefix(aad);
273        let serialized = metadata.encode().unwrap();
274        let parsed = StandardKeyMetadata::decode(&serialized).unwrap();
275
276        assert_eq!(parsed.encryption_key().as_bytes(), key);
277        assert_eq!(parsed.aad_prefix(), Some(aad.as_slice()));
278        assert_eq!(parsed.file_length(), None);
279    }
280
281    #[test]
282    fn test_roundtrip_with_length() {
283        let key = b"0123456789012345";
284        let aad = b"1234567890123456";
285
286        let file_length = 100_000;
287        let metadata = StandardKeyMetadata::try_new(key)
288            .unwrap()
289            .with_aad_prefix(aad)
290            .with_file_length(file_length);
291        let serialized = metadata.encode().unwrap();
292        let parsed = StandardKeyMetadata::decode(&serialized).unwrap();
293
294        assert_eq!(parsed.encryption_key().as_bytes(), key);
295        assert_eq!(parsed.aad_prefix(), Some(aad.as_slice()));
296        assert_eq!(parsed.file_length(), Some(file_length));
297    }
298
299    #[test]
300    fn test_unsupported_version() {
301        let result = StandardKeyMetadata::decode(&[0x02]);
302        assert!(result.is_err());
303        let err = result.unwrap_err();
304        assert_eq!(err.kind(), ErrorKind::FeatureUnsupported);
305    }
306
307    #[test]
308    fn test_empty_buffer() {
309        let result = StandardKeyMetadata::decode(&[]);
310        assert!(result.is_err());
311        assert_eq!(result.unwrap_err().kind(), ErrorKind::DataInvalid);
312    }
313
314    #[test]
315    fn test_roundtrip_without_aad() {
316        let key = b"0123456789012345";
317        let metadata = StandardKeyMetadata::try_new(key).unwrap();
318        let serialized = metadata.encode().unwrap();
319        let parsed = StandardKeyMetadata::decode(&serialized).unwrap();
320
321        assert_eq!(parsed.encryption_key().as_bytes(), key);
322        assert_eq!(parsed.aad_prefix(), None);
323    }
324
325    #[test]
326    fn test_new_rejects_invalid_key_length() {
327        // 24-byte (AES-192) and 32-byte (AES-256) keys are accepted.
328        for len in [16usize, 24, 32] {
329            assert!(StandardKeyMetadata::try_new(&vec![0u8; len]).is_ok());
330        }
331
332        // Invalid lengths are rejected at construction, so an invalid
333        // `StandardKeyMetadata` can never exist.
334        for len in [0usize, 4, 15, 20, 33] {
335            assert!(StandardKeyMetadata::try_new(&vec![0u8; len]).is_err());
336        }
337    }
338
339    #[test]
340    fn test_decode_rejects_invalid_key_length() {
341        // Craft wire bytes carrying an invalid-length DEK directly via the
342        // serde struct (bypassing the validated public constructors) to prove
343        // `decode` still rejects malformed key material off the wire.
344        for len in [0usize, 4, 15, 20, 33] {
345            let serialized = _serde::StandardKeyMetadataV1 {
346                encryption_key: serde_bytes::ByteBuf::from(vec![0u8; len]),
347                aad_prefix: None,
348                file_length: None,
349            }
350            .encode()
351            .unwrap();
352
353            let err = StandardKeyMetadata::decode(&serialized).unwrap_err();
354            assert_eq!(err.kind(), ErrorKind::DataInvalid);
355            assert!(
356                err.to_string()
357                    .contains("Invalid encryption key in key metadata")
358            );
359        }
360    }
361}
iceberg/encryption/key_metadata.rs

iceberg/encryption/
key_metadata.rs
