Fix kernel-mode executive atom implementation (mostly add SEH and tidy up the code...
[reactos.git] / reactos / ntoskrnl / ex / atom.c
index be573da..0326d51 100644 (file)
@@ -13,6 +13,8 @@
 #define NDEBUG\r
 #include <internal/debug.h>\r
 \r
+#define TAG_ATOM TAG('A', 't', 'o', 'm')\r
+\r
 /* GLOBALS ****************************************************************/\r
 \r
 /* \r
@@ -45,79 +47,6 @@ ExpGetGlobalAtomTable(VOID)
     return GlobalAtomTable;\r
 }\r
 \r
-NTSTATUS\r
-NTAPI\r
-RtlpQueryAtomInformation(PRTL_ATOM_TABLE AtomTable,\r
-                         RTL_ATOM Atom,\r
-                         PATOM_BASIC_INFORMATION AtomInformation,\r
-                         ULONG AtomInformationLength,\r
-                         PULONG ReturnLength)\r
-{\r
-    NTSTATUS Status;\r
-    ULONG UsageCount;\r
-    ULONG Flags;\r
-    ULONG NameLength;\r
-\r
-    NameLength = AtomInformationLength - sizeof(ATOM_BASIC_INFORMATION) + sizeof(WCHAR);\r
-    Status = RtlQueryAtomInAtomTable(AtomTable,\r
-                                     Atom,\r
-                                     &UsageCount,\r
-                                     &Flags,\r
-                                     AtomInformation->Name,\r
-                                     &NameLength);\r
-\r
-    if (!NT_SUCCESS(Status)) return Status;\r
-    DPRINT("NameLength: %lu\n", NameLength);\r
-\r
-    if (ReturnLength != NULL)\r
-    {\r
-        *ReturnLength = NameLength + sizeof(ATOM_BASIC_INFORMATION);\r
-    }\r
-\r
-    if (NameLength + sizeof(ATOM_BASIC_INFORMATION) > AtomInformationLength)\r
-    {\r
-        return STATUS_INFO_LENGTH_MISMATCH;\r
-    }\r
-\r
-    AtomInformation->UsageCount = (USHORT)UsageCount;\r
-    AtomInformation->Flags = (USHORT)Flags;\r
-    AtomInformation->NameLength = (USHORT)NameLength;\r
-\r
-    return STATUS_SUCCESS;\r
-}\r
-\r
-NTSTATUS\r
-NTAPI\r
-RtlpQueryAtomTableInformation(PRTL_ATOM_TABLE AtomTable,\r
-                              RTL_ATOM Atom,\r
-                              PATOM_TABLE_INFORMATION AtomInformation,\r
-                              ULONG AtomInformationLength,\r
-                              PULONG ReturnLength)\r
-{\r
-    ULONG Length;\r
-    NTSTATUS Status;\r
-\r
-    Length = sizeof(ATOM_TABLE_INFORMATION);\r
-    DPRINT("RequiredLength: %lu\n", Length);\r
-\r
-    if (ReturnLength) *ReturnLength = Length;\r
-\r
-    if (Length > AtomInformationLength) return STATUS_INFO_LENGTH_MISMATCH;\r
-\r
-    Status = RtlQueryAtomListInAtomTable(AtomTable,\r
-                                         (AtomInformationLength - Length) /\r
-                                         sizeof(RTL_ATOM),\r
-                                         &AtomInformation->NumberOfAtoms,\r
-                                         AtomInformation->Atoms);\r
-    if (NT_SUCCESS(Status))\r
-    {\r
-        ReturnLength += AtomInformation->NumberOfAtoms * sizeof(RTL_ATOM);\r
-        if (ReturnLength != NULL) *ReturnLength = Length;\r
-    }\r
-\r
-    return Status;\r
-}\r
-\r
 /* FUNCTIONS ****************************************************************/\r
 \r
 /*\r
@@ -130,14 +59,96 @@ NtAddAtom(IN PWSTR AtomName,
           OUT PRTL_ATOM Atom)\r
 {\r
     PRTL_ATOM_TABLE AtomTable = ExpGetGlobalAtomTable();\r
+    NTSTATUS Status = STATUS_SUCCESS;\r
+    KPROCESSOR_MODE PreviousMode = ExGetPreviousMode();\r
+    LPWSTR CapturedName = NULL;\r
+    ULONG CapturedSize;\r
+    RTL_ATOM SafeAtom;\r
+    PAGED_CODE();\r
 \r
     /* Check for the table */\r
     if (AtomTable == NULL) return STATUS_ACCESS_DENIED;\r
 \r
-    /* FIXME: SEH! */\r
+    /* Check for valid name */\r
+    if (AtomNameLength > (RTL_MAXIMUM_ATOM_LENGTH * sizeof(WCHAR)))\r
+    {\r
+        /* Fail */\r
+        DPRINT1("Atom name too long\n");\r
+        return STATUS_INVALID_PARAMETER;\r
+    }\r
 \r
-    /* Call the worker function */\r
-    return RtlAddAtomToAtomTable(AtomTable, AtomName, Atom);\r
+    /* Check if we're called from user-mode*/\r
+    if (PreviousMode != KernelMode)\r
+    {\r
+        /* Enter SEH */\r
+        _SEH_TRY\r
+        {\r
+            /* Check if we have a name */\r
+            if (AtomName)\r
+            {\r
+                /* Probe the atom */\r
+                ProbeForRead(AtomName, AtomNameLength, sizeof(WCHAR));\r
+\r
+                /* Allocate an aligned buffer + the null char */\r
+                CapturedSize = ((AtomNameLength + sizeof(WCHAR)) &~\r
+                                (sizeof(WCHAR) -1));\r
+                CapturedName = ExAllocatePoolWithTag(PagedPool,\r
+                                                     CapturedSize,\r
+                                                     TAG_ATOM);\r
+                if (!CapturedName)\r
+                {\r
+                    /* Fail the call */\r
+                    Status = STATUS_INSUFFICIENT_RESOURCES;\r
+                }\r
+                else\r
+                {\r
+                    /* Copy the name and null-terminate it */\r
+                    RtlMoveMemory(CapturedName, AtomName, AtomNameLength);\r
+                    CapturedName[AtomNameLength / sizeof(WCHAR)] = UNICODE_NULL;\r
+                }\r
+\r
+                /* Probe the atom too */\r
+                if (Atom) ProbeForWriteUshort(Atom);\r
+            }\r
+        }\r
+        _SEH_EXCEPT(_SEH_ExSystemExceptionFilter)\r
+        {\r
+            Status = _SEH_GetExceptionCode();\r
+        }\r
+        _SEH_END;\r
+    }\r
+    else\r
+    {\r
+        /* Simplify code and re-use one variable */\r
+        if (AtomName) CapturedName = AtomName;\r
+    }\r
+\r
+    /* Make sure probe worked */\r
+    if (NT_SUCCESS(Status))\r
+    {\r
+        /* Call the runtime function */\r
+        Status = RtlAddAtomToAtomTable(AtomTable, CapturedName, &SafeAtom);\r
+        if (NT_SUCCESS(Status) && (Atom))\r
+        {\r
+            /* Success and caller wants the atom back.. .enter SEH */\r
+            _SEH_TRY\r
+            {\r
+                /* Return the atom */\r
+                *Atom = SafeAtom;\r
+            }\r
+            _SEH_EXCEPT(_SEH_ExSystemExceptionFilter)\r
+            {\r
+                Status = _SEH_GetExceptionCode();\r
+            }\r
+            _SEH_END;\r
+        }\r
+    }\r
+\r
+    /* If we captured anything, free it */\r
+    if ((CapturedName) && (CapturedName != AtomName)) ExFreePool(CapturedName);\r
+\r
+    /* Return to caller */\r
+    return Status;\r
 }\r
 \r
 /*\r
@@ -148,6 +159,7 @@ NTAPI
 NtDeleteAtom(IN RTL_ATOM Atom)\r
 {\r
     PRTL_ATOM_TABLE AtomTable = ExpGetGlobalAtomTable();\r
+    PAGED_CODE();\r
 \r
     /* Check for valid table */\r
     if (AtomTable == NULL) return STATUS_ACCESS_DENIED;\r
