From f0737cc307f6e503cae1fb0063189e6ef8eebdb3 Mon Sep 17 00:00:00 2001
From: Kevin Krakauer <krakauer@google.com>
Date: Tue, 10 May 2022 13:24:59 -0700
Subject: [PATCH] netstack: use checkescape in iptables hot paths

These functions used to allocate even when iptables were disabled. Prevent that
from happening again.

Update to also use gVisor's sync package, as we were using the standard one.

PiperOrigin-RevId: 447812920
---
 pkg/tcpip/stack/iptables.go       | 29 +++++++++++++++++++++++++++++
 pkg/tcpip/stack/iptables_types.go |  2 +-
 2 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/pkg/tcpip/stack/iptables.go b/pkg/tcpip/stack/iptables.go
index fbee4dcd1..5aee1a560 100644
--- a/pkg/tcpip/stack/iptables.go
+++ b/pkg/tcpip/stack/iptables.go
@@ -273,6 +273,10 @@ type checkTable struct {
 //
 // If IPTables should not be skipped, tables will be updated with the
 // specified table.
+//
+// This is called in the hot path even when iptables are disabled, so we ensure
+// it does not allocate.
+// +checkescape:heap,builtin
 func (it *IPTables) shouldSkipOrPopulateTables(tables []checkTable, pkt *PacketBuffer) bool {
 	switch pkt.NetworkProtocolNumber {
 	case header.IPv4ProtocolNumber, header.IPv6ProtocolNumber:
@@ -303,6 +307,11 @@ func (it *IPTables) shouldSkipOrPopulateTables(tables []checkTable, pkt *PacketB
 // must be dropped if false is returned.
 //
 // Precondition: The packet's network and transport header must be set.
+//
+// This is called in the hot path even when iptables are disabled, so we ensure
+// that it does not allocate. Note that called functions (e.g.
+// getConnAndUpdate) can allocate.
+// +checkescape
 func (it *IPTables) CheckPrerouting(pkt *PacketBuffer, addressEP AddressableEndpoint, inNicName string) bool {
 	tables := [...]checkTable{
 		{
@@ -336,6 +345,11 @@ func (it *IPTables) CheckPrerouting(pkt *PacketBuffer, addressEP AddressableEndp
 // must be dropped if false is returned.
 //
 // Precondition: The packet's network and transport header must be set.
+//
+// This is called in the hot path even when iptables are disabled, so we ensure
+// that it does not allocate. Note that called functions (e.g.
+// getConnAndUpdate) can allocate.
+// +checkescape
 func (it *IPTables) CheckInput(pkt *PacketBuffer, inNicName string) bool {
 	tables := [...]checkTable{
 		{
@@ -371,6 +385,11 @@ func (it *IPTables) CheckInput(pkt *PacketBuffer, inNicName string) bool {
 // must be dropped if false is returned.
 //
 // Precondition: The packet's network and transport header must be set.
+//
+// This is called in the hot path even when iptables are disabled, so we ensure
+// that it does not allocate. Note that called functions (e.g.
+// getConnAndUpdate) can allocate.
+// +checkescape
 func (it *IPTables) CheckForward(pkt *PacketBuffer, inNicName, outNicName string) bool {
 	tables := [...]checkTable{
 		{
@@ -398,6 +417,11 @@ func (it *IPTables) CheckForward(pkt *PacketBuffer, inNicName, outNicName string
 // must be dropped if false is returned.
 //
 // Precondition: The packet's network and transport header must be set.
+//
+// This is called in the hot path even when iptables are disabled, so we ensure
+// that it does not allocate. Note that called functions (e.g.
+// getConnAndUpdate) can allocate.
+// +checkescape
 func (it *IPTables) CheckOutput(pkt *PacketBuffer, r *Route, outNicName string) bool {
 	tables := [...]checkTable{
 		{
@@ -435,6 +459,11 @@ func (it *IPTables) CheckOutput(pkt *PacketBuffer, r *Route, outNicName string)
 // must be dropped if false is returned.
 //
 // Precondition: The packet's network and transport header must be set.
+//
+// This is called in the hot path even when iptables are disabled, so we ensure
+// that it does not allocate. Note that called functions (e.g.
+// getConnAndUpdate) can allocate.
+// +checkescape
 func (it *IPTables) CheckPostrouting(pkt *PacketBuffer, r *Route, addressEP AddressableEndpoint, outNicName string) bool {
 	tables := [...]checkTable{
 		{
diff --git a/pkg/tcpip/stack/iptables_types.go b/pkg/tcpip/stack/iptables_types.go
index d77ed1c3a..97f9e1f6c 100644
--- a/pkg/tcpip/stack/iptables_types.go
+++ b/pkg/tcpip/stack/iptables_types.go
@@ -17,8 +17,8 @@ package stack
 import (
 	"fmt"
 	"strings"
-	"sync"
 
+	"gvisor.dev/gvisor/pkg/sync"
 	"gvisor.dev/gvisor/pkg/tcpip"
 	"gvisor.dev/gvisor/pkg/tcpip/header"
 )