Use specified source address for IGMP/MLD packets

This change also considers interfaces and network endpoints enabled up
up to the point all work to disable them are complete. This was needed
so that protocols can perform shutdown work while being disabled (e.g.
sending a packet which requires the endpoint to be enabled to obtain a
source address).

Bug #4682, #4861
Fixes #4888

Startblock:
  has LGTM from peterjohnston
  and then
  add reviewer brunodalbo
PiperOrigin-RevId: 346869702

2 files changed, 11 insertions, 3 deletions

diff --git a/pkg/tcpip/stack/addressable_endpoint_state.go b/pkg/tcpip/stack/addressable_endpoint_state.go
index 6e4f5fa46..cd423bf71 100644
--- a/pkg/tcpip/stack/addressable_endpoint_state.go
+++ b/pkg/tcpip/stack/addressable_endpoint_state.go
@@ -82,12 +82,16 @@ func (a *AddressableEndpointState) ForEachEndpoint(f func(AddressEndpoint) bool)
 }
 
 // ForEachPrimaryEndpoint calls f for each primary address.
-func (a *AddressableEndpointState) ForEachPrimaryEndpoint(f func(AddressEndpoint)) {
+//
+// Once f returns false, f will no longer be called.
+func (a *AddressableEndpointState) ForEachPrimaryEndpoint(f func(AddressEndpoint) bool) {
 	a.mu.RLock()
 	defer a.mu.RUnlock()
 
 	for _, ep := range a.mu.primary {
-		f(ep)
+		if !f(ep) {
+			return
+		}
 	}
 }
 
diff --git a/pkg/tcpip/stack/nic.go b/pkg/tcpip/stack/nic.go
index 5887aa1ed..a6237dd5f 100644
--- a/pkg/tcpip/stack/nic.go
+++ b/pkg/tcpip/stack/nic.go
@@ -172,7 +172,7 @@ func (n *NIC) disable() {
 //
 // n MUST be locked.
 func (n *NIC) disableLocked() {
-	if !n.setEnabled(false) {
+	if !n.Enabled() {
 		return
 	}
 
@@ -184,6 +184,10 @@ func (n *NIC) disableLocked() {
 	for _, ep := range n.networkEndpoints {
 		ep.Disable()
 	}
+
+	if !n.setEnabled(false) {
+		panic("should have only done work to disable the NIC if it was enabled")
+	}
 }
 
 // enable enables n.