From ef989bef6362ad42eec9376a512a5d7be2b7e9e7 Mon Sep 17 00:00:00 2001
From: Neil Alexander <neilalexander@users.noreply.github.com>
Date: Sun, 11 Aug 2024 10:41:58 +0100
Subject: [PATCH] Multicast module state tweaks

---
 src/multicast/multicast.go            | 26 +++++++++++++++-----------
 src/multicast/multicast_darwin_cgo.go |  2 +-
 2 files changed, 16 insertions(+), 12 deletions(-)

diff --git a/src/multicast/multicast.go b/src/multicast/multicast.go
index bf10370f..0fe72dc7 100644
--- a/src/multicast/multicast.go
+++ b/src/multicast/multicast.go
@@ -9,6 +9,7 @@ import (
 	"math/rand"
 	"net"
 	"net/url"
+	"sync/atomic"
 	"time"
 
 	"github.com/Arceliar/phony"
@@ -28,7 +29,7 @@ type Multicast struct {
 	core        *core.Core
 	log         *log.Logger
 	sock        *ipv6.PacketConn
-	_isOpen     bool
+	running     atomic.Bool
 	_listeners  map[string]*listenerInfo
 	_interfaces map[string]*interfaceInfo
 	_timer      *time.Timer
@@ -79,7 +80,7 @@ func New(core *core.Core, log *log.Logger, opts ...SetupOption) (*Multicast, err
 }
 
 func (m *Multicast) _start() error {
-	if m._isOpen {
+	if !m.running.CompareAndSwap(false, true) {
 		return fmt.Errorf("multicast module is already started")
 	}
 	var anyEnabled bool
@@ -87,12 +88,14 @@ func (m *Multicast) _start() error {
 		anyEnabled = anyEnabled || intf.Beacon || intf.Listen
 	}
 	if !anyEnabled {
+		m.running.Store(false)
 		return nil
 	}
 	m.log.Debugln("Starting multicast module")
 	defer m.log.Debugln("Started multicast module")
 	addr, err := net.ResolveUDPAddr("udp", string(m.config._groupAddr))
 	if err != nil {
+		m.running.Store(false)
 		return err
 	}
 	listenString := fmt.Sprintf("[::]:%v", addr.Port)
@@ -101,6 +104,7 @@ func (m *Multicast) _start() error {
 	}
 	conn, err := lc.ListenPacket(context.Background(), "udp6", listenString)
 	if err != nil {
+		m.running.Store(false)
 		return err
 	}
 	m.sock = ipv6.NewPacketConn(conn)
@@ -108,7 +112,6 @@ func (m *Multicast) _start() error {
 		// Windows can't set this flag, so we need to handle it in other ways
 	}
 
-	m._isOpen = true
 	go m.listen()
 	m.Act(nil, m._multicastStarted)
 	m.Act(nil, m._announce)
@@ -118,11 +121,7 @@ func (m *Multicast) _start() error {
 
 // IsStarted returns true if the module has been started.
 func (m *Multicast) IsStarted() bool {
-	var isOpen bool
-	phony.Block(m, func() {
-		isOpen = m._isOpen
-	})
-	return isOpen
+	return m.running.Load()
 }
 
 // Stop stops the multicast module.
@@ -136,8 +135,10 @@ func (m *Multicast) Stop() error {
 }
 
 func (m *Multicast) _stop() error {
+	if !m.running.CompareAndSwap(true, false) {
+		return nil
+	}
 	m.log.Infoln("Stopping multicast module")
-	m._isOpen = false
 	if m.sock != nil {
 		m.sock.Close()
 	}
@@ -233,7 +234,7 @@ func (m *Multicast) AnnounceNow() {
 }
 
 func (m *Multicast) _announce() {
-	if !m._isOpen {
+	if !m.running.Load() {
 		return
 	}
 	m._updateInterfaces()
@@ -250,7 +251,7 @@ func (m *Multicast) _announce() {
 	for name, info := range m._listeners {
 		// Prepare our stop function!
 		stop := func() {
-			info.listener.Close()
+			info.listener.Cancel()
 			delete(m._listeners, name)
 			m.log.Debugln("No longer multicasting on", name)
 		}
@@ -376,6 +377,9 @@ func (m *Multicast) listen() {
 	bs := make([]byte, 2048)
 	hb := make([]byte, 0, blake2b.Size) // Reused to reduce hash allocations
 	for {
+		if !m.running.Load() {
+			return
+		}
 		n, rcm, fromAddr, err := m.sock.ReadFrom(bs)
 		if err != nil {
 			if !m.IsStarted() {
diff --git a/src/multicast/multicast_darwin_cgo.go b/src/multicast/multicast_darwin_cgo.go
index 5c2af7ab..16266381 100644
--- a/src/multicast/multicast_darwin_cgo.go
+++ b/src/multicast/multicast_darwin_cgo.go
@@ -31,7 +31,7 @@ import (
 )
 
 func (m *Multicast) _multicastStarted() {
-	if !m._isOpen {
+	if !m.running.Load() {
 		return
 	}
 	C.StopAWDLBrowsing()