Merge pull request #406 from Washi1337/feature/metadata-stream-preser…

…vation

Metadata stream preservation

docs/dotnet/advanced-pe-image-building.rst

@@ -72,10 +72,34 @@ Some .NET modules are carefully crafted and rely on the raw structure of all met
 
 - RIDs of rows within a metadata table.
 - Indices of blobs within the ``#Blob``, ``#Strings``, ``#US`` or ``#GUID`` heaps.
-
-The default PE image builder for .NET modules (``ManagedPEImageBuilder``) defines a property called ``DotNetDirectoryFactory``, which contains the object responsible for constructing the .NET data directory, can be configured to preserve as much of this structure as possible. With the help of the ``MetadataBuilderFlags`` enum, it is possible to indicate which structures of the metadata directory need to preserved.
-
-Below an example on how to configure the image builder to preserve blob data and all metadata tokens to type references:
+- Unknown or unconventional metadata streams and their order.
+
+The default PE image builder for .NET modules (``ManagedPEImageBuilder``) defines a property called ``DotNetDirectoryFactory``, which contains the object responsible for constructing the .NET data directory, can be configured to preserve as much of this structure as possible. With the help of the ``MetadataBuilderFlags`` enum, it is possible to indicate which structures of the metadata directory need to preserved. The following table provides an overview of all preservation metadata builder flags that can be used and combined:
+
++----------------------------------------+-------------------------------------------------------------------+
+| flag                                   | Description                                                       |
++========================================+===================================================================+
+| ``PreserveXXXIndices``                 | Preserves all row indices of the original ``XXX`` metadata table. |
++----------------------------------------+-------------------------------------------------------------------+
+| ``PreserveTableIndices``               | Preserves all row indices from all original metadata tables.      |
++----------------------------------------+-------------------------------------------------------------------+
+| ``PreserveBlobIndices``                | Preserves all blob indices in the ``#Blob`` stream.               |
++----------------------------------------+-------------------------------------------------------------------+
+| ``PreserveGuidIndices``                | Preserves all GUID indices in the ``#GUID`` stream.               |
++----------------------------------------+-------------------------------------------------------------------+
+| ``PreserveStringIndices``              | Preserves all string indices in the ``#Strings`` stream.          |
++----------------------------------------+-------------------------------------------------------------------+
+| ``PreserveUserStringIndices``          | Preserves all user-string indices in the ``#US`` stream.          |
++----------------------------------------+-------------------------------------------------------------------+
+| ``PreserveUnknownStreams``             | Preserves any of the unknown / unconventional metadata streams.   |
++----------------------------------------+-------------------------------------------------------------------+
+| ``PreserveStreamOrder``                | Preserves the original order of all metadata streams.             |
++----------------------------------------+-------------------------------------------------------------------+
+| ``PreserveAll``                        | Preserves as much of the original metadata as possible.           |
++----------------------------------------+-------------------------------------------------------------------+
+
+
+Below is an example on how to configure the image builder to preserve blob data and all metadata tokens to type references:
 
 .. code-block:: csharp
 
@@ -84,7 +108,6 @@ Below an example on how to configure the image builder to preserve blob data and
                                  | MetadataBuilderFlags.PreserveTypeReferenceIndices;
     imageBuilder.DotNetDirectoryFactory = factory;
 
-If everything is supposed to be preserved as much as possible, then instead of specifying all flags defined in the ``MetadataBuilderFlags`` enum, we can also use ``MetadataBuilderFlags.PreserveAll`` as a shortcut. 
 
 .. warning::
 

src/AsmResolver.DotNet/Builder/DotNetDirectoryFactory.cs

@@ -1,10 +1,13 @@
 using System;
+using System.Linq;
 using AsmResolver.DotNet.Builder.Discovery;
 using AsmResolver.DotNet.Builder.Metadata;
 using AsmResolver.DotNet.Code;
 using AsmResolver.DotNet.Code.Cil;
 using AsmResolver.DotNet.Code.Native;
+using AsmResolver.DotNet.Serialized;
 using AsmResolver.PE.DotNet;
