diff --git a/perf/src/noprotection.rs b/perf/src/noprotection.rs index db7e2ea70..44caa3d37 100644 --- a/perf/src/noprotection.rs +++ b/perf/src/noprotection.rs @@ -16,6 +16,16 @@ impl NoProtectionSession { fn new(tls: Box) -> Self { Self { inner: tls } } + + /// Wraps the provided keys in `NoProtectionPacketKey` to disable packet encryption / decryption + fn wrap_packet_keys( + keys: crypto::KeyPair>, + ) -> crypto::KeyPair> { + crypto::KeyPair { + local: Box::new(NoProtectionPacketKey::new(keys.local)), + remote: Box::new(NoProtectionPacketKey::new(keys.remote)), + } + } } struct NoProtectionPacketKey { @@ -88,17 +98,17 @@ impl crypto::Session for NoProtectionSession { } fn write_handshake(&mut self, buf: &mut Vec) -> Option { - self.inner.write_handshake(buf) + let keys = self.inner.write_handshake(buf)?; + + Some(crypto::Keys { + header: keys.header, + packet: Self::wrap_packet_keys(keys.packet), + }) } fn next_1rtt_keys(&mut self) -> Option>> { let keys = self.inner.next_1rtt_keys()?; - - // use wrapper type to disable packet encryption/decryption - Some(crypto::KeyPair { - local: Box::new(NoProtectionPacketKey::new(keys.local)), - remote: Box::new(NoProtectionPacketKey::new(keys.remote)), - }) + Some(Self::wrap_packet_keys(keys)) } fn is_valid_retry(&self, orig_dst_cid: &ConnectionId, header: &[u8], payload: &[u8]) -> bool {