diff --git a/nftables_test.go b/nftables_test.go index 584d6f9..baefa63 100644 --- a/nftables_test.go +++ b/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) + } + } +} diff --git a/set.go b/set.go index e2f58fe..452ae4e 100644 --- a/set.go +++ b/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) } } diff --git a/userdata/userdata.go b/userdata/userdata.go index b8df006..984d84c 100644 --- a/userdata/userdata.go +++ b/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...)