Implement Counter operation (parsing, interpretation, evaluation, tests).

Counter operation is thread-safe.

PiperOrigin-RevId: 673993509

pkg/tcpip/nftables/nftables.go

@@ -45,6 +45,7 @@ import (
 	"encoding/binary"
 	"fmt"
 	"slices"
+	"sync/atomic"
 
 	"gvisor.dev/gvisor/pkg/abi/linux"
 	"gvisor.dev/gvisor/pkg/tcpip/checksum"
@@ -52,8 +53,8 @@ import (
 	"gvisor.dev/gvisor/pkg/tcpip/stack"
 )
 
-// TODO(b/345684870): Remove unused functions once initial implementation is
-// complete.
+// TODO(b/345684870): Make the nftables package thread-safe! Must be done before
+// the package is used in production.
 
 // Defines general constants for the nftables interpreter.
 const (
@@ -581,6 +582,7 @@ var (
 	_ operation = (*payloadLoad)(nil)
 	_ operation = (*payloadSet)(nil)
 	_ operation = (*bitwise)(nil)
+	_ operation = (*counter)(nil)
 )
 
 // immediate is an operation that sets the data in a register.
@@ -1212,6 +1214,30 @@ func (op bitwise) evaluate(regs *registerSet, pkt *stack.PacketBuffer) {
 	}
 }
 
+// counter is an operation that increments a counter for the packets and number
+// of bytes each time the operation is evaluated.
+type counter struct {
+	// Must be thread-safe because data stored here is updated for each evaluation
+	// and evaluations can happen in parallel for processing multiple packets.
+
+	bytes   atomic.Int64 // Number of bytes that have passed through counter.
+	packets atomic.Int64 // Number of packets that have passed through counter.
+}
+
+// newCounter creates a new counter operation.
+func newCounter(startBytes, startPackets int64) *counter {
+	cntr := &counter{}
+	cntr.bytes.Store(startBytes)
+	cntr.packets.Store(startPackets)
+	return cntr
+}
+
+// evaluate for counter increments the counter for the packet and bytes.
+func (op *counter) evaluate(regs *registerSet, pkt *stack.PacketBuffer) {
+	op.bytes.Add(int64(pkt.Size()))
+	op.packets.Add(1)
+}
+
 //
 // Register and Register-Related Implementations.
 // Note: Registers are represented by type uint8 for the register number.

pkg/tcpip/nftables/nftables_test.go

@@ -2309,6 +2309,80 @@ func TestEvaluateBitwise(t *testing.T) {
 	}
 }
 
