Add bounds checking to affinity map bit operations and introduce InitializeAffinityMap

sdk/xtdk/xtstatus.h

@@ -60,6 +60,7 @@
 #define STATUS_INVALID_PARAMETER                                           ((XTSTATUS) 0xC000000DL)
 #define STATUS_END_OF_FILE                                                 ((XTSTATUS) 0xC0000011L)
 #define STATUS_NO_MEMORY                                                   ((XTSTATUS) 0xC0000017L)
+#define STATUS_BUFFER_TOO_SMALL                                            ((XTSTATUS) 0xC0000023L)
 #define STATUS_PORT_DISCONNECTED                                           ((XTSTATUS) 0xC0000037L)
 #define STATUS_CRC_ERROR                                                   ((XTSTATUS) 0xC000003FL)
 #define STATUS_FLOAT_OVERFLOW                                              ((XTSTATUS) 0xC0000091L)

xtoskrnl/includes/ke/affinity.hh

@@ -29,7 +29,7 @@ namespace KE
             STATIC XTFASTCALL VOID ClearProcessorAffinity(IN OUT PKAFFINITY_MAP AffinityMap,
                                                           IN ULONG CpuNumber);
 
-            STATIC XTAPI VOID CopyAffinity(OUT PKAFFINITY_MAP Destination,
+            STATIC XTAPI XTSTATUS CopyAffinity(OUT PKAFFINITY_MAP Destination,
                                                IN PKAFFINITY_MAP Source);
             STATIC XTAPI XTSTATUS CreateAffinityMap(IN ULONG CpuCount,
                                                     OUT PKAFFINITY_MAP* AffinityMap);
@@ -37,6 +37,8 @@ namespace KE
                                                         IN PKAFFINITY_MAP AffinityMap);
             STATIC XTAPI ULONG FindNextRightSetProcessor(IN ULONG ThreadSeed,
                                                          IN PKAFFINITY_MAP AffinityMap);
+            STATIC XTAPI XTSTATUS InitializeAffinityMap(IN OUT PKAFFINITY_MAP AffinityMap,
+                                                        IN ULONG BufferSize);
             STATIC XTFASTCALL VOID SetAllProcessorsAffinity(IN OUT PKAFFINITY_MAP AffinityMap);
             STATIC XTFASTCALL VOID SetProcessorAffinity(IN OUT PKAFFINITY_MAP AffinityMap,
                                                         IN ULONG CpuNumber);

xtoskrnl/ke/affinity.cc

@@ -26,10 +26,14 @@ XTFASTCALL
 VOID
 KE::Affinity::AtomicSetProcessorAffinity(IN OUT PKAFFINITY_MAP AffinityMap,
                                          IN ULONG CpuNumber)
+{
+    /* Verify that the target processor falls within the allocated map boundaries */
+    if((CpuNumber / 64) < AffinityMap->Size)
     {
         /* Atomically set the target CPU bit */
         RTL::Atomic::Or64((PLONG_PTR)&AffinityMap->Bitmap[CpuNumber / 64], ((KAFFINITY)1 << (CpuNumber % 64)));
     }
+}
 
 /**
  * Computes the memory size required to allocate an affinity map.
@@ -70,7 +74,7 @@ KE::Affinity::CalculateAffinityMapSize(IN ULONG CpuCount,
     if(RequiredBlockCount != NULLPTR)
     {
         /* Return the required logical block count */
-        *RequiredBlockCount = (USHORT)AffinitySize;
+        *RequiredBlockCount = AffinitySize;
     }
 }
 
@@ -91,11 +95,18 @@ XTFASTCALL
 BOOLEAN
 KE::Affinity::CheckProcessorAffinity(IN PKAFFINITY_MAP AffinityMap,
                                      IN ULONG CpuNumber)
+{
+    /* Verify that the target processor falls within the allocated map boundaries */
+    if((CpuNumber / 64) < AffinityMap->Size)
     {
         /* Isolate and test the specific bit corresponding to the target CPU */
         return (AffinityMap->Bitmap[CpuNumber / 64] & ((KAFFINITY)1 << (CpuNumber % 64))) != 0;
     }
 
+    /* Return FALSE if the requested CPU exceeds the map capacity */
+    return FALSE;
+}
+
 /**
  * Clears all processor bindings from the given affinity map.
  *
@@ -131,10 +142,14 @@ XTFASTCALL
 VOID
 KE::Affinity::ClearProcessorAffinity(IN OUT PKAFFINITY_MAP AffinityMap,
                                      IN ULONG CpuNumber)
+{
+    /* Verify that the target processor falls within the allocated map boundaries */
+    if((CpuNumber / 64) < AffinityMap->Size)
     {
         /* Clear the target CPU bit in the affinity map */
         AffinityMap->Bitmap[CpuNumber / 64] &= ~((KAFFINITY)1 << (CpuNumber % 64));
     }
