1use std::path::Path;
20
21use serde::{Deserialize, Serialize};
22
23const CURRENT_VERSION: u32 = 1;
25
26pub const OWNER_BASE: &str = "base";
28
29pub const OWNER_SHARED_NAT: &str = "shared-nat";
33
34#[must_use]
36pub fn owner_for_service(service: &str) -> String {
37 format!("service:{service}")
38}
39
40#[must_use]
48pub fn owner_for_isolation_network(net: &str) -> String {
49 format!("iso:{net}")
50}
51
52#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
54pub struct ManagedNetwork {
55 pub owner: String,
59 pub kind: String,
62 pub name: String,
64 pub id: String,
66 pub subnet: String,
68 #[serde(default, skip_serializing_if = "Option::is_none")]
70 pub wg_port: Option<u16>,
71 #[serde(default, skip_serializing_if = "Option::is_none")]
75 pub wg_private_key: Option<String>,
76 #[serde(default, skip_serializing_if = "Option::is_none")]
78 pub wg_public_key: Option<String>,
79 #[serde(default, skip_serializing_if = "Option::is_none")]
81 pub interface: Option<String>,
82}
83
84#[derive(Debug, Clone, Serialize, Deserialize)]
86pub struct NetworkState {
87 #[serde(default = "default_version")]
89 pub version: u32,
90 #[serde(default)]
92 pub networks: Vec<ManagedNetwork>,
93}
94
95fn default_version() -> u32 {
96 CURRENT_VERSION
97}
98
99impl Default for NetworkState {
100 fn default() -> Self {
101 Self {
102 version: CURRENT_VERSION,
103 networks: Vec::new(),
104 }
105 }
106}
107
108impl NetworkState {
109 #[must_use]
113 pub fn load(path: &Path) -> Self {
114 match std::fs::read(path) {
115 Ok(bytes) => serde_json::from_slice(&bytes).unwrap_or_default(),
116 Err(_) => Self::default(),
117 }
118 }
119
120 pub fn save(&self, path: &Path) -> std::io::Result<()> {
128 if let Some(parent) = path.parent() {
129 std::fs::create_dir_all(parent)?;
130 }
131 let json = serde_json::to_vec_pretty(self).map_err(std::io::Error::other)?;
132 let tmp = path.with_extension("json.tmp");
133 std::fs::write(&tmp, &json)?;
134 std::fs::rename(&tmp, path)?;
135 Ok(())
136 }
137
138 #[must_use]
140 pub fn get(&self, owner: &str) -> Option<&ManagedNetwork> {
141 self.networks.iter().find(|n| n.owner == owner)
142 }
143
144 pub fn upsert(&mut self, net: ManagedNetwork) {
146 if let Some(existing) = self.networks.iter_mut().find(|n| n.owner == net.owner) {
147 *existing = net;
148 } else {
149 self.networks.push(net);
150 }
151 }
152
153 pub fn remove(&mut self, owner: &str) -> Option<ManagedNetwork> {
155 self.networks
156 .iter()
157 .position(|n| n.owner == owner)
158 .map(|pos| self.networks.remove(pos))
159 }
160}
161
162pub const DEDICATED_PORT_BAND: u16 = 256;
168
169#[derive(Debug, Clone)]
182pub struct DedicatedPortAllocator {
183 base: u16,
184 used: std::collections::BTreeSet<u16>,
185}
186
187impl DedicatedPortAllocator {
188 pub fn new(base: u16, in_use: impl IntoIterator<Item = u16>) -> Self {
195 Self {
196 base,
197 used: in_use.into_iter().collect(),
198 }
199 }
200
201 fn band_start(&self) -> u16 {
203 self.base.saturating_add(1)
204 }
205
206 fn band_end(&self) -> u16 {
208 self.base.saturating_add(DEDICATED_PORT_BAND)
209 }
210
211 pub fn allocate(&mut self) -> crate::error::Result<u16> {
217 for port in self.band_start()..=self.band_end() {
218 if !self.used.contains(&port) {
219 self.used.insert(port);
220 return Ok(port);
221 }
222 }
223 Err(crate::error::OverlaydError::Other(format!(
224 "dedicated-overlay port band exhausted ({}..={}, {} ports)",
225 self.band_start(),
226 self.band_end(),
227 DEDICATED_PORT_BAND
228 )))
229 }
230
231 pub fn release(&mut self, port: u16) {
233 self.used.remove(&port);
234 }
235
236 pub fn reserve(&mut self, port: u16) {
239 self.used.insert(port);
240 }
241
242 #[must_use]
244 pub fn is_used(&self, port: u16) -> bool {
245 self.used.contains(&port)
246 }
247}
248
249#[cfg(test)]
250mod tests {
251 use super::*;
252
253 fn sample(owner: &str, id: &str) -> ManagedNetwork {
254 ManagedNetwork {
255 owner: owner.to_string(),
256 kind: "hcn-internal".to_string(),
257 name: "zlayer-overlay".to_string(),
258 id: id.to_string(),
259 subnet: "10.200.0.0/28".to_string(),
260 wg_port: None,
261 wg_private_key: None,
262 wg_public_key: None,
263 interface: None,
264 }
265 }
266
267 #[test]
268 fn upsert_replaces_same_owner_and_get_finds_it() {
269 let mut st = NetworkState::default();
270 st.upsert(sample(OWNER_BASE, "guid-1"));
271 st.upsert(sample(OWNER_BASE, "guid-2")); assert_eq!(st.networks.len(), 1);
273 assert_eq!(st.get(OWNER_BASE).unwrap().id, "guid-2");
274 }
275
276 #[test]
277 fn distinct_owners_coexist_and_remove_targets_one() {
278 let mut st = NetworkState::default();
279 st.upsert(sample(OWNER_BASE, "base-guid"));
280 st.upsert(sample(&owner_for_service("web"), "web-guid"));
281 assert_eq!(st.networks.len(), 2);
282
283 let removed = st.remove(OWNER_BASE).expect("base entry present");
284 assert_eq!(removed.id, "base-guid");
285 assert_eq!(st.networks.len(), 1);
286 assert!(st.get(OWNER_BASE).is_none());
287 assert_eq!(st.get(&owner_for_service("web")).unwrap().id, "web-guid");
288 assert!(st.remove("service:nope").is_none());
289 }
290
291 #[test]
292 fn save_then_load_roundtrips() {
293 let dir = std::env::temp_dir().join(format!("zlayer-netstate-test-{}", std::process::id()));
294 let path = dir.join("agent_network.json");
295 let _ = std::fs::remove_dir_all(&dir);
296
297 let mut st = NetworkState::default();
298 st.upsert(sample(OWNER_BASE, "guid-rt"));
299 st.save(&path).expect("save must succeed");
300
301 let loaded = NetworkState::load(&path);
302 assert_eq!(loaded.version, CURRENT_VERSION);
303 assert_eq!(loaded.networks, st.networks);
304
305 let _ = std::fs::remove_dir_all(&dir);
306 }
307
308 #[test]
309 fn load_missing_file_is_empty_default() {
310 let path = std::env::temp_dir().join("zlayer-netstate-does-not-exist-xyz.json");
311 let _ = std::fs::remove_file(&path);
312 let st = NetworkState::load(&path);
313 assert_eq!(st.version, CURRENT_VERSION);
314 assert!(st.networks.is_empty());
315 }
316
317 #[test]
318 fn dedicated_fields_survive_save_load_roundtrip() {
319 let dir = std::env::temp_dir().join(format!("zlayer-netstate-ded-{}", std::process::id()));
320 let path = dir.join("agent_network.json");
321 let _ = std::fs::remove_dir_all(&dir);
322
323 let mut net = sample(&owner_for_service("web"), "ded-guid");
324 net.wg_port = Some(51823);
325 net.wg_private_key = Some("cHJpdmF0ZS1rZXktYjY0".to_string());
326 net.wg_public_key = Some("cHVibGljLWtleS1iNjQ=".to_string());
327 net.interface = Some("zl-web0".to_string());
328
329 let mut st = NetworkState::default();
330 st.upsert(net.clone());
331 st.save(&path).expect("save must succeed");
332
333 let loaded = NetworkState::load(&path);
334 let got = loaded
335 .get(&owner_for_service("web"))
336 .expect("service entry present");
337 assert_eq!(got.wg_port, Some(51823));
338 assert_eq!(got.wg_private_key.as_deref(), Some("cHJpdmF0ZS1rZXktYjY0"));
339 assert_eq!(got.wg_public_key.as_deref(), Some("cHVibGljLWtleS1iNjQ="));
340 assert_eq!(got.interface.as_deref(), Some("zl-web0"));
341 assert_eq!(got, &net);
342
343 let _ = std::fs::remove_dir_all(&dir);
344 }
345
346 #[test]
347 fn older_marker_without_dedicated_fields_still_loads() {
348 let dir = std::env::temp_dir().join(format!("zlayer-netstate-bc-{}", std::process::id()));
351 let path = dir.join("agent_network.json");
352 let _ = std::fs::remove_dir_all(&dir);
353 std::fs::create_dir_all(&dir).expect("mkdir");
354
355 let legacy = r#"{
356 "version": 1,
357 "networks": [
358 {
359 "owner": "base",
360 "kind": "hcn-internal",
361 "name": "zlayer-overlay",
362 "id": "legacy-guid",
363 "subnet": "10.200.0.0/28"
364 }
365 ]
366 }"#;
367 std::fs::write(&path, legacy).expect("write legacy marker");
368
369 let loaded = NetworkState::load(&path);
370 let got = loaded.get(OWNER_BASE).expect("base entry present");
371 assert_eq!(got.id, "legacy-guid");
372 assert_eq!(got.wg_port, None);
373 assert_eq!(got.wg_private_key, None);
374 assert_eq!(got.wg_public_key, None);
375 assert_eq!(got.interface, None);
376
377 let _ = std::fs::remove_dir_all(&dir);
378 }
379
380 #[test]
381 fn allocate_returns_distinct_ascending_ports() {
382 let mut alloc = DedicatedPortAllocator::new(51820, std::iter::empty());
383 let a = alloc.allocate().expect("port a");
384 let b = alloc.allocate().expect("port b");
385 let c = alloc.allocate().expect("port c");
386 assert_eq!(a, 51821);
387 assert_eq!(b, 51822);
388 assert_eq!(c, 51823);
389 }
390
391 #[test]
392 fn release_then_allocate_reuses_freed_port() {
393 let mut alloc = DedicatedPortAllocator::new(51820, std::iter::empty());
394 let a = alloc.allocate().expect("port a");
395 let b = alloc.allocate().expect("port b");
396 assert_eq!(a, 51821);
397 assert_eq!(b, 51822);
398
399 alloc.release(a);
400 let reused = alloc.allocate().expect("reused port");
402 assert_eq!(reused, 51821);
403 }
404
405 #[test]
406 fn reserved_port_is_skipped_by_allocate() {
407 let mut alloc = DedicatedPortAllocator::new(51820, [51821]);
409 assert!(alloc.is_used(51821));
410 let first = alloc.allocate().expect("first allocation");
411 assert_eq!(first, 51822);
412
413 alloc.reserve(51823);
415 let next = alloc.allocate().expect("next allocation");
416 assert_eq!(next, 51824);
417 }
418
419 #[test]
420 fn band_exhaustion_errors() {
421 let base = 51820u16;
423 let full: Vec<u16> = (base + 1..=base + DEDICATED_PORT_BAND).collect();
424 let mut alloc = DedicatedPortAllocator::new(base, full);
425 let err = alloc.allocate().expect_err("band must be exhausted");
426 assert!(
427 matches!(err, crate::error::OverlaydError::Other(ref m) if m.contains("exhausted")),
428 "unexpected error: {err:?}"
429 );
430 }
431}