Ensure socket buffer draining during Flush (#330)

This commit re-implements response handling for batch messages.
The current implementation uses heuristics to drain the socket
buffer by counting sent messages and making a blocking recvmsg
call for each one. While this generally works, it has several
problems:

  - No kernel guarantees: The kernel makes no guarantee about
    the number of ACKs returned. Error conditions like EPERM,
    ENOBUFS, and ENOMEM can cause unexpected behavior.
  - Vulnerable to kernel bugs: As described in #329
  - Complex code: ACK counting complicates the response
    handling logic in Flush

The new approach follows nft's pattern by leveraging netlink's
synchronous request-response behavior. When a blocking sendmsg
completes, the kernel has already processed the request and
queued any response data in the socket's receive buffer. This
means we can simply check if data is available rather than
blocking indefinitely waiting for a specific number of
responses.

The implementation uses pselect6 (the same syscall nft uses)
to poll the socket for readability.

Fixes: #329

Author: nickgarlis

conn.go

@@ -20,7 +20,6 @@ import (
 	"math"
 	"os"
 	"sync"
-	"syscall"
 
 	"github.com/google/nftables/binaryutil"
 	"github.com/google/nftables/expr"
@@ -262,71 +261,100 @@ func (cc *Conn) flush(genID uint32) error {
 		return err
 	}
 
-	messages, err := conn.SendMessages(batch)
+	sentMsgs, err := conn.SendMessages(batch)
 	if err != nil {
 		return fmt.Errorf("SendMessages: %w", err)
 	}
 
 	var errs error
 
-	// Fetch replies. Each message with the Echo flag triggers a reply of the same
-	// type. Additionally, if the first message of the batch has the Echo flag, we
-	// get a reply of type NFT_MSG_NEWGEN, which we ignore.
-	replyIndex := 0
-	for replyIndex < len(cc.messages) && cc.messages[replyIndex].Header.Flags&netlink.Echo == 0 {
-		replyIndex++
-	}
-	replies, err := conn.Receive()
-	for err == nil && len(replies) != 0 {
-		reply := replies[0]
-		if reply.Header.Type == netlink.Error && reply.Header.Sequence == messages[1].Header.Sequence {
-			// The next message is the acknowledgement for the first message in the
-			// batch; stop looking for replies.
+	seqToMsgMap := cc.getSeqToMsgMap(sentMsgs)
+
+	for {
+		ready, err := cc.isReadReady(conn)
+		if err != nil {
+			return err
+		}
+
+		// Since SendMessages is blocking and netlink communication is synchronous,
+		// the kernel has already processed the request and queued any responses by
+		// the time SendMessages returns. Therefore, if isReadReady returns false on
+		// the first call, it means there are no messages coming at all and we can
+		// safely exit.
+		if !ready {
 			break
-		} else if replyIndex < len(cc.messages) {
-			msg := messages[replyIndex+1]
-			if msg.Header.Sequence == reply.Header.Sequence && msg.Header.Type == reply.Header.Type {
-				// The only messages which set the echo flag are rule create messages.
-				err := cc.messages[replyIndex].rule.handleCreateReply(reply)
-				if err != nil {
-					errs = errors.Join(errs, err)
-				}
-				replyIndex++
-				for replyIndex < len(cc.messages) && cc.messages[replyIndex].Header.Flags&netlink.Echo == 0 {
-					replyIndex++
-				}
-			}
 		}
-		replies = replies[1:]
-		if len(replies) == 0 {
-			replies, err = conn.Receive()
+
+		replies, err := conn.Receive()
+		if err != nil {
+			errs = errors.Join(errs, fmt.Errorf("receive: %w", err))
 		}
-	}
 
-	// Fetch the requested acknowledgement for each message we sent.
-	for i := range cc.messages {
-		if i != 0 {
-			_, err = conn.Receive()
+		if len(replies) == 0 && cc.TestDial != nil {
+			// When using a test dial function, we don't always get a reply for each
+			// sent message. Additionally, there is no buffer to poll for more data,
+			// so we stop here.
+			break
 		}
-		if err != nil {
-			if errors.Is(err, os.ErrPermission) || errors.Is(err, syscall.ENOBUFS) || errors.Is(err, syscall.ENOMEM) {
-				// Kernel will only send one error to user space.
-				return err
+
+		for _, reply := range replies {
+			if err := cc.handleEchoReply(seqToMsgMap, reply); err != nil {
+				errs = errors.Join(errs, err)
 			}
-			errs = errors.Join(errs, err)
 		}
 	}
 
 	if errs != nil {
-		return fmt.Errorf("conn.Receive: %w", errs)
-	}
-	if replyIndex < len(cc.messages) {
-		return fmt.Errorf("missing reply for message %d in batch", replyIndex)
+		return errs
 	}
 
 	return nil
 }
 
+// getSeqToMsgMap returns a map of the cc.messages that were sent, indexed by
+// their sequence number as included in the sent netlink messages. The returned
+// map will not include the batch begin and end messages.
+func (cc *Conn) getSeqToMsgMap(sentMsgs []netlink.Message) map[uint32]netlinkMessage {
+	seqToMsgMap := make(map[uint32]netlinkMessage)
+	for i, msg := range sentMsgs {
+		if i == 0 || i == len(sentMsgs)-1 {
+			// Skip batch begin and end messages.
+			continue
+		}
+		if i-1 >= len(cc.messages) {
+			// Should not happen, but be defensive.
+			break
+		}
+		// Update the header in the original message, as the sequence number
+		// and possibly other fields have been updated by the the underlying
+		// netlink library.
+		cc.messages[i-1].Header = msg.Header
+		seqToMsgMap[msg.Header.Sequence] = cc.messages[i-1]
+	}
+
+	return seqToMsgMap
+}
+
+func (cc *Conn) handleEchoReply(seqToMsgMap map[uint32]netlinkMessage, reply netlink.Message) error {
+	sentMsg, ok := seqToMsgMap[reply.Header.Sequence]
+	if !ok {
+		// We don't have a record of sending this message, ignore.
+		return nil
+	}
+
+	if sentMsg.Header.Flags&netlink.Echo == 0 {
+		return nil
+	}
+
+	switch reply.Header.Type {
+	case newRuleHeaderType:
+		// The only messages which set the echo flag are rule create messages.
+		return sentMsg.rule.handleCreateReply(reply)
+	default:
+		return nil
+	}
+}
+
 // FlushRuleset flushes the entire ruleset. See also
 // https://wiki.nftables.org/wiki-nftables/index.php/Operations_at_ruleset_level
 func (cc *Conn) FlushRuleset() {

internal/nftest/util.go

@@ -0,0 +1,38 @@
+package nftest
+
+import (
+	"fmt"
+	"math"
+	"runtime"
+
+	"golang.org/x/sys/unix"
+)
+
+// AsUnprivileged temporarily drops the effective UID to an unprivileged
+// value (65535) while executing the provided function. It requires the
+// process to be running as root to be able to regain privileges afterwards.
+func AsUnprivileged(fn func() error) error {
+	runtime.LockOSThread()
+	defer runtime.UnlockOSThread()
+
+	targetUID := math.MaxUint16
+	_, euid, suid := unix.Getresuid()
+
+	if euid != 0 && suid != 0 {
+		return fmt.Errorf("must be run as root to regain privileges (euid=%d suid=%d)", euid, suid)
+	}
+
+	// Drop privileges by changing only the effective UID
+	if err := unix.Setresuid(-1, targetUID, -1); err != nil {
+		return fmt.Errorf("failed to drop effective UID to %d: %w", targetUID, err)
+	}
+
+	// Restore when done
+	defer func() {
+		if err := unix.Setresuid(-1, euid, -1); err != nil {
+			panic(fmt.Sprintf("failed to restore euid=%d: %v", euid, err))
+		}
+	}()
+
+	return fn()
+}

nftables_test.go

@@ -34,6 +34,7 @@ import (
 	"github.com/google/nftables/internal/nftest"
 	"github.com/google/nftables/xt"
 	"github.com/mdlayher/netlink"
+	"github.com/vishvananda/netns"
 	"golang.org/x/sys/unix"
 )
 
@@ -7947,3 +7948,82 @@ func TestGetPortID(t *testing.T) {
 		t.Fatalf("conn.GetPortID() returned invalid port ID: %d", pid)
 	}
 }
+
+func TestSocketDrainingOnErrors(t *testing.T) {
+	tests := []struct {
+		name          string
+		setupError    func(conn *nftables.Conn, ns netns.NsHandle) error
+		expectedError error
+		description   string
+	}{
+		{
+			name: "short_circuited_error",
+			setupError: func(conn *nftables.Conn, ns netns.NsHandle) error {
+				// Add multiple tables but trigger EPERM before kernel processes all
+				conn.AddTable(&nftables.Table{Name: "table1", Family: nftables.TableFamilyIPv4})
+				conn.AddTable(&nftables.Table{Name: "table2", Family: nftables.TableFamilyIPv4})
+				conn.AddTable(&nftables.Table{Name: "table3", Family: nftables.TableFamilyIPv4})
+
+				// Drop privileges to trigger immediate EPERM (short-circuit)
+				return nftest.AsUnprivileged(conn.Flush)
+			},
+			expectedError: syscall.EPERM,
+			description:   "kernel returns EPERM immediately without processing all messages",
+		},
+		{
+			name: "non_short_circuited_error",
+			setupError: func(conn *nftables.Conn, ns netns.NsHandle) error {
+				// Use new connection to create an owned table
+				newConn, err := nftables.New(nftables.WithNetNSFd(int(ns)), nftables.AsLasting())
+				if err != nil {
+					return err
+				}
+				defer newConn.CloseLasting()
+
+				ownedTable := &nftables.Table{
+					Name:   "owned-table",
+					Family: nftables.TableFamilyIPv4,
+					Flags:  nftables.TableFlagOwner | nftables.TableFlagPersist,
+				}
+				newConn.AddTable(ownedTable)
+				if err := newConn.Flush(); err != nil {
+					return err
+				}
+
+				// Use old connection to try deleting owned table (will fail)
+				conn.DelTable(ownedTable)
+				conn.AddTable(&nftables.Table{Name: "table2", Family: nftables.TableFamilyIPv4})
+				conn.AddTable(&nftables.Table{Name: "table3", Family: nftables.TableFamilyIPv4})
+
+				return conn.Flush()
+			},
+			expectedError: syscall.EPERM,
+			description:   "kernel processes all messages but returns EPERM for owned table deletion",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			_, newNS := nftest.OpenSystemConn(t, *enableSysTests)
+			conn, err := nftables.New(nftables.WithNetNSFd(int(newNS)), nftables.AsLasting())
+			if err != nil {
+				t.Fatalf("nftables.New() failed: %v", err)
+			}
+			defer nftest.CleanupSystemConn(t, newNS)
+			defer conn.FlushRuleset()
+			defer conn.CloseLasting()
+
+			err = tt.setupError(conn, newNS)
+
+			if err == nil || !errors.Is(err, tt.expectedError) {
+				t.Fatalf("expected %v error, got %v", tt.expectedError, err)
+			}
+
+			// Verify socket is properly drained. If not, this will fail as
+			// there will be leftover messages in the socket buffer.
+			if _, err := conn.ListTables(); err != nil {
+				t.Fatalf("ListTables failed after error - socket not properly drained: %v", err)
+			}
+		})
+	}
+}

socket.go

@@ -0,0 +1,39 @@
+package nftables
+
+import (
+	"fmt"
+
+	"github.com/mdlayher/netlink"
+	"golang.org/x/sys/unix"
+)
+
+// isReadReady reports whether the netlink connection is ready for reading.
+// It uses pselect6 with a zero timeout on the underlying raw connection.
+// This allows for an efficient check of socket readiness without blocking.
+// If the Conn was created with a TestDial function, it assumes readiness.
+func (cc *Conn) isReadReady(conn *netlink.Conn) (bool, error) {
+	if cc.TestDial != nil {
+		return true, nil
+	}
+
+	rawConn, err := conn.SyscallConn()
+	if err != nil {
+		return false, fmt.Errorf("get raw conn: %w", err)
+	}
+
+	var n int
+	var opErr error
+	err = rawConn.Control(func(fd uintptr) {
+		var readfds unix.FdSet
+		readfds.Zero()
+		readfds.Set(int(fd))
+
+		ts := &unix.Timespec{} // zero timeout: immediate return
+		n, opErr = unix.Pselect(int(fd)+1, &readfds, nil, nil, ts, nil)
+	})
+	if err != nil {
+		return false, err
+	}
+
+	return n > 0, opErr
+}