Address several issues with reading specially crafted .NET metadata

.editorconfig

@@ -13,6 +13,7 @@ csharp_preferred_modifier_order = public, private, protected, internal, new, abs
 csharp_style_var_elsewhere = true:suggestion
 csharp_style_var_for_built_in_types = false:suggestion
 csharp_style_var_when_type_is_apparent = true:suggestion
+csharp_new_line_before_open_brace = all
 dotnet_style_parentheses_in_arithmetic_binary_operators = never_if_unnecessary:none
 dotnet_style_parentheses_in_other_binary_operators = never_if_unnecessary:none
 dotnet_style_parentheses_in_relational_binary_operators = never_if_unnecessary:none

src/AsmResolver.DotNet/Builder/DotNetDirectoryFactory.cs

@@ -143,6 +143,9 @@ public virtual DotNetDirectoryBuildResult CreateDotNetDirectory(
                 ReorderMetadataStreams(serializedModule, result.Directory.Metadata!);
             }
 
+            if (result.Directory.Metadata is { IsEncMetadata: true } metadata && metadata.TryGetStream("#JTD", out _))
+                result.Directory.Metadata.GetStream<TablesStream>().ForceLargeColumns = true;
+
             return result;
         }
 

src/AsmResolver.PE/DotNet/Metadata/DefaultMetadataStreamReader.cs

@@ -1,3 +1,4 @@
+using System;
 using AsmResolver.IO;
 using AsmResolver.PE.DotNet.Metadata.Blob;
 using AsmResolver.PE.DotNet.Metadata.Guid;
@@ -25,33 +26,35 @@ public static DefaultMetadataStreamReader Instance
 
         /// <inheritdoc />
         public IMetadataStream ReadStream(MetadataReaderContext context,
+            MetadataStreamReaderFlags flags,
             MetadataStreamHeader header,
             ref BinaryStreamReader reader)
         {
-            switch (header.Name)
-            {
-                case TablesStream.CompressedStreamName:
-                case TablesStream.EncStreamName:
-                    return new SerializedTableStream(context, header.Name, reader);
-
-                case StringsStream.DefaultName:
-                    return new SerializedStringsStream(header.Name, reader);
-
-                case UserStringsStream.DefaultName:
-                    return new SerializedUserStringsStream(header.Name, reader);
-
-                case BlobStream.DefaultName:
-                    return new SerializedBlobStream(header.Name, reader);
+            // The CLR performs a case-insensitive comparison for the names of the streams when ENC metadata is present.
+            var comparisonKind = (flags & MetadataStreamReaderFlags.IsEnc) != 0
+                ? StringComparison.OrdinalIgnoreCase
+                : StringComparison.Ordinal;
 
-                case GuidStream.DefaultName:
-                    return new SerializedGuidStream(header.Name, reader);
-
-                case PdbStream.DefaultName:
-                    return new SerializedPdbStream(header.Name, reader);
-
-                default:
-                    return new CustomMetadataStream(header.Name, DataSegment.FromReader(ref reader));
+            if (string.Equals(header.Name, TablesStream.CompressedStreamName, comparisonKind) ||
+                string.Equals(header.Name, TablesStream.EncStreamName, comparisonKind))
+            {
+                bool forceLargeColumns = (flags & MetadataStreamReaderFlags.IsEnc) != 0 &&
+                                         (flags & MetadataStreamReaderFlags.HasJtdStream) != 0;
+                return new SerializedTableStream(context, header.Name, reader)
+                    { ForceLargeColumns = forceLargeColumns };
             }
+            if (string.Equals(header.Name, StringsStream.DefaultName, comparisonKind))
+                return new SerializedStringsStream(header.Name, reader);
+            if (string.Equals(header.Name, UserStringsStream.DefaultName, comparisonKind))
+                return new SerializedUserStringsStream(header.Name, reader);
+            if (string.Equals(header.Name, BlobStream.DefaultName, comparisonKind))
+                return new SerializedBlobStream(header.Name, reader);
+            if (string.Equals(header.Name, GuidStream.DefaultName, comparisonKind))
+                return new SerializedGuidStream(header.Name, reader);
+            // Always perform a case-sensitive comparison for PdbStream since it is not a stream read by the CLR
+            if (header.Name == PdbStream.DefaultName)
+                return new SerializedPdbStream(header.Name, reader);
+            return new CustomMetadataStream(header.Name, DataSegment.FromReader(ref reader));
         }
     }
 }

src/AsmResolver.PE/DotNet/Metadata/IMetadataStreamReader.cs

