refactor(rest): deduplicate WebSocket monitor controllers (#622)

The black-box monitor tests showed socketv1_controller.go and
socketv2_controller.go were near-identical: same constants (identical
values), upgrader, read/write pumps, Close, and handler bodies, differing
only in the event payload shape and v1 dropping deletions.

Extract the shared machinery into socket_controller.go:

- msgListener[T] generic listener with toStored/toDeleted payload
  transforms; a nil toDeleted preserves the socketv1 behavior of ignoring
  deletions
- socketReadPump / socketWritePump free functions (WriteJSON is
  payload-type agnostic)
- serveMonitor shared handler body (upgrade, gauge tracking, listener
  registration, pump wiring)

The version files shrink to their listener constructors and the four
exported handlers, whose signatures are unchanged.  v1 payloads are now
shaped at enqueue time rather than write time; the JSON on the wire is
identical, pinned by the commit-1 tests, which pass unmodified.

Adds direct unit tests for the extracted listener logic: watches()
filtering, Receive/Delete filtering + payload shaping, nil-toDeleted
deletion dropping, and sequential Close idempotency.

pkg/rest/socket_controller.go

@@ -0,0 +1,225 @@
+package rest
+
+import (
+	"net/http"
+	"time"
+
+	"github.com/gorilla/websocket"
+	"github.com/inbucket/inbucket/v3/pkg/extension/event"
+	"github.com/inbucket/inbucket/v3/pkg/msghub"
+	"github.com/inbucket/inbucket/v3/pkg/rest/model"
+	"github.com/inbucket/inbucket/v3/pkg/server/web"
+	"github.com/inbucket/inbucket/v3/pkg/stringutil"
+	"github.com/rs/zerolog/log"
+)
+
+const (
+	// socketWriteWait is the time allowed to write a message to the peer.
+	socketWriteWait = 10 * time.Second
+
+	// socketPingPeriod is how often pings are sent to the peer. Must be less than
+	// socketPongWait.
+	socketPingPeriod = (socketPongWait * 9) / 10
+
+	// socketPongWait is the time allowed to read the next pong message from the peer.
+	socketPongWait = 60 * time.Second
+
+	// socketMaxMessageSize is the maximum message size allowed from the peer.
+	socketMaxMessageSize = 512
+
+	// socketChanLen is the event queue length of a monitor listener.
+	socketChanLen = 100
+)
+
+// upgrader holds the options for monitor WebSocket connection upgrades.
+var upgrader = websocket.Upgrader{
+	ReadBufferSize:  1024,
+	WriteBufferSize: 1024,
+}
+
+// msgListener handles messages from the msghub for a single monitor connection.  Stored and
+// deleted events are shaped into the API-specific payload by toStored/toDeleted; a nil
+// toDeleted drops deletions (socketv1 API behavior).
+type msgListener[T any] struct {
+	hub       *msghub.Hub
+	c         chan T
+	mailbox   string // Name of mailbox to monitor, "" == all mailboxes
+	toStored  func(msg event.MessageMetadata) T
+	toDeleted func(mailbox string, id string) T
+}
+
+// newMsgListener creates a listener and registers it with the hub.  The optional mailbox
+// parameter restricts events relayed to the client to that mailbox only.
+func newMsgListener[T any](
+	hub *msghub.Hub,
+	mailbox string,
+	toStored func(msg event.MessageMetadata) T,
+	toDeleted func(mailbox string, id string) T,
+) *msgListener[T] {
+	ml := &msgListener[T]{
+		hub:       hub,
+		c:         make(chan T, socketChanLen),
+		mailbox:   mailbox,
+		toStored:  toStored,
+		toDeleted: toDeleted,
+	}
+	hub.AddListener(ml)
+	return ml
+}
+
+// watches reports whether events for the given mailbox should be relayed to the client.
+func (ml *msgListener[T]) watches(mailbox string) bool {
+	return ml.mailbox == "" || ml.mailbox == mailbox
+}
+
+// Receive handles an incoming message.
+func (ml *msgListener[T]) Receive(msg event.MessageMetadata) error {
+	if !ml.watches(msg.Mailbox) {
+		return nil
+	}
+	ml.c <- ml.toStored(msg)
+	return nil
+}
+
+// Delete handles a deleted message.
+func (ml *msgListener[T]) Delete(mailbox string, id string) error {
+	if ml.toDeleted == nil || !ml.watches(mailbox) {
+		return nil
+	}
+	ml.c <- ml.toDeleted(mailbox, id)
+	return nil
+}
+
+// Close removes the listener registration and closes the event queue.
+func (ml *msgListener[T]) Close() {
+	select {
+	case <-ml.c:
+		// Already closed
+	default:
+		ml.hub.RemoveListener(ml)
+		close(ml.c)
+	}
+}
+
+// socketReadPump makes sure the websocket client is still connected, discarding any messages
+// the client sends.  It returns once the connection fails or the peer closes it; closeFn is
+// then invoked to tear the listener down.
+func socketReadPump(conn *websocket.Conn, closeFn func()) {
+	slog := log.With().Str("module", "rest").Str("proto", "WebSocket").
+		Str("remote", conn.RemoteAddr().String()).Logger()
+
+	defer closeFn()
+
+	conn.SetReadLimit(socketMaxMessageSize)
+	if err := conn.SetReadDeadline(time.Now().Add(socketPongWait)); err != nil {
+		slog.Warn().Err(err).Msg("Failed to setup read deadline")
+	}
+	conn.SetPongHandler(func(string) error {
+		slog.Debug().Msg("Got pong")
+		if err := conn.SetReadDeadline(time.Now().Add(socketPongWait)); err != nil {
+			slog.Warn().Err(err).Msg("Failed to set read deadline in pong")
+		}
+		return nil
+	})
+
+	for {
+		if _, _, err := conn.ReadMessage(); err != nil {
+			if websocket.IsUnexpectedCloseError(
+				err,
+				websocket.CloseNormalClosure,
+				websocket.CloseGoingAway,
+				websocket.CloseNoStatusReceived,
+			) {
+				// Unexpected close code
+				slog.Warn().Err(err).Msg("Socket error")
+			} else {
+				slog.Debug().Msg("Closing socket")
+			}
+			break
+		}
+	}
+}
+
+// socketWritePump relays events from c to the peer as JSON messages until c is closed, sending
+// pings while idle to keep the connection alive.  closeFn is invoked on exit.
+func socketWritePump[T any](c chan T, conn *websocket.Conn, closeFn func()) {
+	slog := log.With().Str("module", "rest").Str("proto", "WebSocket").
+		Str("remote", conn.RemoteAddr().String()).Logger()
+
+	ticker := time.NewTicker(socketPingPeriod)
+	defer func() {
+		ticker.Stop()
+		closeFn()
+	}()
+
+	// Handle messages from hub until listener is closed
+	for {
+		select {
+		case event, ok := <-c:
+			if err := conn.SetWriteDeadline(time.Now().Add(socketWriteWait)); err != nil {
+				slog.Warn().Err(err).Msg("Failed to set write deadline for msg")
+			}
+			if !ok {
+				// Listener closed, exit
+				_ = conn.WriteMessage(websocket.CloseMessage, []byte{})
+				return
+			}
+			if conn.WriteJSON(event) != nil {
+				// Write failed
+				return
+			}
+		case <-ticker.C:
+			// Send ping
+			if err := conn.SetWriteDeadline(time.Now().Add(socketWriteWait)); err != nil {
+				slog.Warn().Err(err).Msg("Failed to set write deadline for ping")
+			}
+			if conn.WriteMessage(websocket.PingMessage, []byte{}) != nil {
+				// Write error
+				return
+			}
+			slog.Debug().Msg("Sent ping")
+		}
+	}
+}
+
+// serveMonitor upgrades the connection to a WebSocket, registers a new listener with the hub,
+// and relays events to the client until it disconnects.
+func serveMonitor[T any](
+	w http.ResponseWriter,
+	req *http.Request,
+	hub *msghub.Hub,
+	mailbox string,
+	newListener func(hub *msghub.Hub, mailbox string) *msgListener[T],
+) error {
+	// Upgrade to Websocket.
+	conn, err := upgrader.Upgrade(w, req, nil)
+	if err != nil {
+		return err
+	}
+	web.ExpWebSocketConnectsCurrent.Add(1)
+	defer func() {
+		_ = conn.Close()
+		web.ExpWebSocketConnectsCurrent.Add(-1)
+	}()
+	log.Debug().Str("module", "rest").Str("proto", "WebSocket").
+		Str("remote", conn.RemoteAddr().String()).Msg("Upgraded to WebSocket")
+	// Create, register listener; then interact with conn.
+	ml := newListener(hub, mailbox)
+	go socketWritePump(ml.c, conn, ml.Close)
+	socketReadPump(conn, ml.Close)
+	return nil
+}
+
+// metadataToHeader converts hub event metadata into a JSON message header.
+func metadataToHeader(msg *event.MessageMetadata) *model.JSONMessageHeaderV1 {
+	return &model.JSONMessageHeaderV1{
+		Mailbox:     msg.Mailbox,
+		ID:          msg.ID,
+		From:        stringutil.StringAddress(msg.From),
+		To:          stringutil.StringAddressList(msg.To),
+		Subject:     msg.Subject,
+		Date:        msg.Date,
+		PosixMillis: msg.Date.UnixNano() / 1000000,
+		Size:        msg.Size,
+	}
+}

pkg/rest/socket_controller_test.go

@@ -337,3 +337,74 @@ func TestMonitorUpgradeFailedV2(t *testing.T) {
 	require.NoError(t, err)
 	assert.Equal(t, 400, w.Code)
 }
+
+// assertNoMoreEvents fails the test if any further events are queued on the channel.
+func assertNoMoreEvents[T any](t *testing.T, c chan T) {
+	t.Helper()
+	select {
+	case ev := <-c:
+		t.Fatalf("unexpected event: %v", ev)
+	default:
+	}
+}
+
+func TestMsgListenerWatches(t *testing.T) {
+	testCases := []struct {
+		name    string
+		mailbox string // listener watch; "" watches all mailboxes
+		event   string // mailbox of the incoming event
+		want    bool
+	}{
+		{name: "all watcher matches any mailbox", mailbox: "", event: "anything", want: true},
+		{name: "watched mailbox matches", mailbox: "watched", event: "watched", want: true},
+		{name: "other mailbox filtered", mailbox: "watched", event: "other", want: false},
+	}
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			ml := &msgListener[string]{mailbox: tc.mailbox}
+			assert.Equal(t, tc.want, ml.watches(tc.event))
+		})
+	}
+}
+
+func TestMsgListenerReceiveFiltersAndShapes(t *testing.T) {
+	hub := startMsgHub(t)
+	ml := newMsgListener(hub, "watched",
+		func(msg event.MessageMetadata) string { return "stored:" + msg.Mailbox + "/" + msg.ID },
+		func(mailbox, id string) string { return "deleted:" + mailbox + "/" + id })
+
+	require.NoError(t, ml.Receive(testMetadata("other", "0001")))
+	require.NoError(t, ml.Receive(testMetadata("watched", "0002")))
+	require.NoError(t, ml.Delete("other", "0003"))
+	require.NoError(t, ml.Delete("watched", "0004"))
+
+	// Events for other mailboxes are filtered; payloads are shaped by the transforms.
+	assert.Equal(t, "stored:watched/0002", <-ml.c)
+	assert.Equal(t, "deleted:watched/0004", <-ml.c)
+	assertNoMoreEvents(t, ml.c)
+}
+
+func TestMsgListenerNilDeletedDropsDeletes(t *testing.T) {
+	hub := startMsgHub(t)
+	ml := newMsgListener(hub, "watched",
+		func(msg event.MessageMetadata) string { return "stored:" + msg.ID },
+		nil) // socketv1 API ignores deletions
+
+	require.NoError(t, ml.Receive(testMetadata("watched", "0001")))
+	require.NoError(t, ml.Delete("watched", "0001"))
+
+	assert.Equal(t, "stored:0001", <-ml.c)
+	assertNoMoreEvents(t, ml.c)
+}
+
+func TestMsgListenerCloseIdempotent(t *testing.T) {
+	hub := startMsgHub(t)
+	ml := newMsgListener(hub, "",
+		func(msg event.MessageMetadata) string { return "stored" }, nil)
+
+	ml.Close()
+	ml.Close() // second close must not double-close the event channel
+
+	_, ok := <-ml.c
+	assert.False(t, ok, "event channel should be closed")
+}