+using AsmResolver.PE.DotNet.Metadata;
 using AsmResolver.PE.DotNet.Metadata.Blob;
 using AsmResolver.PE.DotNet.Metadata.Guid;
 using AsmResolver.PE.DotNet.Metadata.Strings;
@@ -122,7 +125,16 @@ public virtual DotNetDirectoryBuildResult CreateDotNetDirectory(
             else if ((module.Attributes & DotNetDirectoryFlags.StrongNameSigned) != 0)
                 buffer.StrongNameSize = 0x80;
 
-            return buffer.CreateDirectory();
+            var result = buffer.CreateDirectory();
+
+            // If we need to preserve streams or stream order, apply the shifts accordingly.
+            if (module is SerializedModuleDefinition serializedModule
+                && (MetadataBuilderFlags & (MetadataBuilderFlags.PreserveUnknownStreams | MetadataBuilderFlags.PreserveStreamOrder)) != 0)
+            {
+                ReorderMetadataStreams(serializedModule, result.Directory.Metadata!);
+            }
+
+            return result;
         }
 
         private MemberDiscoveryResult DiscoverMemberDefinitionsInModule(ModuleDefinition module)
@@ -278,5 +290,61 @@ private static void ImportTables<TMember>(ModuleDefinition module, TableIndex ta
             foreach (var member in module.TokenAllocator.GetAssignees(tableIndex))
                 importAction((TMember) member);
         }
+
+        private void ReorderMetadataStreams(SerializedModuleDefinition serializedModule, IMetadata newMetadata)
+        {
+            IMetadataStream? GetStreamOrNull<TStream>()
+                where TStream : class, IMetadataStream
+            {
+                return newMetadata.TryGetStream(out TStream? stream)
+                    ? stream
+                    : null;
+            }
+
+            var readerContext = serializedModule.ReaderContext;
+            var streamIndices = new (int Index, IMetadataStream? Stream)[]
+            {
+                (readerContext.TablesStreamIndex, GetStreamOrNull<TablesStream>()),
+                (readerContext.BlobStreamIndex, GetStreamOrNull<BlobStream>()),
+                (readerContext.GuidStreamIndex, GetStreamOrNull<GuidStream>()),
+                (readerContext.StringsStreamIndex, GetStreamOrNull<StringsStream>()),
+                (readerContext.UserStringsStreamIndex, GetStreamOrNull<UserStringsStream>()),
+            };
+
+            if ((MetadataBuilderFlags & MetadataBuilderFlags.PreserveUnknownStreams) != 0)
+            {
+                var originalStreams = serializedModule.DotNetDirectory.Metadata!.Streams;
+
+                if ((MetadataBuilderFlags & MetadataBuilderFlags.PreserveStreamOrder) != 0)
+                {
+                    newMetadata.Streams.Clear();
+
+                    for (int i = 0; i < originalStreams.Count; i++)
+                    {
+                        var entry = streamIndices.FirstOrDefault(x => x.Index == i);
+                        newMetadata.Streams.Insert(i, entry.Stream ?? originalStreams[i]);
+                    }
+                }
+                else
+                {
+                    for (int i = 0; i < originalStreams.Count; i++)
+                    {
+                        if (streamIndices.All(x => x.Index != i))
+                            newMetadata.Streams.Add(originalStreams[i]);
+                    }
+                }
+            }
+            else if ((MetadataBuilderFlags & MetadataBuilderFlags.PreserveStreamOrder) != 0)
+            {
+                Array.Sort(streamIndices, (a, b) => a.Index.CompareTo(b.Index));
+                newMetadata.Streams.Clear();
+
+                for (int i = 0; i < streamIndices.Length; i++)
+                {
+                    if (streamIndices[i].Stream is { } stream)
+                        newMetadata.Streams.Add(stream);
+                }
+            }
+        }
     }
 }

src/AsmResolver.DotNet/Builder/MetadataBuilderFlags.cs

@@ -126,11 +126,23 @@ public enum MetadataBuilderFlags
             | PreserveAssemblyReferenceIndices | PreserveMethodSpecificationIndices,
 
         /// <summary>
