Merge dc9df31dfa into 8a8ad2be81

This is an attempt to port the logic that nftables uses to automatically
adjust the recvmsg & sndmsg buffer sizes. The implementation of setting
sndmsg size is not the same as nftables due to some limitations in the
underlying netlink & socket libraries.

- https://git.netfilter.org/nftables/tree/src/mnl.c?id=713592c6008a8c589a00d3d3d2e49709ff2de62c#n262
- https://git.netfilter.org/libmnl/tree/include/libmnl/libmnl.h?id=03da98bcd284d55212bc79e91dfb63da0ef7b937#n20
- https://git.netfilter.org/nftables/tree/src/mnl.c?id=713592c6008a8c589a00d3d3d2e49709ff2de62c#n391
- https://git.netfilter.org/nftables/tree/src/mnl.c?id=713592c6008a8c589a00d3d3d2e49709ff2de62c#n426

conn.go

@@ -250,6 +250,15 @@ func (cc *Conn) Flush() error {
 	}
 	defer func() { _ = closer() }()
 
+	err = cc.setWriteBuffer(conn)
+	if err != nil {
+		return err
+	}
+	err = cc.setReadBuffer(conn)
+	if err != nil {
+		return err
+	}
+
 	messages, err := conn.SendMessages(batch(cc.messages))
 	if err != nil {
 		return fmt.Errorf("SendMessages: %w", err)
@@ -423,3 +432,84 @@ func (cc *Conn) allocateTransactionID() uint32 {
 	}
 	return cc.lastID
 }
+
+// getMessageSize returns the total size of all messages in the buffer.
+func (cc *Conn) getMessageSize() int {
+	var total int
+	for _, msg := range cc.messages {
+		total += len(msg.Data) + 16 // 16 bytes for the header
+	}
+	return total
+}
+
+// setWriteBuffer automatically sets the write buffer of the given connection to
+// the accumulated message size. This is only done if the current write buffer
+// is smaller than the message size.
+//
+// nftables actually handles this differently, it multiplies the number of
+// iovec entries by 2MB. This is not possible to do here as our underlying
+// netlink and socket libraries will only add a single iovec entry and
+// won't expose the number of entries.
+// https://git.netfilter.org/nftables/tree/src/mnl.c?id=713592c6008a8c589a00d3d3d2e49709ff2de62c#n262
+//
+// TODO: Update this function to mimic the behavior of nftables once those
+// limitations are no longer present.
+func (cc *Conn) setWriteBuffer(conn *netlink.Conn) error {
+	messageSize := cc.getMessageSize()
+	writeBuffer, err := conn.WriteBuffer()
+	if err != nil {
+		return err
+	}
+	if writeBuffer < messageSize {
+		return conn.SetWriteBuffer(messageSize)
+	}
+
+	return nil
+}
+
+// getDefaultEchoReadBuffer returns the minimum read buffer size for batches
+// with echo messages.
+//
+// See https://git.netfilter.org/libmnl/tree/include/libmnl/libmnl.h?id=03da98bcd284d55212bc79e91dfb63da0ef7b937#n20
+// and https://git.netfilter.org/nftables/tree/src/mnl.c?id=713592c6008a8c589a00d3d3d2e49709ff2de62c#n391
+func (cc *Conn) getDefaultEchoReadBuffer() int {
+	pageSize := os.Getpagesize()
+	if pageSize < 8192 {
+		return pageSize * 1024
+	}
+
+	return 8192 * 1024
+}
+
+// setReadBuffer  automatically sets the read buffer of the given connection
+// to the required size. This is only done if the current read buffer is smaller
+// than the required size.
+//
+// See https://git.netfilter.org/nftables/tree/src/mnl.c?id=713592c6008a8c589a00d3d3d2e49709ff2de62c#n426
+func (cc *Conn) setReadBuffer(conn *netlink.Conn) error {
+	var bufferSize int
+
+	// If there are any messages with the Echo flag, we initialize the buffer size
+	// to the default echo read buffer size.
+	for _, msg := range cc.messages {
+		if msg.Header.Flags&netlink.Echo == 0 {
+			bufferSize = cc.getDefaultEchoReadBuffer()
+			break
+		}
+	}
+
+	// Just like nftables, we allocate 1024 bytes for each message in the batch.
+	requiredSize := len(cc.messages) * 1024
+	if bufferSize < requiredSize {
+		bufferSize = requiredSize
+	}
+
+	currSize, err := conn.ReadBuffer()
+	if err != nil {
+		return err
+	}
+	if currSize < bufferSize {
+		return conn.SetReadBuffer(bufferSize)
+	}
+	return nil
+}

expr/ct.go

@@ -110,6 +110,12 @@ const (
 	CtStateUDPREPLIED
 )
 
