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