diff --git a/internal/configuration/database/Database.go b/internal/configuration/database/Database.go index 0de9867c..6fca4e77 100644 --- a/internal/configuration/database/Database.go +++ b/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 diff --git a/internal/configuration/database/Database_test.go b/internal/configuration/database/Database_test.go index d44cea85..29e6fd28 100644 --- a/internal/configuration/database/Database_test.go +++ b/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") }) diff --git a/internal/configuration/database/dbabstraction/DbAbstraction.go b/internal/configuration/database/dbabstraction/DbAbstraction.go index aa283e80..fea4ec20 100644 --- a/internal/configuration/database/dbabstraction/DbAbstraction.go +++ b/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) diff --git a/internal/configuration/database/provider/redis/Redis_test.go b/internal/configuration/database/provider/redis/Redis_test.go index 7523af9c..9774ff14 100644 --- a/internal/configuration/database/provider/redis/Redis_test.go +++ b/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() diff --git a/internal/configuration/database/provider/redis/metadata.go b/internal/configuration/database/provider/redis/metadata.go index 568379f5..ba52bca4 100644 --- a/internal/configuration/database/provider/redis/metadata.go +++ b/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 +} diff --git a/internal/configuration/database/provider/sqlite/Sqlite_test.go b/internal/configuration/database/provider/sqlite/Sqlite_test.go index 10443deb..e315d9a7 100644 --- a/internal/configuration/database/provider/sqlite/Sqlite_test.go +++ b/internal/configuration/database/provider/sqlite/Sqlite_test.go @@ -113,6 +113,7 @@ func TestMetaData(t *testing.T) { dbInstance.SaveMetaData(models.File{ Id: "test3", Name: "test3", + DownloadsRemaining: 4, UnlimitedDownloads: false, UnlimitedTime: true, }) @@ -120,6 +121,9 @@ func TestMetaData(t *testing.T) { 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() diff --git a/internal/configuration/database/provider/sqlite/metadata.go b/internal/configuration/database/provider/sqlite/metadata.go index e5a7f731..ef8f8791 100644 --- a/internal/configuration/database/provider/sqlite/metadata.go +++ b/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) diff --git a/internal/storage/FileServing.go b/internal/storage/FileServing.go index 8458562b..9ef1228b 100644 --- a/internal/storage/FileServing.go +++ b/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" diff --git a/internal/storage/FileServing_test.go b/internal/storage/FileServing_test.go index 5e9e0031..d17005e9 100644 --- a/internal/storage/FileServing_test.go +++ b/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) + }) +} diff --git a/internal/webserver/Webserver.go b/internal/webserver/Webserver.go index 47e60a3e..6b707e5b 100644 --- a/internal/webserver/Webserver.go +++ b/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 { diff --git a/internal/webserver/api/Api.go b/internal/webserver/api/Api.go index 796f8b5f..9fedee13 100644 --- a/internal/webserver/api/Api.go +++ b/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, "")