use iroh::endpoint::{RecvStream, SendStream};
use crate::application_crypto::{
is_application_encrypted_payload, open_application_payload, protect_application_payload,
ApplicationCryptoError, APPLICATION_KEY_BYTES, RAW_STREAM_TYPE_ID,
};
use crate::application_crypto_streams::{ApplicationCryptoRecvStream, ApplicationCryptoSendStream};
pub const EXPLICIT_FILE_PROTOCOL_BYTE: u8 = 0x02;
pub fn protect_webrtc_explicit_payload(
key: &[u8; APPLICATION_KEY_BYTES],
payload: &[u8],
) -> Result<Vec<u8>, ApplicationCryptoError> {
protect_application_payload(key, RAW_STREAM_TYPE_ID, payload)
}
pub fn open_webrtc_explicit_payload(
key: &[u8; APPLICATION_KEY_BYTES],
payload: &[u8],
) -> Result<Vec<u8>, ApplicationCryptoError> {
open_application_payload(key, RAW_STREAM_TYPE_ID, payload, true)
}
pub fn open_webrtc_explicit_payload_optional(
key: Option<&[u8; APPLICATION_KEY_BYTES]>,
payload: &[u8],
) -> Result<Vec<u8>, String> {
match key {
Some(key) => open_webrtc_explicit_payload(key, payload).map_err(|error| format!("{error:?}")),
None if is_application_encrypted_payload(payload) => Err(
"encrypted explicit transfer payload rejected because no application crypto key is installed"
.to_string(),
),
None => Ok(payload.to_vec()),
}
}
pub enum ExplicitTransferSender {
Plain(SendStream),
Encrypted(ApplicationCryptoSendStream),
}
impl ExplicitTransferSender {
pub async fn write_protocol_byte(&mut self) -> Result<(), String> {
self.write_all(&[EXPLICIT_FILE_PROTOCOL_BYTE]).await
}
pub async fn write_all(&mut self, buf: &[u8]) -> Result<(), String> {
match self {
Self::Plain(stream) => stream
.write_all(buf)
.await
.map_err(|error| error.to_string()),
Self::Encrypted(stream) => stream
.write_all(buf)
.await
.map_err(|error| error.to_string()),
}
}
pub fn finish(self) -> Result<(), String> {
match self {
Self::Plain(mut stream) => stream.finish().map_err(|error| error.to_string()),
Self::Encrypted(stream) => stream.finish().map_err(|error| error.to_string()),
}
}
}
pub fn explicit_transfer_sender(
stream: SendStream,
key: Option<[u8; APPLICATION_KEY_BYTES]>,
) -> ExplicitTransferSender {
match key {
Some(key) => {
ExplicitTransferSender::Encrypted(ApplicationCryptoSendStream::new(stream, key))
}
None => ExplicitTransferSender::Plain(stream),
}
}
pub struct ExplicitTransferReceiver {
inner: ApplicationCryptoRecvStream,
}
impl ExplicitTransferReceiver {
pub fn new(stream: RecvStream, key: [u8; APPLICATION_KEY_BYTES]) -> Result<Self, String> {
ApplicationCryptoRecvStream::new(stream, key)
.map(|inner| Self { inner })
.map_err(|error| error.to_string())
}
pub async fn read(&mut self, buf: &mut [u8]) -> Result<usize, String> {
self.inner
.read(buf)
.await
.map_err(|error| error.to_string())
}
pub async fn read_exact(&mut self, buf: &mut [u8]) -> Result<(), String> {
let mut offset = 0;
while offset < buf.len() {
let read_bytes = self.read(&mut buf[offset..]).await?;
if read_bytes == 0 {
return Err("unexpected end of encrypted explicit transfer stream".to_string());
}
offset += read_bytes;
}
Ok(())
}
}
pub async fn read_explicit_transfer_header_json(
receiver: &mut ExplicitTransferReceiver,
) -> Result<Vec<u8>, String> {
let mut len_buf = [0u8; 4];
receiver.read_exact(&mut len_buf).await?;
let len = u32::from_be_bytes(len_buf) as usize;
if len == 0 {
return Err("invalid explicit transfer header length: 0 bytes".to_string());
}
let mut header_buf = vec![0u8; len];
receiver.read_exact(&mut header_buf).await?;
Ok(header_buf)
}
pub fn plaintext_header_looks_encrypted(len_buf: [u8; 4], available: &[u8]) -> bool {
let frame_len = u32::from_be_bytes(len_buf) as usize;
if frame_len < 32 || frame_len > crate::application_crypto::DEFAULT_MAX_RAW_FRAME_BYTES {
return false;
}
if available.len() < 4 + 6 {
return false;
}
is_application_encrypted_payload(&available[4..])
}
#[cfg(test)]
mod tests {
use super::*;
use crate::application_crypto::protect_raw_stream_frame;
#[test]
fn webrtc_round_trip_matches_application_payload_type() {
let key = [7u8; APPLICATION_KEY_BYTES];
let payload = br#"{"transfer_id":"tx-1","filename":"a.txt","size":3,"mime_type":"application/octet-stream"}"#;
let protected = protect_webrtc_explicit_payload(&key, payload).expect("protect");
let opened = open_webrtc_explicit_payload(&key, &protected).expect("open");
assert_eq!(opened, payload);
}
#[test]
fn iroh_header_frame_round_trip_uses_raw_stream_type() {
let key = [9u8; APPLICATION_KEY_BYTES];
let header_json = br#"{"transfer_id":"tx-2","filename":"b.bin","size":10,"mime_type":"application/octet-stream"}"#;
let mut header_frame = Vec::with_capacity(4 + header_json.len());
header_frame.extend_from_slice(&(header_json.len() as u32).to_be_bytes());
header_frame.extend_from_slice(header_json);
let wire = protect_raw_stream_frame(&key, &header_frame).expect("frame");
let opened = open_application_payload(&key, RAW_STREAM_TYPE_ID, &wire[4..], true)
.expect("open frame body");
assert_eq!(opened, header_frame);
}
}