Extract IGMPv2 core state machine

The IGMPv2 core state machine can be shared with MLDv1 since they are
almost identical, ignoring specific addresses, constants and packets.

Bug #4682, #4861

PiperOrigin-RevId: 344102615

5 files changed, 737 insertions, 177 deletions

diff --git a/pkg/tcpip/network/ip/BUILD b/pkg/tcpip/network/ip/BUILD
new file mode 100644
index 000000000..6ca200b48
--- /dev/null
+++ b/pkg/tcpip/network/ip/BUILD
@@ -0,0 +1,25 @@
+load("//tools:defs.bzl", "go_library", "go_test")
+
+package(licenses = ["notice"])
+
+go_library(
+    name = "ip",
+    srcs = ["generic_multicast_protocol.go"],
+    visibility = ["//visibility:public"],
+    deps = [
+        "//pkg/sync",
+        "//pkg/tcpip",
+    ],
+)
+
+go_test(
+    name = "ip_test",
+    size = "small",
+    srcs = ["generic_multicast_protocol_test.go"],
+    deps = [
+        ":ip",
+        "//pkg/tcpip",
+        "//pkg/tcpip/faketime",
+        "@com_github_google_go_cmp//cmp:go_default_library",
+    ],
+)
diff --git a/pkg/tcpip/network/ip/generic_multicast_protocol.go b/pkg/tcpip/network/ip/generic_multicast_protocol.go
new file mode 100644
index 000000000..3113f4bbe
--- /dev/null
+++ b/pkg/tcpip/network/ip/generic_multicast_protocol.go
@@ -0,0 +1,346 @@
+// Copyright 2020 The gVisor Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// Package ip holds IPv4/IPv6 common utilities.
+package ip
+
+import (
+	"fmt"
+	"math/rand"
+	"time"
+
+	"gvisor.dev/gvisor/pkg/sync"
+	"gvisor.dev/gvisor/pkg/tcpip"
+)
+
+// hostState is the state a host may be in for a multicast group.
+type hostState int
+
+// The states below are generic across IGMPv2 (RFC 2236 section 6) and MLDv1
+// (RFC 2710 section 5). Even though the states are generic across both IGMPv2
+// and MLDv1, IGMPv2 terminology will be used.
+const (
+	// "'Non-Member' state, when the host does not belong to the group on
+	// the interface.  This is the initial state for all memberships on
+	// all network interfaces; it requires no storage in the host."
+	//
+	// 'Non-Listener' is the MLDv1 term used to describe this state.
+	_ hostState = iota
+
+	// delayingMember is the "'Delaying Member' state, when the host belongs to
+	// the group on the interface and has a report delay timer running for that
+	// membership."
+	//
+	// 'Delaying Listener' is the MLDv1 term used to describe this state.
+	delayingMember
+
+	// idleMember is the "Idle Member" state, when the host belongs to the group
+	// on the interface and does not have a report delay timer running for that
+	// membership.
+	//
+	// 'Idle Listener' is the MLDv1 term used to describe this state.
+	idleMember
+)
+
+// multicastGroupState holds the Generic Multicast Protocol state for a
+// multicast group.
+type multicastGroupState struct {
+	// state contains the host's state for the group.
+	state hostState
+
+	// lastToSendReport is true if we sent the last report for the group. It is
+	// used to track whether there are other hosts on the subnet that are also
+	// members of the group.
+	//
+	// Defined in RFC 2236 section 6 page 9 for IGMPv2 and RFC 2710 section 5 page
+	// 8 for MLDv1.
+	lastToSendReport bool
+
+	// delayedReportJob is used to delay sending responses to membership report
+	// messages in order to reduce duplicate reports from multiple hosts on the
+	// interface.
+	//
+	// Must not be nil.
+	delayedReportJob *tcpip.Job
+}
+
+// MulticastGroupProtocol is a multicast group protocol whose core state machine
+// can be represented by GenericMulticastProtocolState.
+type MulticastGroupProtocol interface {
+	// SendReport sends a multicast report for the specified group address.
+	SendReport(groupAddress tcpip.Address) *tcpip.Error
+
+	// SendLeave sends a multicast leave for the specified group address.
+	SendLeave(groupAddress tcpip.Address) *tcpip.Error
+}
+
+// GenericMulticastProtocolState is the per interface generic multicast protocol
+// state.
+//
+// There is actually no protocol named "Generic Multicast Protocol". Instead,
+// the term used to refer to a generic multicast protocol that applies to both
+// IPv4 and IPv6. Specifically, Generic Multicast Protocol is the core state
+// machine of IGMPv2 as defined by RFC 2236 and MLDv1 as defined by RFC 2710.
+//
+// GenericMulticastProtocolState.Init MUST be called before calling any of
+// the methods on GenericMulticastProtocolState.
+type GenericMulticastProtocolState struct {
+	rand                      *rand.Rand
+	clock                     tcpip.Clock
+	protocol                  MulticastGroupProtocol
+	maxUnsolicitedReportDelay time.Duration
+
+	mu struct {
+		sync.Mutex
+
+		// memberships holds group addresses and their associated state.
+		memberships map[tcpip.Address]multicastGroupState
+	}
+}
+
+// Init initializes the Generic Multicast Protocol state.
+//
+// maxUnsolicitedReportDelay is the maximum time between sending unsolicited
+// reports after joining a group.
+func (g *GenericMulticastProtocolState) Init(rand *rand.Rand, clock tcpip.Clock, protocol MulticastGroupProtocol, maxUnsolicitedReportDelay time.Duration) {
+	g.mu.Lock()
+	defer g.mu.Unlock()
+	g.rand = rand
+	g.clock = clock
+	g.protocol = protocol
+	g.maxUnsolicitedReportDelay = maxUnsolicitedReportDelay
+	g.mu.memberships = make(map[tcpip.Address]multicastGroupState)
+}
+
+// JoinGroup handles joining a new group.
+//
+// Returns false if the group has already been joined.
+func (g *GenericMulticastProtocolState) JoinGroup(groupAddress tcpip.Address) bool {
+	g.mu.Lock()
+	defer g.mu.Unlock()
+
+	if _, ok := g.mu.memberships[groupAddress]; ok {
+		// The group has already been joined.
+		return false
+	}
+
+	info := multicastGroupState{
+		// There isn't a job scheduled currently, so it's just idle.
+		state: idleMember,
+		// Joining a group immediately sends a report.
+		lastToSendReport: true,
+		delayedReportJob: tcpip.NewJob(g.clock, &g.mu, func() {
+			info, ok := g.mu.memberships[groupAddress]
+			if !ok {
+				panic(fmt.Sprintf("expected to find group state for group = %s", groupAddress))
+			}
+
+			info.lastToSendReport = g.protocol.SendReport(groupAddress) == nil
+			info.state = idleMember
+			g.mu.memberships[groupAddress] = info
+		}),
+	}
+
+	// As per RFC 2236 section 3 page 5 (for IGMPv2),
+	//
+	//   When a host joins a multicast group, it should immediately transmit an
+	//   unsolicited Version 2 Membership Report for that group" ... "it is
+	//   recommended that it be repeated".
+	//
+	// As per RFC 2710 section 4 page 6 (for MLDv1),
+	//
+	//   When a node starts listening to a multicast address on an interface,
+	//   it should immediately transmit an unsolicited Report for that address
+	//   on that interface, in case it is the first listener on the link. To
+	//   cover the possibility of the initial Report being lost or damaged, it
+	//   is recommended that it be repeated once or twice after short delays
+	//   [Unsolicited Report Interval].
+	//
+	// TODO(gvisor.dev/issue/4901): Support a configurable number of initial
+	// unsolicited reports.
+	info.lastToSendReport = g.protocol.SendReport(groupAddress) == nil
+	g.setDelayTimerForAddressRLocked(groupAddress, &info, g.maxUnsolicitedReportDelay)
+	g.mu.memberships[groupAddress] = info
+	return true
+}
+
+// LeaveGroup handles leaving the group.
+func (g *GenericMulticastProtocolState) LeaveGroup(groupAddress tcpip.Address) {
+	g.mu.Lock()
+	defer g.mu.Unlock()
+
+	info, ok := g.mu.memberships[groupAddress]
+	if !ok {
+		return
+	}
+
+	info.delayedReportJob.Cancel()
+	delete(g.mu.memberships, groupAddress)
+	if info.lastToSendReport {
+		// Okay to ignore the error here as if packet write failed, the multicast
+		// routers will eventually drop our membership anyways. If the interface is
+		// being disabled or removed, the generic multicast protocol's should be
+		// cleared eventually.
+		//
+		// As per RFC 2236 section 3 page 5 (for IGMPv2),
+		//
+		//   When a router receives a Report, it adds the group being reported to
+		//   the list of multicast group memberships on the network on which it
+		//   received the Report and sets the timer for the membership to the
+		//   [Group Membership Interval]. Repeated Reports refresh the timer. If
+		//   no Reports are received for a particular group before this timer has
+		//   expired, the router assumes that the group has no local members and
+		//   that it need not forward remotely-originated multicasts for that
+		//   group onto the attached network.
+		//
+		// As per RFC 2710 section 4 page 5 (for MLDv1),
+		//
+		//   When a router receives a Report from a link, if the reported address
+		//   is not already present in the router's list of multicast address
+		//   having listeners on that link, the reported address is added to the
+		//   list, its timer is set to [Multicast Listener Interval], and its
+		//   appearance is made known to the router's multicast routing component.
+		//   If a Report is received for a multicast address that is already
+		//   present in the router's list, the timer for that address is reset to
+		//   [Multicast Listener Interval]. If an address's timer expires, it is
+		//   assumed that there are no longer any listeners for that address
+		//   present on the link, so it is deleted from the list and its
+		//   disappearance is made known to the multicast routing component.
+		//
+		// The requirement to send a leave message is also optional (it MAY be
+		// skipped):
+		//
+		// As per RFC 2236 section 6 page 8 (for IGMPv2),
+		//
+		//  "send leave" for the group on the interface. If the interface
+		//   state says the Querier is running IGMPv1, this action SHOULD be
+		//   skipped. If the flag saying we were the last host to report is
+		//   cleared, this action MAY be skipped. The Leave Message is sent to
+		//   the ALL-ROUTERS group (224.0.0.2).
+		//
+		// As per RFC 2710 section 5 page 8 (for MLDv1),
+		//
+		//   "send done" for the address on the interface. If the flag saying
+		//   we were the last node to report is cleared, this action MAY be
+		//   skipped. The Done message is sent to the link-scope all-routers
+		//   address (FF02::2).
+		_ = g.protocol.SendLeave(groupAddress)
+	}
+}
+
+// HandleQuery handles a query message with the specified maximum response time.
+//
+// If the group address is unspecified, then reports will be scheduled for all
+// joined groups.
+//
+// Report(s) will be scheduled to be sent after a random duration between 0 and
+// the maximum response time.
+func (g *GenericMulticastProtocolState) HandleQuery(groupAddress tcpip.Address, maxResponseTime time.Duration) {
+	g.mu.Lock()
+	defer g.mu.Unlock()
+
+	// As per RFC 2236 section 2.4 (for IGMPv2),
+	//
+	//   In a Membership Query message, the group address field is set to zero
+	//   when sending a General Query, and set to the group address being
+	//   queried when sending a Group-Specific Query.
+	//
+	// As per RFC 2710 section 3.6 (for MLDv1),
+	//
+	//   In a Query message, the Multicast Address field is set to zero when
+	//   sending a General Query, and set to a specific IPv6 multicast address
+	//   when sending a Multicast-Address-Specific Query.
+	if groupAddress.Unspecified() {
+		// This is a general query as the group address is unspecified.
+		for groupAddress, info := range g.mu.memberships {
+			g.setDelayTimerForAddressRLocked(groupAddress, &info, maxResponseTime)
+			g.mu.memberships[groupAddress] = info
+		}
+	} else if info, ok := g.mu.memberships[groupAddress]; ok {
+		g.setDelayTimerForAddressRLocked(groupAddress, &info, maxResponseTime)
+		g.mu.memberships[groupAddress] = info
+	}
+}
+
+// HandleReport handles a report message.
+//
+// If the report is for a joined group, any active delayed report will be
+// cancelled and the host state for the group transitions to idle.
+func (g *GenericMulticastProtocolState) HandleReport(groupAddress tcpip.Address) {
+	g.mu.Lock()
+	defer g.mu.Unlock()
+
+	// As per RFC 2236 section 3 pages 3-4 (for IGMPv2),
+	//
+	//   If the host receives another host's Report (version 1 or 2) while it has
+	//   a timer running, it stops its timer for the specified group and does not
+	//   send a Report
+	//
+	// As per RFC 2710 section 4 page 6 (for MLDv1),
+	//
+	//   If a node receives another node's Report from an interface for a
+	//   multicast address while it has a timer running for that same address
+	//   on that interface, it stops its timer and does not send a Report for
+	//   that address, thus suppressing duplicate reports on the link.
+	if info, ok := g.mu.memberships[groupAddress]; ok {
+		info.delayedReportJob.Cancel()
+		info.lastToSendReport = false
+		info.state = idleMember
+		g.mu.memberships[groupAddress] = info
+	}
+}
+
+// setDelayTimerForAddressRLocked sets timer to send a delay report.
+//
+// Precondition: g.mu MUST be read locked.
+func (g *GenericMulticastProtocolState) setDelayTimerForAddressRLocked(groupAddress tcpip.Address, info *multicastGroupState, maxResponseTime time.Duration) {
+	// As per RFC 2236 section 3 page 3 (for IGMPv2),
+	//
+	//   If a timer for the group is already unning, it is reset to the random
+	//   value only if the requested Max Response Time is less than the remaining
+	//   value of the running timer.
+	//
+	// As per RFC 2710 section 4 page 5 (for MLDv1),
+	//
+	//   If a timer for any address is already running, it is reset to the new
+	//   random value only if the requested Maximum Response Delay is less than
+	//   the remaining value of the running timer.
+	if info.state == delayingMember {
+		// TODO: Reset the timer if time remaining is greater than maxResponseTime.
+		return
+	}
+	info.state = delayingMember
+	info.delayedReportJob.Cancel()
+	info.delayedReportJob.Schedule(g.calculateDelayTimerDuration(maxResponseTime))
+}
+
+// calculateDelayTimerDuration returns a random time between (0, maxRespTime].
+func (g *GenericMulticastProtocolState) calculateDelayTimerDuration(maxRespTime time.Duration) time.Duration {
+	// As per RFC 2236 section 3 page 3 (for IGMPv2),
+	//
+	//   When a host receives a Group-Specific Query, it sets a delay timer to a
+	//   random value selected from the range (0, Max Response Time]...
+	//
+	// As per RFC 2710 section 4 page 6 (for MLDv1),
+	//
+	//   When a node receives a Multicast-Address-Specific Query, if it is
+	//   listening to the queried Multicast Address on the interface from
+	//   which the Query was received, it sets a delay timer for that address
+	//   to a random value selected from the range [0, Maximum Response Delay],
+	//   as above.
+	if maxRespTime == 0 {
+		return 0
+	}
+	return time.Duration(g.rand.Int63n(int64(maxRespTime)))
+}
diff --git a/pkg/tcpip/network/ip/generic_multicast_protocol_test.go b/pkg/tcpip/network/ip/generic_multicast_protocol_test.go
new file mode 100644
index 000000000..eb48c0d51
--- /dev/null
+++ b/pkg/tcpip/network/ip/generic_multicast_protocol_test.go
@@ -0,0 +1,294 @@
+// Copyright 2020 The gVisor Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package ip_test
+
+import (
+	"math/rand"
+	"testing"
+	"time"
+
+	"github.com/google/go-cmp/cmp"
+	"gvisor.dev/gvisor/pkg/tcpip"
+	"gvisor.dev/gvisor/pkg/tcpip/faketime"
+	"gvisor.dev/gvisor/pkg/tcpip/network/ip"
+)
+
+const (
+	addr1 = tcpip.Address("\x01")
+	addr2 = tcpip.Address("\x02")
+	addr3 = tcpip.Address("\x03")
+)
+
+var _ ip.MulticastGroupProtocol = (*mockMulticastGroupProtocol)(nil)
+
+type mockMulticastGroupProtocol struct {
+	sendReportGroupAddrCount map[tcpip.Address]int
+	sendLeaveGroupAddr       tcpip.Address
+}
+
+func (m *mockMulticastGroupProtocol) init() {
+	m.sendReportGroupAddrCount = make(map[tcpip.Address]int)
+	m.sendLeaveGroupAddr = ""
+}
+
+func (m *mockMulticastGroupProtocol) SendReport(groupAddress tcpip.Address) *tcpip.Error {
+	m.sendReportGroupAddrCount[groupAddress]++
+	return nil
+}
+
+func (m *mockMulticastGroupProtocol) SendLeave(groupAddress tcpip.Address) *tcpip.Error {
+	m.sendLeaveGroupAddr = groupAddress
+	return nil
+}
+
+func checkProtocol(mgp *mockMulticastGroupProtocol, sendReportGroupAddresses []tcpip.Address, sendLeaveGroupAddr tcpip.Address) string {
+	sendReportGroupAddressesMap := make(map[tcpip.Address]int)
+	for _, a := range sendReportGroupAddresses {
+		sendReportGroupAddressesMap[a] = 1
+	}
+
+	diff := cmp.Diff(mockMulticastGroupProtocol{
+		sendReportGroupAddrCount: sendReportGroupAddressesMap,
+		sendLeaveGroupAddr:       sendLeaveGroupAddr,
+	}, *mgp, cmp.AllowUnexported(mockMulticastGroupProtocol{}))
+	mgp.init()
+	return diff
+}
+
+func TestJoinGroup(t *testing.T) {
+	const maxUnsolicitedReportDelay = time.Second
+
+	var g ip.GenericMulticastProtocolState
+	var mgp mockMulticastGroupProtocol
+	mgp.init()
+	clock := faketime.NewManualClock()
+	g.Init(rand.New(rand.NewSource(0)), clock, &mgp, maxUnsolicitedReportDelay)
+
+	// Joining a group should send a report immediately and another after
+	// a random interval between 0 and the maximum unsolicited report delay.
+	if !g.JoinGroup(addr1) {
+		t.Errorf("got g.JoinGroup(%s) = false, want = true", addr1)
+	}
+	if diff := checkProtocol(&mgp, []tcpip.Address{addr1} /* sendReportGroupAddresses */, "" /* sendLeaveGroupAddr */); diff != "" {
+		t.Errorf("mockMulticastGroupProtocol mismatch (-want +got):\n%s", diff)
+	}
+
+	clock.Advance(maxUnsolicitedReportDelay)
+	if diff := checkProtocol(&mgp, []tcpip.Address{addr1} /* sendReportGroupAddresses */, "" /* sendLeaveGroupAddr */); diff != "" {
+		t.Errorf("mockMulticastGroupProtocol mismatch (-want +got):\n%s", diff)
+	}
+
+	// Should have no more messages to send.
+	clock.Advance(time.Hour)
+	if diff := checkProtocol(&mgp, nil /* sendReportGroupAddresses */, "" /* sendLeaveGroupAddr */); diff != "" {
+		t.Errorf("mockMulticastGroupProtocol mismatch (-want +got):\n%s", diff)
+	}
+}
+
+func TestLeaveGroup(t *testing.T) {
+	const maxUnsolicitedReportDelay = time.Second
+
+	var g ip.GenericMulticastProtocolState
+	var mgp mockMulticastGroupProtocol
+	mgp.init()
+	clock := faketime.NewManualClock()
+	g.Init(rand.New(rand.NewSource(1)), clock, &mgp, maxUnsolicitedReportDelay)
+
+	if !g.JoinGroup(addr1) {
+		t.Fatalf("got g.JoinGroup(%s) = false, want = true", addr1)
+	}
+	if diff := checkProtocol(&mgp, []tcpip.Address{addr1} /* sendReportGroupAddresses */, "" /* sendLeaveGroupAddr */); diff != "" {
+		t.Fatalf("mockMulticastGroupProtocol mismatch (-want +got):\n%s", diff)
+	}
+
+	// Leaving a group should send a leave report immediately and cancel any
+	// delayed reports.
+	g.LeaveGroup(addr1)
+	if diff := checkProtocol(&mgp, nil /* sendReportGroupAddresses */, addr1 /* sendLeaveGroupAddr */); diff != "" {
+		t.Errorf("mockMulticastGroupProtocol mismatch (-want +got):\n%s", diff)
+	}
+
+	// Should have no more messages to send.
+	clock.Advance(time.Hour)
+	if diff := checkProtocol(&mgp, nil /* sendReportGroupAddresses */, "" /* sendLeaveGroupAddr */); diff != "" {
+		t.Errorf("mockMulticastGroupProtocol mismatch (-want +got):\n%s", diff)
+	}
+}
+
+func TestHandleReport(t *testing.T) {
+	const maxUnsolicitedReportDelay = time.Second
+
+	tests := []struct {
+		name             string
+		reportAddr       tcpip.Address
+		expectReportsFor []tcpip.Address
+	}{
+		{
+			name:             "Unpecified empty",
+			reportAddr:       "",
+			expectReportsFor: []tcpip.Address{addr1, addr2},
+		},
+		{
+			name:             "Unpecified any",
+			reportAddr:       "\x00",
+			expectReportsFor: []tcpip.Address{addr1, addr2},
+		},
+		{
+			name:             "Specified",
+			reportAddr:       addr1,
+			expectReportsFor: []tcpip.Address{addr2},
+		},
+		{
+			name:             "Specified other",
+			reportAddr:       addr3,
+			expectReportsFor: []tcpip.Address{addr1, addr2},
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			var g ip.GenericMulticastProtocolState
+			var mgp mockMulticastGroupProtocol
+			mgp.init()
+			clock := faketime.NewManualClock()
+			g.Init(rand.New(rand.NewSource(2)), clock, &mgp, maxUnsolicitedReportDelay)
+
+			if !g.JoinGroup(addr1) {
+				t.Fatalf("got g.JoinGroup(%s) = false, want = true", addr1)
+			}
+			if diff := checkProtocol(&mgp, []tcpip.Address{addr1} /* sendReportGroupAddresses */, "" /* sendLeaveGroupAddr */); diff != "" {
+				t.Fatalf("mockMulticastGroupProtocol mismatch (-want +got):\n%s", diff)
+			}
+			if !g.JoinGroup(addr2) {
+				t.Fatalf("got g.JoinGroup(%s) = false, want = true", addr2)
+			}
+			if diff := checkProtocol(&mgp, []tcpip.Address{addr2} /* sendReportGroupAddresses */, "" /* sendLeaveGroupAddr */); diff != "" {
+				t.Fatalf("mockMulticastGroupProtocol mismatch (-want +got):\n%s", diff)
+			}
+
+			// Receiving a report for a group we have a timer scheduled for should
+			// cancel our delayed report timer for the group.
+			g.HandleReport(test.reportAddr)
+			if len(test.expectReportsFor) != 0 {
+				clock.Advance(maxUnsolicitedReportDelay)
+				if diff := checkProtocol(&mgp, test.expectReportsFor /* sendReportGroupAddresses */, "" /* sendLeaveGroupAddr */); diff != "" {
+					t.Errorf("mockMulticastGroupProtocol mismatch (-want +got):\n%s", diff)
+				}
+			}
+
+			// Should have no more messages to send.
+			clock.Advance(time.Hour)
+			if diff := checkProtocol(&mgp, nil /* sendReportGroupAddresses */, "" /* sendLeaveGroupAddr */); diff != "" {
+				t.Errorf("mockMulticastGroupProtocol mismatch (-want +got):\n%s", diff)
+			}
+		})
+	}
+}
+
+func TestHandleQuery(t *testing.T) {
+	const maxUnsolicitedReportDelay = time.Second
+
+	tests := []struct {
+		name             string
+		queryAddr        tcpip.Address
+		maxDelay         time.Duration
+		expectReportsFor []tcpip.Address
+	}{
+		{
+			name:             "Unpecified empty",
+			queryAddr:        "",
+			maxDelay:         0,
+			expectReportsFor: []tcpip.Address{addr1, addr2},
+		},
+		{
+			name:             "Unpecified any",
+			queryAddr:        "\x00",
+			maxDelay:         1,
+			expectReportsFor: []tcpip.Address{addr1, addr2},
+		},
+		{
+			name:             "Specified",
+			queryAddr:        addr1,
+			maxDelay:         2,
+			expectReportsFor: []tcpip.Address{addr1},
+		},
+		{
+			name:             "Specified other",
+			queryAddr:        addr3,
+			maxDelay:         3,
+			expectReportsFor: nil,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			var g ip.GenericMulticastProtocolState
+			var mgp mockMulticastGroupProtocol
+			mgp.init()
+			clock := faketime.NewManualClock()
+			g.Init(rand.New(rand.NewSource(3)), clock, &mgp, maxUnsolicitedReportDelay)
+
+			if !g.JoinGroup(addr1) {
+				t.Fatalf("got g.JoinGroup(%s) = false, want = true", addr1)
+			}
+			if diff := checkProtocol(&mgp, []tcpip.Address{addr1} /* sendReportGroupAddresses */, "" /* sendLeaveGroupAddr */); diff != "" {
+				t.Fatalf("mockMulticastGroupProtocol mismatch (-want +got):\n%s", diff)
+			}
+			if !g.JoinGroup(addr2) {
+				t.Fatalf("got g.JoinGroup(%s) = false, want = true", addr2)
+			}
+			if diff := checkProtocol(&mgp, []tcpip.Address{addr2} /* sendReportGroupAddresses */, "" /* sendLeaveGroupAddr */); diff != "" {
+				t.Fatalf("mockMulticastGroupProtocol mismatch (-want +got):\n%s", diff)
+			}
+			clock.Advance(maxUnsolicitedReportDelay)
+			if diff := checkProtocol(&mgp, []tcpip.Address{addr1, addr2} /* sendReportGroupAddresses */, "" /* sendLeaveGroupAddr */); diff != "" {
+				t.Fatalf("mockMulticastGroupProtocol mismatch (-want +got):\n%s", diff)
+			}
+
+			// Receiving a query should make us schedule a new delayed report if it
+			// is a query directed at us or a general query.
+			g.HandleQuery(test.queryAddr, test.maxDelay)
+			if len(test.expectReportsFor) != 0 {
+				clock.Advance(test.maxDelay)
+				if diff := checkProtocol(&mgp, test.expectReportsFor /* sendReportGroupAddresses */, "" /* sendLeaveGroupAddr */); diff != "" {
+					t.Errorf("mockMulticastGroupProtocol mismatch (-want +got):\n%s", diff)
+				}
+			}
+
+			// Should have no more messages to send.
+			clock.Advance(time.Hour)
+			if diff := checkProtocol(&mgp, nil /* sendReportGroupAddresses */, "" /* sendLeaveGroupAddr */); diff != "" {
+				t.Errorf("mockMulticastGroupProtocol mismatch (-want +got):\n%s", diff)
+			}
+		})
+	}
+}
+
+func TestDoubleJoinGroup(t *testing.T) {
+	var g ip.GenericMulticastProtocolState
+	var mgp mockMulticastGroupProtocol
+	mgp.init()
+	clock := faketime.NewManualClock()
+	g.Init(rand.New(rand.NewSource(4)), clock, &mgp, time.Second)
+
+	if !g.JoinGroup(addr1) {
+		t.Fatalf("got g.JoinGroup(%s) = false, want = true", addr1)
+	}
+
+	// Joining the same group twice should fail.
+	if g.JoinGroup(addr1) {
+		t.Errorf("got g.JoinGroup(%s) = true, want = false", addr1)
+	}
+}
diff --git a/pkg/tcpip/network/ipv4/BUILD b/pkg/tcpip/network/ipv4/BUILD
index 68b1ea1cd..32f53f217 100644
--- a/pkg/tcpip/network/ipv4/BUILD
+++ b/pkg/tcpip/network/ipv4/BUILD
@@ -18,6 +18,7 @@ go_library(
         "//pkg/tcpip/header/parse",
         "//pkg/tcpip/network/fragmentation",
         "//pkg/tcpip/network/hash",
+        "//pkg/tcpip/network/ip",
         "//pkg/tcpip/stack",
     ],
 )
