Compare commits

..

1 Commits

Author SHA1 Message Date
Antonio Ojea a84c81bda2
Merge 1e48c1007e into 385f80f4ef 2025-02-26 18:00:08 +01:00
2 changed files with 74 additions and 115 deletions

View File

@ -3865,58 +3865,7 @@ func TestIP6SetAddElements(t *testing.T) {
t.Errorf("c.GetSetElements(portSet) failed: %v", err)
}
if len(elements) != 2 {
t.Fatalf("len(portSetElements) = %d, want 2", len(elements))
}
}
func TestSetElementBatching(t *testing.T) {
// Create a new network namespace to test these operations,
// and tear down the namespace at test completion.
c, newNS := nftest.OpenSystemConn(t, *enableSysTests)
defer nftest.CleanupSystemConn(t, newNS)
// Clear all rules at the beginning + end of the test.
c.FlushRuleset()
defer c.FlushRuleset()
filter := c.AddTable(&nftables.Table{
Family: nftables.TableFamilyIPv4,
Name: "filter",
})
portSet := &nftables.Set{
Table: filter,
Name: "ports",
KeyType: nftables.TypeInetService,
}
// The 5000 elements will need to be split into 3 batches to make each batch
// fit into a message.
elements := make([]nftables.SetElement, 5000)
for i := range elements {
elements[i].Key = binaryutil.BigEndian.PutUint16(uint16(i))
elements[i].Comment = "0123456789"
}
if err := c.AddSet(portSet, elements); err != nil {
t.Errorf("c.AddSet(portSet) failed: %v", err)
}
if err := c.Flush(); err != nil {
t.Errorf("c.Flush() failed: %v", err)
}
gotElements, err := c.GetSetElements(portSet)
if err != nil {
t.Errorf("c.GetSetElements(portSet) failed: %v", err)
}
if len(gotElements) != len(elements) {
t.Errorf("len(gotElements) = %d, want %d", len(gotElements), len(elements))
}
gotNumbers := make([]bool, len(elements))
for _, element := range gotElements {
gotNumbers[binaryutil.BigEndian.Uint16(element.Key)] = true
}
for i := range gotNumbers {
if !gotNumbers[i] {
t.Errorf("Missing element %d", i)
break
}
t.Fatalf("len(portSetElements) = %d, want 2", len(sets))
}
}

124
set.go
View File

