From b077a41f83ac66b064cd90441d45685314e19a11 Mon Sep 17 00:00:00 2001
From: Patrick Bennett <patrick+git@bennettfamily.org>
Date: Fri, 21 Apr 2023 00:10:11 -0400
Subject: [PATCH] Fix race condition on status in Relay.Publish method and
 failure to send

A race-condition exists between setting of the (unprotected) status and the callback which sets the status upon receiving an OK.
The message is sent which can receive an OK in separate goroutine (setting status) prior to the status being set to 'sent.'
The OK can be received prior to the status being set.

This fix also sets the status to PublishStatusFailed if the WriteJSON call fails.
---
 relay.go      |  3 ++-
 relay_test.go | 22 +++++++++++++++++++++-
 2 files changed, 23 insertions(+), 2 deletions(-)

diff --git a/relay.go b/relay.go
index dacfd8b..2d84d47 100644
--- a/relay.go
+++ b/relay.go
@@ -305,10 +305,11 @@ func (r *Relay) Publish(ctx context.Context, event Event) (Status, error) {
 	// publish event
 	message := []any{"EVENT", event}
 	debugLog("{%s} sending %v\n", r.URL, message)
+	status = PublishStatusSent
 	if err := r.Connection.WriteJSON(message); err != nil {
+		status = PublishStatusFailed
 		return status, err
 	}
-	status = PublishStatusSent
 
 	sub := r.PrepareSubscription(ctx)
 	sub.SetLabel("publish-check")
diff --git a/relay_test.go b/relay_test.go
index 48b4e18..287505d 100644
--- a/relay_test.go
+++ b/relay_test.go
@@ -87,7 +87,27 @@ func TestPublishBlocked(t *testing.T) {
 	rl := mustRelayConnect(ws.URL)
 	status, _ := rl.Publish(context.Background(), textNote)
 	if status != PublishStatusFailed {
-		t.Errorf("published status is %d, not %d", status, PublishStatusSucceeded)
+		t.Errorf("published status is %d, not %d", status, PublishStatusFailed)
+	}
+}
+
+func TestPublishWriteFailed(t *testing.T) {
+	// test note to be sent over websocket
+	textNote := Event{Kind: 1, Content: "hello"}
+	textNote.ID = textNote.GetID()
+
+	// fake relay server
+	ws := newWebsocketServer(func(conn *websocket.Conn) {
+		// reject receive - force send error
+		conn.Close()
+	})
+	defer ws.Close()
+
+	// connect a client and send a text note
+	rl := mustRelayConnect(ws.URL)
+	status, _ := rl.Publish(context.Background(), textNote)
+	if status != PublishStatusFailed {
+		t.Errorf("published status is %d, not %d", status, PublishStatusFailed)
 	}
 }