Merge a0423c9897 into 0420ffbf57

gen.go

@@ -3,7 +3,6 @@ package nftables
 import (
 	"encoding/binary"
 	"fmt"
-
 	"github.com/mdlayher/netlink"
 	"golang.org/x/sys/unix"
 )
@@ -14,7 +13,7 @@ type GenMsg struct {
 	ProcComm string // [16]byte - max 16bytes - kernel TASK_COMM_LEN
 }
 
-const genHeaderType = netlink.HeaderType((unix.NFNL_SUBSYS_NFTABLES << 8) | unix.NFT_MSG_NEWGEN)
+var genHeaderType = netlink.HeaderType((unix.NFNL_SUBSYS_NFTABLES << 8) | unix.NFT_MSG_NEWGEN)
 
 func genFromMsg(msg netlink.Message) (*GenMsg, error) {
 	if got, want := msg.Header.Type, genHeaderType; got != want {

nftables_test.go

@@ -622,14 +622,6 @@ func TestMasqMarshalUnmarshal(t *testing.T) {
 		Table: filter,
 		Chain: postrouting,
 		Exprs: []expr.Any{
-			&expr.Immediate{
-				Register: min,
-				Data:     binaryutil.BigEndian.PutUint16(4070),
-			},
-			&expr.Immediate{
-				Register: max,
-				Data:     binaryutil.BigEndian.PutUint16(4090),
-			},
 			&expr.Masq{
 				ToPorts:     true,
 				RegProtoMin: min,
@@ -660,13 +652,13 @@ func TestMasqMarshalUnmarshal(t *testing.T) {
 	}
 
 	rule := rules[0]
-	if got, want := len(rule.Exprs), 3; got != want {
+	if got, want := len(rule.Exprs), 1; got != want {
 		t.Fatalf("unexpected number of exprs: got %d, want %d", got, want)
 	}
 
-	me, ok := rule.Exprs[2].(*expr.Masq)
+	me, ok := rule.Exprs[0].(*expr.Masq)
 	if !ok {
-		t.Fatalf("unexpected expression type: got %T, want *expr.Masq", rule.Exprs[2])
+		t.Fatalf("unexpected expression type: got %T, want *expr.Masq", rule.Exprs[0])
 	}
 
 	if got, want := me.ToPorts, true; got != want {
@@ -3865,58 +3857,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))
+		t.Fatalf("len(portSetElements) = %d, want 2", len(sets))
-	}
-}
-
-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
-		}
 	}
 }
 

obj.go

@@ -25,7 +25,7 @@ import (
 	"golang.org/x/sys/unix"
 )
 
-const (
+var (
 	newObjHeaderType = netlink.HeaderType((unix.NFNL_SUBSYS_NFTABLES << 8) | unix.NFT_MSG_NEWOBJ)
 	delObjHeaderType = netlink.HeaderType((unix.NFNL_SUBSYS_NFTABLES << 8) | unix.NFT_MSG_DELOBJ)
 )

rule.go

@@ -25,7 +25,7 @@ import (
 	"golang.org/x/sys/unix"
 )
 
-const (
+var (
 	newRuleHeaderType = netlink.HeaderType((unix.NFNL_SUBSYS_NFTABLES << 8) | unix.NFT_MSG_NEWRULE)
 	delRuleHeaderType = netlink.HeaderType((unix.NFNL_SUBSYS_NFTABLES << 8) | unix.NFT_MSG_DELRULE)
 )

set.go

@@ -18,7 +18,6 @@ import (
 	"encoding/binary"
 	"errors"
 	"fmt"
-	"math"
 	"strings"
 	"time"
 
@@ -167,9 +166,7 @@ var (
 		TypeTimeDay,
 		TypeCGroupV2,
 	}
-)
 
-const (
 	// ctLabelBitSize is defined in https://git.netfilter.org/nftables/tree/src/ct.c.
 	ctLabelBitSize uint32 = 128
 
@@ -380,31 +377,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)
+
+	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)...),
+	})
+
+	return nil
 }
 
-// SetDeleteElements deletes data points from an nftables set.
+func (s *Set) makeElemList(vals []SetElement, id uint32) ([]netlink.Attribute, error) {
-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)
-}
-
-// 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
-	}
 	var elements []netlink.Attribute
-	batchSize := 0
-	var batches [][]netlink.Attribute
 
 	for i, v := range vals {
 		item := make([]netlink.Attribute, 0)
@@ -416,14 +406,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 +433,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 +441,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 +471,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)
+			return nil, 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
 		}
 		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(elements)
-		encodedElem, err := netlink.MarshalAttributes(batch)
+	if err != nil {
-		if err != nil {
+		return nil, fmt.Errorf("marshal elements: %v", err)
-			return 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")},
-			{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
+
+	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},
+	}, nil
 }
 
 // AddSet adds the specified Set.
@@ -692,7 +662,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 +697,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()
@@ -729,7 +737,7 @@ func (cc *Conn) FlushSet(s *Set) {
 	})
 }
 
-const (
+var (
 	newSetHeaderType = netlink.HeaderType((unix.NFNL_SUBSYS_NFTABLES << 8) | unix.NFT_MSG_NEWSET)
 	delSetHeaderType = netlink.HeaderType((unix.NFNL_SUBSYS_NFTABLES << 8) | unix.NFT_MSG_DELSET)
 )
@@ -829,7 +837,7 @@ func parseSetDatatype(magic uint32) (SetDatatype, error) {
 	return dt, nil
 }
 
-const (
+var (
 	newElemHeaderType = netlink.HeaderType((unix.NFNL_SUBSYS_NFTABLES << 8) | unix.NFT_MSG_NEWSETELEM)
 	delElemHeaderType = netlink.HeaderType((unix.NFNL_SUBSYS_NFTABLES << 8) | unix.NFT_MSG_DELSETELEM)
 )
@@ -967,8 +975,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_TABLE, Data: []byte(s.Table.Name + "\x00")},
-		{Type: unix.NFTA_SET_ELEM_LIST_SET, Data: []byte(s.Name + "\x00")},
+		{Type: unix.NFTA_SET_NAME, Data: []byte(s.Name + "\x00")},
 	})
 	if err != nil {
 		return nil, err

table.go

@@ -21,7 +21,7 @@ import (
 	"golang.org/x/sys/unix"
 )
 
-const (
+var (
 	newTableHeaderType = netlink.HeaderType((unix.NFNL_SUBSYS_NFTABLES << 8) | unix.NFT_MSG_NEWTABLE)
 	delTableHeaderType = netlink.HeaderType((unix.NFNL_SUBSYS_NFTABLES << 8) | unix.NFT_MSG_DELTABLE)
 )