Fix NativeAOT ThunksPool thunk data block size handling. (#110732)

* Fix NativeAOT ThunksPool thunk data block size handling.

#88710 made a change in TunkPool.cs
moving away from using a page size define to calling ThunkBlockSize
to get to end of thunk data block where a common stub address get stored.

This change is not equivalent on platforms where the thunk blocks are laid out
in pair where a stub thunk blocks are followed by a data thunk block and all
gets mapped from file. This is the schema used on Windows platforms.

In that layout schema the ThunkBlockSize is 2 * page size meaning that the
calculation getting to the end of the thunk data block will move to the end
of next thunk stub that is RX memory and storing the common stub address
at that location will trigger an AV.

This works on iOS since it reports its ThunkBlockSize as one page
but that is not totally correct since it uses 2 pages, just that
they are allocated in the same way as FEATURE_RX_THUNKS, all thunk stubs
blocks followed by all thunk data blocks. The reason why this works
is because it only maps the thunk stubs from file, reporting a
ThunkBlockSize that is inline with what gets map:ed from file,
but then there is a special handling in PalAllocateThunksFromTemplate on
iOS that virutal alloc template size * 2, mapping the first
template size bytes from the file and the rest are kept as its
thunk data blocks.

This commit adds a new function returning the size of the thunk data
block and use that when calculating the end of the data block instead
of using the ThunkBlockSize since its reflects the size of each block getting
mapped from file, on Windows platforms that is stub+data, 2 * page size.

* Switch to ThunkDataBlockSizeMask in one place.

* Include pointer size slot in data block size calculation.

Co-authored-by: Jan Kotas <[email protected]>

Author: lateralusX

src/coreclr/nativeaot/System.Private.CoreLib/src/System/Runtime/ThunkPool.cs

@@ -35,6 +35,7 @@
 //
 
 using System.Diagnostics;
+using System.Numerics;
 
 namespace System.Runtime
 {
@@ -44,8 +45,10 @@ internal static class Constants
         public static readonly int ThunkCodeSize = RuntimeImports.RhpGetThunkSize();
         public static readonly int NumThunksPerBlock = RuntimeImports.RhpGetNumThunksPerBlock();
         public static readonly int NumThunkBlocksPerMapping = RuntimeImports.RhpGetNumThunkBlocksPerMapping();
-        public static readonly uint ThunkBlockSize = (uint)RuntimeImports.RhpGetThunkBlockSize();
-        public static readonly nuint ThunkBlockSizeMask = ThunkBlockSize - 1;
+        public static readonly uint ThunkCodeBlockSize = BitOperations.RoundUpToPowerOf2((uint)(ThunkCodeSize * NumThunksPerBlock));
+        public static readonly nuint ThunkCodeBlockSizeMask = ThunkCodeBlockSize - 1;
+        public static readonly uint ThunkDataBlockSize = BitOperations.RoundUpToPowerOf2((uint)(ThunkDataSize * NumThunksPerBlock + IntPtr.Size));
+        public static readonly nuint ThunkDataBlockSizeMask = ThunkDataBlockSize - 1;
     }
 
     internal class ThunksHeap
@@ -97,11 +100,11 @@ private unsafe ThunksHeap(IntPtr commonStubAddress)
                 IntPtr thunkDataBlock = RuntimeImports.RhpGetThunkDataBlockAddress(thunkStubsBlock);
 
                 // Address of the first thunk data cell should be at the beginning of the thunks data block (page-aligned)
-                Debug.Assert(((nuint)(nint)thunkDataBlock % Constants.ThunkBlockSize) == 0);
+                Debug.Assert(((nuint)(nint)thunkDataBlock % Constants.ThunkDataBlockSize) == 0);
 
                 // Update the last pointer value in the thunks data section with the value of the common stub address
-                *(IntPtr*)(thunkDataBlock + (int)(Constants.ThunkBlockSize - IntPtr.Size)) = commonStubAddress;
-                Debug.Assert(*(IntPtr*)(thunkDataBlock + (int)(Constants.ThunkBlockSize - IntPtr.Size)) == commonStubAddress);
+                *(IntPtr*)(thunkDataBlock + (int)(Constants.ThunkDataBlockSize - IntPtr.Size)) = commonStubAddress;
+                Debug.Assert(*(IntPtr*)(thunkDataBlock + (int)(Constants.ThunkDataBlockSize - IntPtr.Size)) == commonStubAddress);
 
                 // Set the head and end of the linked list
                 _nextAvailableThunkPtr = thunkDataBlock;
@@ -153,11 +156,11 @@ private unsafe bool ExpandHeap()
                 IntPtr thunkDataBlock = RuntimeImports.RhpGetThunkDataBlockAddress(thunkStubsBlock);
 
                 // Address of the first thunk data cell should be at the beginning of the thunks data block (page-aligned)
-                Debug.Assert(((nuint)(nint)thunkDataBlock % Constants.ThunkBlockSize) == 0);
+                Debug.Assert(((nuint)(nint)thunkDataBlock % Constants.ThunkDataBlockSize) == 0);
 
                 // Update the last pointer value in the thunks data section with the value of the common stub address
-                *(IntPtr*)(thunkDataBlock + (int)(Constants.ThunkBlockSize - IntPtr.Size)) = _commonStubAddress;
-                Debug.Assert(*(IntPtr*)(thunkDataBlock + (int)(Constants.ThunkBlockSize - IntPtr.Size)) == _commonStubAddress);
+                *(IntPtr*)(thunkDataBlock + (int)(Constants.ThunkDataBlockSize - IntPtr.Size)) = _commonStubAddress;
+                Debug.Assert(*(IntPtr*)(thunkDataBlock + (int)(Constants.ThunkDataBlockSize - IntPtr.Size)) == _commonStubAddress);
 
                 // Link the last entry in the old list to the first entry in the new list
                 *((IntPtr*)_lastThunkPtr) = thunkDataBlock;
@@ -210,7 +213,7 @@ public unsafe IntPtr AllocateThunk()
             *((IntPtr*)(nextAvailableThunkPtr + IntPtr.Size)) = IntPtr.Zero;
 #endif
 
-            int thunkIndex = (int)(((nuint)(nint)nextAvailableThunkPtr) - ((nuint)(nint)nextAvailableThunkPtr & ~Constants.ThunkBlockSizeMask));
+            int thunkIndex = (int)(((nuint)(nint)nextAvailableThunkPtr) - ((nuint)(nint)nextAvailableThunkPtr & ~Constants.ThunkDataBlockSizeMask));
             Debug.Assert((thunkIndex % Constants.ThunkDataSize) == 0);
             thunkIndex /= Constants.ThunkDataSize;
 
@@ -266,7 +269,7 @@ private static IntPtr TryGetThunkDataAddress(IntPtr thunkAddress)
             nuint thunkAddressValue = (nuint)(nint)ClearThumbBit(thunkAddress);
 
             // Compute the base address of the thunk's mapping
-            nuint currentThunksBlockAddress = thunkAddressValue & ~Constants.ThunkBlockSizeMask;
+            nuint currentThunksBlockAddress = thunkAddressValue & ~Constants.ThunkCodeBlockSizeMask;
 
             // Make sure the thunk address is valid by checking alignment
             if ((thunkAddressValue - currentThunksBlockAddress) % (nuint)Constants.ThunkCodeSize != 0)