@ -18,7 +18,6 @@ import (
"encoding/binary"
"errors"
"fmt"
"math"
"strings"
"time"
@ -380,31 +379,24 @@ func (cc *Conn) SetAddElements(s *Set, vals []SetElement) error {
if s.Anonymous {
return errors.New("anonymous sets cannot be updated")
}
return cc.appendElemList(s, vals, unix.NFT_MSG_NEWSETELEM)
}
// SetDeleteElements deletes data points from an nftables set.
func (cc *Conn) SetDeleteElements(s *Set, vals []SetElement) error {
cc.mu.Lock()
defer cc.mu.Unlock()
if s.Anonymous {
return errors.New("anonymous sets cannot be updated")
}
return cc.appendElemList(s, vals, unix.NFT_MSG_DELSETELEM)
elements, err := s.makeElemList(vals, s.ID)
if err != nil {
return err
}
cc.messages = append(cc.messages, netlink.Message{
Header: netlink.Header{
Type: netlink.HeaderType((unix.NFNL_SUBSYS_NFTABLES << 8) | unix.NFT_MSG_NEWSETELEM),
Flags: netlink.Request | netlink.Acknowledge | netlink.Create,
},
Data: append(extraHeader(uint8(s.Table.Family), 0), cc.marshalAttr(elements)...),
})
// maxElemBatchSize is the maximum size in bytes of encoded set elements which
// are sent in one netlink message. The size field of a netlink attribute is a
// uint16, and 1024 bytes is more than enough for the per-message headers.
const maxElemBatchSize = math.MaxUint16 - 1024
func (cc *Conn) appendElemList(s *Set, vals []SetElement, hdrType uint16) error {
if len(vals) == 0 {
return nil
}
func (s *Set) makeElemList(vals []SetElement, id uint32) ([]netlink.Attribute, error) {
var elements []netlink.Attribute
batchSize := 0
var batches [][]netlink.Attribute
for i, v := range vals {
item := make([]netlink.Attribute, 0)
@ -416,14 +408,14 @@ func (cc *Conn) appendElemList(s *Set, vals []SetElement, hdrType uint16) error
encodedKey, err := netlink.MarshalAttributes([]netlink.Attribute{{Type: unix.NFTA_DATA_VALUE, Data: v.Key}})
if err != nil {
return fmt.Errorf("marshal key %d: %v", i, err)
return nil, fmt.Errorf("marshal key %d: %v", i, err)
}
item = append(item, netlink.Attribute{Type: unix.NFTA_SET_ELEM_KEY | unix.NLA_F_NESTED, Data: encodedKey})
if len(v.KeyEnd) > 0 {
encodedKeyEnd, err := netlink.MarshalAttributes([]netlink.Attribute{{Type: unix.NFTA_DATA_VALUE, Data: v.KeyEnd}})
if err != nil {
return fmt.Errorf("marshal key end %d: %v", i, err)
return nil, fmt.Errorf("marshal key end %d: %v", i, err)
}
item = append(item, netlink.Attribute{Type: NFTA_SET_ELEM_KEY_END | unix.NLA_F_NESTED, Data: encodedKeyEnd})
}
@ -443,7 +435,7 @@ func (cc *Conn) appendElemList(s *Set, vals []SetElement, hdrType uint16) error
{Type: unix.NFTA_DATA_VALUE, Data: binaryutil.BigEndian.PutUint32(uint32(v.VerdictData.Kind))},
})
if err != nil {
return fmt.Errorf("marshal item %d: %v", i, err)
return nil, fmt.Errorf("marshal item %d: %v", i, err)
}
encodedVal = append(encodedVal, encodedKind...)
if len(v.VerdictData.Chain) != 0 {
@ -451,21 +443,21 @@ func (cc *Conn) appendElemList(s *Set, vals []SetElement, hdrType uint16) error
{Type: unix.NFTA_SET_ELEM_DATA, Data: []byte(v.VerdictData.Chain + "\x00")},
})
if err != nil {
return fmt.Errorf("marshal item %d: %v", i, err)
return nil, fmt.Errorf("marshal item %d: %v", i, err)
}
encodedVal = append(encodedVal, encodedChain...)
}
encodedVerdict, err := netlink.MarshalAttributes([]netlink.Attribute{
{Type: unix.NFTA_SET_ELEM_DATA | unix.NLA_F_NESTED, Data: encodedVal}})
if err != nil {
return fmt.Errorf("marshal item %d: %v", i, err)
return nil, fmt.Errorf("marshal item %d: %v", i, err)
}
item = append(item, netlink.Attribute{Type: unix.NFTA_SET_ELEM_DATA | unix.NLA_F_NESTED, Data: encodedVerdict})
case len(v.Val) > 0:
// Since v.Val's length is not 0 then, v is a regular map element, need to add to the attributes
encodedVal, err := netlink.MarshalAttributes([]netlink.Attribute{{Type: unix.NFTA_DATA_VALUE, Data: v.Val}})
if err != nil {
return fmt.Errorf("marshal item %d: %v", i, err)
return nil, fmt.Errorf("marshal item %d: %v", i, err)
}
item = append(item, netlink.Attribute{Type: unix.NFTA_SET_ELEM_DATA | unix.NLA_F_NESTED, Data: encodedVal})
@ -481,42 +473,22 @@ func (cc *Conn) appendElemList(s *Set, vals []SetElement, hdrType uint16) error
encodedItem, err := netlink.MarshalAttributes(item)
if err != nil {
return fmt.Errorf("marshal item %d: %v", i, err)
}
itemSize := unix.NLA_HDRLEN + len(encodedItem)
if batchSize+itemSize > maxElemBatchSize {
batches = append(batches, elements)
elements = nil
batchSize = 0
return nil, fmt.Errorf("marshal item %d: %v", i, err)
}
elements = append(elements, netlink.Attribute{Type: uint16(i+1) | unix.NLA_F_NESTED, Data: encodedItem})
batchSize += itemSize
}
batches = append(batches, elements)
for _, batch := range batches {
encodedElem, err := netlink.MarshalAttributes(batch)
encodedElem, err := netlink.MarshalAttributes(elements)
if err != nil {
return fmt.Errorf("marshal elements: %v", err)
return nil, fmt.Errorf("marshal elements: %v", err)
}
message := []netlink.Attribute{
{Type: unix.NFTA_SET_ELEM_LIST_SET, Data: []byte(s.Name + "\x00")},
{Type: unix.NFTA_SET_ELEM_LIST_SET_ID, Data: binaryutil.BigEndian.PutUint32(s.ID)},
{Type: unix.NFTA_SET_ELEM_LIST_TABLE, Data: []byte(s.Table.Name + "\x00")},
return []netlink.Attribute{
{Type: unix.NFTA_SET_NAME, Data: []byte(s.Name + "\x00")},
{Type: unix.NFTA_LOOKUP_SET_ID, Data: binaryutil.BigEndian.PutUint32(id)},
{Type: unix.NFTA_SET_TABLE, Data: []byte(s.Table.Name + "\x00")},
{Type: unix.NFTA_SET_ELEM_LIST_ELEMENTS | unix.NLA_F_NESTED, Data: encodedElem},
}
cc.messages = append(cc.messages, netlink.Message{
Header: netlink.Header{
Type: netlink.HeaderType((unix.NFNL_SUBSYS_NFTABLES << 8) | hdrType),
Flags: netlink.Request | netlink.Acknowledge | netlink.Create,
},
Data: append(extraHeader(uint8(s.Table.Family), 0), cc.marshalAttr(message)...),
})
}
return nil
}, nil
}
// AddSet adds the specified Set.
@ -692,7 +664,22 @@ func (cc *Conn) AddSet(s *Set, vals []SetElement) error {
})
// Set the values of the set if initial values were provided.
return cc.appendElemList(s, vals, unix.NFT_MSG_NEWSETELEM)
if len(vals) > 0 {
hdrType := unix.NFT_MSG_NEWSETELEM
elements, err := s.makeElemList(vals, s.ID)
if err != nil {
return err
}
cc.messages = append(cc.messages, netlink.Message{
Header: netlink.Header{
Type: netlink.HeaderType((unix.NFNL_SUBSYS_NFTABLES << 8) | hdrType),
Flags: netlink.Request | netlink.Acknowledge | netlink.Create,
},
Data: append(extraHeader(uint8(s.Table.Family), 0), cc.marshalAttr(elements)...),
})
}
return nil
}
// DelSet deletes a specific set, along with all elements it contains.
@ -712,6 +699,29 @@ func (cc *Conn) DelSet(s *Set) {
})
}
// SetDeleteElements deletes data points from an nftables set.
func (cc *Conn) SetDeleteElements(s *Set, vals []SetElement) error {
cc.mu.Lock()
defer cc.mu.Unlock()
if s.Anonymous {
return errors.New("anonymous sets cannot be updated")
}
elements, err := s.makeElemList(vals, s.ID)
if err != nil {
return err
}
cc.messages = append(cc.messages, netlink.Message{
Header: netlink.Header{
Type: netlink.HeaderType((unix.NFNL_SUBSYS_NFTABLES << 8) | unix.NFT_MSG_DELSETELEM),
Flags: netlink.Request | netlink.Acknowledge | netlink.Create,
},
Data: append(extraHeader(uint8(s.Table.Family), 0), cc.marshalAttr(elements)...),
})
return nil
}
// FlushSet deletes all data points from an nftables set.
func (cc *Conn) FlushSet(s *Set) {
cc.mu.Lock()
@ -967,8 +977,8 @@ func (cc *Conn) GetSetElements(s *Set) ([]SetElement, error) {
defer func() { _ = closer() }()
data, err := netlink.MarshalAttributes([]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")},
{Type: unix.NFTA_SET_TABLE, Data: []byte(s.Table.Name + "\x00")},
{Type: unix.NFTA_SET_NAME, Data: []byte(s.Name + "\x00")},
})
if err != nil {
return nil, err