1use std::{ops::Deref, sync::Arc};
2
3use object_rainbow::{
4 Address, ByteNode, Error, FailFuture, Fetch, FetchBytes, FullHash, Hash, ListHashes, Node,
5 Object, Parse, ParseSliceExtra, PointInput, PointVisitor, Resolve, Singular, SingularFetch,
6 Tagged, ToOutput, Topological, Traversible, length_prefixed::Lp,
7};
8use object_rainbow_point::{ExtractResolve, IntoPoint, Point};
9
10pub trait Key: 'static + Sized + Send + Sync + Clone {
11 fn encrypt(&self, data: &[u8]) -> Vec<u8>;
12 fn decrypt(&self, data: &[u8]) -> object_rainbow::Result<Vec<u8>>;
13}
14
15type Resolution<K> = Arc<Lp<Vec<Point<Encrypted<K, Vec<u8>>>>>>;
16
17#[derive(ToOutput, Clone)]
18struct Unkeyed<T>(T);
19
20impl<
21 T: Parse<I::WithExtra<Extra>>,
22 K: 'static + Clone,
23 Extra: 'static + Clone,
24 I: PointInput<Extra = (K, Extra)>,
25> Parse<I> for Unkeyed<T>
26{
27 fn parse(input: I) -> object_rainbow::Result<Self> {
28 Ok(Self(T::parse(input.map_extra(|(_, extra)| extra))?))
29 }
30}
31
32#[derive(ToOutput, Parse)]
33struct EncryptedInner<K, T> {
34 resolution: Resolution<K>,
35 decrypted: Unkeyed<Arc<T>>,
36}
37
38impl<K, T> Clone for EncryptedInner<K, T> {
39 fn clone(&self) -> Self {
40 Self {
41 resolution: self.resolution.clone(),
42 decrypted: self.decrypted.clone(),
43 }
44 }
45}
46
47type ResolutionIter<'a, K> = std::slice::Iter<'a, Point<Encrypted<K, Vec<u8>>>>;
48
49struct IterateResolution<'a, 'r, K, V> {
50 resolution: &'r mut ResolutionIter<'a, K>,
51 visitor: &'a mut V,
52}
53
54struct Visited<K, P> {
55 decrypted: P,
56 encrypted: Point<Encrypted<K, Vec<u8>>>,
57}
58
59impl<K, P> FetchBytes for Visited<K, P> {
60 fn fetch_bytes(&'_ self) -> FailFuture<'_, ByteNode> {
61 self.encrypted.fetch_bytes()
62 }
63
64 fn fetch_data(&'_ self) -> FailFuture<'_, Vec<u8>> {
65 self.encrypted.fetch_data()
66 }
67
68 fn fetch_bytes_local(&self) -> object_rainbow::Result<Option<ByteNode>> {
69 self.encrypted.fetch_bytes_local()
70 }
71
72 fn fetch_data_local(&self) -> Option<Vec<u8>> {
73 self.encrypted.fetch_data_local()
74 }
75}
76
77impl<K, P: Send + Sync> Singular for Visited<K, P> {
78 fn hash(&self) -> Hash {
79 self.encrypted.hash()
80 }
81}
82
83impl<K: Key, P: Fetch<T: Traversible>> Fetch for Visited<K, P> {
84 type T = Encrypted<K, P::T>;
85
86 fn fetch_full(&'_ self) -> FailFuture<'_, Node<Self::T>> {
87 Box::pin(async move {
88 let (
89 Encrypted {
90 key,
91 inner:
92 EncryptedInner {
93 resolution,
94 decrypted: _,
95 },
96 },
97 resolve,
98 ) = self.encrypted.fetch_full().await?;
99 let decrypted = self.decrypted.fetch().await?;
100 let decrypted = Unkeyed(Arc::new(decrypted));
101 Ok((
102 Encrypted {
103 key,
104 inner: EncryptedInner {
105 resolution,
106 decrypted,
107 },
108 },
109 resolve,
110 ))
111 })
112 }
113
114 fn fetch(&'_ self) -> FailFuture<'_, Self::T> {
115 Box::pin(async move {
116 let Encrypted {
117 key,
118 inner:
119 EncryptedInner {
120 resolution,
121 decrypted: _,
122 },
123 } = self.encrypted.fetch().await?;
124 let decrypted = self.decrypted.fetch().await?;
125 let decrypted = Unkeyed(Arc::new(decrypted));
126 Ok(Encrypted {
127 key,
128 inner: EncryptedInner {
129 resolution,
130 decrypted,
131 },
132 })
133 })
134 }
135
136 fn try_fetch_local(&self) -> object_rainbow::Result<Option<Node<Self::T>>> {
137 let Some((
138 Encrypted {
139 key,
140 inner:
141 EncryptedInner {
142 resolution,
143 decrypted: _,
144 },
145 },
146 resolve,
147 )) = self.encrypted.try_fetch_local()?
148 else {
149 return Ok(None);
150 };
151 let Some((decrypted, _)) = self.decrypted.try_fetch_local()? else {
152 return Ok(None);
153 };
154 let decrypted = Unkeyed(Arc::new(decrypted));
155 Ok(Some((
156 Encrypted {
157 key,
158 inner: EncryptedInner {
159 resolution,
160 decrypted,
161 },
162 },
163 resolve,
164 )))
165 }
166
167 fn fetch_local(&self) -> Option<Self::T> {
168 let Encrypted {
169 key,
170 inner:
171 EncryptedInner {
172 resolution,
173 decrypted: _,
174 },
175 } = self.encrypted.fetch_local()?;
176 let decrypted = Unkeyed(Arc::new(self.decrypted.fetch_local()?));
177 Some(Encrypted {
178 key,
179 inner: EncryptedInner {
180 resolution,
181 decrypted,
182 },
183 })
184 }
185}
186
187impl<'a, K: Key, V: PointVisitor> PointVisitor for IterateResolution<'a, '_, K, V> {
188 fn visit<T: Traversible>(&mut self, decrypted: &(impl 'static + SingularFetch<T = T> + Clone)) {
189 let decrypted = decrypted.clone();
190 let encrypted = self.resolution.next().expect("length mismatch").clone();
191 let point = Point::from_fetch(
192 encrypted.hash(),
193 Arc::new(Visited {
194 decrypted,
195 encrypted,
196 }),
197 );
198 self.visitor.visit(&point);
199 }
200}
201
202impl<K, T> ListHashes for EncryptedInner<K, T> {
203 fn list_hashes(&self, f: &mut impl FnMut(Hash)) {
204 self.resolution.list_hashes(f);
205 }
206
207 fn topology_hash(&self) -> Hash {
208 self.resolution.0.data_hash()
209 }
210
211 fn point_count(&self) -> usize {
212 self.resolution.len()
213 }
214}
215
216impl<K: Key, T: Topological> Topological for EncryptedInner<K, T> {
217 fn traverse(&self, visitor: &mut impl PointVisitor) {
218 let resolution = &mut self.resolution.iter();
219 self.decrypted.0.traverse(&mut IterateResolution {
220 resolution,
221 visitor,
222 });
223 assert!(resolution.next().is_none());
224 }
225}
226
227pub struct Encrypted<K, T> {
228 key: K,
229 inner: EncryptedInner<K, T>,
230}
231
232impl<K, T: Clone> Encrypted<K, T> {
233 pub fn into_inner(self) -> T {
234 Arc::unwrap_or_clone(self.inner.decrypted.0)
235 }
236}
237
238impl<K, T> Deref for Encrypted<K, T> {
239 type Target = T;
240
241 fn deref(&self) -> &Self::Target {
242 self.inner.decrypted.0.as_ref()
243 }
244}
245
246impl<K: Clone, T> Clone for Encrypted<K, T> {
247 fn clone(&self) -> Self {
248 Self {
249 key: self.key.clone(),
250 inner: self.inner.clone(),
251 }
252 }
253}
254
255impl<K, T> ListHashes for Encrypted<K, T> {
256 fn list_hashes(&self, f: &mut impl FnMut(Hash)) {
257 self.inner.list_hashes(f);
258 }
259
260 fn topology_hash(&self) -> Hash {
261 self.inner.topology_hash()
262 }
263
264 fn point_count(&self) -> usize {
265 self.inner.point_count()
266 }
267}
268
269impl<K: Key, T: Topological> Topological for Encrypted<K, T> {
270 fn traverse(&self, visitor: &mut impl PointVisitor) {
271 self.inner.traverse(visitor);
272 }
273}
274
275impl<K: Key, T: ToOutput> ToOutput for Encrypted<K, T> {
276 fn to_output(&self, output: &mut dyn object_rainbow::Output) {
277 let source = self.inner.vec();
278 output.write(&self.key.encrypt(&source));
279 }
280}
281
282#[derive(Clone)]
283struct Decrypt<K> {
284 resolution: Resolution<K>,
285}
286
287impl<K: Key> Decrypt<K> {
288 async fn resolve_bytes(
289 &self,
290 address: Address,
291 ) -> object_rainbow::Result<(Vec<u8>, Resolution<K>)> {
292 let Encrypted {
293 key: _,
294 inner:
295 EncryptedInner {
296 resolution,
297 decrypted,
298 },
299 } = self
300 .resolution
301 .get(address.index)
302 .ok_or(Error::AddressOutOfBounds)?
303 .clone()
304 .fetch()
305 .await?;
306 let data = Arc::unwrap_or_clone(decrypted.0);
307 Ok((data, resolution))
308 }
309}
310
311impl<K: Key> Resolve for Decrypt<K> {
312 fn resolve(&'_ self, address: Address) -> FailFuture<'_, ByteNode> {
313 Box::pin(async move {
314 let (data, resolution) = self.resolve_bytes(address).await?;
315 Ok((data, Arc::new(Decrypt { resolution }) as _))
316 })
317 }
318
319 fn resolve_data(&'_ self, address: Address) -> FailFuture<'_, Vec<u8>> {
320 Box::pin(async move {
321 let (data, _) = self.resolve_bytes(address).await?;
322 Ok(data)
323 })
324 }
325
326 fn try_resolve_local(&self, address: Address) -> object_rainbow::Result<Option<ByteNode>> {
327 let Some((
328 Encrypted {
329 key: _,
330 inner:
331 EncryptedInner {
332 resolution,
333 decrypted,
334 },
335 },
336 _,
337 )) = self
338 .resolution
339 .get(address.index)
340 .ok_or(Error::AddressOutOfBounds)?
341 .clone()
342 .try_fetch_local()?
343 else {
344 return Ok(None);
345 };
346 let data = Arc::unwrap_or_clone(decrypted.0);
347 Ok(Some((data, Arc::new(Decrypt { resolution }) as _)))
348 }
349}
350
351impl<
352 K: Key,
353 T: Object<Extra>,
354 Extra: 'static + Send + Sync + Clone,
355 I: PointInput<Extra = (K, Extra)>,
356> Parse<I> for Encrypted<K, T>
357{
358 fn parse(input: I) -> object_rainbow::Result<Self> {
359 let with_key = input.extra().clone();
360 let resolve = input.resolve().clone();
361 let source = with_key.0.decrypt(&input.parse_all()?)?;
362 let EncryptedInner {
363 resolution,
364 decrypted,
365 } = EncryptedInner::<K, Vec<u8>>::parse_slice_extra(&source, &resolve, &with_key)?;
366 let decrypted = T::parse_slice_extra(
367 &decrypted.0,
368 &(Arc::new(Decrypt {
369 resolution: resolution.clone(),
370 }) as _),
371 &with_key.1,
372 )?;
373 let decrypted = Unkeyed(Arc::new(decrypted));
374 let inner = EncryptedInner {
375 resolution,
376 decrypted,
377 };
378 Ok(Self {
379 key: with_key.0,
380 inner,
381 })
382 }
383}
384
385impl<K, T> Tagged for Encrypted<K, T> {}
386
387type Extracted<K> = Vec<
388 std::pin::Pin<
389 Box<dyn Future<Output = Result<Point<Encrypted<K, Vec<u8>>>, Error>> + Send + 'static>,
390 >,
391>;
392
393struct ExtractResolution<'a, K> {
394 extracted: &'a mut Extracted<K>,
395 key: &'a K,
396}
397
398struct Untyped<K, T> {
399 key: (K, ()),
400 encrypted: Point<Encrypted<K, T>>,
401}
402
403impl<K, T> FetchBytes for Untyped<K, T> {
404 fn fetch_bytes(&'_ self) -> FailFuture<'_, ByteNode> {
405 self.encrypted.fetch_bytes()
406 }
407
408 fn fetch_data(&'_ self) -> FailFuture<'_, Vec<u8>> {
409 self.encrypted.fetch_data()
410 }
411
412 fn fetch_bytes_local(&self) -> object_rainbow::Result<Option<ByteNode>> {
413 self.encrypted.fetch_bytes_local()
414 }
415
416 fn fetch_data_local(&self) -> Option<Vec<u8>> {
417 self.encrypted.fetch_data_local()
418 }
419}
420
421impl<K: Send + Sync, T> Singular for Untyped<K, T> {
422 fn hash(&self) -> Hash {
423 self.encrypted.hash()
424 }
425}
426
427impl<K: Key, T: FullHash> Fetch for Untyped<K, T> {
428 type T = Encrypted<K, Vec<u8>>;
429
430 fn fetch_full(&'_ self) -> FailFuture<'_, Node<Self::T>> {
431 Box::pin(async move {
432 let (data, resolve) = self.fetch_bytes().await?;
433 let encrypted = Self::T::parse_slice_extra(&data, &resolve, &self.key)?;
434 Ok((encrypted, resolve))
435 })
436 }
437
438 fn fetch(&'_ self) -> FailFuture<'_, Self::T> {
439 Box::pin(async move {
440 let (data, resolve) = self.fetch_bytes().await?;
441 let encrypted = Self::T::parse_slice_extra(&data, &resolve, &self.key)?;
442 Ok(encrypted)
443 })
444 }
445
446 fn try_fetch_local(&self) -> object_rainbow::Result<Option<Node<Self::T>>> {
447 let Some((data, resolve)) = self.fetch_bytes_local()? else {
448 return Ok(None);
449 };
450 let encrypted = Self::T::parse_slice_extra(&data, &resolve, &self.key)?;
451 Ok(Some((encrypted, resolve)))
452 }
453
454 fn fetch_local(&self) -> Option<Self::T> {
455 let Encrypted {
456 key,
457 inner:
458 EncryptedInner {
459 resolution,
460 decrypted,
461 },
462 } = self.encrypted.fetch_local()?;
463 let decrypted = Unkeyed(Arc::new(decrypted.vec()));
464 Some(Encrypted {
465 key,
466 inner: EncryptedInner {
467 resolution,
468 decrypted,
469 },
470 })
471 }
472}
473
474impl<K: Key> PointVisitor for ExtractResolution<'_, K> {
475 fn visit<T: Traversible>(&mut self, decrypted: &(impl 'static + SingularFetch<T = T> + Clone)) {
476 let decrypted = decrypted.clone();
477 let key = self.key.clone();
478 self.extracted.push(Box::pin(async move {
479 let encrypted = encrypt_point(key.clone(), decrypted).await?;
480 let encrypted = Point::from_fetch(
481 encrypted.hash(),
482 Arc::new(Untyped {
483 key: (key, ()),
484 encrypted,
485 }),
486 );
487 Ok(encrypted)
488 }));
489 }
490}
491
492pub async fn encrypt_point<K: Key, T: Traversible>(
493 key: K,
494 decrypted: impl 'static + SingularFetch<T = T>,
495) -> object_rainbow::Result<Point<Encrypted<K, T>>> {
496 if let Some((address, decrypt)) = decrypted.extract_resolve::<Decrypt<K>>() {
497 let encrypted = decrypt
498 .resolution
499 .get(address.index)
500 .ok_or(Error::AddressOutOfBounds)?
501 .clone();
502 let point = Point::from_fetch(
503 encrypted.hash(),
504 Arc::new(Visited {
505 decrypted,
506 encrypted,
507 }),
508 );
509 return Ok(point);
510 }
511 let decrypted = decrypted.fetch().await?;
512 let encrypted = encrypt(key.clone(), decrypted).await?;
513 let point = encrypted.point();
514 Ok(point)
515}
516
517pub async fn encrypt<K: Key, T: Traversible>(
518 key: K,
519 decrypted: T,
520) -> object_rainbow::Result<Encrypted<K, T>> {
521 let mut futures = Vec::with_capacity(decrypted.point_count());
522 decrypted.traverse(&mut ExtractResolution {
523 extracted: &mut futures,
524 key: &key,
525 });
526 let resolution = futures_util::future::try_join_all(futures).await?;
527 let resolution = Arc::new(Lp(resolution));
528 let decrypted = Unkeyed(Arc::new(decrypted));
529 let inner = EncryptedInner {
530 resolution,
531 decrypted,
532 };
533 Ok(Encrypted { key, inner })
534}