Added restore of partition tables

proprietary/DiskImage/Partition/GPT.cs

@@ -443,6 +443,29 @@ public class GPT : IPartitionTable
         throw new InvalidOperationException("No protective MBR available.");
     }
 
+    public async Task<Stream> GetPartitionTableDataAsync(CancellationToken cancellationToken)
+    {
+        if (!m_parsed)
+            throw new InvalidOperationException("GPT header not parsed.");
+
+        if (m_rawDisk == null)
+            throw new InvalidOperationException("No raw disk available for reading GPT data.");
+
+        // Calculate the total size needed:
+        // - Protective MBR (512 bytes)
+        // - GPT Header (1 sector)
+        // - Partition Entries (m_partitionEntryLba sectors)
+        long partitionEntriesEnd = m_partitionEntryLba * m_bytesPerSector + (m_numPartitionEntries * m_partitionEntrySize);
+        long totalSize = partitionEntriesEnd;
+
+        // Read all the data
+        using var stream = await m_rawDisk.ReadBytesAsync(0, (int)totalSize, cancellationToken).ConfigureAwait(false);
+        var buffer = new byte[totalSize];
+        await stream.ReadAtLeastAsync(buffer, (int)totalSize, cancellationToken: cancellationToken).ConfigureAwait(false);
+
+        return new MemoryStream(buffer, writable: false);
+    }
+
     private async Task<bool> VerifyBackupHeaderAsync(CancellationToken token)
     {
         if (m_rawDisk == null)

proprietary/DiskImage/Partition/Interfaces.cs

@@ -93,4 +93,11 @@ public interface IPartitionTable : IDisposable
     /// Gets the protective MBR (for GPT disks).
     /// </summary>
     Task<Stream> GetProtectiveMbrAsync(CancellationToken cancellationToken);
+
+    /// <summary>
+    /// Gets the raw partition table data for backup/restore purposes.
+    /// For MBR: returns the MBR sector (512 bytes).
+    /// For GPT: returns the protective MBR + GPT header + partition entries.
+    /// </summary>
+    Task<Stream> GetPartitionTableDataAsync(CancellationToken cancellationToken);
 }

proprietary/DiskImage/Partition/MBR.cs

@@ -246,6 +246,15 @@ public class MBR : IPartitionTable
         throw new NotSupportedException("MBR partition tables do not have a protective MBR. Use GetPartitionAsync instead.");
     }
 
+    public Task<Stream> GetPartitionTableDataAsync(CancellationToken cancellationToken)
+    {
+        if (!m_parsed || m_mbrBytes == null)
+            throw new InvalidOperationException("MBR not parsed.");
+
+        // Return the MBR sector (512 bytes)
+        return Task.FromResult<Stream>(new MemoryStream(m_mbrBytes, writable: false));
+    }
+
     public void Dispose()
     {
         Dispose(true);

proprietary/DiskImage/Partition/PartitionTableFactory.cs

@@ -45,7 +45,7 @@ public static class PartitionTableFactory
         {
             PartitionTableType.GPT => await CreateGPTAsync(disk, mbrBytes, cancellationToken).ConfigureAwait(false),
             PartitionTableType.MBR => await CreateMBRAsync(disk, mbrBytes, cancellationToken).ConfigureAwait(false),
-            _ => null
+            _ => new UnknownPartitionTable(disk)
         };
     }
 
@@ -77,7 +77,7 @@ public static class PartitionTableFactory
         {
             PartitionTableType.GPT => await CreateGPTAsync(bytes, sectorSize, cancellationToken).ConfigureAwait(false),
             PartitionTableType.MBR => await CreateMBRAsync(mbrBytes, sectorSize, cancellationToken).ConfigureAwait(false),
-            _ => null
+            _ => new UnknownPartitionTable(null)
         };
     }
 

proprietary/DiskImage/Partition/Unknown.cs

@@ -71,6 +71,13 @@ public class UnknownPartitionTable : IPartitionTable
         throw new NotSupportedException("Unknown partition table does not have a protective MBR.");
     }
 
