enclavia_protocol/
noise.rs

1use snow::{Builder, HandshakeState, TransportState};
2
3#[cfg(feature = "async-transport")]
4use ciborium::{de::from_reader, ser::into_writer};
5#[cfg(feature = "async-transport")]
6use serde::{Deserialize, Serialize};
7#[cfg(feature = "async-transport")]
8use std::io::Cursor;
9#[cfg(feature = "async-transport")]
10use tracing::{debug, info, instrument, trace};
11
12pub const NOISE_PATTERN: &str = "Noise_NN_25519_ChaChaPoly_BLAKE2s";
13
14pub struct NoiseHandshake {
15    state: HandshakeState,
16}
17
18impl NoiseHandshake {
19    pub fn initiator() -> Result<Self, snow::Error> {
20        let builder: Builder<'_> = Builder::new(NOISE_PATTERN.parse()?);
21        let state = builder.build_initiator()?;
22        Ok(Self { state })
23    }
24
25    pub fn responder() -> Result<Self, snow::Error> {
26        let builder: Builder<'_> = Builder::new(NOISE_PATTERN.parse()?);
27        let state = builder.build_responder()?;
28        Ok(Self { state })
29    }
30
31    pub fn write_message(
32        &mut self,
33        payload: &[u8],
34        message: &mut [u8],
35    ) -> Result<usize, snow::Error> {
36        self.state.write_message(payload, message)
37    }
38
39    pub fn read_message(
40        &mut self,
41        message: &[u8],
42        payload: &mut [u8],
43    ) -> Result<usize, snow::Error> {
44        self.state.read_message(message, payload)
45    }
46
47    pub fn into_transport_mode(self) -> Result<NoiseTransport, snow::Error> {
48        let transport = self.state.into_transport_mode()?;
49        Ok(NoiseTransport { state: transport })
50    }
51
52    pub fn get_handshake_hash(&self) -> &[u8] {
53        self.state.get_handshake_hash()
54    }
55}
56
57pub struct NoiseTransport {
58    state: TransportState,
59}
60
61impl NoiseTransport {
62    pub fn write_message(
63        &mut self,
64        payload: &[u8],
65        message: &mut [u8],
66    ) -> Result<usize, snow::Error> {
67        self.state.write_message(payload, message)
68    }
69
70    pub fn read_message(
71        &mut self,
72        message: &[u8],
73        payload: &mut [u8],
74    ) -> Result<usize, snow::Error> {
75        self.state.read_message(message, payload)
76    }
77}
78
79// --- Async transport layer (requires tokio) ---
80
81#[cfg(feature = "async-transport")]
82#[instrument(skip(stream))]
83pub async fn perform_handshake_as_initiator(
84    stream: &mut (impl tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin),
85) -> Result<(NoiseTransport, Vec<u8>), Box<dyn std::error::Error>> {
86    info!("Starting Noise handshake as initiator");
87    let mut handshake = NoiseHandshake::initiator()?;
88    let mut buffer = vec![0u8; 65535];
89
90    debug!("Sending first handshake message");
91    let len = handshake.write_message(&[], &mut buffer)?;
92    tokio::io::AsyncWriteExt::write_all(stream, &buffer[..len]).await?;
93
94    debug!("Waiting for handshake response");
95    let len = tokio::io::AsyncReadExt::read(stream, &mut buffer).await?;
96    let mut payload = vec![0u8; 65535];
97    handshake.read_message(&buffer[..len], &mut payload)?;
98
99    let handshake_hash = handshake.get_handshake_hash().to_vec();
100
101    info!("Handshake completed successfully as initiator");
102    Ok((handshake.into_transport_mode()?, handshake_hash))
103}
104
105#[cfg(feature = "async-transport")]
106#[instrument(skip(stream))]
107pub async fn perform_handshake_as_responder(
108    stream: &mut (impl tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin),
109) -> Result<(NoiseTransport, Vec<u8>), Box<dyn std::error::Error>> {
110    info!("Starting Noise handshake as responder");
111    let mut handshake = NoiseHandshake::responder()?;
112    let mut buffer = vec![0u8; 65535];
113
114    debug!("Waiting for first handshake message");
115    let len = tokio::io::AsyncReadExt::read(stream, &mut buffer).await?;
116    let mut payload = vec![0u8; 65535];
117    handshake.read_message(&buffer[..len], &mut payload)?;
118
119    debug!("Sending handshake response");
120    let len = handshake.write_message(&[], &mut buffer)?;
121    tokio::io::AsyncWriteExt::write_all(stream, &buffer[..len]).await?;
122
123    let handshake_hash = handshake.get_handshake_hash().to_vec();
124
125    info!("Handshake completed successfully as responder");
126    Ok((handshake.into_transport_mode()?, handshake_hash))
127}
128
129/// Perform handshake as initiator and return a CborTransport ready for CBOR messaging.
130#[cfg(feature = "async-transport")]
131#[instrument(skip(stream))]
132pub async fn perform_cbor_handshake_as_initiator<S>(
133    mut stream: S,
134) -> Result<(CborTransport<S>, Vec<u8>), Box<dyn std::error::Error>>
135where
136    S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin,
137{
138    let transport = perform_handshake_as_initiator(&mut stream).await?;
139    Ok((CborTransport::new(transport.0, stream), transport.1))
140}
141
142/// Perform handshake as responder and return a CborTransport ready for CBOR messaging.
143#[cfg(feature = "async-transport")]
144#[instrument(skip(stream))]
145pub async fn perform_cbor_handshake_as_responder<S>(
146    mut stream: S,
147) -> Result<(CborTransport<S>, Vec<u8>), Box<dyn std::error::Error>>
148where
149    S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin,
150{
151    let transport = perform_handshake_as_responder(&mut stream).await?;
152    Ok((CborTransport::new(transport.0, stream), transport.1))
153}
154
155/// A wrapper around NoiseTransport that provides CBOR message sending/receiving
156/// with length-prefixed framing (4-byte big-endian length prefix + encrypted payload).
157#[cfg(feature = "async-transport")]
158pub struct CborTransport<S> {
159    transport: NoiseTransport,
160    stream: S,
161    read_buffer: Vec<u8>,
162    write_buffer: Vec<u8>,
163}
164
165#[cfg(feature = "async-transport")]
166impl<S> CborTransport<S>
167where
168    S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin,
169{
170    pub fn new(transport: NoiseTransport, stream: S) -> Self {
171        Self {
172            transport,
173            stream,
174            read_buffer: vec![0u8; 65535],
175            write_buffer: vec![0u8; 65535],
176        }
177    }
178
179    #[instrument(skip(self, message))]
180    pub async fn send<T: Serialize>(
181        &mut self,
182        message: &T,
183    ) -> Result<(), Box<dyn std::error::Error>> {
184        let mut cbor_bytes = Vec::new();
185        into_writer(message, &mut cbor_bytes)?;
186        trace!(cbor_len = cbor_bytes.len(), "Serialized CBOR message");
187
188        let encrypted_len = self
189            .transport
190            .write_message(&cbor_bytes, &mut self.write_buffer)?;
191        trace!(encrypted_len = encrypted_len, "Encrypted message");
192
193        let length_bytes = (encrypted_len as u32).to_be_bytes();
194        tokio::io::AsyncWriteExt::write_all(&mut self.stream, &length_bytes).await?;
195        tokio::io::AsyncWriteExt::write_all(&mut self.stream, &self.write_buffer[..encrypted_len])
196            .await?;
197        tokio::io::AsyncWriteExt::flush(&mut self.stream).await?;
198
199        trace!("CBOR message sent successfully");
200        Ok(())
201    }
202
203    #[instrument(skip(self))]
204    pub async fn receive<T: for<'de> Deserialize<'de>>(
205        &mut self,
206    ) -> Result<T, Box<dyn std::error::Error>> {
207        let mut length_bytes = [0u8; 4];
208        tokio::io::AsyncReadExt::read_exact(&mut self.stream, &mut length_bytes).await?;
209        let encrypted_len = u32::from_be_bytes(length_bytes) as usize;
210
211        if encrypted_len > self.read_buffer.len() {
212            return Err(format!(
213                "Message too large: {} bytes (max: {})",
214                encrypted_len,
215                self.read_buffer.len()
216            )
217            .into());
218        }
219
220        trace!(encrypted_len = encrypted_len, "Reading encrypted message");
221
222        tokio::io::AsyncReadExt::read_exact(
223            &mut self.stream,
224            &mut self.read_buffer[..encrypted_len],
225        )
226        .await?;
227
228        let mut payload = vec![0u8; 65535];
229        let payload_len = self
230            .transport
231            .read_message(&self.read_buffer[..encrypted_len], &mut payload)?;
232        trace!(payload_len = payload_len, "Decrypted message");
233
234        let mut cursor = Cursor::new(&payload[..payload_len]);
235        let message: T = from_reader(&mut cursor)?;
236        trace!("CBOR message received and deserialized successfully");
237        Ok(message)
238    }
239
240    pub fn transport(&self) -> &NoiseTransport {
241        &self.transport
242    }
243
244    pub fn transport_mut(&mut self) -> &mut NoiseTransport {
245        &mut self.transport
246    }
247}
enclavia_protocol/noise.rs

enclavia_protocol/
noise.rs