-        /// Indicates any kind of index into a blob or tables stream should be preserved whenever possible during the
-        /// construction of the metadata directory.
+        /// Indicates unconventional / spurious metadata streams present in the .NET metadata directory should be
+        /// preserved when possible.
+        /// </summary>
+        PreserveUnknownStreams = 0x20000,
+
+        /// <summary>
+        /// Indicates unconventional metadata stream order in the .NET metadata directory should be preserved when
+        /// possible.
+        /// </summary>
+        PreserveStreamOrder = 0x40000,
+
+        /// <summary>
+        /// Indicates any kind of index into a blob or tables stream, as well as unknown spurious metadata streams
+        /// should be preserved whenever possible during the construction of the metadata directory.
         /// </summary>
         PreserveAll = PreserveBlobIndices | PreserveGuidIndices | PreserveStringIndices | PreserveUserStringIndices
-                      | PreserveTableIndices,
+                      | PreserveTableIndices | PreserveUnknownStreams | PreserveStreamOrder,
 
         /// <summary>
         /// <para>

src/AsmResolver.DotNet/Serialized/ModuleReaderContext.cs

@@ -46,18 +46,23 @@ public ModuleReaderContext(IPEImage image, SerializedModuleDefinition parentModu
                 {
                     case TablesStream tablesStream when TablesStream is null:
                         TablesStream = tablesStream;
+                        TablesStreamIndex = i;
                         break;
                     case BlobStream blobStream when BlobStream is null || !isEncMetadata:
                         BlobStream = blobStream;
+                        BlobStreamIndex = i;
                         break;
                     case GuidStream guidStream when GuidStream is null || !isEncMetadata:
                         GuidStream = guidStream;
+                        GuidStreamIndex = i;
                         break;
                     case StringsStream stringsStream when StringsStream is null || !isEncMetadata:
                         StringsStream = stringsStream;
+                        StringsStreamIndex = i;
                         break;
                     case UserStringsStream userStringsStream when UserStringsStream is null || !isEncMetadata:
                         UserStringsStream = userStringsStream;
+                        UserStringsStreamIndex = i;
                         break;
                 }
             }
@@ -96,6 +101,14 @@ public TablesStream TablesStream
             get;
         }
 
+        /// <summary>
+        /// Gets the original index of the tables stream.
+        /// </summary>
+        public int TablesStreamIndex
+        {
+            get;
+        } = -1;
+
         /// <summary>
         /// Gets the main blob stream in the metadata directory.
         /// </summary>
@@ -104,6 +117,14 @@ public BlobStream? BlobStream
             get;
         }
 
+        /// <summary>
+        /// Gets the original index of the blob stream.
+        /// </summary>
+        public int BlobStreamIndex
+        {
+            get;
+        } = -1;
+
         /// <summary>
         /// Gets the main GUID stream in the metadata directory.
         /// </summary>
@@ -112,6 +133,14 @@ public GuidStream? GuidStream
             get;
         }
 
+        /// <summary>
+        /// Gets the original index of the GUID stream.
+        /// </summary>
+        public int GuidStreamIndex
+        {
+            get;
+        } = -1;
+
         /// <summary>
         /// Gets the main strings stream in the metadata directory.
         /// </summary>
@@ -120,6 +149,14 @@ public StringsStream? StringsStream
             get;
         }
 
+        /// <summary>
+        /// Gets the original index of the strings stream.
+        /// </summary>
+        public int StringsStreamIndex
+        {
+            get;
+        } = -1;
+
         /// <summary>
         /// Gets the main user-strings stream in the metadata directory.
         /// </summary>
@@ -128,6 +165,14 @@ public UserStringsStream? UserStringsStream
             get;
         }
 
+        /// <summary>
+        /// Gets the original index of the user-strings stream.
+        /// </summary>
+        public int UserStringsStreamIndex
+        {
+            get;
+        } = -1;
+
         /// <summary>
         /// Gets the reader parameters.
         /// </summary>

