From 007bf05b17cef2bf848c3e57f899156be149c6e3 Mon Sep 17 00:00:00 2001
From: Spencer Hance <shance@google.com>
Date: Wed, 5 Aug 2020 03:59:31 -0700
Subject: [PATCH] Add Custom health check port to firewall rule

---
 pkg/firewalls/controller.go      | 22 +++++++++++++++++-
 pkg/firewalls/controller_test.go | 38 ++++++++++++++++++++++++++++++++
 pkg/utils/utils.go               |  5 +++++
 3 files changed, 64 insertions(+), 1 deletion(-)

diff --git a/pkg/firewalls/controller.go b/pkg/firewalls/controller.go
index f026d3a63d..974f9aad2e 100644
--- a/pkg/firewalls/controller.go
+++ b/pkg/firewalls/controller.go
@@ -20,6 +20,7 @@ import (
 	"errors"
 	"fmt"
 	"reflect"
+	"strconv"
 	"time"
 
 	apiv1 "k8s.io/api/core/v1"
@@ -181,8 +182,15 @@ func (fwc *FirewallController) sync(key string) error {
 		}
 	}
 
+	var additionalPorts []string
+	if flags.F.EnableBackendConfigHealthCheck {
+		hcPorts := fwc.getCustomHealthCheckPorts(gceSvcPorts)
+		additionalPorts = append(additionalPorts, hcPorts...)
+	}
+	additionalPorts = append(additionalPorts, negPorts...)
+
 	// Ensure firewall rule for the cluster and pass any NEG endpoint ports.
-	if err := fwc.firewallPool.Sync(nodeNames, negPorts, additionalRanges); err != nil {
+	if err := fwc.firewallPool.Sync(nodeNames, additionalPorts, additionalRanges); err != nil {
 		if fwErr, ok := err.(*FirewallXPNError); ok {
 			// XPN: Raise an event on each ingress
 			for _, ing := range gceIngresses {
@@ -217,3 +225,15 @@ func (fwc *FirewallController) ilbFirewallSrcRange(gceIngresses []*v1beta1.Ingre
 
 	return "", ErrNoILBIngress
 }
+
+func (fwc *FirewallController) getCustomHealthCheckPorts(svcPorts []utils.ServicePort) []string {
+	var result []string
+
+	for _, svcPort := range svcPorts {
+		if svcPort.BackendConfig != nil && svcPort.BackendConfig.Spec.HealthCheck != nil && svcPort.BackendConfig.Spec.HealthCheck.Port != nil {
+			result = append(result, strconv.FormatInt(*svcPort.BackendConfig.Spec.HealthCheck.Port, 10))
+		}
+	}
+
+	return result
+}
diff --git a/pkg/firewalls/controller_test.go b/pkg/firewalls/controller_test.go
index f2c0524261..c443da2dae 100644
--- a/pkg/firewalls/controller_test.go
+++ b/pkg/firewalls/controller_test.go
@@ -20,10 +20,12 @@ import (
 	"testing"
 	"time"
 
+	"github.com/google/go-cmp/cmp"
 	api_v1 "k8s.io/api/core/v1"
 	"k8s.io/api/networking/v1beta1"
 	"k8s.io/apimachinery/pkg/types"
 	"k8s.io/client-go/kubernetes/fake"
+	v1 "k8s.io/ingress-gce/pkg/apis/backendconfig/v1"
 	backendconfigclient "k8s.io/ingress-gce/pkg/backendconfig/client/clientset/versioned/fake"
 	test "k8s.io/ingress-gce/pkg/test"
 	"k8s.io/ingress-gce/pkg/utils"
@@ -102,3 +104,39 @@ func TestFirewallCreateDelete(t *testing.T) {
 		t.Fatalf("cloud.GetFirewall(%v) = _, %v, want _, 404 error", ruleName, err)
 	}
 }
+
+func TestGetCustomHealthCheckPorts(t *testing.T) {
+	t.Parallel()
+
+	testCases := []struct {
+		desc     string
+		svcPorts []utils.ServicePort
+		expect   []string
+	}{
+		{
+			desc:     "One service port with custom port",
+			svcPorts: []utils.ServicePort{utils.ServicePort{BackendConfig: &v1.BackendConfig{Spec: v1.BackendConfigSpec{HealthCheck: &v1.HealthCheckConfig{Port: utils.NewInt64Pointer(8000)}}}}},
+			expect:   []string{"8000"},
+		},
+		{
+			desc: "Two service ports with custom port",
+			svcPorts: []utils.ServicePort{utils.ServicePort{BackendConfig: &v1.BackendConfig{Spec: v1.BackendConfigSpec{HealthCheck: &v1.HealthCheckConfig{Port: utils.NewInt64Pointer(8000)}}}},
+				utils.ServicePort{BackendConfig: &v1.BackendConfig{Spec: v1.BackendConfigSpec{HealthCheck: &v1.HealthCheckConfig{Port: utils.NewInt64Pointer(9000)}}}}},
+			expect: []string{"8000", "9000"},
+		},
+		{
+			desc:   "No service ports",
+			expect: nil,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.desc, func(t *testing.T) {
+			fwc := newFirewallController()
+			result := fwc.getCustomHealthCheckPorts(tc.svcPorts)
+			if diff := cmp.Diff(tc.expect, result); diff != "" {
+				t.Errorf("unexpected diff of ports (-want +got): \n%s", diff)
+			}
+		})
+	}
+}
diff --git a/pkg/utils/utils.go b/pkg/utils/utils.go
index 8cc65dd3e9..1a41d9a83f 100644
--- a/pkg/utils/utils.go
+++ b/pkg/utils/utils.go
@@ -563,3 +563,8 @@ func MakeL4ILBServiceDescription(svcName, ip string, version meta.Version) (stri
 func NewStringPointer(s string) *string {
 	return &s
 }
+
+// NewInt64Pointer returns a pointer to the provided int64 literal
+func NewInt64Pointer(i int64) *int64 {
+	return &i
+}