Skip to main content

fraiseql_server/encryption/
schema.rs

1// Phase 12.3 Cycle 6: Schema Detection (GREEN)
2//! Schema detection for automatically identifying and managing encrypted fields.
3//!
4//! Supports multiple encryption marks (`#[encrypted]`, `#[sensitive]`, `#[encrypt(key="...")]`),
5//! key reference management, and schema evolution tracking.
6
7use std::collections::HashMap;
8
9use crate::secrets_manager::SecretsError;
10
11/// Encryption mark type used in struct annotations
12#[derive(Debug, Clone, Copy, PartialEq, Eq)]
13pub enum EncryptionMark {
14    /// Basic `#[encrypted]` mark
15    Encrypted,
16    /// Alternative `#[sensitive]` mark
17    Sensitive,
18    /// Explicit `#[encrypt(...)]` with configuration
19    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/// Metadata about an encrypted field in schema
33#[derive(Debug, Clone)]
34pub struct SchemaFieldInfo {
35    /// Field name in struct
36    pub field_name:    String,
37    /// Field type (e.g., "String", "Uuid", "`DateTime<Utc>`")
38    pub field_type:    String,
39    /// Whether field is marked for encryption
40    pub is_encrypted:  bool,
41    /// Key reference path for encryption (e.g., "encryption/email")
42    pub key_reference: String,
43    /// Encryption algorithm hint
44    pub algorithm:     String,
45    /// Whether field can be NULL
46    pub nullable:      bool,
47    /// Which encryption mark was used
48    pub mark:          Option<EncryptionMark>,
49}
50
51impl SchemaFieldInfo {
52    /// Create new field info
53    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    /// Set algorithm hint
71    pub fn with_algorithm(mut self, algorithm: impl Into<String>) -> Self {
72        self.algorithm = algorithm.into();
73        self
74    }
75
76    /// Mark as nullable
77    pub fn with_nullable(mut self, nullable: bool) -> Self {
78        self.nullable = nullable;
79        self
80    }
81
82    /// Set encryption mark
83    pub fn with_mark(mut self, mark: EncryptionMark) -> Self {
84        self.mark = Some(mark);
85        self
86    }
87}
88
89/// Schema information for a struct type
90#[derive(Debug, Clone)]
91pub struct StructSchema {
92    /// Type name (e.g., "User")
93    pub type_name:        String,
94    /// All fields in struct (including non-encrypted)
95    pub all_fields:       Vec<SchemaFieldInfo>,
96    /// Only encrypted fields (subset of all_fields)
97    pub encrypted_fields: Vec<SchemaFieldInfo>,
98    /// Schema version for evolution tracking
99    pub version:          u32,
100}
101
102impl StructSchema {
103    /// Create new struct schema
104    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    /// Add field to schema
114    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    /// Add multiple fields
122    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    /// Set schema version for evolution tracking
130    pub fn with_version(mut self, version: u32) -> Self {
131        self.version = version;
132        self
133    }
134
135    /// Get field by name
136    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    /// Get encrypted field by name
141    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    /// Check if field is encrypted
146    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    /// Get list of encrypted field names
151    pub fn encrypted_field_names(&self) -> Vec<&str> {
152        self.encrypted_fields.iter().map(|f| f.field_name.as_str()).collect()
153    }
154
155    /// Internal filter helper to reduce duplication
156    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    /// Get fields that are marked as nullable
164    pub fn nullable_encrypted_fields(&self) -> Vec<&SchemaFieldInfo> {
165        self.filter_fields(|f| f.is_encrypted && f.nullable)
166    }
167
168    /// Get fields requiring specific encryption key
169    pub fn fields_for_key(&self, key_ref: &str) -> Vec<&SchemaFieldInfo> {
170        self.filter_fields(|f| f.key_reference == key_ref)
171    }
172
173    /// Count encrypted fields
174    pub fn encrypted_field_count(&self) -> usize {
175        self.encrypted_fields.len()
176    }
177
178    /// Count total fields
179    pub fn total_field_count(&self) -> usize {
180        self.all_fields.len()
181    }
182
183    /// Validate schema configuration
184    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        // Validate each encrypted field has key reference
192        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
205/// Registry for managing schemas of different types
206pub struct SchemaRegistry {
207    /// Map of type name to schema
208    schemas:               HashMap<String, StructSchema>,
209    /// Default key reference for fields without explicit key
210    default_key_reference: String,
211}
212
213impl SchemaRegistry {
214    /// Create new schema registry
215    pub fn new() -> Self {
216        Self {
217            schemas:               HashMap::new(),
218            default_key_reference: "encryption/default".to_string(),
219        }
220    }
221
222    /// Set default key reference
223    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    /// Register schema
229    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    /// Get schema by type name
236    pub fn get(&self, type_name: &str) -> Option<&StructSchema> {
237        self.schemas.get(type_name)
238    }
239
240    /// Get encrypted fields for type
241    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    /// Check if type has encrypted fields
253    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    /// Get list of all registered types
260    pub fn list_types(&self) -> Vec<&str> {
261        self.schemas.keys().map(|s| s.as_str()).collect()
262    }
263
264    /// Get list of all types that have encrypted fields
265    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    /// Get all encryption keys used across all schemas
274    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    /// Validate all registered schemas
287    pub fn validate_all(&self) -> Result<(), SecretsError> {
288        for schema in self.schemas.values() {
289            schema.validate()?;
290        }
291        Ok(())
292    }
293
294    /// Unregister schema
295    pub fn unregister(&mut self, type_name: &str) -> Option<StructSchema> {
296        self.schemas.remove(type_name)
297    }
298
299    /// Clear all schemas
300    pub fn clear(&mut self) {
301        self.schemas.clear();
302    }
303
304    /// Count registered schemas
305    pub fn count(&self) -> usize {
306        self.schemas.len()
307    }
308
309    /// Count total encrypted fields across all schemas
310    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}