Simplify buffer checks for NtQueryInformation...() and NtSetInformation...() using universal macros and information class tables
Modified: trunk/reactos/ntoskrnl/include/internal/ob.h
Modified: trunk/reactos/ntoskrnl/ps/process.c

Modified: trunk/reactos/ntoskrnl/include/internal/ob.h
--- trunk/reactos/ntoskrnl/include/internal/ob.h	2005-01-21 21:34:49 UTC (rev 13195)
+++ trunk/reactos/ntoskrnl/include/internal/ob.h	2005-01-21 21:37:32 UTC (rev 13196)
@@ -339,5 +339,169 @@
                            IN KPROCESSOR_MODE AccessMode,
                            IN BOOLEAN CaptureIfKernel);
 
+/* object information classes */
 
+#define ICIF_QUERY               0x1
+#define ICIF_SET                 0x2
+#define ICIF_QUERY_SIZE_VARIABLE 0x4
+#define ICIF_SET_SIZE_VARIABLE   0x8
+#define ICIF_SIZE_VARIABLE (ICIF_QUERY_SIZE_VARIABLE | ICIF_SET_SIZE_VARIABLE)
+
+typedef struct _INFORMATION_CLASS_INFO
+{
+  ULONG RequiredSizeQUERY;
+  ULONG RequiredSizeSET;
+  ULONG AlignmentSET;
+  ULONG AlignmentQUERY;
+  ULONG Flags;
+} INFORMATION_CLASS_INFO, *PINFORMATION_CLASS_INFO;
+
+#define ICI_SQ_SAME(Size, Alignment, Flags)                                    \
+  { Size, Size, Alignment, Alignment, Flags }
+
+#define ICI_SQ(SizeQuery, SizeSet, AlignmentQuery, AlignmentSet, Flags)        \
+  { SizeQuery, SizeSet, AlignmentQuery, AlignmentSet, Flags }
+
+#define CheckInfoClass(Class, BufferLen, ClassList, StatusVar, Mode)           \
+  do {                                                                         \
+  if((Class) >= 0 && (Class) < sizeof(ClassList) / sizeof(ClassList[0]))       \
+  {                                                                            \
+    if(!(ClassList[Class].Flags & ICIF_##Mode))                                \
+    {                                                                          \
+      *(StatusVar) = STATUS_INVALID_INFO_CLASS;                                \
+    }                                                                          \
+    else if(ClassList[Class].RequiredSize##Mode > 0 &&                         \
+            (BufferLen) != ClassList[Class].RequiredSize##Mode)                \
+    {                                                                          \
+      if(!(ClassList[Class].Flags & ICIF_##Mode##_SIZE_VARIABLE) ||            \
+         ((ClassList[Class].Flags & ICIF_##Mode##_SIZE_VARIABLE) &&            \
+          (BufferLen) < ClassList[Class].RequiredSize##Mode))                  \
+      {                                                                        \
+        *(StatusVar) = STATUS_INFO_LENGTH_MISMATCH;                            \
+      }                                                                        \
+    }                                                                          \
+  }                                                                            \
+  else                                                                         \
+  {                                                                            \
+    *(StatusVar) = STATUS_INVALID_INFO_CLASS;                                  \
+  }                                                                            \
+  } while(0)
+
+#define GetInfoClassAlignment(Class, ClassList, AlignmentVar, Mode)            \
+  do {                                                                         \
+  if((Class) >= 0 && (Class) < sizeof(ClassList) / sizeof(ClassList[0]))       \
+  {                                                                            \
+    *(AlignmentVar) = ClassList[Class].Alignment##Mode;                        \
+  }                                                                            \
+  else                                                                         \
+  {                                                                            \
+    *(AlignmentVar) = sizeof(ULONG);                                           \
+  }                                                                            \
+  } while(0)
+
+#define ProbeQueryInfoBuffer(Buffer, BufferLen, Alignment, RetLen, PrevMode, StatusVar) \
+  do {                                                                         \
+  if(PrevMode == UserMode)                                                     \
+  {                                                                            \
+    _SEH_TRY                                                                   \
+    {                                                                          \
+      ProbeForWrite(Buffer,                                                    \
+                    BufferLen,                                                 \
+                    Alignment);                                                \
+      if(RetLen != NULL)                                                       \
+      {                                                                        \
+        ProbeForWrite(RetLen,                                                  \
+                      sizeof(ULONG),                                           \
+                      1);                                                      \
+      }                                                                        \
+    }                                                                          \
+    _SEH_HANDLE                                                                \
+    {                                                                          \
+      *(StatusVar) = _SEH_GetExceptionCode();                                  \
+    }                                                                          \
+    _SEH_END;                                                                  \
+                                                                               \
+    if(!NT_SUCCESS(*(StatusVar)))                                              \
+    {                                                                          \
+      DPRINT1("ProbeQueryInfoBuffer failed: 0x%x\n", *(StatusVar));            \
+      return *(StatusVar);                                                     \
+    }                                                                          \
+  }                                                                            \
+  } while(0)
+
+#define ProbeSetInfoBuffer(Buffer, BufferLen, Alignment, PrevMode, StatusVar) \
+  do {                                                                         \
+  if(PrevMode == UserMode)                                                     \
+  {                                                                            \
+    _SEH_TRY                                                                   \
+    {                                                                          \
+      ProbeForRead(Buffer,                                                     \
+                   BufferLen,                                                  \
+                   Alignment);                                                 \
+    }                                                                          \
+    _SEH_HANDLE                                                                \
+    {                                                                          \
+      *(StatusVar) = _SEH_GetExceptionCode();                                  \
+    }                                                                          \
+    _SEH_END;                                                                  \
+                                                                               \
+    if(!NT_SUCCESS(*(StatusVar)))                                              \
+    {                                                                          \
+      DPRINT1("ProbeAllInfoBuffer failed: 0x%x\n", *(StatusVar));              \
+      return *(StatusVar);                                                     \
+    }                                                                          \
+  }                                                                            \
+  } while(0)
+
+#define DefaultSetInfoBufferCheck(Class, ClassList, Buffer, BufferLen, PrevMode, StatusVar) \
+  do {                                                                         \
+  ULONG _Alignment;                                                            \
+  /* get the preferred alignment for the information class or return */        \
+  /* default alignment in case the class doesn't exist */                      \
+  GetInfoClassAlignment(Class,                                                 \
+                        ClassList,                                             \
+                        &_Alignment,                                           \
+                        SET);                                                  \
+                                                                               \
+  /* probe the ENTIRE buffers and return on failure */                         \
+  ProbeSetInfoBuffer(Buffer,                                                   \
+                     BufferLen,                                                \
+                     _Alignment,                                               \
+                     PrevMode,                                                 \
+                     StatusVar);                                               \
+                                                                               \
+  /* validate information class index and check buffer size */                 \
+  CheckInfoClass(Class,                                                        \
+                 BufferLen,                                                    \
+                 ClassList,                                                    \
+                 StatusVar,                                                    \
+                 SET);                                                         \
+  } while(0)
+
+#define DefaultQueryInfoBufferCheck(Class, ClassList, Buffer, BufferLen, RetLen, PrevMode, StatusVar) \
+  do {                                                                         \
+    ULONG _Alignment;                                                          \
+   /* get the preferred alignment for the information class or return */       \
+   /* alignment in case the class doesn't exist */                             \
+   GetInfoClassAlignment(Class,                                                \
+                         ClassList,                                            \
+                         &_Alignment,                                          \
+                         QUERY);                                               \
+                                                                               \
+   /* probe the ENTIRE buffers and return on failure */                        \
+   ProbeQueryInfoBuffer(Buffer,                                                \
+                        BufferLen,                                             \
+                        _Alignment,                                            \
+                        RetLen,                                                \
+                        PrevMode,                                              \
+                        StatusVar);                                            \
+                                                                               \
+   /* validate information class index and check buffer size */                \
+   CheckInfoClass(Class,                                                       \
+                  BufferLen,                                                   \
+                  ClassList,                                                   \
+                  StatusVar,                                                   \
+                  QUERY);                                                      \
+  } while(0)
+
 #endif /* __INCLUDE_INTERNAL_OBJMGR_H */

Modified: trunk/reactos/ntoskrnl/ps/process.c
--- trunk/reactos/ntoskrnl/ps/process.c	2005-01-21 21:34:49 UTC (rev 13195)
+++ trunk/reactos/ntoskrnl/ps/process.c	2005-01-21 21:37:32 UTC (rev 13196)
@@ -33,6 +33,44 @@
 					   STANDARD_RIGHTS_EXECUTE | SYNCHRONIZE,
 					   PROCESS_ALL_ACCESS};
 
+static const INFORMATION_CLASS_INFO PsProcessInfoClass[] =
+{
+  ICI_SQ_SAME( sizeof(PROCESS_BASIC_INFORMATION),     sizeof(ULONG), ICIF_QUERY ),                     /* ProcessBasicInformation */
+  ICI_SQ_SAME( sizeof(QUOTA_LIMITS),                  sizeof(ULONG), ICIF_QUERY | ICIF_SET ),          /* ProcessQuotaLimits */
+  ICI_SQ_SAME( sizeof(IO_COUNTERS),                   sizeof(ULONG), ICIF_QUERY ),                     /* ProcessIoCounters */
+  ICI_SQ_SAME( sizeof(VM_COUNTERS),                   sizeof(ULONG), ICIF_QUERY ),                     /* ProcessVmCounters */
+  ICI_SQ_SAME( sizeof(KERNEL_USER_TIMES),             sizeof(ULONG), ICIF_QUERY ),                     /* ProcessTimes */
+  ICI_SQ_SAME( sizeof(KPRIORITY),                     sizeof(ULONG), ICIF_SET ),                       /* ProcessBasePriority */
+  ICI_SQ_SAME( sizeof(ULONG),                         sizeof(ULONG), ICIF_SET ),                       /* ProcessRaisePriority */
+  ICI_SQ_SAME( sizeof(HANDLE),                        sizeof(ULONG), ICIF_QUERY | ICIF_SET ),          /* ProcessDebugPort */
+  ICI_SQ_SAME( sizeof(HANDLE),                        sizeof(ULONG), ICIF_SET ),                       /* ProcessExceptionPort */
+  ICI_SQ_SAME( sizeof(PROCESS_ACCESS_TOKEN),          sizeof(ULONG), ICIF_SET ),                       /* ProcessAccessToken */
+  ICI_SQ_SAME( 0 /* FIXME */,                         sizeof(ULONG), ICIF_QUERY | ICIF_SET ),          /* ProcessLdtInformation */
+  ICI_SQ_SAME( 0 /* FIXME */,                         sizeof(ULONG), ICIF_SET ),                       /* ProcessLdtSize */
+  ICI_SQ_SAME( sizeof(ULONG),                         sizeof(ULONG), ICIF_QUERY | ICIF_SET ),          /* ProcessDefaultHardErrorMode */
+  ICI_SQ_SAME( 0 /* FIXME */,                         sizeof(ULONG), ICIF_SET ),                       /* ProcessIoPortHandlers */
+  ICI_SQ_SAME( sizeof(POOLED_USAGE_AND_LIMITS),       sizeof(ULONG), ICIF_QUERY ),                     /* ProcessPooledUsageAndLimits */
+  ICI_SQ_SAME( sizeof(PROCESS_WS_WATCH_INFORMATION),  sizeof(ULONG), ICIF_QUERY | ICIF_SET ),          /* ProcessWorkingSetWatch */
+  ICI_SQ_SAME( 0 /* FIXME */,                         sizeof(ULONG), ICIF_SET ),                       /* ProcessUserModeIOPL */
+  ICI_SQ_SAME( sizeof(BOOLEAN),                       sizeof(ULONG), ICIF_SET ),                       /* ProcessEnableAlignmentFaultFixup */
+  ICI_SQ_SAME( sizeof(PROCESS_PRIORITY_CLASS),        sizeof(ULONG), ICIF_SET ),                       /* ProcessPriorityClass */
+  ICI_SQ_SAME( sizeof(ULONG),                         sizeof(ULONG), ICIF_QUERY ),                     /* ProcessWx86Information */
+  ICI_SQ_SAME( sizeof(ULONG),                         sizeof(ULONG), ICIF_QUERY ),                     /* ProcessHandleCount */
+  ICI_SQ_SAME( sizeof(KAFFINITY),                     sizeof(ULONG), ICIF_SET ),                       /* ProcessAffinityMask */
+  ICI_SQ_SAME( sizeof(ULONG),                         sizeof(ULONG), ICIF_QUERY | ICIF_SET ),          /* ProcessPriorityBoost */
+
+  ICI_SQ(/*Q*/ sizeof(((PPROCESS_DEVICEMAP_INFORMATION)0x0)->Query),                                   /* ProcessDeviceMap */
+         /*S*/ sizeof(((PPROCESS_DEVICEMAP_INFORMATION)0x0)->Set),
+                                                /*Q*/ sizeof(ULONG),
+                                                /*S*/ sizeof(ULONG),
+                                                                     ICIF_QUERY | ICIF_SET ),
+
+  ICI_SQ_SAME( sizeof(PROCESS_SESSION_INFORMATION),   sizeof(ULONG), ICIF_QUERY | ICIF_SET ),          /* ProcessSessionInformation */
+  ICI_SQ_SAME( sizeof(BOOLEAN),                       sizeof(ULONG), ICIF_SET ),                       /* ProcessForegroundInformation */
+  ICI_SQ_SAME( sizeof(ULONG),                         sizeof(ULONG), ICIF_QUERY ),                     /* ProcessWow64Information */
+  ICI_SQ_SAME( sizeof(UNICODE_STRING),                sizeof(ULONG), ICIF_QUERY | ICIF_SIZE_VARIABLE), /* ProcessImageFileName */
+};
+
 #define MAX_PROCESS_NOTIFY_ROUTINE_COUNT    8
 #define MAX_LOAD_IMAGE_NOTIFY_ROUTINE_COUNT  8
 
@@ -1133,39 +1171,22 @@
 			  OUT PULONG ReturnLength  OPTIONAL)
 {
    PEPROCESS Process;
-   NTSTATUS Status;
    KPROCESSOR_MODE PreviousMode;
-
+   NTSTATUS Status = STATUS_SUCCESS;
+   
    PreviousMode = ExGetPreviousMode();
    
-   /* check for valid buffers */
-   if(PreviousMode == UserMode)
+   DefaultQueryInfoBufferCheck(ProcessInformationClass,
+                               PsProcessInfoClass,
+                               ProcessInformation,
+                               ProcessInformationLength,
+                               ReturnLength,
+                               PreviousMode,
+                               &Status);
+   if(!NT_SUCCESS(Status))
    {
-     _SEH_TRY
-     {
-       /* probe with 32bit alignment */
-       ProbeForWrite(ProcessInformation,
-                     ProcessInformationLength,
-                     sizeof(ULONG));
-       if(ReturnLength)
-       {
-         ProbeForWrite(ReturnLength,
-                       sizeof(ULONG),
-                       1);
-       }
-
-       Status = STATUS_SUCCESS;
-     }
-     _SEH_HANDLE
-     {
-       Status = _SEH_GetExceptionCode();
-     }
-     _SEH_END;
-     
-     if(!NT_SUCCESS(Status))
-     {
-       return Status;
-     }
+     DPRINT1("NtQueryInformationProcess() failed, Status: 0x%x\n", Status);
+     return Status;
    }
 
    /*
@@ -1669,46 +1690,37 @@
 			IN ULONG ProcessInformationLength)
 {
    PEPROCESS Process;
-   NTSTATUS Status;
    KPROCESSOR_MODE PreviousMode;
    ACCESS_MASK Access;
+   NTSTATUS Status = STATUS_SUCCESS;
    
    PreviousMode = ExGetPreviousMode();
    
-   /* check for valid buffers */
-   if(PreviousMode == UserMode)
+   DefaultSetInfoBufferCheck(ProcessInformationClass,
+                             PsProcessInfoClass,
+                             ProcessInformation,
+                             ProcessInformationLength,
+                             PreviousMode,
+                             &Status);
+   if(!NT_SUCCESS(Status))
    {
-     _SEH_TRY
-     {
-       /* probe with 32bit alignment */
-       ProbeForRead(ProcessInformation,
-                    ProcessInformationLength,
-                    sizeof(ULONG));
-       Status = STATUS_SUCCESS;
-     }
-     _SEH_HANDLE
-     {
-       Status = _SEH_GetExceptionCode();
-     }
-     _SEH_END;
-
-     if(!NT_SUCCESS(Status))
-     {
-       return Status;
-     }
+     DPRINT1("NtSetInformationProcess() failed, Status: 0x%x\n", Status);
+     return Status;
    }
    
-   Access = PROCESS_SET_INFORMATION;
-   
    switch(ProcessInformationClass)
    {
      case ProcessSessionInformation:
-       Access |= PROCESS_SET_SESSIONID;
+       Access = PROCESS_SET_INFORMATION | PROCESS_SET_SESSIONID;
        break;
      case ProcessExceptionPort:
      case ProcessDebugPort:
-       Access |= PROCESS_SET_PORT;
+       Access = PROCESS_SET_INFORMATION | PROCESS_SET_PORT;
        break;
+
+     default:
+       Access = PROCESS_SET_INFORMATION;
+       break;
    }
    
    Status = ObReferenceObjectByHandle(ProcessHandle,