test/AsmResolver.DotNet.Tests/Builder/ManagedPEImageBuilderTest.cs

@@ -2,6 +2,7 @@
 using System.Linq;
 using AsmResolver.DotNet.Builder;
 using AsmResolver.PE;
+using AsmResolver.PE.DotNet.Metadata;
 using Xunit;
 
 namespace AsmResolver.DotNet.Tests.Builder
@@ -79,5 +80,72 @@ public void ConstructPEImageFromExistingModuleWithPreservation()
             var newModule = ModuleDefinition.FromImage(result);
             Assert.Equal(module.Name, newModule.Name);
         }
+
+        [Fact]
+        public void PreserveUnknownStreams()
+        {
+            // Prepare a PE image with an extra unconventional stream.
+            var image = PEImage.FromBytes(Properties.Resources.HelloWorld);
+            byte[] data = { 1, 2, 3, 4 };
+            image.DotNetDirectory!.Metadata!.Streams.Add(new CustomMetadataStream("#Custom", data));
+
+            // Load and rebuild.
+            var module = ModuleDefinition.FromImage(image);
+            var newImage = module.ToPEImage(new ManagedPEImageBuilder(MetadataBuilderFlags.PreserveUnknownStreams));
+
+            // Verify unconventional stream is still present.
+            var newStream = Assert.IsAssignableFrom<CustomMetadataStream>(
+                newImage.DotNetDirectory!.Metadata!.GetStream("#Custom"));
+            Assert.Equal(data, Assert.IsAssignableFrom<IReadableSegment>(newStream.Contents).ToArray());
+        }
+
+        [Fact]
+        public void PreserveStreamOrder()
+        {
+            // Prepare a PE image with an unconventional stream order.
+            var image = PEImage.FromBytes(Properties.Resources.HelloWorld);
+            var streams = image.DotNetDirectory!.Metadata!.Streams;
+            for (int i = 0; i < streams.Count / 2; i++)
+                (streams[i], streams[streams.Count - i - 1]) = (streams[streams.Count - i - 1], streams[i]);
+
+            // Load and rebuild.
+            var module = ModuleDefinition.FromImage(image);
+            var newImage = module.ToPEImage(new ManagedPEImageBuilder(MetadataBuilderFlags.PreserveStreamOrder));
+
+            // Verify order is still the same.
+            Assert.Equal(
+                streams.Select(x => x.Name),
+                newImage.DotNetDirectory!.Metadata!.Streams.Select(x => x.Name));
+        }
+
+        [Fact]
+        public void PreserveUnknownStreamsAndStreamOrder()
+        {
+            // Prepare a PE image with an unconventional stream order and custom stream.
+            var image = PEImage.FromBytes(Properties.Resources.HelloWorld);
+            var streams = image.DotNetDirectory!.Metadata!.Streams;
+
+            for (int i = 0; i < streams.Count / 2; i++)
+                (streams[i], streams[streams.Count - i - 1]) = (streams[streams.Count - i - 1], streams[i]);
+
+            byte[] data = { 1, 2, 3, 4 };
+            image.DotNetDirectory!.Metadata!.Streams.Insert(streams.Count / 2,
+                new CustomMetadataStream("#Custom", data));
+
+            // Load and rebuild.
+            var module = ModuleDefinition.FromImage(image);
+            var newImage = module.ToPEImage(new ManagedPEImageBuilder(
+                MetadataBuilderFlags.PreserveStreamOrder | MetadataBuilderFlags.PreserveUnknownStreams));
+
+            // Verify order is still the same.
+            Assert.Equal(
+                streams.Select(x => x.Name),
+                newImage.DotNetDirectory!.Metadata!.Streams.Select(x => x.Name));
+
+            // Verify unconventional stream is still present.
+            var newStream = Assert.IsAssignableFrom<CustomMetadataStream>(
+                newImage.DotNetDirectory!.Metadata!.GetStream("#Custom"));
+            Assert.Equal(data, Assert.IsAssignableFrom<IReadableSegment>(newStream.Contents).ToArray());
+        }
     }
 }