ossa_core/
util.rs

1pub mod merkle_tree;
2
3use bytes::{Bytes, BytesMut};
4use futures;
5use futures::task::{Context, Poll};
6use ossa_typeable::Typeable;
7use rand::{rngs::OsRng, TryRngCore};
8use serde::{Deserialize, Serialize};
9use sha2::{Digest, Sha256};
10use std::fmt::{self, Debug, Display};
11use std::marker::PhantomData;
12use std::ops::{Add, Range};
13use std::pin::Pin;
14use tokio::sync::mpsc::{Receiver, Sender};
15
16use crate::network::protocol::ProtocolError;
17use crate::store::Nonce;
18
19/// Generate a random nonce.
20pub(crate) fn generate_nonce() -> Nonce {
21    let mut nonce = [0; 32];
22    OsRng
23        .try_fill_bytes(&mut nonce)
24        .expect("Nonce generation failed");
25    nonce
26}
27
28#[test]
29fn test() {
30    let n1 = generate_nonce();
31    let n2 = generate_nonce();
32    // println!("Nonce: {:?}", n1);
33    // println!("Nonce: {:?}", n2);
34    assert!(n1 != n2);
35}
36
37// TODO: Remove this trait
38pub trait Hash: PartialEq + AsRef<[u8]> + Copy {
39    type HashState;
40
41    fn new() -> Self::HashState;
42    fn update(state: &mut Self::HashState, data: impl AsRef<[u8]>);
43    fn finalize(state: Self::HashState) -> Self;
44}
45
46#[derive(Clone, Copy, Deserialize, Eq, Ord, PartialOrd, Serialize, Typeable)]
47pub struct Sha256Hash(pub [u8; 32]);
48
49impl Debug for Sha256Hash {
50    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> {
51        write!(f, "0x")?;
52        for b in self.0 {
53            write!(f, "{:02X}", b)?;
54        }
55        Ok(())
56    }
57}
58
59impl Display for Sha256Hash {
60    // bs58
61    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> Result<(), std::fmt::Error> {
62        let base58 = bs58::encode(self.0).into_string();
63        write!(f, "{base58}")?;
64        Ok(())
65    }
66}
67
68impl AsRef<[u8]> for Sha256Hash {
69    fn as_ref(&self) -> &[u8] {
70        self.0.as_ref()
71    }
72}
73
74impl PartialEq for Sha256Hash {
75    fn eq(&self, other: &Self) -> bool {
76        self.0 == other.0
77    }
78}
79
80impl Hash for Sha256Hash {
81    type HashState = Sha256;
82
83    fn new() -> Self::HashState {
84        Sha256::new()
85    }
86
87    fn update(state: &mut Self::HashState, data: impl AsRef<[u8]>) {
88        state.update(data)
89    }
90
91    fn finalize(state: Self::HashState) -> Self {
92        Sha256Hash(state.finalize().into())
93    }
94}
95
96#[derive(Debug)]
97pub enum Sha256HashParseError {
98    Base58(bs58::decode::Error),
99    Hex(hex::FromHexError),
100}
101impl std::str::FromStr for Sha256Hash {
102    type Err = Sha256HashParseError;
103    fn from_str(s: &str) -> Result<Self, Sha256HashParseError> {
104        // "0xA01BE6E4A62BA7D0988FD6F1FE5DC964FD818628A96B472AA3472E6EFFB9A74F"
105        let mut store_id = [0; 32];
106        if s.len() == 66 {
107            // Parse as Hex (with 0x).
108            hex::decode_to_slice(&s[2..], &mut store_id).map_err(Sha256HashParseError::Hex)?;
109        } else if s.len() == 64 {
110            // Parse as Hex.
111            hex::decode_to_slice(s, &mut store_id).map_err(Sha256HashParseError::Hex)?;
112        } else {
113            bs58::decode(s)
114                .onto(&mut store_id)
115                .map_err(Sha256HashParseError::Base58)?;
116        }
117        Ok(Sha256Hash(store_id))
118    }
119}
120
121// TODO: Generalize the error and stream.
122pub trait Stream<T>:
123      futures::Stream<Item=Result<T,ProtocolError>> // Result<BytesMut,std::io::Error>>
124    + futures::Sink<T, Error=ProtocolError>
125    + Unpin
126    + Send // JP: This is needed for async_recursion. Not sure if this makes sense in practice.
127    + Sync // JP: This is needed for async_recursion. Not sure if this makes sense in practice.
128{}
129
130// TODO: Move this somewhere else, or remove this since we have MuxStream now?
131pub struct TypedStream<S, T> {
132    stream: S,
133    phantom: PhantomData<fn(T)>,
134}
135
136impl<S, T> TypedStream<S, T> {
137    pub fn new(stream: S) -> TypedStream<S, T> {
138        TypedStream {
139            stream,
140            phantom: PhantomData,
141        }
142    }
143
144    pub fn finalize(self) -> S {
145        self.stream
146    }
147}
148
149impl<S, T> futures::Stream for TypedStream<S, T>
150where
151    S: futures::Stream<Item = Result<BytesMut, std::io::Error>> + Unpin,
152    T: for<'a> Deserialize<'a>,
153{
154    type Item = Result<T, ProtocolError>;
155    fn poll_next(
156        mut self: Pin<&mut Self>,
157        ctx: &mut Context<'_>,
158    ) -> Poll<std::option::Option<<Self as futures::Stream>::Item>> {
159        let p = futures::Stream::poll_next(Pin::new(&mut self.stream), ctx);
160        p.map(|o| {
161            o.map(|t| match t {
162                Ok(bytes) => serde_cbor::from_slice(&bytes).map_err(|err| {
163                    // log::error!("Failed to parse type {}: {}", type_name::<T>(), err);
164                    ProtocolError::DeserializationError(err)
165                }),
166                Err(err) => Err(ProtocolError::StreamReceiveError(err)),
167            })
168        })
169    }
170}
171
172impl<S, T> futures::Sink<T> for TypedStream<S, T>
173where
174    S: futures::Sink<Bytes, Error = std::io::Error> + Unpin,
175    T: Serialize,
176{
177    type Error = ProtocolError;
178
179    fn poll_ready(
180        mut self: Pin<&mut Self>,
181        ctx: &mut Context<'_>,
182    ) -> Poll<Result<(), <Self as futures::Sink<T>>::Error>> {
183        let p = Pin::new(&mut self.stream).poll_ready(ctx);
184        p.map(|r| {
185            r.map_err(|e| {
186                // log::error!("Send error: {:?}", e);
187                ProtocolError::StreamSendError(e)
188            })
189        })
190    }
191
192    fn start_send(mut self: Pin<&mut Self>, x: T) -> Result<(), <Self as futures::Sink<T>>::Error> {
193        // TODO: to_writer instead?
194        // serde_cbor::to_writer(&stream, &response);
195        match serde_cbor::to_vec(&x) {
196            Err(err) => Err(ProtocolError::SerializationError(err)),
197            Ok(cbor) => {
198                let p = Pin::new(&mut self.stream).start_send(cbor.into());
199                p.map_err(|e| ProtocolError::StreamSendError(e))
200            }
201        }
202    }
203
204    fn poll_flush(
205        mut self: Pin<&mut Self>,
206        ctx: &mut Context<'_>,
207    ) -> Poll<Result<(), <Self as futures::Sink<T>>::Error>> {
208        let p = Pin::new(&mut self.stream).poll_flush(ctx);
209        p.map_err(|e| ProtocolError::StreamSendError(e))
210    }
211
212    fn poll_close(
213        mut self: Pin<&mut Self>,
214        ctx: &mut Context<'_>,
215    ) -> Poll<Result<(), <Self as futures::Sink<T>>::Error>> {
216        let p = Pin::new(&mut self.stream).poll_close(ctx);
217        p.map_err(|e| ProtocolError::StreamSendError(e))
218    }
219}
220
221impl<S, T> Stream<T> for TypedStream<S, T>
222where
223    S: futures::Stream<Item = Result<BytesMut, std::io::Error>> + Send,
224    S: futures::Sink<Bytes, Error = std::io::Error>,
225    S: Unpin,
226    S: Sync,
227    T: for<'a> Deserialize<'a> + Serialize,
228{
229}
230
231pub struct Channel<T> {
232    send: Sender<T>,
233    recv: Receiver<T>,
234}
235
236impl<T> Channel<T> {
237    pub fn new_pair(capacity: usize) -> (Channel<T>, Channel<T>) {
238        let (send1, recv1) = tokio::sync::mpsc::channel(capacity);
239        let (send2, recv2) = tokio::sync::mpsc::channel(capacity);
240        let c1 = Channel {
241            send: send1,
242            recv: recv2,
243        };
244        let c2 = Channel {
245            send: send2,
246            recv: recv1,
247        };
248        (c1, c2)
249    }
250}
251
252impl<T> Stream<T> for Channel<T>
253where
254    Channel<T>: Sync,
255    Channel<T>: Send,
256{
257}
258
259impl<T> futures::Stream for Channel<T> {
260    type Item = Result<T, ProtocolError>;
261
262    fn poll_next(
263        mut self: Pin<&mut Self>,
264        ctx: &mut Context<'_>,
265    ) -> Poll<Option<Result<T, ProtocolError>>> {
266        todo!()
267        // let p = futures::Stream::poll_next(Pin::new(&mut self.recv), ctx);
268        // p.map(|o| o.map(|t| Ok(t)))
269    }
270
271    fn size_hint(&self) -> (usize, Option<usize>) {
272        todo!()
273        // self.recv.size_hint()
274    }
275}
276
277impl<T> futures::Sink<T> for Channel<T> {
278    type Error = ProtocolError;
279
280    fn poll_ready(
281        mut self: Pin<&mut Self>,
282        ctx: &mut Context<'_>,
283    ) -> Poll<Result<(), <Self as futures::Sink<T>>::Error>> {
284        todo!()
285        // let p = Pin::new(&mut self.send).poll_ready(ctx);
286        // p.map(|r| {
287        //     r.map_err(|e| {
288        //         log::error!("Send error: {:?}", e);
289        //         ProtocolError::StreamSendError(std::io::Error::other("poll_ready error"))
290        //     })
291        // })
292    }
293
294    fn start_send(mut self: Pin<&mut Self>, x: T) -> Result<(), <Self as futures::Sink<T>>::Error> {
295        todo!()
296        // let p = Pin::new(&mut self.send).start_send(x);
297        // p.map_err(|e| {
298        //     log::error!("Send error: {:?}", e);
299        //     ProtocolError::StreamSendError(std::io::Error::other("start_send error"))
300        // })
301    }
302
303    fn poll_flush(
304        mut self: Pin<&mut Self>,
305        ctx: &mut Context<'_>,
306    ) -> Poll<Result<(), <Self as futures::Sink<T>>::Error>> {
307        todo!()
308        // let p = Pin::new(&mut self.send).poll_flush(ctx);
309        // p.map(|r| {
310        //     r.map_err(|e| {
311        //         log::error!("Send error: {:?}", e);
312        //         ProtocolError::StreamSendError(std::io::Error::other("poll_flush error"))
313        //     })
314        // })
315    }
316
317    fn poll_close(
318        mut self: Pin<&mut Self>,
319        ctx: &mut Context<'_>,
320    ) -> Poll<Result<(), <Self as futures::Sink<T>>::Error>> {
321        todo!()
322        // let p = Pin::new(&mut self.send).poll_close(ctx);
323        // p.map(|r| {
324        //     r.map_err(|e| {
325        //         log::error!("Send error: {:?}", e);
326        //         ProtocolError::StreamSendError(std::io::Error::other("poll_close error"))
327        //     })
328        // })
329    }
330}
331
332#[cfg(test)]
333use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender};
334#[cfg(test)]
335use tokio_stream::wrappers::UnboundedReceiverStream;
336
337#[cfg(test)]
338pub struct UnboundChannel<T> {
339    send: UnboundedSender<T>,
340    recv: UnboundedReceiverStream<T>, // UnboundedReceiver<T>,
341}
342
343#[cfg(test)]
344impl<T> UnboundChannel<T> {
345    pub fn new_pair() -> (UnboundChannel<T>, UnboundChannel<T>) {
346        let (send1, recv1) = tokio::sync::mpsc::unbounded_channel();
347        let (send2, recv2) = tokio::sync::mpsc::unbounded_channel();
348        let c1 = UnboundChannel {
349            send: send1,
350            recv: UnboundedReceiverStream::new(recv2),
351        };
352        let c2 = UnboundChannel {
353            send: send2,
354            recv: UnboundedReceiverStream::new(recv1),
355        };
356        (c1, c2)
357    }
358}
359
360#[cfg(test)]
361impl<T> Stream<T> for UnboundChannel<T>
362where
363    UnboundChannel<T>: Sync,
364    UnboundChannel<T>: Send,
365{
366}
367
368#[cfg(test)]
369impl<T> futures::Stream for UnboundChannel<T> {
370    type Item = Result<T, ProtocolError>;
371
372    fn poll_next(
373        mut self: Pin<&mut Self>,
374        ctx: &mut Context<'_>,
375    ) -> Poll<Option<Result<T, ProtocolError>>> {
376        // todo!()
377        let p = futures::Stream::poll_next(Pin::new(&mut self.recv), ctx);
378        p.map(|o| o.map(|t| Ok(t)))
379    }
380
381    fn size_hint(&self) -> (usize, Option<usize>) {
382        // todo!()
383        self.recv.size_hint()
384    }
385}
386
387#[cfg(test)]
388use tracing::error;
389
390#[cfg(test)]
391impl<T> futures::Sink<T> for UnboundChannel<T> {
392    type Error = ProtocolError;
393
394    fn poll_ready(
395        mut self: Pin<&mut Self>,
396        ctx: &mut Context<'_>,
397    ) -> Poll<Result<(), <Self as futures::Sink<T>>::Error>> {
398        Poll::Ready(Ok(()))
399        // todo!()
400        // let p = Pin::new(&mut self.send).poll_ready(ctx);
401        // p.map(|r| {
402        //     r.map_err(|e| {
403        //         log::error!("Send error: {:?}", e);
404        //         ProtocolError::StreamSendError(std::io::Error::other("poll_ready error"))
405        //     })
406        // })
407    }
408
409    fn start_send(mut self: Pin<&mut Self>, x: T) -> Result<(), <Self as futures::Sink<T>>::Error> {
410        let p = Pin::new(&mut self.send).send(x);
411        p.map_err(|e| {
412            error!("Send error: {:?}", e);
413            ProtocolError::StreamSendError(std::io::Error::other("start_send error"))
414        })
415    }
416
417    fn poll_flush(
418        mut self: Pin<&mut Self>,
419        ctx: &mut Context<'_>,
420    ) -> Poll<Result<(), <Self as futures::Sink<T>>::Error>> {
421        Poll::Ready(Ok(()))
422        // todo!()
423        // let p = Pin::new(&mut self.send).poll_flush(ctx);
424        // p.map(|r| {
425        //     r.map_err(|e| {
426        //         log::error!("Send error: {:?}", e);
427        //         ProtocolError::StreamSendError(std::io::Error::other("poll_flush error"))
428        //     })
429        // })
430    }
431
432    fn poll_close(
433        mut self: Pin<&mut Self>,
434        ctx: &mut Context<'_>,
435    ) -> Poll<Result<(), <Self as futures::Sink<T>>::Error>> {
436        Poll::Ready(Ok(()))
437        // todo!()
438        // let p = Pin::new(&mut self.send).poll_close(ctx);
439        // p.map(|r| {
440        //     r.map_err(|e| {
441        //         log::error!("Send error: {:?}", e);
442        //         ProtocolError::StreamSendError(std::io::Error::other("poll_close error"))
443        //     })
444        // })
445    }
446}
447
448/*
449pub(crate) fn merkle_root<H: Hash>(hashes: &[H]) -> H {
450    let mut h = H::new();
451    for hash in hashes.iter() {
452        H::update(&mut h, hash);
453    }
454    H::finalize(h)
455}
456
457pub(crate) fn validate_piece<H: Hash>(piece: &[u8], expected_hash: &H) -> bool {
458    let mut h = H::new();
459    H::update(&mut h, piece);
460    let h = H::finalize(h);
461
462    &h == expected_hash
463}
464*/
465
466pub(crate) struct CompressConsecutive<I, T> {
467    current: Option<Range<T>>,
468    inner: I,
469}
470
471/// Assumes inputs are sorted.
472pub(crate) fn compress_consecutive_into_ranges<I, T>(i: I) -> CompressConsecutive<I, T> {
473    CompressConsecutive {
474        current: None,
475        inner: i,
476    }
477}
478
479impl<I, T> Iterator for CompressConsecutive<I, T>
480where
481    I: Iterator<Item = T>,
482    T: Add<u64, Output = T> + PartialEq + Copy,
483{
484    type Item = Range<T>;
485
486    fn next(&mut self) -> Option<Self::Item> {
487        while let Some(x) = self.inner.next() {
488            match &self.current {
489                Some(r) => {
490                    // Consecutive.
491                    if r.end == x {
492                        self.current = Some(Range {
493                            start: r.start,
494                            end: x + 1u64,
495                        });
496                    } else {
497                        let new = Some(Range {
498                            start: x,
499                            end: x + 1u64,
500                        });
501                        let result = std::mem::replace(&mut self.current, new);
502                        return result;
503                    }
504                }
505                None => {
506                    self.current = Some(Range {
507                        start: x,
508                        end: x + 1u64,
509                    });
510                }
511            };
512        }
513
514        // We're done, so return what we have.
515        let result = std::mem::replace(&mut self.current, None);
516        result
517    }
518}
519
520/// Check if the input is a power of two (inclusive of 0).
521pub(crate) fn is_power_of_two(x: u64) -> bool {
522    0 == (x & (x.wrapping_sub(1)))
523}
524
525mod test {
526    use super::*;
527
528    #[test]
529    fn test_single_number() {
530        let numbers = vec![5];
531        let result: Vec<_> = compress_consecutive_into_ranges(numbers.into_iter()).collect();
532        assert_eq!(result, vec![Range { start: 5, end: 6 }]);
533    }
534
535    #[test]
536    fn test_consecutive_numbers() {
537        let numbers = vec![1, 2, 3, 4, 5];
538        let result: Vec<_> = compress_consecutive_into_ranges(numbers.into_iter()).collect();
539        assert_eq!(result, vec![Range { start: 1, end: 6 }]);
540    }
541
542    #[test]
543    fn test_split() {
544        let numbers = vec![1, 4];
545        let result: Vec<_> = compress_consecutive_into_ranges(numbers.into_iter()).collect();
546        assert_eq!(
547            result,
548            vec![Range { start: 1, end: 2 }, Range { start: 4, end: 5 },]
549        );
550    }
551
552    #[test]
553    fn test_mixed_numbers() {
554        let numbers = vec![1, 2, 3, 7, 8, 10, 11, 12, 15];
555        let result: Vec<_> = compress_consecutive_into_ranges(numbers.into_iter()).collect();
556        assert_eq!(
557            result,
558            vec![
559                Range { start: 1, end: 4 },
560                Range { start: 7, end: 9 },
561                Range { start: 10, end: 13 },
562                Range { start: 15, end: 16 },
563            ]
564        );
565    }
566}
ossa_core/util.rs

ossa_core/
util.rs