pkg/rest/socketv1_controller.go

@@ -2,220 +2,37 @@ package rest
 
 import (
 	"net/http"
-	"time"
 
-	"github.com/gorilla/websocket"
 	"github.com/inbucket/inbucket/v3/pkg/extension/event"
 	"github.com/inbucket/inbucket/v3/pkg/msghub"
 	"github.com/inbucket/inbucket/v3/pkg/rest/model"
 	"github.com/inbucket/inbucket/v3/pkg/server/web"
-	"github.com/inbucket/inbucket/v3/pkg/stringutil"
-	"github.com/rs/zerolog/log"
 )
 
-const (
+// newMsgListenerV1 creates a v1 API listener which relays message headers only; deletions are
-	// Time allowed to write a message to the peer.
+// ignored by the socketv1 API.
-	writeWaitV1 = 10 * time.Second
+func newMsgListenerV1(hub *msghub.Hub, mailbox string) *msgListener[*model.JSONMessageHeaderV1] {
-
+	return newMsgListener(hub, mailbox,
-	// Send pings to peer with this period. Must be less than pongWait.
+		func(msg event.MessageMetadata) *model.JSONMessageHeaderV1 {
-	pingPeriodV1 = (pongWaitV1 * 9) / 10
+			return metadataToHeader(&msg)
-
+		},
-	// Time allowed to read the next pong message from the peer.
+		nil)
-	pongWaitV1 = 60 * time.Second
-
-	// Maximum message size allowed from peer.
-	maxMessageSizeV1 = 512
-)
-
-// options for gorilla connection upgrader
-var upgraderV1 = websocket.Upgrader{
-	ReadBufferSize:  1024,
-	WriteBufferSize: 1024,
-}
-
-// msgListenerV1 handles messages from the msghub
-type msgListenerV1 struct {
-	hub     *msghub.Hub                // Global message hub
-	c       chan event.MessageMetadata // Queue of messages from Receive()
-	mailbox string                     // Name of mailbox to monitor, "" == all mailboxes
-}
-
-// newMsgListenerV1 creates a listener and registers it.  Optional mailbox parameter will restrict
-// messages sent to WebSocket to that mailbox only.
-func newMsgListenerV1(hub *msghub.Hub, mailbox string) *msgListenerV1 {
-	ml := &msgListenerV1{
-		hub:     hub,
-		c:       make(chan event.MessageMetadata, 100),
-		mailbox: mailbox,
-	}
-	hub.AddListener(ml)
-	return ml
-}
-
-// Receive handles an incoming message.
-func (ml *msgListenerV1) Receive(msg event.MessageMetadata) error {
-	if ml.mailbox != "" && ml.mailbox != msg.Mailbox {
-		// Did not match the watched mailbox name.
-		return nil
-	}
-	ml.c <- msg
-	return nil
-}
-
-// Delete handles a deleted message.
-func (ml *msgListenerV1) Delete(mailbox string, id string) error {
-	// Deletes are ignored in socketv1 API.
-	return nil
-}
-
-// WSReader makes sure the websocket client is still connected, discards any messages from client
-func (ml *msgListenerV1) WSReader(conn *websocket.Conn) {
-	slog := log.With().Str("module", "rest").Str("proto", "WebSocket").
-		Str("remote", conn.RemoteAddr().String()).Logger()
-
-	defer ml.Close()
-
-	conn.SetReadLimit(maxMessageSizeV1)
-	if err := conn.SetReadDeadline(time.Now().Add(pongWaitV1)); err != nil {
-		slog.Warn().Err(err).Msg("Failed to setup read deadline")
-	}
-	conn.SetPongHandler(func(string) error {
-		slog.Debug().Msg("Got pong")
-		if err := conn.SetReadDeadline(time.Now().Add(pongWaitV1)); err != nil {
-			slog.Warn().Err(err).Msg("Failed to set read deadline in pong")
-		}
-		return nil
-	})
-
-	for {
-		if _, _, err := conn.ReadMessage(); err != nil {
-			if websocket.IsUnexpectedCloseError(
-				err,
-				websocket.CloseNormalClosure,
-				websocket.CloseGoingAway,
-				websocket.CloseNoStatusReceived,
-			) {
-				// Unexpected close code
-				slog.Warn().Err(err).Msg("Socket error")
-			} else {
-				slog.Debug().Msg("Closing socket")
-			}
-			break
-		}
-	}
-}
-
-// WSWriter makes sure the websocket client is still connected
-func (ml *msgListenerV1) WSWriter(conn *websocket.Conn) {
-	slog := log.With().Str("module", "rest").Str("proto", "WebSocket").
-		Str("remote", conn.RemoteAddr().String()).Logger()
-
-	ticker := time.NewTicker(pingPeriodV1)
-	defer func() {
-		ticker.Stop()
-		ml.Close()
-	}()
-
-	// Handle messages from hub until msgListener is closed
-	for {
-		select {
-		case msg, ok := <-ml.c:
-			if err := conn.SetWriteDeadline(time.Now().Add(writeWaitV1)); err != nil {
-				slog.Warn().Err(err).Msg("Failed to set write deadline for msg")
-			}
-			if !ok {
-				// msgListener closed, exit
-				_ = conn.WriteMessage(websocket.CloseMessage, []byte{})
-				return
-			}
-			if conn.WriteJSON(metadataToHeader(&msg)) != nil {
-				// Write failed
-				return
-			}
-		case <-ticker.C:
-			// Send ping
-			if err := conn.SetWriteDeadline(time.Now().Add(writeWaitV1)); err != nil {
-				slog.Warn().Err(err).Msg("Failed to set write deadline for ping")
-			}
-			if conn.WriteMessage(websocket.PingMessage, []byte{}) != nil {
-				// Write error
-				return
-			}
-			slog.Debug().Msg("Sent ping")
-		}
-	}
-}
-
-// Close removes the listener registration
-func (ml *msgListenerV1) Close() {
-	select {
-	case <-ml.c:
-		// Already closed
-	default:
-		ml.hub.RemoveListener(ml)
-		close(ml.c)
-	}
 }
 
 // MonitorAllMessagesV1 is a web handler which upgrades the connection to a websocket and notifies
 // the client of all messages received.
 func MonitorAllMessagesV1(
-	w http.ResponseWriter, req *http.Request, ctx *web.Context) (err error) {
+	w http.ResponseWriter, req *http.Request, ctx *web.Context) error {
-	// Upgrade to Websocket.
+	return serveMonitor(w, req, ctx.MsgHub, "", newMsgListenerV1)
-	conn, err := upgraderV1.Upgrade(w, req, nil)
-	if err != nil {
-		return err
-	}
-	web.ExpWebSocketConnectsCurrent.Add(1)
-	defer func() {
-		_ = conn.Close()
-		web.ExpWebSocketConnectsCurrent.Add(-1)
-	}()
-	log.Debug().Str("module", "rest").Str("proto", "WebSocket").
-		Str("remote", conn.RemoteAddr().String()).Msg("Upgraded to WebSocket")
-	// Create, register listener; then interact with conn.
-	ml := newMsgListenerV1(ctx.MsgHub, "")
-	go ml.WSWriter(conn)
-	ml.WSReader(conn)
-	return nil
 }
 
 // MonitorMailboxMessagesV1 is a web handler which upgrades the connection to a websocket and
 // notifies the client of messages received by a particular mailbox.
 func MonitorMailboxMessagesV1(
-	w http.ResponseWriter, req *http.Request, ctx *web.Context) (err error) {
+	w http.ResponseWriter, req *http.Request, ctx *web.Context) error {
 	name, err := ctx.Manager.MailboxForAddress(ctx.Vars["name"])
 	if err != nil {
 		return err
 	}
-	// Upgrade to Websocket.
+	return serveMonitor(w, req, ctx.MsgHub, name, newMsgListenerV1)
-	conn, err := upgraderV1.Upgrade(w, req, nil)
-	if err != nil {
-		return err
-	}
-	web.ExpWebSocketConnectsCurrent.Add(1)
-	defer func() {
-		_ = conn.Close()
-		web.ExpWebSocketConnectsCurrent.Add(-1)
-	}()
-	log.Debug().Str("module", "rest").Str("proto", "WebSocket").
-		Str("remote", conn.RemoteAddr().String()).Msg("Upgraded to WebSocket")
-	// Create, register listener; then interact with conn.
-	ml := newMsgListenerV1(ctx.MsgHub, name)
-	go ml.WSWriter(conn)
-	ml.WSReader(conn)
-	return nil
-}
-
-func metadataToHeader(msg *event.MessageMetadata) *model.JSONMessageHeaderV1 {
-	return &model.JSONMessageHeaderV1{
-		Mailbox:     msg.Mailbox,
-		ID:          msg.ID,
-		From:        stringutil.StringAddress(msg.From),
-		To:          stringutil.StringAddressList(msg.To),
-		Subject:     msg.Subject,
-		Date:        msg.Date,
-		PosixMillis: msg.Date.UnixNano() / 1000000,
-		Size:        msg.Size,
-	}
 }

pkg/rest/socketv2_controller.go

@@ -2,224 +2,48 @@ package rest
 
 import (
 	"net/http"
-	"time"
 
-	"github.com/gorilla/websocket"
 	"github.com/inbucket/inbucket/v3/pkg/extension/event"
 	"github.com/inbucket/inbucket/v3/pkg/msghub"
 	"github.com/inbucket/inbucket/v3/pkg/rest/model"
 	"github.com/inbucket/inbucket/v3/pkg/server/web"
-	"github.com/rs/zerolog/log"
 )
 
-const (
+// newMsgListenerV2 creates a v2 API listener which relays message-stored and message-deleted
-	// Time allowed to write a message to the peer.
+// events.
-	writeWaitV2 = 10 * time.Second
+func newMsgListenerV2(hub *msghub.Hub, mailbox string) *msgListener[*model.JSONMonitorEventV2] {
-
+	return newMsgListener(hub, mailbox,
-	// Send pings to peer with this period. Must be less than pongWait.
+		func(msg event.MessageMetadata) *model.JSONMonitorEventV2 {
-	pingPeriodV2 = (pongWaitV2 * 9) / 10
+			return &model.JSONMonitorEventV2{
-
+				Variant: "message-stored",
-	// Time allowed to read the next pong message from the peer.
+				Header:  metadataToHeader(&msg),
-	pongWaitV2 = 60 * time.Second
+			}
-
-	// Maximum message size allowed from peer.
-	maxMessageSizeV2 = 512
-)
-
-// options for gorilla connection upgrader
-var upgraderV2 = websocket.Upgrader{
-	ReadBufferSize:  1024,
-	WriteBufferSize: 1024,
-}
-
-// msgListenerV2 handles messages from the msghub
-type msgListenerV2 struct {
-	hub     *msghub.Hub                    // Global message hub.
-	c       chan *model.JSONMonitorEventV2 // Queue of incoming events.
-	mailbox string                         // Name of mailbox to monitor, "" == all mailboxes.
-}
-
-// newMsgListenerV2 creates a listener and registers it.  Optional mailbox parameter will restrict
-// messages sent to WebSocket to that mailbox only.
-func newMsgListenerV2(hub *msghub.Hub, mailbox string) *msgListenerV2 {
-	ml := &msgListenerV2{
-		hub:     hub,
-		c:       make(chan *model.JSONMonitorEventV2, 100),
-		mailbox: mailbox,
-	}
-	hub.AddListener(ml)
-	return ml
-}
-
-// Receive handles an incoming message.
-func (ml *msgListenerV2) Receive(msg event.MessageMetadata) error {
-	if ml.mailbox != "" && ml.mailbox != msg.Mailbox {
-		// Did not match the watched mailbox name.
-		return nil
-	}
-
-	// Enqueue for websocket.
-	ml.c <- &model.JSONMonitorEventV2{
-		Variant: "message-stored",
-		Header:  metadataToHeader(&msg),
-	}
-
-	return nil
-}
-
-// Delete handles a deleted message.
-func (ml *msgListenerV2) Delete(mailbox string, id string) error {
-	if ml.mailbox != "" && ml.mailbox != mailbox {
-		// Did not match watched mailbox name.
-		return nil
-	}
-
-	// Enqueue for websocket.
-	ml.c <- &model.JSONMonitorEventV2{
-		Variant: "message-deleted",
-		Identifier: &model.JSONMessageIDV2{
-			Mailbox: mailbox,
-			ID:      id,
 		},
-	}
+		func(mailbox string, id string) *model.JSONMonitorEventV2 {
-
+			return &model.JSONMonitorEventV2{
-	return nil
+				Variant: "message-deleted",
-}
+				Identifier: &model.JSONMessageIDV2{
-
+					Mailbox: mailbox,
-// WSReader makes sure the websocket client is still connected, discards any messages from client
+					ID:      id,
-func (ml *msgListenerV2) WSReader(conn *websocket.Conn) {
+				},
-	slog := log.With().Str("module", "rest").Str("proto", "WebSocket").
-		Str("remote", conn.RemoteAddr().String()).Logger()
-	defer ml.Close()
-
-	conn.SetReadLimit(maxMessageSizeV2)
-	if err := conn.SetReadDeadline(time.Now().Add(pongWaitV2)); err != nil {
-		slog.Warn().Err(err).Msg("Failed to setup read deadline")
-	}
-	conn.SetPongHandler(func(string) error {
-		slog.Debug().Msg("Got pong")
-		if err := conn.SetReadDeadline(time.Now().Add(pongWaitV2)); err != nil {
-			slog.Warn().Err(err).Msg("Failed to set read deadline in pong")
-		}
-		return nil
-	})
-
-	for {
-		if _, _, err := conn.ReadMessage(); err != nil {
-			if websocket.IsUnexpectedCloseError(
-				err,
-				websocket.CloseNormalClosure,
-				websocket.CloseGoingAway,
-				websocket.CloseNoStatusReceived,
-			) {
-				// Unexpected close code
-				slog.Warn().Err(err).Msg("Socket error")
-			} else {
-				slog.Debug().Msg("Closing socket")
 			}
-			break
+		})
-		}
-	}
-}
-
-// WSWriter makes sure the websocket client is still connected
-func (ml *msgListenerV2) WSWriter(conn *websocket.Conn) {
-	slog := log.With().Str("module", "rest").Str("proto", "WebSocket").
-		Str("remote", conn.RemoteAddr().String()).Logger()
-
-	ticker := time.NewTicker(pingPeriodV2)
-	defer func() {
-		ticker.Stop()
-		ml.Close()
-	}()
-
-	// Handle messages from hub until msgListener is closed
-	for {
-		select {
-		case event, ok := <-ml.c:
-			if err := conn.SetWriteDeadline(time.Now().Add(writeWaitV2)); err != nil {
-				slog.Warn().Err(err).Msg("Failed to set write deadline for msg")
-			}
-			if !ok {
-				// msgListener closed, exit
-				_ = conn.WriteMessage(websocket.CloseMessage, []byte{})
-				return
-			}
-			if conn.WriteJSON(event) != nil {
-				// Write failed
-				return
-			}
-		case <-ticker.C:
-			// Send ping
-			if err := conn.SetWriteDeadline(time.Now().Add(writeWaitV2)); err != nil {
-				slog.Warn().Err(err).Msg("Failed to set write deadline for ping")
-			}
-			if conn.WriteMessage(websocket.PingMessage, []byte{}) != nil {
-				// Write error
-				return
-			}
-			slog.Debug().Msg("Sent ping")
-		}
-	}
-}
-
-// Close removes the listener registration
-func (ml *msgListenerV2) Close() {
-	select {
-	case <-ml.c:
-		// Already closed
-	default:
-		ml.hub.RemoveListener(ml)
-		close(ml.c)
-	}
 }
 
 // MonitorAllMessagesV2 is a web handler which upgrades the connection to a websocket and notifies
 // the client of all messages received.
 func MonitorAllMessagesV2(
-	w http.ResponseWriter, req *http.Request, ctx *web.Context) (err error) {
+	w http.ResponseWriter, req *http.Request, ctx *web.Context) error {
-	// Upgrade to Websocket.
+	return serveMonitor(w, req, ctx.MsgHub, "", newMsgListenerV2)
-	conn, err := upgraderV2.Upgrade(w, req, nil)
-	if err != nil {
-		return err
-	}
-	web.ExpWebSocketConnectsCurrent.Add(1)
-	defer func() {
-		_ = conn.Close()
-		web.ExpWebSocketConnectsCurrent.Add(-1)
-	}()
-	log.Debug().Str("module", "rest").Str("proto", "WebSocket").
-		Str("remote", conn.RemoteAddr().String()).Msg("Upgraded to WebSocket")
-	// Create, register listener; then interact with conn.
-	ml := newMsgListenerV2(ctx.MsgHub, "")
-	go ml.WSWriter(conn)
-	ml.WSReader(conn)
-	return nil
 }
 
 // MonitorMailboxMessagesV2 is a web handler which upgrades the connection to a websocket and
 // notifies the client of messages received by a particular mailbox.
 func MonitorMailboxMessagesV2(
-	w http.ResponseWriter, req *http.Request, ctx *web.Context) (err error) {
+	w http.ResponseWriter, req *http.Request, ctx *web.Context) error {
 	name, err := ctx.Manager.MailboxForAddress(ctx.Vars["name"])
 	if err != nil {
 		return err
 	}
-	// Upgrade to Websocket.
+	return serveMonitor(w, req, ctx.MsgHub, name, newMsgListenerV2)
-	conn, err := upgraderV2.Upgrade(w, req, nil)
-	if err != nil {
-		return err
-	}
-	web.ExpWebSocketConnectsCurrent.Add(1)
-	defer func() {
-		_ = conn.Close()
-		web.ExpWebSocketConnectsCurrent.Add(-1)
-	}()
-	log.Debug().Str("module", "rest").Str("proto", "WebSocket").
-		Str("remote", conn.RemoteAddr().String()).Msg("Upgraded to WebSocket")
-	// Create, register listener; then interact with conn.
-	ml := newMsgListenerV2(ctx.MsgHub, name)
-	go ml.WSWriter(conn)
-	ml.WSReader(conn)
-	return nil
 }