diff --git a/rust/connlib/clients/android/src/lib.rs b/rust/connlib/clients/android/src/lib.rs index 94cdefd086..4855add95e 100644 --- a/rust/connlib/clients/android/src/lib.rs +++ b/rust/connlib/clients/android/src/lib.rs @@ -4,8 +4,8 @@ // ecosystem, so it's used here for consistency. use connlib_client_shared::{ - file_logger, keypair, Callbacks, Cidrv4, Cidrv6, Error, LoginUrl, LoginUrlError, - ResourceDescription, Session, Sockets, + callbacks::ResourceDescription, file_logger, keypair, Callbacks, Cidrv4, Cidrv6, Error, + LoginUrl, LoginUrlError, Session, Sockets, }; use jni::{ objects::{GlobalRef, JClass, JObject, JString, JValue}, diff --git a/rust/connlib/clients/apple/src/lib.rs b/rust/connlib/clients/apple/src/lib.rs index 8df91c862e..afa93c0c90 100644 --- a/rust/connlib/clients/apple/src/lib.rs +++ b/rust/connlib/clients/apple/src/lib.rs @@ -2,8 +2,8 @@ #![allow(clippy::unnecessary_cast, improper_ctypes, non_camel_case_types)] use connlib_client_shared::{ - file_logger, keypair, Callbacks, Cidrv4, Cidrv6, Error, LoginUrl, ResourceDescription, Session, - Sockets, + callbacks::ResourceDescription, file_logger, keypair, Callbacks, Cidrv4, Cidrv6, Error, + LoginUrl, Session, Sockets, }; use secrecy::SecretString; use std::{ diff --git a/rust/connlib/clients/shared/src/eventloop.rs b/rust/connlib/clients/shared/src/eventloop.rs index 9aa2654d93..09a40a40a7 100644 --- a/rust/connlib/clients/shared/src/eventloop.rs +++ b/rust/connlib/clients/shared/src/eventloop.rs @@ -269,6 +269,7 @@ where gateway_id, resource_id, relays, + site_id, .. }) => { let should_accept = self @@ -280,10 +281,12 @@ where return; } - match self - .tunnel - .create_or_reuse_connection(resource_id, gateway_id, relays) - { + match self.tunnel.create_or_reuse_connection( + resource_id, + gateway_id, + relays, + site_id, + ) { Ok(firezone_tunnel::Request::NewConnection(connection_request)) => { // TODO: keep track for the response let _id = self.portal.send( @@ -321,7 +324,7 @@ where tracing::debug!(resource_id = %offline_resource, "Resource is offline"); - self.tunnel.cleanup_connection(offline_resource); + self.tunnel.set_resource_offline(offline_resource); } ErrorReply::Disabled => { diff --git a/rust/connlib/clients/shared/src/lib.rs b/rust/connlib/clients/shared/src/lib.rs index 44a0935b60..7cc4ccfaac 100644 --- a/rust/connlib/clients/shared/src/lib.rs +++ b/rust/connlib/clients/shared/src/lib.rs @@ -1,7 +1,7 @@ //! Main connlib library for clients. pub use connlib_shared::messages::client::ResourceDescription; pub use connlib_shared::{ - keypair, Callbacks, Cidrv4, Cidrv6, Error, LoginUrl, LoginUrlError, StaticSecret, + callbacks, keypair, Callbacks, Cidrv4, Cidrv6, Error, LoginUrl, LoginUrlError, StaticSecret, }; pub use firezone_tunnel::Sockets; pub use tracing_appender::non_blocking::WorkerGuard; diff --git a/rust/connlib/clients/shared/src/messages.rs b/rust/connlib/clients/shared/src/messages.rs index f9f1339d49..c6fe3a102a 100644 --- a/rust/connlib/clients/shared/src/messages.rs +++ b/rust/connlib/clients/shared/src/messages.rs @@ -1,6 +1,7 @@ use connlib_shared::messages::{ - client::ResourceDescription, GatewayId, GatewayResponse, Interface, Key, Relay, RelaysPresence, - RequestConnection, ResourceId, ReuseConnection, + client::{ResourceDescription, SiteId}, + GatewayId, GatewayResponse, Interface, Key, Relay, RelaysPresence, RequestConnection, + ResourceId, ReuseConnection, }; use serde::{Deserialize, Serialize}; use std::{collections::HashSet, net::IpAddr}; @@ -25,6 +26,8 @@ pub struct ConnectionDetails { pub resource_id: ResourceId, pub gateway_id: GatewayId, pub gateway_remote_ip: IpAddr, + #[serde(rename = "gateway_group_id")] + pub site_id: SiteId, } #[derive(Debug, Deserialize, Clone, PartialEq)] @@ -101,8 +104,7 @@ mod test { use super::*; use chrono::DateTime; use connlib_shared::messages::{ - client::ResourceDescriptionCidr, - client::{GatewayGroup, ResourceDescriptionDns}, + client::{ResourceDescriptionCidr, ResourceDescriptionDns, Site}, DnsServer, IpDnsServer, Stun, Turn, }; use phoenix_channel::{OutboundRequestId, PhoenixMessage}; @@ -234,7 +236,7 @@ mod test { address: "172.172.0.0/16".parse().unwrap(), name: "172.172.0.0/16".to_string(), address_description: "cidr resource".to_string(), - gateway_groups: vec![GatewayGroup { + sites: vec![Site { name: "test".to_string(), id: "bf56f32d-7b2c-4f5d-a784-788977d014a4".parse().unwrap(), }], @@ -244,7 +246,7 @@ mod test { address: "gitlab.mycorp.com".to_string(), name: "gitlab.mycorp.com".to_string(), address_description: "dns resource".to_string(), - gateway_groups: vec![GatewayGroup { + sites: vec![Site { name: "test".to_string(), id: "bf56f32d-7b2c-4f5d-a784-788977d014a4".parse().unwrap(), }], @@ -307,7 +309,7 @@ mod test { address: "172.172.0.0/16".parse().unwrap(), name: "172.172.0.0/16".to_string(), address_description: "cidr resource".to_string(), - gateway_groups: vec![GatewayGroup { + sites: vec![Site { name: "test".to_string(), id: "bf56f32d-7b2c-4f5d-a784-788977d014a4".parse().unwrap(), }], @@ -317,7 +319,7 @@ mod test { address: "gitlab.mycorp.com".to_string(), name: "gitlab.mycorp.com".to_string(), address_description: "dns resource".to_string(), - gateway_groups: vec![GatewayGroup { + sites: vec![Site { name: "test".to_string(), id: "bf56f32d-7b2c-4f5d-a784-788977d014a4".parse().unwrap(), }], @@ -536,6 +538,7 @@ mod test { gateway_id: "73037362-715d-4a83-a749-f18eadd970e6".parse().unwrap(), gateway_remote_ip: "172.28.0.1".parse().unwrap(), resource_id: "f16ecfa0-a94f-4bfd-a2ef-1cc1f2ef3da3".parse().unwrap(), + site_id: "bf56f32d-7b2c-4f5d-a784-788977d014a4".parse().unwrap(), relays: vec![ Relay::Stun(Stun { id: "c9cb8892-e355-41e6-a882-b6d6c38beb66".parse().unwrap(), @@ -573,6 +576,7 @@ mod test { "resource_id": "f16ecfa0-a94f-4bfd-a2ef-1cc1f2ef3da3", "gateway_id": "73037362-715d-4a83-a749-f18eadd970e6", "gateway_remote_ip": "172.28.0.1", + "gateway_group_id": "bf56f32d-7b2c-4f5d-a784-788977d014a4", "relays": [ { "id": "c9cb8892-e355-41e6-a882-b6d6c38beb66", diff --git a/rust/connlib/shared/Cargo.toml b/rust/connlib/shared/Cargo.toml index 22ae16b943..b9fc3bbf03 100644 --- a/rust/connlib/shared/Cargo.toml +++ b/rust/connlib/shared/Cargo.toml @@ -7,6 +7,7 @@ edition = "2021" # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html [features] mock = [] +proptest = ["dep:proptest", "dep:itertools"] [dependencies] anyhow = "1.0.82" @@ -34,13 +35,14 @@ libc = "0.2" snownet = { workspace = true } phoenix-channel = { workspace = true } proptest = { version = "1.4.0", optional = true } +itertools = { version = "0.12", optional = true } # Needed for Android logging until tracing is working log = "0.4" [dev-dependencies] -itertools = "0.12" tempfile = "3.10.1" +itertools = "0.12" mutants = "0.0.3" # Needed to mark functions as exempt from `cargo-mutants` testing tokio = { version = "1.36", features = ["macros", "rt"] } diff --git a/rust/connlib/shared/src/callbacks.rs b/rust/connlib/shared/src/callbacks.rs index 084537a133..19a493b348 100644 --- a/rust/connlib/shared/src/callbacks.rs +++ b/rust/connlib/shared/src/callbacks.rs @@ -1,9 +1,12 @@ -use crate::messages::client::ResourceDescription; -use ip_network::{Ipv4Network, Ipv6Network}; -use serde::Serialize; +use ip_network::{IpNetwork, Ipv4Network, Ipv6Network}; +use serde::{Deserialize, Serialize}; +use std::borrow::Cow; use std::fmt::Debug; use std::net::{IpAddr, Ipv4Addr, Ipv6Addr}; +use crate::messages::client::Site; +use crate::messages::ResourceId; + // Avoids having to map types for Windows type RawFd = i32; @@ -39,6 +42,123 @@ impl From for Cidrv6 { } } +#[derive(Debug, Serialize, Deserialize, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)] +pub enum Status { + Unknown, + Online, + Offline, +} + +#[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Eq, Hash)] +#[serde(tag = "type", rename_all = "snake_case")] +pub enum ResourceDescription { + Dns(ResourceDescriptionDns), + Cidr(ResourceDescriptionCidr), +} + +impl ResourceDescription { + pub fn name(&self) -> &str { + match self { + ResourceDescription::Dns(r) => &r.name, + ResourceDescription::Cidr(r) => &r.name, + } + } + + pub fn status(&self) -> Status { + match self { + ResourceDescription::Dns(r) => r.status, + ResourceDescription::Cidr(r) => r.status, + } + } + + pub fn id(&self) -> ResourceId { + match self { + ResourceDescription::Dns(r) => r.id, + ResourceDescription::Cidr(r) => r.id, + } + } + + /// What the GUI clients should paste to the clipboard, e.g. `https://github.com/firezone` + pub fn pastable(&self) -> Cow<'_, str> { + match self { + ResourceDescription::Dns(r) => Cow::from(&r.address), + ResourceDescription::Cidr(r) => Cow::from(r.address.to_string()), + } + } +} + +impl From for crate::messages::client::ResourceDescription { + fn from(value: ResourceDescription) -> Self { + match value { + ResourceDescription::Dns(r) => { + crate::messages::client::ResourceDescription::Dns(r.into()) + } + ResourceDescription::Cidr(r) => { + crate::messages::client::ResourceDescription::Cidr(r.into()) + } + } + } +} + +#[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Eq, Hash)] +pub struct ResourceDescriptionDns { + /// Resource's id. + pub id: ResourceId, + /// Internal resource's domain name. + pub address: String, + /// Name of the resource. + /// + /// Used only for display. + pub name: String, + + pub address_description: String, + pub sites: Vec, + + pub status: Status, +} + +impl From for crate::messages::client::ResourceDescriptionDns { + fn from(r: ResourceDescriptionDns) -> Self { + crate::messages::client::ResourceDescriptionDns { + id: r.id, + address: r.address, + address_description: r.address_description, + name: r.name, + sites: r.sites, + } + } +} + +/// Description of a resource that maps to a CIDR. +#[derive(Debug, Serialize, Deserialize, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)] +pub struct ResourceDescriptionCidr { + /// Resource's id. + pub id: ResourceId, + /// CIDR that this resource points to. + pub address: IpNetwork, + /// Name of the resource. + /// + /// Used only for display. + pub name: String, + + pub address_description: String, + pub sites: Vec, + + pub status: Status, +} + +impl From for crate::messages::client::ResourceDescriptionCidr { + fn from(r: ResourceDescriptionCidr) -> Self { + crate::messages::client::ResourceDescriptionCidr { + id: r.id, + address: r.address, + address_description: r.address_description, + name: r.name, + sites: r.sites, + } + } +} + /// Traits that will be used by connlib to callback the client upper layers. pub trait Callbacks: Clone + Send + Sync { /// Called when the tunnel address is set. diff --git a/rust/connlib/shared/src/lib.rs b/rust/connlib/shared/src/lib.rs index 87185dd2e2..9f7130d7bf 100644 --- a/rust/connlib/shared/src/lib.rs +++ b/rust/connlib/shared/src/lib.rs @@ -3,7 +3,7 @@ //! This includes types provided by external crates, i.e. [boringtun] to make sure that //! we are using the same version across our own crates. -mod callbacks; +pub mod callbacks; pub mod error; pub mod messages; diff --git a/rust/connlib/shared/src/messages.rs b/rust/connlib/shared/src/messages.rs index d687bedbf8..0d1b725b37 100644 --- a/rust/connlib/shared/src/messages.rs +++ b/rust/connlib/shared/src/messages.rs @@ -50,6 +50,13 @@ impl ClientId { } } +impl GatewayId { + #[cfg(feature = "proptest")] + pub(crate) fn from_u128(v: u128) -> Self { + Self(Uuid::from_u128(v)) + } +} + #[derive(Hash, Debug, Deserialize, Serialize, Clone, Copy, PartialEq, Eq)] pub struct ClientId(Uuid); @@ -315,7 +322,7 @@ mod tests { use super::{ client::ResourceDescription, - client::{GatewayGroup, ResourceDescriptionDns}, + client::{ResourceDescriptionDns, Site}, ResourceId, }; @@ -325,7 +332,7 @@ mod tests { name: name.to_string(), address: "unused.example.com".to_string(), address_description: "test description".to_string(), - gateway_groups: vec![GatewayGroup { + sites: vec![Site { name: "test".to_string(), id: "99ba0c1e-5189-4cfc-a4db-fd6cb1c937fd".parse().unwrap(), }], diff --git a/rust/connlib/shared/src/messages/client.rs b/rust/connlib/shared/src/messages/client.rs index 6025bc604d..0b3f4139ed 100644 --- a/rust/connlib/shared/src/messages/client.rs +++ b/rust/connlib/shared/src/messages/client.rs @@ -1,11 +1,13 @@ //! Client related messages that are needed within connlib -use std::{borrow::Cow, str::FromStr}; +use std::{collections::HashSet, str::FromStr}; use ip_network::IpNetwork; use serde::{Deserialize, Serialize}; use uuid::Uuid; +use crate::callbacks::Status; + use super::ResourceId; /// Description of a resource that maps to a DNS record. @@ -21,7 +23,21 @@ pub struct ResourceDescriptionDns { pub name: String, pub address_description: String, - pub gateway_groups: Vec, + #[serde(rename = "gateway_groups")] + pub sites: Vec, +} + +impl ResourceDescriptionDns { + fn with_status(self, status: Status) -> crate::callbacks::ResourceDescriptionDns { + crate::callbacks::ResourceDescriptionDns { + id: self.id, + address: self.address, + name: self.name, + address_description: self.address_description, + sites: self.sites, + status, + } + } } /// Description of a resource that maps to a CIDR. @@ -37,16 +53,30 @@ pub struct ResourceDescriptionCidr { pub name: String, pub address_description: String, - pub gateway_groups: Vec, + #[serde(rename = "gateway_groups")] + pub sites: Vec, +} + +impl ResourceDescriptionCidr { + fn with_status(self, status: Status) -> crate::callbacks::ResourceDescriptionCidr { + crate::callbacks::ResourceDescriptionCidr { + id: self.id, + address: self.address, + name: self.name, + address_description: self.address_description, + sites: self.sites, + status, + } + } } #[derive(Debug, Deserialize, Serialize, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)] -pub struct GatewayGroup { +pub struct Site { pub name: String, pub id: SiteId, } -#[derive(Debug, Deserialize, Serialize, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)] +#[derive(Debug, Deserialize, Serialize, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)] pub struct SiteId(Uuid); impl FromStr for SiteId { @@ -79,19 +109,18 @@ impl ResourceDescription { } } - /// What the GUI clients should show as the user-friendly display name, e.g. `Firezone GitHub` - pub fn name(&self) -> &str { + pub fn sites(&self) -> HashSet<&Site> { match self { - ResourceDescription::Dns(r) => &r.name, - ResourceDescription::Cidr(r) => &r.name, + ResourceDescription::Dns(r) => HashSet::from_iter(r.sites.iter()), + ResourceDescription::Cidr(r) => HashSet::from_iter(r.sites.iter()), } } - /// What the GUI clients should paste to the clipboard, e.g. `https://github.com/firezone` - pub fn pastable(&self) -> Cow<'_, str> { + /// What the GUI clients should show as the user-friendly display name, e.g. `Firezone GitHub` + pub fn name(&self) -> &str { match self { - ResourceDescription::Dns(r) => Cow::from(&r.address), - ResourceDescription::Cidr(r) => Cow::from(r.address.to_string()), + ResourceDescription::Dns(r) => &r.name, + ResourceDescription::Cidr(r) => &r.name, } } @@ -106,6 +135,17 @@ impl ResourceDescription { _ => true, } } + + pub fn with_status(self, status: Status) -> crate::callbacks::ResourceDescription { + match self { + ResourceDescription::Dns(r) => { + crate::callbacks::ResourceDescription::Dns(r.with_status(status)) + } + ResourceDescription::Cidr(r) => { + crate::callbacks::ResourceDescription::Cidr(r.with_status(status)) + } + } + } } impl PartialOrd for ResourceDescription { diff --git a/rust/connlib/shared/src/proptest.rs b/rust/connlib/shared/src/proptest.rs index 79953e1468..df6a6c9e3a 100644 --- a/rust/connlib/shared/src/proptest.rs +++ b/rust/connlib/shared/src/proptest.rs @@ -1,64 +1,110 @@ use crate::messages::{ client::ResourceDescriptionCidr, - client::{GatewayGroup, ResourceDescriptionDns, SiteId}, - ClientId, ResourceId, + client::{ResourceDescription, ResourceDescriptionDns, Site, SiteId}, + ClientId, GatewayId, ResourceId, }; use ip_network::{IpNetwork, Ipv4Network, Ipv6Network}; +use itertools::Itertools; use proptest::{ arbitrary::{any, any_with}, collection, sample, - strategy::Strategy, + strategy::{Just, Strategy}, }; use std::net::{Ipv4Addr, Ipv6Addr}; -pub fn dns_resource() -> impl Strategy { +// Generate resources sharing 1 site +pub fn resources_sharing_site() -> impl Strategy, Site)> { + (collection::vec(sites(), 1..=100), site()).prop_flat_map(|(sites, site)| { + ( + sites + .iter() + .map(|sites| { + let mut sites = sites.clone(); + sites.push(site.clone()); + resource(sites.clone()) + }) + .collect_vec(), + Just(site), + ) + }) +} + +// Generate resources sharing all sites +pub fn resources_sharing_all_sites() -> impl Strategy> { + sites().prop_flat_map(|sites| collection::vec(resource(sites), 1..=100)) +} + +pub fn resource(sites: Vec) -> impl Strategy { + any::().prop_flat_map(move |is_dns| { + if is_dns { + dns_resource_with_sites(sites.clone()) + .prop_map(ResourceDescription::Dns) + .boxed() + } else { + cidr_resource_with_sites(8, sites.clone()) + .prop_map(ResourceDescription::Cidr) + .boxed() + } + }) +} + +pub fn dns_resource_with_sites(sites: Vec) -> impl Strategy { ( resource_id(), resource_name(), dns_resource_address(), - gateway_groups(), address_description(), ) - .prop_map(|(id, name, address, gateway_groups, address_description)| { - ResourceDescriptionDns { + .prop_map( + move |(id, name, address, address_description)| ResourceDescriptionDns { id, address, name, - gateway_groups, + sites: sites.clone(), address_description, - } - }) + }, + ) } -pub fn cidr_resource(host_mask_bits: usize) -> impl Strategy { +pub fn cidr_resource_with_sites( + host_mask_bits: usize, + sites: Vec, +) -> impl Strategy { ( resource_id(), resource_name(), ip_network(host_mask_bits), - gateway_groups(), address_description(), ) - .prop_map(|(id, name, address, gateway_groups, address_description)| { - ResourceDescriptionCidr { + .prop_map( + move |(id, name, address, address_description)| ResourceDescriptionCidr { id, address, name, - gateway_groups, + sites: sites.clone(), address_description, - } - }) + }, + ) +} + +pub fn dns_resource() -> impl Strategy { + sites().prop_flat_map(dns_resource_with_sites) +} + +pub fn cidr_resource(host_mask_bits: usize) -> impl Strategy { + sites().prop_flat_map(move |sites| cidr_resource_with_sites(host_mask_bits, sites)) } pub fn address_description() -> impl Strategy { any_with::("[a-z]{4,10}".into()) } -pub fn gateway_groups() -> impl Strategy> { - collection::vec(gateway_group(), 1..=10) +pub fn sites() -> impl Strategy> { + collection::vec(site(), 1..=10) } -pub fn gateway_group() -> impl Strategy { - (any_with::("[a-z]{4,10}".into()), any::()).prop_map(|(name, id)| GatewayGroup { +pub fn site() -> impl Strategy { + (any_with::("[a-z]{4,10}".into()), any::()).prop_map(|(name, id)| Site { name, id: SiteId::from_u128(id), }) @@ -68,6 +114,10 @@ pub fn resource_id() -> impl Strategy + Clone { any::().prop_map(ResourceId::from_u128) } +pub fn gateway_id() -> impl Strategy + Clone { + any::().prop_map(GatewayId::from_u128) +} + pub fn client_id() -> impl Strategy { any::().prop_map(ClientId::from_u128) } diff --git a/rust/connlib/tunnel/src/client.rs b/rust/connlib/tunnel/src/client.rs index 0621686d00..2ae8eed67c 100644 --- a/rust/connlib/tunnel/src/client.rs +++ b/rust/connlib/tunnel/src/client.rs @@ -1,13 +1,15 @@ use crate::peer_store::PeerStore; use crate::{dns, dns::DnsQuery}; use bimap::BiMap; +use connlib_shared::callbacks::Status; use connlib_shared::error::{ConnlibError as Error, ConnlibError}; +use connlib_shared::messages::client::{Site, SiteId}; use connlib_shared::messages::{ client::ResourceDescription, client::ResourceDescriptionCidr, client::ResourceDescriptionDns, Answer, ClientPayload, DnsServer, DomainResponse, GatewayId, Interface as InterfaceConfig, IpDnsServer, Key, Offer, Relay, RelayId, RequestConnection, ResourceId, ReuseConnection, }; -use connlib_shared::{Callbacks, Dname, PublicKey, StaticSecret}; +use connlib_shared::{callbacks, Callbacks, Dname, PublicKey, StaticSecret}; use domain::base::Rtype; use ip_network::{IpNetwork, Ipv4Network, Ipv6Network}; use ip_network_table::IpNetworkTable; @@ -174,6 +176,15 @@ where self.role_state.on_connection_failed(id); } + pub fn set_resource_offline(&mut self, id: ResourceId) { + self.role_state.set_resource_offline(id); + + self.role_state.on_connection_failed(id); + + self.callbacks + .on_update_resources(self.role_state.resources()); + } + pub fn add_ice_candidate(&mut self, conn_id: GatewayId, ice_candidate: String) { self.role_state .node @@ -191,10 +202,12 @@ where resource_id: ResourceId, gateway_id: GatewayId, relays: Vec, + site_id: SiteId, ) -> connlib_shared::Result { self.role_state.create_or_reuse_connection( resource_id, gateway_id, + site_id, stun(&relays, |addr| self.io.sockets_ref().can_handle(addr)), turn(&relays), ) @@ -277,6 +290,9 @@ pub struct ClientState { next_dns_refresh: Option, system_resolvers: Vec, + + gateways_site: HashMap, + sites_status: HashMap, } #[derive(Debug, Clone, PartialEq, Eq)] @@ -306,11 +322,51 @@ impl ClientState { next_dns_refresh: Default::default(), node: ClientNode::new(private_key), system_resolvers: Default::default(), + sites_status: Default::default(), + gateways_site: Default::default(), + } + } + + pub(crate) fn resources(&self) -> Vec { + self.resource_ids + .values() + .sorted() + .cloned() + .map(|r| { + let status = self.resource_status(&r); + r.with_status(status) + }) + .collect_vec() + } + + fn resource_status(&self, resource: &ResourceDescription) -> Status { + if resource.sites().iter().any(|s| { + self.sites_status + .get(&s.id) + .is_some_and(|s| *s == Status::Online) + }) { + return Status::Online; + } + + if resource.sites().iter().all(|s| { + self.sites_status + .get(&s.id) + .is_some_and(|s| *s == Status::Offline) + }) { + return Status::Offline; } + + Status::Unknown } - fn resources(&self) -> Vec { - self.resource_ids.values().sorted().cloned().collect_vec() + fn set_resource_offline(&mut self, id: ResourceId) { + let Some(resource) = self.resource_ids.get(&id).cloned() else { + return; + }; + + for Site { id, .. } in resource.sites() { + self.sites_status.insert(*id, Status::Offline); + } } pub(crate) fn encapsulate<'s>( @@ -434,11 +490,14 @@ impl ClientState { &mut self, resource_id: ResourceId, gateway_id: GatewayId, + site_id: SiteId, allowed_stun_servers: HashSet, allowed_turn_servers: HashSet<(RelayId, RelaySocket, String, String, String)>, ) -> connlib_shared::Result { tracing::trace!("create_or_reuse_connection"); + self.gateways_site.insert(gateway_id, site_id); + let desc = self .resource_ids .get(&resource_id) @@ -746,6 +805,7 @@ impl ClientState { } pub fn cleanup_connected_gateway(&mut self, gateway_id: &GatewayId) { + self.update_site_status_by_gateway(gateway_id, Status::Unknown); self.peers.remove(gateway_id); self.dns_resources_internal_ips.retain(|resource, _| { !self @@ -853,11 +913,24 @@ impl ClientState { conn_id: connection, candidate, }), - snownet::Event::ConnectionEstablished { .. } => {} + snownet::Event::ConnectionEstablished(id) => { + self.update_site_status_by_gateway(&id, Status::Online); + } } } } + fn update_site_status_by_gateway(&mut self, gateway_id: &GatewayId, status: Status) { + // Note: we can do this because in theory we shouldn't have multiple gateways for the same site + // connected at the same time. + self.sites_status.insert( + *self.gateways_site.get(gateway_id).expect( + "if we're updating a site status there should be an associated site to a gateway", + ), + status, + ); + } + pub(crate) fn poll_event(&mut self) -> Option { self.buffered_events.pop_front() } @@ -965,6 +1038,7 @@ impl ClientState { // If there's no allowed ip left we remove the whole peer because there's no point on keeping it around if peer.allowed_ips.is_empty() { self.peers.remove(&gateway_id); + self.update_site_status_by_gateway(&gateway_id, Status::Unknown); // TODO: should we have a Node::remove_connection? } } @@ -1433,8 +1507,13 @@ mod proptests { ]); assert_eq!( - hashset(client_state.resources().iter()), - hashset(&[ + hashset( + client_state + .resources() + .into_iter() + .map_into::() + ), + hashset([ ResourceDescription::Cidr(resource1.clone()), ResourceDescription::Dns(resource2.clone()) ]) @@ -1443,8 +1522,13 @@ mod proptests { client_state.add_resources(&[ResourceDescription::Cidr(resource3.clone())]); assert_eq!( - hashset(client_state.resources().iter()), - hashset(&[ + hashset( + client_state + .resources() + .into_iter() + .map_into::() + ), + hashset([ ResourceDescription::Cidr(resource1), ResourceDescription::Dns(resource2), ResourceDescription::Cidr(resource3) @@ -1468,8 +1552,13 @@ mod proptests { client_state.add_resources(&[ResourceDescription::Cidr(updated_resource.clone())]); assert_eq!( - hashset(client_state.resources().iter()), - hashset(&[ResourceDescription::Cidr(updated_resource),]) + hashset( + client_state + .resources() + .into_iter() + .map_into::() + ), + hashset([ResourceDescription::Cidr(updated_resource),]) ); assert_eq!( hashset(client_state.routes()), @@ -1490,14 +1579,19 @@ mod proptests { id: resource.id, name: resource.name, address_description: resource.address_description, - gateway_groups: resource.gateway_groups, + sites: resource.sites, }; client_state.add_resources(&[ResourceDescription::Cidr(dns_as_cidr_resource.clone())]); assert_eq!( - hashset(client_state.resources().iter()), - hashset(&[ResourceDescription::Cidr(dns_as_cidr_resource),]) + hashset( + client_state + .resources() + .into_iter() + .map_into::() + ), + hashset([ResourceDescription::Cidr(dns_as_cidr_resource),]) ); assert_eq!( hashset(client_state.routes()), @@ -1519,8 +1613,13 @@ mod proptests { client_state.remove_resources(&[dns_resource.id]); assert_eq!( - hashset(client_state.resources().iter()), - hashset(&[ResourceDescription::Cidr(cidr_resource.clone())]) + hashset( + client_state + .resources() + .into_iter() + .map_into::() + ), + hashset([ResourceDescription::Cidr(cidr_resource.clone())]) ); assert_eq!( hashset(client_state.routes()), @@ -1552,8 +1651,13 @@ mod proptests { ]); assert_eq!( - hashset(client_state.resources().iter()), - hashset(&[ + hashset( + client_state + .resources() + .into_iter() + .map_into::() + ), + hashset([ ResourceDescription::Dns(dns_resource2), ResourceDescription::Cidr(cidr_resource2.clone()), ]) @@ -1563,4 +1667,100 @@ mod proptests { expected_routes(vec![cidr_resource2.address]) ); } + + #[test_strategy::proptest] + fn setting_gateway_online_sets_all_related_resources_online( + #[strategy(resources_sharing_site())] resource_config_online: ( + Vec, + Site, + ), + #[strategy(resources_sharing_site())] resource_config_unknown: ( + Vec, + Site, + ), + #[strategy(gateway_id())] first_resource_gateway_id: GatewayId, + ) { + let (resources_online, site) = resource_config_online; + let (resources_unknown, _) = resource_config_unknown; + let mut client_state = ClientState::for_test(); + client_state.add_resources(&resources_online); + client_state.add_resources(&resources_unknown); + client_state.resources_gateways.insert( + resources_online.first().unwrap().id(), + first_resource_gateway_id, + ); + client_state + .gateways_site + .insert(first_resource_gateway_id, site.id); + + client_state.update_site_status_by_gateway(&first_resource_gateway_id, Status::Online); + + for resource in resources_online { + assert_eq!(client_state.resource_status(&resource), Status::Online); + } + + for resource in resources_unknown { + assert_eq!(client_state.resource_status(&resource), Status::Unknown); + } + } + + #[test_strategy::proptest] + fn disconnecting_gateway_sets_related_resources_unknown( + #[strategy(resources_sharing_site())] resource_config: (Vec, Site), + #[strategy(gateway_id())] first_resource_gateway_id: GatewayId, + ) { + let (resources, site) = resource_config; + let mut client_state = ClientState::for_test(); + client_state.add_resources(&resources); + client_state + .resources_gateways + .insert(resources.first().unwrap().id(), first_resource_gateway_id); + client_state + .gateways_site + .insert(first_resource_gateway_id, site.id); + + client_state.update_site_status_by_gateway(&first_resource_gateway_id, Status::Online); + client_state.update_site_status_by_gateway(&first_resource_gateway_id, Status::Unknown); + + for resource in resources { + assert_eq!(client_state.resource_status(&resource), Status::Unknown); + } + } + + #[test_strategy::proptest] + fn setting_resource_offline_doesnt_set_all_related_resources_offline( + #[strategy(resources_sharing_site())] resource_config_online: ( + Vec, + Site, + ), + ) { + let (mut resources, _) = resource_config_online; + let mut client_state = ClientState::for_test(); + client_state.add_resources(&resources); + let resource_offline = resources.pop().unwrap(); + + client_state.set_resource_offline(resource_offline.id()); + + assert_eq!( + client_state.resource_status(&resource_offline), + Status::Offline + ); + for resource in resources { + assert_eq!(client_state.resource_status(&resource), Status::Unknown); + } + } + + #[test_strategy::proptest] + fn setting_resource_offline_set_all_resources_sharing_all_groups_offline( + #[strategy(resources_sharing_all_sites())] resources: Vec, + ) { + let mut client_state = ClientState::for_test(); + client_state.add_resources(&resources); + + client_state.set_resource_offline(resources.first().unwrap().id()); + + for resource in resources { + assert_eq!(client_state.resource_status(&resource), Status::Offline); + } + } } diff --git a/rust/connlib/tunnel/src/lib.rs b/rust/connlib/tunnel/src/lib.rs index 608e5ad2dd..5d4451eb13 100644 --- a/rust/connlib/tunnel/src/lib.rs +++ b/rust/connlib/tunnel/src/lib.rs @@ -119,6 +119,8 @@ where )? { Poll::Ready(io::Input::Timeout(timeout)) => { self.role_state.handle_timeout(timeout); + self.callbacks + .on_update_resources(self.role_state.resources()); continue; } Poll::Ready(io::Input::Device(packet)) => { diff --git a/rust/gui-client/src-tauri/src/client/gui.rs b/rust/gui-client/src-tauri/src/client/gui.rs index 12b4ba394b..53988e2913 100644 --- a/rust/gui-client/src-tauri/src/client/gui.rs +++ b/rust/gui-client/src-tauri/src/client/gui.rs @@ -9,7 +9,7 @@ use crate::client::{ Failure, }; use anyhow::{bail, Context, Result}; -use connlib_client_shared::ResourceDescription; +use connlib_client_shared::callbacks::ResourceDescription; use connlib_shared::messages::ResourceId; use secrecy::{ExposeSecret, SecretString}; use std::{path::PathBuf, str::FromStr, sync::Arc, time::Duration}; diff --git a/rust/gui-client/src-tauri/src/client/gui/system_tray_menu.rs b/rust/gui-client/src-tauri/src/client/gui/system_tray_menu.rs index c94daeccd3..65b7c88912 100644 --- a/rust/gui-client/src-tauri/src/client/gui/system_tray_menu.rs +++ b/rust/gui-client/src-tauri/src/client/gui/system_tray_menu.rs @@ -3,7 +3,7 @@ //! "Notification Area" is Microsoft's official name instead of "System tray": //! -use connlib_client_shared::ResourceDescription; +use connlib_client_shared::callbacks::ResourceDescription; use std::str::FromStr; use tauri::{CustomMenuItem, SystemTrayMenu, SystemTrayMenuItem, SystemTraySubmenu}; diff --git a/rust/gui-client/src-tauri/src/client/tunnel-wrapper/in_proc.rs b/rust/gui-client/src-tauri/src/client/tunnel-wrapper/in_proc.rs index c1ae0d0648..4b71dda98b 100644 --- a/rust/gui-client/src-tauri/src/client/tunnel-wrapper/in_proc.rs +++ b/rust/gui-client/src-tauri/src/client/tunnel-wrapper/in_proc.rs @@ -10,8 +10,8 @@ use anyhow::{Context, Result}; use arc_swap::ArcSwap; -use connlib_client_shared::{ResourceDescription, Sockets}; -use connlib_shared::{keypair, LoginUrl}; +use connlib_client_shared::Sockets; +use connlib_shared::{callbacks::ResourceDescription, keypair, LoginUrl}; use secrecy::SecretString; use std::{ net::{IpAddr, Ipv4Addr, Ipv6Addr}, diff --git a/rust/gui-client/src-tauri/src/client/tunnel-wrapper/ipc.rs b/rust/gui-client/src-tauri/src/client/tunnel-wrapper/ipc.rs index 73a2f1428a..3b7b66f30e 100644 --- a/rust/gui-client/src-tauri/src/client/tunnel-wrapper/ipc.rs +++ b/rust/gui-client/src-tauri/src/client/tunnel-wrapper/ipc.rs @@ -1,6 +1,7 @@ use anyhow::{Context, Result}; use arc_swap::ArcSwap; -use connlib_client_shared::{Callbacks, ResourceDescription}; +use connlib_client_shared::Callbacks; +use connlib_shared::callbacks::ResourceDescription; use firezone_headless_client::{imp::sock_path, IpcClientMsg, IpcServerMsg}; use futures::{SinkExt, StreamExt}; use secrecy::{ExposeSecret, SecretString}; diff --git a/rust/headless-client/src/imp_linux.rs b/rust/headless-client/src/imp_linux.rs index 37502e25c7..9fa6de8582 100644 --- a/rust/headless-client/src/imp_linux.rs +++ b/rust/headless-client/src/imp_linux.rs @@ -3,9 +3,9 @@ use super::{Cli, IpcClientMsg, IpcServerMsg, FIREZONE_GROUP, TOKEN_ENV_KEY}; use anyhow::{bail, Context as _, Result}; use clap::Parser; -use connlib_client_shared::{file_logger, Callbacks, ResourceDescription, Sockets}; +use connlib_client_shared::{file_logger, Callbacks, Sockets}; use connlib_shared::{ - keypair, + callbacks, keypair, linux::{etc_resolv_conf, get_dns_control_from_env, DnsControlMethod}, LoginUrl, }; @@ -240,8 +240,8 @@ impl Callbacks for CallbackHandlerIpc { None } - fn on_update_resources(&self, resources: Vec) { - tracing::info!(len = resources.len(), "New resource list"); + fn on_update_resources(&self, resources: Vec) { + tracing::debug!(len = resources.len(), "New resource list"); self.cb_tx .try_send(IpcServerMsg::OnUpdateResources(resources)) .expect("Should be able to send OnUpdateResources"); diff --git a/rust/headless-client/src/lib.rs b/rust/headless-client/src/lib.rs index f57ec8bbed..2379d647a5 100644 --- a/rust/headless-client/src/lib.rs +++ b/rust/headless-client/src/lib.rs @@ -10,9 +10,8 @@ use anyhow::{Context, Result}; use clap::Parser; -use connlib_client_shared::{ - file_logger, keypair, Callbacks, LoginUrl, ResourceDescription, Session, Sockets, -}; +use connlib_client_shared::{file_logger, keypair, Callbacks, LoginUrl, Session, Sockets}; +use connlib_shared::callbacks; use firezone_cli_utils::setup_global_subscriber; use secrecy::SecretString; use std::{future, net::IpAddr, path::PathBuf, task::Poll}; @@ -133,7 +132,7 @@ pub enum IpcClientMsg { pub enum IpcServerMsg { Ok, OnDisconnect, - OnUpdateResources(Vec), + OnUpdateResources(Vec), TunnelReady, } @@ -274,9 +273,8 @@ impl Callbacks for CallbackHandler { .expect("should be able to tell the main thread that we disconnected"); } - fn on_update_resources(&self, resources: Vec) { + fn on_update_resources(&self, resources: Vec) { // See easily with `export RUST_LOG=firezone_headless_client=debug` - tracing::debug!(len = resources.len(), "Printing the resource list one time"); for resource in &resources { tracing::debug!(?resource); }