Fix race condition when checking remaining downloads

internal/configuration/database/Database.go

@@ -221,11 +221,16 @@ func DeleteMetaData(id string) {
 	db.DeleteMetaData(id)
 }
 
-// IncreaseDownloadCount increases the download count of a file, preventing race conditions
+// IncreaseDownloadCount increases the download count of a file atomically
 func IncreaseDownloadCount(id string, decreaseRemainingDownloads bool) {
 	db.IncreaseDownloadCount(id, decreaseRemainingDownloads)
 }
 
+// GetDownloadsRemaining returns the remaining downloads of a file that does not implement UnlimitedDownloads
+func GetDownloadsRemaining(id string) int {
+	return db.GetDownloadsRemaining(id)
+}
+
 // Session Section
 
 // GetSession returns the session with the given ID or false if not a valid ID

internal/configuration/database/Database_test.go

@@ -213,6 +213,7 @@ func TestMetaData(t *testing.T) {
 		UnlimitedTime:      true,
 	}
 	runAllTypesNoOutput(t, func() { SaveMetaData(file) })
+	runAllTypesCompareOutput(t, func() any { return GetDownloadsRemaining(file.Id) }, 2)
 	runAllTypesCompareTwoOutputs(t, func() (any, any) { return GetMetaDataById("testid") }, file, true)
 	runAllTypesCompareOutput(t, func() any { return GetAllMetadata() }, map[string]models.File{"testid": file})
 	runAllTypesNoOutput(t, func() { DeleteMetaData("testid") })

internal/configuration/database/dbabstraction/DbAbstraction.go

@@ -73,6 +73,8 @@ type Database interface {
 	DeleteMetaData(id string)
 	// IncreaseDownloadCount increases the download count of a file, preventing race conditions
 	IncreaseDownloadCount(id string, decreaseRemainingDownloads bool)
+	// GetDownloadsRemaining returns the remaining downloads of a file that does not implement UnlimitedDownloads
+	GetDownloadsRemaining(id string) int
 
 	// GetSession returns the session with the given ID or false if not a valid ID
 	GetSession(id string) (models.Session, bool)

internal/configuration/database/provider/redis/Redis_test.go

@@ -540,12 +540,16 @@ func TestMetaData(t *testing.T) {
 		Id:                 "test3",
 		Name:               "test3",
 		UnlimitedDownloads: false,
+		DownloadsRemaining: 4,
 		UnlimitedTime:      true,
 	})
 	file, ok = dbInstance.GetMetaDataById("test3")
 	test.IsEqualBool(t, ok, true)
 	test.IsEqualBool(t, file.UnlimitedDownloads, false)
 	test.IsEqualBool(t, file.UnlimitedTime, true)
+	test.IsEqualInt(t, file.DownloadsRemaining, 4)
+	remaining := dbInstance.GetDownloadsRemaining(file.Id)
+	test.IsEqualInt(t, remaining, 4)
 	dbInstance.Close()
 	defer test.ExpectPanic(t)
 	_ = dbInstance.GetAllMetadata()

internal/configuration/database/provider/redis/metadata.go