+// TestEvaluateCounter tests that the Counter operation correctly increments the
+// counter for the number of bytes and packets as it encounters packets.
+// Note: relies on expected behavior of Comparison and Payload Load operations.
+func TestEvaluateCounter(t *testing.T) {
+	// Creates a counter operation.
+	counter := newCounter(0, 0)
+	// Defines the packets to be used in the test.
+	desiredIpv4Address := tcpip.AddrFrom4(arbitraryIPv4AddrB)
+	countedIPv4Pkt := func() *stack.PacketBuffer {
+		fields := arbitraryIPv4Fields()
+		fields.SrcAddr = desiredIpv4Address
+		return makeIPv4Packet(header.IPv4MinimumSize, fields)
+	}
+	uncountedIPv4Pkt := func() *stack.PacketBuffer {
+		fields := arbitraryIPv4Fields()
+		fields.SrcAddr = tcpip.AddrFrom4(arbitraryIPv4AddrB2)
+		return makeIPv4Packet(header.IPv4MinimumSize, fields)
+	}
+	pkts := []*stack.PacketBuffer{countedIPv4Pkt(), uncountedIPv4Pkt(), countedIPv4Pkt(), countedIPv4Pkt(),
+		uncountedIPv4Pkt(), countedIPv4Pkt(), uncountedIPv4Pkt(), uncountedIPv4Pkt(), uncountedIPv4Pkt(), countedIPv4Pkt()}
+	t.Run("counter increment tests", func(t *testing.T) {
+		// Sets up an NFTables object with a base chain with policy accept.
+		nf := NewNFTables()
+		tab, err := nf.AddTable(arbitraryFamily, "test", "test table", false)
+		if err != nil {
+			t.Fatalf("unexpected error for AddTable: %v", err)
+		}
+		bc, err := tab.AddChain("base_chain", nil, "test chain", false)
+		if err != nil {
+			t.Fatalf("unexpected error for AddChain: %v", err)
+		}
+		bc.SetBaseChainInfo(arbitraryInfoPolicyAccept)
+
+		// Creates a rule that filters for the desired IPv4 address and adds the
+		// counter to the end of the rule. So, the counter should only increment for
+		// packets that satisfy the comparison.
+		rule := &Rule{}
+		rule.addOperation(mustCreatePayloadLoad(t, linux.NFT_PAYLOAD_NETWORK_HEADER, ipv4SrcAddrOffset, ipv4SrcAddrLen, linux.NFT_REG_1))
+		rule.addOperation(mustCreateComparison(t, linux.NFT_REG_1, linux.NFT_CMP_EQ, desiredIpv4Address.AsSlice()))
+		rule.addOperation(counter)
+		if err := bc.RegisterRule(rule, -1); err != nil {
+			t.Fatalf("unexpected error for RegisterRule: %v", err)
+		}
+
+		// Runs evaluation for each packet and checks whether the counter has
+		// incremented correctly.
+		prevBytes := counter.bytes.Load()
+		prevPackets := counter.packets.Load()
+		for i, pkt := range pkts {
+			_, err := nf.EvaluateHook(arbitraryFamily, arbitraryHook, pkt)
+			if err != nil {
+				t.Fatalf("unexpected error for EvaluateHook for packet %d: %v", i, err)
+			}
+			// Checks whether the counter should have incremented for the packet.
+			expectedDBytes, expectedDPackets := int64(0), int64(0)
+			if pkt.Network().SourceAddress() == desiredIpv4Address {
+				expectedDBytes, expectedDPackets = int64(pkt.Size()), 1
+			}
+			// Checks that the counter incremented correctly.
+			newBytes := counter.bytes.Load()
+			newPackets := counter.packets.Load()
+			if dBytes := newBytes - prevBytes; dBytes != expectedDBytes {
+				t.Fatalf("counter bytes incremented by %d for packet %d, expected %d", dBytes, i, expectedDBytes)
+			}
+			if dPackets := newPackets - prevPackets; dPackets != expectedDPackets {
+				t.Fatalf("counter packets incremented by %d for packet %d, expected %d", dPackets, i, expectedDPackets)
+			}
+			// Updates the previous values for the next packet.
+			prevBytes = newBytes
+			prevPackets = newPackets
+		}
+	})
+}
+
 // TestLoopCheckOnRegisterAndUnregister tests the loop checking and accompanying
 // logic on registering and unregistering rules.
 func TestLoopCheckOnRegisterAndUnregister(t *testing.T) {

pkg/tcpip/nftables/nftinterp.go

@@ -151,6 +151,8 @@ func InterpretOperation(line string, lnIdx int) (operation, error) {
 		// Assumes the bitwise operation is a boolean because interpretation of
 		// non-boolean operations is not supported from the nft binary debug output.
 		return InterpretBitwiseBool(line, lnIdx)
+	case "counter":
+		return InterpretCounter(line, lnIdx)
 	default:
 		return nil, &SyntaxError{lnIdx, 1, fmt.Sprintf("unrecognized operation type: %s", tokens[1])}
 	}
@@ -612,6 +614,60 @@ func InterpretBitwiseBool(line string, lnIdx int) (operation, error) {
 	return bitwiseBool, nil
 }
 
+// InterpretCounter creates a new Counter operation from the given string.
+func InterpretCounter(line string, lnIdx int) (operation, error) {
+	tokens := strings.Fields(line)
+
+	// Requires exactly 7 tokens:
+	// 		"[", "counter", "pkts", initial packets, "bytes", initial bytes, "]".
+	if len(tokens) != 7 {
+		return nil, &SyntaxError{lnIdx, 0, fmt.Sprintf("incorrect number of tokens for counter operation, should be exactly 7, got %d", len(tokens))}
+	}
+
+	if err := checkOperationBrackets(tokens, lnIdx); err != nil {
+		return nil, err
+	}
+
+	tkIdx := 1
+
+	// First token should be "counter".
+	if err := consumeToken("counter", tokens, lnIdx, tkIdx); err != nil {
+		return nil, err
+	}
+	tkIdx++
+
+	// Second token should be "pkts".
+	if err := consumeToken("pkts", tokens, lnIdx, tkIdx); err != nil {
+		return nil, err
+	}
+	tkIdx++
+
+	// Third token should be int64 representing initial packets.
+	initialPkts, err := strconv.ParseInt(tokens[tkIdx], 10, 64)
+	if err != nil {
+		return nil, &SyntaxError{lnIdx, tkIdx, fmt.Sprintf("could not parse int64 initial packets: '%s'", tokens[tkIdx])}
+	}
+	tkIdx++
+
+	// Fourth token should be "bytes".
+	if err := consumeToken("bytes", tokens, lnIdx, tkIdx); err != nil {
+		return nil, err
+	}
+	tkIdx++
+
+	// Fifth token should be int64 representing initial bytes.
+	initialBytes, err := strconv.ParseInt(tokens[tkIdx], 10, 64)
+	if err != nil {
+		return nil, &SyntaxError{lnIdx, tkIdx, fmt.Sprintf("could not parse int64 initial bytes: '%s'", tokens[tkIdx])}
+	}
+	tkIdx++
+
+	// Create the operation with the specified arguments.
+	cntr := newCounter(initialPkts, initialBytes)
+
+	return cntr, nil
+}
+
 //
 // Interpreter Helper Functions.
 //

pkg/tcpip/nftables/nftinterp_test.go

@@ -771,6 +771,43 @@ func checkBitwiseOp(tname string, expected operation, actual operation) error {
 	return nil
 }
 
+// TestInterpretCounterOps tests interpretation of counter operations.
+// Note: test cases are pretty simple because the counter operation is
+// essentially always called with 0 initial bytes and packets.
+func TestInterpretCounterOps(t *testing.T) {
+	for _, test := range []interpretOperationTestAction{
+		{ // cmd: add rule ip tab ch counter
+			tname:    "counter with 0 initial bytes and packets",
+			opStr:    "[ counter pkts 0 bytes 0 ]",
+			expected: newCounter(0, 0),
+		},
+		{
+			tname:    "counter with non-zero initial bytes and packets",
+			opStr:    "[ counter pkts 4561 bytes 39 ]",
+			expected: newCounter(4561, 39),
+		},
+	} {
+		t.Run(test.tname, func(t *testing.T) { checkOp(t, test, checkCounterOp) })
+	}
+}
+
+// checkCounterOp checks that the given operation is a counter operation and
+// that it matches the expected counter operation.
+func checkCounterOp(tname string, expected operation, actual operation) error {
+	expectedCntr := expected.(*counter)
+	cntr, ok := actual.(*counter)
+	if !ok {
+		return fmt.Errorf("expected operation type to be Counter for %s, got %T", tname, actual)
+	}
+	if bytes, expectedBytes := cntr.bytes.Load(), expectedCntr.bytes.Load(); bytes != expectedBytes {
+		return fmt.Errorf("expected bytes counter to be %d for %s, got %d", expectedBytes, tname, bytes)
+	}
+	if pkts, expectedPkts := cntr.packets.Load(), expectedCntr.packets.Load(); pkts != expectedPkts {
+		return fmt.Errorf("expected packets counter to be %d for %s, got %d", expectedPkts, tname, pkts)
+	}
+	return nil
+}
+
 // TestInterpretRule tests the interpretation of basic and general rules as a
 // list of operations.
 func TestInterpretRule(t *testing.T) {