Plumb IP_ADD_MEMBERSHIP and IP_DROP_MEMBERSHIP to netstack.

Also includes a few fixes for IPv4 multicast support. IPv6 support is coming in
a followup CL.

PiperOrigin-RevId: 233008638
Change-Id: If7dae6222fef43fda48033f0292af77832d95e82

2 files changed, 17 insertions, 5 deletions

diff --git a/pkg/tcpip/transport/udp/endpoint.go b/pkg/tcpip/transport/udp/endpoint.go
index b2a27a7cb..d46bf0ade 100644
--- a/pkg/tcpip/transport/udp/endpoint.go
+++ b/pkg/tcpip/transport/udp/endpoint.go
@@ -99,6 +99,7 @@ type endpoint struct {
 	effectiveNetProtos []tcpip.NetworkProtocolNumber
 }
 
+// +stateify savable
 type multicastMembership struct {
 	nicID         tcpip.NICID
 	multicastAddr tcpip.Address
@@ -412,6 +413,8 @@ func (e *endpoint) SetSockOpt(opt interface{}) *tcpip.Error {
 			nicID = e.stack.CheckLocalAddress(nicID, e.netProto, v.InterfaceAddr)
 		}
 		if nicID == 0 {
+			// TODO: Allow adding memberships without
+			// specifing an interface.
 			return tcpip.ErrNoRoute
 		}
 
@@ -766,9 +769,11 @@ func (e *endpoint) bindLocked(addr tcpip.FullAddress, commit func() *tcpip.Error
 		}
 	}
 
+	nicid := addr.NIC
 	if len(addr.Addr) != 0 {
 		// A local address was specified, verify that it's valid.
-		if e.stack.CheckLocalAddress(addr.NIC, netProto, addr.Addr) == 0 {
+		nicid = e.stack.CheckLocalAddress(addr.NIC, netProto, addr.Addr)
+		if nicid == 0 {
 			return tcpip.ErrBadLocalAddress
 		}
 	}
@@ -777,21 +782,21 @@ func (e *endpoint) bindLocked(addr tcpip.FullAddress, commit func() *tcpip.Error
 		LocalPort:    addr.Port,
 		LocalAddress: addr.Addr,
 	}
-	id, err = e.registerWithStack(addr.NIC, netProtos, id)
+	id, err = e.registerWithStack(nicid, netProtos, id)
 	if err != nil {
 		return err
 	}
 	if commit != nil {
 		if err := commit(); err != nil {
 			// Unregister, the commit failed.
-			e.stack.UnregisterTransportEndpoint(addr.NIC, netProtos, ProtocolNumber, id, e)
+			e.stack.UnregisterTransportEndpoint(nicid, netProtos, ProtocolNumber, id, e)
 			e.stack.ReleasePort(netProtos, ProtocolNumber, id.LocalAddress, id.LocalPort)
 			return err
 		}
 	}
 
 	e.id = id
-	e.regNICID = addr.NIC
+	e.regNICID = nicid
 	e.effectiveNetProtos = netProtos
 
 	// Mark endpoint as bound.
@@ -815,7 +820,8 @@ func (e *endpoint) Bind(addr tcpip.FullAddress, commit func() *tcpip.Error) *tcp
 		return err
 	}
 
-	e.bindNICID = addr.NIC
+	// Save the effective NICID generated by bindLocked.
+	e.bindNICID = e.regNICID
 
 	return nil
 }
diff --git a/pkg/tcpip/transport/udp/endpoint_state.go b/pkg/tcpip/transport/udp/endpoint_state.go
index db1e281ad..4d8210294 100644
--- a/pkg/tcpip/transport/udp/endpoint_state.go
+++ b/pkg/tcpip/transport/udp/endpoint_state.go
@@ -103,4 +103,10 @@ func (e *endpoint) afterLoad() {
 	if err != nil {
 		panic(*err)
 	}
+
+	for _, m := range e.multicastMemberships {
+		if err := e.stack.JoinGroup(e.netProto, m.nicID, m.multicastAddr); err != nil {
+			panic(err)
+		}
+	}
 }