@@ -166,14 +178,96 @@ NtFindAtom(IN PWSTR AtomName,
            OUT PRTL_ATOM Atom)\r
 {\r
     PRTL_ATOM_TABLE AtomTable = ExpGetGlobalAtomTable();\r
+    NTSTATUS Status = STATUS_SUCCESS;\r
+    KPROCESSOR_MODE PreviousMode = ExGetPreviousMode();\r
+    LPWSTR CapturedName = NULL;\r
+    ULONG CapturedSize;\r
+    RTL_ATOM SafeAtom;\r
+    PAGED_CODE();\r
 \r
-    /* Check for valid table */\r
+    /* Check for the table */\r
     if (AtomTable == NULL) return STATUS_ACCESS_DENIED;\r
 \r
-    /* FIXME: SEH!!! */\r
+    /* Check for valid name */\r
+    if (AtomNameLength > (RTL_MAXIMUM_ATOM_LENGTH * sizeof(WCHAR)))\r
+    {\r
+        /* Fail */\r
+        DPRINT1("Atom name too long\n");\r
+        return STATUS_INVALID_PARAMETER;\r
+    }\r
+\r
+    /* Check if we're called from user-mode*/\r
+    if (PreviousMode != KernelMode)\r
+    {\r
+        /* Enter SEH */\r
+        _SEH_TRY\r
+        {\r
+            /* Check if we have a name */\r
+            if (AtomName)\r
+            {\r
+                /* Probe the atom */\r
+                ProbeForRead(AtomName, AtomNameLength, sizeof(WCHAR));\r
+\r
+                /* Allocate an aligned buffer + the null char */\r
+                CapturedSize = ((AtomNameLength + sizeof(WCHAR)) &~\r
+                                (sizeof(WCHAR) -1));\r
+                CapturedName = ExAllocatePoolWithTag(PagedPool,\r
+                                                     CapturedSize,\r
+                                                     TAG_ATOM);\r
+                if (!CapturedName)\r
+                {\r
+                    /* Fail the call */\r
+                    Status = STATUS_INSUFFICIENT_RESOURCES;\r
+                }\r
+                else\r
+                {\r
+                    /* Copy the name and null-terminate it */\r
+                    RtlMoveMemory(CapturedName, AtomName, AtomNameLength);\r
+                    CapturedName[AtomNameLength / sizeof(WCHAR)] = UNICODE_NULL;\r
+                }\r
+\r
+                /* Probe the atom too */\r
+                if (Atom) ProbeForWriteUshort(Atom);\r
+            }\r
+        }\r
+        _SEH_EXCEPT(_SEH_ExSystemExceptionFilter)\r
+        {\r
+            Status = _SEH_GetExceptionCode();\r
+        }\r
+        _SEH_END;\r
+    }\r
+    else\r
+    {\r
+        /* Simplify code and re-use one variable */\r
+        if (AtomName) CapturedName = AtomName;\r
+    }\r
 \r
-    /* Call worker function */\r
-    return RtlLookupAtomInAtomTable(AtomTable, AtomName, Atom);\r
+    /* Make sure probe worked */\r
+    if (NT_SUCCESS(Status))\r
+    {\r
+        /* Call the runtime function */\r
+        Status = RtlLookupAtomInAtomTable(AtomTable, CapturedName, &SafeAtom);\r
+        if (NT_SUCCESS(Status) && (Atom))\r
+        {\r
+            /* Success and caller wants the atom back.. .enter SEH */\r
+            _SEH_TRY\r
+            {\r
+                /* Return the atom */\r
+                *Atom = SafeAtom;\r
+            }\r
+            _SEH_EXCEPT(_SEH_ExSystemExceptionFilter)\r
+            {\r
+                Status = _SEH_GetExceptionCode();\r
+            }\r
+            _SEH_END;\r
+        }\r
+    }\r
+\r
+    /* If we captured anything, free it */\r
+    if ((CapturedName) && (CapturedName != AtomName)) ExFreePool(CapturedName);\r
+\r
+    /* Return to caller */\r
+    return Status;\r
 }\r
 \r
 /*\r
@@ -188,7 +282,10 @@ NtQueryInformationAtom(RTL_ATOM Atom,
                        PULONG ReturnLength)\r
 {\r
     PRTL_ATOM_TABLE AtomTable = ExpGetGlobalAtomTable();\r
+    PATOM_BASIC_INFORMATION BasicInformation = AtomInformation;\r
+    PATOM_TABLE_INFORMATION TableInformation = AtomInformation;\r
     NTSTATUS Status;\r
+    ULONG Flags, UsageCount, NameLength;\r
 \r
     /* Check for valid table */\r
     if (AtomTable == NULL) return STATUS_ACCESS_DENIED;\r
