diff --git a/pkg/tcpip/nftables/nftables.go b/pkg/tcpip/nftables/nftables.go index 43061bdcd..af20ca6bc 100644 --- a/pkg/tcpip/nftables/nftables.go +++ b/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. diff --git a/pkg/tcpip/nftables/nftables_test.go b/pkg/tcpip/nftables/nftables_test.go index ff62d6b96..ef81ae2c7 100644 --- a/pkg/tcpip/nftables/nftables_test.go +++ b/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) { diff --git a/pkg/tcpip/nftables/nftinterp.go b/pkg/tcpip/nftables/nftinterp.go index 31bdc8bad..2169c98fb 100644 --- a/pkg/tcpip/nftables/nftinterp.go +++ b/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. // diff --git a/pkg/tcpip/nftables/nftinterp_test.go b/pkg/tcpip/nftables/nftinterp_test.go index d9d04c528..d476caf28 100644 --- a/pkg/tcpip/nftables/nftinterp_test.go +++ b/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) {