Merge pull request #346 from Washi1337/feature/imgbase-in-update-offsets

Refactor ISegment::UpdateOffsets

docs/peimage/tls.rst

@@ -6,7 +6,7 @@ Executables that use multiple threads might require static (non-stack) memory th
 All code relevant to the TLS data directory of a PE resides in the following namespace:
 
 .. code-block:: csharp
-    
+
     using AsmResolver.PE.Tls;
 
 
@@ -21,7 +21,7 @@ The PE file format defines a segment of memory within the TLS data directory tha
 
 
 .. code-block:: csharp
-    
+
     var indexSegment = new DataSegment(new byte[8]);
 
     var directory = new TlsDirectory
@@ -48,24 +48,20 @@ Next to static initialization data, it is also possible to specify a list of fun
 Creating new TLS directories
 ----------------------------
 
-Since the TLS data directory stores its data using virtual addresses (VA) rather than relative virtual addresses (RVA), AsmResolver requires the image base as well as the pointer size. This is done through the ``ImageBase`` and ``Is32Bit`` properties. By default, the following values are assumed:
+Adding a new TLS directory to an image can be done using the parameterless constructor of the ``TlsDirectory`` class:
 
 .. code-block:: csharp
 
-    var directory = new TlsDirectory();
-    directory.ImageBase = 0x00400000;
-    directory.Is32Bit = true;
-
-
-Typically, you should make sure they are in sync with the values found in the file and optional header of the final PE file. Upon reading from an existing PE file, these two properties are initialized to the values stored in these two headers.
+    var tlsDirectory = new TlsDirectory();
+    image.TlsDirectory = tlsDirectory;
 
-When building a relocatable PE file, you might also need to add base address relocations to the VAs inside the TLS directory. To quickly get all the base relocations required, use the ``GetRequiredBaseRelocations`` method:
+A TLS directory references all its sub segments using virtual addresses (VA) rather than relative addresses (RVA). This means that constructing a relocatable PE image with a TLS directory requires base relocation entries to be registered that let the Windows PE loader rebase all virtual addresses used in the directory when necessary. To quickly register all the required base relocations, you can call the ``GetRequiredBaseRelocations`` method and add all returned entries to the base relocation directory of the PE image:
 
 .. code-block:: csharp
 
     using AsmResolver.PE.Relocations;
 
     IPEImage image = ...;
 
-    foreach (var relocation in image.TlsDirectory.GetRequiredBaseRelocations())
-        image.Relocations.Add(relocation);
+    foreach (var relocation in tlsDirectory.GetRequiredBaseRelocations())
+        image.Relocations.Add(relocation);

src/AsmResolver.DotNet/Builder/ManagedPEImageBuilder.cs

@@ -70,7 +70,7 @@ public PEImageBuildResult CreateImage(ModuleDefinition module)
                 };
 
                 // Construct new .NET directory.
-                var symbolProvider = new NativeSymbolsProvider(image.ImageBase);
+                var symbolProvider = new NativeSymbolsProvider();
                 var result = DotNetDirectoryFactory.CreateDotNetDirectory(
                     module,
                     symbolProvider,

src/AsmResolver.DotNet/Builder/VTableFixups/VTableFixupsBuffer.cs

@@ -61,7 +61,7 @@ public void MapTokenToExport(UnmanagedExportInfo exportInfo, MetadataToken token
             vtableFixup.Tokens.Add(token);
             var vtableSymbol = new Symbol(vtableFixup.Tokens.GetReferenceToIndex(vtableFixup.Tokens.Count - 1));
 
-            var thunkStub = _targetPlatform.CreateThunkStub(0x00400000, vtableSymbol);
+            var thunkStub = _targetPlatform.CreateThunkStub(vtableSymbol);
 
             // Register exported symbol.
             var stubReference = thunkStub.Segment.ToReference();

src/AsmResolver.DotNet/Bundles/BundleManifest.cs

@@ -445,7 +445,7 @@ private void WriteFileContents(IBinaryStreamWriter writer, uint alignment)
                 if (file.Type == BundleFileType.Assembly)
                     writer.Align(alignment);
 
-                file.Contents.UpdateOffsets(writer.Offset, (uint) writer.Offset);
+                file.Contents.UpdateOffsets(new RelocationParameters(writer.Offset, (uint) writer.Offset));
                 file.Contents.Write(writer);
             }
         }

