diff --git a/agent/internal/filetransfer/transfer.go b/agent/internal/filetransfer/transfer.go index f93f392cdb..83d0926eac 100644 --- a/agent/internal/filetransfer/transfer.go +++ b/agent/internal/filetransfer/transfer.go @@ -99,6 +99,17 @@ func (m *Manager) HandleTransfer(payload map[string]any) map[string]any { m.transfers[transferID] = transfer m.mu.Unlock() + // The entry only needs to live while the transfer is in flight (it is how + // CancelTransfer reaches the shared Transfer). Nothing reads a transfer + // back after it reaches a terminal state, so always remove it on return — + // otherwise the map grows for the life of the agent process (#2388). A + // late cancel after removal hits CancelTransfer's not-found no-op branch. + defer func() { + m.mu.Lock() + delete(m.transfers, transferID) + m.mu.Unlock() + }() + // Process transfer var err error if direction == "upload" { @@ -108,8 +119,10 @@ func (m *Manager) HandleTransfer(payload map[string]any) map[string]any { } if err != nil { + m.mu.Lock() transfer.Status = "failed" transfer.Error = err.Error() + m.mu.Unlock() m.reportProgress(transfer) return map[string]any{ "status": "failed", @@ -117,8 +130,10 @@ func (m *Manager) HandleTransfer(payload map[string]any) map[string]any { } } + m.mu.Lock() transfer.Status = "completed" transfer.Progress = 100 + m.mu.Unlock() m.reportProgress(transfer) return map[string]any{ @@ -184,7 +199,9 @@ func (m *Manager) upload(transfer *Transfer) error { } uploaded += int64(n) + m.mu.Lock() transfer.Progress = int((uploaded * 100) / totalSize) + m.mu.Unlock() m.reportProgress(transfer) chunkNum++ } @@ -288,7 +305,9 @@ func (m *Manager) download(transfer *Transfer) error { } downloaded += int64(n) if totalSize > 0 { + m.mu.Lock() transfer.Progress = int((downloaded * 100) / totalSize) + m.mu.Unlock() m.reportProgress(transfer) } } @@ -304,12 +323,16 @@ func (m *Manager) download(transfer *Transfer) error { } func (m *Manager) reportProgress(transfer *Transfer) { + // Snapshot mutable fields under the lock — CancelTransfer may flip Status + // on the shared pointer concurrently. + m.mu.RLock() data := map[string]any{ "transferId": transfer.ID, "status": transfer.Status, "progress": transfer.Progress, "error": transfer.Error, } + m.mu.RUnlock() body, err := json.Marshal(data) if err != nil { diff --git a/agent/internal/filetransfer/transfer_download_test.go b/agent/internal/filetransfer/transfer_download_test.go index cf3203c85c..19c87d4b9f 100644 --- a/agent/internal/filetransfer/transfer_download_test.go +++ b/agent/internal/filetransfer/transfer_download_test.go @@ -62,6 +62,7 @@ func TestHandleTransferDownloadSuccess(t *testing.T) { if string(data) != body { t.Fatalf("expected %q, got %q", body, string(data)) } + assertNoTransfers(t, m) } func TestHandleTransferDownloadDirectoryTraversal(t *testing.T) { @@ -117,6 +118,7 @@ func TestHandleTransferDownloadServerError(t *testing.T) { if result["status"] != "failed" { t.Fatalf("expected failed, got %v", result["status"]) } + assertNoTransfers(t, m) } // ---------- CancelTransfer ---------- @@ -144,6 +146,84 @@ func TestCancelTransferSetsStatus(t *testing.T) { } } +// TestCancelTransferMidFlight cancels through the real HandleTransfer path +// while the first chunk upload is blocked in the server handler, then keeps +// hammering CancelTransfer concurrently while the transfer winds down so the +// mutex guards on Status/Progress/Error and the reportProgress snapshot are +// genuinely exercised under -race. +func TestCancelTransferMidFlight(t *testing.T) { + firstChunk := make(chan struct{}) + release := make(chan struct{}) + var once sync.Once + + ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method == "POST" { + once.Do(func() { + close(firstChunk) + <-release + }) + } + w.WriteHeader(http.StatusOK) + })) + defer ts.Close() + + tmpDir := t.TempDir() + srcFile := filepath.Join(tmpDir, "big.dat") + // Two chunks, so a cancellation check runs after the first chunk uploads. + if err := os.WriteFile(srcFile, make([]byte, ChunkSize+100), 0644); err != nil { + t.Fatal(err) + } + + cfg := &Config{ + ServerURL: ts.URL, + AuthToken: secmem.NewSecureString("tok"), + AgentID: "agent-1", + } + m := NewManager(cfg) + + done := make(chan map[string]any, 1) + go func() { + done <- m.HandleTransfer(map[string]any{ + "transferId": "t-cancel-mid", + "direction": "upload", + "remotePath": "/remote/dest.dat", + "localPath": srcFile, + }) + }() + + <-firstChunk + m.CancelTransfer("t-cancel-mid") + + // Race the manager's terminal-state writes against concurrent cancels. + stopCancel := make(chan struct{}) + var cancelWg sync.WaitGroup + cancelWg.Add(1) + go func() { + defer cancelWg.Done() + for { + select { + case <-stopCancel: + return + default: + m.CancelTransfer("t-cancel-mid") + } + } + }() + + close(release) + result := <-done + close(stopCancel) + cancelWg.Wait() + + if result["status"] != "failed" { + t.Fatalf("expected failed after mid-flight cancel, got %v", result["status"]) + } + if result["error"] != "transfer cancelled" { + t.Fatalf("expected 'transfer cancelled', got %v", result["error"]) + } + assertNoTransfers(t, m) +} + func TestCancelTransferNonexistentIsNoop(t *testing.T) { cfg := &Config{ ServerURL: "https://example.com", @@ -189,6 +269,7 @@ func TestConcurrentHandleTransfer(t *testing.T) { }(i) } wg.Wait() + assertNoTransfers(t, m) } // ---------- reportProgress ---------- diff --git a/agent/internal/filetransfer/transfer_test.go b/agent/internal/filetransfer/transfer_test.go index 5bb5b6bb4c..725ec0b2b4 100644 --- a/agent/internal/filetransfer/transfer_test.go +++ b/agent/internal/filetransfer/transfer_test.go @@ -10,6 +10,18 @@ import ( "github.com/breeze-rmm/agent/internal/secmem" ) +// assertNoTransfers verifies the transfers map is empty — HandleTransfer must +// remove its entry on every terminal path (#2388). +func assertNoTransfers(t *testing.T, m *Manager) { + t.Helper() + m.mu.RLock() + n := len(m.transfers) + m.mu.RUnlock() + if n != 0 { + t.Fatalf("expected transfers map to be empty after HandleTransfer, got %d entries", n) + } +} + // ---------- NewManager ---------- func TestNewManagerInitializesFields(t *testing.T) { @@ -122,6 +134,7 @@ func TestHandleTransferUploadSuccess(t *testing.T) { if progressReports == 0 { t.Fatal("expected at least one progress report") } + assertNoTransfers(t, m) } func TestHandleTransferUploadDirectoryTraversal(t *testing.T) { @@ -213,6 +226,7 @@ func TestHandleTransferUploadChunkServerError(t *testing.T) { if result["status"] != "failed" { t.Fatalf("expected failed, got %v", result["status"]) } + assertNoTransfers(t, m) } // ---------- Multi-chunk upload ----------