Drop connection on reply tuple conflict

Updates #6850.

PiperOrigin-RevId: 410368440

pkg/tcpip/stack/conntrack.go

@@ -19,6 +19,7 @@ import (
 	"fmt"
 	"math/rand"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"gvisor.dev/gvisor/pkg/tcpip"
@@ -108,6 +109,17 @@ const (
 	manipPerformedNoop
 )
 
+type finalizeResult uint32
+
+const (
+	// A finalizeResult must be explicitly set so we don't make use of the zero
+	// value.
+	_ finalizeResult = iota
+
+	finalizeResultSuccess
+	finalizeResultConflict
+)
+
 // conn is a tracked connection.
 //
 // +stateify savable
@@ -120,11 +132,13 @@ type conn struct {
 	// reply is the tuple in reply direction.
 	reply tuple
 
-	mu sync.RWMutex `state:"nosave"`
-	// Indicates that the connection has been finalized and may handle replies.
+	finalizeOnce sync.Once
+	// Holds a finalizeResult.
 	//
-	// +checklocks:mu
-	finalized bool
+	// +checkatomics
+	finalizeResult uint32
+
+	mu sync.RWMutex `state:"nosave"`
 	// sourceManip indicates the source manipulation type.
 	//
 	// +checklocks:mu
@@ -505,51 +519,66 @@ func (bkt *bucket) connForTIDRLocked(tid tupleID, now tcpip.MonotonicTime) *tupl
 	return nil
 }
 
-func (ct *ConnTrack) finalize(cn *conn) {
-	tid := cn.reply.id()
-	id := ct.bucket(tid)
-
+func (ct *ConnTrack) finalize(cn *conn) finalizeResult {
 	ct.mu.RLock()
-	bkt := &ct.buckets[id]
+	buckets := ct.buckets
 	ct.mu.RUnlock()
 
+	{
+		tid := cn.reply.id()
+		id := ct.bucket(tid)
+
+		bkt := &buckets[id]
+		bkt.mu.Lock()
+		if bkt.connForTIDRLocked(tid, ct.clock.NowMonotonic()) == nil {
+			bkt.tuples.PushFront(&cn.reply)
+			bkt.mu.Unlock()
+			return finalizeResultSuccess
+		}
+		bkt.mu.Unlock()
+	}
+
+	// Another connection for the reply already exists. Remove the original and
+	// let the caller know we failed.
+	//
+	// TODO(https://gvisor.dev/issue/6850): Investigate handling this clash
+	// better.
+
+	tid := cn.original.id()
+	id := ct.bucket(tid)
+	bkt := &buckets[id]
 	bkt.mu.Lock()
 	defer bkt.mu.Unlock()
-
-	if t := bkt.connForTIDRLocked(tid, ct.clock.NowMonotonic()); t != nil {
-		// Another connection for the reply already exists. We can't do much about
-		// this so we leave the connection cn represents in a state where it can
-		// send packets but its responses will be mapped to some other connection.
-		// This may be okay if the connection only expects to send packets without
-		// any responses.
-		//
-		// TODO(https://gvisor.dev/issue/6850): Investigate handling this clash
-		// better.
-		return
-	}
-
-	bkt.tuples.PushFront(&cn.reply)
+	bkt.tuples.Remove(&cn.original)
+	return finalizeResultConflict
 }
 
-func (cn *conn) finalize() {
-	{
-		cn.mu.RLock()
-		finalized := cn.finalized
-		cn.mu.RUnlock()
-		if finalized {
-			return
-		}
-	}
+func (cn *conn) getFinalizeResult() finalizeResult {
+	return finalizeResult(atomic.LoadUint32(&cn.finalizeResult))
+}
 
-	cn.mu.Lock()
-	finalized := cn.finalized
-	cn.finalized = true
-	cn.mu.Unlock()
-	if finalized {
-		return
-	}
+// finalize attempts to finalize the connection and returns true iff the
+// connection was successfully finalized.
+//
+// If the connection failed to finalize, the caller should drop the packet
+// associated with the connection.
+//
+// If multiple goroutines attempt to finalize at the same time, only one
+// goroutine will perform the work to finalize the connection, but all
+// goroutines will block until the finalizing goroutine finishes finalizing.
+func (cn *conn) finalize() bool {
+	cn.finalizeOnce.Do(func() {
+		atomic.StoreUint32(&cn.finalizeResult, uint32(cn.ct.finalize(cn)))
+	})
 
-	cn.ct.finalize(cn)
+	switch res := cn.getFinalizeResult(); res {
+	case finalizeResultSuccess:
+		return true
+	case finalizeResultConflict:
+		return false
+	default:
+		panic(fmt.Sprintf("unhandled result = %d", res))
+	}
 }
 
 func (cn *conn) maybePerformNoopNAT(dnat bool) {
@@ -919,9 +948,7 @@ func (ct *ConnTrack) reapTupleLocked(reapingTuple *tuple, bktID int, bkt *bucket
 	}
 
 	otherTupleBktID := ct.bucket(otherTuple.id())
-	reapingTuple.conn.mu.RLock()
-	replyTupleInserted := reapingTuple.conn.finalized
-	reapingTuple.conn.mu.RUnlock()
+	replyTupleInserted := reapingTuple.conn.getFinalizeResult() == finalizeResultSuccess
 
 	// To maintain lock order, we can only reap both tuples if the tuple for the
 	// other direction appears later in the table.

pkg/tcpip/stack/iptables.go

@@ -312,9 +312,9 @@ func (it *IPTables) CheckInput(pkt *PacketBuffer, inNicName string) bool {
 	}
 
 	if t := pkt.tuple; t != nil {
-		t.conn.finalize()
+		pkt.tuple = nil
+		return t.conn.finalize()
 	}
-	pkt.tuple = nil
 	return true
 }
 
@@ -388,9 +388,9 @@ func (it *IPTables) CheckPostrouting(pkt *PacketBuffer, r *Route, addressEP Addr
 	}
 
 	if t := pkt.tuple; t != nil {
-		t.conn.finalize()
+		pkt.tuple = nil
+		return t.conn.finalize()
 	}
-	pkt.tuple = nil
 	return true
 }
 

pkg/tcpip/stack/iptables_test.go

@@ -18,15 +18,16 @@ import (
 	"math/rand"
 	"testing"
 
+	"gvisor.dev/gvisor/pkg/tcpip"
 	"gvisor.dev/gvisor/pkg/tcpip/faketime"
 	"gvisor.dev/gvisor/pkg/tcpip/header"
 	"gvisor.dev/gvisor/pkg/tcpip/testutil"
 )
 
 const (
-	nattedDstPort = 1
-	srcPort       = 2
-	dstPort       = 3
+	nattedPort = 1
+	srcPort    = 2
+	dstPort    = 3
 
 	// The network protocol used for these tests doesn't matter as the tests are
 	// not targetting anything protocol specific.
@@ -35,12 +36,12 @@ const (
 )
 
 var (
-	nattedDstAddr = testutil.MustParse6("a::1")
-	srcAddr       = testutil.MustParse6("b::2")
-	dstAddr       = testutil.MustParse6("c::3")
+	nattedAddr = testutil.MustParse6("a::1")
+	srcAddr    = testutil.MustParse6("b::2")
+	dstAddr    = testutil.MustParse6("c::3")
 )
 
-func v6PacketBuffer() *PacketBuffer {
+func v6PacketBufferWithSrcAddr(srcAddr tcpip.Address) *PacketBuffer {
 	pkt := NewPacketBuffer(PacketBufferOptions{
 		ReserveHeaderBytes: header.IPv6MinimumSize + header.UDPMinimumSize,
 	})
@@ -67,6 +68,10 @@ func v6PacketBuffer() *PacketBuffer {
 	return pkt
 }
 
+func v6PacketBuffer() *PacketBuffer {
+	return v6PacketBufferWithSrcAddr(srcAddr)
+}
+
 // TestNATedConnectionReap tests that NATed connections are properly reaped.
 func TestNATedConnectionReap(t *testing.T) {
 	clock := faketime.NewManualClock()
@@ -76,7 +81,7 @@ func TestNATedConnectionReap(t *testing.T) {
 		Rules: []Rule{
 			// Prerouting
 			{
-				Target: &DNATTarget{NetworkProtocol: netProto, Addr: nattedDstAddr, Port: nattedDstPort},
+				Target: &DNATTarget{NetworkProtocol: netProto, Addr: nattedAddr, Port: nattedPort},
 			},
 			{
 				Target: &AcceptTarget{},
@@ -317,3 +322,109 @@ func TestNATAlwaysPerformed(t *testing.T) {
 		})
 	}
 }
+
+func TestNATConflict(t *testing.T) {
+	otherSrcAddr := testutil.MustParse6("d::4")
+
+	tests := []struct {
+		name          string
+		checkIPTables func(*testing.T, *IPTables, *PacketBuffer, bool)
+	}{
+		{
+			name: "Prerouting and Input",
+			checkIPTables: func(t *testing.T, iptables *IPTables, pkt *PacketBuffer, lastHookOK bool) {
+				t.Helper()
+
+				if !iptables.CheckPrerouting(pkt, nil /* addressEP */, "" /* inNicName */) {
+					t.Fatal("got ipt.CheckPrerouting(...) = false, want = true")
+				}
+				if got := iptables.CheckInput(pkt, "" /* inNicName */); got != lastHookOK {
+					t.Fatalf("got ipt.CheckInput(...) = %t, want = %t", got, lastHookOK)
+				}
+			},
+		},
+		{
+			name: "Output and Postrouting",
+			checkIPTables: func(t *testing.T, iptables *IPTables, pkt *PacketBuffer, lastHookOK bool) {
+				t.Helper()
+
+				// Output and Postrouting hooks depends on a route but if the route is
+				// local, we don't need anything else from it.
+				r := Route{
+					routeInfo: routeInfo{
+						Loop: PacketLoop,
+					},
+				}
+				if !iptables.CheckOutput(pkt, &r, "" /* outNicName */) {
+					t.Fatal("got iptables.CheckOutput(...) = false, want = true")
+				}
+				if got := iptables.CheckPostrouting(pkt, &r, nil /* addressEP */, "" /* outNicName */); got != lastHookOK {
+					t.Fatalf("got iptables.CheckPostrouting(...) = %t, want = %t", got, lastHookOK)
+				}
+			},
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+
+			clock := faketime.NewManualClock()
+			iptables := DefaultTables(clock, rand.New(rand.NewSource(0 /* seed */)))
+
+			table := Table{
+				Rules: []Rule{
+					// Prerouting
+					{
+						Target: &AcceptTarget{},
+					},
+
+					// Input
+					{
+						Target: &SNATTarget{NetworkProtocol: header.IPv6ProtocolNumber, Addr: nattedAddr, Port: nattedPort},
+					},
+					{
+						Target: &AcceptTarget{},
+					},
+
+					// Forward
+					{
+						Target: &AcceptTarget{},
+					},
+
+					// Output
+					{
+						Target: &AcceptTarget{},
+					},
+
+					// Postrouting
+					{
+						Target: &SNATTarget{NetworkProtocol: header.IPv6ProtocolNumber, Addr: nattedAddr, Port: nattedPort},
+					},
+					{
+						Target: &AcceptTarget{},
+					},
+				},
+				BuiltinChains: [NumHooks]int{
+					Prerouting:  0,
+					Input:       1,
+					Forward:     3,
+					Output:      4,
+					Postrouting: 5,
+				},
+			}
+			if err := iptables.ReplaceTable(NATID, table, ipv6); err != nil {
+				t.Fatalf("ipt.ReplaceTable(%d, _, true): %s", NATID, err)
+			}
+
+			// Create and finalize the connection.
+			test.checkIPTables(t, iptables, v6PacketBufferWithSrcAddr(srcAddr), true /* lastHookOK */)
+
+			// A packet from a different source that get NATed to the same tuple as
+			// the connection created above should be dropped when finalizing.
+			test.checkIPTables(t, iptables, v6PacketBufferWithSrcAddr(otherSrcAddr), false /* lastHookOK */)
+
+			// A packet from the original source should be NATed as normal.
+			test.checkIPTables(t, iptables, v6PacketBufferWithSrcAddr(srcAddr), true /* lastHookOK */)
+		})
+	}
+}