1 files changed, 51 insertions, 41 deletions

diff --git a/tests/server_shell.rs b/tests/server_shell.rs
index 3813013..c139233 100644
--- a/tests/server_shell.rs
+++ b/tests/server_shell.rs
@@ -1,14 +1,12 @@
 use anyhow::Result;
 use assert_cmd::{cargo::cargo_bin, Command};
 use get_port::{tcp::TcpPort, Ops, Range};
+use once_cell::sync::Lazy;
 use rexpect::session::{spawn_command, PtySession};
-use std::{io, path, sync::Once};
+use std::{io, path, sync::Mutex};
 use tempfile::TempDir;
 use thiserror::Error;
 
-static BUILD_DOCKER: Once = Once::new();
-static mut BUILD_DOCKER_RESULT: Result<(), DockerBuildError> = Ok(());
-
 struct TestContext {
     workdir: TempDir,
     ssh_port: u16,
@@ -37,53 +35,65 @@ impl From<io::Error> for DockerBuildError {
     }
 }
 
-fn build_docker_image() -> Result<(), DockerBuildError> {
-    let build_docker = || {
-        let mut command = std::process::Command::new("docker");
-
-        let absolute_shell_path = cargo_bin(env!("CARGO_PKG_NAME"));
-        let relative_shell_path = absolute_shell_path.strip_prefix(std::fs::canonicalize(".")?)?;
-        command.args([
-            "build",
-            "-t",
-            "shackle-server",
-            "--build-arg",
-            &format!("SHELL={}", relative_shell_path.display()),
-            "./",
-        ]);
-
-        let status = command.status()?;
-        if status.success() {
-            Ok(())
-        } else {
-            Err(DockerBuildError::CliError)
-        }
-    };
-
-    // Based on this example: https://doc.rust-lang.org/std/sync/struct.Once.html#examples-1
-    // Could be replaced by https://doc.rust-lang.org/std/cell/struct.LazyCell.html once that is in stable rust
-    unsafe {
-        BUILD_DOCKER.call_once(|| {
-            BUILD_DOCKER_RESULT = build_docker();
-        });
-        BUILD_DOCKER_RESULT.clone()
+static BUILD_DOCKER_RESULT: Lazy<Result<(), DockerBuildError>> = Lazy::new(|| {
+    let mut command = std::process::Command::new("docker");
+
+    let absolute_shell_path = cargo_bin(env!("CARGO_PKG_NAME"));
+    let relative_shell_path = absolute_shell_path.strip_prefix(std::fs::canonicalize(".")?)?;
+    command.args([
+        "build",
+        "-t",
+        "shackle-server",
+        "--build-arg",
+        &format!("SHELL={}", relative_shell_path.display()),
+        "./",
+    ]);
+
+    let status = command.status()?;
+    if status.success() {
+        Ok(())
+    } else {
+        Err(DockerBuildError::CliError)
     }
+});
+
+fn build_docker_image() -> Result<(), DockerBuildError> {
+    BUILD_DOCKER_RESULT.clone()
 }
 
-fn spawn_ssh_server() -> Result<TestContext> {
-    build_docker_image()?;
+#[derive(Error, Debug, Clone)]
+pub enum PortAssignmentError {
+    #[error("Mutex Error: `{0}`")]
+    MutexError(String),
+    #[error("Couldn't find an available port")]
+    NoMorePorts,
+}
 
-    let workdir = tempfile::tempdir()?;
+static LAST_USED_PORT: Lazy<Mutex<u16>> = Lazy::new(|| Mutex::new(2022));
+fn find_unused_port() -> Result<u16, PortAssignmentError> {
+    let mut last_used = LAST_USED_PORT
+        .lock()
+        .map_err(|e| PortAssignmentError::MutexError(e.to_string()))?;
 
-    // TODO: Put the used ports in an Arc<BTreeSet>
-    let ssh_port = TcpPort::in_range(
+    let port = TcpPort::in_range(
         "127.0.0.1",
         Range {
-            min: 2022,
+            min: *last_used + 1,
             max: 3022,
         },
     )
-    .unwrap();
+    .ok_or(PortAssignmentError::NoMorePorts)?;
+
+    *last_used = port;
+
+    Ok(port)
+}
+
+fn spawn_ssh_server() -> Result<TestContext> {
+    build_docker_image()?;
+
+    let workdir = tempfile::tempdir()?;
+    let ssh_port = find_unused_port()?;
 
     let mut command = std::process::Command::new("docker");
     command.args([