src/AsmResolver.DotNet/Code/Native/INativeSymbolsProvider.cs

@@ -8,14 +8,6 @@ namespace AsmResolver.DotNet.Code.Native
     /// </summary>
     public interface INativeSymbolsProvider
     {
-        /// <summary>
-        /// Gets or sets the image base the final PE image is using.
-        /// </summary>
-        ulong ImageBase
-        {
-            get;
-        }
-
         /// <summary>
         /// Adds a single symbol to the prototype.
         /// </summary>

src/AsmResolver.DotNet/Code/Native/NativeMethodBodySerializer.cs

@@ -18,7 +18,7 @@ public ISegmentReference SerializeMethodBody(MethodBodySerializationContext cont
             var provider = context.SymbolsProvider;
 
             // Create new raw code segment containing the native code.
-            var segment = new CodeSegment(provider.ImageBase, nativeMethodBody.Code);
+            var segment = new CodeSegment(nativeMethodBody.Code);
 
             // Process fixups.
             for (int i = 0; i < nativeMethodBody.AddressFixups.Count; i++)

src/AsmResolver.DotNet/Code/Native/NativeSymbolsProvider.cs

@@ -22,21 +22,6 @@ public class NativeSymbolsProvider : INativeSymbolsProvider
         private uint _maxExportedOrdinal = 0;
         private readonly List<ExportedSymbol> _floatingExportedSymbols = new();
 
-        /// <summary>
-        /// Creates a new instance of the <see cref="NativeSymbolsProvider"/> class.
-        /// </summary>
-        /// <param name="imageBase">The image base of the final PE image.</param>
-        public NativeSymbolsProvider(ulong imageBase)
-        {
-            ImageBase = imageBase;
-        }
-
-        /// <inheritdoc />
-        public ulong ImageBase
-        {
-            get;
-        }
-
         /// <inheritdoc />
         public ISymbol ImportSymbol(ISymbol symbol)
         {

src/AsmResolver.PE.File/PEFile.cs

@@ -371,12 +371,15 @@ public void UpdateHeaders()
 
             FileHeader.NumberOfSections = (ushort) Sections.Count;
 
-            FileHeader.UpdateOffsets(
+            var relocation = new RelocationParameters(OptionalHeader.ImageBase, 0, 0,
+                OptionalHeader.Magic == OptionalHeaderMagic.Pe32);
+
+            FileHeader.UpdateOffsets(relocation.WithOffsetRva(
                 DosHeader.NextHeaderOffset + 4,
-                DosHeader.NextHeaderOffset + 4);
-            OptionalHeader.UpdateOffsets(
+                DosHeader.NextHeaderOffset + 4));
+            OptionalHeader.UpdateOffsets(relocation.WithOffsetRva(
                 FileHeader.Offset + FileHeader.GetPhysicalSize(),
-                FileHeader.Rva + FileHeader.GetVirtualSize());
+                FileHeader.Rva + FileHeader.GetVirtualSize()));
 
             FileHeader.SizeOfOptionalHeader = (ushort) OptionalHeader.GetPhysicalSize();
             OptionalHeader.SizeOfHeaders = (uint) (OptionalHeader.Offset
@@ -391,27 +394,30 @@ public void UpdateHeaders()
             OptionalHeader.SizeOfImage = lastSection.Rva
                                          + lastSection.GetVirtualSize().Align(OptionalHeader.SectionAlignment);
 
-            EofData?.UpdateOffsets(lastSection.Offset + lastSection.GetPhysicalSize(), OptionalHeader.SizeOfImage);
+            EofData?.UpdateOffsets(relocation.WithOffsetRva(
+                lastSection.Offset + lastSection.GetPhysicalSize(),
+                OptionalHeader.SizeOfImage));
         }
 
         /// <summary>
         /// Aligns all sections according to the file and section alignment properties in the optional header.
         /// </summary>
         public void AlignSections()
         {
-            uint currentFileOffset = OptionalHeader.SizeOfHeaders;
+            var relocation = new RelocationParameters(
+                OptionalHeader.ImageBase,
+                OptionalHeader.SizeOfHeaders.Align(OptionalHeader.FileAlignment),
+                OptionalHeader.SizeOfHeaders.Align(OptionalHeader.SectionAlignment),
+                OptionalHeader.Magic == OptionalHeaderMagic.Pe32);
 
             for (int i = 0; i < Sections.Count; i++)
             {
                 var section = Sections[i];
 
-                uint rva = i > 0
-                    ? Sections[i - 1].Rva + Sections[i - 1].GetVirtualSize()
-                    : OptionalHeader.SizeOfHeaders.Align(OptionalHeader.SectionAlignment);
-
-                currentFileOffset = currentFileOffset.Align(OptionalHeader.FileAlignment);
-                section.UpdateOffsets(currentFileOffset, rva.Align(OptionalHeader.SectionAlignment));
-                currentFileOffset += section.GetPhysicalSize();
+                section.UpdateOffsets(relocation);
+                relocation.Advance(
+                    section.GetPhysicalSize().Align(OptionalHeader.FileAlignment),
+                    section.GetVirtualSize().Align(OptionalHeader.SectionAlignment));
             }
         }
 

src/AsmResolver.PE.File/PESection.cs

@@ -216,7 +216,7 @@ public ISegment? Contents
         public bool CanUpdateOffsets => true;
 
         /// <inheritdoc />
-        public void UpdateOffsets(ulong newOffset, uint newRva) => Contents?.UpdateOffsets(newOffset, newRva);
+        public void UpdateOffsets(in RelocationParameters parameters) => Contents?.UpdateOffsets(parameters);
 
         /// <inheritdoc />
         public uint GetPhysicalSize() => Contents?.GetPhysicalSize() ?? 0;

src/AsmResolver.PE.File/PESegmentReference.cs

@@ -32,18 +32,12 @@ public uint Rva
             get;
         }
 
-        /// <inheritdoc />
-        bool IOffsetProvider.CanUpdateOffsets => false;
-
         /// <inheritdoc />
         public bool CanRead => _peFile.TryGetSectionContainingRva(Rva, out _);
 
         /// <inheritdoc />
         public bool IsBounded => false;
 
-        /// <inheritdoc />
-        void IOffsetProvider.UpdateOffsets(ulong newOffset, uint newRva) => throw new InvalidOperationException();
-
         /// <inheritdoc />
         public BinaryStreamReader CreateReader() => _peFile.CreateReaderAtRva(Rva);
 

src/AsmResolver.PE.File/SerializedPEFile.cs

@@ -12,6 +12,8 @@ public class SerializedPEFile : PEFile
     {
         private readonly List<SectionHeader> _sectionHeaders;
         private readonly BinaryStreamReader _reader;
+        private readonly ulong _originalImageBase;
+        private readonly bool _is32Bit;
 
         /// <summary>
         /// Reads a PE file from an input stream.
@@ -36,6 +38,8 @@ public SerializedPEFile(in BinaryStreamReader reader, PEMappingMode mode)
             // Read NT headers.
             FileHeader = FileHeader.FromReader(ref _reader);
             OptionalHeader = OptionalHeader.FromReader(ref _reader);
+            _originalImageBase = OptionalHeader.ImageBase;
+            _is32Bit = OptionalHeader.Magic == OptionalHeaderMagic.Pe32;
 
             // Read section headers.
             _reader.Offset = OptionalHeader.Offset + FileHeader.SizeOfOptionalHeader;
@@ -70,7 +74,12 @@ protected override IList<PESection> GetSections()
                     physicalContents = new DataSourceSegment(_reader.DataSource, offset, header.VirtualAddress, size);
 
                 var virtualSegment = new VirtualSegment(physicalContents, header.VirtualSize);
-                virtualSegment.UpdateOffsets(offset, header.VirtualAddress);
+                virtualSegment.UpdateOffsets(new RelocationParameters(
+                    _originalImageBase,
+                    offset,
+                    header.VirtualAddress,
+                    _is32Bit));
+
                 result.Add(new PESection(header, virtualSegment));
             }
 

src/AsmResolver.PE.Win32Resources/Version/FixedVersionInfo.cs

@@ -24,7 +24,7 @@ public class FixedVersionInfo : SegmentBase
         public static FixedVersionInfo FromReader(ref BinaryStreamReader reader)
         {
             var result = new FixedVersionInfo();
-            result.UpdateOffsets(reader.Offset, reader.Rva);
+            result.UpdateOffsets(new RelocationParameters(reader.Offset, reader.Rva));
 
             uint signature = reader.ReadUInt32();
             if (signature != Signature)

src/AsmResolver.PE/Code/CodeSegment.cs

@@ -56,6 +56,13 @@ public IList<AddressFixup> AddressFixups
             get;
         } = new List<AddressFixup>();
 
+        /// <inheritdoc />
+        public override void UpdateOffsets(in RelocationParameters parameters)
+        {
+            ImageBase = parameters.ImageBase;
+            base.UpdateOffsets(in parameters);
+        }
+
         /// <inheritdoc />
         public override uint GetPhysicalSize() => (uint) Code.Length;
 

src/AsmResolver.PE/Debug/Builder/DebugDirectoryBuffer.cs

@@ -41,7 +41,7 @@ public void AddEntry(DebugDataEntry entry)
         public bool CanUpdateOffsets => true;
 
         /// <inheritdoc />
-        public void UpdateOffsets(ulong newOffset, uint newRva) => _headers.UpdateOffsets(newOffset, newRva);
+        public void UpdateOffsets(in RelocationParameters parameters) => _headers.UpdateOffsets(parameters);
 
         /// <inheritdoc />
         public uint GetPhysicalSize() => _headers.GetPhysicalSize();

src/AsmResolver.PE/Debug/CustomDebugDataSegment.cs

@@ -45,10 +45,10 @@ public ISegment? Contents
         public bool CanUpdateOffsets => Contents?.CanUpdateOffsets ?? false;
 
         /// <inheritdoc />
-        public void UpdateOffsets(ulong newOffset, uint newRva)
+        public void UpdateOffsets(in RelocationParameters parameters)
         {
             if (Contents != null)
-                Contents.UpdateOffsets(newOffset, newRva);
+                Contents.UpdateOffsets(parameters);
             else
                 throw new ArgumentNullException(nameof(Contents));
         }

src/AsmResolver.PE/DotNet/Builder/DotNetSegmentBuffer.cs

@@ -73,7 +73,7 @@ private void AddIfPresent(ISegment? segment)
         }
 
         /// <inheritdoc />
-        public void UpdateOffsets(ulong newOffset, uint newRva) => _segments.UpdateOffsets(newOffset, newRva);
+        public void UpdateOffsets(in RelocationParameters parameters) => _segments.UpdateOffsets(parameters);
 
         /// <inheritdoc />
         public uint GetPhysicalSize() => _segments.GetPhysicalSize();

src/AsmResolver.PE/DotNet/Builder/ManagedPEFileBuilder.cs

@@ -262,7 +262,7 @@ private static void CreateImportDirectory(IPEImage image, ManagedPEBuilderContex
                 if (entrypointSymbol is null)
                     throw new InvalidOperationException("Entrypoint symbol was required but not imported.");
 
-                context.Bootstrapper = context.Platform.CreateThunkStub(image.ImageBase, entrypointSymbol);
+                context.Bootstrapper = context.Platform.CreateThunkStub(entrypointSymbol);
             }
         }
 