@@ -11,9 +11,10 @@ public interface IMetadataStreamReader
         /// Reads the contents of a metadata stream.
         /// </summary>
         /// <param name="context">The reader context.</param>
+        /// <param name="flags">Flags describing the currently read metadata.</param>
         /// <param name="header">The header of the metadata stream.</param>
         /// <param name="reader">The input stream to read from.</param>
         /// <returns>The read metadata stream.</returns>
-        IMetadataStream ReadStream(MetadataReaderContext context, MetadataStreamHeader header, ref BinaryStreamReader reader);
+        IMetadataStream ReadStream(MetadataReaderContext context, MetadataStreamReaderFlags flags, MetadataStreamHeader header, ref BinaryStreamReader reader);
     }
 }

src/AsmResolver.PE/DotNet/Metadata/MetadataStreamList.cs

@@ -13,24 +13,28 @@ public class MetadataStreamList : LazyList<IMetadataStream>
         private readonly MetadataStreamHeader[] _streamHeaders;
         private readonly IMetadata _owner;
         private readonly BinaryStreamReader _directoryReader;
+        private readonly MetadataStreamReaderFlags _streamReaderFlags;
 
         /// <summary>
         /// Prepares a new lazy-initialized metadata stream list.
         /// </summary>
         /// <param name="owner">The owner of the metadata stream list.</param>
         /// <param name="context">The reader context.</param>
+        /// <param name="streamReaderFlags">Flags describing the currently read metadata.</param>
         /// <param name="streamHeaders">The stream headers.</param>
         /// <param name="directoryReader">The input stream containing the metadata directory.</param>
         public MetadataStreamList(
             IMetadata owner,
             MetadataReaderContext context,
+            MetadataStreamReaderFlags streamReaderFlags,
             MetadataStreamHeader[] streamHeaders,
             in BinaryStreamReader directoryReader)
         {
             _context = context ?? throw new ArgumentNullException(nameof(context));
             _streamHeaders = streamHeaders;
             _owner = owner;
             _directoryReader = directoryReader;
+            _streamReaderFlags = streamReaderFlags;
         }
 
         /// <inheritdoc />
@@ -42,7 +46,7 @@ protected override void Initialize()
             foreach (var header in _streamHeaders)
             {
                 var streamReader = _directoryReader.ForkAbsolute(_directoryReader.Offset + header.Offset, header.Size);
-                var stream = _context.MetadataStreamReader.ReadStream(_context, header, ref streamReader);
+                var stream = _context.MetadataStreamReader.ReadStream(_context, _streamReaderFlags, header, ref streamReader);
                 Items.Add(stream);
             }
         }

src/AsmResolver.PE/DotNet/Metadata/MetadataStreamReaderFlags.cs

@@ -0,0 +1,27 @@
+using System;
+
+namespace AsmResolver.PE.DotNet.Metadata
+{
+    /// <summary>
+    /// Special flags used to provide information about metadata to <see cref="IMetadataStreamReader"/>.
+    /// </summary>
+    [Flags]
+    public enum MetadataStreamReaderFlags
+    {
+        /// <summary>
+        /// No special flags are set.
+        /// </summary>
+        None = 0,
+        /// <summary>
+        /// EnC metadata is being used.
+        /// </summary>
+        IsEnc = 1,
+        /// <summary>
+        /// A #JTD stream is present.
+        /// </summary>
+        /// <remarks>
+        /// If the #JTD stream is present and EnC metadata is being used, all indices in the tables stream are 4 bytes.
+        /// </remarks>
+        HasJtdStream = 2,
+    }
+}

src/AsmResolver.PE/DotNet/Metadata/SerializedMetadata.cs

@@ -15,6 +15,7 @@ public class SerializedMetadata : Metadata
         private readonly MetadataReaderContext _context;
         private readonly BinaryStreamReader _streamContentsReader;
         private readonly MetadataStreamHeader[] _streamHeaders;
+        private readonly bool _hasJtdStream;
 
         /// <summary>
         /// Reads a metadata directory from an input stream.
