1use std::collections::HashMap;
8
9use crate::secrets_manager::SecretsError;
10
11#[derive(Debug, Clone, Copy, PartialEq, Eq)]
13pub enum EncryptionMark {
14 Encrypted,
16 Sensitive,
18 Encrypt,
20}
21
22impl std::fmt::Display for EncryptionMark {
23 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
24 match self {
25 Self::Encrypted => write!(f, "encrypted"),
26 Self::Sensitive => write!(f, "sensitive"),
27 Self::Encrypt => write!(f, "encrypt"),
28 }
29 }
30}
31
32#[derive(Debug, Clone)]
34pub struct SchemaFieldInfo {
35 pub field_name: String,
37 pub field_type: String,
39 pub is_encrypted: bool,
41 pub key_reference: String,
43 pub algorithm: String,
45 pub nullable: bool,
47 pub mark: Option<EncryptionMark>,
49}
50
51impl SchemaFieldInfo {
52 pub fn new(
54 field_name: impl Into<String>,
55 field_type: impl Into<String>,
56 is_encrypted: bool,
57 key_reference: impl Into<String>,
58 ) -> Self {
59 Self {
60 field_name: field_name.into(),
61 field_type: field_type.into(),
62 is_encrypted,
63 key_reference: key_reference.into(),
64 algorithm: "aes256-gcm".to_string(),
65 nullable: false,
66 mark: None,
67 }
68 }
69
70 pub fn with_algorithm(mut self, algorithm: impl Into<String>) -> Self {
72 self.algorithm = algorithm.into();
73 self
74 }
75
76 pub fn with_nullable(mut self, nullable: bool) -> Self {
78 self.nullable = nullable;
79 self
80 }
81
82 pub fn with_mark(mut self, mark: EncryptionMark) -> Self {
84 self.mark = Some(mark);
85 self
86 }
87}
88
89#[derive(Debug, Clone)]
91pub struct StructSchema {
92 pub type_name: String,
94 pub all_fields: Vec<SchemaFieldInfo>,
96 pub encrypted_fields: Vec<SchemaFieldInfo>,
98 pub version: u32,
100}
101
102impl StructSchema {
103 pub fn new(type_name: impl Into<String>) -> Self {
105 Self {
106 type_name: type_name.into(),
107 all_fields: Vec::new(),
108 encrypted_fields: Vec::new(),
109 version: 1,
110 }
111 }
112
113 pub fn add_field(&mut self, field: SchemaFieldInfo) {
115 if field.is_encrypted {
116 self.encrypted_fields.push(field.clone());
117 }
118 self.all_fields.push(field);
119 }
120
121 pub fn with_fields(mut self, fields: Vec<SchemaFieldInfo>) -> Self {
123 for field in fields {
124 self.add_field(field);
125 }
126 self
127 }
128
129 pub fn with_version(mut self, version: u32) -> Self {
131 self.version = version;
132 self
133 }
134
135 pub fn get_field(&self, field_name: &str) -> Option<&SchemaFieldInfo> {
137 self.all_fields.iter().find(|f| f.field_name == field_name)
138 }
139
140 pub fn get_encrypted_field(&self, field_name: &str) -> Option<&SchemaFieldInfo> {
142 self.encrypted_fields.iter().find(|f| f.field_name == field_name)
143 }
144
145 pub fn is_field_encrypted(&self, field_name: &str) -> bool {
147 self.encrypted_fields.iter().any(|f| f.field_name == field_name)
148 }
149
150 pub fn encrypted_field_names(&self) -> Vec<&str> {
152 self.encrypted_fields.iter().map(|f| f.field_name.as_str()).collect()
153 }
154
155 fn filter_fields<F>(&self, predicate: F) -> Vec<&SchemaFieldInfo>
157 where
158 F: Fn(&&SchemaFieldInfo) -> bool,
159 {
160 self.all_fields.iter().filter(predicate).collect()
161 }
162
163 pub fn nullable_encrypted_fields(&self) -> Vec<&SchemaFieldInfo> {
165 self.filter_fields(|f| f.is_encrypted && f.nullable)
166 }
167
168 pub fn fields_for_key(&self, key_ref: &str) -> Vec<&SchemaFieldInfo> {
170 self.filter_fields(|f| f.key_reference == key_ref)
171 }
172
173 pub fn encrypted_field_count(&self) -> usize {
175 self.encrypted_fields.len()
176 }
177
178 pub fn total_field_count(&self) -> usize {
180 self.all_fields.len()
181 }
182
183 pub fn validate(&self) -> Result<(), SecretsError> {
185 if self.type_name.is_empty() {
186 return Err(SecretsError::ValidationError(
187 "Schema type name cannot be empty".to_string(),
188 ));
189 }
190
191 for field in &self.encrypted_fields {
193 if field.key_reference.is_empty() {
194 return Err(SecretsError::ValidationError(format!(
195 "Encrypted field '{}' missing key reference",
196 field.field_name
197 )));
198 }
199 }
200
201 Ok(())
202 }
203}
204
205pub struct SchemaRegistry {
207 schemas: HashMap<String, StructSchema>,
209 default_key_reference: String,
211}
212
213impl SchemaRegistry {
214 pub fn new() -> Self {
216 Self {
217 schemas: HashMap::new(),
218 default_key_reference: "encryption/default".to_string(),
219 }
220 }
221
222 pub fn with_default_key(mut self, key_reference: impl Into<String>) -> Self {
224 self.default_key_reference = key_reference.into();
225 self
226 }
227
228 pub fn register(&mut self, schema: StructSchema) -> Result<(), SecretsError> {
230 schema.validate()?;
231 self.schemas.insert(schema.type_name.clone(), schema);
232 Ok(())
233 }
234
235 pub fn get(&self, type_name: &str) -> Option<&StructSchema> {
237 self.schemas.get(type_name)
238 }
239
240 pub fn get_encrypted_fields(
242 &self,
243 type_name: &str,
244 ) -> Result<Vec<&SchemaFieldInfo>, SecretsError> {
245 self.get(type_name)
246 .map(|schema| schema.encrypted_fields.iter().collect())
247 .ok_or_else(|| {
248 SecretsError::ValidationError(format!("Schema '{}' not registered", type_name))
249 })
250 }
251
252 pub fn has_encrypted_fields(&self, type_name: &str) -> bool {
254 self.get(type_name)
255 .map(|schema| !schema.encrypted_fields.is_empty())
256 .unwrap_or(false)
257 }
258
259 pub fn list_types(&self) -> Vec<&str> {
261 self.schemas.keys().map(|s| s.as_str()).collect()
262 }
263
264 pub fn types_with_encryption(&self) -> Vec<&str> {
266 self.schemas
267 .iter()
268 .filter(|(_, schema)| !schema.encrypted_fields.is_empty())
269 .map(|(name, _)| name.as_str())
270 .collect()
271 }
272
273 pub fn all_encryption_keys(&self) -> Vec<String> {
275 let mut keys = std::collections::HashSet::new();
276 for schema in self.schemas.values() {
277 for field in &schema.encrypted_fields {
278 keys.insert(field.key_reference.clone());
279 }
280 }
281 let mut sorted: Vec<_> = keys.into_iter().collect();
282 sorted.sort();
283 sorted
284 }
285
286 pub fn validate_all(&self) -> Result<(), SecretsError> {
288 for schema in self.schemas.values() {
289 schema.validate()?;
290 }
291 Ok(())
292 }
293
294 pub fn unregister(&mut self, type_name: &str) -> Option<StructSchema> {
296 self.schemas.remove(type_name)
297 }
298
299 pub fn clear(&mut self) {
301 self.schemas.clear();
302 }
303
304 pub fn count(&self) -> usize {
306 self.schemas.len()
307 }
308
309 pub fn total_encrypted_fields(&self) -> usize {
311 self.schemas.values().map(|schema| schema.encrypted_fields.len()).sum()
312 }
313}
314
315impl Default for SchemaRegistry {
316 fn default() -> Self {
317 Self::new()
318 }
319}
320
321#[cfg(test)]
322mod tests {
323 use super::*;
324
325 #[test]
326 fn test_encryption_mark_display() {
327 assert_eq!(EncryptionMark::Encrypted.to_string(), "encrypted");
328 assert_eq!(EncryptionMark::Sensitive.to_string(), "sensitive");
329 assert_eq!(EncryptionMark::Encrypt.to_string(), "encrypt");
330 }
331
332 #[test]
333 fn test_field_info_creation() {
334 let field = SchemaFieldInfo::new("email", "String", true, "encryption/email");
335 assert_eq!(field.field_name, "email");
336 assert_eq!(field.field_type, "String");
337 assert!(field.is_encrypted);
338 assert_eq!(field.key_reference, "encryption/email");
339 assert_eq!(field.algorithm, "aes256-gcm");
340 }
341
342 #[test]
343 fn test_field_info_with_algorithm() {
344 let field = SchemaFieldInfo::new("email", "String", true, "encryption/email")
345 .with_algorithm("aes256-gcm");
346 assert_eq!(field.algorithm, "aes256-gcm");
347 }
348
349 #[test]
350 fn test_field_info_with_nullable() {
351 let field = SchemaFieldInfo::new("email", "Option<String>", true, "encryption/email")
352 .with_nullable(true);
353 assert!(field.nullable);
354 }
355
356 #[test]
357 fn test_field_info_with_mark() {
358 let field = SchemaFieldInfo::new("email", "String", true, "encryption/email")
359 .with_mark(EncryptionMark::Encrypted);
360 assert_eq!(field.mark, Some(EncryptionMark::Encrypted));
361 }
362
363 #[test]
364 fn test_struct_schema_creation() {
365 let schema = StructSchema::new("User");
366 assert_eq!(schema.type_name, "User");
367 assert!(schema.all_fields.is_empty());
368 assert!(schema.encrypted_fields.is_empty());
369 assert_eq!(schema.version, 1);
370 }
371
372 #[test]
373 fn test_struct_schema_add_field() {
374 let mut schema = StructSchema::new("User");
375 let email = SchemaFieldInfo::new("email", "String", true, "encryption/email");
376 schema.add_field(email);
377 assert_eq!(schema.all_fields.len(), 1);
378 assert_eq!(schema.encrypted_fields.len(), 1);
379 }
380
381 #[test]
382 fn test_struct_schema_mixed_fields() {
383 let mut schema = StructSchema::new("User");
384 let encrypted = SchemaFieldInfo::new("email", "String", true, "encryption/email");
385 let unencrypted = SchemaFieldInfo::new("name", "String", false, "");
386 schema.add_field(encrypted);
387 schema.add_field(unencrypted);
388 assert_eq!(schema.all_fields.len(), 2);
389 assert_eq!(schema.encrypted_fields.len(), 1);
390 }
391
392 #[test]
393 fn test_struct_schema_get_field() {
394 let mut schema = StructSchema::new("User");
395 let email = SchemaFieldInfo::new("email", "String", true, "encryption/email");
396 schema.add_field(email);
397 assert!(schema.get_field("email").is_some());
398 assert!(schema.get_field("phone").is_none());
399 }
400
401 #[test]
402 fn test_struct_schema_is_field_encrypted() {
403 let mut schema = StructSchema::new("User");
404 let encrypted = SchemaFieldInfo::new("email", "String", true, "encryption/email");
405 let unencrypted = SchemaFieldInfo::new("name", "String", false, "");
406 schema.add_field(encrypted);
407 schema.add_field(unencrypted);
408 assert!(schema.is_field_encrypted("email"));
409 assert!(!schema.is_field_encrypted("name"));
410 }
411
412 #[test]
413 fn test_struct_schema_encrypted_field_names() {
414 let mut schema = StructSchema::new("User");
415 let email = SchemaFieldInfo::new("email", "String", true, "encryption/email");
416 let phone = SchemaFieldInfo::new("phone", "String", true, "encryption/phone");
417 let name = SchemaFieldInfo::new("name", "String", false, "");
418 schema.add_field(email);
419 schema.add_field(phone);
420 schema.add_field(name);
421 let names = schema.encrypted_field_names();
422 assert_eq!(names.len(), 2);
423 assert!(names.contains(&"email"));
424 assert!(names.contains(&"phone"));
425 }
426
427 #[test]
428 fn test_struct_schema_validate_success() {
429 let mut schema = StructSchema::new("User");
430 let email = SchemaFieldInfo::new("email", "String", true, "encryption/email");
431 schema.add_field(email);
432 assert!(schema.validate().is_ok());
433 }
434
435 #[test]
436 fn test_struct_schema_validate_empty_type_name() {
437 let schema = StructSchema::new("");
438 let result = schema.validate();
439 assert!(result.is_err());
440 }
441
442 #[test]
443 fn test_struct_schema_validate_missing_key_reference() {
444 let mut schema = StructSchema::new("User");
445 let email = SchemaFieldInfo::new("email", "String", true, "");
446 schema.add_field(email);
447 let result = schema.validate();
448 assert!(result.is_err());
449 }
450
451 #[test]
452 fn test_struct_schema_with_version() {
453 let schema = StructSchema::new("User").with_version(2);
454 assert_eq!(schema.version, 2);
455 }
456
457 #[test]
458 fn test_schema_registry_creation() {
459 let registry = SchemaRegistry::new();
460 assert_eq!(registry.count(), 0);
461 }
462
463 #[test]
464 fn test_schema_registry_register() {
465 let mut registry = SchemaRegistry::new();
466 let mut schema = StructSchema::new("User");
467 let email = SchemaFieldInfo::new("email", "String", true, "encryption/email");
468 schema.add_field(email);
469 assert!(registry.register(schema).is_ok());
470 assert_eq!(registry.count(), 1);
471 }
472
473 #[test]
474 fn test_schema_registry_get() {
475 let mut registry = SchemaRegistry::new();
476 let mut schema = StructSchema::new("User");
477 let email = SchemaFieldInfo::new("email", "String", true, "encryption/email");
478 schema.add_field(email);
479 registry.register(schema).unwrap();
480 assert!(registry.get("User").is_some());
481 assert!(registry.get("Product").is_none());
482 }
483
484 #[test]
485 fn test_schema_registry_has_encrypted_fields() {
486 let mut registry = SchemaRegistry::new();
487 let mut schema = StructSchema::new("User");
488 let email = SchemaFieldInfo::new("email", "String", true, "encryption/email");
489 schema.add_field(email);
490 registry.register(schema).unwrap();
491 assert!(registry.has_encrypted_fields("User"));
492 assert!(!registry.has_encrypted_fields("Product"));
493 }
494
495 #[test]
496 fn test_schema_registry_list_types() {
497 let mut registry = SchemaRegistry::new();
498 let mut user_schema = StructSchema::new("User");
499 let email = SchemaFieldInfo::new("email", "String", true, "encryption/email");
500 user_schema.add_field(email);
501 registry.register(user_schema).unwrap();
502
503 let mut product_schema = StructSchema::new("Product");
504 let name = SchemaFieldInfo::new("name", "String", false, "");
505 product_schema.add_field(name);
506 registry.register(product_schema).unwrap();
507
508 let types = registry.list_types();
509 assert_eq!(types.len(), 2);
510 assert!(types.contains(&"User"));
511 assert!(types.contains(&"Product"));
512 }
513
514 #[test]
515 fn test_schema_registry_unregister() {
516 let mut registry = SchemaRegistry::new();
517 let mut schema = StructSchema::new("User");
518 let email = SchemaFieldInfo::new("email", "String", true, "encryption/email");
519 schema.add_field(email);
520 registry.register(schema).unwrap();
521 assert_eq!(registry.count(), 1);
522
523 registry.unregister("User");
524 assert_eq!(registry.count(), 0);
525 }
526
527 #[test]
528 fn test_schema_registry_clear() {
529 let mut registry = SchemaRegistry::new();
530 let mut schema = StructSchema::new("User");
531 let email = SchemaFieldInfo::new("email", "String", true, "encryption/email");
532 schema.add_field(email);
533 registry.register(schema).unwrap();
534 assert_eq!(registry.count(), 1);
535
536 registry.clear();
537 assert_eq!(registry.count(), 0);
538 }
539
540 #[test]
541 fn test_schema_registry_default_instance() {
542 let registry = SchemaRegistry::default();
543 assert_eq!(registry.count(), 0);
544 }
545
546 #[test]
547 fn test_struct_schema_nullable_encrypted_fields() {
548 let mut schema = StructSchema::new("User");
549 let email =
550 SchemaFieldInfo::new("email", "String", true, "encryption/email").with_nullable(true);
551 let phone =
552 SchemaFieldInfo::new("phone", "String", true, "encryption/phone").with_nullable(false);
553 let name = SchemaFieldInfo::new("name", "String", false, "").with_nullable(true);
554 schema.add_field(email);
555 schema.add_field(phone);
556 schema.add_field(name);
557 let nullable = schema.nullable_encrypted_fields();
558 assert_eq!(nullable.len(), 1);
559 assert_eq!(nullable[0].field_name, "email");
560 }
561
562 #[test]
563 fn test_struct_schema_fields_for_key() {
564 let mut schema = StructSchema::new("User");
565 let email = SchemaFieldInfo::new("email", "String", true, "encryption/email");
566 let phone = SchemaFieldInfo::new("phone", "String", true, "encryption/email");
567 let ssn = SchemaFieldInfo::new("ssn", "String", true, "encryption/ssn");
568 schema.add_field(email);
569 schema.add_field(phone);
570 schema.add_field(ssn);
571 let email_fields = schema.fields_for_key("encryption/email");
572 assert_eq!(email_fields.len(), 2);
573 }
574
575 #[test]
576 fn test_struct_schema_encrypted_field_count() {
577 let mut schema = StructSchema::new("User");
578 let email = SchemaFieldInfo::new("email", "String", true, "encryption/email");
579 let phone = SchemaFieldInfo::new("phone", "String", true, "encryption/phone");
580 schema.add_field(email);
581 schema.add_field(phone);
582 assert_eq!(schema.encrypted_field_count(), 2);
583 }
584
585 #[test]
586 fn test_struct_schema_total_field_count() {
587 let mut schema = StructSchema::new("User");
588 let email = SchemaFieldInfo::new("email", "String", true, "encryption/email");
589 let name = SchemaFieldInfo::new("name", "String", false, "");
590 schema.add_field(email);
591 schema.add_field(name);
592 assert_eq!(schema.total_field_count(), 2);
593 }
594
595 #[test]
596 fn test_schema_registry_types_with_encryption() {
597 let mut registry = SchemaRegistry::new();
598 let mut user_schema = StructSchema::new("User");
599 let email = SchemaFieldInfo::new("email", "String", true, "encryption/email");
600 user_schema.add_field(email);
601 registry.register(user_schema).unwrap();
602
603 let mut product_schema = StructSchema::new("Product");
604 let name = SchemaFieldInfo::new("name", "String", false, "");
605 product_schema.add_field(name);
606 registry.register(product_schema).unwrap();
607
608 let encrypted_types = registry.types_with_encryption();
609 assert_eq!(encrypted_types.len(), 1);
610 assert_eq!(encrypted_types[0], "User");
611 }
612
613 #[test]
614 fn test_schema_registry_all_encryption_keys() {
615 let mut registry = SchemaRegistry::new();
616 let mut user_schema = StructSchema::new("User");
617 let email = SchemaFieldInfo::new("email", "String", true, "encryption/email");
618 let phone = SchemaFieldInfo::new("phone", "String", true, "encryption/phone");
619 user_schema.add_field(email);
620 user_schema.add_field(phone);
621 registry.register(user_schema).unwrap();
622
623 let keys = registry.all_encryption_keys();
624 assert_eq!(keys.len(), 2);
625 assert!(keys.contains(&"encryption/email".to_string()));
626 assert!(keys.contains(&"encryption/phone".to_string()));
627 }
628
629 #[test]
630 fn test_schema_registry_validate_all() {
631 let mut registry = SchemaRegistry::new();
632 let mut schema = StructSchema::new("User");
633 let email = SchemaFieldInfo::new("email", "String", true, "encryption/email");
634 schema.add_field(email);
635 registry.register(schema).unwrap();
636 assert!(registry.validate_all().is_ok());
637 }
638
639 #[test]
640 fn test_schema_registry_total_encrypted_fields() {
641 let mut registry = SchemaRegistry::new();
642 let mut user_schema = StructSchema::new("User");
643 let email = SchemaFieldInfo::new("email", "String", true, "encryption/email");
644 let phone = SchemaFieldInfo::new("phone", "String", true, "encryption/phone");
645 user_schema.add_field(email);
646 user_schema.add_field(phone);
647 registry.register(user_schema).unwrap();
648
649 let mut product_schema = StructSchema::new("Product");
650 let sku = SchemaFieldInfo::new("sku", "String", true, "encryption/sku");
651 product_schema.add_field(sku);
652 registry.register(product_schema).unwrap();
653
654 assert_eq!(registry.total_encrypted_fields(), 3);
655 }
656}