@@ -3,10 +3,11 @@ package redis
 import (
 	"bytes"
 	"encoding/gob"
+	"strings"
+
 	"github.com/forceu/gokapi/internal/helper"
 	"github.com/forceu/gokapi/internal/models"
 	redigo "github.com/gomodule/redigo/redis"
-	"strings"
 )
 
 const (
@@ -86,10 +87,19 @@ func (p DatabaseProvider) DeleteMetaData(id string) {
 	p.deleteKey(prefixMetaData + id)
 }
 
-// IncreaseDownloadCount increases the download count of a file, preventing race conditions
+// IncreaseDownloadCount increases the download count of a file atomically
 func (p DatabaseProvider) IncreaseDownloadCount(id string, decreaseRemainingDownloads bool) {
 	if decreaseRemainingDownloads {
 		p.decreaseHashmapIntField(prefixMetaData+id, "DownloadsRemaining")
 	}
 	p.increaseHashmapIntField(prefixMetaData+id, "DownloadCount")
 }
+
+// GetDownloadsRemaining returns the remaining downloads of a file that does not implement UnlimitedDownloads
+func (p DatabaseProvider) GetDownloadsRemaining(id string) int {
+	file, ok := p.GetMetaDataById(id)
+	if !ok {
+		return 0
+	}
+	return file.DownloadsRemaining
+}

internal/configuration/database/provider/sqlite/Sqlite_test.go

@@ -113,13 +113,17 @@ func TestMetaData(t *testing.T) {
 	dbInstance.SaveMetaData(models.File{
 		Id:                 "test3",
 		Name:               "test3",
+		DownloadsRemaining: 4,
 		UnlimitedDownloads: false,
 		UnlimitedTime:      true,
 	})
 	file, ok = dbInstance.GetMetaDataById("test3")
 	test.IsEqualBool(t, ok, true)
 	test.IsEqualBool(t, file.UnlimitedDownloads, false)
 	test.IsEqualBool(t, file.UnlimitedTime, true)
+	test.IsEqualInt(t, file.DownloadsRemaining, 4)
+	remaining := dbInstance.GetDownloadsRemaining(file.Id)
+	test.IsEqualInt(t, remaining, 4)
 	dbInstance.Close()
 	defer test.ExpectPanic(t)
 	_ = dbInstance.GetAllMetadata()

internal/configuration/database/provider/sqlite/metadata.go

@@ -150,7 +150,7 @@ func (p DatabaseProvider) SaveMetaData(file models.File) {
 	helper.Check(err)
 }
 
-// IncreaseDownloadCount increases the download count of a file, preventing race conditions
+// IncreaseDownloadCount increases the download count of a file atomically
 func (p DatabaseProvider) IncreaseDownloadCount(id string, decreaseRemainingDownloads bool) {
 	if decreaseRemainingDownloads {
 		_, err := p.sqliteDb.Exec(`UPDATE FileMetaData SET DownloadCount = DownloadCount + 1,
@@ -162,6 +162,21 @@ func (p DatabaseProvider) IncreaseDownloadCount(id string, decreaseRemainingDown
 	}
 }
 
+// GetDownloadsRemaining returns the remaining downloads of a file that does not implement UnlimitedDownloads
+func (p DatabaseProvider) GetDownloadsRemaining(id string) int {
+	var downloadsRemaining int
+	row := p.sqliteDb.QueryRow("SELECT DownloadsRemaining FROM FileMetaData WHERE Id = ?", id)
+	err := row.Scan(&downloadsRemaining)
+	if err != nil {
+		if errors.Is(err, sql.ErrNoRows) {
+			return 0
+		}
+		helper.Check(err)
+		return downloadsRemaining
+	}
+	return downloadsRemaining
+}
+
 // DeleteMetaData deletes information about a file
 func (p DatabaseProvider) DeleteMetaData(id string) {
 	_, err := p.sqliteDb.Exec("DELETE FROM FileMetaData WHERE Id = ?", id)

internal/storage/FileServing.go

@@ -32,6 +32,7 @@ import (
 	"github.com/forceu/gokapi/internal/storage/filesystem"
 	"github.com/forceu/gokapi/internal/storage/filesystem/s3filesystem/aws"
 	"github.com/forceu/gokapi/internal/storage/processingstatus"
+	"github.com/forceu/gokapi/internal/webserver/api/mutex/apimutex"
 	"github.com/forceu/gokapi/internal/webserver/downloadstatus"
 	"github.com/forceu/gokapi/internal/webserver/headers"
 	"github.com/forceu/gokapi/internal/webserver/sse"
@@ -618,13 +619,26 @@ func GetFileByHotlink(id string) (models.File, bool) {
 }
 
 // ServeFile subtracts a download allowance and serves the file to the browser
-func ServeFile(file models.File, w http.ResponseWriter, r *http.Request, forceDownload, increaseCounter, forceDecryption bool) {
+// Returns false if the file expired during the request (most likely race condition due to parallel downloads, requires recheckExpiry)
+func ServeFile(file models.File, w http.ResponseWriter, r *http.Request, forceDownload, increaseCounter, forceDecryption, recheckExpiry bool) bool {
+	apimutex.Lock(apimutex.TypeMetaData, file.Id)
+	if recheckExpiry {
+		if !file.UnlimitedDownloads {
+			file.DownloadsRemaining = database.GetDownloadsRemaining(file.Id)
+		}
+		if IsExpiredFile(file, time.Now().Unix()) {
+			apimutex.Unlock(apimutex.TypeMetaData, file.Id)
+			return false
+		}
+	}
 	if increaseCounter {
 		file.DownloadsRemaining = file.DownloadsRemaining - 1
 		file.DownloadCount = file.DownloadCount + 1
 		database.IncreaseDownloadCount(file.Id, !file.UnlimitedDownloads)
 		go sse.PublishDownloadCount(file)
 	}
+	apimutex.Unlock(apimutex.TypeMetaData, file.Id)
+
 	logging.LogDownload(file, r, configuration.Get().SaveIp)
 	go serverstats.AddTraffic(uint64(file.SizeBytes))
 
@@ -638,19 +652,19 @@ func ServeFile(file models.File, w http.ResponseWriter, r *http.Request, forceDo
 		if isBlocking {
 			downloadstatus.SetComplete(statusId)
 		}
-		return
+		return true
 	}
 	fileHandler, _, err := getFileHandler(file, configuration.Get().DataDir)
 	defer fileHandler.Close()
 	if err != nil {
 		fmt.Println(err)
 		_, _ = w.Write([]byte("Error getting file handler"))
-		return
+		return true
 	}
 	if file.Encryption.IsEncrypted && !file.RequiresClientDecryption() {
 		if !encryption.IsCorrectKey(file.Encryption, fileHandler) {
 			_, _ = w.Write([]byte("Internal error - Error decrypting file, source data might be damaged or an incorrect key has been used"))
-			return
+			return true
 		}
 	}
 	statusId := downloadstatus.SetDownload(file)
@@ -660,12 +674,13 @@ func ServeFile(file models.File, w http.ResponseWriter, r *http.Request, forceDo
 		if err != nil {
 			_, _ = w.Write([]byte("Error decrypting file"))
 			fmt.Println(err)
-			return
+			return true
 		}
 	} else {
 		http.ServeContent(w, r, file.Name, time.Now(), fileHandler)
 	}
 	downloadstatus.SetComplete(statusId)
+	return true
 }
 
 // Returns the filename if unique or a new filename in the format "Name (x).ext"

internal/storage/FileServing_test.go

@@ -9,6 +9,7 @@ import (
 	"net/textproto"
 	"os"
 	"testing"
+	"testing/synctest"
 	"time"
 
 	"github.com/forceu/gokapi/internal/configuration"
@@ -573,7 +574,7 @@ func TestServeFile(t *testing.T) {
 	test.IsEqualBool(t, result, true)
 	r := httptest.NewRequest("GET", "/", nil)
 	w := httptest.NewRecorder()
-	ServeFile(file, w, r, true, true, false)
+	ServeFile(file, w, r, true, true, false, false)
 	_, result = GetFile(idNewFile)
 	test.IsEqualBool(t, result, false)
 
@@ -594,7 +595,7 @@ func TestServeFile(t *testing.T) {
 		w = httptest.NewRecorder()
 		file, result = GetFile("awsTest1234567890123")
 		test.IsEqualBool(t, result, true)
-		ServeFile(file, w, r, false, true, false)
+		ServeFile(file, w, r, false, true, false, false)
 		if aws.IsMockApi {
 			test.ResponseBodyContains(t, w, "https://redirect.url")
 		} else {
@@ -619,7 +620,7 @@ func TestServeFile(t *testing.T) {
 	file.Encryption.Nonce = nonce
 	r = httptest.NewRequest("GET", "/", nil)
 	w = httptest.NewRecorder()
-	ServeFile(file, w, r, true, true, false)
+	ServeFile(file, w, r, true, true, false, false)
 	test.ResponseBodyContains(t, w, "Error decrypting file")
 }
 
@@ -893,3 +894,59 @@ func TestReplaceFile(t *testing.T) {
 	_, ok = GetFile(newFile.Id)
 	test.IsEqualBool(t, ok, false)
 }
+func TestParallelDownloads(t *testing.T) {
+	const allowedDownloads = 5
+
+	singleDownloadFile := models.File{
+		Id:                 "only5downloads",
+		Name:               "only5downloads.txt",
+		Size:               "1KB",
+		SHA1:               "replacetest1",
+		ContentType:        "text/plain",
+		DownloadsRemaining: allowedDownloads,
+		UnlimitedDownloads: false,
+		UnlimitedTime:      true,
+		AwsBucket:          "",
+		SizeBytes:          1024,
+		Encryption: models.EncryptionInfo{
+			IsEncrypted:         false,
+			IsEndToEndEncrypted: false,
+			DecryptionKey:       nil,
+			Nonce:               nil,
+		},
+	}
+	database.SaveMetaData(singleDownloadFile)
+
+	synctest.Test(t, func(t *testing.T) {
+		const workers = 50
+		results := make(chan bool, workers)
+
+		for i := 0; i < workers; i++ {
+			go func() {
+				w := httptest.NewRecorder()
+				r := httptest.NewRequest("GET", "/"+singleDownloadFile.Id, nil)
+
+				// The mutex inside ServeFile should serialize the decrement logic.
+				success := ServeFile(singleDownloadFile, w, r, false, true, false, true)
+				results <- success
+			}()
+		}
+
+		synctest.Wait()
+		close(results)
+
+		var successCount int
+		var failureCount int
+
+		for res := range results {
+			if res {
+				successCount++
+			} else {
+				failureCount++
+			}
+		}
+
+		test.IsEqualInt(t, successCount, allowedDownloads)
+		test.IsEqualInt(t, failureCount, workers-allowedDownloads)
+	})
+}

internal/webserver/Webserver.go

@@ -637,7 +637,13 @@ func showHotlink(w http.ResponseWriter, r *http.Request) {
 		_, _ = w.Write(imageExpiredPicture)
 		return
 	}
-	storage.ServeFile(file, w, r, false, true, false)
+	validFile := storage.ServeFile(file, w, r, false, true, false, true)
+	if !validFile {
+		// Only called if the file has already expired during the expiry check of storage.ServeFile()
+		w.Header().Set("Content-Type", "image/svg+xml")
+		_, _ = w.Write(imageExpiredPicture)
+		return
+	}
 }
 
 // Checks if a file is associated with the GET parameter from the current URL
@@ -1057,7 +1063,7 @@ func downloadPresigned(w http.ResponseWriter, r *http.Request) {
 	if len(files) == 1 {
 		file := files[0]
 		forceDecryption := file.Encryption.IsEncrypted && !file.Encryption.IsEndToEndEncrypted
-		storage.ServeFile(file, w, r, true, false, forceDecryption)
+		storage.ServeFile(file, w, r, true, false, forceDecryption, false)
 		return
 	}
 	storage.ServeFilesAsZip(files, presignedUrl.Filename, w, r)
@@ -1066,6 +1072,7 @@ func downloadPresigned(w http.ResponseWriter, r *http.Request) {
 func serveFile(id string, isRootUrl bool, w http.ResponseWriter, r *http.Request) {
 	addNoCacheHeader(w)
 	savedFile, ok := storage.GetFile(id)
+
 	if !ok || savedFile.IsFileRequest() {
 		if isRootUrl {
 			redirectOnIncorrectId(w, r, "error")
@@ -1084,7 +1091,16 @@ func serveFile(id string, isRootUrl bool, w http.ResponseWriter, r *http.Request
 			return
 		}
 	}
-	storage.ServeFile(savedFile, w, r, true, true, false)
+	validFile := storage.ServeFile(savedFile, w, r, true, true, false, true)
+	if !validFile {
+		// Only called if the file has already expired during the expiry check of storage.ServeFile()
+		if isRootUrl {
+			redirectOnIncorrectId(w, r, "error")
+		} else {
+			redirectOnIncorrectId(w, r, "../../error")
+		}
+		return
+	}
 }
 
 func requireLogin(next http.HandlerFunc, isUiCall, isPwChangeView bool) http.HandlerFunc {

internal/webserver/api/Api.go

@@ -626,7 +626,7 @@ func apiDownloadSingle(w http.ResponseWriter, r requestParser, user models.User)
 	}
 	if !request.PresignUrl {
 		forceDecryption := file.Encryption.IsEncrypted && !file.Encryption.IsEndToEndEncrypted
-		storage.ServeFile(file, w, request.WebRequest, true, request.IncreaseCounter, forceDecryption)
+		storage.ServeFile(file, w, request.WebRequest, true, request.IncreaseCounter, forceDecryption, false)
 		return
 	}
 	createAndOutputPresignedUrl([]string{file.Id}, w, "")