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