Skip to content

Commit

Permalink
Merge pull request #300 from balena-io-modules/refine_migrator
Browse files Browse the repository at this point in the history
Refine migrator internals
  • Loading branch information
flowzone-app[bot] authored Jun 22, 2023
2 parents 18ffec9 + 77f3a3c commit bc260b2
Show file tree
Hide file tree
Showing 2 changed files with 106 additions and 65 deletions.
51 changes: 35 additions & 16 deletions lib/diskpart.ts
Original file line number Diff line number Diff line change
Expand Up @@ -267,7 +267,7 @@ export const setPartitionOnlineStatus = async (
*
* @param {string} device - device path
* @param {string} label - volume/partition label
* @return {number} identifier the volume, or '' if not found
* @return {number} identifier for the first matching volume, or '' if not found
* @example
* findVolume('\\\\.\\PhysicalDrive0', 'flash-boot')
* .then(...)
Expand All @@ -281,13 +281,15 @@ export const findVolume = async (

/* Retrieves diskpart output formatted like the example below.
*
* Volume ### Ltr Label Fs Type Size Status Info
* ---------- --- ----------- ----- ---------- ------- --------- --------
* Volume 0 C NTFS Partition 45 GB Healthy Boot
* Volume 1 flash-boot FAT Partition 41 MB Healthy
* Volume 2 RAW Partition 3793 MB Healthy
* Volume 3 FAT32 Partition 100 MB Healthy System
* Volume 4 NTFS Partition 530 MB Healthy Hidden
* DISKPART> list volume
*
* Volume ### Ltr Label Fs Type Size Status Info
* ---------- --- ----------- ----- ---------- ------- --------- --------
* Volume 0 C NTFS Partition 45 GB Healthy Boot
* Volume 1 flash-boot FAT Partition 41 MB Healthy
* Volume 2 RAW Partition 3793 MB Healthy
* Volume 3 FAT32 Partition 100 MB Healthy System
* Volume 4 NTFS Partition 530 MB Healthy Hidden
*/
if (platform() !== 'win32') {
throw new Error("findVolume() not available on this platform")
Expand All @@ -303,16 +305,33 @@ export const findVolume = async (
throw(`findVolume: ${error}${error.stdout ? `\n${error.stdout}` : ''}`);
}

let labelPos = -1;
// Look for 'Label' in column headings; then compare text on subsequent rows
// at that position for the expected label.
// Search for label in a language independent way, based on columns.
enum Columns { Volume = 0, Letter, Label, Fs }
let colOffsets:number[] = []

for (let line of listText.split('\n')) {
if (labelPos < 0) {
labelPos = line.indexOf('Label');
if (!colOffsets.length && line.indexOf('-----') >= 0) {
// Collect the line position for each column into colOffsets.
let linePos = 0
for (let i = 0; i <= Columns.Fs; i++) {
// Expecting even first column preceded by space(s).
linePos = line.indexOf(' -', linePos)
if (linePos == -1) {
throw Error(`findVolume: Only found ${i} columns.`)
}
linePos += 1 // advance past space
colOffsets.push(linePos)
}
debug(`colOffsets: ${colOffsets}`)
} else {
const volMatch = line.match(/Volume\s+(\d+)/);
if (volMatch && (line.substring(labelPos, labelPos + label.length) == label)) {
return volMatch[1]
// Look for the label in the expected column and collect the volume ID if found.
if (line.substring(colOffsets[Columns.Label], colOffsets[Columns.Fs]).trim() == label) {
const volText = line.substring(colOffsets[Columns.Volume],colOffsets[Columns.Letter])
// Assumes only a space before the number at the end of the field.
const volMatch = volText.match(/\s+(\d+)\s*$/)
if (volMatch && volMatch[1]) {
return volMatch[1]
}
}
}
}
Expand Down
120 changes: 71 additions & 49 deletions lib/migrator/index.ts
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@ import * as checkDiskSpace from 'check-disk-space';
import { GPTPartition, MBRPartition } from 'partitioninfo';

import * as diskpart from '../diskpart';
import { File } from '../source-destination';
import { File, BlockDevice } from '../source-destination';
import {
copyPartitionFromImageToDevice,
calcRequiredPartitionSize,
Expand Down Expand Up @@ -40,6 +40,25 @@ function formatMB(bytes: number): string {
return (bytes / (1024 * 1024)).toFixed(2)
}

/**
* A storage device that is the target for the migration.
* Collects useful attributes for the device from multiple sources.
*/
interface TargetDevice {
/** (Windows) name for the device, like '\\.\PhysicalDrive0' */
name: string;
/** Etcher BlockDevice representation */
etcher: BlockDevice
}

/** A partition on the target device. */
interface TargetPartition {
/** Etcher representation */
etcher?: GPTPartition | MBRPartition | null;
/** Identifier for Windows volume */
volumeId: string
}

/** Options for migrate(): */
export interface MigrateOptions {
// don't perform these tasks; comma separated list like 'bootloader,reboot'
Expand Down Expand Up @@ -103,15 +122,18 @@ export const migrate = async (

// Define objects for image file source for partitions, storage device target,
// and the target's partition table.
const source = new File({ path: imagePath });
const targetDevice = await getTargetBlockDevice(windowsPartition)
let currentPartitions = await targetDevice.getPartitionTable()
if (currentPartitions === undefined) {
const sourceFile = new File({ path: imagePath });
const targetDevice:TargetDevice = {
name: deviceName,
etcher: await getTargetBlockDevice(windowsPartition)
}
let etcherPartitions = await targetDevice.etcher.getPartitionTable()
if (etcherPartitions === undefined) {
throw Error("Can't read partition table");
}
// Log existing partitions for debugging
console.log("\nPartitions on target:")
for (const p of currentPartitions.partitions) {
for (const p of etcherPartitions.partitions) {
// Satisfy TypeScript that p is not an MBRPartition even though we tested above on the table
if (!('guid' in p)) {
continue
Expand All @@ -121,31 +143,39 @@ export const migrate = async (

// Prepare to check for the balenaOS boot and rootA partitions already present.
// If partitions not present, determine required partition sizes and free space.
// Calculations are in units of bytes. However, on Windows, required sizes are
// rounded up to the nearest MB due to tool limitations.
let targetBootPartition: GPTPartition | MBRPartition | null
let targetRootAPartition: GPTPartition | MBRPartition | null
// Calculations are in units of bytes. However, on Windows, required sizes are in MB.
let bootPartition:TargetPartition = {
volumeId: '',
}
let rootAPartition:TargetPartition = {
volumeId: '',
}
// Values are rounded up to the nearest MB due to tool limitations.
let requiredBootSize = 0
let requiredRootASize = 0

// Look for boot partition on a FAT16 filesystem
targetBootPartition = await findFilesystemLabel(currentPartitions, targetDevice,
bootPartition.etcher = await findFilesystemLabel(etcherPartitions, targetDevice.etcher,
BOOT_PARTITION_LABEL, 'fat16')
if (targetBootPartition) {
console.log(`Boot partition already exists at index ${targetBootPartition.index}`)
if (bootPartition.etcher) {
console.log(`Boot partition already exists at index ${bootPartition.etcher.index}`)
bootPartition.volumeId = await diskpart.findVolume(targetDevice.name, BOOT_PARTITION_LABEL)
console.log(`flasherBootPartition volume: ${bootPartition.volumeId}`)
} else {
console.log("Boot partition not found on target")
requiredBootSize = await calcRequiredPartitionSize(source, BOOT_PARTITION_INDEX);
requiredBootSize = await calcRequiredPartitionSize(sourceFile, BOOT_PARTITION_INDEX);
console.log(`Require ${requiredBootSize} (${formatMB(requiredBootSize)} MB) for boot partition`);
}
// Look for rootA partition on an ext4 filesystem
targetRootAPartition = await findFilesystemLabel(currentPartitions, targetDevice,
rootAPartition.etcher = await findFilesystemLabel(etcherPartitions, targetDevice.etcher,
ROOTA_PARTITION_LABEL, 'ext4')
if (targetRootAPartition) {
console.log(`RootA partition already exists at index ${targetRootAPartition.index}`)
if (rootAPartition.etcher) {
console.log(`RootA partition already exists at index ${rootAPartition.etcher.index}`)
rootAPartition.volumeId = await diskpart.findVolume(targetDevice.name, ROOTA_PARTITION_LABEL)
console.log(`flasherRootAPartition volume: ${rootAPartition.volumeId}`)
} else {
console.log("RootA partition not found on target")
requiredRootASize = await calcRequiredPartitionSize(source, ROOTA_PARTITION_INDEX);
requiredRootASize = await calcRequiredPartitionSize(sourceFile, ROOTA_PARTITION_INDEX);
console.log(`Require ${requiredRootASize} (${formatMB(requiredRootASize)} MB) for rootA partition`)
}
const requiredFreeSize = requiredBootSize + requiredRootASize;
Expand All @@ -154,8 +184,8 @@ export const migrate = async (
// Shrink amount must be for *all* of required space to ensure it is contiguous.
// IOW, don't assume the shrink will merge with any existing unallocated space.
if (requiredFreeSize) {
const unallocSpace = (await diskpart.getUnallocatedSize(deviceName)) * 1024;
console.log(`Found ${unallocSpace} (${formatMB(unallocSpace)} MB) not allocated on disk ${deviceName}`)
const unallocSpace = (await diskpart.getUnallocatedSize(targetDevice.name)) * 1024;
console.log(`Found ${unallocSpace} (${formatMB(unallocSpace)} MB) not allocated on disk ${targetDevice.name}`)

if (unallocSpace < requiredFreeSize) {
// must force upper case
Expand All @@ -177,52 +207,44 @@ export const migrate = async (
if (tasks.includes('copy')) {
// create partitions
console.log("") //force newline
let volumeIds = ['', '']
if (!targetBootPartition) {
if (!bootPartition.etcher) {
console.log("Create flasherBootPartition");
await diskpart.createPartition(deviceName, requiredBootSize / (1024 * 1024));
const afterFirstPartitions = await targetDevice.getPartitionTable()
const firstNewPartition = findNewPartitions(currentPartitions, afterFirstPartitions);
await diskpart.createPartition(targetDevice.name, requiredBootSize / (1024 * 1024));
const afterFirstPartitions = await targetDevice.etcher.getPartitionTable()
const firstNewPartition = findNewPartitions(etcherPartitions, afterFirstPartitions);
if (firstNewPartition.length !== 1) {
throw Error(`Found ${firstNewPartition.length} new partitions for flasher boot, but expected 1`)
}
targetBootPartition = firstNewPartition[0];
console.log(`Created new partition for boot at offset ${targetBootPartition.offset} with size ${targetBootPartition.size}`);
currentPartitions = afterFirstPartitions
bootPartition.etcher = firstNewPartition[0];
console.log(`Created new partition for boot at offset ${bootPartition.etcher.offset} with size ${bootPartition.etcher.size}`);
etcherPartitions = afterFirstPartitions
}
volumeIds[0] = await diskpart.findVolume(deviceName, BOOT_PARTITION_LABEL)
console.log(`flasherBootPartition volume: ${volumeIds[0]}`)

if (!targetRootAPartition) {
if (!rootAPartition.etcher) {
console.log("Create flasherRootAPartition");
await diskpart.createPartition(deviceName, requiredRootASize / (1024 * 1024));
const afterSecondPartitions = await targetDevice.getPartitionTable()
const secondNewPartition = findNewPartitions(currentPartitions, afterSecondPartitions)
await diskpart.createPartition(targetDevice.name, requiredRootASize / (1024 * 1024));
const afterSecondPartitions = await targetDevice.etcher.getPartitionTable()
const secondNewPartition = findNewPartitions(etcherPartitions, afterSecondPartitions)
if (secondNewPartition.length !== 1) {
throw Error(`Found ${secondNewPartition.length} new partitions for flasher rootA, but expected 1`)
}
targetRootAPartition = secondNewPartition[0];
console.log(`Created new partition for data at offset ${targetRootAPartition.offset} with size ${targetRootAPartition.size}`);
currentPartitions = afterSecondPartitions
rootAPartition.etcher = secondNewPartition[0];
console.log(`Created new partition for data at offset ${rootAPartition.etcher.offset} with size ${rootAPartition.etcher.size}`);
etcherPartitions = afterSecondPartitions
}
volumeIds[1] = await diskpart.findVolume(deviceName, ROOTA_PARTITION_LABEL)
console.log(`flasherRootAPartition volume: ${volumeIds[1]}`)

// copy partition data
// Use volume ID to take volume offine. At present really only necessary
// when overwriting boot partition because Windows recognizes the filesystem
// and will not allow overwriting it. No need to bring a volume back online.
console.log("Copy flasherBootPartition from image to disk");
if (volumeIds[0]) {
await diskpart.setPartitionOnlineStatus(volumeIds[0], false)
if (bootPartition.volumeId) {
// Must ensure volume offline before overwrite.
await diskpart.setPartitionOnlineStatus(bootPartition.volumeId, false)
}
await copyPartitionFromImageToDevice(source, 1, targetDevice, targetBootPartition!.offset);
// Sets volume online only if a new partition.
await copyPartitionFromImageToDevice(sourceFile, 1, targetDevice.etcher, bootPartition.etcher!.offset);
console.log("Copy complete")
console.log("Copy flasherRootAPartition from image to disk");
if (volumeIds[1]) {
await diskpart.setPartitionOnlineStatus(volumeIds[1], false)
}
await copyPartitionFromImageToDevice(source, 2, targetDevice, targetRootAPartition!.offset);
// We never set rootA partition online, so no need to offline.
await copyPartitionFromImageToDevice(sourceFile, 2, targetDevice.etcher, rootAPartition.etcher!.offset);
console.log("Copy complete")
} else {
console.log(`\nSkip task: create and copy partitions`)
Expand Down

0 comments on commit bc260b2

Please sign in to comment.