Merge pull request #241 from marquiz/devel/fixes-237

nfd-master: fix --label-whitelist

pkg/nfd-master/nfd-master-internal_test.go

@@ -17,6 +17,8 @@ limitations under the License.
 package nfdmaster
 
 import (
+	"regexp"
+	"sort"
 	"strings"
 	"testing"
 
@@ -54,6 +56,7 @@ func TestUpdateNodeFeatures(t *testing.T) {
 		for k, _ := range fakeFeatureLabels {
 			fakeFeatureLabelNames = append(fakeFeatureLabelNames, k)
 		}
+		sort.Strings(fakeFeatureLabelNames)
 		fakeAnnotations["feature-labels"] = strings.Join(fakeFeatureLabelNames, ",")
 
 		mockAPIHelper := new(apihelper.MockAPIHelpers)
@@ -78,8 +81,8 @@ func TestUpdateNodeFeatures(t *testing.T) {
 					So(mockNode.Labels[labelNs+k], ShouldEqual, v)
 				}
 				So(len(mockNode.Annotations), ShouldEqual, len(fakeAnnotations))
-				for k, v := range fakeFeatureLabels {
-					So(mockNode.Labels[labelNs+k], ShouldEqual, v)
+				for k, v := range fakeAnnotations {
+					So(mockNode.Annotations[annotationNs+k], ShouldEqual, v)
 				}
 			})
 		})
@@ -182,9 +185,18 @@ func TestSetLabels(t *testing.T) {
 		mockHelper := &apihelper.MockAPIHelpers{}
 		mockClient := &k8sclient.Clientset{}
 		mockNode := newMockNode()
-		mockServer := labelerServer{args: Args{}, apiHelper: mockHelper}
+		mockServer := labelerServer{args: Args{LabelWhiteList: regexp.MustCompile("")}, apiHelper: mockHelper}
 		mockCtx := context.Background()
-		mockReq := &labeler.SetLabelsRequest{NodeName: workerName, NfdVersion: workerVer, Labels: map[string]string{"feature-1": "val-1"}}
+		mockLabels := map[string]string{"feature-1": "val-1", "feature-2": "val-2", "feature-3": "val-3"}
+		mockReq := &labeler.SetLabelsRequest{NodeName: workerName, NfdVersion: workerVer, Labels: mockLabels}
+
+		mockLabelNames := make([]string, 0, len(mockLabels))
+		for k := range mockLabels {
+			mockLabelNames = append(mockLabelNames, k)
+		}
+		sort.Strings(mockLabelNames)
+		expectedAnnotations := map[string]string{"worker.version": workerVer}
+		expectedAnnotations["feature-labels"] = strings.Join(mockLabelNames, ",")
 
 		Convey("When node update succeeds", func() {
 			mockHelper.On("GetClient").Return(mockClient, nil)
@@ -194,6 +206,35 @@ func TestSetLabels(t *testing.T) {
 			Convey("No error should be returned", func() {
 				So(err, ShouldBeNil)
 			})
+			Convey("Node object should have updated with labels and annotations", func() {
+				So(len(mockNode.Labels), ShouldEqual, len(mockLabels))
+				for k, v := range mockLabels {
+					So(mockNode.Labels[labelNs+k], ShouldEqual, v)
+				}
+				So(len(mockNode.Annotations), ShouldEqual, len(expectedAnnotations))
+				for k, v := range expectedAnnotations {
+					So(mockNode.Annotations[annotationNs+k], ShouldEqual, v)
+				}
+			})
+		})
+
+		Convey("When --label-whitelist is specified", func() {
+			mockServer.args.LabelWhiteList = regexp.MustCompile("^f.*2$")
+			mockHelper.On("GetClient").Return(mockClient, nil)
+			mockHelper.On("GetNode", mockClient, workerName).Return(mockNode, nil)
+			mockHelper.On("UpdateNode", mockClient, mockNode).Return(nil)
+			_, err := mockServer.SetLabels(mockCtx, mockReq)
+			Convey("Error is nil", func() {
+				So(err, ShouldBeNil)
+			})
+			Convey("Node object should only have whitelisted labels", func() {
+				So(len(mockNode.Labels), ShouldEqual, 1)
+				So(mockNode.Labels, ShouldResemble, map[string]string{labelNs + "feature-2": "val-2"})
+
+				a := map[string]string{annotationNs + "worker.version": workerVer, annotationNs + "feature-labels": "feature-2"}
+				So(len(mockNode.Annotations), ShouldEqual, len(a))
+				So(mockNode.Annotations, ShouldResemble, a)
+			})
 		})
 
 		mockErr := errors.New("mock-error")

pkg/nfd-master/nfd-master.go

@@ -203,14 +203,15 @@ func updateMasterNode(helper apihelper.APIHelpers) error {
 }
 
 // Filter labels if whitelist has been defined
-func filterFeatureLabels(labels *Labels, labelWhiteList *regexp.Regexp) {
-	for name := range *labels {
+func filterFeatureLabels(labels Labels, labelWhiteList *regexp.Regexp) Labels {
+	for name := range labels {
 		// Skip if label doesn't match labelWhiteList
 		if !labelWhiteList.MatchString(name) {
 			stderrLogger.Printf("%s does not match the whitelist (%s) and will not be published.", name, labelWhiteList.String())
-			delete(*labels, name)
+			delete(labels, name)
 		}
 	}
+	return labels
 }
 
 // Implement LabelerServer
@@ -246,17 +247,19 @@ func (s *labelerServer) SetLabels(c context.Context, r *pb.SetLabelsRequest) (*p
 	}
 	stdoutLogger.Printf("REQUEST Node: %s NFD-version: %s Labels: %s", r.NodeName, r.NfdVersion, r.Labels)
 
+	labels := filterFeatureLabels(r.Labels, s.args.LabelWhiteList)
+
 	if !s.args.NoPublish {
 		// Advertise NFD worker version and label names as annotations
-		keys := make([]string, 0, len(r.Labels))
-		for k, _ := range r.Labels {
+		keys := make([]string, 0, len(labels))
+		for k, _ := range labels {
 			keys = append(keys, k)
 		}
 		sort.Strings(keys)
 		annotations := Annotations{"worker.version": r.NfdVersion,
 			"feature-labels": strings.Join(keys, ",")}
 
-		err := updateNodeFeatures(s.apiHelper, r.NodeName, r.Labels, annotations)
+		err := updateNodeFeatures(s.apiHelper, r.NodeName, labels, annotations)
 		if err != nil {
 			stderrLogger.Printf("failed to advertise labels: %s", err.Error())
 			return &pb.SetLabelsReply{}, err

pkg/nfd-worker/nfd-worker_test.go

@@ -19,6 +19,7 @@ package nfdworker_test
 import (
 	"fmt"
 	"os"
+	"regexp"
 	"testing"
 	"time"
 
@@ -37,6 +38,7 @@ func setupTest(args nfdmaster.Args) testContext {
 	// Fixed port and no-publish, for convenience
 	args.NoPublish = true
 	args.Port = 8192
+	args.LabelWhiteList = regexp.MustCompile("")
 	m, err := nfdmaster.NewNfdMaster(args)
 	if err != nil {
 		fmt.Printf("Test setup failed: %v\n", err)