Skip to main content

calimero_network_primitives/
stream.rs

1use core::pin::Pin;
2use core::task::{Context, Poll};
3
4use futures_util::{Sink as FuturesSink, SinkExt, Stream as FuturesStream, StreamExt};
5use libp2p::{Stream as P2pStream, StreamProtocol};
6use tokio::io::BufStream;
7use tokio_util::codec::Framed;
8use tokio_util::compat::{Compat, FuturesAsyncReadCompatExt};
9
10mod codec;
11
12use codec::MessageCodec;
13pub use codec::{CodecError, Message};
14
15pub const MAX_MESSAGE_SIZE: usize = 8 * 1_024 * 1_024;
16
17pub const CALIMERO_STREAM_PROTOCOL: StreamProtocol = StreamProtocol::new("/calimero/stream/0.0.2");
18pub const CALIMERO_BLOB_PROTOCOL: StreamProtocol = StreamProtocol::new("/calimero/blob/0.0.2");
19
20#[derive(Debug)]
21pub struct Stream {
22    inner: Framed<BufStream<Compat<P2pStream>>, MessageCodec>,
23}
24
25impl Stream {
26    #[must_use]
27    pub fn new(stream: P2pStream) -> Self {
28        let stream = BufStream::new(stream.compat());
29        let stream = Framed::new(stream, MessageCodec::new(MAX_MESSAGE_SIZE));
30        Self { inner: stream }
31    }
32}
33
34impl FuturesStream for Stream {
35    type Item = Result<Message<'static>, CodecError>;
36
37    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
38        self.inner.poll_next_unpin(cx)
39    }
40}
41
42impl<'a> FuturesSink<Message<'a>> for Stream {
43    type Error = CodecError;
44
45    fn poll_ready(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
46        self.inner.poll_ready_unpin(cx)
47    }
48
49    fn start_send(mut self: Pin<&mut Self>, item: Message<'a>) -> Result<(), Self::Error> {
50        self.inner.start_send_unpin(item)
51    }
52
53    fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
54        self.inner.poll_flush_unpin(cx)
55    }
56
57    fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
58        self.inner.poll_close_unpin(cx)
59    }
60}