@@ -75,12 +76,26 @@ public SerializedMetadata(MetadataReaderContext context, ref BinaryStreamReader
 
             // Eagerly read stream headers to determine if we are EnC metadata.
             _streamHeaders = new MetadataStreamHeader[numberOfStreams];
+
+            bool? isEncMetadata = null;
             for (int i = 0; i < numberOfStreams; i++)
             {
                 _streamHeaders[i] = MetadataStreamHeader.FromReader(ref directoryReader);
-                if (_streamHeaders[i].Name == TablesStream.EncStreamName)
-                    IsEncMetadata = true;
+                string name = _streamHeaders[i].Name;
+                if (isEncMetadata is null)
+                {
+                    if (name == TablesStream.CompressedStreamName)
+                        isEncMetadata = false;
+                    else if (name == TablesStream.EncStreamName)
+                        isEncMetadata = true;
+                }
+                if (name == TablesStream.UncompressedStreamName)
+                    isEncMetadata = true;
+                else if (name == TablesStream.MinimalStreamName)
+                    _hasJtdStream = true;
             }
+
+            IsEncMetadata = isEncMetadata ?? false;
         }
 
         /// <inheritdoc />
@@ -89,8 +104,15 @@ protected override IList<IMetadataStream> GetStreams()
             if (_streamHeaders.Length == 0)
                 return base.GetStreams();
 
+            var flags = MetadataStreamReaderFlags.None;
+            if (IsEncMetadata)
+                flags |= MetadataStreamReaderFlags.IsEnc;
+            if (_hasJtdStream)
+                flags |= MetadataStreamReaderFlags.HasJtdStream;
+
             return new MetadataStreamList(this,
                 _context,
+                flags,
                 _streamHeaders,
                 _streamContentsReader);
         }

src/AsmResolver.PE/DotNet/Metadata/Tables/IMetadataTable.cs

@@ -18,14 +18,6 @@ TableLayout Layout
             get;
         }
 
-        /// <summary>
-        /// Gets the size of an index into this table.
-        /// </summary>
-        IndexSize IndexSize
-        {
-            get;
-        }
-
         /// <summary>
         /// Gets or sets the row at the provided index.
         /// </summary>

src/AsmResolver.PE/DotNet/Metadata/Tables/IndexEncoder.cs

