From d6304b087aed7fb8c8d97bb5791f5fb9657249fe Mon Sep 17 00:00:00 2001
From: Leander Beernaert <leander.beernaert@proton.ch>
Date: Tue, 16 May 2023 16:51:20 +0200
Subject: [PATCH] fix(GODT-2627): Fix crash on closed channel

See `TestSendHasher_DualAddDoesNotCauseCrash`'s comment for more
details.
---
 internal/user/send_recorder.go      | 21 +++++++++++++++------
 internal/user/send_recorder_test.go | 26 +++++++++++++++++++++++++-
 2 files changed, 40 insertions(+), 7 deletions(-)

diff --git a/internal/user/send_recorder.go b/internal/user/send_recorder.go
index fae307e0..228bbfe6 100644
--- a/internal/user/send_recorder.go
+++ b/internal/user/send_recorder.go
@@ -46,10 +46,18 @@ func newSendRecorder(expiry time.Duration) *sendRecorder {
 }
 
 type sendEntry struct {
-	msgID  string
-	toList []string
-	exp    time.Time
-	waitCh chan struct{}
+	msgID        string
+	toList       []string
+	exp          time.Time
+	waitCh       chan struct{}
+	waitChClosed bool
+}
+
+func (s *sendEntry) closeWaitChannel() {
+	if !s.waitChClosed {
+		close(s.waitCh)
+		s.waitChClosed = true
+	}
 }
 
 // tryInsertWait tries to insert the given message into the send recorder.
@@ -142,6 +150,7 @@ func (h *sendRecorder) hasEntry(hash string) bool {
 	return false
 }
 
+// addMessageID should be called after a message has been successfully sent.
 func (h *sendRecorder) addMessageID(hash, msgID string) {
 	h.entriesLock.Lock()
 	defer h.entriesLock.Unlock()
@@ -153,7 +162,7 @@ func (h *sendRecorder) addMessageID(hash, msgID string) {
 		logrus.Warn("Cannot add message ID to send hash entry, it may have expired")
 	}
 
-	close(entry.waitCh)
+	entry.closeWaitChannel()
 }
 
 func (h *sendRecorder) removeOnFail(hash string) {
@@ -165,7 +174,7 @@ func (h *sendRecorder) removeOnFail(hash string) {
 		return
 	}
 
-	close(entry.waitCh)
+	entry.closeWaitChannel()
 
 	delete(h.entries, hash)
 }
diff --git a/internal/user/send_recorder_test.go b/internal/user/send_recorder_test.go
index 075a3b52..0aabe90b 100644
--- a/internal/user/send_recorder_test.go
+++ b/internal/user/send_recorder_test.go
@@ -194,6 +194,31 @@ func TestSendHasher_HasEntry_SendSuccess(t *testing.T) {
 	require.Equal(t, "abc", messageID)
 }
 
+func TestSendHasher_DualAddDoesNotCauseCrash(t *testing.T) {
+	// There may be a rare case where one 2 smtp connections attempt to send the same message, but if the first message
+	// is stuck long enough for it to expire, the second connection will remove it from the list and cause it to be
+	// inserted as a new entry. The two clients end up sending the message twice and calling the `addMessageID` x2,
+	// resulting in a crash.
+	h := newSendRecorder(sendEntryExpiry)
+
+	// Insert a message into the hasher.
+	hash, ok, err := testTryInsert(h, literal1, time.Now().Add(time.Second))
+	require.NoError(t, err)
+	require.True(t, ok)
+	require.NotEmpty(t, hash)
+
+	// Simulate successfully sending the message. We call this method twice as it possible for multiple SMTP connections
+	// to attempt to send the same message.
+	h.addMessageID(hash, "abc")
+	h.addMessageID(hash, "abc")
+
+	// The message was already sent; we should find it in the hasher.
+	messageID, ok, err := testHasEntry(h, literal1, time.Now().Add(time.Second))
+	require.NoError(t, err)
+	require.True(t, ok)
+	require.Equal(t, "abc", messageID)
+}
+
 func TestSendHasher_HasEntry_SendFail(t *testing.T) {
 	h := newSendRecorder(sendEntryExpiry)
 
@@ -264,7 +289,6 @@ Content-Disposition: attachment; filename="attname.txt"
 attachment
 --longrandomstring--
 `
-
 const literal2 = `From: Sender <sender@pm.me>
 To: Receiver <receiver@pm.me>
 Content-Type: multipart/mixed; boundary=longrandomstring