From df1ed79431b3954006696e9b8ec56077befd0a0e Mon Sep 17 00:00:00 2001 From: nickgarlis Date: Tue, 18 Nov 2025 19:20:23 +0100 Subject: [PATCH] Add ResetSetElements method ResetSetElements resets the stateful expressions (e.g., counters) of all elements in the specified set. The reset is applied immediately (no Flush is required) and the returned elements reflect their state prior to the reset. This update also adds a reset flag to the internal getSetElements method, along with a safeguard for situations where the user provides no valid set identifier. --- nftables_test.go | 70 ++++++++++++++++++++++++++++++++++++++++++++++++ set.go | 28 ++++++++++++++----- 2 files changed, 92 insertions(+), 6 deletions(-) diff --git a/nftables_test.go b/nftables_test.go index 43ff4c4..4e915a0 100644 --- a/nftables_test.go +++ b/nftables_test.go @@ -8491,3 +8491,73 @@ func TestSetElementCounter(t *testing.T) { t.Fatalf("got counter %v, want %v", gotCounter, counter) } } + +func TestResetSetElements(t *testing.T) { + conn, newNS := nftest.OpenSystemConn(t, *enableSysTests) + defer nftest.CleanupSystemConn(t, newNS) + defer conn.FlushRuleset() + + table := conn.AddTable(&nftables.Table{ + Name: "test-table", + Family: nftables.TableFamilyIPv4, + }) + set := &nftables.Set{ + Name: "test-set", + Table: table, + KeyType: nftables.TypeInetService, + } + elements := []nftables.SetElement{ + { + Key: binaryutil.BigEndian.PutUint16(80), + Counter: &expr.Counter{ + Bytes: 1024, + Packets: 10, + }, + }, + { + Key: binaryutil.BigEndian.PutUint16(443), + Counter: &expr.Counter{ + Bytes: 2048, + Packets: 20, + }, + }, + } + 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) + } + + got, err := conn.ResetSetElements(set) + if err != nil { + t.Fatalf("failed to reset set elements: %v", err) + } + if len(got) != len(elements) { + t.Fatalf("got %d elements, want %d", len(got), len(elements)) + } + + for i, ge := range got { + want := elements[i] + if !bytes.Equal(ge.Key, want.Key) { + t.Errorf("element %d: got key %v, want %v", i, ge.Key, want.Key) + } + if !reflect.DeepEqual(ge.Counter, want.Counter) { + t.Errorf("element %d: got counter %v, want %v", i, ge.Counter, want.Counter) + } + } + + resetEls, err := conn.GetSetElements(set) + if err != nil { + t.Fatalf("failed to get set elements after reset: %v", err) + } + for i, re := range resetEls { + want := elements[i] + if !bytes.Equal(re.Key, want.Key) { + t.Errorf("element %d: got key %v, want %v", i, re.Key, want.Key) + } + if re.Counter.Bytes != 0 || re.Counter.Packets != 0 { + t.Errorf("element %d: got counter %v, want zeroed counter", i, re.Counter) + } + } +} diff --git a/set.go b/set.go index b220a82..38997e1 100644 --- a/set.go +++ b/set.go @@ -1043,7 +1043,9 @@ func (cc *Conn) GetSetByName(t *Table, name string) (*Set, error) { // getSetElements retrieves elements from a set. // If e is empty, all elements are retrieved. Otherwise, only the specified // elements are retrieved if they exist in the set. -func (cc *Conn) getSetElements(s *Set, e []SetElement) ([]SetElement, error) { +// If reset is true, the stateful expressions (e.g., counters) of the elements +// being retrieved are reset. +func (cc *Conn) getSetElements(s *Set, e []SetElement, reset bool) ([]SetElement, error) { conn, closer, err := cc.netlinkConn() if err != nil { return nil, err @@ -1053,13 +1055,14 @@ func (cc *Conn) getSetElements(s *Set, e []SetElement) ([]SetElement, error) { flags := netlink.Request attrs := []netlink.Attribute{ {Type: unix.NFTA_SET_ELEM_LIST_TABLE, Data: []byte(s.Table.Name + "\x00")}, - {Type: unix.NFTA_SET_ELEM_LIST_SET, Data: []byte(s.Name + "\x00")}, } if s.Name != "" { attrs = append(attrs, netlink.Attribute{Type: unix.NFTA_SET_ELEM_LIST_SET, Data: []byte(s.Name + "\x00")}) - } else { + } else if s.ID > 0 { attrs = append(attrs, netlink.Attribute{Type: unix.NFTA_SET_ELEM_LIST_SET_ID, Data: binaryutil.BigEndian.PutUint32(s.ID)}) + } else { + return nil, fmt.Errorf("set must either have a valid name or ID") } if len(e) > 0 { @@ -1086,9 +1089,14 @@ func (cc *Conn) getSetElements(s *Set, e []SetElement) ([]SetElement, error) { return nil, err } + msgType := nftMsgGetSetElem + if reset { + msgType = nftMsgGetSetElemReset + } + message := netlink.Message{ Header: netlink.Header{ - Type: nftMsgGetSetElem.HeaderType(), + Type: msgType.HeaderType(), Flags: flags, }, Data: append(extraHeader(uint8(s.Table.Family), 0), data...), @@ -1116,10 +1124,18 @@ func (cc *Conn) getSetElements(s *Set, e []SetElement) ([]SetElement, error) { // GetSetElements returns the elements in the specified set. func (cc *Conn) GetSetElements(s *Set) ([]SetElement, error) { - return cc.getSetElements(s, nil) + return cc.getSetElements(s, nil, false) } // FindSetElements returns the specified elements in the set. func (cc *Conn) FindSetElements(s *Set, e []SetElement) ([]SetElement, error) { - return cc.getSetElements(s, e) + return cc.getSetElements(s, e, false) +} + +// ResetSetElements resets the stateful expressions (e.g., counters) of all +// elements in the specified set. The reset is applied immediately +// (no Flush is required). The returned elements reflect their state prior to +// the reset. Requires a kernel version >= 6.3. +func (cc *Conn) ResetSetElements(s *Set) ([]SetElement, error) { + return cc.getSetElements(s, nil, true) }