@@ -37,6 +37,9 @@ public IndexSize IndexSize
         {
             get
             {
+                if (_tableStream.ForceLargeColumns)
+                    return IndexSize.Long;
+
                 uint maxCount = 0;
                 foreach (var table in _tables)
                     maxCount = Math.Max(maxCount, _tableStream.GetTableRowCount(table));

src/AsmResolver.PE/DotNet/Metadata/Tables/MetadataTable.cs

@@ -57,9 +57,6 @@ public TableLayout Layout
             private set;
         }
 
-        /// <inheritdoc />
-        public IndexSize IndexSize => Count > 0xFFFF ? IndexSize.Long : IndexSize.Short;
-
         /// <inheritdoc cref="IMetadataTable" />
         public TRow this[int index]
         {

src/AsmResolver.PE/DotNet/Metadata/Tables/SerializedTableStream.cs

@@ -142,7 +142,7 @@ private IndexSize[] InitializeIndexSizes()
 
             // Add index sizes for each table:
             foreach (uint t in _rowCounts)
-                result.Add(t > 0xFFFF ? IndexSize.Long : IndexSize.Short);
+                result.Add(ForceLargeColumns || t > 0xFFFF ? IndexSize.Long : IndexSize.Short);
 
             // Add index sizes for each coded index:
             result.AddRange(new[]
@@ -212,6 +212,9 @@ private IndexSize GetCodedIndexSize(params TableIndex[] tables)
             if (_combinedRowCounts is null)
                 throw new InvalidOperationException("Serialized tables stream is not fully initialized yet.");
 
+            if (ForceLargeColumns)
+                return IndexSize.Long;
+
             int tableIndexBitCount = (int) Math.Ceiling(Math.Log(tables.Length, 2));
             int maxSmallTableMemberCount = ushort.MaxValue >> tableIndexBitCount;
 

src/AsmResolver.PE/DotNet/Metadata/Tables/TablesStream.cs

@@ -216,6 +216,19 @@ public uint[]? ExternalRowCounts
             set;
         }
 
+        /// <summary>
+        /// Gets or sets a value whether to force large columns in the tables stream.
+        /// </summary>
+        /// <remarks>
+        /// This value is typically <c>false</c>, it is intended to be <c>true</c> for cases when
+        /// EnC metadata is used and a stream with the <see cref="MinimalStreamName"/> name is present.
+        /// </remarks>
+        public bool ForceLargeColumns
+        {
+            get;
+            set;
+        }
+
         /// <summary>
         /// Gets a collection of all tables in the tables stream.
         /// </summary>
@@ -258,6 +271,9 @@ public uint GetTableRowCount(TableIndex table)
         /// </remarks>
         public IndexSize GetTableIndexSize(TableIndex table)
         {
+            if (ForceLargeColumns)
+                return IndexSize.Long;
+
             return GetTableRowCount(table) > 0xFFFF
                 ? IndexSize.Long
                 : IndexSize.Short;
@@ -496,7 +512,12 @@ public virtual MetadataTable<TRow> GetTable<TRow>(TableIndex index)
             return (MetadataTable<TRow>) (Tables[(int) index] ?? throw new ArgumentOutOfRangeException(nameof(index)));
         }
 
-        private IndexSize GetStreamIndexSize(int bitIndex) => (IndexSize) (((((int) Flags >> bitIndex) & 1) + 1) * 2);
+        private IndexSize GetStreamIndexSize(int bitIndex)
+        {
+            if (ForceLargeColumns)
+                return IndexSize.Long;
+            return (IndexSize)(((((int)Flags >> bitIndex) & 1) + 1) * 2);
+        }
 
         private void SetStreamIndexSize(int bitIndex, IndexSize newSize)
         {

test/AsmResolver.PE.Tests/DotNet/Metadata/MetadataTest.cs

@@ -139,7 +139,12 @@ public void PreserveMetadataNoChange()
         private void AssertCorrectStreamIsSelected<TStream>(byte[] assembly, bool isEnC)
             where TStream : class, IMetadataStream
         {
-            var peImage = PEImage.FromBytes(assembly);
+            AssertCorrectStreamIsSelected<TStream>(PEImage.FromBytes(assembly), isEnC);
+        }
+
+        private void AssertCorrectStreamIsSelected<TStream>(IPEImage peImage, bool isEnC)
+            where TStream : class, IMetadataStream
+        {
             var metadata = peImage.DotNetDirectory!.Metadata!;
 
             var allStreams = metadata.Streams
@@ -198,5 +203,51 @@ public void SelectFirstUserStringsStreamInEnCMetadata()
         {
             AssertCorrectStreamIsSelected<UserStringsStream>(Properties.Resources.HelloWorld_DoubleUserStringsStream_EnC, true);
         }
+
+        [Fact]
+        public void SchemaStreamShouldForceEnCMetadata()
+        {
+            var peImage = PEImage.FromBytes(Properties.Resources.HelloWorld_SchemaStream);
+            AssertCorrectStreamIsSelected<BlobStream>(peImage, true);
+            AssertCorrectStreamIsSelected<GuidStream>(peImage, true);
+            AssertCorrectStreamIsSelected<StringsStream>(peImage, true);
+            AssertCorrectStreamIsSelected<UserStringsStream>(peImage, true);
+        }
+
+        [Fact]
+        public void UseCaseInsensitiveComparisonForHeapNamesInEnCMetadata()
+        {
+            var peImage = PEImage.FromBytes(Properties.Resources.HelloWorld_LowerCaseHeapsWithEnC);
+            var metadata = peImage.DotNetDirectory!.Metadata!;
+
+            Assert.True(metadata.TryGetStream(out BlobStream? _));
+            Assert.True(metadata.TryGetStream(out GuidStream? _));
+            Assert.True(metadata.TryGetStream(out StringsStream? _));
+            Assert.True(metadata.TryGetStream(out UserStringsStream? _));
+        }
+
+        [Fact]
+        public void UseLargeTableIndicesWhenJTDStreamIsPresentInEnCMetadata()
+        {
+            var peImage = PEImage.FromBytes(Properties.Resources.HelloWorld_JTDStream);
+            var metadata = peImage.DotNetDirectory!.Metadata!;
+
+            var tablesStream = metadata.GetStream<TablesStream>();
+
+            Assert.True(tablesStream.ForceLargeColumns);
+
+            for (var i = TableIndex.Module; i <= TableIndex.GenericParamConstraint; i++)
+                Assert.Equal(IndexSize.Long, tablesStream.GetTableIndexSize(i));
+
+            for (var i = TableIndex.Document; i <= TableIndex.CustomDebugInformation; i++)
+                Assert.Equal(IndexSize.Long, tablesStream.GetTableIndexSize(i));
+
+            for (var i = CodedIndex.TypeDefOrRef; i <= CodedIndex.HasCustomDebugInformation; i++)
+                Assert.Equal(IndexSize.Long, tablesStream.GetIndexEncoder(i).IndexSize);
+
+            Assert.Equal(IndexSize.Long, tablesStream.StringIndexSize);
+            Assert.Equal(IndexSize.Long, tablesStream.GuidIndexSize);
+            Assert.Equal(IndexSize.Long, tablesStream.BlobIndexSize);
+        }
     }
 }