Drop ConnTrack.handlePacket

Move the hook specific logic to the IPTables hook functions.

This lets us avoid having to perform checks on the hook to determine
what action to take.

Later changes will drop the need for handlePacket's return value,
reducing the value of this function that all hooks call into.

PiperOrigin-RevId: 400298023

2 files changed, 85 insertions, 42 deletions

diff --git a/pkg/tcpip/stack/conntrack.go b/pkg/tcpip/stack/conntrack.go
index bd47f734f..79bc001c7 100644
--- a/pkg/tcpip/stack/conntrack.go
+++ b/pkg/tcpip/stack/conntrack.go
@@ -403,29 +403,6 @@ func insertConn(tupleBkt *bucket, replyBkt *bucket, conn *conn) {
 	}
 }
 
-// handlePacket will manipulate the port and address of the packet if the
-// connection exists. Returns whether, after the packet traverses the tables,
-// it should create a new entry in the table.
-func (ct *ConnTrack) handlePacket(pkt *PacketBuffer, hook Hook, r *Route) bool {
-	switch hook {
-	case Prerouting, Input, Output, Postrouting:
-	default:
-		return false
-	}
-
-	if conn, dir := ct.connFor(pkt); conn != nil {
-		conn.handlePacket(pkt, hook, dir, r)
-		return false
-	}
-
-	// Connection not found for the packet.
-	//
-	// If this is the last hook in the data path for this packet (Input if
-	// incoming, Postrouting if outgoing), indicate that a connection should be
-	// inserted by the end of this hook.
-	return hook == Input || hook == Postrouting
-}
-
 func (cn *conn) handlePacket(pkt *PacketBuffer, hook Hook, dir direction, r *Route) {
 	if pkt.NatDone {
 		return
diff --git a/pkg/tcpip/stack/iptables.go b/pkg/tcpip/stack/iptables.go
index a20bef3c5..1021e484a 100644
--- a/pkg/tcpip/stack/iptables.go
+++ b/pkg/tcpip/stack/iptables.go
@@ -271,7 +271,17 @@ const (
 //
 // Precondition: The packet's network and transport header must be set.
 func (it *IPTables) CheckPrerouting(pkt *PacketBuffer, addressEP AddressableEndpoint, inNicName string) bool {
-	return it.check(Prerouting, pkt, nil /* route */, addressEP, inNicName, "" /* outNicName */)
+	const hook = Prerouting
+
+	if it.shouldSkip(pkt.NetworkProtocolNumber) {
+		return true
+	}
+
+	if conn, dir := it.connections.connFor(pkt); conn != nil {
+		conn.handlePacket(pkt, hook, dir, nil /* route */)
+	}
+
+	return it.check(hook, pkt, nil /* route */, addressEP, inNicName, "" /* outNicName */)
 }
 
 // CheckInput performs the input hook on the packet.
@@ -281,7 +291,26 @@ func (it *IPTables) CheckPrerouting(pkt *PacketBuffer, addressEP AddressableEndp
 //
 // Precondition: The packet's network and transport header must be set.
 func (it *IPTables) CheckInput(pkt *PacketBuffer, inNicName string) bool {
-	return it.check(Input, pkt, nil /* route */, nil /* addressEP */, inNicName, "" /* outNicName */)
+	const hook = Input
+
+	if it.shouldSkip(pkt.NetworkProtocolNumber) {
+		return true
+	}
+
+	shouldTrack := true
+	if conn, dir := it.connections.connFor(pkt); conn != nil {
+		conn.handlePacket(pkt, hook, dir, nil /* route */)
+		shouldTrack = false
+	}
+
+	if !it.check(hook, pkt, nil /* route */, nil /* addressEP */, inNicName, "" /* outNicName */) {
+		return false
+	}
+
+	// This is the last hook a packet will perform so if the packet's
+	// connection is not tracked, we may need to add a no-op entry.
+	it.maybeinsertNoopConn(pkt, hook, shouldTrack)
+	return true
 }
 
 // CheckForward performs the forward hook on the packet.
@@ -291,6 +320,10 @@ func (it *IPTables) CheckInput(pkt *PacketBuffer, inNicName string) bool {
 //
 // Precondition: The packet's network and transport header must be set.
 func (it *IPTables) CheckForward(pkt *PacketBuffer, inNicName, outNicName string) bool {
+	if it.shouldSkip(pkt.NetworkProtocolNumber) {
+		return true
+	}
+
 	return it.check(Forward, pkt, nil /* route */, nil /* addressEP */, inNicName, outNicName)
 }
 
@@ -301,7 +334,17 @@ func (it *IPTables) CheckForward(pkt *PacketBuffer, inNicName, outNicName string
 //
 // Precondition: The packet's network and transport header must be set.
 func (it *IPTables) CheckOutput(pkt *PacketBuffer, r *Route, outNicName string) bool {
-	return it.check(Output, pkt, r, nil /* addressEP */, "" /* inNicName */, outNicName)
+	const hook = Output
+
+	if it.shouldSkip(pkt.NetworkProtocolNumber) {
+		return true
+	}
+
+	if conn, dir := it.connections.connFor(pkt); conn != nil {
+		conn.handlePacket(pkt, hook, dir, r)
+	}
+
+	return it.check(hook, pkt, r, nil /* addressEP */, "" /* inNicName */, outNicName)
 }
 
 // CheckPostrouting performs the postrouting hook on the packet.
@@ -311,7 +354,41 @@ func (it *IPTables) CheckOutput(pkt *PacketBuffer, r *Route, outNicName string)
 //
 // Precondition: The packet's network and transport header must be set.
 func (it *IPTables) CheckPostrouting(pkt *PacketBuffer, r *Route, addressEP AddressableEndpoint, outNicName string) bool {
-	return it.check(Postrouting, pkt, r, addressEP, "" /* inNicName */, outNicName)
+	const hook = Postrouting
+
+	if it.shouldSkip(pkt.NetworkProtocolNumber) {
+		return true
+	}
+
+	shouldTrack := true
+	if conn, dir := it.connections.connFor(pkt); conn != nil {
+		conn.handlePacket(pkt, hook, dir, r)
+		shouldTrack = false
+	}
+
+	if !it.check(hook, pkt, r, addressEP, "" /* inNicName */, outNicName) {
+		return false
+	}
+
+	// This is the last hook a packet will perform so if the packet's
+	// connection is not tracked, we may need to add a no-op entry.
+	it.maybeinsertNoopConn(pkt, hook, shouldTrack)
+	return true
+}
+
+func (it *IPTables) shouldSkip(netProto tcpip.NetworkProtocolNumber) bool {
+	switch netProto {
+	case header.IPv4ProtocolNumber, header.IPv6ProtocolNumber:
+	default:
+		// IPTables only supports IPv4/IPv6.
+		return true
+	}
+
+	it.mu.RLock()
+	defer it.mu.RUnlock()
+	// Many users never configure iptables. Spare them the cost of rule
+	// traversal if rules have never been set.
+	return !it.modified
 }
 
 // check runs pkt through the rules for hook. It returns true when the packet
@@ -320,20 +397,8 @@ func (it *IPTables) CheckPostrouting(pkt *PacketBuffer, r *Route, addressEP Addr
 //
 // Precondition: The packet's network and transport header must be set.
 func (it *IPTables) check(hook Hook, pkt *PacketBuffer, r *Route, addressEP AddressableEndpoint, inNicName, outNicName string) bool {
-	if pkt.NetworkProtocolNumber != header.IPv4ProtocolNumber && pkt.NetworkProtocolNumber != header.IPv6ProtocolNumber {
-		return true
-	}
-	// Many users never configure iptables. Spare them the cost of rule
-	// traversal if rules have never been set.
 	it.mu.RLock()
 	defer it.mu.RUnlock()
-	if !it.modified {
-		return true
-	}
-
-	// Packets are manipulated only if connection and matching
-	// NAT rule exists.
-	shouldTrack := it.connections.handlePacket(pkt, hook, r)
 
 	// Go through each table containing the hook.
 	priorities := it.priorities[hook]
@@ -377,6 +442,10 @@ func (it *IPTables) check(hook Hook, pkt *PacketBuffer, r *Route, addressEP Addr
 		}
 	}
 
+	return true
+}
+
+func (it *IPTables) maybeinsertNoopConn(pkt *PacketBuffer, hook Hook, shouldTrack bool) {
 	// If this connection should be tracked, try to add an entry for it. If
 	// traversing the nat table didn't end in adding an entry,
 	// maybeInsertNoop will add a no-op entry for the connection. This is
@@ -390,9 +459,6 @@ func (it *IPTables) check(hook Hook, pkt *PacketBuffer, r *Route, addressEP Addr
 	if shouldTrack {
 		it.connections.maybeInsertNoop(pkt)
 	}
-
-	// Every table returned Accept.
-	return true
 }
 
 // beforeSave is invoked by stateify.