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

cliquenet/
connection.rs
