From 4fb77f8edad121a31e8c85d0e2279d9aea6aaaf3 Mon Sep 17 00:00:00 2001 From: Brian Schwind Date: Wed, 3 Apr 2024 13:45:52 +0900 Subject: [PATCH] Report wireguard endpoint as a candidate when an endpoint override is in place (#305) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Use our Endpoints type alias * Add the recent wireguard endpoint to NAT candidates if a peer has an endpoint override * Simplify logic in the inject_endpoints() function Co-authored-by: Matěj Laitl * Specify mock wireguard endpoints for developer 1 and 2 in the test data * Add a test for verifying the wireguard endpoint is returned in the list of NAT candidates * Remove FromStr usage * Appease clippy --------- Co-authored-by: Matěj Laitl --- hostsfile/src/lib.rs | 1 + server/src/api/mod.rs | 17 ++++++-- server/src/api/user.rs | 89 ++++++++++++++++++++++++++++++++++++++++++ server/src/main.rs | 2 +- server/src/test.rs | 37 ++++++++++++------ 5 files changed, 129 insertions(+), 17 deletions(-) diff --git a/hostsfile/src/lib.rs b/hostsfile/src/lib.rs index a509dea..514cc5c 100644 --- a/hostsfile/src/lib.rs +++ b/hostsfile/src/lib.rs @@ -198,6 +198,7 @@ impl HostsBuilder { let hosts_file = OpenOptions::new() .create(true) + .truncate(false) .read(true) .write(true) .open(hosts_path)?; diff --git a/server/src/api/mod.rs b/server/src/api/mod.rs index f7db98f..da0a12c 100644 --- a/server/src/api/mod.rs +++ b/server/src/api/mod.rs @@ -6,12 +6,21 @@ pub mod admin; pub mod user; /// Inject the collected endpoints from the WG interface into a list of peers. -/// This is essentially what adds NAT holepunching functionality. +/// This is essentially what adds NAT holepunching functionality. If a peer +/// already has an endpoint specified (by calling the override-endpoint) API, +/// the relatively recent wireguard endpoint will be added to the list of NAT +/// candidates, so other peers have a better chance of connecting. pub fn inject_endpoints(session: &Session, peers: &mut Vec) { for peer in peers { - if peer.contents.endpoint.is_none() { - if let Some(endpoint) = session.context.endpoints.read().get(&peer.public_key) { - peer.contents.endpoint = Some(endpoint.to_owned().into()); + let endpoints = session.context.endpoints.read(); + if let Some(wg_endpoint) = endpoints.get(&peer.public_key) { + if peer.contents.endpoint.is_none() { + peer.contents.endpoint = Some(wg_endpoint.to_owned().into()); + } else { + // The peer already has an endpoint specified, but it might be stale. + // If there is an endpoint reported from wireguard, we should add it + // to the list of candidates so others can try to connect using it. + peer.contents.candidates.push(wg_endpoint.to_owned().into()); } } } diff --git a/server/src/api/user.rs b/server/src/api/user.rs index 42f73fa..3385482 100644 --- a/server/src/api/user.rs +++ b/server/src/api/user.rs @@ -464,4 +464,93 @@ mod tests { assert_eq!(peer.candidates, candidates); Ok(()) } + + #[tokio::test] + async fn test_endpoint_in_candidates() -> Result<(), Error> { + // We want to verify that the current wireguard endpoint always shows up + // either in the peer.endpoint field, or the peer.candidates field (in the + // case that the peer has specified an endpoint override). + let server = test::Server::new()?; + + let peer = DatabasePeer::get(&server.db().lock(), test::DEVELOPER1_PEER_ID)?; + assert_eq!(peer.candidates, vec![]); + + // Specify one NAT candidate. At this point, we have an unspecified + // endpoint and one NAT candidate specified. + let candidates = vec!["1.1.1.1:51820".parse::().unwrap()]; + assert_eq!( + server + .form_request( + test::DEVELOPER1_PEER_IP, + "PUT", + "/v1/user/candidates", + &candidates + ) + .await + .status(), + StatusCode::NO_CONTENT + ); + + let res = server + .request(test::DEVELOPER1_PEER_IP, "GET", "/v1/user/state") + .await; + + assert_eq!(res.status(), StatusCode::OK); + + let whole_body = hyper::body::aggregate(res).await?; + let State { peers, .. } = serde_json::from_reader(whole_body.reader())?; + + let developer_1 = peers + .into_iter() + .find(|p| p.id == test::DEVELOPER1_PEER_ID) + .unwrap(); + assert_eq!( + developer_1.endpoint, + Some(test::DEVELOPER1_PEER_ENDPOINT.parse().unwrap()) + ); + assert_eq!(developer_1.candidates, candidates); + + // Now, explicitly set an endpoint with the override-endpoint API + // and check that the original wireguard endpoint still shows up + // in the list of NAT candidates. + assert_eq!( + server + .form_request( + test::DEVELOPER1_PEER_IP, + "PUT", + "/v1/user/endpoint", + &EndpointContents::Set("1.2.3.4:51820".parse().unwrap()) + ) + .await + .status(), + StatusCode::NO_CONTENT + ); + + let res = server + .request(test::DEVELOPER1_PEER_IP, "GET", "/v1/user/state") + .await; + + assert_eq!(res.status(), StatusCode::OK); + + let whole_body = hyper::body::aggregate(res).await?; + let State { peers, .. } = serde_json::from_reader(whole_body.reader())?; + + let developer_1 = peers + .into_iter() + .find(|p| p.id == test::DEVELOPER1_PEER_ID) + .unwrap(); + + // The peer endpoint should be the one we just specified in the override-endpoint request. + assert_eq!(developer_1.endpoint, Some("1.2.3.4:51820".parse().unwrap())); + + // The list of candidates should now contain the one we specified at the beginning of the + // test, and the wireguard-reported endpoint of the peer. + let nat_candidate_1 = "1.1.1.1:51820".parse().unwrap(); + let nat_candidate_2 = test::DEVELOPER1_PEER_ENDPOINT.parse().unwrap(); + assert_eq!(developer_1.candidates.len(), 2); + assert!(developer_1.candidates.contains(&nat_candidate_1)); + assert!(developer_1.candidates.contains(&nat_candidate_2)); + + Ok(()) + } } diff --git a/server/src/main.rs b/server/src/main.rs index aee0fd1..d34949c 100644 --- a/server/src/main.rs +++ b/server/src/main.rs @@ -147,7 +147,7 @@ pub type Endpoints = Arc>>; #[derive(Clone)] pub struct Context { pub db: Db, - pub endpoints: Arc>>, + pub endpoints: Endpoints, pub interface: InterfaceName, pub backend: Backend, pub public_key: Key, diff --git a/server/src/test.rs b/server/src/test.rs index 912cb80..e221035 100644 --- a/server/src/test.rs +++ b/server/src/test.rs @@ -10,7 +10,7 @@ use parking_lot::{Mutex, RwLock}; use rusqlite::Connection; use serde::Serialize; use shared::{Cidr, CidrContents, Error, PeerContents}; -use std::{collections::HashMap, net::SocketAddr, path::PathBuf, sync::Arc}; +use std::{net::SocketAddr, path::PathBuf, sync::Arc}; use tempfile::TempDir; use wireguard_control::{Backend, InterfaceName, Key, KeyPair}; @@ -27,7 +27,9 @@ mod v4 { pub const ADMIN_PEER_IP: &str = "10.80.1.1"; pub const WG_MANAGE_PEER_IP: &str = ADMIN_PEER_IP; pub const DEVELOPER1_PEER_IP: &str = "10.80.64.2"; + pub const DEVELOPER1_PEER_ENDPOINT: &str = "169.10.26.8:14720"; pub const DEVELOPER2_PEER_IP: &str = "10.80.64.3"; + pub const DEVELOPER2_PEER_ENDPOINT: &str = "169.55.140.9:5833"; pub const USER1_PEER_IP: &str = "10.80.128.2"; pub const USER2_PEER_IP: &str = "10.80.129.2"; pub const EXPERIMENT_SUBCIDR_PEER_IP: &str = "10.81.0.1"; @@ -48,7 +50,9 @@ mod v6 { pub const ADMIN_PEER_IP: &str = "fd00:1337::1:0:0:1"; pub const WG_MANAGE_PEER_IP: &str = ADMIN_PEER_IP; pub const DEVELOPER1_PEER_IP: &str = "fd00:1337::2:0:0:1"; + pub const DEVELOPER1_PEER_ENDPOINT: &str = "[1001:db8::1]:14720"; pub const DEVELOPER2_PEER_IP: &str = "fd00:1337::2:0:0:2"; + pub const DEVELOPER2_PEER_ENDPOINT: &str = "[2001:db8::1]:5833"; pub const USER1_PEER_IP: &str = "fd00:1337::3:0:0:1"; pub const USER2_PEER_IP: &str = "fd00:1337::3:0:0:2"; pub const EXPERIMENT_SUBCIDR_PEER_IP: &str = "fd00:1337::4:0:0:1"; @@ -114,21 +118,19 @@ impl Server { DEVELOPER_CIDR_ID, create_cidr(&db, "developer", DEVELOPER_CIDR)?.id ); + + let developer_1 = developer_peer_contents("developer1", DEVELOPER1_PEER_IP)?; + let developer_1_public_key = developer_1.public_key.clone(); assert_eq!( DEVELOPER1_PEER_ID, - DatabasePeer::create( - &db, - developer_peer_contents("developer1", DEVELOPER1_PEER_IP)? - )? - .id + DatabasePeer::create(&db, developer_1,)?.id ); + + let developer_2 = developer_peer_contents("developer2", DEVELOPER2_PEER_IP)?; + let developer_2_public_key = developer_2.public_key.clone(); assert_eq!( DEVELOPER2_PEER_ID, - DatabasePeer::create( - &db, - developer_peer_contents("developer2", DEVELOPER2_PEER_IP)? - )? - .id + DatabasePeer::create(&db, developer_2)?.id ); assert_eq!(USER_CIDR_ID, create_cidr(&db, "user", USER_CIDR)?.id); assert_eq!( @@ -141,7 +143,18 @@ impl Server { ); let db = Arc::new(Mutex::new(db)); - let endpoints = Arc::new(RwLock::new(HashMap::new())); + + let endpoints = [ + ( + developer_1_public_key, + DEVELOPER1_PEER_ENDPOINT.parse().unwrap(), + ), + ( + developer_2_public_key, + DEVELOPER2_PEER_ENDPOINT.parse().unwrap(), + ), + ]; + let endpoints = Arc::new(RwLock::new(endpoints.into())); Ok(Self { conf,