fix: incorrect parsing of exposedPorts in readiness check

lifecycle.go

@@ -198,6 +198,34 @@ var defaultLogConsumersHook = func(cfg *LogConsumerConfig) ContainerLifecycleHoo
 	}
 }
 
+func checkPortsMapped(exposedAndMappedPorts nat.PortMap, exposedPorts []string) error {
+	portMap, _, err := nat.ParsePortSpecs(exposedPorts)
+	if err != nil {
+		return fmt.Errorf("parse exposed ports: %w", err)
+	}
+
+	for exposedPort := range portMap {
+		// having entries in exposedAndMappedPorts, where the key is the exposed port,
+		// and the value is the mapped port, means that the port has been already mapped.
+		if _, ok := exposedAndMappedPorts[exposedPort]; ok {
+			continue
+		}
+
+		// check if the port is mapped with the protocol (default is TCP)
+		if strings.Contains(string(exposedPort), "/") {
+			return fmt.Errorf("port %s is not mapped yet", exposedPort)
+		}
+
+		// Port didn't have a type, default to tcp and retry.
+		exposedPort += "/tcp"
+		if _, ok := exposedAndMappedPorts[exposedPort]; !ok {
+			return fmt.Errorf("port %s is not mapped yet", exposedPort)
+		}
+	}
+
+	return nil
+}
+
 // defaultReadinessHook is a hook that will wait for the container to be ready
 var defaultReadinessHook = func() ContainerLifecycleHooks {
 	return ContainerLifecycleHooks{
@@ -222,26 +250,7 @@ var defaultReadinessHook = func() ContainerLifecycleHooks {
 							return err
 						}
 
-						exposedAndMappedPorts := jsonRaw.NetworkSettings.Ports
-
-						for _, exposedPort := range dockerContainer.exposedPorts {
-							portMap := nat.Port(exposedPort)
-							// having entries in exposedAndMappedPorts, where the key is the exposed port,
-							// and the value is the mapped port, means that the port has been already mapped.
-							if _, ok := exposedAndMappedPorts[portMap]; !ok {
-								// check if the port is mapped with the protocol (default is TCP)
-								if !strings.Contains(exposedPort, "/") {
-									portMap = nat.Port(fmt.Sprintf("%s/tcp", exposedPort))
-									if _, ok := exposedAndMappedPorts[portMap]; !ok {
-										return fmt.Errorf("port %s is not mapped yet", exposedPort)
-									}
-								} else {
-									return fmt.Errorf("port %s is not mapped yet", exposedPort)
-								}
-							}
-						}
-
-						return nil
+						return checkPortsMapped(jsonRaw.NetworkSettings.Ports, dockerContainer.exposedPorts)
 					},
 					b,
 					func(err error, duration time.Duration) {

lifecycle_test.go

@@ -456,6 +456,73 @@ func TestMergePortBindings(t *testing.T) {
 	}
 }
 
+func TestPortMappingCheck(t *testing.T) {
+	makePortMap := func(ports ...string) nat.PortMap {
+		out := make(nat.PortMap)
+		for _, port := range ports {
+			// We don't care about the actual binding in this test
+			out[nat.Port(port)] = nil
+		}
+		return out
+	}
+
+	tests := map[string]struct {
+		exposedAndMappedPorts nat.PortMap
+		exposedPorts          []string
+		expectError           bool
+	}{
+		"no-protocol": {
+			exposedAndMappedPorts: makePortMap("1024/tcp"),
+			exposedPorts:          []string{"1024"},
+		},
+		"protocol": {
+			exposedAndMappedPorts: makePortMap("1024/tcp"),
+			exposedPorts:          []string{"1024/tcp"},
+		},
+		"protocol-target-port": {
+			exposedAndMappedPorts: makePortMap("1024/tcp"),
+			exposedPorts:          []string{"1024:1024/tcp"},
+		},
+		"target-port": {
+			exposedAndMappedPorts: makePortMap("1024/tcp"),
+			exposedPorts:          []string{"1024:1024"},
+		},
+		"multiple-ports": {
+			exposedAndMappedPorts: makePortMap("1024/tcp", "1025/tcp", "1026/tcp"),
+			exposedPorts:          []string{"1024", "25:1025/tcp", "1026:1026"},
+		},
+		"only-ipv4": {
+			exposedAndMappedPorts: makePortMap("1024/tcp"),
+			exposedPorts:          []string{"0.0.0.0::1024/tcp"},
+		},
+		"no-mapped-ports": {
+			exposedAndMappedPorts: makePortMap(),
+			exposedPorts:          []string{"1024"},
+			expectError:           true,
+		},
+		"wrong-mapped-port": {
+			exposedAndMappedPorts: makePortMap("1023/tcp"),
+			exposedPorts:          []string{"1024"},
+			expectError:           true,
+		},
+		"subset-mapped-ports": {
+			exposedAndMappedPorts: makePortMap("1024/tcp", "1025/tcp"),
+			exposedPorts:          []string{"1024", "1025", "1026"},
+			expectError:           true,
+		},
+	}
+	for name, tt := range tests {
+		t.Run(name, func(t *testing.T) {
+			err := checkPortsMapped(tt.exposedAndMappedPorts, tt.exposedPorts)
+			if tt.expectError {
+				require.Error(t, err)
+				return
+			}
+			require.NoError(t, err)
+		})
+	}
+}
+
 func TestLifecycleHooks(t *testing.T) {
 	tests := []struct {
 		name  string