diff --git a/pkg/tcpip/network/ipv4/igmp.go b/pkg/tcpip/network/ipv4/igmp.go
index 18fe2fd2f..c9bf117de 100644
--- a/pkg/tcpip/network/ipv4/igmp.go
+++ b/pkg/tcpip/network/ipv4/igmp.go
@@ -17,11 +17,13 @@ package ipv4
 import (
 	"fmt"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"gvisor.dev/gvisor/pkg/tcpip"
 	"gvisor.dev/gvisor/pkg/tcpip/buffer"
 	"gvisor.dev/gvisor/pkg/tcpip/header"
+	"gvisor.dev/gvisor/pkg/tcpip/network/ip"
 	"gvisor.dev/gvisor/pkg/tcpip/stack"
 )
 
@@ -29,7 +31,7 @@ const (
 	// igmpV1PresentDefault is the initial state for igmpV1Present in the
 	// igmpState. As per RFC 2236 Page 9 says "No IGMPv1 Router Present ... is
 	// the initial state."
-	igmpV1PresentDefault = false
+	igmpV1PresentDefault = 0
 
 	// v1RouterPresentTimeout from RFC 2236 Section 8.11, Page 18
 	// See note on igmpState.igmpV1Present for more detail.
@@ -49,6 +51,8 @@ const (
 	UnsolicitedReportIntervalMax = 10 * time.Second
 )
 
+var _ ip.MulticastGroupProtocol = (*igmpState)(nil)
+
 // igmpState is the per-interface IGMP state.
 //
 // igmpState.init() MUST be called after creating an IGMP state.
@@ -56,22 +60,23 @@ type igmpState struct {
 	// The IPv4 endpoint this igmpState is for.
 	ep *endpoint
 
+	// igmpV1Present is for maintaining compatibility with IGMPv1 Routers, from
+	// RFC 2236 Section 4 Page 6: "The IGMPv1 router expects Version 1
+	// Membership Reports in response to its Queries, and will not pay
+	// attention to Version 2 Membership Reports.  Therefore, a state variable
+	// MUST be kept for each interface, describing whether the multicast
+	// Querier on that interface is running IGMPv1 or IGMPv2.  This variable
+	// MUST be based upon whether or not an IGMPv1 query was heard in the last
+	// [Version 1 Router Present Timeout] seconds".
+	//
+	// Must be accessed with atomic operations. Holds a value of 1 when true, 0
+	// when false.
+	igmpV1Present uint32
+
 	mu struct {
 		sync.RWMutex
 
-		// memberships contains the map of host groups to their state, timer, and
-		// flag info.
-		memberships map[tcpip.Address]membershipInfo
-
-		// igmpV1Present is for maintaining compatibility with IGMPv1 Routers, from
-		// RFC 2236 Section 4 Page 6: "The IGMPv1 router expects Version 1
-		// Membership Reports in response to its Queries, and will not pay
-		// attention to Version 2 Membership Reports.  Therefore, a state variable
-		// MUST be kept for each interface, describing whether the multicast
-		// Querier on that interface is running IGMPv1 or IGMPv2.  This variable
-		// MUST be based upon whether or not an IGMPv1 query was heard in the last
-		// [Version 1 Router Present Timeout] seconds"
-		igmpV1Present bool
+		genericMulticastProtocol ip.GenericMulticastProtocolState
 
 		// igmpV1Job is scheduled when this interface receives an IGMPv1 style
 		// message, upon expiration the igmpV1Present flag is cleared.
@@ -80,43 +85,26 @@ type igmpState struct {
 	}
 }
 
-// membershipInfo holds the IGMPv2 state for a particular multicast address.
-type membershipInfo struct {
-	// state contains the current IGMP state for this member.
-	state hostState
-
-	// lastToSendReport is true if this was "the last host to send a report from
-	// this group."
-	// RFC 2236, Section 6, Page 9. This is used to track whether or not there
-	// are other hosts on this subnet that belong to this group - RFC 2236
-	// Section 3, Page 5.
-	lastToSendReport bool
-
-	// delayedReportJob is used to delay sending responses to IGMP messages in
-	// order to reduce duplicate reports from multiple hosts on the interface.
-	// Must not be nil.
-	delayedReportJob *tcpip.Job
+// SendReport implements ip.MulticastGroupProtocol.
+func (igmp *igmpState) SendReport(groupAddress tcpip.Address) *tcpip.Error {
+	igmpType := header.IGMPv2MembershipReport
+	if igmp.v1Present() {
+		igmpType = header.IGMPv1MembershipReport
+	}
+	return igmp.writePacket(groupAddress, groupAddress, igmpType)
 }
 
-type hostState int
-
-// From RFC 2236, Section 6, Page 7.
-const (
-	// "'Non-Member' state, when the host does not belong to the group on
-	// the interface.  This is the initial state for all memberships on
-	// all network interfaces; it requires no storage in the host."
-	_ hostState = iota
-
-	// delayingMember is the "'Delaying Member' state, when the host belongs to
-	// the group on the interface and has a report delay timer running for that
-	// membership."
-	delayingMember
-
-	// idleMember is the "Idle Member" state, when the host belongs to the group
-	// on the interface and does not have a report delay timer running for that
-	// membership.
-	idleMember
-)
+// SendLeave implements ip.MulticastGroupProtocol.
+func (igmp *igmpState) SendLeave(groupAddress tcpip.Address) *tcpip.Error {
+	// As per RFC 2236 Section 6, Page 8: "If the interface state says the
+	// Querier is running IGMPv1, this action SHOULD be skipped. If the flag
+	// saying we were the last host to report is cleared, this action MAY be
+	// skipped."
+	if igmp.v1Present() {
+		return nil
+	}
+	return igmp.writePacket(header.IPv4AllRoutersGroup, groupAddress, header.IGMPLeaveGroup)
+}
 
 // init sets up an igmpState struct, and is required to be called before using
 // a new igmpState.
@@ -124,10 +112,10 @@ func (igmp *igmpState) init(ep *endpoint) {
 	igmp.mu.Lock()
 	defer igmp.mu.Unlock()
 	igmp.ep = ep
-	igmp.mu.memberships = make(map[tcpip.Address]membershipInfo)
-	igmp.mu.igmpV1Present = igmpV1PresentDefault
+	igmp.mu.genericMulticastProtocol.Init(ep.protocol.stack.Rand(), ep.protocol.stack.Clock(), igmp, UnsolicitedReportIntervalMax)
+	igmp.igmpV1Present = igmpV1PresentDefault
 	igmp.mu.igmpV1Job = igmp.ep.protocol.stack.NewJob(&igmp.mu, func() {
-		igmp.mu.igmpV1Present = false
+		igmp.setV1Present(false)
 	})
 }
 
@@ -188,6 +176,18 @@ func (igmp *igmpState) handleIGMP(pkt *stack.PacketBuffer) {
 	}
 }
 
+func (igmp *igmpState) v1Present() bool {
+	return atomic.LoadUint32(&igmp.igmpV1Present) == 1
+}
+
+func (igmp *igmpState) setV1Present(v bool) {
+	if v {
+		atomic.StoreUint32(&igmp.igmpV1Present, 1)
+	} else {
+		atomic.StoreUint32(&igmp.igmpV1Present, 0)
+	}
+}
+
 func (igmp *igmpState) handleMembershipQuery(groupAddress tcpip.Address, maxRespTime time.Duration) {
 	igmp.mu.Lock()
 	defer igmp.mu.Unlock()
@@ -198,56 +198,22 @@ func (igmp *igmpState) handleMembershipQuery(groupAddress tcpip.Address, maxResp
 	if maxRespTime == 0 {
 		igmp.mu.igmpV1Job.Cancel()
 		igmp.mu.igmpV1Job.Schedule(v1RouterPresentTimeout)
-		igmp.mu.igmpV1Present = true
+		igmp.setV1Present(true)
 		maxRespTime = v1MaxRespTime
 	}
 
-	// IPv4Any is the General Query Address.
-	if groupAddress == header.IPv4Any {
-		for membershipAddress, info := range igmp.mu.memberships {
-			igmp.setDelayTimerForAddressRLocked(membershipAddress, &info, maxRespTime)
-			igmp.mu.memberships[membershipAddress] = info
-		}
-	} else if info, ok := igmp.mu.memberships[groupAddress]; ok {
-		igmp.setDelayTimerForAddressRLocked(groupAddress, &info, maxRespTime)
-		igmp.mu.memberships[groupAddress] = info
-	}
-}
-
-// setDelayTimerForAddressRLocked modifies the passed info only and does not
-// modify IGMP state directly.
-//
-// Precondition: igmp.mu MUST be read locked.
-func (igmp *igmpState) setDelayTimerForAddressRLocked(groupAddress tcpip.Address, info *membershipInfo, maxRespTime time.Duration) {
-	if info.state == delayingMember {
-		// As per RFC 2236 Section 3, page 3: "If a timer for the group is already
-		// running, it is reset to the random value only if the requested Max
-		// Response Time is less than the remaining value of the running timer.
-		// TODO: Reset the timer if time remaining is greater than maxRespTime.
-		return
-	}
-	info.state = delayingMember
-	info.delayedReportJob.Cancel()
-	info.delayedReportJob.Schedule(igmp.calculateDelayTimerDuration(maxRespTime))
+	igmp.mu.genericMulticastProtocol.HandleQuery(groupAddress, maxRespTime)
 }
 
 func (igmp *igmpState) handleMembershipReport(groupAddress tcpip.Address) {
 	igmp.mu.Lock()
 	defer igmp.mu.Unlock()
-
-	// As per RFC 2236 Section 3, pages 3-4: "If the host receives another host's
-	// Report (version 1 or 2) while it has a timer running, it stops its timer
-	// for the specified group and does not send a Report"
-	if info, ok := igmp.mu.memberships[groupAddress]; ok {
-		info.delayedReportJob.Cancel()
-		info.lastToSendReport = false
-		igmp.mu.memberships[groupAddress] = info
-	}
+	igmp.mu.genericMulticastProtocol.HandleReport(groupAddress)
 }
 
 // writePacket assembles and sends an IGMP packet with the provided fields,
 // incrementing the provided stat counter on success.
-func (igmp *igmpState) writePacket(destAddress tcpip.Address, groupAddress tcpip.Address, igmpType header.IGMPType) {
+func (igmp *igmpState) writePacket(destAddress tcpip.Address, groupAddress tcpip.Address, igmpType header.IGMPType) *tcpip.Error {
 	igmpData := header.IGMP(buffer.NewView(header.IGMPReportMinimumSize))
 	igmpData.SetType(igmpType)
 	igmpData.SetGroupAddress(groupAddress)
@@ -275,56 +241,19 @@ func (igmp *igmpState) writePacket(destAddress tcpip.Address, groupAddress tcpip
 	sent := igmp.ep.protocol.stack.Stats().IGMP.PacketsSent
 	if err := igmp.ep.nic.WritePacketToRemote(header.EthernetAddressFromMulticastIPv4Address(destAddress), nil /* gso */, header.IPv4ProtocolNumber, pkt); err != nil {
 		sent.Dropped.Increment()
-	} else {
-		switch igmpType {
-		case header.IGMPv1MembershipReport:
-			sent.V1MembershipReport.Increment()
-		case header.IGMPv2MembershipReport:
-			sent.V2MembershipReport.Increment()
-		case header.IGMPLeaveGroup:
-			sent.LeaveGroup.Increment()
-		default:
-			panic(fmt.Sprintf("unrecognized igmp type = %d", igmpType))
-		}
+		return err
 	}
-}
-
-// sendReport sends a Host Membership Report in response to a query or after
-// this host joins a new group on this interface.
-//
-// Precondition: igmp.mu MUST be locked.
-func (igmp *igmpState) sendReportLocked(groupAddress tcpip.Address) {
-	igmpType := header.IGMPv2MembershipReport
-	if igmp.mu.igmpV1Present {
-		igmpType = header.IGMPv1MembershipReport
-	}
-	igmp.writePacket(groupAddress, groupAddress, igmpType)
-
-	// Update the state of the membership for this group. If the group no longer
-	// exists, do nothing since this report must have been a race with a remove
-	// or is in the process of being added.
-	info, ok := igmp.mu.memberships[groupAddress]
-	if !ok {
-		return
-	}
-	info.state = idleMember
-	info.lastToSendReport = true
-	igmp.mu.memberships[groupAddress] = info
-}
-
-// sendLeave sends a Leave Group report to the IPv4 All Routers Group.
-//
-// Precondition: igmp.mu MUST be read locked.
-func (igmp *igmpState) sendLeaveRLocked(groupAddress tcpip.Address) {
-	// As per RFC 2236 Section 6, Page 8: "If the interface state says the
-	// Querier is running IGMPv1, this action SHOULD be skipped. If the flag
-	// saying we were the last host to report is cleared, this action MAY be
-	// skipped."
-	if igmp.mu.igmpV1Present || !igmp.mu.memberships[groupAddress].lastToSendReport {
-		return
+	switch igmpType {
+	case header.IGMPv1MembershipReport:
+		sent.V1MembershipReport.Increment()
+	case header.IGMPv2MembershipReport:
+		sent.V2MembershipReport.Increment()
+	case header.IGMPLeaveGroup:
+		sent.LeaveGroup.Increment()
+	default:
+		panic(fmt.Sprintf("unrecognized igmp type = %d", igmpType))
 	}
-
-	igmp.writePacket(header.IPv4AllRoutersGroup, groupAddress, header.IGMPLeaveGroup)
+	return nil
 }
 
 // joinGroup handles adding a new group to the membership map, setting up the
@@ -336,28 +265,11 @@ func (igmp *igmpState) sendLeaveRLocked(groupAddress tcpip.Address) {
 func (igmp *igmpState) joinGroup(groupAddress tcpip.Address) *tcpip.Error {
 	igmp.mu.Lock()
 	defer igmp.mu.Unlock()
-	if _, ok := igmp.mu.memberships[groupAddress]; ok {
-		// The group already exists in the membership map.
-		return tcpip.ErrDuplicateAddress
-	}
 
-	info := membershipInfo{
-		// There isn't a Job scheduled currently, so it's just idle.
-		state: idleMember,
-		// Joining a group immediately sends a report.
-		lastToSendReport: true,
-		delayedReportJob: igmp.ep.protocol.stack.NewJob(&igmp.mu, func() {
-			igmp.sendReportLocked(groupAddress)
-		}),
+	// JoinGroup returns false if we have already joined the group.
+	if !igmp.mu.genericMulticastProtocol.JoinGroup(groupAddress) {
+		return tcpip.ErrDuplicateAddress
 	}
-
-	// As per RFC 2236 Section 3, Page 5: "When a host joins a multicast group,
-	// it should immediately transmit an unsolicited Version 2 Membership Report
-	// for that group" ... "it is recommended that it be repeated"
-	igmp.sendReportLocked(groupAddress)
-	igmp.setDelayTimerForAddressRLocked(groupAddress, &info, UnsolicitedReportIntervalMax)
-	igmp.mu.memberships[groupAddress] = info
-
 	return nil
 }
 
@@ -370,23 +282,5 @@ func (igmp *igmpState) joinGroup(groupAddress tcpip.Address) *tcpip.Error {
 func (igmp *igmpState) leaveGroup(groupAddress tcpip.Address) {
 	igmp.mu.Lock()
 	defer igmp.mu.Unlock()
-	info, ok := igmp.mu.memberships[groupAddress]
-	if !ok {
-		return
-	}
-
-	// Clean up the state of the group before sending the leave message and
-	// removing it from the map.
-	info.delayedReportJob.Cancel()
-	info.state = idleMember
-	igmp.mu.memberships[groupAddress] = info
-
-	igmp.sendLeaveRLocked(groupAddress)
-	delete(igmp.mu.memberships, groupAddress)
-}
-
-// RFC 2236 Section 3, Page 3: The response time is set to a "random value...
-// selected from the range (0, Max Response Time]".
-func (igmp *igmpState) calculateDelayTimerDuration(maxRespTime time.Duration) time.Duration {
-	return time.Duration(igmp.ep.protocol.stack.Rand().Int63n(int64(maxRespTime)))
+	igmp.mu.genericMulticastProtocol.LeaveGroup(groupAddress)
 }