1use std::net::SocketAddr;
2use std::path::PathBuf;
3
4use std::sync::Arc;
5use tokio::task::JoinSet;
6use unb_core::validate_node_identifier;
7use unb_runtime::Pipe;
8
9use crate::host::HostConfig;
10use crate::{
11 ConnectionStatus, Endpoint, EndpointSet, Hosting, Node, PeerConnection, TransportKind,
12};
13use unb_runtime::WsError;
14
15pub async fn connect_unix(path: impl AsRef<std::path::Path>) -> Result<Pipe, WsError> {
16 Ok(Pipe::Piped {
17 pipe: unb_transport::unix::connect(path).await?,
18 initiator: true,
19 })
20}
21
22pub async fn accept_unix(listener: &unb_transport::unix::UnixListener) -> Result<Pipe, WsError> {
23 Ok(Pipe::Piped {
24 pipe: listener.accept().await?,
25 initiator: false,
26 })
27}
28
29#[derive(Clone, Debug)]
30pub struct ParentLink {
31 pub node: String,
32 pub path: PathBuf,
33}
34
35impl ParentLink {
36 pub fn unix(node: impl Into<String>, path: impl Into<PathBuf>) -> ParentLink {
37 ParentLink {
38 node: node.into(),
39 path: path.into(),
40 }
41 }
42
43 fn validate(&self) -> Result<(), WsError> {
44 validate_node_identifier(&self.node)
45 .map_err(|error| WsError::Connect(error.to_string()))?;
46 if self.path.as_os_str().is_empty() {
47 return Err(WsError::Connect("parent Unix path is empty".into()));
48 }
49 Ok(())
50 }
51}
52
53#[derive(Default)]
54pub struct TopologyConfig {
55 pub host: Option<HostConfig>,
56 pub parent: Option<ParentLink>,
57}
58
59impl TopologyConfig {
60 pub fn new() -> TopologyConfig {
61 TopologyConfig::default()
62 }
63
64 pub fn host(host: HostConfig) -> TopologyConfig {
65 TopologyConfig {
66 host: Some(host),
67 parent: None,
68 }
69 }
70
71 pub fn parent(parent: ParentLink) -> TopologyConfig {
72 TopologyConfig {
73 host: None,
74 parent: Some(parent),
75 }
76 }
77
78 pub fn with_parent(mut self, parent: ParentLink) -> TopologyConfig {
79 self.parent = Some(parent);
80 self
81 }
82
83 pub fn validate(&self) -> Result<(), WsError> {
84 if self.host.is_none() && self.parent.is_none() {
85 return Err(WsError::Connect(
86 "topology requires a host or parent".into(),
87 ));
88 }
89 if let Some(host) = &self.host {
90 host.validate()
91 .map_err(|error| WsError::Connect(error.to_string()))?;
92 }
93 if let Some(parent) = &self.parent {
94 parent.validate()?;
95 }
96 Ok(())
97 }
98}
99
100pub struct UnbTopology {
101 hosting: Option<Hosting>,
102 parent: Option<PeerConnection>,
103}
104
105impl UnbTopology {
106 pub fn is_finished(&self) -> bool {
107 self.hosting.as_ref().is_some_and(Hosting::is_finished)
108 || self.parent.as_ref().is_some_and(|parent| {
109 matches!(parent.status(), ConnectionStatus::Disconnected { .. })
110 })
111 }
112
113 pub fn health(&self) -> crate::host::HealthStatus {
114 let hosting = self.hosting.as_ref().map(Hosting::health);
115 let parent_link_ready = match &self.parent {
116 None => true,
117 Some(parent) => parent.status() == ConnectionStatus::Connected,
118 };
119 crate::host::HealthStatus {
120 process_alive: true,
121 websocket_bound: hosting.is_some_and(|health| health.websocket_bound),
122 websocket_addr: hosting.and_then(|health| health.websocket_addr),
123 webtransport_bound: hosting.is_some_and(|health| health.webtransport_bound),
124 webtransport_addr: hosting.and_then(|health| health.webtransport_addr),
125 listeners_running: hosting.is_some_and(|health| health.listeners_running),
126 parent_link_ready,
127 child_link_ready: true,
128 }
129 }
130
131 pub fn websocket_addr(&self) -> Option<SocketAddr> {
132 self.hosting.as_ref().and_then(Hosting::websocket_addr)
133 }
134
135 pub fn webtransport_addr(&self) -> Option<SocketAddr> {
136 self.hosting.as_ref().and_then(Hosting::webtransport_addr)
137 }
138
139 pub async fn shutdown(self) -> Result<(), WsError> {
140 if let Some(parent) = &self.parent {
141 parent.disconnect();
142 }
143 if let Some(hosting) = &self.hosting {
144 hosting.cancel();
145 }
146 let hosting = async {
147 match self.hosting {
148 Some(hosting) => hosting
149 .shutdown()
150 .await
151 .map_err(|error| WsError::Connect(error.to_string())),
152 None => Ok(()),
153 }
154 };
155 hosting.await
156 }
157
158 pub async fn wait(&mut self) -> Result<(), WsError> {
159 let host_failure = |error: crate::host::HostError| WsError::Connect(error.to_string());
160 match (&mut self.hosting, &mut self.parent) {
161 (Some(hosting), Some(parent)) => tokio::select! {
162 biased;
163 result = hosting.wait() => result.map_err(host_failure),
164 () = wait_for_parent_termination(parent.clone()) => Ok(()),
165 },
166 (Some(hosting), None) => hosting.wait().await.map_err(host_failure),
167 (None, Some(parent)) => {
168 wait_for_parent_termination(parent.clone()).await;
169 Ok(())
170 }
171 (None, None) => Ok(()),
172 }
173 }
174}
175
176async fn wait_for_parent_termination(parent: PeerConnection) {
177 loop {
178 let changed = parent.changed();
179 if matches!(parent.status(), ConnectionStatus::Disconnected { .. }) {
180 return;
181 }
182 let _ = changed.await;
183 }
184}
185
186impl Node {
187 pub async fn start_topology(
188 self: &Arc<Self>,
189 config: TopologyConfig,
190 ) -> Result<UnbTopology, WsError> {
191 config.validate()?;
192 let hosting = match config.host {
193 Some(host) => Some(
194 host.start(self)
195 .await
196 .map_err(|error| WsError::Connect(error.to_string()))?,
197 ),
198 None => None,
199 };
200 let parent = match config.parent {
201 Some(parent) => {
202 let endpoint = Endpoint {
203 kind: TransportKind::Unix,
204 address: parent.path.to_string_lossy().into_owned(),
205 cert_hash: None,
206 };
207 match self
208 .connect_expected(EndpointSet::from(endpoint), &parent.node)
209 .await
210 {
211 Ok(connection) if connection.peer() == parent.node => Some(connection),
212 Ok(connection) => {
213 let actual = connection.peer().to_owned();
214 connection.disconnect();
215 if let Some(hosting) = hosting {
216 let _ = hosting.shutdown().await;
217 }
218 return Err(WsError::Connect(format!(
219 "peer identity mismatch: expected {:?}, got {:?}",
220 parent.node, actual
221 )));
222 }
223 Err(error) => {
224 if let Some(hosting) = hosting {
225 let _ = hosting.shutdown().await;
226 }
227 return Err(WsError::Connect(error.to_string()));
228 }
229 }
230 }
231 None => None,
232 };
233 Ok(UnbTopology { hosting, parent })
234 }
235}
236
237pub struct UnixHosting {
238 cancellation: unb_runtime::CancellationToken,
239 task: tokio::task::JoinHandle<Result<(), WsError>>,
240}
241
242impl UnixHosting {
243 pub fn is_finished(&self) -> bool {
244 self.task.is_finished()
245 }
246
247 pub fn cancel(&self) {
248 self.cancellation.cancel();
249 }
250
251 pub async fn wait(&mut self) -> Result<(), WsError> {
252 (&mut self.task)
253 .await
254 .map_err(|error| WsError::Connect(error.to_string()))?
255 }
256
257 pub async fn shutdown(self) -> Result<(), WsError> {
258 self.cancellation.cancel();
259 self.task
260 .await
261 .map_err(|error| WsError::Connect(error.to_string()))?
262 }
263}
264
265impl Node {
266 pub fn host_unix_child(
267 self: &Arc<Self>,
268 path: impl AsRef<std::path::Path>,
269 expected_child: impl Into<String>,
270 ) -> Result<UnixHosting, WsError> {
271 let expected_child = expected_child.into();
272 validate_node_identifier(&expected_child)
273 .map_err(|error| WsError::Connect(error.to_string()))?;
274 let listener = unb_transport::unix::UnixListener::bind(path)?;
275 let cancellation = self.cancellation().child_token();
276 let task_cancellation = cancellation.clone();
277 let node = self.clone();
278 let task = tokio::spawn(async move {
279 let mut connections = JoinSet::new();
280 loop {
281 tokio::select! {
282 biased;
283 () = task_cancellation.cancelled() => break,
284 result = connections.join_next(), if !connections.is_empty() => {
285 result.expect("connection task exists").map_err(|error| WsError::Connect(error.to_string()))?;
286 }
287 pipe = accept_unix(&listener) => {
288 let pipe = pipe?;
289 let connection = tokio::select! {
290 biased;
291 () = task_cancellation.cancelled() => break,
292 connection = node.connect_transport(&expected_child, pipe) => connection,
293 };
294 if let Ok(connection) = connection {
295 let cancellation = task_cancellation.clone();
296 connections.spawn(async move {
297 tokio::select! {
298 biased;
299 () = cancellation.cancelled() => {
300 connection.shutdown();
301 connection.closed().await;
302 }
303 () = connection.closed() => {}
304 }
305 });
306 }
307 }
308 }
309 }
310 while let Some(result) = connections.join_next().await {
311 result.map_err(|error| WsError::Connect(error.to_string()))?;
312 }
313 Ok(())
314 });
315 Ok(UnixHosting { cancellation, task })
316 }
317
318 pub fn host_unix(self: &Arc<Self>, listener: unb_transport::unix::UnixListener) -> UnixHosting {
319 let cancellation = self.cancellation().child_token();
320 let task_cancellation = cancellation.clone();
321 let node = self.clone();
322 let task = tokio::spawn(async move {
323 let mut connections = JoinSet::new();
324 loop {
325 tokio::select! {
326 biased;
327 () = task_cancellation.cancelled() => break,
328 result = connections.join_next(), if !connections.is_empty() => {
329 result.expect("connection task exists").map_err(|error| WsError::Connect(error.to_string()))?;
330 }
331 pipe = accept_unix(&listener) => {
332 let pipe = pipe?;
333 let connection = tokio::select! {
334 biased;
335 () = task_cancellation.cancelled() => break,
336 connection = node.serve_transport(pipe) => connection,
337 };
338 let cancellation = task_cancellation.clone();
339 connections.spawn(async move {
340 tokio::select! {
341 biased;
342 () = cancellation.cancelled() => {
343 connection.shutdown();
344 connection.closed().await;
345 }
346 () = connection.closed() => {}
347 }
348 });
349 }
350 }
351 }
352 while let Some(result) = connections.join_next().await {
353 result.map_err(|error| WsError::Connect(error.to_string()))?;
354 }
355 Ok(())
356 });
357 UnixHosting { cancellation, task }
358 }
359}