openrtc 0.2.1

//! Encrypted explicit file-transfer helpers shared by native hosts.
//!
//! Iroh bi-stream wire format (after the plaintext `0x02` protocol byte):
//! - each logical write is one length-framed `ORTCE1` envelope (`RAW_STREAM_TYPE_ID = 3`)
//! - the first decrypted frame is `4-byte JSON length || header JSON`
//! - subsequent decrypted frames are file chunks
//!
//! WebRTC datachannel wire format keeps the legacy `0x02` frame layout but encrypts
//! the header/chunk payloads with the same `ORTCE1` envelope (no extra length frame).

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);
    }
}