+const (
+	// https://git.netfilter.org/libnftnl/tree/src/expr/ct.c?id=116e95aa7b6358c917de8c69f6f173874030b46b#n31
+	CtDirOriginal = iota
+	CtDirReply
+)
+
 // https://git.netfilter.org/libnftnl/tree/src/obj/ct_timeout.c?id=116e95aa7b6358c917de8c69f6f173874030b46b#n57
 var CtStateUDPTimeoutDefaults CtStatePolicyTimeout = map[uint16]uint32{
 	CtStateUDPUNREPLIED: 30,
@@ -122,6 +128,7 @@ type Ct struct {
 	SourceRegister bool
 	Key            CtKey
 	Direction      uint32
+	OptDirection   bool
 }
 
 func (e *Ct) marshal(fam byte) ([]byte, error) {
@@ -165,10 +172,16 @@ func (e *Ct) marshalData(fam byte) ([]byte, error) {
 	exprData = append(exprData, regData...)
 
 	switch e.Key {
+	case CtKeyPKTS, CtKeyBYTES, CtKeyAVGPKT, CtKeyL3PROTOCOL, CtKeyPROTOCOL:
+		if !e.OptDirection {
+			break
+		}
+
+		fallthrough
 	case CtKeySRC, CtKeyDST, CtKeyPROTOSRC, CtKeyPROTODST, CtKeySRCIP, CtKeyDSTIP, CtKeySRCIP6, CtKeyDSTIP6:
 		regData, err = netlink.MarshalAttributes(
 			[]netlink.Attribute{
-				{Type: unix.NFTA_CT_DIRECTION, Data: binaryutil.BigEndian.PutUint32(e.Direction)},
+				{Type: unix.NFTA_CT_DIRECTION, Data: []byte{uint8(e.Direction)}},
 			},
 		)
 		if err != nil {
@@ -186,6 +199,8 @@ func (e *Ct) unmarshal(fam byte, data []byte) error {
 		return err
 	}
 	ad.ByteOrder = binary.BigEndian
+
+	var hasDirection bool
 	for ad.Next() {
 		switch ad.Type() {
 		case unix.NFTA_CT_KEY:
@@ -193,12 +208,19 @@ func (e *Ct) unmarshal(fam byte, data []byte) error {
 		case unix.NFTA_CT_DREG:
 			e.Register = ad.Uint32()
 		case unix.NFTA_CT_DIRECTION:
-			e.Direction = ad.Uint32()
+			e.Direction = uint32(ad.Uint8())
+			hasDirection = true
 		case unix.NFTA_CT_SREG:
 			e.SourceRegister = true
 			e.Register = ad.Uint32()
 		}
 	}
+
+	switch e.Key {
+	case CtKeyPKTS, CtKeyBYTES, CtKeyAVGPKT, CtKeyL3PROTOCOL, CtKeyPROTOCOL:
+		e.OptDirection = hasDirection
+	}
+
 	return ad.Err()
 }
 

expr/ct_test.go

@@ -78,6 +78,22 @@ func TestCt(t *testing.T) {
 				Direction: 1,
 			},
 		},
+		{
+			name: "Unmarshal Ct packets direction original case",
+			ct: Ct{
+				Register:     1,
+				Key:          CtKeyPKTS,
+				Direction:    CtDirOriginal,
+				OptDirection: true,
+			},
+		},
+		{
+			name: "Unmarshal Ct bytes without direction case",
+			ct: Ct{
+				Register: 1,
+				Key:      CtKeyBYTES,
+			},
+		},
 	}
 
 	for _, tt := range tests {

nftables_test.go

@@ -7396,3 +7396,37 @@ func TestSetElementComment(t *testing.T) {
 		}
 	}
 }
+
+func TestAutoBufferSize(t *testing.T) {
+	conn, newNS := nftest.OpenSystemConn(t, *enableSysTests)
+	defer nftest.CleanupSystemConn(t, newNS)
+	conn.FlushRuleset()
+	defer conn.FlushRuleset()
+
+	table := conn.AddTable(&nftables.Table{
+		Family: nftables.TableFamilyIPv4,
+		Name:   "test-table",
+	})
+
+	chain := conn.AddChain(&nftables.Chain{
+		Name:  "test-chain",
+		Table: table,
+	})
+
+	for i := 0; i < 4096; i++ {
+		conn.AddRule(&nftables.Rule{
+			Table: table,
+			Chain: chain,
+			Exprs: []expr.Any{
+				&expr.Verdict{
+					Kind: expr.VerdictAccept,
+				},
+			},
+		})
+	}
+
+	err := conn.Flush()
+	if err != nil {
+		t.Fatalf("failed to flush: %v", err)
+	}
+}