From 90d7bfd9a470a1f401297387bed24753fedfeb48 Mon Sep 17 00:00:00 2001
From: Snider <snider@host.uk.com>
Date: Fri, 17 Apr 2026 20:38:57 +0100
Subject: [PATCH] Harden websocket and marketplace failure paths

---
 pkg/display/events.go               | 14 ++++++++++++++
 pkg/display/events_test.go          | 13 +++++++++++++
 pkg/marketplace/marketplace.go      | 12 +++++++++++-
 pkg/marketplace/marketplace_test.go | 27 +++++++++++++++++++++++++++
 4 files changed, 65 insertions(+), 1 deletion(-)

diff --git a/pkg/display/events.go b/pkg/display/events.go
index 6246fa84..fa77e0cd 100644
--- a/pkg/display/events.go
+++ b/pkg/display/events.go
@@ -266,6 +266,15 @@ func (em *WSEventManager) HandleWebSocket(w http.ResponseWriter, r *http.Request
 		}
 		return
 	}
+	em.mu.RLock()
+	closed := em.closed
+	em.mu.RUnlock()
+	if closed {
+		if w != nil {
+			http.Error(w, http.StatusText(http.StatusServiceUnavailable), http.StatusServiceUnavailable)
+		}
+		return
+	}
 	if !trustedWebSocketOrigin(r) {
 		http.Error(w, http.StatusText(http.StatusForbidden), http.StatusForbidden)
 		return
@@ -277,6 +286,11 @@ func (em *WSEventManager) HandleWebSocket(w http.ResponseWriter, r *http.Request
 	}
 
 	em.mu.Lock()
+	if em.closed {
+		em.mu.Unlock()
+		_ = conn.Close()
+		return
+	}
 	em.clients[conn] = &clientState{
 		subscriptions: make(map[string]*Subscription),
 	}
diff --git a/pkg/display/events_test.go b/pkg/display/events_test.go
index 0d241ae9..b762482e 100644
--- a/pkg/display/events_test.go
+++ b/pkg/display/events_test.go
@@ -180,6 +180,19 @@ func TestWSEventManager_HandleWebSocket_NilReceiverFailsClosed(t *testing.T) {
 	assert.Equal(t, http.StatusServiceUnavailable, recorder.Code)
 }
 
+func TestWSEventManager_HandleWebSocket_RejectsAfterClose(t *testing.T) {
+	em := NewWSEventManager()
+	em.Close()
+
+	req := httptest.NewRequest(http.MethodGet, "http://127.0.0.1/events", nil)
+	req.RemoteAddr = "127.0.0.1:12345"
+	recorder := httptest.NewRecorder()
+
+	em.HandleWebSocket(recorder, req)
+
+	assert.Equal(t, http.StatusServiceUnavailable, recorder.Code)
+}
+
 func TestEvents_trustedWebSocketOrigin_Good(t *testing.T) {
 	tests := []struct {
 		name string
diff --git a/pkg/marketplace/marketplace.go b/pkg/marketplace/marketplace.go
index 1ce0559d..76fa316e 100644
--- a/pkg/marketplace/marketplace.go
+++ b/pkg/marketplace/marketplace.go
@@ -153,6 +153,13 @@ func (i Installer) Install(ctx context.Context, manifest Manifest) (string, erro
 	if err != nil {
 		return "", err
 	}
+	cleanupTarget := true
+	defer func() {
+		if cleanupTarget {
+			_ = os.RemoveAll(targetDir)
+		}
+	}()
+
 	rel, err := filepath.Rel(rootResolved, targetAbs)
 	if err != nil {
 		return "", err
@@ -160,7 +167,9 @@ func (i Installer) Install(ctx context.Context, manifest Manifest) (string, erro
 	if rel == ".." || strings.HasPrefix(rel, ".."+string(filepath.Separator)) {
 		return "", errors.New("install path escapes install dir")
 	}
-	_ = os.RemoveAll(targetDir)
+	if err := os.RemoveAll(targetDir); err != nil {
+		return "", err
+	}
 	args := []string{"clone", "--depth", "1"}
 	if manifest.Ref != "" {
 		args = append(args, "--branch", manifest.Ref)
@@ -177,6 +186,7 @@ func (i Installer) Install(ctx context.Context, manifest Manifest) (string, erro
 	if err := writeInstalledManifest(targetDir, manifest); err != nil {
 		return "", err
 	}
+	cleanupTarget = false
 	return targetDir, nil
 }
 
diff --git a/pkg/marketplace/marketplace_test.go b/pkg/marketplace/marketplace_test.go
index 9c93e621..42df732f 100644
--- a/pkg/marketplace/marketplace_test.go
+++ b/pkg/marketplace/marketplace_test.go
@@ -266,6 +266,33 @@ func TestMarketplace_Install_Ugly(t *testing.T) {
 	assert.NotContains(t, err.Error(), "token:")
 }
 
+func TestMarketplace_Install_CleansUpOnCloneFailure(t *testing.T) {
+	scriptDir := t.TempDir()
+	scriptPath := filepath.Join(scriptDir, "git")
+	targetRoot := t.TempDir()
+	script := "#!/bin/sh\nlast=''\nfor arg in \"$@\"; do last=\"$arg\"; done\nmkdir -p \"$last\"\ntouch \"$last/partial\"\nexit 1\n"
+	require.NoError(t, os.WriteFile(scriptPath, []byte(script), 0o755))
+
+	installer := Installer{
+		GitBinary:  scriptPath,
+		InstallDir: targetRoot,
+	}
+
+	manifest := signedManifest(t, Manifest{
+		Name:       "core-ui",
+		Version:    "1.2.3",
+		Repository: "https://example.com/core-ui.git",
+		Ref:        "main",
+	})
+	_, err := installer.Install(context.Background(), manifest)
+	require.Error(t, err)
+
+	targetDir := filepath.Join(targetRoot, "core-ui")
+	_, statErr := os.Stat(targetDir)
+	assert.Error(t, statErr)
+	assert.True(t, os.IsNotExist(statErr))
+}
+
 func TestMarketplace_Verify_Good(t *testing.T) {
 	manifest := signedManifest(t, Manifest{
 		Name:       "core-ui",