dbus: keep match rules across subscribe socket re-dials

port 2.0

core/internal/server/dbus/handlers.go

@@ -1,6 +1,7 @@
 package dbus
 
 import (
+	"context"
 	"fmt"
 
 	"github.com/duxlabs-dev/dux-go/ipc"
@@ -43,7 +44,7 @@ func extractObjectParams(p map[string]any, requirePath bool) (objectParams, erro
 	return objectParams{bus: bus, dest: dest, path: path, iface: iface}, nil
 }
 
-func HandleRequest(conn *ipc.ConnWriter, req ipc.Request, m *Manager, clientID string) {
+func HandleRequest(ctx context.Context, conn *ipc.ConnWriter, req ipc.Request, m *Manager, clientID string) {
 	switch req.Method {
 	case "dbus.call":
 		handleCall(conn, req, m)
@@ -58,7 +59,7 @@ func HandleRequest(conn *ipc.ConnWriter, req ipc.Request, m *Manager, clientID s
 	case "dbus.listNames":
 		handleListNames(conn, req, m)
 	case "dbus.subscribe":
-		handleSubscribe(conn, req, m, clientID)
+		handleSubscribe(ctx, conn, req, m, clientID)
 	case "dbus.unsubscribe":
 		handleUnsubscribe(conn, req, m)
 	default:
@@ -200,7 +201,7 @@ func handleListNames(conn *ipc.ConnWriter, req ipc.Request, m *Manager) {
 	conn.Respond(req.ID, result)
 }
 
-func handleSubscribe(conn *ipc.ConnWriter, req ipc.Request, m *Manager, clientID string) {
+func handleSubscribe(ctx context.Context, conn *ipc.ConnWriter, req ipc.Request, m *Manager, clientID string) {
 	if id := params.StringOpt(req.Params, "clientId", ""); id != "" {
 		clientID = id
 	}
@@ -216,12 +217,13 @@ func handleSubscribe(conn *ipc.ConnWriter, req ipc.Request, m *Manager, clientID
 	iface := params.StringOpt(req.Params, "interface", "")
 	member := params.StringOpt(req.Params, "member", "")
 
-	result, err := m.Subscribe(clientID, bus, sender, path, iface, member)
+	result, err := m.Subscribe(conn, clientID, bus, sender, path, iface, member)
 	if err != nil {
 		conn.RespondError(req.ID, err.Error())
 		return
 	}
 
+	m.ReleaseOnClose(ctx, conn)
 	conn.Respond(req.ID, result)
 }
 

core/internal/server/dbus/manager.go

@@ -1,12 +1,14 @@
 package dbus
 
 import (
+	"context"
 	"crypto/rand"
 	"encoding/hex"
 	"fmt"
 	"strings"
 
 	"github.com/duxlabs-dev/dux-go/dbusutil"
+	"github.com/duxlabs-dev/dux-go/ipc"
 	"github.com/duxlabs-dev/dux-shell/core/internal/log"
 	"github.com/godbus/dbus/v5"
 )
@@ -188,7 +190,7 @@ func (m *Manager) ListNames(bus string) (*ListNamesResult, error) {
 	return &ListNamesResult{Names: names}, nil
 }
 
-func (m *Manager) Subscribe(clientID, bus, sender, path, iface, member string) (*SubscribeResult, error) {
+func (m *Manager) Subscribe(owner *ipc.ConnWriter, clientID, bus, sender, path, iface, member string) (*SubscribeResult, error) {
 	conn, err := m.getConn(bus)
 	if err != nil {
 		return nil, err
@@ -223,6 +225,7 @@ func (m *Manager) Subscribe(clientID, bus, sender, path, iface, member string) (
 		Interface: iface,
 		Member:    member,
 		ClientID:  clientID,
+		Owner:     owner,
 	}
 	m.subscriptions.Store(subID, sub)
 
@@ -267,10 +270,23 @@ func (m *Manager) Unsubscribe(subID string) error {
 	return nil
 }
 
-func (m *Manager) UnsubscribeClient(clientID string) {
+// Match rules live as long as the request connection that added them, not the
+// subscribe socket, which the shell re-dials whenever its service list changes.
+func (m *Manager) ReleaseOnClose(ctx context.Context, owner *ipc.ConnWriter) {
+	if _, loaded := m.owners.LoadOrStore(owner, struct{}{}); loaded {
+		return
+	}
+	go func() {
+		<-ctx.Done()
+		m.owners.Delete(owner)
+		m.unsubscribeOwner(owner)
+	}()
+}
+
+func (m *Manager) unsubscribeOwner(owner *ipc.ConnWriter) {
 	var toDelete []string
 	m.subscriptions.Range(func(subID string, sub *signalSubscription) bool {
-		if sub.ClientID == clientID {
+		if sub.Owner == owner {
 			toDelete = append(toDelete, subID)
 		}
 		return true
@@ -281,20 +297,23 @@ func (m *Manager) UnsubscribeClient(clientID string) {
 	}
 }
 
+// A re-dialed subscribe socket can register before the old one tears down, so
+// the old release must only drop the channel it created. Channels are never
+// closed here because dispatchSignal may be mid-send on another goroutine.
 func (m *Manager) SubscribeSignals(clientID string) chan SignalEvent {
 	ch := make(chan SignalEvent, 64)
-	existing, loaded := m.signalSubscribers.LoadOrStore(clientID, ch)
-	if loaded {
-		return existing
-	}
+	m.subscribersMu.Lock()
+	m.signalSubscribers.Store(clientID, ch)
+	m.subscribersMu.Unlock()
 	return ch
 }
 
-func (m *Manager) UnsubscribeSignals(clientID string) {
-	if ch, ok := m.signalSubscribers.LoadAndDelete(clientID); ok {
-		close(ch)
+func (m *Manager) UnsubscribeSignals(clientID string, ch chan SignalEvent) {
+	m.subscribersMu.Lock()
+	defer m.subscribersMu.Unlock()
+	if current, ok := m.signalSubscribers.Load(clientID); ok && current == ch {
+		m.signalSubscribers.Delete(clientID)
 	}
-	m.UnsubscribeClient(clientID)
 }
 
 func (m *Manager) processSystemSignals() {

core/internal/server/dbus/types.go

@@ -1,6 +1,9 @@
 package dbus
 
 import (
+	"sync"
+
+	"github.com/duxlabs-dev/dux-go/ipc"
 	"github.com/duxlabs-dev/dux-go/syncmap"
 	"github.com/godbus/dbus/v5"
 )
@@ -12,7 +15,9 @@ type Manager struct {
 	sessionSignals chan *dbus.Signal
 
 	subscriptions     syncmap.Map[string, *signalSubscription]
+	subscribersMu     sync.Mutex
 	signalSubscribers syncmap.Map[string, chan SignalEvent]
+	owners            syncmap.Map[*ipc.ConnWriter, struct{}]
 }
 
 type signalSubscription struct {
@@ -22,6 +27,7 @@ type signalSubscription struct {
 	Interface string
 	Member    string
 	ClientID  string
+	Owner     *ipc.ConnWriter
 }
 
 type SignalEvent struct {

core/internal/server/router.go

@@ -203,13 +203,13 @@ func newRequestMux(readiness *systemd.Readiness) *ipc.Mux {
 		evdev.HandleRequest(conn, req, evdevManager)
 	}))
 
-	mux.HandlePrefix("dbus.", requestHandler(func(conn *ipc.ConnWriter, req ipc.Request) {
+	mux.HandlePrefix("dbus.", func(ctx context.Context, conn *ipc.ConnWriter, req ipc.Request, _ *ipc.Subscriber) {
 		if dbusManager == nil {
 			conn.RespondError(req.ID, "dbus manager not initialized")
 			return
 		}
-		serverDbus.HandleRequest(conn, req, dbusManager, dbusClientID)
-	}))
+		serverDbus.HandleRequest(ctx, conn, req, dbusManager, dbusClientID)
+	})
 
 	mux.HandlePrefix("clipboard.", requestHandler(func(conn *ipc.ConnWriter, req ipc.Request) {
 		if clipboardManager == nil {

core/internal/server/server.go

@@ -795,7 +795,7 @@ func handleSubscribe(ctx context.Context, conn *ipc.ConnWriter, req ipc.Request)
 		mgr := dbusManager
 		events := mgr.SubscribeSignals(dbusClient)
 		forwardSubscription(&wg, eventChan, stopChan, "dbus", events, func() {
-			mgr.UnsubscribeSignals(dbusClient)
+			mgr.UnsubscribeSignals(dbusClient, events)
 		}, nil)
 	}