Update connection state on iptables entrypoint

We only need to update the state of the connection once per packet since
the same packet will never change the state of a connection as it
performs subsequent hooks.

Put more concretely, TCP segments' flags, seq/ack number, etc. will
not change between the {Prerouting, Output}, {Forward} and
{Input, Postrouting} hooks.

PiperOrigin-RevId: 409480126

pkg/tcpip/stack/conntrack.go

@@ -149,9 +149,13 @@ func (cn *conn) timedOut(now tcpip.MonotonicTime) bool {
 }
 
 // update the connection tracking state.
-//
-// +checklocks:cn.stateMu
-func (cn *conn) updateLocked(pkt *PacketBuffer, reply bool) {
+func (cn *conn) update(pkt *PacketBuffer, reply bool) {
+	cn.stateMu.Lock()
+	defer cn.stateMu.Unlock()
+
+	// Mark the connection as having been used recently so it isn't reaped.
+	cn.lastUsed = cn.ct.clock.NowMonotonic()
+
 	if pkt.TransportProtocolNumber != header.TCPProtocolNumber {
 		return
 	}
@@ -394,60 +398,72 @@ func (ct *ConnTrack) init() {
 	ct.buckets = make([]bucket, numBuckets)
 }
 
-func (ct *ConnTrack) getConnOrMaybeInsertNoop(pkt *PacketBuffer) *tuple {
-	tid, isICMPError, ok := getTupleID(pkt)
-	if !ok {
-		return nil
+// getConnAndUpdate attempts to get a connection or creates one if no
+// connection exists for the packet and packet's protocol is trackable.
+//
+// If the packet's protocol is trackable, the connection's state is updated to
+// match the contents of the packet.
+func (ct *ConnTrack) getConnAndUpdate(pkt *PacketBuffer) *tuple {
+	// Get or (maybe) create a connection.
+	t := func() *tuple {
+		tid, isICMPError, ok := getTupleID(pkt)
+		if !ok {
+			return nil
+		}
+
+		bktID := ct.bucket(tid)
+
+		ct.mu.RLock()
+		bkt := &ct.buckets[bktID]
+		ct.mu.RUnlock()
+
+		now := ct.clock.NowMonotonic()
+		if t := bkt.connForTID(tid, now); t != nil {
+			return t
+		}
+
+		if isICMPError {
+			// Do not create a noop entry in response to an ICMP error.
+			return nil
+		}
+
+		bkt.mu.Lock()
+		defer bkt.mu.Unlock()
+
+		// Make sure a connection wasn't added between when we last checked the
+		// bucket and acquired the bucket's write lock.
+		if t := bkt.connForTIDRLocked(tid, now); t != nil {
+			return t
+		}
+
+		// This is the first packet we're seeing for the connection. Create an entry
+		// for this new connection.
+		conn := &conn{
+			ct:       ct,
+			original: tuple{tupleID: tid},
+			reply:    tuple{tupleID: tid.reply(), reply: true},
+			lastUsed: now,
+		}
+		conn.original.conn = conn
+		conn.reply.conn = conn
+
+		// For now, we only map an entry for the packet's original tuple as NAT may be
+		// performed on this connection. Until the packet goes through all the hooks
+		// and its final address/port is known, we cannot know what the response
+		// packet's addresses/ports will look like.
+		//
+		// This is okay because the destination cannot send its response until it
+		// receives the packet; the packet will only be received once all the hooks
+		// have been performed.
+		//
+		// See (*conn).finalize.
+		bkt.tuples.PushFront(&conn.original)
+		return &conn.original
+	}()
+	if t != nil {
+		t.conn.update(pkt, t.reply)
 	}
-
-	bktID := ct.bucket(tid)
-
-	ct.mu.RLock()
-	bkt := &ct.buckets[bktID]
-	ct.mu.RUnlock()
-
-	now := ct.clock.NowMonotonic()
-	if t := bkt.connForTID(tid, now); t != nil {
-		return t
-	}
-
-	if isICMPError {
-		// Do not create a noop entry in response to an ICMP error.
-		return nil
-	}
-
-	bkt.mu.Lock()
-	defer bkt.mu.Unlock()
-
-	// Make sure a connection wasn't added between when we last checked the
-	// bucket and acquired the bucket's write lock.
-	if t := bkt.connForTIDRLocked(tid, now); t != nil {
-		return t
-	}
-
-	// This is the first packet we're seeing for the connection. Create an entry
-	// for this new connection.
-	conn := &conn{
-		ct:       ct,
-		original: tuple{tupleID: tid},
-		reply:    tuple{tupleID: tid.reply(), reply: true},
-		lastUsed: now,
-	}
-	conn.original.conn = conn
-	conn.reply.conn = conn
-
-	// For now, we only map an entry for the packet's original tuple as NAT may be
-	// performed on this connection. Until the packet goes through all the hooks
-	// and its final address/port is known, we cannot know what the response
-	// packet's addresses/ports will look like.
-	//
-	// This is okay because the destination cannot send its response until it
-	// receives the packet; the packet will only be received once all the hooks
-	// have been performed.
-	//
-	// See (*conn).finalize.
-	bkt.tuples.PushFront(&conn.original)
-	return &conn.original
+	return t
 }
 
 func (ct *ConnTrack) connForTID(tid tupleID) *tuple {
@@ -670,13 +686,6 @@ func (cn *conn) handlePacket(pkt *PacketBuffer, hook Hook, rt *Route) bool {
 
 	reply := pkt.tuple.reply
 
-	cn.stateMu.Lock()
-	// Mark the connection as having been used recently so it isn't reaped.
-	cn.lastUsed = cn.ct.clock.NowMonotonic()
-	// Update connection state.
-	cn.updateLocked(pkt, reply)
-	cn.stateMu.Unlock()
-
 	tid, performManip := func() (tupleID, bool) {
 		cn.mu.RLock()
 		defer cn.mu.RUnlock()

pkg/tcpip/stack/conntrack_test.go

@@ -35,7 +35,7 @@ func TestReap(t *testing.T) {
 	// the connection won't be considered established. Thus the timeout for
 	// reaping is unestablishedTimeout.
 	pkt1 := genTCPPacket()
-	pkt1.tuple = ct.getConnOrMaybeInsertNoop(pkt1)
+	pkt1.tuple = ct.getConnAndUpdate(pkt1)
 	// We set rt.routeInfo.Loop to avoid a panic when handlePacket calls
 	// rt.RequiresTXTransportChecksum.
 	var rt Route
@@ -49,7 +49,7 @@ func TestReap(t *testing.T) {
 	// lastUsed, but per #6748 didn't.
 	clock.Advance(unestablishedTimeout / 2)
 	pkt2 := genTCPPacket()
-	pkt2.tuple = ct.getConnOrMaybeInsertNoop(pkt2)
+	pkt2.tuple = ct.getConnAndUpdate(pkt2)
 	if pkt2.tuple.conn.handlePacket(pkt2, Output, &rt) {
 		t.Fatal("handlePacket() shouldn't perform any NAT")
 	}

pkg/tcpip/stack/iptables.go

@@ -274,7 +274,7 @@ func (it *IPTables) CheckPrerouting(pkt *PacketBuffer, addressEP AddressableEndp
 		return true
 	}
 
-	pkt.tuple = it.connections.getConnOrMaybeInsertNoop(pkt)
+	pkt.tuple = it.connections.getConnAndUpdate(pkt)
 
 	for _, check := range [...]checkTableFn{
 		it.checkMangleRLocked,
@@ -349,7 +349,7 @@ func (it *IPTables) CheckOutput(pkt *PacketBuffer, r *Route, outNicName string)
 		return true
 	}
 
-	pkt.tuple = it.connections.getConnOrMaybeInsertNoop(pkt)
+	pkt.tuple = it.connections.getConnAndUpdate(pkt)
 
 	for _, check := range [...]checkTableFn{
 		it.checkMangleRLocked,