chie_crypto/
searchable.rs1use blake3::Hasher;
50use rand::RngCore;
51use serde::{Deserialize, Serialize};
52use std::collections::{HashMap, HashSet};
53use thiserror::Error;
54
55#[derive(Error, Debug)]
56pub enum SearchableError {
57 #[error("Invalid trapdoor")]
58 InvalidTrapdoor,
59 #[error("Serialization error: {0}")]
60 Serialization(String),
61 #[error("Empty keyword")]
62 EmptyKeyword,
63}
64
65pub type SearchableResult<T> = Result<T, SearchableError>;
66
67pub struct SearchableEncryption {
69 master_key: [u8; 32],
70}
71
72impl SearchableEncryption {
73 pub fn new() -> Self {
75 let mut master_key = [0u8; 32];
76 rand::rngs::OsRng.fill_bytes(&mut master_key);
77 Self { master_key }
78 }
79
80 pub fn with_key(key: [u8; 32]) -> Self {
82 Self { master_key: key }
83 }
84
85 pub fn generate_trapdoor(&self, keyword: &[u8]) -> Trapdoor {
89 let token = self.compute_keyword_token(keyword);
90 Trapdoor { token }
91 }
92
93 fn encrypt_keyword(&self, keyword: &[u8]) -> Vec<u8> {
95 self.compute_keyword_token(keyword)
96 }
97
98 fn compute_keyword_token(&self, keyword: &[u8]) -> Vec<u8> {
100 let mut hasher = Hasher::new();
101 hasher.update(&self.master_key);
102 hasher.update(b"keyword-token");
103 hasher.update(keyword);
104 hasher.finalize().as_bytes()[..].to_vec()
105 }
106
107 pub fn export_key(&self) -> [u8; 32] {
109 self.master_key
110 }
111}
112
113impl Default for SearchableEncryption {
114 fn default() -> Self {
115 Self::new()
116 }
117}
118
119#[derive(Clone, Serialize, Deserialize)]
121pub struct Trapdoor {
122 token: Vec<u8>,
123}
124
125impl Trapdoor {
126 pub fn to_bytes(&self) -> SearchableResult<Vec<u8>> {
128 crate::codec::encode(self).map_err(|e| SearchableError::Serialization(e.to_string()))
129 }
130
131 pub fn from_bytes(bytes: &[u8]) -> SearchableResult<Self> {
133 crate::codec::decode(bytes).map_err(|e| SearchableError::Serialization(e.to_string()))
134 }
135}
136
137pub type DocumentId = u64;
139
140pub struct EncryptedIndex {
142 index: HashMap<Vec<u8>, HashSet<DocumentId>>,
144 sse_key: [u8; 32],
146}
147
148impl EncryptedIndex {
149 pub fn new(sse: &SearchableEncryption) -> Self {
151 Self {
152 index: HashMap::new(),
153 sse_key: sse.master_key,
154 }
155 }
156
157 pub fn add_document(&mut self, doc_id: DocumentId, keywords: &[Vec<u8>]) {
163 let sse = SearchableEncryption::with_key(self.sse_key);
164
165 for keyword in keywords {
166 let encrypted_keyword = sse.encrypt_keyword(keyword);
167 self.index
168 .entry(encrypted_keyword)
169 .or_default()
170 .insert(doc_id);
171 }
172 }
173
174 pub fn remove_document(&mut self, doc_id: DocumentId) {
176 for doc_set in self.index.values_mut() {
178 doc_set.remove(&doc_id);
179 }
180
181 self.index.retain(|_, doc_set| !doc_set.is_empty());
183 }
184
185 pub fn search(&self, trapdoor: &Trapdoor) -> Vec<DocumentId> {
189 self.index
190 .get(&trapdoor.token)
191 .map(|doc_set| doc_set.iter().copied().collect())
192 .unwrap_or_default()
193 }
194
195 pub fn keyword_count(&self) -> usize {
197 self.index.len()
198 }
199
200 pub fn document_count(&self) -> usize {
202 let mut all_docs: HashSet<DocumentId> = HashSet::new();
203 for doc_set in self.index.values() {
204 all_docs.extend(doc_set);
205 }
206 all_docs.len()
207 }
208}
209
210pub struct MultiKeywordSearch<'a> {
212 index: &'a EncryptedIndex,
213}
214
215impl<'a> MultiKeywordSearch<'a> {
216 pub fn new(index: &'a EncryptedIndex) -> Self {
218 Self { index }
219 }
220
221 pub fn search_and(&self, trapdoors: &[Trapdoor]) -> Vec<DocumentId> {
223 if trapdoors.is_empty() {
224 return Vec::new();
225 }
226
227 let mut result: HashSet<DocumentId> =
229 self.index.search(&trapdoors[0]).into_iter().collect();
230
231 for trapdoor in &trapdoors[1..] {
233 let docs: HashSet<DocumentId> = self.index.search(trapdoor).into_iter().collect();
234 result.retain(|doc_id| docs.contains(doc_id));
235 }
236
237 result.into_iter().collect()
238 }
239
240 pub fn search_or(&self, trapdoors: &[Trapdoor]) -> Vec<DocumentId> {
242 let mut result = HashSet::new();
243
244 for trapdoor in trapdoors {
245 let docs = self.index.search(trapdoor);
246 result.extend(docs);
247 }
248
249 result.into_iter().collect()
250 }
251}
252
253pub struct EncryptedIndexBuilder {
255 index: HashMap<Vec<u8>, HashSet<DocumentId>>,
256 sse_key: [u8; 32],
257}
258
259impl EncryptedIndexBuilder {
260 pub fn new(sse: &SearchableEncryption) -> Self {
262 Self {
263 index: HashMap::new(),
264 sse_key: sse.master_key,
265 }
266 }
267
268 pub fn add_document(mut self, doc_id: DocumentId, keywords: &[Vec<u8>]) -> Self {
270 let sse = SearchableEncryption::with_key(self.sse_key);
271
272 for keyword in keywords {
273 let encrypted_keyword = sse.encrypt_keyword(keyword);
274 self.index
275 .entry(encrypted_keyword)
276 .or_default()
277 .insert(doc_id);
278 }
279
280 self
281 }
282
283 pub fn build(self) -> EncryptedIndex {
285 EncryptedIndex {
286 index: self.index,
287 sse_key: self.sse_key,
288 }
289 }
290}
291
292#[cfg(test)]
293mod tests {
294 use super::*;
295
296 #[test]
297 fn test_searchable_basic() {
298 let sse = SearchableEncryption::new();
299 let mut index = EncryptedIndex::new(&sse);
300
301 index.add_document(1, &[b"rust".to_vec(), b"crypto".to_vec()]);
302 index.add_document(2, &[b"crypto".to_vec()]);
303
304 let trapdoor = sse.generate_trapdoor(b"crypto");
305 let results = index.search(&trapdoor);
306
307 assert_eq!(results.len(), 2);
308 assert!(results.contains(&1));
309 assert!(results.contains(&2));
310 }
311
312 #[test]
313 fn test_no_matches() {
314 let sse = SearchableEncryption::new();
315 let mut index = EncryptedIndex::new(&sse);
316
317 index.add_document(1, &[b"rust".to_vec()]);
318
319 let trapdoor = sse.generate_trapdoor(b"python");
320 let results = index.search(&trapdoor);
321
322 assert!(results.is_empty());
323 }
324
325 #[test]
326 fn test_remove_document() {
327 let sse = SearchableEncryption::new();
328 let mut index = EncryptedIndex::new(&sse);
329
330 index.add_document(1, &[b"keyword".to_vec()]);
331 index.add_document(2, &[b"keyword".to_vec()]);
332
333 let trapdoor = sse.generate_trapdoor(b"keyword");
334 assert_eq!(index.search(&trapdoor).len(), 2);
335
336 index.remove_document(1);
337 let results = index.search(&trapdoor);
338 assert_eq!(results.len(), 1);
339 assert!(results.contains(&2));
340 }
341
342 #[test]
343 fn test_keyword_count() {
344 let sse = SearchableEncryption::new();
345 let mut index = EncryptedIndex::new(&sse);
346
347 index.add_document(1, &[b"key1".to_vec(), b"key2".to_vec()]);
348 index.add_document(2, &[b"key2".to_vec(), b"key3".to_vec()]);
349
350 assert_eq!(index.keyword_count(), 3);
351 }
352
353 #[test]
354 fn test_document_count() {
355 let sse = SearchableEncryption::new();
356 let mut index = EncryptedIndex::new(&sse);
357
358 index.add_document(1, &[b"key1".to_vec()]);
359 index.add_document(2, &[b"key2".to_vec()]);
360 index.add_document(3, &[b"key3".to_vec()]);
361
362 assert_eq!(index.document_count(), 3);
363 }
364
365 #[test]
366 fn test_multi_keyword_and() {
367 let sse = SearchableEncryption::new();
368 let mut index = EncryptedIndex::new(&sse);
369
370 index.add_document(1, &[b"rust".to_vec(), b"crypto".to_vec()]);
371 index.add_document(2, &[b"crypto".to_vec()]);
372 index.add_document(3, &[b"rust".to_vec(), b"crypto".to_vec()]);
373
374 let trapdoors = vec![
375 sse.generate_trapdoor(b"rust"),
376 sse.generate_trapdoor(b"crypto"),
377 ];
378
379 let search = MultiKeywordSearch::new(&index);
380 let results = search.search_and(&trapdoors);
381
382 assert_eq!(results.len(), 2);
383 assert!(results.contains(&1));
384 assert!(results.contains(&3));
385 }
386
387 #[test]
388 fn test_multi_keyword_or() {
389 let sse = SearchableEncryption::new();
390 let mut index = EncryptedIndex::new(&sse);
391
392 index.add_document(1, &[b"rust".to_vec()]);
393 index.add_document(2, &[b"go".to_vec()]);
394 index.add_document(3, &[b"python".to_vec()]);
395
396 let trapdoors = vec![sse.generate_trapdoor(b"rust"), sse.generate_trapdoor(b"go")];
397
398 let search = MultiKeywordSearch::new(&index);
399 let results = search.search_or(&trapdoors);
400
401 assert_eq!(results.len(), 2);
402 assert!(results.contains(&1));
403 assert!(results.contains(&2));
404 }
405
406 #[test]
407 fn test_trapdoor_serialization() {
408 let sse = SearchableEncryption::new();
409 let trapdoor = sse.generate_trapdoor(b"keyword");
410
411 let bytes = trapdoor.to_bytes().unwrap();
412 let deserialized = Trapdoor::from_bytes(&bytes).unwrap();
413
414 let mut index = EncryptedIndex::new(&sse);
415 index.add_document(1, &[b"keyword".to_vec()]);
416
417 let results1 = index.search(&trapdoor);
418 let results2 = index.search(&deserialized);
419 assert_eq!(results1, results2);
420 }
421
422 #[test]
423 fn test_deterministic_trapdoor() {
424 let sse = SearchableEncryption::new();
425
426 let td1 = sse.generate_trapdoor(b"keyword");
427 let td2 = sse.generate_trapdoor(b"keyword");
428
429 let bytes1 = td1.to_bytes().unwrap();
430 let bytes2 = td2.to_bytes().unwrap();
431 assert_eq!(bytes1, bytes2);
432 }
433
434 #[test]
435 fn test_index_builder() {
436 let sse = SearchableEncryption::new();
437
438 let index = EncryptedIndexBuilder::new(&sse)
439 .add_document(1, &[b"rust".to_vec()])
440 .add_document(2, &[b"crypto".to_vec()])
441 .build();
442
443 assert_eq!(index.document_count(), 2);
444 assert_eq!(index.keyword_count(), 2);
445 }
446
447 #[test]
448 fn test_sse_default() {
449 let sse = SearchableEncryption::default();
450 let trapdoor = sse.generate_trapdoor(b"test");
451 assert!(!trapdoor.to_bytes().unwrap().is_empty());
452 }
453
454 #[test]
455 fn test_export_import_key() {
456 let sse1 = SearchableEncryption::new();
457 let key = sse1.export_key();
458 let sse2 = SearchableEncryption::with_key(key);
459
460 let td1 = sse1.generate_trapdoor(b"test");
461 let td2 = sse2.generate_trapdoor(b"test");
462
463 assert_eq!(td1.to_bytes().unwrap(), td2.to_bytes().unwrap());
464 }
465
466 #[test]
467 fn test_empty_trapdoors_and() {
468 let sse = SearchableEncryption::new();
469 let index = EncryptedIndex::new(&sse);
470 let search = MultiKeywordSearch::new(&index);
471
472 let results = search.search_and(&[]);
473 assert!(results.is_empty());
474 }
475
476 #[test]
477 fn test_empty_trapdoors_or() {
478 let sse = SearchableEncryption::new();
479 let index = EncryptedIndex::new(&sse);
480 let search = MultiKeywordSearch::new(&index);
481
482 let results = search.search_or(&[]);
483 assert!(results.is_empty());
484 }
485}