add support for comments in set elements (google#293)

Author: shiningw

nftables_test.go

@@ -7816,3 +7816,74 @@ func TestNftablesDeadlock(t *testing.T) {
 		})
 	}
 }
+func TestSetElementComment(t *testing.T) {
+	// Create a new network namespace to test these operations
+	conn, newNS := nftest.OpenSystemConn(t, *enableSysTests)
+	defer nftest.CleanupSystemConn(t, newNS)
+	conn.FlushRuleset()
+	defer conn.FlushRuleset()
+
+	// Add a new table
+	table := &nftables.Table{
+		Family: nftables.TableFamilyIPv4,
+		Name:   "filter",
+	}
+	conn.AddTable(table)
+
+	// Create a new set
+	set := &nftables.Set{
+		Name:    "test-set",
+		Table:   table,
+		KeyType: nftables.TypeIPAddr,
+	}
+
+	// Create set elements with comments
+	elements := []nftables.SetElement{
+		{
+			Key:     net.ParseIP("192.0.2.1").To4(),
+			Comment: "First IP address",
+		},
+		{
+			Key:     net.ParseIP("192.0.2.2").To4(),
+			Comment: "Second IP address",
+		},
+	}
+
+	// Add the set with elements
+	if err := conn.AddSet(set, elements); err != nil {
+		t.Fatalf("failed to add set: %v", err)
+	}
+	if err := conn.Flush(); err != nil {
+		t.Fatalf("failed to flush: %v", err)
+	}
+
+	// Get the set elements back and verify comments
+	gotElements, err := conn.GetSetElements(set)
+	if err != nil {
+		t.Fatalf("failed to get set elements: %v", err)
+	}
+
+	if got, want := len(gotElements), len(elements); got != want {
+		t.Fatalf("got %d elements, want %d", got, want)
+	}
+
+	// Create maps to compare elements by their IP addresses
+	wantMap := make(map[string]string)
+	for _, elem := range elements {
+		wantMap[string(elem.Key)] = elem.Comment
+	}
+
+	gotMap := make(map[string]string)
+	for _, elem := range gotElements {
+		gotMap[string(elem.Key)] = elem.Comment
+	}
+
+	// Compare the comments for each IP
+	for ip, wantComment := range wantMap {
+		if gotComment, ok := gotMap[ip]; !ok {
+			t.Errorf("IP %s not found in retrieved elements", ip)
+		} else if gotComment != wantComment {
+			t.Errorf("for IP %s: got comment %q, want comment %q", ip, gotComment, wantComment)
+		}
+	}
+}

set.go

@@ -288,6 +288,7 @@ type SetElement struct {
 	Expires time.Duration
 
 	Counter *expr.Counter
+	Comment string
 }
 
 func (s *SetElement) decode(fam byte) func(b []byte) error {
@@ -322,6 +323,12 @@ func (s *SetElement) decode(fam byte) func(b []byte) error {
 				s.Timeout = time.Millisecond * time.Duration(ad.Uint64())
 			case unix.NFTA_SET_ELEM_EXPIRATION:
 				s.Expires = time.Millisecond * time.Duration(ad.Uint64())
+			case unix.NFTA_SET_ELEM_USERDATA:
+				userData := ad.Bytes()
+				// Try to extract comment from userdata if present
+				if comment, ok := userdata.GetString(userData, userdata.NFTNL_UDATA_SET_ELEM_COMMENT); ok {
+					s.Comment = comment
+				}
 			case unix.NFTA_SET_ELEM_EXPR:
 				elems, err := parseexprfunc.ParseExprBytesFunc(fam, ad)
 				if err != nil {
@@ -454,6 +461,12 @@ func (s *Set) makeElemList(vals []SetElement, id uint32) ([]netlink.Attribute, e
 			// If niether of previous cases matche, it means 'e' is an element of a regular Set, no need to add to the attributes
 		}
 
+		// Add comment to userdata if present
+		if len(v.Comment) > 0 {
+			userData := userdata.AppendString(nil, userdata.NFTNL_UDATA_SET_ELEM_COMMENT, v.Comment)
+			item = append(item, netlink.Attribute{Type: unix.NFTA_SET_ELEM_USERDATA, Data: userData})
+		}
+
 		encodedItem, err := netlink.MarshalAttributes(item)
 		if err != nil {
 			return nil, fmt.Errorf("marshal item %d: %v", i, err)
@@ -807,6 +820,7 @@ func elementsFromMsg(fam byte, msg netlink.Message) ([]SetElement, error) {
 		b := ad.Bytes()
 		if ad.Type() == unix.NFTA_SET_ELEM_LIST_ELEMENTS {
 			ad, err := netlink.NewAttributeDecoder(b)
+
 			if err != nil {
 				return nil, err
 			}
@@ -818,6 +832,7 @@ func elementsFromMsg(fam byte, msg netlink.Message) ([]SetElement, error) {
 				case unix.NFTA_LIST_ELEM:
 					ad.Do(elem.decode(fam))
 				}
+
 				elements = append(elements, elem)
 			}
 		}

userdata/userdata.go

@@ -46,6 +46,12 @@ const (
 	NFTNL_UDATA_SET_MAX
 )
 
+// Set element userdata types
+const (
+	NFTNL_UDATA_SET_ELEM_COMMENT Type = iota
+	NFTNL_UDATA_SET_ELEM_FLAGS
+)
+
 func Append(udata []byte, typ Type, data []byte) []byte {
 	udata = append(udata, byte(typ), byte(len(data)))
 	udata = append(udata, data...)