microsandbox_server/
port.rs

1//! Port management for the microsandbox server.
2//!
3//! This module handles port assignment and management for sandboxes:
4//! - Assigns truly available ports obtained from the OS
5//! - Tracks assigned ports for fast lookup
6//! - Persists port assignments to disk
7//! - Loads existing port assignments on startup
8//! - Handles port uniqueness with bidirectional mapping
9//!
10//! The module provides:
11//! - Port manager for tracking assigned ports
12//! - Functions for assigning and releasing ports
13//! - File-based persistence of port assignments
14
15use microsandbox_utils::PORTAL_PORTS_FILE;
16use once_cell::sync::Lazy;
17use serde::{Deserialize, Serialize};
18use std::{
19    collections::HashMap,
20    net::{IpAddr, Ipv4Addr, SocketAddr, TcpListener},
21    path::{Path, PathBuf},
22};
23use tokio::{fs, sync::Mutex};
24use tracing::{debug, info, warn};
25
26use crate::{MicrosandboxServerError, MicrosandboxServerResult};
27
28//--------------------------------------------------------------------------------------------------
29// Constants
30//--------------------------------------------------------------------------------------------------
31
32/// The localhost IP address used for all portal connections
33pub const LOCALHOST_IP: IpAddr = IpAddr::V4(Ipv4Addr::LOCALHOST);
34
35/// Lock to ensure only one thread gets a port at a time
36static PORT_ASSIGNMENT_LOCK: Lazy<Mutex<()>> = Lazy::new(|| Mutex::new(()));
37
38//--------------------------------------------------------------------------------------------------
39// Types
40//--------------------------------------------------------------------------------------------------
41
42/// Port mapping for sandbox instances - bidirectional for fast lookups
43#[derive(Debug, Clone, Default)]
44pub struct BiPortMapping {
45    /// Maps sandbox identifiers (namespace/sandboxname) to assigned port numbers
46    sandbox_to_port: HashMap<String, u16>,
47
48    /// Maps port numbers to sandbox identifiers for fast reverse lookup
49    port_to_sandbox: HashMap<u16, String>,
50}
51
52/// Serializable version of the port mapping for file storage
53#[derive(Debug, Clone, Serialize, Deserialize, Default)]
54pub struct PortMapping {
55    /// Maps sandbox identifiers (namespace/sandboxname) to assigned port numbers
56    pub mappings: HashMap<String, u16>,
57}
58
59/// Port manager for handling sandbox port assignments
60#[derive(Debug)]
61pub struct PortManager {
62    /// The port mappings data
63    mappings: BiPortMapping,
64
65    /// Path to the port mappings file
66    file_path: PathBuf,
67}
68
69//--------------------------------------------------------------------------------------------------
70// Methods
71//--------------------------------------------------------------------------------------------------
72
73impl BiPortMapping {
74    /// Create a new bidirectional port mapping
75    pub fn new() -> Self {
76        Self {
77            sandbox_to_port: HashMap::new(),
78            port_to_sandbox: HashMap::new(),
79        }
80    }
81
82    /// Insert a mapping, handling any existing mappings for the port or sandbox
83    pub fn insert(&mut self, sandbox_key: String, port: u16) {
84        // Check if this port is already assigned to a different sandbox
85        if let Some(existing_sandbox) = self.port_to_sandbox.get(&port) {
86            if existing_sandbox != &sandbox_key {
87                // Port is already assigned to a different sandbox - remove that mapping
88                warn!(
89                    "Port {} was already assigned to sandbox {}, reassigning to {}",
90                    port, existing_sandbox, sandbox_key
91                );
92                self.sandbox_to_port.remove(existing_sandbox);
93            }
94        }
95
96        // Check if this sandbox already has a different port
97        if let Some(existing_port) = self.sandbox_to_port.get(&sandbox_key) {
98            if *existing_port != port {
99                // Sandbox had a different port - remove that mapping
100                self.port_to_sandbox.remove(existing_port);
101            }
102        }
103
104        // Insert the new mapping in both directions
105        self.sandbox_to_port.insert(sandbox_key.clone(), port);
106        self.port_to_sandbox.insert(port, sandbox_key);
107    }
108
109    /// Remove a mapping by sandbox key
110    pub fn remove_by_sandbox(&mut self, sandbox_key: &str) -> Option<u16> {
111        if let Some(port) = self.sandbox_to_port.remove(sandbox_key) {
112            self.port_to_sandbox.remove(&port);
113            Some(port)
114        } else {
115            None
116        }
117    }
118
119    /// Remove a mapping by port
120    pub fn remove_by_port(&mut self, port: u16) -> Option<String> {
121        if let Some(sandbox_key) = self.port_to_sandbox.remove(&port) {
122            self.sandbox_to_port.remove(&sandbox_key);
123            Some(sandbox_key)
124        } else {
125            None
126        }
127    }
128
129    /// Get port by sandbox key
130    pub fn get_port(&self, sandbox_key: &str) -> Option<u16> {
131        self.sandbox_to_port.get(sandbox_key).copied()
132    }
133
134    /// Get sandbox key by port
135    pub fn get_sandbox(&self, port: u16) -> Option<&String> {
136        self.port_to_sandbox.get(&port)
137    }
138
139    /// Convert to serializable format
140    pub fn to_port_mapping(&self) -> PortMapping {
141        PortMapping {
142            mappings: self.sandbox_to_port.clone(),
143        }
144    }
145
146    /// Load from serializable format
147    pub fn from_port_mapping(mapping: PortMapping) -> Self {
148        let mut result = Self::new();
149
150        for (sandbox_key, port) in mapping.mappings {
151            result.insert(sandbox_key, port);
152        }
153
154        result
155    }
156}
157
158impl PortManager {
159    /// Create a new port manager
160    pub async fn new(namespace_dir: impl AsRef<Path>) -> MicrosandboxServerResult<Self> {
161        let file_path = namespace_dir.as_ref().join(PORTAL_PORTS_FILE);
162        let mappings = Self::load_mappings(&file_path).await?;
163
164        Ok(Self {
165            mappings,
166            file_path,
167        })
168    }
169
170    /// Load port mappings from file
171    async fn load_mappings(file_path: &Path) -> MicrosandboxServerResult<BiPortMapping> {
172        if file_path.exists() {
173            let contents = fs::read_to_string(file_path).await.map_err(|e| {
174                MicrosandboxServerError::ConfigError(format!(
175                    "Failed to read port mappings file: {}",
176                    e
177                ))
178            })?;
179
180            let port_mapping: PortMapping = serde_json::from_str(&contents).map_err(|e| {
181                MicrosandboxServerError::ConfigError(format!(
182                    "Failed to parse port mappings file: {}",
183                    e
184                ))
185            })?;
186
187            Ok(BiPortMapping::from_port_mapping(port_mapping))
188        } else {
189            debug!("No port mappings file found, creating a new one");
190            Ok(BiPortMapping::new())
191        }
192    }
193
194    /// Save port mappings to file
195    async fn save_mappings(&self) -> MicrosandboxServerResult<()> {
196        let port_mapping = self.mappings.to_port_mapping();
197        let contents = serde_json::to_string_pretty(&port_mapping).map_err(|e| {
198            MicrosandboxServerError::ConfigError(format!(
199                "Failed to serialize port mappings: {}",
200                e
201            ))
202        })?;
203
204        // Create parent directory if it doesn't exist
205        if let Some(parent) = self.file_path.parent() {
206            if !parent.exists() {
207                fs::create_dir_all(parent).await.map_err(|e| {
208                    MicrosandboxServerError::ConfigError(format!(
209                        "Failed to create directory for port mappings file: {}",
210                        e
211                    ))
212                })?;
213            }
214        }
215
216        fs::write(&self.file_path, contents).await.map_err(|e| {
217            MicrosandboxServerError::ConfigError(format!(
218                "Failed to write port mappings file: {}",
219                e
220            ))
221        })
222    }
223
224    /// Assign a port to a sandbox
225    pub async fn assign_port(&mut self, key: &str) -> MicrosandboxServerResult<u16> {
226        // Check if port is already assigned
227        if let Some(port) = self.mappings.get_port(key) {
228            // Verify this port is still available
229            if self.verify_port_availability(port) {
230                return Ok(port);
231            } else {
232                // Port is no longer available, so we need to assign a new one
233                warn!("Previously assigned port {} for sandbox {} is no longer available, reassigning", port, key);
234                self.mappings.remove_by_sandbox(key);
235            }
236        }
237
238        // Get a lock to ensure only one thread gets a port at a time
239        let _lock = PORT_ASSIGNMENT_LOCK.lock().await;
240
241        // Get a truly available port from the OS
242        let port = self.get_available_port_from_os()?;
243
244        // Save the mapping
245        self.mappings.insert(key.to_string(), port);
246        self.save_mappings().await?;
247
248        info!("Assigned port {} to sandbox {}", port, key);
249        Ok(port)
250    }
251
252    /// Release a port assignment
253    pub async fn release_port(&mut self, key: &str) -> MicrosandboxServerResult<()> {
254        if self.mappings.remove_by_sandbox(key).is_some() {
255            self.save_mappings().await?;
256            info!("Released port for sandbox {}", key);
257        }
258
259        Ok(())
260    }
261
262    /// Get a port for a sandbox if assigned
263    pub fn get_port(&self, key: &str) -> Option<u16> {
264        self.mappings.get_port(key)
265    }
266
267    /// Verify that a port is still available (not bound by something else)
268    fn verify_port_availability(&self, port: u16) -> bool {
269        let addr = SocketAddr::new(LOCALHOST_IP, port);
270        match TcpListener::bind(addr) {
271            Ok(_) => true,   // We could bind, so it's available
272            Err(_) => false, // We couldn't bind, so it's not available
273        }
274    }
275
276    /// Get an available port from the OS
277    fn get_available_port_from_os(&self) -> MicrosandboxServerResult<u16> {
278        // Bind to port 0 to let the OS assign an available port
279        let addr = SocketAddr::new(LOCALHOST_IP, 0);
280        let listener = TcpListener::bind(addr).map_err(|e| {
281            MicrosandboxServerError::ConfigError(format!(
282                "Failed to bind to address to get available port: {}",
283                e
284            ))
285        })?;
286
287        // Get the port assigned by the OS
288        let port = listener
289            .local_addr()
290            .map_err(|e| {
291                MicrosandboxServerError::ConfigError(format!(
292                    "Failed to get local address from socket: {}",
293                    e
294                ))
295            })?
296            .port();
297
298        debug!("OS assigned port {}", port);
299
300        // The listener will be dropped here, releasing the port
301        // We return the port value to be used by the caller
302
303        Ok(port)
304    }
305}
microsandbox_server/port.rs

microsandbox_server/
port.rs
