fix(tasks): persist total to DB immediately when updated at runtime

Add Task.UpdateTotal() that writes only the total field to DB,
avoiding race conditions with Done() which updates cur independently.
Download handler now uses UpdateTotal() after discovering chunk count
so that task progress queries return accurate Cur/Total during transfer.

Author: h3zh1

server/internal/core/task.go

@@ -26,6 +26,9 @@ var (
 	taskDBUpdateCur = func(taskID string, cur int) error {
 		return db.UpdateTaskCur(taskID, cur)
 	}
+	taskDBUpdateTotal = func(taskID string, total int) error {
+		return db.UpdateTaskTotal(taskID, total)
+	}
 	taskDBUpdateFinish = func(taskID string) error {
 		return db.UpdateTaskFinish(taskID)
 	}
@@ -216,6 +219,18 @@ func (t *Task) Publish(op string, spite *implantpb.Spite, msg string) {
 		Callee:    t.Callee,
 	})
 }
+
+// UpdateTotal sets the task's total count and persists only the total field to DB.
+// This avoids racing with Done() which updates cur independently.
+func (t *Task) UpdateTotal(total int) {
+	t.progressMu.Lock()
+	t.Total = total
+	t.progressMu.Unlock()
+	if err := taskDBUpdateTotal(t.TaskID(), total); err != nil {
+		logs.Log.Warnf("task %s: update total failed: %v", t.TaskID(), err)
+	}
+}
+
 func (t *Task) Done(spite *implantpb.Spite, msg string) {
 	t.progressMu.Lock()
 	t.Cur++

server/internal/db/session_helper.go

@@ -278,6 +278,13 @@ func UpdateTaskCur(taskID string, cur int) error {
 	return taskModel.UpdateCur(Session(), cur)
 }
 
+func UpdateTaskTotal(taskID string, total int) error {
+	taskModel := &models.Task{
+		ID: taskID,
+	}
+	return taskModel.UpdateTotal(Session(), total)
+}
+
 func UpdateTaskFinish(taskID string) error {
 	task, err := NewTaskQuery().WhereID(taskID).First()
 	if err != nil {

server/rpc/rpc-file.go

@@ -307,7 +307,7 @@ func (rpc *Server) Download(ctx context.Context, req *implantpb.DownloadRequest)
 		}
 		total := downloadChunkCount(int(resp.GetDownloadResponse().Size), greq.Session.GetPacketLength())
 		downloadAbs := resp.GetDownloadResponse()
-		greq.Task.Total = total
+		greq.Task.UpdateTotal(total)
 
 		finalPath, err := fileutils.SafeJoin(configs.ContextPath, filepath.Join(greq.Session.ID, consts.DownloadPath, downloadAbs.Checksum))
 		if err != nil {