Change download count atomically to prevent race condition

internal/configuration/database/Database.go

@@ -200,6 +200,11 @@ func DeleteMetaData(id string) {
 	db.DeleteMetaData(id)
 }
 
+// IncreaseDownloadCount increases the download count of a file, preventing race conditions
+func IncreaseDownloadCount(id string, decreaseRemainingDownloads bool) {
+	db.IncreaseDownloadCount(id, decreaseRemainingDownloads)
+}
+
 // Session Section
 
 // GetSession returns the session with the given ID or false if not a valid ID

internal/configuration/database/Database_test.go

@@ -158,6 +158,24 @@ func TestMetaData(t *testing.T) {
 	runAllTypesCompareOutput(t, func() any { return GetAllMetaDataIds() }, []string{})
 	runAllTypesCompareOutput(t, func() any { return GetAllMetadata() }, map[string]models.File{})
 	runAllTypesCompareTwoOutputs(t, func() (any, any) { return GetMetaDataById("testid") }, models.File{}, false)
+
+	increasedDownload := file
+	increasedDownload.DownloadCount = increasedDownload.DownloadCount + 1
+
+	runAllTypesCompareTwoOutputs(t, func() (any, any) {
+		SaveMetaData(file)
+		IncreaseDownloadCount(file.Id, false)
+		return GetMetaDataById(file.Id)
+	}, increasedDownload, true)
+
+	increasedDownload.DownloadCount = increasedDownload.DownloadCount + 1
+	increasedDownload.DownloadsRemaining = increasedDownload.DownloadsRemaining - 1
+
+	runAllTypesCompareTwoOutputs(t, func() (any, any) {
+		IncreaseDownloadCount(file.Id, true)
+		return GetMetaDataById(file.Id)
+	}, increasedDownload, true)
+	runAllTypesNoOutput(t, func() { DeleteMetaData(file.Id) })
 }
 
 func TestUpgrade(t *testing.T) {

internal/configuration/database/dbabstraction/DbAbstraction.go

@@ -72,6 +72,8 @@ type Database interface {
 	SaveMetaData(file models.File)
 	// DeleteMetaData deletes information about a file
 	DeleteMetaData(id string)
+	// IncreaseDownloadCount increases the download count of a file, preventing race conditions
+	IncreaseDownloadCount(id string, decreaseRemainingDownloads bool)
 
 	// 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.go

@@ -1,6 +1,8 @@
 package redis
 
 import (
+	"bytes"
+	"encoding/gob"
 	"errors"
 	"fmt"
 	"github.com/forceu/gokapi/internal/helper"
@@ -17,7 +19,7 @@ type DatabaseProvider struct {
 	dbPrefix string
 }
 
-const DatabaseSchemeVersion = 2
+const DatabaseSchemeVersion = 3
 
 // New returns an instance
 func New(dbConfig models.DbConnection) (DatabaseProvider, error) {
@@ -91,8 +93,35 @@ func newPool(config models.DbConnection) *redigo.Pool {
 
 // Upgrade migrates the DB to a new Gokapi version, if required
 func (p DatabaseProvider) Upgrade(currentDbVersion int) {
-	// Currently no upgrade necessary
-	return
+	// < 1.9.6
+	if currentDbVersion < 3 {
+		allFiles := getAllLegacyMetaDataAndDelete(p)
+		for _, file := range allFiles {
+			p.SaveMetaData(file)
+		}
+	}
+}
+
+func getAllLegacyMetaDataAndDelete(p DatabaseProvider) map[string]models.File {
+	result := make(map[string]models.File)
+	allMetaData := p.getAllValuesWithPrefix(prefixMetaData)
+	for _, metaData := range allMetaData {
+		content, err := redigo.Bytes(metaData, nil)
+		helper.Check(err)
+		file := legacyDbToMetaData(content)
+		result[file.Id] = file
+		p.deleteKey(prefixMetaData + file.Id)
+	}
+	return result
+}
+
+func legacyDbToMetaData(input []byte) models.File {
+	var result models.File
+	buf := bytes.NewBuffer(input)
+	dec := gob.NewDecoder(buf)
+	err := dec.Decode(&result)
+	helper.Check(err)
+	return result
 }
 
 const keyDbVersion = "dbversion"
@@ -260,6 +289,20 @@ func (p DatabaseProvider) deleteKey(id string) {
 	helper.Check(err)
 }
 
+func (p DatabaseProvider) increaseHashmapIntField(id string, field string) {
+	conn := p.pool.Get()
+	defer conn.Close()
+	_, err := conn.Do("HINCRBY", p.dbPrefix+id, field, 1)
+	helper.Check(err)
+}
+
+func (p DatabaseProvider) decreaseHashmapIntField(id string, field string) {
+	conn := p.pool.Get()
+	defer conn.Close()
+	_, err := conn.Do("HINCRBY", p.dbPrefix+id, field, -1)
+	helper.Check(err)
+}
+
 func (p DatabaseProvider) runEval(cmd string) {
 	conn := p.pool.Get()
 	defer conn.Close()

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

@@ -241,6 +241,49 @@ func TestApiKeys(t *testing.T) {
 	test.IsEqualBool(t, key.LastUsed == 10, true)
 }
 
+func TestDatabaseProvider_IncreaseDownloadCount(t *testing.T) {
+	newFile := models.File{
+		Id:                 "newFileId",
+		Name:               "newFileName",
+		Size:               "3GB",
+		SHA1:               "newSHA1",
+		PasswordHash:       "newPassword",
+		HotlinkId:          "newHotlink",
+		ContentType:        "newContent",
+		AwsBucket:          "newAws",
+		ExpireAt:           123456,
+		SizeBytes:          456789,
+		DownloadsRemaining: 11,
+		DownloadCount:      2,
+		Encryption: models.EncryptionInfo{
+			IsEncrypted:         true,
+			IsEndToEndEncrypted: true,
+			DecryptionKey:       []byte("newDecryptionKey"),
+			Nonce:               []byte("newDecryptionNonce"),
+		},
+		UnlimitedDownloads: true,
+		UnlimitedTime:      true,
+	}
+	dbInstance.SaveMetaData(newFile)
+	dbInstance.IncreaseDownloadCount(newFile.Id, false)
+	retrievedFile, ok := dbInstance.GetMetaDataById(newFile.Id)
+	test.IsEqualBool(t, ok, true)
+	test.IsEqualInt(t, retrievedFile.DownloadCount, 3)
+	test.IsEqualInt(t, retrievedFile.DownloadsRemaining, 11)
+	newFile.DownloadCount = 3
+	test.IsEqual(t, retrievedFile, newFile)
+
+	dbInstance.IncreaseDownloadCount(newFile.Id, true)
+	retrievedFile, ok = dbInstance.GetMetaDataById(newFile.Id)
+	test.IsEqualBool(t, ok, true)
+	test.IsEqualInt(t, retrievedFile.DownloadCount, 4)
+	test.IsEqualInt(t, retrievedFile.DownloadsRemaining, 10)
+	newFile.DownloadCount = 4
+	newFile.DownloadsRemaining = 10
+	test.IsEqual(t, retrievedFile, newFile)
+	dbInstance.DeleteMetaData(newFile.Id)
+}
+
 func TestE2EConfig(t *testing.T) {
 	e2econfig := models.E2EInfoEncrypted{
 		Version:        1,

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

@@ -1,8 +1,6 @@
 package redis
 
 import (
-	"bytes"
-	"encoding/gob"
 	"github.com/forceu/gokapi/internal/helper"
 	"github.com/forceu/gokapi/internal/models"
 	redigo "github.com/gomodule/redigo/redis"
@@ -13,28 +11,29 @@ const (
 	prefixMetaData = "fmeta:"
 )
 
-func dbToMetaData(input []byte) models.File {
-	var result models.File
-	buf := bytes.NewBuffer(input)
-	dec := gob.NewDecoder(buf)
-	err := dec.Decode(&result)
-	helper.Check(err)
-	return result
-}
-
 // GetAllMetadata returns a map of all available files
 func (p DatabaseProvider) GetAllMetadata() map[string]models.File {
 	result := make(map[string]models.File)
-	allMetaData := p.getAllValuesWithPrefix(prefixMetaData)
-	for _, metaData := range allMetaData {
-		content, err := redigo.Bytes(metaData, nil)
+	maps := p.getAllHashesWithPrefix(prefixMetaData)
+	for k, v := range maps {
+		file, err := newDbToMetadata(k, v)
 		helper.Check(err)
-		file := dbToMetaData(content)
 		result[file.Id] = file
 	}
 	return result
 }
 
+func newDbToMetadata(id string, input []any) (models.File, error) {
+	var result models.File
+	err := redigo.ScanStruct(input, &result)
+	if err != nil {
+		return models.File{}, err
+	}
+	result.Id = strings.Replace(id, prefixMetaData, "", 1)
+	err = result.RedisToFile()
+	return result, err
+}
+
 // GetAllMetaDataIds returns all Ids that contain metadata
 func (p DatabaseProvider) GetAllMetaDataIds() []string {
 	result := make([]string, 0)
@@ -46,23 +45,31 @@ func (p DatabaseProvider) GetAllMetaDataIds() []string {
 
 // GetMetaDataById returns a models.File from the ID passed or false if the id is not valid
 func (p DatabaseProvider) GetMetaDataById(id string) (models.File, bool) {
-	input, ok := p.getKeyBytes(prefixMetaData + id)
+	result, ok := p.getHashMap(prefixMetaData + id)
 	if !ok {
 		return models.File{}, false
 	}
-	return dbToMetaData(input), true
+	file, err := newDbToMetadata(id, result)
+	helper.Check(err)
+	return file, true
 }
 
 // SaveMetaData stores the metadata of a file to the disk
 func (p DatabaseProvider) SaveMetaData(file models.File) {
-	var buf bytes.Buffer
-	enc := gob.NewEncoder(&buf)
-	err := enc.Encode(file)
+	err := file.FileToRedis()
 	helper.Check(err)
-	p.setKey(prefixMetaData+file.Id, buf.Bytes())
+	p.setHashMap(p.buildArgs(prefixMetaData + file.Id).AddFlat(file))
 }
 
 // DeleteMetaData deletes information about a file
 func (p DatabaseProvider) DeleteMetaData(id string) {
 	p.deleteKey(prefixMetaData + id)
 }
+
+// IncreaseDownloadCount increases the download count of a file, preventing race conditions
+func (p DatabaseProvider) IncreaseDownloadCount(id string, decreaseRemainingDownloads bool) {
+	if decreaseRemainingDownloads {
+		p.decreaseHashmapIntField(prefixMetaData+id, "DownloadsRemaining")
+	}
+	p.increaseHashmapIntField(prefixMetaData+id, "DownloadCount")
+}

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

@@ -207,15 +207,6 @@ func (p DatabaseProvider) createNewDatabase() error {
 			"ValidUntil"	INTEGER NOT NULL,
 			PRIMARY KEY("Id")
 		) WITHOUT ROWID;
-		CREATE TABLE "UploadConfig" (
-			"id"	INTEGER NOT NULL UNIQUE,
-			"Downloads"	INTEGER,
-			"TimeExpiry"	INTEGER,
-			"Password"	TEXT,
-			"UnlimitedDownloads"	INTEGER,
-			"UnlimitedTime"	INTEGER,
-			PRIMARY KEY("id")
-		);
 		CREATE TABLE "UploadStatus" (
 			"ChunkId"	TEXT NOT NULL UNIQUE,
 			"CurrentStatus"	INTEGER NOT NULL,

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

@@ -176,6 +176,49 @@ func TestHotlink(t *testing.T) {
 	test.IsEqualInt(t, len(hotlinks), 3)
 }
 
+func TestDatabaseProvider_IncreaseDownloadCount(t *testing.T) {
+	newFile := models.File{
+		Id:                 "newFileId",
+		Name:               "newFileName",
+		Size:               "3GB",
+		SHA1:               "newSHA1",
+		PasswordHash:       "newPassword",
+		HotlinkId:          "newHotlink",
+		ContentType:        "newContent",
+		AwsBucket:          "newAws",
+		ExpireAt:           123456,
+		SizeBytes:          456789,
+		DownloadsRemaining: 11,
+		DownloadCount:      2,
+		Encryption: models.EncryptionInfo{
+			IsEncrypted:         true,
+			IsEndToEndEncrypted: true,
+			DecryptionKey:       []byte("newDecryptionKey"),
+			Nonce:               []byte("newDecryptionNonce"),
+		},
+		UnlimitedDownloads: true,
+		UnlimitedTime:      true,
+	}
+	dbInstance.SaveMetaData(newFile)
+	dbInstance.IncreaseDownloadCount(newFile.Id, false)
+	retrievedFile, ok := dbInstance.GetMetaDataById(newFile.Id)
+	test.IsEqualBool(t, ok, true)
+	test.IsEqualInt(t, retrievedFile.DownloadCount, 3)
+	test.IsEqualInt(t, retrievedFile.DownloadsRemaining, 11)
+	newFile.DownloadCount = 3
+	test.IsEqual(t, retrievedFile, newFile)
+
+	dbInstance.IncreaseDownloadCount(newFile.Id, true)
+	retrievedFile, ok = dbInstance.GetMetaDataById(newFile.Id)
+	test.IsEqualBool(t, ok, true)
+	test.IsEqualInt(t, retrievedFile.DownloadCount, 4)
+	test.IsEqualInt(t, retrievedFile.DownloadsRemaining, 10)
+	newFile.DownloadCount = 4
+	newFile.DownloadsRemaining = 10
+	test.IsEqual(t, retrievedFile, newFile)
+	dbInstance.DeleteMetaData(newFile.Id)
+}
+
 func TestApiKey(t *testing.T) {
 	dbInstance.SaveApiKey(models.ApiKey{
 		Id:           "newkey",

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

@@ -158,6 +158,18 @@ func (p DatabaseProvider) SaveMetaData(file models.File) {
 	helper.Check(err)
 }
 
+// IncreaseDownloadCount increases the download count of a file, preventing race conditions
+func (p DatabaseProvider) IncreaseDownloadCount(id string, decreaseRemainingDownloads bool) {
+	if decreaseRemainingDownloads {
+		_, err := p.sqliteDb.Exec(`UPDATE FileMetaData SET DownloadCount = DownloadCount + 1,
+                        DownloadsRemaining = DownloadsRemaining - 1 WHERE id = ?`, id)
+		helper.Check(err)
+	} else {
+		_, err := p.sqliteDb.Exec(`UPDATE FileMetaData SET DownloadCount = DownloadCount + 1 WHERE id = ?`, id)
+		helper.Check(err)
+	}
+}
+
 // DeleteMetaData deletes information about a file
 func (p DatabaseProvider) DeleteMetaData(id string) {
 	_, err := p.sqliteDb.Exec("DELETE FROM FileMetaData WHERE Id = ?", id)