+    public Task<Stream> GetPartitionTableDataAsync(CancellationToken cancellationToken)
+    {
+        // For unknown partition tables, return empty data
+        // The disk will be treated as one big partition with unknown filesystem
+        return Task.FromResult<Stream>(new MemoryStream(Array.Empty<byte>(), writable: false));
+    }
+
     public void Dispose()
     {
         if (m_disposed)

proprietary/DiskImage/RestoreProvider.cs

@@ -225,12 +225,21 @@ public sealed class RestoreProvider : IRestoreDestinationProviderModule, IDispos
         var processedCount = 0;
 
         // Group items by type for ordered restoration
+        var partitionTableItems = GetMetadataByType("partition_table");
         var diskItems = GetMetadataByType("disk");
         var partitionItems = GetMetadataByType("partition");
         var blockItems = GetMetadataByType("block");
         var fileItems = GetMetadataByType("file");
 
-        // Restore disk-level items first (partition tables)
+        // Restore partition table items first (if available)
+        if (!_skipPartitionTable && partitionTableItems.Count > 0)
+        {
+            await RestorePartitionTableItems(partitionTableItems, cancel);
+            processedCount += partitionTableItems.Count;
+            progressCallback?.Invoke(processedCount / (double)totalItems);
+        }
+
+        // Restore disk-level items (full disk image)
         if (!_skipPartitionTable && diskItems.Count > 0)
         {
             await RestoreDiskItems(diskItems, cancel);
@@ -272,7 +281,56 @@ public sealed class RestoreProvider : IRestoreDestinationProviderModule, IDispos
     }
 
     /// <summary>
-    /// Restores disk-level items (partition tables, boot sectors).
+    /// Restores partition table items (MBR or GPT).
+    /// </summary>
+    private async Task RestorePartitionTableItems(List<KeyValuePair<string, Dictionary<string, string?>>> items, CancellationToken cancel)
+    {
+        foreach (var item in items)
+        {
+            if (cancel.IsCancellationRequested)
+                break;
+
+            var path = item.Key;
+            var metadata = item.Value;
+
+            if (!_temporaryFiles.TryGetValue(path, out var tempFile))
+                continue;
+
+            try
+            {
+                // Extract metadata for logging
+                metadata.TryGetValue("partition_table:Type", out var tableType);
+                metadata.TryGetValue("partition_table:Size", out var tableSizeStr);
+                metadata.TryGetValue("disk:SectorSize", out var sectorSizeStr);
+
+                Log.WriteInformationMessage(LOGTAG, "RestorePartitionTableMetadata",
+                    $"Restoring partition table (Type: {tableType}, Size: {tableSizeStr}, SectorSize: {sectorSizeStr})");
+
+                // Read the data from temp file
+                byte[] data;
+                using (var stream = SystemIO.IO_OS.FileOpenRead(tempFile))
+                {
+                    data = new byte[stream.Length];
+                    await stream.ReadExactlyAsync(data, cancel);
+                }
+
+                // Write partition table to disk starting at offset 0
+                // For MBR: this is the 512-byte MBR sector
+                // For GPT: this includes protective MBR + GPT header + partition entries
+                await _targetDisk!.WriteBytesAsync(0, data, cancel);
+
+                Log.WriteInformationMessage(LOGTAG, "RestorePartitionTable", $"Restored partition table: {path} (Type: {tableType})");
+            }
+            catch (Exception ex)
+            {
+                Log.WriteErrorMessage(LOGTAG, "RestorePartitionTableFailed", ex, $"Failed to restore partition table: {path}");
+                throw;
+            }
+        }
+    }
+
+    /// <summary>
+    /// Restores disk-level items (full disk image).
     /// </summary>
     private async Task RestoreDiskItems(List<KeyValuePair<string, Dictionary<string, string?>>> items, CancellationToken cancel)
     {
@@ -318,9 +376,10 @@ public sealed class RestoreProvider : IRestoreDestinationProviderModule, IDispos
                 }
 
                 // Write to disk (offset 0 for partition table/boot sector)
-                // TODO in the future, write the correct GPT table: primary at offset 0, secondary at end of disk. The header should also be updated to have correct offsets and CRCs.
                 await _targetDisk!.WriteBytesAsync(0, data, cancel);
 
+                // TODO for GPT disks, we may also need to write the backup GPT header at the end of the disk. The fields in the header should also reflect this. Furthermore, if the GUID of the disk has changed, we may need to update these as well.
+
                 Log.WriteInformationMessage(LOGTAG, "RestoreDiskItem", $"Restored disk item: {path} (Table: {partitionTableType})");
             }
             catch (Exception ex)

proprietary/DiskImage/SourceItems/DiskSourceEntry.cs

@@ -48,21 +48,22 @@ internal class DiskSourceEntry(SourceProvider provider, IRawDisk disk)
         {
             metadata["disk:PartitionTableType"] = table.TableType.ToString();
 
-            // For GPT disks, include the protective MBR info
-            if (table.TableType == PartitionTableType.GPT)
+            // Get the full partition table data for restore
+            try
             {
-                try
+                using var tableData = await table.GetPartitionTableDataAsync(cancellationToken);
+                metadata["disk:PartitionTableDataSize"] = tableData.Length.ToString();
+
+                // For GPT disks, also record specific info
+                if (table.TableType == PartitionTableType.GPT)
                 {
-                    using var protectiveMbr = await table.GetProtectiveMbrAsync(cancellationToken);
-                    using var ms = new MemoryStream();
-                    await protectiveMbr.CopyToAsync(ms, cancellationToken);
-                    metadata["disk:ProtectiveMbrSize"] = ms.Length.ToString();
-                }
-                catch
-                {
-                    // Protective MBR not available
+                    metadata["disk:HasProtectiveMbr"] = "true";
                 }
             }
+            catch
+            {
+                // Partition table data not available
+            }
         }
 
         return metadata;

