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