Skip to content

Commit 4712e32

Browse files
0xMimir0xMimir
committed
Added error handling to noise
1 parent 29cdba5 commit 4712e32

File tree

3 files changed

+99
-76
lines changed

3 files changed

+99
-76
lines changed

‎p2p/src/network/noise/p2p_network_noise_effects.rs‎

Lines changed: 8 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -52,11 +52,14 @@ impl P2pNetworkNoiseAction {
5252
} else {
5353
None
5454
};
55-
let handshake_error = if let Some(P2pNetworkNoiseStateInner::Error(error)) = &state.inner {
56-
Some(error)
57-
} else {
58-
None
59-
};
55+
56+
if let Some(P2pNetworkNoiseStateInner::Error(error)) = &state.inner {
57+
store.dispatch(P2pNetworkSchedulerAction::Error {
58+
addr: *self.addr(),
59+
error: error.clone().into(),
60+
});
61+
return
62+
}
6063

6164
let middle_initiator =
6265
matches!(&state.inner, Some(P2pNetworkNoiseStateInner::Initiator(..)))
@@ -93,14 +96,6 @@ impl P2pNetworkNoiseAction {
9396
}
9497
}
9598
P2pNetworkNoiseAction::IncomingChunk { addr, .. } => {
96-
if let Some(error) = handshake_error {
97-
store.dispatch(P2pNetworkSchedulerAction::Error {
98-
addr,
99-
error: error.clone().into(),
100-
});
101-
return;
102-
}
103-
10499
if let Some((peer_id, true)) = handshake_done {
105100
let addr = *self.addr();
106101
store.dispatch(P2pConnectionIncomingAction::FinalizePendingLibp2p {

‎p2p/src/network/noise/p2p_network_noise_reducer.rs‎

Lines changed: 69 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -68,13 +68,16 @@ impl P2pNetworkNoiseState {
6868
let mut offset = 0;
6969
loop {
7070
let buf = &self.buffer[offset..];
71-
if buf.len() >= 2 {
72-
let len = u16::from_be_bytes(buf[..2].try_into().expect("cannot fail"));
71+
let len = buf
72+
.get(..2)
73+
.map(|buf| Some(u16::from_be_bytes(buf.try_into().ok()?)))
74+
.flatten();
75+
76+
if let Some(len) = len {
7377
let full_len = 2 + len as usize;
7478
if buf.len() >= full_len {
7579
self.incoming_chunks.push_back(buf[..full_len].to_vec());
7680
offset += full_len;
77-
7881
continue;
7982
}
8083
}
@@ -188,55 +191,65 @@ impl P2pNetworkNoiseState {
188191
} => {
189192
let aead = ChaCha20Poly1305::new(&send_key.0.into());
190193
let chunk_max_size = u16::MAX as usize - 19;
191-
let chunks = data
192-
.chunks(chunk_max_size)
193-
.map(|data| {
194-
let mut chunk = Vec::with_capacity(18 + data.len());
195-
chunk.extend_from_slice(&((data.len() + 16) as u16).to_be_bytes());
196-
chunk.extend_from_slice(data);
194+
let mut chunks = vec![];
195+
196+
for data in data.chunks(chunk_max_size) {
197+
let mut chunk = Vec::with_capacity(18 + data.len());
198+
chunk.extend_from_slice(&((data.len() + 16) as u16).to_be_bytes());
199+
chunk.extend_from_slice(data);
200+
201+
let mut nonce = GenericArray::default();
202+
nonce[4..].clone_from_slice(&send_nonce.to_le_bytes());
203+
*send_nonce += 1;
204+
205+
let tag = aead.encrypt_in_place_detached(
206+
&nonce,
207+
&[],
208+
&mut chunk[2..(2 + data.len())],
209+
);
197210

198-
let mut nonce = GenericArray::default();
199-
nonce[4..].clone_from_slice(&send_nonce.to_le_bytes());
200-
*send_nonce += 1;
211+
let tag = match tag {
212+
Ok(tag) => tag,
213+
Err(_) => {
214+
*state =
215+
P2pNetworkNoiseStateInner::Error(NoiseError::Encryption);
216+
return;
217+
}
218+
};
201219

202-
let tag = aead
203-
.encrypt_in_place_detached(
204-
&nonce,
205-
&[],
206-
&mut chunk[2..(2 + data.len())],
207-
)
208-
.expect("cannot fail");
209-
chunk.extend_from_slice(&tag);
210-
chunk.into()
211-
})
212-
.collect();
220+
chunk.extend_from_slice(&tag);
221+
chunks.push(chunk.into());
222+
}
213223
self.outgoing_chunks.push_back(chunks);
214224
}
215225
P2pNetworkNoiseStateInner::Initiator(i) => {
216-
if let (Some((chunk, (send_key, recv_key))), Some(remote_pk)) =
217-
(i.generate(data), i.remote_pk.clone())
218-
{
219-
self.outgoing_chunks.push_back(vec![chunk.into()]);
220-
let remote_peer_id = remote_pk.peer_id();
226+
match (i.generate(data), i.remote_pk.clone()) {
227+
(Ok(Some((chunk, (send_key, recv_key)))), Some(remote_pk)) => {
228+
self.outgoing_chunks.push_back(vec![chunk.into()]);
229+
let remote_peer_id = remote_pk.peer_id();
221230

222-
if self
223-
.expected_peer_id
224-
.is_some_and(|expected_per_id| expected_per_id != remote_peer_id)
225-
{
226-
*state = P2pNetworkNoiseStateInner::Error(dbg!(
227-
NoiseError::RemotePeerIdMismatch
228-
));
229-
} else {
230-
*state = P2pNetworkNoiseStateInner::Done {
231-
incoming: false,
232-
send_key,
233-
recv_key,
234-
recv_nonce: 0,
235-
send_nonce: 0,
236-
remote_pk,
237-
remote_peer_id,
238-
};
231+
if self.expected_peer_id.is_some_and(|expected_per_id| {
232+
expected_per_id != remote_peer_id
233+
}) {
234+
*state = P2pNetworkNoiseStateInner::Error(dbg!(
235+
NoiseError::RemotePeerIdMismatch
236+
));
237+
} else {
238+
*state = P2pNetworkNoiseStateInner::Done {
239+
incoming: false,
240+
send_key,
241+
recv_key,
242+
recv_nonce: 0,
243+
send_nonce: 0,
244+
remote_pk,
245+
remote_peer_id,
246+
};
247+
}
239248
}
249+
(Err(error), Some(_)) => {
250+
*state = P2pNetworkNoiseStateInner::Error(error);
251+
}
252+
_ => (),
240253
}
241254
}
242255
P2pNetworkNoiseStateInner::Responder(r) => {
@@ -268,9 +281,17 @@ impl P2pNetworkNoiseState {
268281
nonce[4..].clone_from_slice(&send_nonce.to_le_bytes());
269282
*send_nonce += 1;
270283

271-
let tag = aead
272-
.encrypt_in_place_detached(&nonce, &[], &mut chunk[2..(2 + data.len())])
273-
.expect("cannot fail");
284+
let tag = match aead.encrypt_in_place_detached(
285+
&nonce,
286+
&[],
287+
&mut chunk[2..(2 + data.len())],
288+
) {
289+
Ok(tag) => tag,
290+
Err(_) => {
291+
self.inner = Some(P2pNetworkNoiseStateInner::Error(NoiseError::Encryption));
292+
return;
293+
}
294+
};
274295
chunk.extend_from_slice(&tag);
275296

276297
self.outgoing_chunks.push_back(vec![chunk.into()]);

‎p2p/src/network/noise/p2p_network_noise_state.rs‎

Lines changed: 22 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -149,21 +149,22 @@ impl NoiseState {
149149
.map(|()| self.hash.0 = hash.into())
150150
}
151151

152-
pub fn encrypt<const NONCE: u64>(&mut self, data: &mut [u8]) -> [u8; 16] {
152+
pub fn encrypt<const NONCE: u64>(&mut self, data: &mut [u8]) -> Result<[u8; 16], NoiseError> {
153153
let mut nonce = GenericArray::default();
154154
nonce[4..].clone_from_slice(&NONCE.to_le_bytes());
155155

156156
let tag = ChaCha20Poly1305::new(GenericArray::from_slice(&self.aead_key.0))
157157
.encrypt_in_place_detached(&nonce, &self.hash.0, data)
158-
.expect("data length must be sufficiently small");
158+
.map_err(|_| NoiseError::Encryption)?;
159+
159160
let hash = Sha256::default()
160161
.chain(self.hash.0)
161162
.chain(&*data)
162163
.chain(tag)
163164
.finalize_fixed();
164165
self.hash.0 = hash.into();
165166

166-
tag.into()
167+
Ok(tag.into())
167168
}
168169

169170
pub fn finish(&self) -> (DataSized<32>, DataSized<32>) {
@@ -196,6 +197,8 @@ pub enum NoiseError {
196197
SelfConnection,
197198
#[error("remote peer id doesn't match expected peer id")]
198199
RemotePeerIdMismatch,
200+
#[error("failed to encrypt data")]
201+
Encryption,
199202
}
200203

201204
pub struct ResponderOutput {
@@ -205,7 +208,7 @@ pub struct ResponderOutput {
205208
}
206209

207210
impl P2pNetworkNoiseStateInitiator {
208-
pub fn generate(&mut self, data: &[u8]) -> Option<(Vec<u8>, (DataSized<32>, DataSized<32>))> {
211+
pub fn generate(&mut self, data: &[u8]) -> Result<Option<(Vec<u8>, (DataSized<32>, DataSized<32>))>, NoiseError> {
209212
let Self {
210213
i_spk,
211214
i_ssk,
@@ -215,10 +218,13 @@ impl P2pNetworkNoiseStateInitiator {
215218
..
216219
} = self;
217220

218-
let r_epk = r_epk.as_ref()?;
221+
let r_epk = match r_epk.as_ref(){
222+
Some(r_epk) => r_epk,
223+
None => return Ok(None)
224+
};
219225

220226
let mut i_spk_bytes = i_spk.0.to_bytes();
221-
let tag = noise.encrypt::<1>(&mut i_spk_bytes);
227+
let tag = noise.encrypt::<1>(&mut i_spk_bytes)?;
222228
noise.mix_secret(&*i_ssk * r_epk);
223229
let mut payload = payload.0.to_vec();
224230
// if handshake is optimized by early mux negotiation
@@ -227,7 +233,7 @@ impl P2pNetworkNoiseStateInitiator {
227233
payload.push(data.len() as u8);
228234
payload.extend_from_slice(data);
229235
}
230-
let payload_tag = noise.encrypt::<0>(&mut payload);
236+
let payload_tag = noise.encrypt::<0>(&mut payload)?;
231237

232238
let mut chunk = vec![0; 2];
233239
chunk.extend_from_slice(&i_spk_bytes);
@@ -237,7 +243,7 @@ impl P2pNetworkNoiseStateInitiator {
237243
let l = (chunk.len() - 2) as u16;
238244
chunk[..2].clone_from_slice(&l.to_be_bytes());
239245

240-
Some((chunk, noise.finish()))
246+
Ok(Some((chunk, noise.finish())))
241247
}
242248

243249
pub fn consume<'a>(
@@ -258,9 +264,9 @@ impl P2pNetworkNoiseStateInitiator {
258264
if len < 200 {
259265
return Err(ChunkTooShort);
260266
}
261-
let r_epk = Pk::from_bytes(msg[..32].try_into().expect("cannot fail"));
262-
let mut r_spk_bytes =
263-
<[u8; 32]>::try_from(&msg[32..64]).expect("cannot fail, checked above");
267+
let r_epk = Pk::from_bytes(msg[..32].try_into().map_err(|_| ChunkTooShort)?);
268+
let mut r_spk_bytes = <[u8; 32]>::try_from(&msg[32..64]).map_err(|_| ChunkTooShort)?;
269+
264270
let tag = &msg[64..80];
265271

266272
noise.mix_hash(r_epk.0.as_bytes());
@@ -324,7 +330,7 @@ impl P2pNetworkNoiseStateResponder {
324330
payload.push(data.len() as u8);
325331
payload.extend_from_slice(data);
326332
}
327-
let payload_tag = noise.encrypt::<0>(&mut payload);
333+
let payload_tag = noise.encrypt::<0>(&mut payload).ok()?;
328334

329335
buffer.extend_from_slice(&payload);
330336
buffer.extend_from_slice(&payload_tag);
@@ -360,7 +366,7 @@ impl P2pNetworkNoiseStateResponder {
360366
if len < 32 {
361367
return Err(ChunkTooShort);
362368
}
363-
let i_epk = Pk::from_bytes(msg[..32].try_into().expect("cannot fail"));
369+
let i_epk = Pk::from_bytes(msg[..32].try_into().map_err(|_| ChunkTooShort)?);
364370

365371
let r_epk = r_esk.pk();
366372

@@ -370,7 +376,7 @@ impl P2pNetworkNoiseStateResponder {
370376
noise.mix_hash(b"");
371377
noise.mix_hash(r_epk.0.as_bytes());
372378
noise.mix_secret(&*r_esk * &i_epk);
373-
let tag = noise.encrypt::<0>(&mut r_spk_bytes);
379+
let tag = noise.encrypt::<0>(&mut r_spk_bytes)?;
374380
noise.mix_secret(&*r_ssk * &i_epk);
375381
r_ssk.zeroize();
376382

@@ -389,7 +395,8 @@ impl P2pNetworkNoiseStateResponder {
389395
}
390396

391397
// TODO: refactor obscure arithmetics
392-
let mut i_spk_bytes = <[u8; 32]>::try_from(&msg[..32]).expect("cannot fail");
398+
let mut i_spk_bytes =
399+
<[u8; 32]>::try_from(&msg[..32]).map_err(|_| ChunkTooShort)?;
393400
let (tag, msg) = msg[32..].split_at_mut(16);
394401
let len = msg.len();
395402
let (remote_payload, payload_tag) = msg.split_at_mut(len - 16);

0 commit comments

Comments
 (0)