libp2p_noise/io/
framed.rs

1// Copyright 2020 Parity Technologies (UK) Ltd.
2//
3// Permission is hereby granted, free of charge, to any person obtaining a
4// copy of this software and associated documentation files (the "Software"),
5// to deal in the Software without restriction, including without limitation
6// the rights to use, copy, modify, merge, publish, distribute, sublicense,
7// and/or sell copies of the Software, and to permit persons to whom the
8// Software is furnished to do so, subject to the following conditions:
9//
10// The above copyright notice and this permission notice shall be included in
11// all copies or substantial portions of the Software.
12//
13// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
14// OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
15// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
16// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
17// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
18// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
19// DEALINGS IN THE SOFTWARE.
20
21//! Provides a [`Codec`] type implementing the [`Encoder`] and [`Decoder`] traits.
22//!
23//! Alongside a [`asynchronous_codec::Framed`] this provides a [Sink](futures::Sink)
24//! and [Stream](futures::Stream) for length-delimited Noise protocol messages.
25
26use std::{io, mem::size_of};
27
28use asynchronous_codec::{Decoder, Encoder};
29use bytes::{Buf, Bytes, BytesMut};
30use quick_protobuf::{BytesReader, MessageRead, MessageWrite, Writer};
31
32use super::handshake::proto;
33use crate::{protocol::PublicKey, Error};
34
35/// Max. size of a noise message.
36const MAX_NOISE_MSG_LEN: usize = 65535;
37/// Space given to the encryption buffer to hold key material.
38const EXTRA_ENCRYPT_SPACE: usize = 1024;
39/// Max. length for Noise protocol message payloads.
40pub(crate) const MAX_FRAME_LEN: usize = MAX_NOISE_MSG_LEN - EXTRA_ENCRYPT_SPACE;
41static_assertions::const_assert! {
42    MAX_FRAME_LEN + EXTRA_ENCRYPT_SPACE <= MAX_NOISE_MSG_LEN
43}
44
45/// Codec holds the noise session state `S` and acts as a medium for
46/// encoding and decoding length-delimited session messages.
47pub(crate) struct Codec<S> {
48    session: S,
49
50    // We reuse write and encryption buffers across multiple messages to avoid reallocations.
51    // We cannot reuse read and decryption buffers because we cannot return borrowed data.
52    write_buffer: BytesMut,
53    encrypt_buffer: BytesMut,
54}
55
56impl<S> Codec<S> {
57    pub(crate) fn new(session: S) -> Self {
58        Codec {
59            session,
60            write_buffer: BytesMut::default(),
61            encrypt_buffer: BytesMut::default(),
62        }
63    }
64}
65
66impl Codec<snow::HandshakeState> {
67    /// Checks if the session was started in the `initiator` role.
68    pub(crate) fn is_initiator(&self) -> bool {
69        self.session.is_initiator()
70    }
71
72    /// Checks if the session was started in the `responder` role.
73    pub(crate) fn is_responder(&self) -> bool {
74        !self.session.is_initiator()
75    }
76
77    /// Converts the underlying Noise session from the [`snow::HandshakeState`] to a
78    /// [`snow::TransportState`] once the handshake is complete, including the static
79    /// DH [`PublicKey`] of the remote if received.
80    ///
81    /// If the Noise protocol session state does not permit transitioning to
82    /// transport mode because the handshake is incomplete, an error is returned.
83    ///
84    /// An error is also returned if the remote's static DH key is not present or
85    /// cannot be parsed, as that indicates a fatal handshake error for the noise
86    /// `XX` pattern, which is the only handshake protocol libp2p currently supports.
87    pub(crate) fn into_transport(self) -> Result<(PublicKey, Codec<snow::TransportState>), Error> {
88        let dh_remote_pubkey = self.session.get_remote_static().ok_or_else(|| {
89            Error::Io(io::Error::new(
90                io::ErrorKind::Other,
91                "expect key to always be present at end of XX session",
92            ))
93        })?;
94
95        let dh_remote_pubkey = PublicKey::from_slice(dh_remote_pubkey)?;
96        let codec = Codec::new(self.session.into_transport_mode()?);
97
98        Ok((dh_remote_pubkey, codec))
99    }
100}
101
102impl Encoder for Codec<snow::HandshakeState> {
103    type Error = io::Error;
104    type Item<'a> = &'a proto::NoiseHandshakePayload;
105
106    fn encode(&mut self, item: Self::Item<'_>, dst: &mut BytesMut) -> Result<(), Self::Error> {
107        let item_size = item.get_size();
108
109        self.write_buffer.resize(item_size, 0);
110        let mut writer = Writer::new(&mut self.write_buffer[..item_size]);
111        item.write_message(&mut writer)
112            .expect("Protobuf encoding to succeed");
113
114        encrypt(
115            &self.write_buffer[..item_size],
116            dst,
117            &mut self.encrypt_buffer,
118            |item, buffer| self.session.write_message(item, buffer),
119        )?;
120
121        Ok(())
122    }
123}
124
125impl Decoder for Codec<snow::HandshakeState> {
126    type Error = io::Error;
127    type Item = proto::NoiseHandshakePayload;
128
129    fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
130        let cleartext = match decrypt(src, |ciphertext, decrypt_buffer| {
131            self.session.read_message(ciphertext, decrypt_buffer)
132        })? {
133            None => return Ok(None),
134            Some(cleartext) => cleartext,
135        };
136
137        let mut reader = BytesReader::from_bytes(&cleartext[..]);
138        let pb =
139            proto::NoiseHandshakePayload::from_reader(&mut reader, &cleartext).map_err(|_| {
140                io::Error::new(
141                    io::ErrorKind::InvalidData,
142                    "Failed decoding handshake payload",
143                )
144            })?;
145
146        Ok(Some(pb))
147    }
148}
149
150impl Encoder for Codec<snow::TransportState> {
151    type Error = io::Error;
152    type Item<'a> = &'a [u8];
153
154    fn encode(&mut self, item: Self::Item<'_>, dst: &mut BytesMut) -> Result<(), Self::Error> {
155        encrypt(item, dst, &mut self.encrypt_buffer, |item, buffer| {
156            self.session.write_message(item, buffer)
157        })
158    }
159}
160
161impl Decoder for Codec<snow::TransportState> {
162    type Error = io::Error;
163    type Item = Bytes;
164
165    fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Self::Item>, Self::Error> {
166        decrypt(src, |ciphertext, decrypt_buffer| {
167            self.session.read_message(ciphertext, decrypt_buffer)
168        })
169    }
170}
171
172/// Encrypts the given cleartext to `dst`.
173///
174/// This is a standalone function to allow us reusing the `encrypt_buffer` and to use to across
175/// different session states of the noise protocol.
176fn encrypt(
177    cleartext: &[u8],
178    dst: &mut BytesMut,
179    encrypt_buffer: &mut BytesMut,
180    encrypt_fn: impl FnOnce(&[u8], &mut [u8]) -> Result<usize, snow::Error>,
181) -> io::Result<()> {
182    tracing::trace!("Encrypting {} bytes", cleartext.len());
183
184    encrypt_buffer.resize(cleartext.len() + EXTRA_ENCRYPT_SPACE, 0);
185    let n = encrypt_fn(cleartext, encrypt_buffer).map_err(into_io_error)?;
186
187    tracing::trace!("Outgoing ciphertext has {n} bytes");
188
189    encode_length_prefixed(&encrypt_buffer[..n], dst);
190
191    Ok(())
192}
193
194/// Encrypts the given ciphertext.
195///
196/// This is a standalone function so we can use it across different session states of the noise
197/// protocol. In case `ciphertext` does not contain enough bytes to decrypt the entire frame,
198/// `Ok(None)` is returned.
199fn decrypt(
200    ciphertext: &mut BytesMut,
201    decrypt_fn: impl FnOnce(&[u8], &mut [u8]) -> Result<usize, snow::Error>,
202) -> io::Result<Option<Bytes>> {
203    let Some(ciphertext) = decode_length_prefixed(ciphertext) else {
204        return Ok(None);
205    };
206
207    tracing::trace!("Incoming ciphertext has {} bytes", ciphertext.len());
208
209    let mut decrypt_buffer = BytesMut::zeroed(ciphertext.len());
210    let n = decrypt_fn(&ciphertext, &mut decrypt_buffer).map_err(into_io_error)?;
211
212    tracing::trace!("Decrypted cleartext has {n} bytes");
213
214    Ok(Some(decrypt_buffer.split_to(n).freeze()))
215}
216
217fn into_io_error(err: snow::Error) -> io::Error {
218    io::Error::new(io::ErrorKind::InvalidData, err)
219}
220
221const U16_LENGTH: usize = size_of::<u16>();
222
223fn encode_length_prefixed(src: &[u8], dst: &mut BytesMut) {
224    dst.reserve(U16_LENGTH + src.len());
225    dst.extend_from_slice(&(src.len() as u16).to_be_bytes());
226    dst.extend_from_slice(src);
227}
228
229fn decode_length_prefixed(src: &mut BytesMut) -> Option<Bytes> {
230    if src.len() < size_of::<u16>() {
231        return None;
232    }
233
234    let mut len_bytes = [0u8; U16_LENGTH];
235    len_bytes.copy_from_slice(&src[..U16_LENGTH]);
236    let len = u16::from_be_bytes(len_bytes) as usize;
237
238    if src.len() - U16_LENGTH >= len {
239        // Skip the length header we already read.
240        src.advance(U16_LENGTH);
241        Some(src.split_to(len).freeze())
242    } else {
243        None
244    }
245}