Fix round-tripping of x86/x64 code sections

filipnavara · Nov 7, 2021 · 9b55e68 · 9b55e68
1 parent a55957c
commit 9b55e68
Show file tree

Hide file tree

Showing 3 changed files with 27 additions and 6 deletions.
diff --git a/Melanzana.MachO.Tests/RoundtripTests.cs b/Melanzana.MachO.Tests/RoundtripTests.cs
@@ -31,12 +31,9 @@ public void BasicRoundtrip()
             TestRoundtrip(aOutStream);
         }
 
-        //[Fact]
+        [Fact]
         public void FatRoundtrip()
         {
-            // FIXME: This doesn't work yet because the x86-64 object file pads the code in __TEXT section with
-            // 0x90 byte that we don't roundtrip since it's not part of the section. The executable still works
-            // after rewriting though.
             var aFatOutStream = typeof(RoundtripTests).Assembly.GetManifestResourceStream("Melanzana.MachO.Tests.Data.a.fat.out")!;
             TestRoundtrip(aFatOutStream);
         }

diff --git a/Melanzana.MachO/MachWriter.cs b/Melanzana.MachO/MachWriter.cs
@@ -365,6 +365,8 @@ public static void Write(Stream stream, MachObjectFile objectFile)
                     }
                     else
                     {
+                        byte paddingByte = 0;
+
                         foreach (var section in segment.Sections)
                         {
                             if (section.IsInFile)
@@ -374,13 +376,34 @@ public static void Write(Stream stream, MachObjectFile objectFile)
                                 if (section.FileOffset > currentOffset)
                                 {
                                     ulong paddingSize = section.FileOffset - currentOffset;
-                                    stream.WritePadding((long)paddingSize);
+                                    stream.WritePadding((long)paddingSize, paddingByte);
                                     currentOffset += paddingSize;
                                 }
 
                                 using var sectionStream = section.GetReadStream();
                                 sectionStream.CopyTo(stream);
                                 currentOffset += (ulong)sectionStream.Length;
+
+                                bool isCodeSection =
+                                    section.Type == MachSectionType.Regular &&
+                                    section.Attributes.HasFlag(MachSectionAttributes.SomeInstructions) &&
+                                    section.Attributes.HasFlag(MachSectionAttributes.PureInstructions);
+
+                                if (isCodeSection)
+                                {
+                                    // Padding of code sections is done with NOP bytes if possible
+                                    paddingByte = objectFile.CpuType switch
+                                    {
+                                        // TODO: Arm32 / Thumb
+                                        MachCpuType.X86_64 => (byte)0x90,
+                                        MachCpuType.X86 => (byte)0x90,
+                                        _ => (byte)0,
+                                    };
+                                }
+                                else
+                                {
+                                    paddingByte = 0;
+                                }
                             }
                         }
                     }

diff --git a/Melanzana.Streams/StreamExtensions.cs b/Melanzana.Streams/StreamExtensions.cs
@@ -19,9 +19,10 @@ public static void ReadFully(this Stream stream, Span<byte> buffer)
                 throw new EndOfStreamException();
         }
 
-        public static void WritePadding(this Stream stream, long paddingSize)
+        public static void WritePadding(this Stream stream, long paddingSize, byte paddingByte = 0)
         {
             Span<byte> paddingBuffer = stackalloc byte[4096];
+            paddingBuffer.Fill(paddingByte);
             while (paddingSize > 0)
             {
                 long chunkSize = paddingSize > paddingBuffer.Length ? paddingBuffer.Length : paddingSize;