Add endpoints to map only if registerEndpoint succeeds.

epsByNIC.registerEndpoint can add a multiportEndpoint to its map of
nic->multiportEndpoint even if multiport.Endpoint.singleRegisterEndpoint
failed. Same for transportDemuxer.singleRegisterEndpoint which ends up adding an
entry to nic->epsByNIC even if epsByNIC.registerEndpoint fails.

These breaks an invariant which the code assumes that a
multiportEndpoint/endpointsByNIC always have at least one valid entry.

PiperOrigin-RevId: 380310115

2 files changed, 66 insertions, 6 deletions

diff --git a/pkg/tcpip/stack/transport_demuxer.go b/pkg/tcpip/stack/transport_demuxer.go
index 21af0cf37..dda57e225 100644
--- a/pkg/tcpip/stack/transport_demuxer.go
+++ b/pkg/tcpip/stack/transport_demuxer.go
@@ -216,10 +216,17 @@ func (epsByNIC *endpointsByNIC) registerEndpoint(d *transportDemuxer, netProto t
 			netProto:   netProto,
 			transProto: transProto,
 		}
-		epsByNIC.endpoints[bindToDevice] = multiPortEp
 	}
 
-	return multiPortEp.singleRegisterEndpoint(t, flags)
+	if err := multiPortEp.singleRegisterEndpoint(t, flags); err != nil {
+		return err
+	}
+	// Only add this newly created multiportEndpoint if the singleRegisterEndpoint
+	// succeeded.
+	if !ok {
+		epsByNIC.endpoints[bindToDevice] = multiPortEp
+	}
+	return nil
 }
 
 func (epsByNIC *endpointsByNIC) checkEndpoint(flags ports.Flags, bindToDevice tcpip.NICID) tcpip.Error {
@@ -406,7 +413,6 @@ func (ep *multiPortEndpoint) handlePacketAll(id TransportEndpointID, pkt *Packet
 func (ep *multiPortEndpoint) singleRegisterEndpoint(t TransportEndpoint, flags ports.Flags) tcpip.Error {
 	ep.mu.Lock()
 	defer ep.mu.Unlock()
-
 	bits := flags.Bits() & ports.MultiBindFlagMask
 
 	if len(ep.endpoints) != 0 {
@@ -469,17 +475,21 @@ func (d *transportDemuxer) singleRegisterEndpoint(netProto tcpip.NetworkProtocol
 
 	eps.mu.Lock()
 	defer eps.mu.Unlock()
-
 	epsByNIC, ok := eps.endpoints[id]
 	if !ok {
 		epsByNIC = &endpointsByNIC{
 			endpoints: make(map[tcpip.NICID]*multiPortEndpoint),
 			seed:      d.stack.Seed(),
 		}
+	}
+	if err := epsByNIC.registerEndpoint(d, netProto, protocol, ep, flags, bindToDevice); err != nil {
+		return err
+	}
+	// Only add this newly created epsByNIC if registerEndpoint succeeded.
+	if !ok {
 		eps.endpoints[id] = epsByNIC
 	}
-
-	return epsByNIC.registerEndpoint(d, netProto, protocol, ep, flags, bindToDevice)
+	return nil
 }
 
 func (d *transportDemuxer) singleCheckEndpoint(netProto tcpip.NetworkProtocolNumber, protocol tcpip.TransportProtocolNumber, id TransportEndpointID, flags ports.Flags, bindToDevice tcpip.NICID) tcpip.Error {
diff --git a/pkg/tcpip/stack/transport_demuxer_test.go b/pkg/tcpip/stack/transport_demuxer_test.go
index 0972c94de..45b09110d 100644
--- a/pkg/tcpip/stack/transport_demuxer_test.go
+++ b/pkg/tcpip/stack/transport_demuxer_test.go
@@ -203,6 +203,56 @@ func TestTransportDemuxerRegister(t *testing.T) {
 	}
 }
 
+func TestTransportDemuxerRegisterMultiple(t *testing.T) {
+	type test struct {
+		flags ports.Flags
+		want  tcpip.Error
+	}
+	for _, subtest := range []struct {
+		name  string
+		tests []test
+	}{
+		{"zeroFlags", []test{
+			{ports.Flags{}, nil},
+			{ports.Flags{}, &tcpip.ErrPortInUse{}},
+		}},
+		{"multibindFlags", []test{
+			// Allow multiple registrations same TransportEndpointID with multibind flags.
+			{ports.Flags{LoadBalanced: true, MostRecent: true}, nil},
+			{ports.Flags{LoadBalanced: true, MostRecent: true}, nil},
+			// Disallow registration w/same ID for a non-multibindflag.
+			{ports.Flags{TupleOnly: true}, &tcpip.ErrPortInUse{}},
+		}},
+	} {
+		t.Run(subtest.name, func(t *testing.T) {
+			s := stack.New(stack.Options{
+				NetworkProtocols:   []stack.NetworkProtocolFactory{ipv4.NewProtocol},
+				TransportProtocols: []stack.TransportProtocolFactory{udp.NewProtocol},
+			})
+			var eps []tcpip.Endpoint
+			for idx, test := range subtest.tests {
+				var wq waiter.Queue
+				ep, err := s.NewEndpoint(udp.ProtocolNumber, ipv4.ProtocolNumber, &wq)
+				if err != nil {
+					t.Fatal(err)
+				}
+				eps = append(eps, ep)
+				tEP, ok := ep.(stack.TransportEndpoint)
+				if !ok {
+					t.Fatalf("%T does not implement stack.TransportEndpoint", ep)
+				}
+				id := stack.TransportEndpointID{LocalPort: 1}
+				if got, want := s.RegisterTransportEndpoint([]tcpip.NetworkProtocolNumber{ipv4.ProtocolNumber}, udp.ProtocolNumber, id, tEP, test.flags, 0), test.want; got != want {
+					t.Fatalf("test index: %d, s.RegisterTransportEndpoint(ipv4.ProtocolNumber, udp.ProtocolNumber, _, _, %+v, 0) = %s, want %s", idx, test.flags, got, want)
+				}
+			}
+			for _, ep := range eps {
+				ep.Close()
+			}
+		})
+	}
+}
+
 // TestBindToDeviceDistribution injects varied packets on input devices and checks that
 // the distribution of packets received matches expectations.
 func TestBindToDeviceDistribution(t *testing.T) {