1use std::fmt;
22
23use super::SecureKey;
24use crate::{Error, ErrorKind, Result};
25
26#[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 pub fn try_new(encryption_key: &[u8]) -> Result<Self> {
59 Ok(Self::from(SecureKey::new(encryption_key)?))
60 }
61
62 pub fn with_aad_prefix(mut self, aad_prefix: &[u8]) -> Self {
64 self.aad_prefix = Some(aad_prefix.into());
65 self
66 }
67
68 pub fn with_file_length(mut self, length: u64) -> Self {
70 self.file_length = Some(length);
71 self
72 }
73
74 pub fn encryption_key(&self) -> &SecureKey {
76 &self.encryption_key
77 }
78
79 pub fn aad_prefix(&self) -> Option<&[u8]> {
81 self.aad_prefix.as_deref()
82 }
83
84 pub fn file_length(&self) -> Option<u64> {
86 self.file_length
87 }
88
89 pub fn encode(&self) -> Result<Box<[u8]>> {
91 _serde::StandardKeyMetadataV1::from(self).encode()
92 }
93
94 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 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 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 #[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 for len in [16usize, 24, 32] {
311 assert!(StandardKeyMetadata::try_new(&vec![0u8; len]).is_ok());
312 }
313
314 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 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}