fix: address audit findings in force_rebalance code paths

pkg/backup/download.go

@@ -700,7 +700,9 @@ func (b *Backuper) downloadTableData(ctx context.Context, remoteBackup metadata.
 					}
 					if hardlinkExistsFiles {
 						ext := "." + config.ArchiveExtensions[remoteBackup.DataFormat]
-						partName := strings.TrimPrefix(strings.TrimSuffix(archiveFile, ext), capturedDisk+"_")
+						// Archive files are named with the original disk prefix (e.g. "default_part1.tar.gz"),
+						// not the rebalanced disk name, so use `disk` for the trim
+						partName := strings.TrimPrefix(strings.TrimSuffix(archiveFile, ext), disk+"_")
 						var foundPart *metadata.Part
 						var idx int
 						for i, part := range capturedParts {
@@ -967,21 +969,24 @@ func (b *Backuper) downloadDiffParts(ctx context.Context, remoteBackup metadata.
 	for disk, parts := range table.Parts {
 		diskPath, diskExists := b.DiskToPathMap[disk]
 		for i, part := range parts {
+			// Use per-iteration variable to avoid corrupting the outer loop variable
+			activeDisk := disk
+			activeDiskPath := diskPath
 			if !diskExists && part.RebalancedDisk == "" {
 				return 0, errors.Errorf("downloadDiffParts: table: `%s`.`%s`, disk: %s, part.Name: %s, part.RebalancedDisk: `%s` not rebalanced", table.Table, table.Database, disk, part.Name, part.RebalancedDisk)
 			}
 			if part.RebalancedDisk != "" {
-				diskPath, diskExists = b.DiskToPathMap[part.RebalancedDisk]
+				activeDiskPath, diskExists = b.DiskToPathMap[part.RebalancedDisk]
 				if !diskExists {
 					return 0, errors.Errorf("downloadDiffParts: table: `%s`.`%s`, disk: %s, part.Name: %s, part.RebalancedDisk: `%s` not rebalanced", table.Table, table.Database, disk, part.Name, part.RebalancedDisk)
 				}
-				disk = part.RebalancedDisk
-				if b.shouldDiskNameSkipByNameOrType(disk, disks) {
+				activeDisk = part.RebalancedDisk
+				if b.shouldDiskNameSkipByNameOrType(activeDisk, disks) {
 					log.Warn().Str("database", table.Database).Str("table", table.Table).Str("rebalancedDisk", part.RebalancedDisk).Msg("skipped")
 					continue
 				}
 			}
-			newPath := path.Join(diskPath, "backup", remoteBackup.BackupName, "shadow", dbAndTableDir, disk, part.Name)
+			newPath := path.Join(activeDiskPath, "backup", remoteBackup.BackupName, "shadow", dbAndTableDir, activeDisk, part.Name)
 			if checkErr := b.checkNewPath(newPath, part); checkErr != nil {
 				return 0, errors.WithMessage(checkErr, "checkNewPath")
 			}

pkg/filesystemhelper/filesystemhelper.go

@@ -153,12 +153,15 @@ func HardlinkBackupPartsToStorage(backupName string, backupTable metadata.TableM
 					// ClickHouse creates store/{uuid}/detached/ on ALL disks of the
 					// storage policy at table load time (MergeTreeData constructor).
 					// Build the path directly so hardlinks stay on the same filesystem.
-					if rebalancedDiskPath, hasDisk := diskMap[part.RebalancedDisk]; hasDisk && backupTable.UUID != "" {
-						dstParentDir = filepath.Join(rebalancedDiskPath, "store", backupTable.UUID[:3], backupTable.UUID)
-						dstParentDirExists = true
-					} else {
-						return errors.Errorf("can't build store path for rebalanced disk %s (uuid=%s)", part.RebalancedDisk, backupTable.UUID)
+					rebalancedDiskPath, hasDisk := diskMap[part.RebalancedDisk]
+					if !hasDisk {
+						return errors.Errorf("rebalanced disk %s not found in diskMap", part.RebalancedDisk)
 					}
+					if backupTable.UUID == "" {
+						return errors.Errorf("table UUID is empty, can't build store path for rebalanced disk %s", part.RebalancedDisk)
+					}
+					dstParentDir = filepath.Join(rebalancedDiskPath, "store", backupTable.UUID[:3], backupTable.UUID)
+					dstParentDirExists = true
 				}
 			}
 			backupDiskPath := diskMap[activeDisk]

pkg/metadata/table_metadata.go

@@ -42,6 +42,7 @@ func (tm *TableMetadata) Save(location string, metadataOnly bool) (uint64, error
 		newTM.Checksums = tm.Checksums
 		newTM.Size = tm.Size
 		newTM.TotalBytes = tm.TotalBytes
+		newTM.RebalancedFiles = tm.RebalancedFiles
 		newTM.MetadataOnly = false
 	}
 	if err := os.MkdirAll(path.Dir(location), 0750); err != nil {