Fixed race condition that could lead to crash #222

internal/storage/FileServing.go

@@ -716,11 +716,7 @@ func DeleteFile(keyId string, deleteSource bool) bool {
 	item.ExpireAt = 0
 	item.UnlimitedTime = false
 	database.SaveMetaData(item)
-	for _, status := range downloadstatus.GetAll() {
-		if status.FileId == item.Id {
-			downloadstatus.SetComplete(status.Id)
-		}
-	}
+	downloadstatus.SetAllComplete(item.Id)
 	if deleteSource {
 		go CleanUp(false)
 	}

internal/webserver/downloadstatus/DownloadStatus.go

@@ -3,29 +3,35 @@ package downloadstatus
 import (
 	"github.com/forceu/gokapi/internal/helper"
 	"github.com/forceu/gokapi/internal/models"
+	"sync"
 	"time"
 )
 
-var status = make(map[string]models.DownloadStatus)
+var statusMap = make(map[string]models.DownloadStatus)
+var statusMutex sync.RWMutex
 
 // SetDownload creates a new DownloadStatus struct and returns its Id
 func SetDownload(file models.File) string {
 	newStatus := newDownloadStatus(file)
-	status[newStatus.Id] = newStatus
+	statusMutex.Lock()
+	statusMap[newStatus.Id] = newStatus
+	statusMutex.Unlock()
 	return newStatus.Id
 }
 
 // SetComplete removes the download object
-func SetComplete(id string) {
-	delete(status, id)
+func SetComplete(downloadStatusId string) {
+	statusMutex.Lock()
+	delete(statusMap, downloadStatusId)
+	statusMutex.Unlock()
 }
 
 // Clean removes all expires status objects
 func Clean() {
 	now := time.Now().Unix()
-	for _, item := range status {
+	for _, item := range statusMap {
 		if item.ExpireAt < now {
-			delete(status, item.Id)
+			SetComplete(item.Id)
 		}
 	}
 }
@@ -42,22 +48,33 @@ func newDownloadStatus(file models.File) models.DownloadStatus {
 
 // IsCurrentlyDownloading returns true if file is currently being downloaded
 func IsCurrentlyDownloading(file models.File) bool {
-	for _, statusField := range status {
-		if statusField.FileId == file.Id {
-			if statusField.ExpireAt > time.Now().Unix() {
-				return true
+	isDownloading := false
+	statusMutex.RLock()
+	for _, status := range statusMap {
+		if status.FileId == file.Id {
+			if status.ExpireAt > time.Now().Unix() {
+				isDownloading = true
+				break
 			}
 		}
 	}
-	return false
+	statusMutex.RUnlock()
+	return isDownloading
 }
 
-// GetAll returns all download states
-func GetAll() map[string]models.DownloadStatus {
-	return status
+func SetAllComplete(fileId string) {
+	statusMutex.Lock()
+	for _, status := range statusMap {
+		if status.FileId == fileId {
+			delete(statusMap, status.Id)
+		}
+	}
+	statusMutex.Unlock()
 }
 
 // DeleteAll removes all download status
 func DeleteAll() {
-	status = make(map[string]models.DownloadStatus)
+	statusMutex.Lock()
+	statusMap = make(map[string]models.DownloadStatus)
+	statusMutex.Unlock()
 }

internal/webserver/downloadstatus/DownloadStatus_test.go

@@ -34,47 +34,73 @@ func TestNewDownloadStatus(t *testing.T) {
 
 func TestSetDownload(t *testing.T) {
 	statusId = SetDownload(testFile)
-	newStatus := status[statusId]
+	newStatus := statusMap[statusId]
 	test.IsNotEmpty(t, newStatus.Id)
 	test.IsEqualString(t, newStatus.Id, statusId)
 	test.IsEqualString(t, newStatus.FileId, testFile.Id)
 	test.IsEqualBool(t, newStatus.ExpireAt > time.Now().Unix(), true)
 }
 
 func TestSetComplete(t *testing.T) {
-	newStatus := status[statusId]
+	newStatus := statusMap[statusId]
 	test.IsNotEmpty(t, newStatus.Id)
 	SetComplete(statusId)
-	newStatus = status[statusId]
+	newStatus = statusMap[statusId]
 	test.IsEmpty(t, newStatus.Id)
 }
 
 func TestIsCurrentlyDownloading(t *testing.T) {
-	statusId = SetDownload(testFile)
+	test.IsEqualBool(t, IsCurrentlyDownloading(testFile), false)
+	statusIdFirst := SetDownload(testFile)
+	firstStatus := statusMap[statusIdFirst]
+	test.IsEqualBool(t, IsCurrentlyDownloading(testFile), true)
+	statusIdSecond := SetDownload(testFile)
+	secondStatus := statusMap[statusIdSecond]
 	test.IsEqualBool(t, IsCurrentlyDownloading(testFile), true)
+
+	firstStatus.ExpireAt = 0
+	statusMap[firstStatus.Id] = firstStatus
+	test.IsEqualBool(t, IsCurrentlyDownloading(testFile), true)
+	secondStatus.ExpireAt = 0
+	statusMap[secondStatus.Id] = secondStatus
+	test.IsEqualBool(t, IsCurrentlyDownloading(testFile), false)
+
+	statusId = SetDownload(testFile)
 	test.IsEqualBool(t, IsCurrentlyDownloading(models.File{Id: "notDownloading"}), false)
 }
 func TestClean(t *testing.T) {
-	test.IsEqualInt(t, len(status), 1)
+	test.IsEqualInt(t, len(statusMap), 3)
 	Clean()
-	test.IsEqualInt(t, len(status), 1)
-	newStatus := status[statusId]
+	test.IsEqualInt(t, len(statusMap), 1)
+	newStatus := statusMap[statusId]
 	newStatus.ExpireAt = 1
-	status[statusId] = newStatus
-	test.IsEqualInt(t, len(status), 1)
+	statusMap[statusId] = newStatus
+	test.IsEqualInt(t, len(statusMap), 1)
 	Clean()
-	test.IsEqualInt(t, len(status), 0)
-}
-
-func TestGetAll(t *testing.T) {
-	statusId = SetDownload(testFile)
-	test.IsEqualInt(t, len(GetAll()), len(status))
+	test.IsEqualInt(t, len(statusMap), 0)
 }
 
 func TestDeleteAll(t *testing.T) {
 	statusId = SetDownload(testFile)
-	test.IsEqualBool(t, len(GetAll()) != 0, true)
+	test.IsEqualBool(t, len(statusMap) != 0, true)
 	DeleteAll()
-	test.IsEqualInt(t, len(GetAll()), 0)
+	test.IsEqualInt(t, len(statusMap), 0)
+}
 
+func TestSetAllComplete(t *testing.T) {
+	test.IsEqualInt(t, len(statusMap), 0)
+	SetDownload(models.File{Id: "stillDownloading"})
+	SetDownload(models.File{Id: "stillDownloading"})
+	status1 := SetDownload(models.File{Id: "stillDownloading"})
+	SetDownload(models.File{Id: "fileToBeDeleted"})
+	SetDownload(models.File{Id: "fileToBeDeleted"})
+	SetDownload(models.File{Id: "fileToBeDeleted"})
+	status2 := SetDownload(models.File{Id: "fileToBeDeleted"})
+	test.IsEqualInt(t, len(statusMap), 7)
+	SetAllComplete("fileToBeDeleted")
+	test.IsEqualInt(t, len(statusMap), 3)
+	_, ok := statusMap[status1]
+	test.IsEqualBool(t, ok, true)
+	_, ok = statusMap[status2]
+	test.IsEqualBool(t, ok, false)
 }