cliquenet/
connection.rs

1use std::{
2    cmp::min,
3    io,
4    iter::{once, repeat},
5    net::SocketAddr,
6    sync::Arc,
7    time::Duration,
8};
9
10use rand::RngExt;
11use snow::{Builder, HandshakeState, TransportState, params::NoiseParams};
12use socket2::{SockRef, TcpKeepalive};
13use tokio::{
14    io::{AsyncReadExt, AsyncWriteExt},
15    net::TcpStream,
16    time::sleep,
17    try_join,
18};
19use tracing::{debug, warn};
20
21use crate::{
22    Config, NetAddr, Version,
23    error::NetworkError,
24    msg::{Header, MAX_NOISE_MESSAGE_SIZE, hello::Hello},
25    util::until,
26    x25519::PublicKey,
27};
28
29const MAX_NOISE_HANDSHAKE_SIZE: usize = 1024;
30
31type Result<T> = std::result::Result<T, NetworkError>;
32
33pub struct Connection {
34    pub key: PublicKey,
35    pub addr: SocketAddr,
36    pub stream: TcpStream,
37    pub state: TransportState,
38}
39
40type Prologue = Vec<u8>;
41
42impl Connection {
43    pub async fn accept(conf: Arc<Config>, mut stream: TcpStream) -> Result<Self> {
44        let node = conf.keypair.public_key();
45        let addr = stream.peer_addr()?;
46
47        configure_socket(&conf, &node, &addr, &stream);
48
49        until(conf.handshake_timeout, async move {
50            let (version, prologue) =
51                select_version(&node, &addr, &conf, &mut stream, false).await?;
52
53            debug!(name = %conf.name, %node, %addr, %version, "negotiated version");
54
55            let noise_proto = conf
56                .noise_protocols
57                .get(&version)
58                .expect("selected version has noise config");
59
60            let hs = Builder::new(noise_proto.noise_params())
61                .local_private_key(&conf.keypair.secret_key().as_bytes())
62                .expect("valid private key")
63                .prologue(&prologue)
64                .expect("1st time we set the prologue")
65                .build_responder()
66                .expect("valid noise params yield valid handshake state");
67
68            let state = on_handshake(&mut stream, hs).await?;
69
70            if let Some(key) = remote_static_key(&state) {
71                Ok(Self {
72                    key,
73                    addr,
74                    stream,
75                    state,
76                })
77            } else {
78                warn!(name = %conf.name, %node, %addr, "invalid static key");
79                Err(NetworkError::InvalidHandshakeMessage)
80            }
81        })
82        .await
83    }
84
85    pub async fn connect(conf: Arc<Config>, peer: PublicKey, addr: NetAddr) -> Self {
86        let mut delays = once({
87            if conf.random_connect_delay {
88                Duration::from_millis(rand::rng().random_range(0..1000))
89            } else {
90                Duration::ZERO
91            }
92        })
93        .chain(
94            conf.connect_retry_delays
95                .iter()
96                .map(|&d| Duration::from_secs(d.into())),
97        )
98        .chain(repeat({
99            let d = *conf.connect_retry_delays.last();
100            Duration::from_secs(d.into())
101        }));
102
103        let node = conf.keypair.public_key();
104
105        let mut backoff = None;
106
107        loop {
108            if let Some(d) = backoff.take() {
109                sleep(d).await;
110            } else {
111                sleep(delays.next().expect("delays iterator is infinite")).await;
112            }
113
114            debug!(name = %conf.name, %node, %peer, %addr, "connecting");
115
116            match try_connect(&conf, &peer, &addr).await {
117                Ok(mut conn) => {
118                    let hello_exchange = until(conf.handshake_timeout, async {
119                        conn.send_hello(Hello::Ok).await?;
120                        conn.recv_hello().await
121                    });
122                    match hello_exchange.await {
123                        Ok(h) if h.is_ok() => break conn,
124                        Ok(h) => {
125                            warn!(
126                                name = %conf.name,
127                                %node,
128                                %peer,
129                                remote = %conn.key,
130                                %addr,
131                                "hello response was not ok"
132                            );
133                            backoff = h.backoff_duration();
134                        },
135                        Err(err) => {
136                            warn!(
137                                name = %conf.name,
138                                %node,
139                                %peer,
140                                remote = %conn.key,
141                                %addr,
142                                %err,
143                                "failed to exchange hello"
144                            );
145                        },
146                    }
147                },
148                Err(err) => {
149                    warn!(name = %conf.name, %node, %peer, %addr, %err, "connect/handshake error")
150                },
151            }
152        }
153    }
154
155    /// Send a `Hello` frame.
156    pub async fn send_hello(&mut self, h: Hello) -> Result<()> {
157        let mut b = [0u8; 64];
158        let n = self
159            .state
160            .write_message(h.to_bytes().as_ref(), &mut b[Header::SIZE..])?;
161        let h = Header::data(n as u16);
162        send_frame(&mut self.stream, h, &mut b[..Header::SIZE + n]).await?;
163        Ok(())
164    }
165
166    /// Read a `Hello` frame.
167    pub async fn recv_hello(&mut self) -> Result<Hello> {
168        let mut a = [0u8; 64];
169        let h = recv_frame(&mut self.stream, &mut a).await?;
170        let mut b = [0u8; 64];
171        let n = self.state.read_message(&a[..h.len().into()], &mut b)?;
172        let h = Hello::from_bytes(&b[..n]).ok_or(NetworkError::InvalidHello)?;
173        Ok(h)
174    }
175}
176
177async fn try_connect(conf: &Config, peer: &PublicKey, addr: &NetAddr) -> Result<Connection> {
178    let new_handshake_state = |prologue: &Prologue, params: NoiseParams| {
179        Builder::new(params)
180            .local_private_key(conf.keypair.secret_key().as_slice())
181            .expect("valid private key")
182            .remote_public_key(peer.as_slice())
183            .expect("valid remote pub key")
184            .prologue(prologue)
185            .expect("1st time we set the prologue")
186            .build_initiator()
187            .expect("valid noise params yield valid handshake state")
188    };
189
190    let mut stream = match addr {
191        NetAddr::Inet(i, p) => until(conf.connect_timeout, TcpStream::connect((*i, *p))).await?,
192        NetAddr::Name(h, p) => until(conf.connect_timeout, TcpStream::connect((&**h, *p))).await?,
193    };
194
195    let node = conf.keypair.public_key();
196    let addr = stream.peer_addr()?;
197
198    debug!(name = %conf.name, %node, %peer, %addr, "tcp connection established");
199
200    configure_socket(conf, &node, &addr, &stream);
201
202    until(conf.handshake_timeout, async move {
203        let (version, prologue) = select_version(&node, &addr, conf, &mut stream, true).await?;
204
205        debug!(name = %conf.name, %node, %peer, %addr, %version, "negotiated version");
206
207        let noise_proto = conf
208            .noise_protocols
209            .get(&version)
210            .expect("selected version has noise config");
211
212        let hshake = new_handshake_state(&prologue, noise_proto.noise_params());
213        let state = handshake(&mut stream, hshake).await?;
214        match remote_static_key(&state) {
215            Some(key) if key == *peer => Ok(Connection {
216                key,
217                addr,
218                stream,
219                state,
220            }),
221            Some(key) => {
222                warn!(name = %conf.name, %node, %peer, remote = %key, %addr, "static key mismatch");
223                Err(NetworkError::InvalidHandshakeMessage)
224            },
225            None => {
226                warn!(name = %conf.name, %node, %peer, %addr, "invalid static key");
227                Err(NetworkError::InvalidHandshakeMessage)
228            },
229        }
230    })
231    .await
232}
233
234fn configure_socket(conf: &Config, node: &PublicKey, addr: &SocketAddr, stream: &TcpStream) {
235    if let Err(err) = stream.set_nodelay(true) {
236        warn!(name = %conf.name, %node, %addr, %err, "failed to enable no_delay option")
237    }
238
239    let k = TcpKeepalive::new()
240        .with_time(conf.keep_alive_after)
241        .with_interval(conf.keep_alive_interval)
242        .with_retries(conf.keep_alive_retries.into());
243
244    if let Err(err) = SockRef::from(stream).set_tcp_keepalive(&k) {
245        warn!(name = %conf.name, %node, %addr, %err, "failed to enable tcp keepalive");
246    }
247}
248
249fn remote_static_key(state: &TransportState) -> Option<PublicKey> {
250    let k = state.get_remote_static()?;
251    PublicKey::try_from(k).ok()
252}
253
254/// Select a version from the range that both sides support.
255///
256/// This will be the minimum of the max. supported ones from both sides.
257async fn select_version(
258    node: &PublicKey,
259    addr: &SocketAddr,
260    conf: &Config,
261    stream: &mut TcpStream,
262    is_initiator: bool,
263) -> Result<(Version, Prologue)> {
264    const INIT_PAYLOAD_LEN: usize = 4;
265
266    let our_min = *conf.noise_protocols.first().0;
267    let our_max = *conf.noise_protocols.last().0;
268
269    let mut send_buf = [0u8; INIT_PAYLOAD_LEN];
270    let mut recv_buf = [0u8; INIT_PAYLOAD_LEN];
271
272    send_buf[0..2].copy_from_slice(&u16::from(our_min).to_be_bytes());
273    send_buf[2..4].copy_from_slice(&u16::from(our_max).to_be_bytes());
274
275    let (mut r, mut w) = stream.split();
276    try_join!(w.write_all(&send_buf), r.read_exact(&mut recv_buf))?;
277
278    let their_min = Version::from(u16::from_be_bytes([recv_buf[0], recv_buf[1]]));
279    let their_max = Version::from(u16::from_be_bytes([recv_buf[2], recv_buf[3]]));
280
281    let selected = min(our_max, their_max);
282
283    if selected < their_min || selected < our_min {
284        warn!(
285            name = %conf.name,
286            %node,
287            %addr,
288            local  = ?(u16::from(our_min), u16::from(our_max)),
289            remote = ?(u16::from(their_min), u16::from(their_max)),
290            "incompatible versions"
291        );
292        return Err(NetworkError::IncompatibleVersions);
293    }
294
295    // Construct the prologue so that both sides end up with the same value.
296    // We include the sent and received version ranges to ensure no one has
297    // tampered with those values as they were sent in plain text.
298    let mut prologue = Vec::new();
299    prologue.extend_from_slice(conf.name.as_bytes());
300    if is_initiator {
301        prologue.extend_from_slice(&send_buf);
302        prologue.extend_from_slice(&recv_buf);
303    } else {
304        prologue.extend_from_slice(&recv_buf);
305        prologue.extend_from_slice(&send_buf);
306    }
307
308    Ok((selected, prologue))
309}
310
311/// Perform a noise handshake as initiator with the remote party.
312async fn handshake(stream: &mut TcpStream, mut hs: HandshakeState) -> Result<TransportState> {
313    let mut a = [0u8; MAX_NOISE_HANDSHAKE_SIZE];
314    let n = hs.write_message(&[], &mut a[Header::SIZE..])?;
315    let h = Header::data(n as u16);
316    send_frame(stream, h, &mut a[..Header::SIZE + n]).await?;
317    let mut b = [0u8; MAX_NOISE_HANDSHAKE_SIZE];
318    let h = recv_frame(stream, &mut b).await?;
319    if !h.is_data() || h.is_partial() {
320        return Err(NetworkError::InvalidHandshakeMessage);
321    }
322    hs.read_message(&b[..h.len().into()], &mut a)?;
323    Ok(hs.into_transport_mode()?)
324}
325
326/// Perform a noise handshake as responder with a remote party.
327async fn on_handshake(stream: &mut TcpStream, mut hs: HandshakeState) -> Result<TransportState> {
328    let mut a = [0u8; MAX_NOISE_HANDSHAKE_SIZE];
329    let h = recv_frame(stream, &mut a).await?;
330    if !h.is_data() || h.is_partial() {
331        return Err(NetworkError::InvalidHandshakeMessage);
332    }
333    let mut b = [0u8; MAX_NOISE_HANDSHAKE_SIZE];
334    hs.read_message(&a[..h.len().into()], &mut b)?;
335    let n = hs.write_message(&[], &mut b[Header::SIZE..])?;
336    let h = Header::data(n as u16);
337    send_frame(stream, h, &mut b[..Header::SIZE + n]).await?;
338    Ok(hs.into_transport_mode()?)
339}
340
341/// Read a single frame (header + payload) from the remote.
342async fn recv_frame<R, const N: usize>(stream: &mut R, buf: &mut [u8; N]) -> io::Result<Header>
343where
344    R: AsyncReadExt + Unpin,
345{
346    let h = {
347        let n = stream.read_u32().await?;
348        Header::unvalidated(n)
349    };
350    let n = h.len().into();
351    if n > N {
352        return Err(io::ErrorKind::InvalidInput.into());
353    }
354    stream.read_exact(&mut buf[..n]).await?;
355    Ok(h)
356}
357
358/// Write a single frame (header + payload) to the remote.
359///
360/// The header is serialised into the first 4 bytes of `msg`. It is the
361/// caller's responsibility to ensure there is room at the beginning.
362async fn send_frame<W>(stream: &mut W, hdr: Header, msg: &mut [u8]) -> io::Result<()>
363where
364    W: AsyncWriteExt + Unpin,
365{
366    debug_assert!(msg.len() <= MAX_NOISE_MESSAGE_SIZE);
367    msg[..Header::SIZE].copy_from_slice(&hdr.to_bytes());
368    stream.write_all(msg).await?;
369    Ok(())
370}
371
372#[cfg(test)]
373mod tests {
374    use std::net::Ipv4Addr;
375
376    use tokio::net::{TcpListener, TcpStream};
377
378    use super::{Prologue, Result, select_version};
379    use crate::{Config, NetAddr, NetworkError, Version, noise::Protocol, x25519::Keypair};
380
381    fn config<I, V>(versions: I) -> Config
382    where
383        I: IntoIterator<Item = V>,
384        V: Into<Version>,
385    {
386        Config::builder()
387            .name("test")
388            .keypair(Keypair::generate().unwrap())
389            .bind(NetAddr::from((Ipv4Addr::LOCALHOST, 0u16)))
390            .parties([])
391            .noise_protocols(
392                versions
393                    .into_iter()
394                    .map(|v| (v.into(), Protocol::IK_25519_AesGcm_Blake2s)),
395            )
396            .build()
397    }
398
399    async fn negotiate(
400        a: &Config,
401        b: &Config,
402    ) -> (Result<(Version, Prologue)>, Result<(Version, Prologue)>) {
403        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
404        let port = listener.local_addr().unwrap().port();
405        tokio::join!(
406            async {
407                let mut s = TcpStream::connect(("127.0.0.1", port)).await.unwrap();
408                select_version(&a.public_key(), &s.peer_addr().unwrap(), a, &mut s, true).await
409            },
410            async {
411                let (mut s, _) = listener.accept().await.unwrap();
412                select_version(&b.public_key(), &s.peer_addr().unwrap(), b, &mut s, false).await
413            },
414        )
415    }
416
417    #[tokio::test]
418    async fn picks_min_of_maxes() {
419        let (ra, rb) = negotiate(&config([1, 2, 3]), &config([1, 2])).await;
420        let (va, pa) = ra.unwrap();
421        let (vb, pb) = rb.unwrap();
422        assert_eq!(va, 2.into());
423        assert_eq!(vb, 2.into());
424        assert_eq!(pa, pb);
425    }
426
427    #[tokio::test]
428    async fn higher_overlap_takes_higher() {
429        let (ra, rb) = negotiate(&config([2, 3, 4]), &config([1, 2, 3])).await;
430        let (va, pa) = ra.unwrap();
431        let (vb, pb) = rb.unwrap();
432        assert_eq!(va, 3.into());
433        assert_eq!(vb, 3.into());
434        assert_eq!(pa, pb);
435    }
436
437    #[tokio::test]
438    async fn single_version_match() {
439        let (ra, rb) = negotiate(&config([1]), &config([1])).await;
440        let (va, pa) = ra.unwrap();
441        let (vb, pb) = rb.unwrap();
442        assert_eq!(va, 1.into());
443        assert_eq!(vb, 1.into());
444        assert_eq!(pa, pb);
445    }
446
447    #[tokio::test]
448    async fn disjoint_ranges_fail() {
449        let (ra, rb) = negotiate(&config([1]), &config([2])).await;
450        assert!(matches!(ra, Err(NetworkError::IncompatibleVersions)));
451        assert!(matches!(rb, Err(NetworkError::IncompatibleVersions)));
452    }
453}
cliquenet/connection.rs

cliquenet/
connection.rs
