Update download count in real time #206

internal/models/UploadStatus.go

@@ -1,9 +1,5 @@
 package models
 
-import (
-	"encoding/json"
-)
-
 // UploadStatus contains information about the current status of a file upload
 type UploadStatus struct {
 	// ChunkId is the identifier for the chunk
@@ -13,9 +9,3 @@ type UploadStatus struct {
 	// See processingstatus for definition
 	CurrentStatus int `json:"currentstatus"`
 }
-
-// ToJson returns the struct as a Json byte array
-func (u *UploadStatus) ToJson() ([]byte, error) {
-	return json.Marshal(u)
-
-}

internal/storage/FileServing.go

@@ -22,6 +22,7 @@ import (
 	"github.com/forceu/gokapi/internal/storage/processingstatus"
 	"github.com/forceu/gokapi/internal/webserver/downloadstatus"
 	"github.com/forceu/gokapi/internal/webserver/headers"
+	"github.com/forceu/gokapi/internal/webserver/sse"
 	"github.com/jinzhu/copier"
 	"io"
 	"log"
@@ -524,6 +525,7 @@ func ServeFile(file models.File, w http.ResponseWriter, r *http.Request, forceDo
 	file.DownloadCount = file.DownloadCount + 1
 	database.SaveMetaData(file)
 	logging.AddDownload(&file, r, configuration.Get().SaveIp)
+	go sse.PublishDownloadCount(file)
 
 	if !file.IsLocalStorage() {
 		// If non-blocking, we are not setting a download complete status as there is no reliable way to

internal/storage/processingstatus/ProcessingStatus.go

@@ -2,7 +2,6 @@ package processingstatus
 
 import (
 	"github.com/forceu/gokapi/internal/configuration/database"
-	"github.com/forceu/gokapi/internal/helper"
 	"github.com/forceu/gokapi/internal/models"
 	"github.com/forceu/gokapi/internal/webserver/sse"
 )
@@ -13,12 +12,6 @@ const StatusHashingOrEncrypting = 0
 // StatusUploading indicates that the file has been processed, but is now moved to the data filesystem
 const StatusUploading = 1
 
-func passNewStatus(newStatus models.UploadStatus) {
-	status, err := newStatus.ToJson()
-	helper.Check(err)
-	sse.PublishNewStatus(string(status))
-}
-
 // Set sets the status for an id
 func Set(id string, status int) {
 	newStatus := models.UploadStatus{
@@ -30,5 +23,5 @@ func Set(id string, status int) {
 		return
 	}
 	database.SaveUploadStatus(newStatus)
-	go passNewStatus(newStatus)
+	go sse.PublishNewStatus(newStatus)
 }

internal/webserver/sse/Sse.go

@@ -1,8 +1,10 @@
 package sse
 
 import (
+	"encoding/json"
 	"github.com/forceu/gokapi/internal/configuration/database"
 	"github.com/forceu/gokapi/internal/helper"
+	"github.com/forceu/gokapi/internal/models"
 	"io"
 	"net/http"
 	"sync"
@@ -32,14 +34,50 @@ func removeListener(id string) {
 	mutex.Unlock()
 }
 
-func PublishNewStatus(reply string) {
+type eventFileDownload struct {
+	Event         string `json:"event"`
+	FileId        string `json:"file_id"`
+	DownloadCount int    `json:"download_count"`
+}
+type eventUploadStatus struct {
+	Event        string `json:"event"`
+	ChunkId      string `json:"chunk_id"`
+	UploadStatus int    `json:"upload_status"`
+}
+
+type eventData interface {
+	eventUploadStatus | eventFileDownload
+}
+
+func PublishNewStatus(uploadStatus models.UploadStatus) {
+	event := eventUploadStatus{
+		Event:        "uploadStatus",
+		ChunkId:      uploadStatus.ChunkId,
+		UploadStatus: uploadStatus.CurrentStatus,
+	}
+	publishMessage(event)
+}
+
+func publishMessage[d eventData](data d) {
+	message, err := json.Marshal(data)
+	helper.Check(err)
+
 	mutex.RLock()
 	for _, channel := range listeners {
-		go channel.Reply("event: message\ndata: " + reply + "\n\n")
+		go channel.Reply("event: message\ndata: " + string(message) + "\n\n")
 	}
 	mutex.RUnlock()
 }
 
+func PublishDownloadCount(file models.File) {
+	event := eventFileDownload{
+		Event:         "download",
+		FileId:        file.Id,
+		DownloadCount: file.DownloadCount,
+	}
+	publishMessage(event)
+}
+
 func Shutdown() {
 	mutex.RLock()
 	for _, channel := range listeners {
@@ -69,9 +107,7 @@ func GetStatusSSE(w http.ResponseWriter, r *http.Request) {
 
 	allStatus := database.GetAllUploadStatus()
 	for _, status := range allStatus {
-		jsonOutput, err := status.ToJson()
-		helper.Check(err)
-		_, _ = io.WriteString(w, "event: message\ndata: "+string(jsonOutput)+"\n\n")
+		PublishNewStatus(status)
 	}
 	w.(http.Flusher).Flush()
 	for {

internal/webserver/sse/Sse_test.go

@@ -2,6 +2,7 @@ package sse
 
 import (
 	"github.com/forceu/gokapi/internal/configuration"
+	"github.com/forceu/gokapi/internal/models"
 	"github.com/forceu/gokapi/internal/test"
 	"github.com/forceu/gokapi/internal/test/testconfiguration"
 	"io"
@@ -45,10 +46,20 @@ func TestPublishNewStatus(t *testing.T) {
 	channel := listener{Reply: func(reply string) { replyChannel <- reply }, Shutdown: func() {}}
 	addListener("test_id", channel)
 
-	go PublishNewStatus("test_status")
+	go PublishNewStatus(models.UploadStatus{
+		ChunkId:       "testChunkId",
+		CurrentStatus: 4,
+	})
 	receivedStatus := <-replyChannel
-	test.IsEqualString(t, receivedStatus, "event: message\ndata: test_status\n\n")
-	removeListener("test_status")
+	test.IsEqualString(t, receivedStatus, "event: message\ndata: {\"event\":\"uploadStatus\",\"chunk_id\":\"testChunkId\",\"upload_status\":4}\n\n")
+
+	go PublishDownloadCount(models.File{
+		Id:            "testFileId",
+		DownloadCount: 3,
+	})
+	receivedStatus = <-replyChannel
+	test.IsEqualString(t, receivedStatus, "event: message\ndata: {\"event\":\"download\",\"file_id\":\"testFileId\",\"download_count\":3}\n\n")
+	removeListener("test_id")
 }
 
 func TestShutdown(t *testing.T) {
@@ -89,19 +100,22 @@ func TestGetStatusSSE(t *testing.T) {
 	body, err := io.ReadAll(rr.Body)
 	test.IsNil(t, err)
 
-	test.IsEqualString(t, string(body), "event: message\ndata: {\"chunkid\":\"validstatus_0\",\"currentstatus\":0}\n\n"+
-		"event: message\ndata: {\"chunkid\":\"validstatus_1\",\"currentstatus\":1}\n\n")
+	test.IsEqualString(t, string(body), "event: message\ndata: {\"event\":\"uploadStatus\",\"chunk_id\":\"validstatus_0\",\"upload_status\":0}\n\n"+
+		"event: message\ndata: {\"event\":\"uploadStatus\",\"chunk_id\":\"validstatus_1\",\"upload_status\":1}\n\n")
 
 	// Test ping message
 	time.Sleep(3 * time.Second)
 	body, err = io.ReadAll(rr.Body)
 	test.IsNil(t, err)
 	test.IsEqualString(t, string(body), "event: ping\n\n")
 
-	PublishNewStatus("testcontent")
+	PublishNewStatus(models.UploadStatus{
+		ChunkId:       "secondChunkId",
+		CurrentStatus: 1,
+	})
 	time.Sleep(1 * time.Second)
 	body, err = io.ReadAll(rr.Body)
 	test.IsNil(t, err)
-	test.IsEqualString(t, string(body), "event: message\ndata: testcontent\n\n")
+	test.IsEqualString(t, string(body), "event: message\ndata: {\"event\":\"uploadStatus\",\"chunk_id\":\"secondChunkId\",\"upload_status\":1}\n\n")
 	Shutdown()
 }

internal/webserver/web/static/js/admin.js

@@ -576,7 +576,7 @@ function checkBoxChanged(checkBox, correspondingInput) {
     }
 }
 
-function parseData(data) {
+function parseRowData(data) {
     if (!data) return {
         "Result": "error"
     };
@@ -588,15 +588,39 @@ function parseData(data) {
     };
 }
 
+
+function parseSseData(data) {
+    let eventData;
+    try {
+        eventData = JSON.parse(data);
+    } catch (e) {
+        console.error("Failed to parse event data:", e);
+        return;
+    }
+    switch (eventData.event) {
+        case "download":
+            setNewDownloadCount(eventData.file_id, eventData.download_count);
+            return;
+        case "uploadStatus":
+            setProgressStatus(eventData.chunk_id, eventData.upload_status);
+            return;
+        default:
+            console.error("Unknown event", eventData);
+    }
+}
+
+function setNewDownloadCount(id, downloadCount) {
+    let downloadCell = document.getElementById("cell-downloads-" + id);
+    if (downloadCell != null) {
+        downloadCell.innerHTML = downloadCount;
+    }
+}
+
+
 function registerChangeHandler() {
     const source = new EventSource("./uploadStatus")
     source.onmessage = (event) => {
-        try {
-            let eventData = JSON.parse(event.data);
-            setProgressStatus(eventData.chunkid, eventData.currentstatus);
-        } catch (e) {
-            console.error("Failed to parse event data:", e);
-        }
+        parseSseData(event.data);
     }
     source.onerror = (error) => {
 
@@ -676,7 +700,7 @@ function removeFileStatus(chunkId) {
 
 
 function addRow(jsonText) {
-    let jsonObject = parseData(jsonText);
+    let jsonObject = parseRowData(jsonText);
     if (jsonObject.Result !== "OK") {
         alert("Failed to upload file!");
         location.reload();
@@ -700,6 +724,7 @@ function addRow(jsonText) {
     }
     cellFilename.innerText = item.Name;
     cellFilename.id = "cell-name-" + item.Id;
+    cellDownloadCount.id = "cell-downloads-" + item.Id;
     cellFileSize.innerText = item.Size;
     if (item.UnlimitedDownloads) {
         cellRemainingDownloads.innerText = "Unlimited";

internal/webserver/web/templates/html_admin.tmpl

@@ -77,7 +77,7 @@
 			{{ else }}
 						<td>{{ .ExpireAtString }}</td>
 			{{ end }}
-						<td>{{ .DownloadCount }}</td>
+						<td id="cell-downloads-{{ .Id }}">{{ .DownloadCount }}</td>
 						<td><a id="url-href-{{ .Id }}" target="_blank" href="{{ .UrlDownload }}">{{ .Id }}</a>{{ if .IsPasswordProtected }}  <i title="Password protected" class="bi bi-key"></i>{{ end }}</td>
 						<td><button id="url-button-{{ .Id }}" type="button" onclick="showToast()" data-clipboard-text="{{ .UrlDownload }}" class="copyurl btn btn-outline-light btn-sm"><i class="bi bi-copy"></i> URL</button> 
 			{{ if ne .UrlHotlink "" }}