From 483a1e1dfb25e9a06750b3846f03550cfe034673 Mon Sep 17 00:00:00 2001
From: Biao Li <biaol@spotify.com>
Date: Thu, 2 Apr 2020 15:33:21 -0400
Subject: [PATCH] validate duplicate ports with protocols

---
 .../com/spotify/helios/common/JobValidator.java     | 12 ++++++++----
 .../com/spotify/helios/common/JobValidatorTest.java | 13 ++++++++++++-
 .../java/com/spotify/helios/agent/AgentTest.java    |  5 ++++-
 .../helios/system/ConfigFileJobCreationTest.java    |  4 +++-
 4 files changed, 27 insertions(+), 7 deletions(-)

diff --git a/helios-client/src/main/java/com/spotify/helios/common/JobValidator.java b/helios-client/src/main/java/com/spotify/helios/common/JobValidator.java
index 096eb94b3..b74b8a4c4 100644
--- a/helios-client/src/main/java/com/spotify/helios/common/JobValidator.java
+++ b/helios-client/src/main/java/com/spotify/helios/common/JobValidator.java
@@ -107,13 +107,17 @@ public Set<String> validate(final Job job) {
     errors.addAll(validateJobHostName(job.getHostname()));
 
     // Check that there's not external port collision
-    final Set<Integer> externalPorts = Sets.newHashSet();
+    final Set<String> externalPortsWithProto = Sets.newHashSet();
     for (final PortMapping mapping : job.getPorts().values()) {
       final Integer externalMappedPort = mapping.getExternalPort();
-      if (externalPorts.contains(externalMappedPort) && externalMappedPort != null) {
-        errors.add(format("Duplicate external port mapping: %s", externalMappedPort));
+      String externalMappedPortWithProto =
+          String.format(
+              "%s:%d", mapping.getProtocol(), externalMappedPort != null ? externalMappedPort : 0);
+      if (externalPortsWithProto.contains(externalMappedPortWithProto)
+          && externalMappedPort != null) {
+        errors.add(format("Duplicate external port mapping: %s", externalMappedPortWithProto));
       }
-      externalPorts.add(externalMappedPort);
+      externalPortsWithProto.add(externalMappedPortWithProto);
     }
 
     // Verify port mappings
diff --git a/helios-client/src/test/java/com/spotify/helios/common/JobValidatorTest.java b/helios-client/src/test/java/com/spotify/helios/common/JobValidatorTest.java
index 7543f711e..aa28346e4 100644
--- a/helios-client/src/test/java/com/spotify/helios/common/JobValidatorTest.java
+++ b/helios-client/src/test/java/com/spotify/helios/common/JobValidatorTest.java
@@ -207,7 +207,18 @@ public void testPortMappingCollisionFails() throws Exception {
             "2", PortMapping.of(2, 1)))
         .build();
 
-    assertEquals(ImmutableSet.of("Duplicate external port mapping: 1"), validator.validate(job));
+    assertEquals(
+        ImmutableSet.of("Duplicate external port mapping: tcp:1"), validator.validate(job));
+
+    final Job job2 = Job.newBuilder()
+        .setName("foo")
+        .setVersion("1")
+        .setImage("bar")
+        .setPorts(ImmutableMap.of("tcp-proto", PortMapping.of(1, 1, PortMapping.TCP),
+            "udp-proto", PortMapping.of(1, 1,PortMapping.UDP)))
+        .build();
+
+    assertEquals(0, validator.validate(job2).size());
   }
 
   @Test
diff --git a/helios-services/src/test/java/com/spotify/helios/agent/AgentTest.java b/helios-services/src/test/java/com/spotify/helios/agent/AgentTest.java
index e94d13dca..6a1795d2b 100644
--- a/helios-services/src/test/java/com/spotify/helios/agent/AgentTest.java
+++ b/helios-services/src/test/java/com/spotify/helios/agent/AgentTest.java
@@ -110,7 +110,10 @@ public class AgentTest {
       .setName("foo")
       .setVersion("17")
       .setPorts(ImmutableMap.of("p1", PortMapping.of(4711),
-          "p2", PortMapping.of(4712, 12345)))
+          "p2", PortMapping.of(4712, 12345),
+          "p-tcp", PortMapping.of(3100, 4100, PortMapping.TCP),
+          "p-udp", PortMapping.of(3100, 4100, PortMapping.UDP)
+          ))
       .build();
 
   private static final Map<String, Integer> FOO_PORT_ALLOCATION = ImmutableMap.of("p1", 30000,
diff --git a/helios-system-tests/src/main/java/com/spotify/helios/system/ConfigFileJobCreationTest.java b/helios-system-tests/src/main/java/com/spotify/helios/system/ConfigFileJobCreationTest.java
index f0c56ad75..34a6cf817 100644
--- a/helios-system-tests/src/main/java/com/spotify/helios/system/ConfigFileJobCreationTest.java
+++ b/helios-system-tests/src/main/java/com/spotify/helios/system/ConfigFileJobCreationTest.java
@@ -51,7 +51,9 @@ public void test() throws Exception {
     final String image = BUSYBOX;
     final Map<String, PortMapping> ports = ImmutableMap.of(
         "foo", PortMapping.of(4711),
-        "bar", PortMapping.of(5000, externalPort));
+        "bar", PortMapping.of(5000, externalPort),
+        "p-tcp", PortMapping.of(1900, 2900, PortMapping.TCP),
+        "p-udp", PortMapping.of(1900, 2900, PortMapping.UDP));
     final Map<ServiceEndpoint, ServicePorts> registration = ImmutableMap.of(
         ServiceEndpoint.of("foo-service", "tcp"), ServicePorts.of("foo"),
         ServiceEndpoint.of("bar-service", "http"), ServicePorts.of("bar"));