refactor NLB securityGroup handling

staging/src/k8s.io/legacy-cloud-providers/aws/aws.go

@@ -3579,7 +3579,7 @@ func (c *Cloud) EnsureLoadBalancer(ctx context.Context, clusterName string, apiS
 			sourceRangeCidrs = append(sourceRangeCidrs, "0.0.0.0/0")
 		}
 
-		err = c.updateInstanceSecurityGroupsForNLB(v2Mappings, instances, loadBalancerName, sourceRangeCidrs)
+		err = c.updateInstanceSecurityGroupsForNLB(loadBalancerName, instances, sourceRangeCidrs, v2Mappings)
 		if err != nil {
 			klog.Warningf("Error opening ingress rules for the load balancer to the instances: %q", err)
 			return nil, err
@@ -4158,99 +4158,7 @@ func (c *Cloud) EnsureLoadBalancerDeleted(ctx context.Context, clusterName strin
 			}
 		}
 
-		{
-			var matchingGroups []*ec2.SecurityGroup
-			{
-				// Server side filter
-				describeRequest := &ec2.DescribeSecurityGroupsInput{}
-				describeRequest.Filters = []*ec2.Filter{
-					newEc2Filter("ip-permission.protocol", "tcp"),
-				}
-				response, err := c.ec2.DescribeSecurityGroups(describeRequest)
-				if err != nil {
-					return fmt.Errorf("Error querying security groups for NLB: %q", err)
-				}
-				for _, sg := range response {
-					if !c.tagging.hasClusterTag(sg.Tags) {
-						continue
-					}
-					matchingGroups = append(matchingGroups, sg)
-				}
-
-				// client-side filter out groups that don't have IP Rules we've
-				// annotated for this service
-				matchingGroups = filterForIPRangeDescription(matchingGroups, loadBalancerName)
-			}
-
-			{
-				clientRule := fmt.Sprintf("%s=%s", NLBClientRuleDescription, loadBalancerName)
-				mtuRule := fmt.Sprintf("%s=%s", NLBMtuDiscoveryRuleDescription, loadBalancerName)
-				healthRule := fmt.Sprintf("%s=%s", NLBHealthCheckRuleDescription, loadBalancerName)
-
-				for i := range matchingGroups {
-					removes := []*ec2.IpPermission{}
-					for j := range matchingGroups[i].IpPermissions {
-
-						v4rangesToRemove := []*ec2.IpRange{}
-						v6rangesToRemove := []*ec2.Ipv6Range{}
-
-						// Find IpPermission that contains k8s description
-						// If we removed the whole IpPermission, it could contain other non-k8s specified ranges
-						for k := range matchingGroups[i].IpPermissions[j].IpRanges {
-							description := aws.StringValue(matchingGroups[i].IpPermissions[j].IpRanges[k].Description)
-							if description == clientRule || description == mtuRule || description == healthRule {
-								v4rangesToRemove = append(v4rangesToRemove, matchingGroups[i].IpPermissions[j].IpRanges[k])
-							}
-						}
-
-						// Find IpPermission that contains k8s description
-						// If we removed the whole IpPermission, it could contain other non-k8s specified rangesk
-						for k := range matchingGroups[i].IpPermissions[j].Ipv6Ranges {
-							description := aws.StringValue(matchingGroups[i].IpPermissions[j].Ipv6Ranges[k].Description)
-							if description == clientRule || description == mtuRule || description == healthRule {
-								v6rangesToRemove = append(v6rangesToRemove, matchingGroups[i].IpPermissions[j].Ipv6Ranges[k])
-							}
-						}
-
-						// ipv4 and ipv6 removals cannot be included in the same permission
-						if len(v4rangesToRemove) > 0 {
-							// create a new *IpPermission to not accidentally remove UserIdGroupPairs
-							removedPermission := &ec2.IpPermission{
-								FromPort:   matchingGroups[i].IpPermissions[j].FromPort,
-								IpProtocol: matchingGroups[i].IpPermissions[j].IpProtocol,
-								IpRanges:   v4rangesToRemove,
-								ToPort:     matchingGroups[i].IpPermissions[j].ToPort,
-							}
-							removes = append(removes, removedPermission)
-						}
-						if len(v6rangesToRemove) > 0 {
-							// create a new *IpPermission to not accidentally remove UserIdGroupPairs
-							removedPermission := &ec2.IpPermission{
-								FromPort:   matchingGroups[i].IpPermissions[j].FromPort,
-								IpProtocol: matchingGroups[i].IpPermissions[j].IpProtocol,
-								Ipv6Ranges: v6rangesToRemove,
-								ToPort:     matchingGroups[i].IpPermissions[j].ToPort,
-							}
-							removes = append(removes, removedPermission)
-						}
-
-					}
-					if len(removes) > 0 {
-						changed, err := c.removeSecurityGroupIngress(aws.StringValue(matchingGroups[i].GroupId), removes)
-						if err != nil {
-							return err
-						}
-						if !changed {
-							klog.Warning("Revoking ingress was not needed; concurrent change? groupId=", *matchingGroups[i].GroupId)
-						}
-					}
-
-				}
-
-			}
-
-		}
-		return nil
+		return c.updateInstanceSecurityGroupsForNLB(loadBalancerName, nil, nil, nil)
 	}
 
 	lb, err := c.describeLoadBalancer(loadBalancerName)

staging/src/k8s.io/legacy-cloud-providers/aws/aws_loadbalancer_test.go

