Add ResetSetElements method (#345)

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.

Author: nickgarlis

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)
+		}
+	}
+}

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)
 }