@@ -198,23 +295,70 @@ NtQueryInformationAtom(RTL_ATOM Atom,
     /* Choose class */\r
     switch (AtomInformationClass)\r
     {\r
+        /* Caller requested info about an atom */\r
         case AtomBasicInformation:\r
-            Status = RtlpQueryAtomInformation(AtomTable,\r
-                                              Atom,\r
-                                              AtomInformation,\r
-                                              AtomInformationLength,\r
-                                              ReturnLength);\r
+\r
+            /* Size check */\r
+            *ReturnLength = FIELD_OFFSET(ATOM_BASIC_INFORMATION, Name);\r
+            if (*ReturnLength > AtomInformationLength)\r
+            {\r
+                /* Fail */\r
+                DPRINT1("Buffer too small\n");\r
+                return STATUS_INFO_LENGTH_MISMATCH;\r
+            }\r
+\r
+            /* Prepare query */\r
+            UsageCount = 0;\r
+            NameLength = AtomInformationLength - *ReturnLength;\r
+            BasicInformation->Name[0] = UNICODE_NULL;\r
+\r
+            /* Query the data */\r
+            Status = RtlQueryAtomInAtomTable(AtomTable,\r
+                                             Atom,\r
+                                             &UsageCount,\r
+                                             &Flags,\r
+                                             BasicInformation->Name,\r
+                                             &NameLength);\r
+            if (NT_SUCCESS(Status))\r
+            {\r
+                /* Return data */\r
+                BasicInformation->UsageCount = (USHORT)UsageCount;\r
+                BasicInformation->Flags = (USHORT)Flags;\r
+                BasicInformation->NameLength = (USHORT)NameLength;\r
+                *ReturnLength += NameLength + sizeof(WCHAR);\r
+            }\r
             break;\r
 \r
+        /* Caller requested info about an Atom Table */\r
         case AtomTableInformation:\r
-            Status = RtlpQueryAtomTableInformation(AtomTable,\r
-                                                   Atom,\r
-                                                   AtomInformation,\r
-                                                   AtomInformationLength,\r
-                                                   ReturnLength);\r
+\r
+            /* Size check */\r
+            *ReturnLength = FIELD_OFFSET(ATOM_TABLE_INFORMATION, Atoms);\r
+            if (*ReturnLength > AtomInformationLength)\r
+            {\r
+                /* Fail */\r
+                DPRINT1("Buffer too small\n");\r
+                return STATUS_INFO_LENGTH_MISMATCH;\r
+            }\r
+\r
+            /* Query the data */\r
+            Status = RtlQueryAtomListInAtomTable(AtomTable,\r
+                                                 (AtomInformationLength - *ReturnLength) /\r
+                                                 sizeof(RTL_ATOM),\r
+                                                 &TableInformation->NumberOfAtoms,\r
+                                                 TableInformation->Atoms);\r
+            if (NT_SUCCESS(Status))\r
+            {\r
+                /* Update the return length */\r
+                *ReturnLength += TableInformation->NumberOfAtoms *\r
+                                 sizeof(RTL_ATOM);\r
+            }\r
             break;\r
 \r
+        /* Caller was on crack */\r
         default:\r
+\r
+            /* Unrecognized class */\r
             Status = STATUS_INVALID_INFO_CLASS;\r
     }\r
 \r