digitalocean · akrantz01 · Feb 13, 2023
diff --git a/cloud-controller-manager/do/firewall_controller.go b/cloud-controller-manager/do/firewall_controller.go
@@ -283,35 +283,40 @@ func (fm *firewallManager) createReconciledFirewallRequest(serviceList []*v1.Ser
 			continue
 		}
 
-		if svc.Spec.Type == v1.ServiceTypeNodePort {
-			// this is a nodeport service so we should check for existing inbound rules on all ports.
-			for _, servicePort := range svc.Spec.Ports {
-				// In the odd case that a failure is asynchronous causing the NodePort to be set to zero.
-				if servicePort.NodePort == 0 {
-					klog.Warning("NodePort on the service is set to zero")
-					continue
-				}
-				var protocol string
-				switch servicePort.Protocol {
-				case v1.ProtocolTCP:
-					protocol = "tcp"
-				case v1.ProtocolUDP:
-					protocol = "udp"
-				default:
-					klog.Warningf("unsupported service protocol %v, skipping service port %v", servicePort.Protocol, servicePort.Name)
-					continue
-				}
-
-				nodePortInboundRules = append(nodePortInboundRules,
-					godo.InboundRule{
-						Protocol:  protocol,
-						PortRange: strconv.Itoa(int(servicePort.NodePort)),
-						Sources: &godo.Sources{
-							Addresses: []string{"0.0.0.0/0", "::/0"},
-						},
-					},
-				)
+		if !(svc.Spec.Type == v1.ServiceTypeNodePort || svc.Spec.Type == v1.ServiceTypeLoadBalancer) {
+			continue
+		}
+
+		sources, found := getSources(svc)
+		if !found {
+			continue
+		}
+
+		// this is a nodeport service so we should check for existing inbound rules on all ports.
+		for _, servicePort := range svc.Spec.Ports {
+			// In the odd case that a failure is asynchronous causing the NodePort to be set to zero.
+			if servicePort.NodePort == 0 {
+				klog.Warning("NodePort on the service is set to zero")
+				continue
 			}
+			var protocol string
+			switch servicePort.Protocol {
+			case v1.ProtocolTCP:
+				protocol = "tcp"
+			case v1.ProtocolUDP:
+				protocol = "udp"
+			default:
+				klog.Warningf("unsupported service protocol %v, skipping service port %v", servicePort.Protocol, servicePort.Name)
+				continue
+			}
+
+			nodePortInboundRules = append(nodePortInboundRules,
+				godo.InboundRule{
+					Protocol:  protocol,
+					PortRange: strconv.Itoa(int(servicePort.NodePort)),
+					Sources:   &sources,
+				},
+			)
 		}
 	}
 	return &godo.FirewallRequest{
@@ -334,6 +339,32 @@ func isManaged(service *v1.Service) (bool, error) {
 	return !found || val, nil
 }
 
+// getSources returns the sources to allow traffic from. For NodePort services, this
+// allows traffic from anywhere. For LoadBalancer services, this allows services from
+// the managed load balancer.
+func getSources(service *v1.Service) (godo.Sources, bool) {
+	var sources godo.Sources
+
+	switch service.Spec.Type {
+	case v1.ServiceTypeNodePort:
+		sources.Addresses = []string{"0.0.0.0/0", "::/0"}
+	case v1.ServiceTypeLoadBalancer:
+		if service.Spec.AllocateLoadBalancerNodePorts != nil && !*service.Spec.AllocateLoadBalancerNodePorts {
+			klog.Warning("NodePorts must be allocated for LoadBalancer services")
+			return godo.Sources{}, false
+		}
+
+		uid, ok := service.Annotations[annDOLoadBalancerID]
+		if !ok {
+			klog.Warning("No managed load balancer found for service")
+			return godo.Sources{}, false
+		}
+		sources.LoadBalancerUIDs = []string{uid}
+	}
+
+	return sources, true
+}
+
 func (fm *firewallManager) executeInstrumentedFirewallOperationGetByID(ctx context.Context, fwID string) (*godo.Firewall, *godo.Response, error) {
 	return fm.executeInstrumentedFirewallOperation(ctx, firewallOperationGetByID, func(ctx context.Context) (*godo.Firewall, *godo.Response, error) {
 		return fm.client.Firewalls.Get(ctx, fwID)

diff --git a/cloud-controller-manager/do/firewall_controller_test.go b/cloud-controller-manager/do/firewall_controller_test.go
@@ -630,6 +630,212 @@ func TestFirewallController_createReconciledFirewallRequest(t *testing.T) {
 				},
 			},
 		},
+		{
+			name: "reconcile firewall based on added loadbalancer service",
+			firewallRequest: &godo.FirewallRequest{
+				Name: testWorkerFWName,
+				InboundRules: []godo.InboundRule{
+					{
+						Protocol:  "tcp",
+						PortRange: "31000",
+						Sources: &godo.Sources{
+							LoadBalancerUIDs: []string{"test-lb-id"},
+						},
+					},
+					{
+						Protocol:  "udp",
+						PortRange: "31000",
+						Sources: &godo.Sources{
+							LoadBalancerUIDs: []string{"test-lb-id"},
+						},
+					},
+				},
+				OutboundRules: testOutboundRules,
+				Tags:          testWorkerFWTags,
+			},
+			serviceList: []*v1.Service{
+				{
+					ObjectMeta: metav1.ObjectMeta{
+						Name: "test",
+						UID:  "abc123",
+						Annotations: map[string]string{
+							annDOLoadBalancerID: "test-lb-id",
+						},
+					},
+					Spec: v1.ServiceSpec{
+						Type: v1.ServiceTypeLoadBalancer,
+						Ports: []v1.ServicePort{
+							{
+								Name:     "port1",
+								Protocol: v1.ProtocolTCP,
+								NodePort: int32(31000),
+							},
+							{
+								Name:     "port2",
+								Protocol: v1.ProtocolUDP,
+								NodePort: int32(31000),
+							},
+						},
+					},
+				},
+			},
+		},
+		{
+			name: "reconcile firewall when loadbalancer does not allocate nodeports",
+			firewallRequest: &godo.FirewallRequest{
+				Name:          testWorkerFWName,
+				InboundRules:  nil,
+				OutboundRules: testOutboundRules,
+				Tags:          testWorkerFWTags,
+			},
+			serviceList: []*v1.Service{
+				{
+					ObjectMeta: metav1.ObjectMeta{
+						Name: "test",
+						UID:  "abc123",
+						Annotations: map[string]string{
+							annDOLoadBalancerID: "test-lb-id",
+						},
+					},
+					Spec: v1.ServiceSpec{
+						Type:                          v1.ServiceTypeLoadBalancer,
+						AllocateLoadBalancerNodePorts: boolPtr(false),
+					},
+				},
+			},
+		},
+		{
+			name: "reconcile firewall when loadbalancer id not assigned",
+			firewallRequest: &godo.FirewallRequest{
+				Name:          testWorkerFWName,
+				InboundRules:  nil,
+				OutboundRules: testOutboundRules,
+				Tags:          testWorkerFWTags,
+			},
+			serviceList: []*v1.Service{
+				{
+					ObjectMeta: metav1.ObjectMeta{
+						Name: "test",
+						UID:  "abc123",
+					},
+					Spec: v1.ServiceSpec{
+						Type: v1.ServiceTypeLoadBalancer,
+						Ports: []v1.ServicePort{
+							{
+								Name:     "port1",
+								Protocol: v1.ProtocolTCP,
+								NodePort: int32(31000),
+							},
+							{
+								Name:     "port2",
+								Protocol: v1.ProtocolUDP,
+								NodePort: int32(31000),
+							},
+						},
+					},
+				},
+			},
+		},
+		{
+			name: "reconcile firewall when multiple loadbalancers added",
+			firewallRequest: &godo.FirewallRequest{
+				Name: testWorkerFWName,
+				InboundRules: []godo.InboundRule{
+					{
+						Protocol:  "tcp",
+						PortRange: "30000",
+						Sources: &godo.Sources{
+							LoadBalancerUIDs: []string{"lb-test-1"},
+						},
+					},
+					{
+						Protocol:  "udp",
+						PortRange: "31000",
+						Sources: &godo.Sources{
+							LoadBalancerUIDs: []string{"lb-test-2"},
+						},
+					},
+				},
+				OutboundRules: testOutboundRules,
+				Tags:          testWorkerFWTags,
+			},
+			serviceList: []*v1.Service{
+				{
+					ObjectMeta: metav1.ObjectMeta{
+						Name: "test1",
+						UID:  "abc123",
+						Annotations: map[string]string{
+							annDOLoadBalancerID: "lb-test-1",
+						},
+					},
+					Spec: v1.ServiceSpec{
+						Type: v1.ServiceTypeLoadBalancer,
+						Ports: []v1.ServicePort{
+							{
+								Name:     "port1",
+								Protocol: v1.ProtocolTCP,
+								NodePort: int32(30000),
+							},
+						},
+					},
+				},
+				{
+					ObjectMeta: metav1.ObjectMeta{
+						Name: "test2",
+						UID:  "abc456",
+						Annotations: map[string]string{
+							annDOLoadBalancerID: "lb-test-2",
+						},
+					},
+					Spec: v1.ServiceSpec{
+						Type: v1.ServiceTypeLoadBalancer,
+						Ports: []v1.ServicePort{
+							{
+								Name:     "port2",
+								Protocol: v1.ProtocolUDP,
+								NodePort: int32(31000),
+							},
+						},
+					},
+				},
+				{
+					ObjectMeta: metav1.ObjectMeta{
+						Name: "test3",
+						UID:  "abc789",
+					},
+					Spec: v1.ServiceSpec{
+						Type: v1.ServiceTypeLoadBalancer,
+						Ports: []v1.ServicePort{
+							{
+								Name:     "port3",
+								Protocol: v1.ProtocolTCP,
+								NodePort: int32(32000),
+							},
+						},
+					},
+				},
+				{
+					ObjectMeta: metav1.ObjectMeta{
+						Name: "test4",
+						UID:  "abc0",
+						Annotations: map[string]string{
+							annDOLoadBalancerID: "lb-test-3",
+						},
+					},
+					Spec: v1.ServiceSpec{
+						Type:                          v1.ServiceTypeLoadBalancer,
+						AllocateLoadBalancerNodePorts: boolPtr(false),
+						Ports: []v1.ServicePort{
+							{
+								Name:     "port4",
+								Protocol: v1.ProtocolUDP,
+								NodePort: int32(34000),
+							},
+						},
+					},
+				},
+			},
+		},
 	}
 
 	for _, test := range testcases {

diff --git a/cloud-controller-manager/do/firewall_fixture_test.go b/cloud-controller-manager/do/firewall_fixture_test.go
@@ -287,3 +287,7 @@ func mustCopy(v interface{}) interface{} {
 	w, err := copystructure.Copy(v)
 	return copystructure.Must(w, err)
 }
+
+func boolPtr(value bool) *bool {
+	return &value
+}