@@ -17,11 +17,9 @@ limitations under the License.
 package aws
 
 import (
-	"fmt"
 	"testing"
 
 	"github.com/aws/aws-sdk-go/aws"
-	"github.com/aws/aws-sdk-go/service/ec2"
 	"github.com/aws/aws-sdk-go/service/elb"
 	"github.com/stretchr/testify/assert"
 )
@@ -165,66 +163,6 @@ func TestIsNLB(t *testing.T) {
 	}
 }
 
-func TestSecurityGroupFiltering(t *testing.T) {
-	grid := []struct {
-		in          []*ec2.SecurityGroup
-		name        string
-		expected    int
-		description string
-	}{
-		{
-			in: []*ec2.SecurityGroup{
-				{
-					IpPermissions: []*ec2.IpPermission{
-						{
-							IpRanges: []*ec2.IpRange{
-								{
-									Description: aws.String("an unmanaged"),
-								},
-							},
-						},
-					},
-				},
-			},
-			name:        "unmanaged",
-			expected:    0,
-			description: "An environment without managed LBs should have %d, but found %d SecurityGroups",
-		},
-		{
-			in: []*ec2.SecurityGroup{
-				{
-					IpPermissions: []*ec2.IpPermission{
-						{
-							IpRanges: []*ec2.IpRange{
-								{
-									Description: aws.String("an unmanaged"),
-								},
-								{
-									Description: aws.String(fmt.Sprintf("%s=%s", NLBClientRuleDescription, "managedlb")),
-								},
-								{
-									Description: aws.String(fmt.Sprintf("%s=%s", NLBHealthCheckRuleDescription, "managedlb")),
-								},
-							},
-						},
-					},
-				},
-			},
-			name:        "managedlb",
-			expected:    1,
-			description: "Found %d, but should have %d Security Groups",
-		},
-	}
-
-	for _, g := range grid {
-		actual := len(filterForIPRangeDescription(g.in, g.name))
-		if actual != g.expected {
-			t.Errorf(g.description, actual, g.expected)
-		}
-	}
-
-}
-
 func TestSyncElbListeners(t *testing.T) {
 	tests := []struct {
 		name                 string

staging/src/k8s.io/legacy-cloud-providers/aws/sets_ippermissions.go

@@ -20,12 +20,19 @@ import (
 	"encoding/json"
 	"fmt"
 
+	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/service/ec2"
 )
 
 // IPPermissionSet maps IP strings of strings to EC2 IpPermissions
 type IPPermissionSet map[string]*ec2.IpPermission
 
+// IPPermissionPredicate is an predicate to test whether IPPermission matches some condition.
+type IPPermissionPredicate interface {
+	// Test checks whether specified IPPermission matches condition.
+	Test(perm *ec2.IpPermission) bool
+}
+
 // NewIPPermissionSet creates a new IPPermissionSet
 func NewIPPermissionSet(items ...*ec2.IpPermission) IPPermissionSet {
 	s := make(IPPermissionSet)
@@ -90,6 +97,23 @@ func (s IPPermissionSet) Insert(items ...*ec2.IpPermission) {
 	}
 }
 
+// Delete delete permission from the set.
+func (s IPPermissionSet) Delete(items ...*ec2.IpPermission) {
+	for _, p := range items {
+		k := keyForIPPermission(p)
+		delete(s, k)
+	}
+}
+
+// DeleteIf delete permission from the set if permission matches predicate.
+func (s IPPermissionSet) DeleteIf(predicate IPPermissionPredicate) {
+	for k, p := range s {
+		if predicate.Test(p) {
+			delete(s, k)
+		}
+	}
+}
+
 // List returns the contents as a slice.  Order is not defined.
 func (s IPPermissionSet) List() []*ec2.IpPermission {
 	res := make([]*ec2.IpPermission, 0, len(s))
@@ -146,3 +170,47 @@ func keyForIPPermission(p *ec2.IpPermission) string {
 	}
 	return string(v)
 }
+
+var _ IPPermissionPredicate = IPPermissionMatchDesc{}
+
+// IPPermissionMatchDesc checks whether specific IPPermission contains description.
+type IPPermissionMatchDesc struct {
+	Description string
+}
+
+// Test whether specific IPPermission contains description.
+func (p IPPermissionMatchDesc) Test(perm *ec2.IpPermission) bool {
+	for _, v4Range := range perm.IpRanges {
+		if aws.StringValue(v4Range.Description) == p.Description {
+			return true
+		}
+	}
+	for _, v6Range := range perm.Ipv6Ranges {
+		if aws.StringValue(v6Range.Description) == p.Description {
+			return true
+		}
+	}
+	for _, prefixListID := range perm.PrefixListIds {
+		if aws.StringValue(prefixListID.Description) == p.Description {
+			return true
+		}
+	}
+	for _, group := range perm.UserIdGroupPairs {
+		if aws.StringValue(group.Description) == p.Description {
+			return true
+		}
+	}
+	return false
+}
+
+var _ IPPermissionPredicate = IPPermissionNotMatch{}
+
+// IPPermissionNotMatch is the *not* operator for Predicate
+type IPPermissionNotMatch struct {
+	Predicate IPPermissionPredicate
+}
+
+// Test whether specific IPPermission not match the embed predicate.
+func (p IPPermissionNotMatch) Test(perm *ec2.IpPermission) bool {
+	return !p.Predicate.Test(perm)
+}