use std::io;
use iroh::endpoint::{ClosedStream, ReadError, RecvStream, SendStream, WriteError};
use crate::application_crypto::{
protect_raw_stream_frame, ApplicationCryptoFrameDecoder, ApplicationCryptoStreamError,
APPLICATION_KEY_BYTES,
};
#[derive(Debug)]
pub struct ApplicationCryptoSendStream {
inner: SendStream,
key: [u8; APPLICATION_KEY_BYTES],
}
impl ApplicationCryptoSendStream {
pub fn new(inner: SendStream, key: [u8; APPLICATION_KEY_BYTES]) -> Self {
Self { inner, key }
}
pub async fn write_all(&mut self, buf: &[u8]) -> io::Result<()> {
let frame = protect_raw_stream_frame(&self.key, buf).map_err(crypto_io_error)?;
self.inner.write_all(&frame).await.map_err(map_write_error)
}
pub fn finish(mut self) -> io::Result<()> {
self.inner.finish().map_err(map_finish_error)
}
pub fn into_inner(self) -> SendStream {
self.inner
}
}
#[derive(Debug)]
pub struct ApplicationCryptoRecvStream {
inner: RecvStream,
decoder: ApplicationCryptoFrameDecoder,
read_buf: Vec<u8>,
finished: bool,
}
impl ApplicationCryptoRecvStream {
pub fn new(inner: RecvStream, key: [u8; APPLICATION_KEY_BYTES]) -> io::Result<Self> {
Ok(Self {
inner,
decoder: ApplicationCryptoFrameDecoder::new(&key).map_err(crypto_io_error)?,
read_buf: Vec::new(),
finished: false,
})
}
pub fn new_with_prefix(
inner: RecvStream,
key: [u8; APPLICATION_KEY_BYTES],
prefix: &[u8],
) -> io::Result<Self> {
let mut decoder = ApplicationCryptoFrameDecoder::new(&key).map_err(crypto_io_error)?;
let mut read_buf = Vec::new();
if !prefix.is_empty() {
for frame in decoder.push(prefix).map_err(crypto_stream_io_error)? {
read_buf.extend_from_slice(&frame);
}
}
Ok(Self {
inner,
decoder,
read_buf,
finished: false,
})
}
pub async fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
if !self.read_buf.is_empty() {
let copied = self.read_buf.len().min(buf.len());
buf[..copied].copy_from_slice(&self.read_buf[..copied]);
self.read_buf.drain(..copied);
return Ok(copied);
}
if self.finished {
return Ok(0);
}
loop {
let mut chunk = vec![0u8; 16 * 1024];
let read_bytes = match self.inner.read(&mut chunk).await {
Ok(Some(0)) => 0,
Ok(Some(read_bytes)) => read_bytes,
Ok(None) => 0,
Err(error) => return Err(map_read_error(error)),
};
if read_bytes == 0 {
self.finished = true;
self.decoder.finish().map_err(crypto_stream_io_error)?;
if self.read_buf.is_empty() {
return Ok(0);
}
let copied = self.read_buf.len().min(buf.len());
buf[..copied].copy_from_slice(&self.read_buf[..copied]);
self.read_buf.drain(..copied);
return Ok(copied);
}
chunk.truncate(read_bytes);
let opened = self.decoder.push(&chunk).map_err(crypto_stream_io_error)?;
for frame in opened {
self.read_buf.extend_from_slice(&frame);
}
if !self.read_buf.is_empty() {
let copied = self.read_buf.len().min(buf.len());
buf[..copied].copy_from_slice(&self.read_buf[..copied]);
self.read_buf.drain(..copied);
return Ok(copied);
}
}
}
pub fn into_inner(self) -> RecvStream {
self.inner
}
}
#[derive(Debug)]
pub enum PeerSendStream {
Plain(SendStream),
Encrypted(ApplicationCryptoSendStream),
}
impl PeerSendStream {
pub fn plain(inner: SendStream) -> Self {
Self::Plain(inner)
}
pub fn encrypted(inner: SendStream, key: [u8; APPLICATION_KEY_BYTES]) -> Self {
Self::Encrypted(ApplicationCryptoSendStream::new(inner, key))
}
pub fn is_encrypted(&self) -> bool {
matches!(self, Self::Encrypted(_))
}
pub async fn write_all(&mut self, buf: &[u8]) -> io::Result<()> {
match self {
Self::Plain(stream) => stream.write_all(buf).await.map_err(map_write_error),
Self::Encrypted(stream) => stream.write_all(buf).await,
}
}
pub fn finish(self) -> io::Result<()> {
match self {
Self::Plain(mut stream) => stream.finish().map_err(map_finish_error),
Self::Encrypted(stream) => stream.finish(),
}
}
pub fn into_plain(self) -> SendStream {
match self {
Self::Plain(stream) => stream,
Self::Encrypted(stream) => stream.into_inner(),
}
}
}
#[derive(Debug)]
pub enum PeerRecvStream {
Plain(RecvStream),
Encrypted(ApplicationCryptoRecvStream),
}
impl PeerRecvStream {
pub fn plain(inner: RecvStream) -> Self {
Self::Plain(inner)
}
pub fn encrypted(inner: RecvStream, key: [u8; APPLICATION_KEY_BYTES]) -> io::Result<Self> {
Ok(Self::Encrypted(ApplicationCryptoRecvStream::new(
inner, key,
)?))
}
pub fn encrypted_with_prefix(
inner: RecvStream,
key: [u8; APPLICATION_KEY_BYTES],
prefix: &[u8],
) -> io::Result<Self> {
Ok(Self::Encrypted(
ApplicationCryptoRecvStream::new_with_prefix(inner, key, prefix)?,
))
}
pub fn is_encrypted(&self) -> bool {
matches!(self, Self::Encrypted(_))
}
pub async fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
match self {
Self::Plain(stream) => read_recv_stream(stream, buf).await,
Self::Encrypted(stream) => stream.read(buf).await,
}
}
pub fn into_plain(self) -> RecvStream {
match self {
Self::Plain(stream) => stream,
Self::Encrypted(stream) => stream.into_inner(),
}
}
}
pub fn wrap_peer_streams(
key: Option<[u8; APPLICATION_KEY_BYTES]>,
send: SendStream,
recv: RecvStream,
) -> io::Result<(PeerSendStream, PeerRecvStream)> {
match key {
Some(key) => Ok((
PeerSendStream::encrypted(send, key),
PeerRecvStream::encrypted(recv, key)?,
)),
None => Ok((PeerSendStream::plain(send), PeerRecvStream::plain(recv))),
}
}
pub fn wrap_peer_send_stream(
key: Option<[u8; APPLICATION_KEY_BYTES]>,
send: SendStream,
) -> PeerSendStream {
match key {
Some(key) => PeerSendStream::encrypted(send, key),
None => PeerSendStream::plain(send),
}
}
async fn read_recv_stream(stream: &mut RecvStream, buf: &mut [u8]) -> io::Result<usize> {
match stream.read(buf).await {
Ok(Some(0)) => Ok(0),
Ok(Some(read_bytes)) => Ok(read_bytes),
Ok(None) => Ok(0),
Err(error) => Err(map_read_error(error)),
}
}
fn map_write_error(error: WriteError) -> io::Error {
io::Error::new(io::ErrorKind::Other, error.to_string())
}
fn map_read_error(error: ReadError) -> io::Error {
io::Error::new(io::ErrorKind::Other, error.to_string())
}
fn map_finish_error(error: ClosedStream) -> io::Error {
io::Error::new(io::ErrorKind::NotConnected, error.to_string())
}
fn crypto_io_error(error: crate::application_crypto::ApplicationCryptoError) -> io::Error {
io::Error::new(io::ErrorKind::InvalidData, format!("{error:?}"))
}
fn crypto_stream_io_error(error: ApplicationCryptoStreamError) -> io::Error {
io::Error::new(io::ErrorKind::InvalidData, format!("{error:?}"))
}