Files
duplicati/Duplicati/UnitTest/DiskImage/WindowsDiskImageHelper.cs
T

754 lines
28 KiB
C#

// Copyright (C) 2025, The Duplicati Team
// https://duplicati.com, hello@duplicati.com
//
// Permission is hereby granted, free of charge, to any person obtaining a
// copy of this software and associated documentation files (the "Software"),
// to deal in the Software without restriction, including without limitation
// the rights to use, copy, modify, merge, publish, distribute, sublicense,
// and/or sell copies of the Software, and to permit persons to whom the
// Software is furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
// OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
// DEALINGS IN THE SOFTWARE.
using System;
using System.Diagnostics;
using System.IO;
using System.Linq;
using System.Security.Principal;
using System.Text;
using System.Threading;
using NUnit.Framework;
#nullable enable
namespace Duplicati.UnitTest.DiskImage
{
/// <summary>
/// Manages a persistent PowerShell session for executing commands.
/// This avoids the overhead of starting a new PowerShell process for each command.
/// </summary>
internal sealed class PowerShellSession : IDisposable
{
private Process? _process;
private readonly object _lock = new();
private bool _disposed;
private readonly StringBuilder _outputBuffer = new();
private readonly StringBuilder _errorBuffer = new();
private readonly ManualResetEventSlim _outputAvailable = new(false);
/// <summary>
/// Gets a value indicating whether the PowerShell session is running.
/// </summary>
/// <value><c>true</c> if the session is running; otherwise, <c>false</c>.</value>
public bool IsRunning => _process != null && !_process.HasExited;
/// <summary>
/// Ensures the PowerShell session is started.
/// </summary>
/// <exception cref="ObjectDisposedException">Thrown if the session has been disposed.</exception>
public void EnsureStarted()
{
lock (_lock)
{
if (_disposed)
{
throw new ObjectDisposedException(nameof(PowerShellSession));
}
if (_process != null && !_process.HasExited)
{
return;
}
StartProcess();
}
}
/// <summary>
/// Starts the PowerShell process with elevated privileges.
/// </summary>
/// <exception cref="InvalidOperationException">Thrown if the PowerShell process cannot be started.</exception>
private void StartProcess()
{
var startInfo = new ProcessStartInfo
{
FileName = "powershell.exe",
Arguments = "-NoProfile -ExecutionPolicy Bypass -Command \"& { while ($true) { $cmd = Read-Host; if ($cmd -eq 'EXIT_SESSION') { break }; try { $output = Invoke-Expression $cmd 2>&1; $output | Out-String -Stream; Write-Output \"___ENDOFCOMMAND___\" } catch { $_.Exception.Message; Write-Output \"___ENDOFCOMMAND___\" } } }\"",
RedirectStandardInput = true,
RedirectStandardOutput = true,
RedirectStandardError = true,
UseShellExecute = false,
CreateNoWindow = true,
Verb = "runas" // Request elevation
};
_process = Process.Start(startInfo);
if (_process == null)
{
throw new InvalidOperationException("Failed to start PowerShell process");
}
// Clear any startup output
_outputBuffer.Clear();
_errorBuffer.Clear();
_outputAvailable.Reset();
// Attach event handlers for async reading
_process.OutputDataReceived += OnOutputDataReceived;
_process.ErrorDataReceived += OnErrorDataReceived;
_process.BeginOutputReadLine();
_process.BeginErrorReadLine();
}
/// <summary>
/// Handles the output data received from the PowerShell process.
/// </summary>
/// <param name="sender">The source of the event.</param>
/// <param name="e">The event data.</param>
private void OnOutputDataReceived(object sender, DataReceivedEventArgs e)
{
if (e.Data != null)
{
lock (_outputBuffer)
{
_outputBuffer.AppendLine(e.Data);
if (e.Data.Contains("___ENDOFCOMMAND___"))
{
_outputAvailable.Set();
}
}
}
}
/// <summary>
/// Handles the error data received from the PowerShell process.
/// </summary>
/// <param name="sender">The source of the event.</param>
/// <param name="e">The event data.</param>
private void OnErrorDataReceived(object sender, DataReceivedEventArgs e)
{
if (e.Data != null)
{
lock (_errorBuffer)
{
_errorBuffer.AppendLine(e.Data);
}
}
}
/// <summary>
/// Executes a PowerShell script and returns the output.
/// </summary>
/// <param name="script">The PowerShell script to execute.</param>
/// <returns>The output from PowerShell.</returns>
/// <exception cref="InvalidOperationException">Thrown if the PowerShell process is not running.</exception>
/// <exception cref="TimeoutException">Thrown if the command times out.</exception>
public string ExecuteScript(string script)
{
lock (_lock)
{
EnsureStarted();
if (_process == null || _process.HasExited)
{
throw new InvalidOperationException("PowerShell process is not running");
}
// Clear previous output and reset the event
lock (_outputBuffer)
{
_outputBuffer.Clear();
}
lock (_errorBuffer)
{
_errorBuffer.Clear();
}
_outputAvailable.Reset();
// Write the script to stdin using base64 encoding to handle special characters
var encodedScript = Convert.ToBase64String(Encoding.UTF8.GetBytes(script));
var command = $"[System.Text.Encoding]::UTF8.GetString([Convert]::FromBase64String('{encodedScript}')) | Invoke-Expression";
_process.StandardInput.WriteLine(command);
// Wait for output with a timeout
var timeout = TimeSpan.FromSeconds(60);
if (!_outputAvailable.Wait(timeout))
{
throw new TimeoutException($"PowerShell command timed out after {timeout.TotalSeconds} seconds.\nScript:\n{script}");
}
// Give a bit more time for all output to arrive
Thread.Sleep(50);
lock (_outputBuffer)
{
var rawOutput = _outputBuffer.ToString();
_outputBuffer.Clear();
// Remove the end marker and any lines that contain our base64 command
var lines = rawOutput.Split(new[] { '\r', '\n' }, StringSplitOptions.RemoveEmptyEntries);
var filteredLines = lines
.Where(line => !line.Contains("___ENDOFCOMMAND___"))
.Where(line => !line.Contains("FromBase64String"))
.Where(line => !line.Contains(encodedScript.Substring(0, Math.Min(20, encodedScript.Length))))
.ToList();
var output = string.Join(Environment.NewLine, filteredLines).Trim();
lock (_errorBuffer)
{
var error = _errorBuffer.ToString().Trim();
_errorBuffer.Clear();
if (!string.IsNullOrEmpty(error))
{
throw new InvalidOperationException($"PowerShell error: {error}\nScript:\n{script}");
}
}
return output;
}
}
}
/// <summary>
/// Disposes the PowerShell session.
/// </summary>
public void Dispose()
{
lock (_lock)
{
if (_disposed)
{
return;
}
_disposed = true;
if (_process != null)
{
// Unregister event handlers
_process.OutputDataReceived -= OnOutputDataReceived;
_process.ErrorDataReceived -= OnErrorDataReceived;
if (!_process.HasExited)
{
try
{
_process.StandardInput.WriteLine("EXIT_SESSION");
if (!_process.WaitForExit(5000))
{
_process.Kill();
}
}
catch
{
// Ignore errors during cleanup
try { _process.Kill(); } catch { }
}
}
_process.Dispose();
_process = null;
}
_outputAvailable.Dispose();
}
}
}
/// <summary>
/// Windows implementation of <see cref="IDiskImageHelper"/> using PowerShell and VHD files.
/// </summary>
internal class WindowsDiskImageHelper : IDiskImageHelper
{
// Static PowerShell session that persists across method calls
private static PowerShellSession? _session;
private static readonly object _sessionLock = new();
/// <summary>
/// Gets the shared PowerShell session, creating it if necessary.
/// </summary>
/// <returns>The shared <see cref="PowerShellSession"/> instance.</returns>
private static PowerShellSession GetSession()
{
lock (_sessionLock)
{
if (_session == null)
{
_session = new PowerShellSession();
}
return _session;
}
}
/// <summary>
/// Runs a PowerShell script using the persistent session.
/// </summary>
/// <param name="script">The PowerShell script to execute.</param>
/// <returns>The output from PowerShell.</returns>
public static string RunPowerShell(string script)
{
var session = GetSession();
return session.ExecuteScript(script);
}
/// <inheritdoc />
public string CreateAndAttachDisk(string imagePath, long sizeMB)
{
// Ensure the directory exists
var directory = Path.GetDirectoryName(imagePath);
if (!string.IsNullOrEmpty(directory) && !Directory.Exists(directory))
{
Directory.CreateDirectory(directory);
}
// Delete existing VHD if it exists
if (File.Exists(imagePath))
{
try
{
DetachDisk(imagePath);
File.Delete(imagePath);
}
catch
{
// Ignore errors during cleanup
}
}
// Create the VHD using PowerShell
var script = $@"New-VHD -Path '{imagePath}' -SizeBytes {sizeMB}MB -Fixed | Out-Null; Mount-DiskImage -ImagePath '{imagePath}' | Out-Null";
RunPowerShell(script);
// Wait for the disk to be attached and get the disk number
int diskNumber = WaitForDiskAttachment(imagePath, TimeSpan.FromSeconds(30));
if (diskNumber < 0)
{
throw new InvalidOperationException($"Failed to get disk number for VHD: {imagePath}");
}
return $"\\\\.\\PhysicalDrive{diskNumber}";
}
/// <summary>
/// Waits for a VHD to be attached and returns its disk number.
/// </summary>
/// <param name="imagePath">The path to the VHD file.</param>
/// <param name="timeout">Maximum time to wait.</param>
/// <returns>The disk number, or -1 if not found within the timeout.</returns>
private int WaitForDiskAttachment(string imagePath, TimeSpan timeout)
{
var startTime = DateTime.UtcNow;
while (DateTime.UtcNow - startTime < timeout)
{
var diskNumber = this.GetDiskNumber(imagePath);
if (diskNumber >= 0)
{
return diskNumber;
}
Thread.Sleep(100);
}
return -1;
}
/// <inheritdoc />
public void InitializeDisk(int diskNumber, Duplicati.Proprietary.DiskImage.PartitionTableType tableType)
{
if (tableType == Proprietary.DiskImage.PartitionTableType.Unknown)
throw new ArgumentException("Invalid partition table type", nameof(tableType));
var script = $@"
$disk = Get-Disk -Number {diskNumber}
if ($disk.PartitionStyle -ne 'RAW') {{
Clear-Disk -Number {diskNumber} -RemoveData -Confirm:$false
}}
Initialize-Disk -Number {diskNumber} -PartitionStyle {tableType.ToString().ToUpperInvariant()} -PassThru | Out-Null
";
RunPowerShell(script);
// Wait for the initialization to complete
WaitForDiskInitialization(diskNumber, TimeSpan.FromSeconds(30));
}
/// <summary>
/// Waits for a disk to be initialized.
/// </summary>
/// <param name="diskNumber">The disk number.</param>
/// <param name="timeout">Maximum time to wait.</param>
/// <exception cref="TimeoutException">Thrown if initialization times out.</exception>
private static void WaitForDiskInitialization(int diskNumber, TimeSpan timeout)
{
var startTime = DateTime.UtcNow;
while (DateTime.UtcNow - startTime < timeout)
{
var script = $"(Get-Disk -Number {diskNumber}).PartitionStyle";
var result = RunPowerShell(script)?.Trim();
if (result == "GPT" || result == "MBR")
{
return;
}
Thread.Sleep(100);
}
throw new TimeoutException($"Disk {diskNumber} initialization timed out");
}
/// <inheritdoc />
public char CreateAndFormatPartition(int diskNumber, Duplicati.Proprietary.DiskImage.FileSystemType fsType, long sizeMB = 0)
{
var fsTypeUpper = fsType.ToString().ToUpperInvariant();
// Find an available drive letter
char driveLetter = FindAvailableDriveLetter();
var sizeParam = sizeMB > 0 ? $"-Size {sizeMB}MB" : "-UseMaximumSize";
var script = $@"
$partition = New-Partition -DiskNumber {diskNumber} {sizeParam} -AssignDriveLetter
# Wait for the partition to be ready
$timeout = (Get-Date).AddSeconds(30)
while ((Get-Date) -lt $timeout) {{
$vol = Get-Partition -DiskNumber {diskNumber} | Where-Object {{ $_.PartitionNumber -eq $partition.PartitionNumber }} | Get-Volume
if ($vol) {{ break }}
Start-Sleep -Milliseconds 100
}}
# Format the volume
$vol = Format-Volume -Partition $partition -FileSystem {fsTypeUpper} -NewFileSystemLabel 'TestVol' -Confirm:$false
# Get the assigned drive letter
(Get-Partition -DiskNumber {diskNumber} | Where-Object {{ $_.PartitionNumber -eq $partition.PartitionNumber }} | Get-Volume).DriveLetter
";
var result = RunPowerShell(script);
var assignedLetter = result?.Trim();
if (!string.IsNullOrEmpty(assignedLetter) && assignedLetter.Length == 1 && char.IsLetter(assignedLetter[0]))
{
return char.ToUpperInvariant(assignedLetter[0]);
}
// Fallback: try to find the drive letter by checking the disk
return FindDriveLetterForDisk(diskNumber);
}
/// <summary>
/// Finds the drive letter assigned to a disk.
/// </summary>
/// <param name="diskNumber">The disk number.</param>
/// <returns>The drive letter.</returns>
/// <exception cref="InvalidOperationException">Thrown if drive letter cannot be found.</exception>
private static char FindDriveLetterForDisk(int diskNumber)
{
var script = $@"
Get-Partition -DiskNumber {diskNumber} | Get-Volume | Where-Object {{ $_.DriveLetter -ne $null }} | Select-Object -ExpandProperty DriveLetter
";
var result = RunPowerShell(script)?.Trim();
if (!string.IsNullOrEmpty(result) && result.Length >= 1 && char.IsLetter(result[0]))
{
return char.ToUpperInvariant(result[0]);
}
throw new InvalidOperationException($"Could not find drive letter for disk {diskNumber}");
}
/// <inheritdoc />
public void FlushVolume(char driveLetter)
{
try
{
var script = $@"
Write-VolumeCache -DriveLetter {driveLetter}
";
RunPowerShell(script);
}
catch (Exception ex)
{
TestContext.Progress.WriteLine($"Warning: Failed to flush volume cache: {ex.Message}");
}
}
/// <inheritdoc />
public void PopulateTestData(char driveLetter, int fileCount = 10, int fileSizeKB = 10)
{
var drivePath = $"{driveLetter}:\\";
// Create a small text file
File.WriteAllText(Path.Combine(drivePath, "testfile_small.txt"),
"This is a small test file for disk image backup testing.\n" +
"It contains simple text data.\n" +
new string('=', 100));
// Create a medium binary file with random data
var random = new Random(42); // Fixed seed for reproducibility
var mediumData = new byte[fileSizeKB * 1024];
random.NextBytes(mediumData);
File.WriteAllBytes(Path.Combine(drivePath, "testfile_medium.bin"), mediumData);
// Create nested directories with files
var testDir = Path.Combine(drivePath, "testdir");
Directory.CreateDirectory(testDir);
File.WriteAllText(Path.Combine(testDir, "nested_file.txt"),
"This file is in a nested directory.\n" +
$"Created at: {DateTime.UtcNow:O}\n");
var subDir = Path.Combine(testDir, "subdir");
Directory.CreateDirectory(subDir);
var deepData = new byte[1024];
random.NextBytes(deepData);
File.WriteAllBytes(Path.Combine(subDir, "deep_file.bin"), deepData);
// Create additional files if requested
for (int i = 0; i < fileCount - 4; i++)
{
var fileData = new byte[fileSizeKB * 1024];
random.NextBytes(fileData);
File.WriteAllBytes(Path.Combine(drivePath, $"testfile_{i:D3}.bin"), fileData);
}
}
/// <inheritdoc />
public void DetachDisk(string imagePath)
{
try
{
var script = $@"
$image = Get-DiskImage -ImagePath '{imagePath}' -ErrorAction SilentlyContinue
if ($image -and $image.Attached) {{
Dismount-DiskImage -ImagePath '{imagePath}'
}}
";
RunPowerShell(script);
}
catch
{
// Ignore errors during detach - the disk may already be detached
}
}
/// <inheritdoc />
public void UnmountForWriting(string imagePath, char? driveLetter = null)
{
var diskNumber = GetDiskNumber(imagePath);
if (driveLetter == null)
{
try { driveLetter = FindDriveLetterForDisk(diskNumber); } catch { }
}
string script = string.Empty;
if (driveLetter != null && driveLetter != '\0')
{
// Remove the drive letter
script = $@"
Get-Volume -Drive {driveLetter} | Get-Partition | Remove-PartitionAccessPath -AccessPath {driveLetter}:\
";
RunPowerShell(script);
}
// Pull the disk offline to ensure it's not in use
script = $@"
Set-Disk -Number {diskNumber} -IsOffline $true
";
RunPowerShell(script);
// Clear the readonly flag
script = $@"
Set-Disk -Number {diskNumber} -IsReadOnly $false
";
RunPowerShell(script);
}
/// <inheritdoc />
public void BringOnline(string imagePath)
{
var diskNumber = GetDiskNumber(imagePath);
if (diskNumber >= 0)
{
var script = $@"
Set-Disk -Number {diskNumber} -IsOffline $false
Update-Disk -Number {diskNumber}
Update-HostStorageCache
";
RunPowerShell(script);
}
}
/// <inheritdoc />
public char MountForReading(string imagePath, char? driveLetter = null)
{
driveLetter ??= FindAvailableDriveLetter();
try
{
var diskNumber = GetDiskNumber(imagePath);
if (diskNumber >= 0)
{
var script = $@"
Set-Disk -Number {diskNumber} -IsOffline $false
Update-Disk -Number {diskNumber}
Update-HostStorageCache
Get-Partition -DiskNumber {diskNumber} | Set-Partition -NewDriveLetter {driveLetter}
";
RunPowerShell(script);
}
}
catch (Exception ex)
{
TestContext.Progress.WriteLine($"Warning: Failed to mount VHD for reading: {ex.Message}");
}
return driveLetter.Value;
}
/// <inheritdoc />
public int GetDiskNumber(string imagePath)
{
try
{
var script = $@"
$diskImage = Get-DiskImage -ImagePath '{imagePath}' -ErrorAction SilentlyContinue
if ($diskImage -and $diskImage.Attached) {{
Get-Disk | Where-Object {{ $_.Path -like '*{Path.GetFileName(imagePath)}*' -or ($_ | Get-DiskImage).ImagePath -eq '{imagePath}' }} | Select-Object -ExpandProperty Number
}}
";
var result = RunPowerShell(script)?.Trim();
if (!string.IsNullOrEmpty(result) && int.TryParse(result, out int diskNumber))
{
return diskNumber;
}
// Alternative approach: query by disk image path
script = $@"
Get-DiskImage -ImagePath '{imagePath}' | Get-Disk | Select-Object -ExpandProperty Number
";
result = RunPowerShell(script)?.Trim();
if (!string.IsNullOrEmpty(result) && int.TryParse(result, out diskNumber))
{
return diskNumber;
}
}
catch (Exception ex)
{
TestContext.Progress.WriteLine($"Warning: PowerShell GetDiskNumber failed: {ex.Message}");
}
return -1;
}
/// <inheritdoc />
public void CleanupDisk(string imagePath)
{
try
{
DetachDisk(imagePath);
}
catch (Exception ex)
{
Console.WriteLine($"Warning: Failed to detach VHD {imagePath}: {ex.Message}");
}
try
{
if (File.Exists(imagePath))
{
File.Delete(imagePath);
}
}
catch (Exception ex)
{
Console.WriteLine($"Warning: Failed to delete VHD file {imagePath}: {ex.Message}");
}
}
/// <inheritdoc />
public string GetDiskDetails(int diskNumber)
{
var script = $@"
Write-Output ""Disk Details for Disk {diskNumber}:""
Write-Output ""================================""
Get-Disk -Number {diskNumber} | Format-List
Write-Output ''
Write-Output ""Partitions:""
Write-Output ""===========""
Get-Partition -DiskNumber {diskNumber} | Format-Table
";
return RunPowerShell(script);
}
/// <inheritdoc />
public string GetVolumeInfo(char driveLetter)
{
var script = $@"
Get-Volume -DriveLetter {driveLetter} | Format-List
";
return RunPowerShell(script);
}
/// <inheritdoc />
public bool HasRequiredPrivileges()
{
if (!OperatingSystem.IsWindows())
{
return false;
}
try
{
using var identity = WindowsIdentity.GetCurrent();
var principal = new WindowsPrincipal(identity);
return principal.IsInRole(WindowsBuiltInRole.Administrator);
}
catch
{
return false;
}
}
/// <summary>
/// Finds an available drive letter that is not currently in use.
/// </summary>
/// <returns>An available drive letter.</returns>
/// <exception cref="InvalidOperationException">Thrown if no available drive letters are found.</exception>
private static char FindAvailableDriveLetter()
{
var usedDrives = DriveInfo.GetDrives()
.Where(d => d.Name.Length >= 2)
.Select(d => char.ToUpperInvariant(d.Name[0]))
.ToHashSet();
// Try letters from Z down to D (avoid A, B, C which are typically reserved)
for (char c = 'Z'; c >= 'D'; c--)
{
if (!usedDrives.Contains(c))
{
return c;
}
}
throw new InvalidOperationException("No available drive letters found.");
}
/// <inheritdoc />
public void Dispose()
{
lock (_sessionLock)
{
_session?.Dispose();
_session = null;
}
}
}
}