Add udp and sctp support to port mappings

src/clients/cli.rs

@@ -1,5 +1,6 @@
 use crate::core::{
-    env, env::GetEnvValue, logs::LogStream, ports::Ports, Container, Docker, Image, RunArgs,
+    env, env::GetEnvValue, logs::LogStream, ports::Ports, Container, Docker, Image, PortMapping,
+    RunArgs,
 };
 use shiplift::rep::ContainerDetails;
 use std::{
@@ -155,9 +156,14 @@ impl Client {
 
         if let Some(ports) = run_args.ports() {
             for port in &ports {
+                let (local, internal, protocol) = match port {
+                    PortMapping::Tcp { local, internal } => (local, internal, "tcp"),
+                    PortMapping::Udp { local, internal } => (local, internal, "udp"),
+                    PortMapping::Sctp { local, internal } => (local, internal, "sctp"),
+                };
                 command
                     .arg("-p")
-                    .arg(format!("{}:{}", port.local, port.internal));
+                    .arg(format!("{}:{}/{}", local, internal, protocol));
             }
         } else {
             command.arg("-P"); // expose all ports
@@ -474,8 +480,8 @@ mod tests {
 
         assert!(format!("{:?}", command).contains(r#"-d"#));
         assert!(!format!("{:?}", command).contains(r#"-P"#));
-        assert!(format!("{:?}", command).contains(r#""-p" "123:456""#));
-        assert!(format!("{:?}", command).contains(r#""-p" "555:888""#));
+        assert!(format!("{:?}", command).contains(r#""-p" "123:456/tcp""#));
+        assert!(format!("{:?}", command).contains(r#""-p" "555:888/tcp""#));
     }
 
     #[test]

src/clients/http.rs

@@ -1,5 +1,7 @@
 use crate::{
-    core::{env, logs::LogStreamAsync, ports::Ports, ContainerAsync, DockerAsync, RunArgs},
+    core::{
+        env, logs::LogStreamAsync, ports::Ports, ContainerAsync, DockerAsync, PortMapping, RunArgs,
+    },
     Image,
 };
 use async_trait::async_trait;
@@ -100,10 +102,14 @@ impl Http {
 
         // ports
         if let Some(ports) = run_args.ports() {
-            // TODO support UDP?
             for port in &ports {
+                let (local, internal, protocol) = match port {
+                    PortMapping::Tcp { local, internal } => (local, internal, "tcp"),
+                    PortMapping::Udp { local, internal } => (local, internal, "udp"),
+                    PortMapping::Sctp { local, internal } => (local, internal, "sctp"),
+                };
                 // casting u16 to u32
-                options_builder.expose(port.internal as u32, "tcp", port.local as u32);
+                options_builder.expose(*internal as u32, protocol, *local as u32);
             }
         } else {
             options_builder.publish_all_ports();

src/core.rs

@@ -3,7 +3,7 @@ pub use self::{
     container::Container,
     container_async::ContainerAsync,
     docker::RunArgs,
-    image::{Image, Port, WaitFor},
+    image::{Image, PortMapping, WaitFor},
 };
 
 mod container;

src/core/container.rs

@@ -1,5 +1,10 @@
 use crate::{
-    core::{docker::Docker, env::Command, image::WaitFor},
+    core::{
+        docker::Docker,
+        env::Command,
+        image::WaitFor,
+        ports::{MapToHostPort, Ports},
+    },
     Image,
 };
 use std::{fmt, marker::PhantomData};
@@ -128,13 +133,17 @@ impl<'d, I> Container<'d, I> {
     /// This method panics if the given port is not mapped.
     /// Testcontainers is designed to be used in tests only. If a certain port is not mapped, the container
     /// is unlikely to be useful.
-    pub fn get_host_port(&self, internal_port: u16) -> u16 {
+    pub fn get_host_port<T>(&self, internal_port: T) -> T
+    where
+        T: fmt::Debug,
+        Ports: MapToHostPort<T>,
+    {
         self.docker_client
             .ports(&self.id)
-            .map_to_host_port(internal_port)
+            .map_to_host_port(&internal_port)
             .unwrap_or_else(|| {
                 panic!(
-                    "container {} does not expose port {}",
+                    "container {:?} does not expose port {:?}",
                     self.id, internal_port
                 )
             })

src/core/container_async.rs

@@ -1,5 +1,11 @@
 use crate::{
-    core::{env, env::Command, logs::LogStreamAsync, ports::Ports, WaitFor},
+    core::{
+        env,
+        env::Command,
+        logs::LogStreamAsync,
+        ports::{MapToHostPort, Ports},
+        WaitFor,
+    },
     Image,
 };
 use async_trait::async_trait;
@@ -55,14 +61,18 @@ impl<'d, I> ContainerAsync<'d, I> {
     /// This method panics if the given port is not mapped.
     /// Testcontainers is designed to be used in tests only. If a certain port is not mapped, the container
     /// is unlikely to be useful.
-    pub async fn get_host_port(&self, internal_port: u16) -> u16 {
+    pub async fn get_host_port<T>(&self, internal_port: T) -> T
+    where
+        T: fmt::Debug,
+        Ports: MapToHostPort<T>,
+    {
         self.docker_client
             .ports(&self.id)
             .await
-            .map_to_host_port(internal_port)
+            .map_to_host_port(&internal_port)
             .unwrap_or_else(|| {
                 panic!(
-                    "container {} does not expose port {}",
+                    "container {:?} does not expose port {:?}",
                     self.id, internal_port
                 )
             })

src/core/docker.rs

@@ -1,4 +1,4 @@
-use crate::core::{logs::LogStream, ports::Ports, Port};
+use crate::core::{logs::LogStream, ports::Ports, PortMapping};
 
 /// Container run command arguments.
 /// `name` - run image instance with the given name (should be explicitly set to be seen by other containers created in the same docker network).
@@ -8,7 +8,7 @@ use crate::core::{logs::LogStream, ports::Ports, Port};
 pub struct RunArgs {
     name: Option<String>,
     network: Option<String>,
-    ports: Option<Vec<Port>>,
+    ports: Option<Vec<PortMapping>>,
 }
 
 /// Defines operations that we need to perform on docker containers and other entities.
@@ -39,7 +39,7 @@ impl RunArgs {
         }
     }
 
-    pub fn with_mapped_port<P: Into<Port>>(mut self, port: P) -> Self {
+    pub fn with_mapped_port<P: Into<PortMapping>>(mut self, port: P) -> Self {
         let mut ports = self.ports.unwrap_or_default();
         ports.push(port.into());
         self.ports = Some(ports);
@@ -54,7 +54,7 @@ impl RunArgs {
         self.name.clone()
     }
 
-    pub(crate) fn ports(&self) -> Option<Vec<Port>> {
+    pub(crate) fn ports(&self) -> Option<Vec<PortMapping>> {
         self.ports.clone()
     }
 }

src/core/image.rs

@@ -100,9 +100,10 @@ where
 
 /// Represents a port mapping between a local port and the internal port of a container.
 #[derive(Clone, Debug, PartialEq)]
-pub struct Port {
-    pub local: u16,
-    pub internal: u16,
+pub enum PortMapping {
+    Tcp { local: u16, internal: u16 },
+    Udp { local: u16, internal: u16 },
+    Sctp { local: u16, internal: u16 },
 }
 
 /// Represents a condition that needs to be met before a container is considered ready.
@@ -157,8 +158,8 @@ impl WaitFor {
     }
 }
 
-impl From<(u16, u16)> for Port {
+impl From<(u16, u16)> for PortMapping {
     fn from((local, internal): (u16, u16)) -> Self {
-        Port { local, internal }
+        Self::Tcp { local, internal }
     }
 }

src/core/ports.rs

@@ -1,18 +1,33 @@
 use std::collections::HashMap;
 
+#[derive(Debug, PartialEq, Copy, Clone, Eq, Hash)]
+struct Sctp(u16);
+#[derive(Debug, PartialEq, Copy, Clone, Eq, Hash)]
+struct Tcp(u16);
+#[derive(Debug, PartialEq, Copy, Clone, Eq, Hash)]
+struct Udp(u16);
+
 /// The exposed ports of a running container.
 #[derive(Debug, PartialEq, Default)]
 pub struct Ports {
-    mapping: HashMap<u16, u16>,
+    tcp: HashMap<Tcp, Tcp>,
+    udp: HashMap<Udp, Udp>,
+    sctp: HashMap<Sctp, Sctp>,
 }
 
 impl Ports {
     pub fn new(ports: HashMap<String, Option<Vec<HashMap<String, String>>>>) -> Self {
-        let mapping = ports
+        let tcp = HashMap::new();
+        let udp = HashMap::new();
+        let sctp = HashMap::new();
+
+        ports
             .into_iter()
             .filter_map(|(internal, external)| {
-                // internal is '8332/tcp', split off the protocol ...
-                let internal = internal.split('/').next()?;
+                // internal is ')8332/tcp', split off the protocol ...
+                let mut iter = internal.split('/');
+                let internal = iter.next()?;
+                let protocol = iter.next()?.to_string();
 
                 // external is a an optional list of maps: [ { "HostIp": "0.0.0.0", "HostPort": "33078" } ]
                 // get the first entry and get the value of the `HostPort` field
@@ -21,18 +36,60 @@ impl Ports {
                 let internal = parse_port(internal);
                 let external = parse_port(&external);
 
-                log::debug!("Registering port mapping: {} -> {}", internal, external);
+                log::debug!(
+                    "Registering port mapping: {} -> {} / {}",
+                    internal,
+                    external,
+                    protocol
+                );
 
-                Some((internal, external))
+                Some((protocol, internal, external))
+            })
+            .fold(Self { tcp, udp, sctp }, |mut mappings, val| {
+                let (protocol, internal, external) = val;
+                match protocol.as_str() {
+                    "tcp" => {
+                        mappings.tcp.insert(Tcp(internal), Tcp(external));
+                    }
+                    "udp" => {
+                        mappings.udp.insert(Udp(internal), Udp(external));
+                    }
+                    "sctp" => {
+                        mappings.sctp.insert(Sctp(internal), Sctp(external));
+                    }
+                    _ => panic!("Not a valid port mapping."),
+                };
+                mappings
             })
-            .collect::<HashMap<_, _>>();
-
-        Self { mapping }
     }
+}
 
+pub trait MapToHostPort<T> {
     /// Returns the host port for the given internal port.
-    pub fn map_to_host_port(&self, internal_port: u16) -> Option<u16> {
-        self.mapping.get(&internal_port).cloned()
+    fn map_to_host_port(&self, internal_port: &T) -> Option<T>;
+}
+
+impl MapToHostPort<u16> for Ports {
+    fn map_to_host_port(&self, internal_port: &u16) -> Option<u16> {
+        self.map_to_host_port(&Tcp(*internal_port)).map(|p| p.0)
+    }
+}
+
+impl MapToHostPort<Tcp> for Ports {
+    fn map_to_host_port(&self, internal_port: &Tcp) -> Option<Tcp> {
+        self.tcp.get(internal_port).copied()
+    }
+}
+
+impl MapToHostPort<Udp> for Ports {
+    fn map_to_host_port(&self, internal_port: &Udp) -> Option<Udp> {
+        self.udp.get(internal_port).copied()
+    }
+}
+
+impl MapToHostPort<Sctp> for Ports {
+    fn map_to_host_port(&self, internal_port: &Sctp) -> Option<Sctp> {
+        self.sctp.get(internal_port).copied()
     }
 }
 
@@ -281,10 +338,10 @@ mod tests {
             .unwrap_or_default();
 
         let mut expected_ports = Ports::default();
-        expected_ports.mapping.insert(18332, 33076);
-        expected_ports.mapping.insert(18333, 33075);
-        expected_ports.mapping.insert(8332, 33078);
-        expected_ports.mapping.insert(8333, 33077);
+        expected_ports.tcp.insert(Tcp(18332), Tcp(33076));
+        expected_ports.tcp.insert(Tcp(18333), Tcp(33075));
+        expected_ports.tcp.insert(Tcp(8332), Tcp(33078));
+        expected_ports.tcp.insert(Tcp(8333), Tcp(33077));
 
         assert_eq!(parsed_ports, expected_ports)
     }