proprietary/DiskImage/SourceItems/PartitionTableSourceEntry.cs

@@ -0,0 +1,88 @@
+// Copyright (c) 2026 Duplicati Inc. All rights reserved.
+
+using System;
+using System.Collections.Generic;
+using System.IO;
+using System.Runtime.CompilerServices;
+using System.Threading;
+using System.Threading.Tasks;
+using Duplicati.Library.Interface;
+using Duplicati.Proprietary.DiskImage.Disk;
+using Duplicati.Proprietary.DiskImage.Partition;
+
+namespace Duplicati.Proprietary.DiskImage.SourceItems;
+
+/// <summary>
+/// Represents the partition table as a source entry for backup/restore purposes.
+/// This allows the partition table (MBR or GPT) to be backed up and restored separately.
+/// </summary>
+internal class PartitionTableSourceEntry(string parentPath, IPartitionTable table, IRawDisk disk)
+    : DiskImageEntryBase(PathCombine(parentPath, $"partition_table_{table.TableType}"))
+{
+    public override bool IsFolder => false;
+    public override bool IsMetaEntry => true;
+    public override long Size => GetTableSize();
+
+    public override Task<Stream> OpenRead(CancellationToken cancellationToken)
+    {
+        return table.GetPartitionTableDataAsync(cancellationToken);
+    }
+
+    public override async IAsyncEnumerable<ISourceProviderEntry> Enumerate([EnumeratorCancellation] CancellationToken cancellationToken)
+    {
+        // Partition table is a leaf node - no children
+        await Task.CompletedTask;
+        yield break;
+    }
+
+    public override async Task<Dictionary<string, string?>> GetMinorMetadata(CancellationToken cancellationToken)
+    {
+        var metadata = await base.GetMinorMetadata(cancellationToken);
+
+        // Add partition table metadata
+        metadata["diskimage:Type"] = "partition_table";
+        metadata["partition_table:Type"] = table.TableType.ToString();
+        metadata["partition_table:Size"] = Size.ToString();
+        metadata["disk:SectorSize"] = disk.SectorSize.ToString();
+
+        // Add type-specific metadata
+        if (table.TableType == PartitionTableType.MBR)
+        {
+            metadata["partition_table:MbrSize"] = "512";
+        }
+        else if (table.TableType == PartitionTableType.GPT)
+        {
+            metadata["partition_table:HasProtectiveMbr"] = "true";
+            metadata["partition_table:HeaderSize"] = "92"; // GPT header is 92 bytes
+        }
+
+        return metadata;
+    }
+
+    private long GetTableSize()
+    {
+        try
+        {
+            // Try to get the size by reading the table data
+            using var stream = table.GetPartitionTableDataAsync(CancellationToken.None).Result;
+            return stream.Length;
+        }
+        catch
+        {
+            // Fallback to estimated sizes
+            return table.TableType switch
+            {
+                PartitionTableType.MBR => 512,
+                PartitionTableType.GPT => 512 + 512 + (128 * 4), // Protective MBR + Header + 4 entries
+                _ => 0
+            };
+        }
+    }
+
+    private static string PathCombine(string p1, string p2)
+    {
+        if (string.IsNullOrEmpty(p1)) return p2;
+        if (string.IsNullOrEmpty(p2)) return p1;
+        return p1.TrimEnd('/') + "/" + p2.TrimStart('/');
+    }
+}

proprietary/DiskImage/SourceProvider.cs

@@ -8,6 +8,7 @@ using System.Threading;
 using System.Threading.Tasks;
 using Duplicati.Library.Interface;
 using Duplicati.Proprietary.DiskImage.Disk;
+using Duplicati.Proprietary.DiskImage.Partition;
 using Duplicati.Proprietary.DiskImage.SourceItems;
 
 namespace Duplicati.Proprietary.DiskImage;
@@ -74,6 +75,13 @@ public sealed class SourceProvider : ISourceProviderModule, IDisposable
 
         var root = new DiskSourceEntry(this, _disk);
         yield return root;
+
+        // Also enumerate partition table as a separate entry for backup/restore
+        var table = await PartitionTableFactory.CreateAsync(_disk, cancellationToken);
+        if (table != null)
+        {
+            yield return new PartitionTableSourceEntry(root.Path, table, _disk);
+        }
     }
 
     public async Task<ISourceProviderEntry?> GetEntry(string path, bool isFolder, CancellationToken cancellationToken)