Harden context menu websocket payloads

pkg/display/display.go

@@ -669,11 +669,25 @@ func (s *Service) handleWSMessage(msg WSMessage) core.Result {
 		))
 	case "contextmenu:add":
 		name, _ := msg.Data["name"].(string)
-		marshalResult := core.JSONMarshal(msg.Data["menu"])
+		menuValue, ok := msg.Data["menu"]
+		if !ok || menuValue == nil {
+			return core.Result{Value: coreerr.E("display.handleWSMessage", "missing required field \"menu\"", nil), OK: false}
+		}
+		marshalResult := core.JSONMarshal(menuValue)
+		if !marshalResult.OK {
+			if err, ok := marshalResult.Value.(error); ok {
+				return core.Result{Value: coreerr.E("display.handleWSMessage", "failed to marshal menu definition", err), OK: false}
+			}
+			return core.Result{Value: coreerr.E("display.handleWSMessage", "failed to marshal menu definition", nil), OK: false}
+		}
 		var menuDef contextmenu.ContextMenuDef
-		if marshalResult.OK {
-			menuJSON, _ := marshalResult.Value.([]byte)
-			core.JSONUnmarshal(menuJSON, &menuDef)
+		menuJSON, _ := marshalResult.Value.([]byte)
+		unmarshalResult := core.JSONUnmarshal(menuJSON, &menuDef)
+		if !unmarshalResult.OK {
+			if err, ok := unmarshalResult.Value.(error); ok {
+				return core.Result{Value: coreerr.E("display.handleWSMessage", "failed to unmarshal menu definition", err), OK: false}
+			}
+			return core.Result{Value: coreerr.E("display.handleWSMessage", "failed to unmarshal menu definition", nil), OK: false}
 		}
 		return c.Action("contextmenu.add").Run(ctx, core.NewOptions(
 			core.Option{Key: "task", Value: contextmenu.TaskAdd{Name: name, Menu: menuDef}},

pkg/display/display_contextmenu_test.go

@@ -0,0 +1,85 @@
+package display
+
+import (
+	"context"
+	"sync"
+	"testing"
+
+	core "dappco.re/go/core"
+	"forge.lthn.ai/core/gui/pkg/contextmenu"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+type wsContextMenuPlatform struct {
+	mu    sync.Mutex
+	menus map[string]contextmenu.ContextMenuDef
+}
+
+func newWSContextMenuPlatform() *wsContextMenuPlatform {
+	return &wsContextMenuPlatform{
+		menus: make(map[string]contextmenu.ContextMenuDef),
+	}
+}
+
+func (m *wsContextMenuPlatform) Add(name string, menu contextmenu.ContextMenuDef, _ func(string, string, string)) error {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	m.menus[name] = menu
+	return nil
+}
+
+func (m *wsContextMenuPlatform) Remove(name string) error {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	delete(m.menus, name)
+	return nil
+}
+
+func (m *wsContextMenuPlatform) Get(name string) (*contextmenu.ContextMenuDef, bool) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	menu, ok := m.menus[name]
+	if !ok {
+		return nil, false
+	}
+	return &menu, true
+}
+
+func (m *wsContextMenuPlatform) GetAll() map[string]contextmenu.ContextMenuDef {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	out := make(map[string]contextmenu.ContextMenuDef, len(m.menus))
+	for name, menu := range m.menus {
+		out[name] = menu
+	}
+	return out
+}
+
+func newDisplayWithContextMenu(t *testing.T, platform *wsContextMenuPlatform) (*Service, *core.Core) {
+	t.Helper()
+	c := core.New(
+		core.WithService(Register(nil)),
+		core.WithService(contextmenu.Register(platform)),
+		core.WithServiceLock(),
+	)
+	require.True(t, c.ServiceStartup(context.Background(), nil).OK)
+	return core.MustServiceFor[*Service](c, "display"), c
+}
+
+func TestDisplay_handleWSMessage_ContextMenuAdd_MissingMenu(t *testing.T) {
+	platform := newWSContextMenuPlatform()
+	svc, _ := newDisplayWithContextMenu(t, platform)
+
+	result := svc.handleWSMessage(WSMessage{
+		Action: "contextmenu:add",
+		Data: map[string]any{
+			"name": "menu",
+		},
+	})
+
+	require.False(t, result.OK)
+	assert.Contains(t, result.Value.(error).Error(), `missing required field "menu"`)
+	_, ok := platform.Get("menu")
+	assert.False(t, ok)
+}