Merge pull request #633 from worldcoin/Dzejkop/stronger-docker-utils

Strengthen postgres-docker-utils

crates/postgres-docker-utils/src/lib.rs

@@ -1,16 +1,21 @@
+use std::net::SocketAddr;
 use std::process::{Command, Stdio};
 use std::time::Duration;
 
 use anyhow::Context;
 
 pub struct DockerContainerGuard {
-    container_id:   String,
-    container_port: u16,
+    container_id: String,
+    socket_addr:  SocketAddr,
 }
 
 impl DockerContainerGuard {
-    pub fn port(&self) -> u16 {
-        self.container_port
+    pub fn socket_addr(&self) -> SocketAddr {
+        self.socket_addr
+    }
+
+    pub fn address(&self) -> String {
+        self.socket_addr.to_string()
     }
 }
 
@@ -40,13 +45,13 @@ pub async fn setup() -> anyhow::Result<DockerContainerGuard> {
 
     let exposed_port = run_cmd_to_output(&format!("docker container port {container_id} 5432"))
         .context("Fetching container exposed port")?;
-    let container_port = parse_exposed_port(&exposed_port)?;
+    let socket_addr = parse_exposed_port(&exposed_port)?;
 
     std::thread::sleep(Duration::from_secs_f32(2.0));
 
     Ok(DockerContainerGuard {
         container_id,
-        container_port,
+        socket_addr,
     })
 }
 
@@ -76,38 +81,18 @@ fn run_cmd(cmd_str: &str) -> anyhow::Result<()> {
     Ok(())
 }
 
-fn parse_exposed_port(s: &str) -> anyhow::Result<u16> {
+fn parse_exposed_port(s: &str) -> anyhow::Result<SocketAddr> {
     let parts: Vec<_> = s
         .split_whitespace()
         .map(|s| s.trim())
         .filter(|s| !s.is_empty())
         .collect();
 
-    let ports: Vec<_> = parts.into_iter().filter_map(extract_port).collect();
-
-    let mut parsed_port = None;
-
-    for port in ports {
-        let port: u16 = port.parse().with_context(|| format!("Parsing `{port}`"))?;
-
-        if let Some(current) = parsed_port {
-            if current != port {
-                anyhow::bail!(
-                    "Multiple different ports exposed: `{}` and `{}`",
-                    current,
-                    port
-                );
-            }
-        } else {
-            parsed_port = Some(port);
-        }
-    }
-
-    parsed_port.context("No ports parsed")
-}
-
-fn extract_port(s: &str) -> Option<&str> {
-    s.split(':').last()
+    Ok(parts
+        .iter()
+        .map(|p| p.parse::<SocketAddr>())
+        .next()
+        .context("Missing exposed port")??)
 }
 
 #[cfg(test)]
@@ -120,14 +105,8 @@ mod tests {
     #[test_case("   0.0.0.0:55837    " => 55837 ; "ignore whitespace")]
     #[test_case("[::]:12345" => 12345 ; "works with ipv6")]
     #[test_case("0.0.0.0:12345 \n [::]:12345" => 12345 ; "works with multiple ips")]
+    #[test_case("0.0.0.0:12345 \n [::]:54321" => 12345 ; "yields first of multiple ips")]
     fn test_parse_exposed_port(s: &str) -> u16 {
-        parse_exposed_port(s).unwrap()
-    }
-
-    #[test]
-    fn different_ports_result_in_failure() {
-        const S: &str = "0.0.0.0:12345 [::]:54321";
-
-        let _err = parse_exposed_port(S).unwrap_err();
+        parse_exposed_port(s).unwrap().port()
     }
 }

src/database/mod.rs

@@ -369,8 +369,8 @@ impl Database {
     pub async fn get_latest_insertion_timestamp(&self) -> Result<Option<DateTime<Utc>>, Error> {
         let query = sqlx::query(
             r#"
-            SELECT insertion_timestamp 
-            FROM latest_insertion_timestamp 
+            SELECT insertion_timestamp
+            FROM latest_insertion_timestamp
             WHERE Lock = 'X';"#,
         );
 
@@ -835,8 +835,8 @@ mod test {
     // TODO: either use anyhow or eyre
     async fn setup_db() -> anyhow::Result<(Database, DockerContainerGuard)> {
         let db_container = postgres_docker_utils::setup().await?;
-        let port = db_container.port();
-        let url = format!("postgres://postgres:postgres@localhost:{port}/database");
+        let db_socket_addr = db_container.address();
+        let url = format!("postgres://postgres:postgres@{db_socket_addr}/database");
 
         let db = Database::new(Options {
             database:                 SecretUrl::from_str(&url)?,

tests/delete_identities.rs

@@ -35,8 +35,8 @@ async fn delete_identities() -> anyhow::Result<()> {
     let mock_insertion_prover = &insertion_prover_map[&insertion_batch_size];
     let mock_deletion_prover = &deletion_prover_map[&deletion_batch_size];
 
-    let port = db_container.port();
-    let db_url = format!("postgres://postgres:postgres@localhost:{port}/database");
+    let db_socket_addr = db_container.address();
+    let db_url = format!("postgres://postgres:postgres@{db_socket_addr}/database");
 
     let mut options = Options::try_parse_from([
         "signup-sequencer",

tests/delete_padded_identity.rs

@@ -35,8 +35,8 @@ async fn delete_padded_identity() -> anyhow::Result<()> {
     let mock_insertion_prover = &insertion_prover_map[&insertion_batch_size];
     let mock_deletion_prover = &deletion_prover_map[&deletion_batch_size];
 
-    let port = db_container.port();
-    let db_url = format!("postgres://postgres:postgres@localhost:{port}/database");
+    let db_socket_addr = db_container.address();
+    let db_url = format!("postgres://postgres:postgres@{db_socket_addr}/database");
 
     let mut options = Options::try_parse_from([
         "signup-sequencer",

tests/dynamic_batch_sizes.rs

@@ -36,8 +36,8 @@ async fn dynamic_batch_sizes() -> anyhow::Result<()> {
     let first_prover = &insertion_prover_map[&first_batch_size];
     let second_prover = &insertion_prover_map[&second_batch_size];
 
-    let port = db_container.port();
-    let db_url = format!("postgres://postgres:postgres@localhost:{port}/database");
+    let db_socket_addr = db_container.address();
+    let db_url = format!("postgres://postgres:postgres@{db_socket_addr}/database");
 
     // We initially spawn the service with a single prover for batch size 3.
 

tests/insert_identity_and_proofs.rs

@@ -23,8 +23,8 @@ async fn insert_identity_and_proofs() -> anyhow::Result<()> {
 
     let prover_mock = &insertion_prover_map[&batch_size];
 
-    let port = db_container.port();
-    let db_url = format!("postgres://postgres:postgres@localhost:{port}/database");
+    let db_socket_addr = db_container.address();
+    let db_url = format!("postgres://postgres:postgres@{db_socket_addr}/database");
 
     let mut options = Options::try_parse_from([
         "signup-sequencer",

tests/malformed_payload.rs

@@ -22,8 +22,8 @@ async fn malformed_payload() -> anyhow::Result<()> {
 
     let prover_mock = &insertion_prover_map[&batch_size];
 
-    let port = db_container.port();
-    let db_url = format!("postgres://postgres:postgres@localhost:{port}/database");
+    let db_socket_addr = db_container.address();
+    let db_url = format!("postgres://postgres:postgres@{db_socket_addr}/database");
     let mut options = Options::try_parse_from([
         "signup-sequencer",
         "--identity-manager-address",

tests/multi_prover.rs

@@ -38,8 +38,8 @@ async fn multi_prover() -> anyhow::Result<()> {
 
     info!("Running with {prover_arg_string}");
 
-    let port = db_container.port();
-    let db_url = format!("postgres://postgres:postgres@localhost:{port}/database");
+    let db_socket_addr = db_container.address();
+    let db_url = format!("postgres://postgres:postgres@{db_socket_addr}/database");
     let mut options = Options::try_parse_from([
         "signup-sequencer",
         "--identity-manager-address",

tests/recover_identities.rs

@@ -44,8 +44,8 @@ async fn recover_identities() -> anyhow::Result<()> {
     let mock_insertion_prover = &insertion_prover_map[&insertion_batch_size];
     let mock_deletion_prover = &deletion_prover_map[&deletion_batch_size];
 
-    let port = db_container.port();
-    let db_url = format!("postgres://postgres:postgres@localhost:{port}/database");
+    let db_socket_addr = db_container.address();
+    let db_url = format!("postgres://postgres:postgres@{db_socket_addr}/database");
 
     let mut options = Options::try_parse_from([
         "signup-sequencer",

tests/unavailable_prover.rs

@@ -23,8 +23,8 @@ async fn unavailable_prover() -> anyhow::Result<()> {
 
     prover_mock.set_availability(false).await;
 
-    let port = db_container.port();
-    let db_url = format!("postgres://postgres:postgres@localhost:{port}/database");
+    let db_socket_addr = db_container.address();
+    let db_url = format!("postgres://postgres:postgres@{db_socket_addr}/database");
     let mut options = Options::try_parse_from([
         "signup-sequencer",
         "--identity-manager-address",

tests/unreduced_identity.rs

@@ -16,8 +16,8 @@ async fn test_unreduced_identity() -> anyhow::Result<()> {
     let prover_mock = &insertion_prover_map[&batch_size];
     prover_mock.set_availability(false).await;
 
-    let port = db_container.port();
-    let db_url = format!("postgres://postgres:postgres@localhost:{port}/database");
+    let db_socket_addr = db_container.address();
+    let db_url = format!("postgres://postgres:postgres@{db_socket_addr}/database");
     let mut options = Options::try_parse_from([
         "signup-sequencer",
         "--identity-manager-address",

tests/validate_proof_with_age.rs

@@ -28,8 +28,8 @@ async fn validate_proof_with_age() -> anyhow::Result<()> {
 
     let prover_mock = &insertion_prover_map[&batch_size];
 
-    let port = db_container.port();
-    let db_url = format!("postgres://postgres:postgres@localhost:{port}/database");
+    let db_socket_addr = db_container.address();
+    let db_url = format!("postgres://postgres:postgres@{db_socket_addr}/database");
 
     let mut options = Options::try_parse_from([
         "signup-sequencer",

tests/validate_proofs.rs

@@ -26,8 +26,8 @@ async fn validate_proofs() -> anyhow::Result<()> {
 
     let identity_manager = mock_chain.identity_manager.clone();
 
-    let port = db_container.port();
-    let db_url = format!("postgres://postgres:postgres@localhost:{port}/database");
+    let db_socket_addr = db_container.address();
+    let db_url = format!("postgres://postgres:postgres@{db_socket_addr}/database");
 
     let mut options = Options::try_parse_from([
         "signup-sequencer",