src/AsmResolver.PE/DotNet/Builder/MethodBodyTableBuffer.cs

@@ -53,7 +53,7 @@ public void AddCilBody(CilRawMethodBody body)
         public void AddNativeBody(ISegment body, uint alignment) => _nativeBodies.Add(body, alignment);
 
         /// <inheritdoc />
-        public void UpdateOffsets(ulong newOffset, uint newRva) => _segments.UpdateOffsets(newOffset, newRva);
+        public void UpdateOffsets(in RelocationParameters parameters) => _segments.UpdateOffsets(parameters);
 
         /// <inheritdoc />
         public uint GetPhysicalSize()

src/AsmResolver.PE/DotNet/Cil/CilRawFatMethodBody.cs

@@ -153,7 +153,7 @@ public IList<CilExtraSection> ExtraSections
 
             // Create body.
             var body = new CilRawFatMethodBody(flags, maxStack, localVarSigToken, code);
-            body.UpdateOffsets(fileOffset, rva);
+            body.UpdateOffsets(new RelocationParameters(fileOffset, rva));
 
             // Read any extra sections.
             if (body.HasSections)
@@ -172,18 +172,24 @@ public IList<CilExtraSection> ExtraSections
         }
 
         /// <inheritdoc />
-        public override void UpdateOffsets(ulong newOffset, uint newRva)
+        public override void UpdateOffsets(in RelocationParameters parameters)
         {
-            base.UpdateOffsets(newOffset, newRva);
-            Code.UpdateOffsets(newOffset + 12, newRva + 12);
+            base.UpdateOffsets(parameters);
+
+            var current = parameters.WithAdvance(12);
+
+            Code.UpdateOffsets(current);
             if (HasSections)
             {
                 uint codeSize = Code.GetPhysicalSize();
-                newOffset = (Code.Offset + codeSize).Align(4);
-                newRva = (Code.Rva + codeSize).Align(4);
+                current.Advance(codeSize);
+                current.Align(4);
 
                 for (int i = 0; i < ExtraSections.Count; i++)
-                    ExtraSections[i].UpdateOffsets(newOffset, newRva);
+                {
+                    ExtraSections[i].UpdateOffsets(current);
+                    current.Advance(ExtraSections[i].GetPhysicalSize());
+                }
             }
         }
 

