fix(agent): remove filetransfer map entries when HandleTransfer returns (#2388)

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,17 +119,21 @@ 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",
 			"error":  err.Error(),
 		}
 	}
 
+	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 {

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 ----------

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 ----------