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