From 303e789d829d3b5495cfbd894ca53b778a2a4d53 Mon Sep 17 00:00:00 2001
From: MrMelon54 <github@mrmelon54.com>
Date: Tue, 7 Jan 2025 23:54:12 +0000
Subject: [PATCH] Protect upload calls with a mutex for each site

---
 go.mod           |  1 +
 go.sum           |  2 ++
 upload/upload.go | 18 +++++++++++++++++-
 3 files changed, 20 insertions(+), 1 deletion(-)

diff --git a/go.mod b/go.mod
index f1f6b27..bdd51ae 100644
--- a/go.mod
+++ b/go.mod
@@ -4,6 +4,7 @@ go 1.23.4
 
 require (
 	github.com/1f349/mjwt v0.4.1
+	github.com/1f349/syncmap v0.0.3
 	github.com/charmbracelet/log v0.4.0
 	github.com/cloudflare/tableflip v1.2.3
 	github.com/dustin/go-humanize v1.0.1
diff --git a/go.sum b/go.sum
index 4d6bed6..e31a4d7 100644
--- a/go.sum
+++ b/go.sum
@@ -2,6 +2,8 @@ github.com/1f349/mjwt v0.4.1 h1:ooCroMMw2kcL5c9L3sLbdtxI0H4/QC8RfTxiloKr+4Y=
 github.com/1f349/mjwt v0.4.1/go.mod h1:qwnzokkqc7Z9YmKA1m9beI3OZL1GvGYHOQU2rOwoV1M=
 github.com/1f349/rsa-helper v0.0.2 h1:N/fLQqg5wrjIzG6G4zdwa5Xcv9/jIPutCls9YekZr9U=
 github.com/1f349/rsa-helper v0.0.2/go.mod h1:VUQ++1tYYhYrXeOmVFkQ82BegR24HQEJHl5lHbjg7yg=
+github.com/1f349/syncmap v0.0.3 h1:Xc0XJwS+OF8UrbPG/C6eLubNSVFinEZ6SEIx7WmAZ0Y=
+github.com/1f349/syncmap v0.0.3/go.mod h1:MAfbJgAHKVZgkEqxpbgmMP4yga7Ajmo/KMJLoXVCmZk=
 github.com/aymanbagabas/go-osc52/v2 v2.0.1 h1:HwpRHbFMcZLEVr42D4p7XBqjyuxQH5SMiErDT4WkJ2k=
 github.com/aymanbagabas/go-osc52/v2 v2.0.1/go.mod h1:uYgXzlJ7ZpABp8OJ+exZzJJhRNQ2ASbcXHWsFqH8hp8=
 github.com/becheran/wildmatch-go v1.0.0 h1:mE3dGGkTmpKtT4Z+88t8RStG40yN9T+kFEGj2PZFSzA=
diff --git a/upload/upload.go b/upload/upload.go
index 7fee8b7..3ecc55f 100644
--- a/upload/upload.go
+++ b/upload/upload.go
@@ -8,6 +8,7 @@ import (
 	"errors"
 	"fmt"
 	"github.com/1f349/bluebell/database"
+	"github.com/1f349/syncmap"
 	"github.com/dustin/go-humanize"
 	"github.com/julienschmidt/httprouter"
 	"github.com/spf13/afero"
@@ -17,6 +18,7 @@ import (
 	"path/filepath"
 	"slices"
 	"strings"
+	"sync"
 )
 
 var indexBranches = []string{
@@ -82,7 +84,7 @@ type sitesQueries interface {
 }
 
 func New(storage afero.Fs, db sitesQueries) *Handler {
-	return &Handler{storage, db}
+	return &Handler{storageFs: storage, db: db}
 }
 
 const maxFileSize = 1 * humanize.GiByte
@@ -90,6 +92,7 @@ const maxFileSize = 1 * humanize.GiByte
 type Handler struct {
 	storageFs afero.Fs
 	db        sitesQueries
+	mu        syncmap.Map[string, *sync.Mutex]
 }
 
 func (h *Handler) Handle(rw http.ResponseWriter, req *http.Request, params httprouter.Params) {
@@ -144,6 +147,19 @@ func (h *Handler) extractTarGzUpload(fileData io.Reader, site, branch string) er
 		return fmt.Errorf("invalid site: %w", err)
 	}
 
+	key := site + "@" + branch
+
+	// ensure upload mutex is locked
+	actual, _ := h.mu.LoadOrStore(key, new(sync.Mutex))
+	actual.Lock()
+	defer func() {
+		// The mutex is no longer used so delete it here to safe memory in a "lots of
+		// sites" configuration. Delete should happen first to prevent another upload
+		// reusing the mutex.
+		h.mu.Delete(key)
+		actual.Unlock()
+	}()
+
 	siteBranchPath := filepath.Join(site, "@"+branch)
 	siteBranchOldPath := filepath.Join(site, "old@"+branch)