Skip to main content

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