Fix UTs and update SGs for EKS create ENIs

pkg/awsutils/awsutils.go

@@ -72,6 +72,7 @@ const (
 
 	// Stagger cleanup start time to avoid calling EC2 too much. Time in seconds.
 	eniCleanupStartupDelayMax = 300
+	envCustomNetworkCfg       = "AWS_VPC_K8S_CNI_CUSTOM_NETWORK_CFG"
 )
 
 var (
@@ -154,6 +155,12 @@ type APIs interface {
 
 	// GetPrimaryENImac returns the mac address of the primary ENI
 	GetPrimaryENImac() string
+
+	//Setunmanaged ENI
+	SetUnmanagedENIs(eniID []string) error
+
+	//isUnmanagedENI
+	IsUnmanagedENI(eniID string) bool
 }
 
 // EC2InstanceMetadataCache caches instance metadata
@@ -169,6 +176,8 @@ type EC2InstanceMetadataCache struct {
 	primaryENImac    string
 	availabilityZone string
 	region           string
+	accountID        string
+	unmanagedENIs    StringSet
 
 	ec2Metadata ec2metadata.EC2Metadata
 	ec2SVC      ec2wrapper.EC2
@@ -251,6 +260,24 @@ func (ss *StringSet) Difference(other *StringSet) *StringSet {
 	return &StringSet{data: ss.data.Difference(other.data)}
 }
 
+func (ss *StringSet) Has(item string) bool {
+	ss.RLock()
+	defer ss.RUnlock()
+	return ss.data.Has(item)
+}
+
+// UseCustomNetworkCfg returns whether Pods needs to use pod specific configuration or not.
+func UseCustomNetworkCfg() bool {
+	if strValue := os.Getenv(envCustomNetworkCfg); strValue != "" {
+		parsedValue, err := strconv.ParseBool(strValue)
+		if err == nil {
+			return parsedValue
+		}
+		log.Warnf("Failed to parse %s; using default: false, err: %v", envCustomNetworkCfg, err)
+	}
+	return false
+}
+
 // New creates an EC2InstanceMetadataCache
 func New() (*EC2InstanceMetadataCache, error) {
 	//ctx is passed to initWithEC2Metadata func to cancel spawned go-routines when tests are run
@@ -392,15 +419,56 @@ func (cache *EC2InstanceMetadataCache) refreshSGIDs(mac string) error {
 	newSGs := StringSet{}
 	newSGs.Set(sgIDs)
 	addedSGs := newSGs.Difference(&cache.securityGroups)
+	addedSGsCount := 0
 	deletedSGs := cache.securityGroups.Difference(&newSGs)
+	deletedSGsCount := 0
 
 	for _, sg := range addedSGs.SortedList() {
 		log.Infof("Found %s, added to ipamd cache", sg)
+		addedSGsCount++
 	}
 	for _, sg := range deletedSGs.SortedList() {
 		log.Infof("Removed %s from ipamd cache", sg)
+		deletedSGsCount++
 	}
 	cache.securityGroups.Set(sgIDs)
+
+	if !UseCustomNetworkCfg() && (addedSGsCount != 0 || deletedSGsCount != 0) {
+		var sgIDsPtrs []*string
+		sgIDsPtrs = aws.StringSlice(sgIDs)
+
+		allENIs, err := cache.GetAttachedENIs()
+		if err != nil {
+			return errors.Wrap(err, "DescribeAllENIs: failed to get local ENI metadata")
+		}
+
+		var eniIDs []string
+		for _, eni := range allENIs {
+			eniIDs = append(eniIDs, string(eni.ENIID))
+		}
+
+		newENIs := StringSet{}
+		newENIs.Set(eniIDs)
+
+		filteredENIs := newENIs.Difference(&cache.unmanagedENIs)
+
+		//This will update SG for managed ENIs created by EKS.
+		for _, eniID := range filteredENIs.SortedList() {
+			log.Debugf("Update ENI %s", eniID)
+
+			attributeInput := &ec2.ModifyNetworkInterfaceAttributeInput{
+				Groups:             sgIDsPtrs,
+				NetworkInterfaceId: aws.String(eniID),
+			}
+			start := time.Now()
+			_, err = cache.ec2SVC.ModifyNetworkInterfaceAttributeWithContext(context.Background(), attributeInput, userAgent)
+			awsAPILatency.WithLabelValues("ModifyNetworkInterfaceAttribute", fmt.Sprint(err != nil)).Observe(msSince(start))
+			if err != nil {
+				awsAPIErrInc("ModifyNetworkInterfaceAttribute", err)
+				return errors.Wrap(err, "refreshSGIDs: unable to update the ENI's SG")
+			}
+		}
+	}
 	return nil
 }
 
@@ -1347,3 +1415,19 @@ func (cache *EC2InstanceMetadataCache) GetPrimaryENI() string {
 func (cache *EC2InstanceMetadataCache) GetPrimaryENImac() string {
 	return cache.primaryENImac
 }
+
+//SetUnmanagedENIs Set unmanaged ENI set
+func (cache *EC2InstanceMetadataCache) SetUnmanagedENIs(eniID []string) error {
+	if len(eniID) != 0 {
+		cache.unmanagedENIs.Set(eniID)
+	}
+	return nil
+}
+
+//IsUnmanagedENI returns if the eni is unmanaged
+func (cache *EC2InstanceMetadataCache) IsUnmanagedENI(eniID string) bool {
+	if len(eniID) != 0 {
+		return cache.unmanagedENIs.Has(eniID)
+	}
+	return false
+}

pkg/awsutils/awsutils_test.go

@@ -68,7 +68,7 @@ func setup(t *testing.T) (*gomock.Controller,
 func TestInitWithEC2metadata(t *testing.T) {
 	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Millisecond)
 	defer cancel()
-	ctrl, mockMetadata, _ := setup(t)
+	ctrl, mockMetadata, mockEC2 := setup(t)
 	defer ctrl.Finish()
 
 	metadataVPCIPv4CIDRs := "192.168.0.0/16	100.66.0.0/1"
@@ -77,15 +77,19 @@ func TestInitWithEC2metadata(t *testing.T) {
 	mockMetadata.EXPECT().GetMetadata(metadataLocalIP).Return(localIP, nil)
 	mockMetadata.EXPECT().GetMetadata(metadataInstanceID).Return(instanceID, nil)
 	mockMetadata.EXPECT().GetMetadata(metadataInstanceType).Return(instanceType, nil)
-	mockMetadata.EXPECT().GetMetadata(metadataMAC).Return(primaryMAC, nil)
-	mockMetadata.EXPECT().GetMetadata(metadataMACPath).Return(primaryMAC, nil)
-	mockMetadata.EXPECT().GetMetadata(metadataMACPath+primaryMAC+metadataDeviceNum).Return(eni1Device, nil)
-	mockMetadata.EXPECT().GetMetadata(metadataMACPath+primaryMAC+metadataInterface).Return(primaryMAC, nil)
+	mockMetadata.EXPECT().GetMetadata(metadataMAC).Return(primaryMAC, nil).AnyTimes()
+	mockMetadata.EXPECT().GetMetadata(metadataMACPath).Return(primaryMAC, nil).AnyTimes()
+	mockMetadata.EXPECT().GetMetadata(metadataMACPath+primaryMAC+metadataDeviceNum).Return(eni1Device, nil).AnyTimes()
+	mockMetadata.EXPECT().GetMetadata(metadataMACPath+primaryMAC+metadataInterface).Return(primaryMAC, nil).AnyTimes()
 	mockMetadata.EXPECT().GetMetadata(metadataMACPath+primaryMAC+metadataSGs).Return(sgs, nil).AnyTimes()
 	mockMetadata.EXPECT().GetMetadata(metadataMACPath+primaryMAC+metadataSubnetID).Return(subnetID, nil).AnyTimes()
+	mockMetadata.EXPECT().GetMetadata(metadataMACPath+primaryMAC+metadataSubnetCIDR).Return(subnetCIDR, nil)
 	mockMetadata.EXPECT().GetMetadata(metadataMACPath+primaryMAC+metadataVPCcidrs).Return(metadataVPCIPv4CIDRs, nil).AnyTimes()
+	mockMetadata.EXPECT().GetMetadata(metadataMACPath+primaryMAC+metadataIPv4s).Return("", nil)
 
-	ins := &EC2InstanceMetadataCache{ec2Metadata: mockMetadata}
+	mockEC2.EXPECT().ModifyNetworkInterfaceAttributeWithContext(gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, nil)
+
+	ins := &EC2InstanceMetadataCache{ec2Metadata: mockMetadata, ec2SVC: mockEC2}
 	err := ins.initWithEC2Metadata(ctx)
 	assert.NoError(t, err)
 	assert.Equal(t, az, ins.availabilityZone)
@@ -409,12 +413,15 @@ func TestTagEni(t *testing.T) {
 	mockMetadata.EXPECT().GetMetadata(metadataInstanceID).Return(instanceID, nil)
 	mockMetadata.EXPECT().GetMetadata(metadataInstanceType).Return(instanceType, nil)
 	mockMetadata.EXPECT().GetMetadata(metadataMAC).Return(primaryMAC, nil)
-	mockMetadata.EXPECT().GetMetadata(metadataMACPath).Return(primaryMAC, nil)
-	mockMetadata.EXPECT().GetMetadata(metadataMACPath+primaryMAC+metadataDeviceNum).Return(eni1Device, nil)
-	mockMetadata.EXPECT().GetMetadata(metadataMACPath+primaryMAC+metadataInterface).Return(primaryMAC, nil)
+	mockMetadata.EXPECT().GetMetadata(metadataMACPath).Return(primaryMAC, nil).AnyTimes()
+	mockMetadata.EXPECT().GetMetadata(metadataMACPath+primaryMAC+metadataDeviceNum).Return(eni1Device, nil).AnyTimes()
+	mockMetadata.EXPECT().GetMetadata(metadataMACPath+primaryMAC+metadataInterface).Return(primaryMAC, nil).AnyTimes()
 	mockMetadata.EXPECT().GetMetadata(metadataMACPath+primaryMAC+metadataSGs).Return(sgs, nil).AnyTimes()
 	mockMetadata.EXPECT().GetMetadata(metadataMACPath+primaryMAC+metadataSubnetID).Return(subnetID, nil)
+	mockMetadata.EXPECT().GetMetadata(metadataMACPath+primaryMAC+metadataSubnetCIDR).Return(subnetCIDR, nil).AnyTimes()
 	mockMetadata.EXPECT().GetMetadata(metadataMACPath+primaryMAC+metadataVPCcidrs).Return(vpcCIDR, nil).AnyTimes()
+	mockMetadata.EXPECT().GetMetadata(metadataMACPath+primaryMAC+metadataIPv4s).Return("", nil)
+	mockEC2.EXPECT().ModifyNetworkInterfaceAttributeWithContext(gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, nil)
 
 	ins := &EC2InstanceMetadataCache{ec2Metadata: mockMetadata, ec2SVC: mockEC2}
 	err := ins.initWithEC2Metadata(ctx)
@@ -424,6 +431,7 @@ func TestTagEni(t *testing.T) {
 	mockEC2.EXPECT().CreateTagsWithContext(gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, errors.New("tagging failed"))
 	mockEC2.EXPECT().CreateTagsWithContext(gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, errors.New("tagging failed"))
 	mockEC2.EXPECT().CreateTagsWithContext(gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, nil)
+
 	ins.tagENI(eniID, time.Millisecond)
 	assert.NoError(t, err)
 }

pkg/ipamd/ipamd.go

@@ -194,7 +194,6 @@ type IPAMContext struct {
 	networkClient        networkutils.NetworkAPIs
 	maxIPsPerENI         int
 	maxENI               int
-	unmanagedENIs        UnmanagedENISet // a set of ENIs tagged with "node.k8s.amazonaws.com/no_manage"
 	unmanagedENI         int
 	warmENITarget        int
 	warmIPTarget         int
@@ -211,48 +210,24 @@ type IPAMContext struct {
 	myNodeName             string
 }
 
-// UnmanagedENISet keeps a set of ENI IDs for ENIs tagged with "node.k8s.amazonaws.com/no_manage"
-type UnmanagedENISet struct {
-	sync.RWMutex
-	data map[string]bool
-}
-
-func (u *UnmanagedENISet) isUnmanaged(eniID string) bool {
-	val, ok := u.data[eniID]
-	return ok && val
-}
-
-func (u *UnmanagedENISet) reset() {
-	u.Lock()
-	defer u.Unlock()
-	u.data = make(map[string]bool)
-}
-
-func (u *UnmanagedENISet) add(eniID string) {
-	u.Lock()
-	defer u.Unlock()
-	if len(u.data) == 0 {
-		u.data = make(map[string]bool)
-	}
-	u.data[eniID] = true
-}
-
 // setUnmanagedENIs will rebuild the set of ENI IDs for ENIs tagged as "no_manage"
 func (c *IPAMContext) setUnmanagedENIs(tagMap map[string]awsutils.TagMap) {
-	c.unmanagedENIs.reset()
 	if len(tagMap) == 0 {
 		return
 	}
+	log.Infof("Calling setunmanged ENO")
+	var unmanagedENIlist []string
 	for eniID, tags := range tagMap {
 		if tags[eniNoManageTagKey] == "true" {
 			if eniID == c.awsClient.GetPrimaryENI() {
-				log.Debugf("Ignoring no_manage tag on primary ENI %s", eniID)
+				log.Infof("Ignoring no_manage tag on primary ENI %s", eniID)
 			} else {
-				log.Debugf("Marking ENI %s tagged with %s as being unmanaged", eniID, eniNoManageTagKey)
-				c.unmanagedENIs.add(eniID)
+				log.Infof("Marking ENI %s tagged with %s as being unmanaged", eniID, eniNoManageTagKey)
+				unmanagedENIlist = append(unmanagedENIlist, eniID)
 			}
 		}
 	}
+	c.awsClient.SetUnmanagedENIs(unmanagedENIlist)
 }
 
 // ReconcileCooldownCache keep track of recently freed IPs to avoid reading stale EC2 metadata
@@ -1181,8 +1156,8 @@ func (c *IPAMContext) filterUnmanagedENIs(enis []awsutils.ENIMetadata) []awsutil
 	ret := make([]awsutils.ENIMetadata, 0, len(enis))
 	for _, eni := range enis {
 		// If we have unmanaged ENIs, filter them out
-		if c.unmanagedENIs.isUnmanaged(eni.ENIID) {
-			log.Debugf("Skipping ENI %s: tagged with %s", eni.ENIID, eniNoManageTagKey)
+		if c.awsClient.IsUnmanagedENI(eni.ENIID) {
+			log.Infof("Skipping ENI %s: tagged with %s", eni.ENIID, eniNoManageTagKey)
 			numFiltered++
 			continue
 		}

pkg/ipamd/ipamd_test.go

@@ -107,6 +107,8 @@ func TestNodeInit(t *testing.T) {
 	m.awsutils.EXPECT().GetENIIPv4Limit().Return(14, nil)
 	m.awsutils.EXPECT().GetIPv4sFromEC2(eni1.ENIID).AnyTimes().Return(eni1.IPv4Addresses, nil)
 	m.awsutils.EXPECT().GetIPv4sFromEC2(eni2.ENIID).AnyTimes().Return(eni2.IPv4Addresses, nil)
+	m.awsutils.EXPECT().IsUnmanagedENI(eni1.ENIID).Return(false).AnyTimes()
+	m.awsutils.EXPECT().IsUnmanagedENI(eni2.ENIID).Return(false).AnyTimes()
 
 	primaryIP := net.ParseIP(ipaddr01)
 	m.awsutils.EXPECT().GetVPCIPv4CIDRs().AnyTimes().Return(cidrs)
@@ -377,6 +379,7 @@ func TestNodeIPPoolReconcile(t *testing.T) {
 	m.awsutils.EXPECT().GetAttachedENIs().Return(eniMetadata, nil)
 	m.awsutils.EXPECT().GetPrimaryENI().Times(2).Return(primaryENIid)
 	m.awsutils.EXPECT().DescribeAllENIs().Return(eniMetadata, map[string]awsutils.TagMap{}, "", nil)
+	m.awsutils.EXPECT().IsUnmanagedENI(primaryENIid).Return(false).AnyTimes()
 
 	mockContext.nodeIPPoolReconcile(0)
 
@@ -571,19 +574,33 @@ func TestIPAMContext_filterUnmanagedENIs(t *testing.T) {
 	mockAWSUtils.EXPECT().GetPrimaryENI().Times(2).Return(eni1.ENIID)
 
 	tests := []struct {
-		name   string
-		tagMap map[string]awsutils.TagMap
-		enis   []awsutils.ENIMetadata
-		want   []awsutils.ENIMetadata
+		name          string
+		tagMap        map[string]awsutils.TagMap
+		enis          []awsutils.ENIMetadata
+		want          []awsutils.ENIMetadata
+		unmanagedenis []string
 	}{
-		{"No tags at all", nil, allENIs, allENIs},
-		{"Primary ENI unmanaged", eni1TagMap, allENIs, allENIs},
-		{"Secondary ENI unmanaged", eni2TagMap, allENIs, primaryENIonly},
+		{"No tags at all", nil, allENIs, allENIs, nil},
+		{"Primary ENI unmanaged", eni1TagMap, allENIs, allENIs, nil},
+		{"Secondary ENI unmanaged", eni2TagMap, allENIs, primaryENIonly, []string{eni2.ENIID}},
 	}
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			c := &IPAMContext{awsClient: mockAWSUtils}
+			mockAWSUtils.EXPECT().SetUnmanagedENIs(tt.unmanagedenis).Return(nil).AnyTimes()
 			c.setUnmanagedENIs(tt.tagMap)
+
+			mockAWSUtils.EXPECT().IsUnmanagedENI(gomock.Any()).DoAndReturn(
+				func(eni string) (unmanaged bool) {
+					if eni != eni1.ENIID {
+						if _, ok := tt.tagMap[eni]; ok {
+							return true
+						}
+					}
+					return false
+
+				}).AnyTimes()
+
 			if got := c.filterUnmanagedENIs(tt.enis); !reflect.DeepEqual(got, tt.want) {
 				t.Errorf("filterUnmanagedENIs() = %v, want %v", got, tt.want)
 			}