From bfaeba0d5ff7ddd1dfe7880026cb699575f87c67 Mon Sep 17 00:00:00 2001
From: bird_egop <sampletext32@bk.ru>
Date: Sun, 13 Apr 2025 01:39:38 +0300
Subject: [PATCH] Fixed segment override prefix handling for 0x64 (FS) opcode
 with tests

---
 X86Disassembler/X86/InstructionDecoder.cs    | 34 +++++++++++--
 X86DisassemblerTests/SegmentOverrideTests.cs | 51 ++++++++++++++++++++
 2 files changed, 81 insertions(+), 4 deletions(-)
 create mode 100644 X86DisassemblerTests/SegmentOverrideTests.cs

diff --git a/X86Disassembler/X86/InstructionDecoder.cs b/X86Disassembler/X86/InstructionDecoder.cs
index b05cc99..fcff641 100644
--- a/X86Disassembler/X86/InstructionDecoder.cs
+++ b/X86Disassembler/X86/InstructionDecoder.cs
@@ -86,7 +86,7 @@ public class InstructionDecoder
                 _addressSizePrefix = true;
                 _position++;
             }
-            else if (prefix >= 0x26 && prefix <= 0x3E && (prefix & 0x7) == 0x6) // Segment override prefix
+            else if ((prefix >= 0x26 && prefix <= 0x3E && (prefix & 0x7) == 0x6) || prefix == 0x64 || prefix == 0x65) // Segment override prefix
             {
                 _segmentOverridePrefix = true;
                 switch (prefix)
@@ -118,6 +118,21 @@ public class InstructionDecoder
         
         if (_position >= _length)
         {
+            // If we reached the end of the buffer while processing prefixes,
+            // create an instruction with just the prefix information
+            if (_segmentOverridePrefix)
+            {
+                instruction.Mnemonic = _segmentOverride;
+                instruction.Operands = "";
+                
+                // Set the raw bytes
+                int length = _position - startPosition;
+                instruction.RawBytes = new byte[length];
+                Array.Copy(_codeBuffer, startPosition, instruction.RawBytes, 0, length);
+                
+                return instruction;
+            }
+            
             return null;
         }
         
@@ -142,10 +157,21 @@ public class InstructionDecoder
             instruction.Operands = "??";
         }
         
+        // Add segment override prefix to the instruction if present
+        if (_segmentOverridePrefix && !string.IsNullOrEmpty(instruction.Operands))
+        {
+            // If the instruction has memory operands, add the segment override
+            if (instruction.Operands.Contains("["))
+            {
+                // Replace the first '[' with the segment override
+                instruction.Operands = instruction.Operands.Replace("[", $"{_segmentOverride}:[" );
+            }
+        }
+        
         // Set the raw bytes
-        int length = _position - startPosition;
-        instruction.RawBytes = new byte[length];
-        Array.Copy(_codeBuffer, startPosition, instruction.RawBytes, 0, length);
+        int instructionLength = _position - startPosition;
+        instruction.RawBytes = new byte[instructionLength];
+        Array.Copy(_codeBuffer, startPosition, instruction.RawBytes, 0, instructionLength);
         
         return instruction;
     }
diff --git a/X86DisassemblerTests/SegmentOverrideTests.cs b/X86DisassemblerTests/SegmentOverrideTests.cs
new file mode 100644
index 0000000..502ad2c
--- /dev/null
+++ b/X86DisassemblerTests/SegmentOverrideTests.cs
@@ -0,0 +1,51 @@
+namespace X86DisassemblerTests;
+
+using System;
+using Xunit;
+using X86Disassembler.X86;
+
+/// <summary>
+/// Tests for segment override prefixes
+/// </summary>
+public class SegmentOverrideTests
+{
+    /// <summary>
+    /// Tests that the FS segment override prefix (0x64) is correctly recognized
+    /// </summary>
+    [Fact]
+    public void FsSegmentOverride_IsRecognized()
+    {
+        // Arrange
+        // FS segment override prefix (0x64) followed by MOV [0], ESP (89 25 00 00 00 00)
+        byte[] codeBuffer = new byte[] { 0x64, 0x89, 0x25, 0x00, 0x00, 0x00, 0x00 };
+        var decoder = new InstructionDecoder(codeBuffer, codeBuffer.Length);
+        
+        // Act
+        var instruction = decoder.DecodeInstruction();
+        
+        // Assert
+        Assert.NotNull(instruction);
+        Assert.Equal("mov", instruction.Mnemonic);
+        Assert.Equal("esp, dword ptr fs:[0x00000000]", instruction.Operands);
+    }
+    
+    /// <summary>
+    /// Tests that the FS segment override prefix (0x64) is correctly recognized when it's the only byte
+    /// </summary>
+    [Fact]
+    public void FsSegmentOverride_Alone_IsRecognized()
+    {
+        // Arrange
+        // Just the FS segment override prefix (0x64)
+        byte[] codeBuffer = new byte[] { 0x64 };
+        var decoder = new InstructionDecoder(codeBuffer, codeBuffer.Length);
+        
+        // Act
+        var instruction = decoder.DecodeInstruction();
+        
+        // Assert
+        Assert.NotNull(instruction);
+        Assert.Equal("fs", instruction.Mnemonic);
+        Assert.Equal("", instruction.Operands);
+    }
+}