+}
 
 /**
  * Copies the topological layout and processor bindings from a source affinity map to a destination map.
@@ -150,15 +165,20 @@ KE::Affinity::ClearProcessorAffinity(IN OUT PKAFFINITY_MAP AffinityMap,
  * @since XT 1.0
  */
 XTAPI
-VOID
+XTSTATUS
 KE::Affinity::CopyAffinity(OUT PKAFFINITY_MAP Destination,
                            IN PKAFFINITY_MAP Source)
 {
-    USHORT Index;
+    ULONG Index;
+
+    /* Validate that the destination map can accommodate the source topology */
+    if(Destination->Size < Source->Size)
+    {
+        /* Buffer overrun prevention */
+        return STATUS_BUFFER_TOO_SMALL;
+    }
 
     /* Copy map metadata */
-    Destination->Count = Source->Count;
-    Destination->Size = Source->Size;
     Destination->Reserved = Source->Reserved;
 
     /* Copy the active affinity bitmasks */
@@ -167,6 +187,9 @@ KE::Affinity::CopyAffinity(OUT PKAFFINITY_MAP Destination,
         /* Replicate the hardware topology bindings across all active array elements */
         Destination->Bitmap[Index] = Source->Bitmap[Index];
     }
+
+    /* Return success */
+    return STATUS_SUCCESS;
 }
 
 /**
@@ -207,7 +230,6 @@ KE::Affinity::CreateAffinityMap(IN ULONG CpuCount,
 
     /* Initialize the internal metadata required by iteration and validation routines */
     AllocatedMap->Size = BlockCount;
-    AllocatedMap->Count = BlockCount;
 
     /* Return the constructed map to the caller */
     *AffinityMap = AllocatedMap;
@@ -225,7 +247,7 @@ KE::Affinity::CreateAffinityMap(IN ULONG CpuCount,
  * @param AffinityMap
  *        Supplies a pointer to the extended affinity map defining the permitted processors.
  *
- * @return Returns the absolute topological index of the selected processor.
+ * @return This routine returns the absolute topological index of the selected processor.
  *
  * @since XT 1.0
  */
@@ -323,7 +345,7 @@ KE::Affinity::FindNextLeftSetProcessor(IN ULONG ThreadSeed,
  * @param AffinityMap
  *        Supplies a pointer to the extended affinity map defining the permitted processors.
  *
- * @return Returns the absolute topological index of the selected processor.
+ * @return This routine returns the absolute topological index of the selected processor.
  *
  * @since XT 1.0
  */
@@ -404,6 +426,50 @@ KE::Affinity::FindNextRightSetProcessor(IN ULONG ThreadSeed,
     return 0;
 }
 
+/**
+ * Initializes a caller-allocated affinity map structure.
+ *
+ * @param AffinityMap
+ *        Supplies a pointer to the caller-allocated buffer to be initialized.
+ *
+ * @param BufferSize
+ *        Supplies the size, in bytes, of the provided buffer.
+ *
+ * @return This routine returns a status code indicating the success or failure of the operation.
+ *
+ * @since XT 1.0
+ */
+XTAPI
+XTSTATUS
+KE::Affinity::InitializeAffinityMap(IN OUT PKAFFINITY_MAP AffinityMap,
+                                    IN ULONG BufferSize)
+{
+    ULONG Cpus, RequiredBlocks, RequiredSize;
+
+    /* Get the number of available CPUs */
+    Cpus = KE::Processor::GetInstalledCpus();
+
+    /* Query the required size in bytes and the architectural block count */
+    KE::Affinity::CalculateAffinityMapSize(Cpus, &RequiredSize, &RequiredBlocks);
+
+    /* Validate that the provided buffer is large enough to hold the structure */
+    if(BufferSize < RequiredSize)
+    {
+        /* Buffer overrun prevention */
+        return STATUS_BUFFER_TOO_SMALL;
+    }
+
+    /* Zero only the required portion of the memory to clear any garbage data */
+    RTL::Memory::ZeroMemory(AffinityMap, RequiredSize);
+
+    /* Initialize the internal metadata required by iteration routines */
+    AffinityMap->Reserved = 0;
+    AffinityMap->Size = RequiredBlocks;
+
+    /* Return success */
+    return STATUS_SUCCESS;
+}
+
 /**
  * Populates the affinity map with all available processors.
  *
@@ -464,7 +530,11 @@ XTFASTCALL
 VOID
 KE::Affinity::SetProcessorAffinity(IN OUT PKAFFINITY_MAP AffinityMap,
                                    IN ULONG CpuNumber)
+{
+    /* Verify that the target processor falls within the allocated map boundaries */
+    if((CpuNumber / 64) < AffinityMap->Size)
     {
         /* Set the target CPU bit in the affinity map */
         AffinityMap->Bitmap[CpuNumber / 64] |= ((KAFFINITY)1 << (CpuNumber % 64));
     }
+}