From cfcc3793c5a5d96688518ad4425c32321dd32842 Mon Sep 17 00:00:00 2001
From: binwiederhier <pheckel@datto.com>
Date: Fri, 10 Feb 2023 21:44:12 -0500
Subject: [PATCH] Fix 404 race when uploading attachments

---
 server/server.go      | 12 ++++++++++--
 server/server_test.go | 19 +++++++++++++++++++
 2 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/server/server.go b/server/server.go
index a6d5cf72..97c79494 100644
--- a/server/server.go
+++ b/server/server.go
@@ -41,7 +41,6 @@ import (
   - tokens
 - HIGH Self-review
 - MEDIUM: Test for expiring messages after reservation removal
-- MEDIUM: uploading attachments leads to 404 -- race
 - MEDIUM: Test new token endpoints & never-expiring token
 - LOW: UI: Flickering upgrade banner when logging in
 - LOW: Menu item -> popup click should not open page
@@ -563,7 +562,16 @@ func (s *Server) handleFile(w http.ResponseWriter, r *http.Request, v *visitor)
 	//   - and also uses the higher bandwidth limits of a paying user
 	m, err := s.messageCache.Message(messageID)
 	if err == errMessageNotFound {
-		return errHTTPNotFound
+		if s.config.CacheBatchTimeout > 0 {
+			// Strange edge case: If we immediately after upload request the file (the web app does this for images),
+			// and messages are persisted asynchronously, retry fetching from the database
+			m, err = util.Retry(func() (*message, error) {
+				return s.messageCache.Message(messageID)
+			}, s.config.CacheBatchTimeout, 100*time.Millisecond, 300*time.Millisecond, 600*time.Millisecond)
+		}
+		if err != nil {
+			return errHTTPNotFound
+		}
 	} else if err != nil {
 		return err
 	}
diff --git a/server/server_test.go b/server/server_test.go
index bedc55d0..6fb614bf 100644
--- a/server/server_test.go
+++ b/server/server_test.go
@@ -1714,6 +1714,25 @@ func TestServer_PublishAttachmentBandwidthLimitUploadOnly(t *testing.T) {
 	require.Equal(t, 41301, err.Code)
 }
 
+func TestServer_PublishAttachmentAndImmediatelyGetItWithCacheTimeout(t *testing.T) {
+	// This tests the awkward util.Retry in handleFile: Due to the async persisting of messages,
+	// the message is not immediately available when attempting to download it.
+
+	c := newTestConfig(t)
+	c.CacheBatchTimeout = 500 * time.Millisecond
+	c.CacheBatchSize = 10
+	s := newTestServer(t, c)
+	content := "this is an ATTACHMENT"
+	rr := request(t, s, "PUT", "/mytopic?f=myfile.txt", content, nil)
+	m := toMessage(t, rr.Body.String())
+	require.Equal(t, "myfile.txt", m.Attachment.Name)
+
+	path := strings.TrimPrefix(m.Attachment.URL, "http://127.0.0.1:12345")
+	rr = request(t, s, "GET", path, "", nil)
+	require.Equal(t, 200, rr.Code) // Not 404!
+	require.Equal(t, content, rr.Body.String())
+}
+
 func TestServer_PublishAttachmentAccountStats(t *testing.T) {
 	content := util.RandomString(4999) // > 4096