src/AsmResolver.PE/DotNet/Cil/CilRawTinyMethodBody.cs

@@ -58,15 +58,15 @@ public CilRawTinyMethodBody(IReadableSegment code)
 
             uint codeSize = (uint) flag >> 2;
             var methodBody = new CilRawTinyMethodBody(reader.ReadSegment(codeSize));
-            methodBody.UpdateOffsets(fileOffset, rva);
+            methodBody.UpdateOffsets(new RelocationParameters(fileOffset, rva));
             return methodBody;
         }
 
         /// <inheritdoc />
-        public override void UpdateOffsets(ulong newOffset, uint newRva)
+        public override void UpdateOffsets(in RelocationParameters parameters)
         {
-            base.UpdateOffsets(newOffset, newRva);
-            Code.UpdateOffsets(newOffset + 1, newRva + 1);
+            base.UpdateOffsets(parameters);
+            Code.UpdateOffsets(parameters.WithAdvance(1));
         }
 
         /// <inheritdoc />

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

@@ -68,7 +68,7 @@ public BinaryStreamReader CreateReader()
         }
 
         /// <inheritdoc />
-        public void UpdateOffsets(ulong newOffset, uint newRva) => Contents.UpdateOffsets(newOffset, newRva);
+        public void UpdateOffsets(in RelocationParameters parameters) => Contents.UpdateOffsets(parameters);
 
         /// <inheritdoc />
         public uint GetPhysicalSize() => Contents.GetPhysicalSize();

src/AsmResolver.PE/DotNet/SerializedDotNetDirectory.cs

@@ -126,7 +126,7 @@ public SerializedDotNetDirectory(PEReaderContext context, ref BinaryStreamReader
             }
 
             var vtables = new VTableFixupsDirectory();
-            vtables.UpdateOffsets(directoryReader.Offset, directoryReader.Rva);
+            vtables.UpdateOffsets(_context.GetRelocation(directoryReader.Offset, directoryReader.Rva));
 
             for (int i = 0; i < directoryReader.Length / 8; i++)
             {