Fix RtlQueryAtomInAtomTable and add regression tests

svn path=/trunk/; revision=18020

reactos/lib/rtl/atom.c

@@ -565,6 +565,23 @@ RtlPinAtomInAtomTable(IN PRTL_ATOM_TABLE AtomTable,
 
 /*
  * @implemented
+ *
+ * This API is really messed up with regards to NameLength. If you pass in a
+ * valid buffer for AtomName, NameLength should be the size of the buffer
+ * (in bytes, not characters). So if you expect the string to be 6 char long,
+ * you need to allocate a buffer of 7 WCHARs and pass 14 for NameLength.
+ * The AtomName returned is always null terminated. If the NameLength you pass
+ * is smaller than 4 (4 would leave room for 1 character) the function will
+ * return with status STATUS_BUFFER_TOO_SMALL. If you pass more than 4, the
+ * return status will be STATUS_SUCCESS, even if the buffer is not large enough
+ * to hold the complete string. In that case, the string is silently truncated
+ * and made to fit in the provided buffer. On return NameLength is set to the
+ * number of bytes (but EXCLUDING the bytes for the null terminator) copied.
+ * So, if the string is 6 char long, you pass a buffer of 10 bytes, on return
+ * NameLength will be set to 8.
+ * If you pass in a NULL value for AtomName, the length of the string in bytes
+ * (again EXCLUDING the null terminator) is returned in NameLength, at least
+ * on Win2k, XP and ReactOS. NT4 will return 0 in that case.
  */
 NTSTATUS STDCALL
 RtlQueryAtomInAtomTable(PRTL_ATOM_TABLE AtomTable,
@@ -575,53 +592,34 @@ RtlQueryAtomInAtomTable(PRTL_ATOM_TABLE AtomTable,
                         PULONG NameLength)
 {
    ULONG Length;
+   union
+     {
+     /* A RTL_ATOM_TABLE_ENTRY has a "WCHAR Name[1]" entry at the end.
+      * Make sure we reserve enough room to facilitate a 12 character name */
+     RTL_ATOM_TABLE_ENTRY AtomTableEntry;
+     WCHAR StringBuffer[sizeof(RTL_ATOM_TABLE_ENTRY) / sizeof(WCHAR) + 12];
+     } NumberEntry;
    PRTL_ATOM_TABLE_ENTRY Entry;
    NTSTATUS Status = STATUS_SUCCESS;
 
    if (Atom < 0xC000)
      {
-        if (RefCount != NULL)
+        /* Synthesize an entry */
-          {
+        NumberEntry.AtomTableEntry.Atom = Atom;
-             *RefCount = 1;
+        NumberEntry.AtomTableEntry.NameLength = swprintf(NumberEntry.AtomTableEntry.Name,
-          }
+                                                         L"#%lu",
-
+                                                         (ULONG)Atom);
-        if (PinCount != NULL)
+        NumberEntry.AtomTableEntry.ReferenceCount = 1;
-          {
+        NumberEntry.AtomTableEntry.Flags = RTL_ATOM_IS_PINNED;
-             *PinCount = 1;
+        Entry = &NumberEntry.AtomTableEntry;
-          }
-
-        if ((AtomName != NULL) && (NameLength != NULL) && (NameLength > 0))
-          {
-             WCHAR NameString[12];
-             
-             Length = swprintf(NameString, L"#%lu", (ULONG)Atom) * sizeof(WCHAR);
-
-             if (*NameLength < Length + sizeof(WCHAR))
-               {
-                  *NameLength = Length;
-                  Status = STATUS_BUFFER_TOO_SMALL;
-               }
-             else 
-               {
-                  RtlCopyMemory(AtomName,
-                                NameString,
-                                Length);
-                  AtomName[Length / sizeof(WCHAR)] = L'\0';
-                  *NameLength = Length;
-               }
-          }
-        else if (NameLength != NULL)
-          {
-             *NameLength = (Entry->NameLength + 1) * sizeof(WCHAR);
-          }
-
-        return Status;
      }
+   else
+     {
+        RtlpLockAtomTable(AtomTable);
 
-   RtlpLockAtomTable(AtomTable);
+        Entry = RtlpGetAtomEntry(AtomTable,
-
+                                 (ULONG)((USHORT)Atom - 0xC000));
-   Entry = RtlpGetAtomEntry(AtomTable,
+     }
-                            (ULONG)((USHORT)Atom - 0xC000));
 
    if (Entry != NULL && Entry->Atom == (USHORT)Atom)
      {
@@ -637,27 +635,40 @@ RtlQueryAtomInAtomTable(PRTL_ATOM_TABLE AtomTable,
              *PinCount = ((Entry->Flags & RTL_ATOM_IS_PINNED) != 0);
           }
 
-        if ((AtomName != NULL) && (NameLength != NULL))
+        if (NULL != NameLength)
           {
              Length = Entry->NameLength * sizeof(WCHAR);
-
+             if (NULL != AtomName)
-             if (*NameLength < Length + sizeof(WCHAR))
                {
-                  *NameLength = Length;
+                  if (*NameLength < Length + sizeof(WCHAR))
-                  Status = STATUS_BUFFER_TOO_SMALL;
+                    {
+                       if (*NameLength < 4)
+                         {
+                            *NameLength = Length;
+                            Status = STATUS_BUFFER_TOO_SMALL;
+                         }
+                       else
+                         {
+                            Length = *NameLength - sizeof(WCHAR);
+                         }
+                    }
+                  if (NT_SUCCESS(Status))
+                    {
+                       RtlCopyMemory(AtomName,
+                                     Entry->Name,
+                                     Length);
+                       AtomName[Length / sizeof(WCHAR)] = L'\0';
+                       *NameLength = Length;
+                    }
                }
              else
                {
-                  RtlCopyMemory(AtomName,
-                                Entry->Name,
-                                Length);
-                  AtomName[Length / sizeof(WCHAR)] = L'\0';
                   *NameLength = Length;
                }
           }
-        else if (NameLength != NULL)
+        else if (NULL != AtomName)
           {
-             *NameLength = (Entry->NameLength + 1) * sizeof(WCHAR);
+             Status = STATUS_INVALID_PARAMETER;
           }
      }
    else
@@ -665,7 +676,10 @@ RtlQueryAtomInAtomTable(PRTL_ATOM_TABLE AtomTable,
         Status = STATUS_INVALID_HANDLE;
      }
 
-   RtlpUnlockAtomTable(AtomTable);
+   if (NULL != Entry && Entry != &NumberEntry.AtomTableEntry)
+     {
+        RtlpUnlockAtomTable(AtomTable);
+     }
 
    return Status;
 }

reactos/regtests/winetests/ntdll/atom.c

@@ -233,6 +233,19 @@ static void test_NtAtom(void)
         ok(res == STATUS_BUFFER_TOO_SMALL, "Got wrong retval, retval: %lx\n", res);
         ok((strlenW(testAtom1) * sizeof(WCHAR)) == Len, "Got wrong length %lx\n", Len);
 
+        res = pRtlQueryAtomInAtomTable(AtomTable, Atom1, NULL, NULL, NULL, &Len);
+        ok(!res, "Failed to retrieve atom length, retval: %lx\n", res);
+        ok(Len == strlenW(testAtom1) * sizeof(WCHAR), "Invalid atom length got %lu expected %u\n",
+           Len, strlenW(testAtom1) * sizeof(WCHAR));
+
+        Len = strlenW(testAtom1) * sizeof(WCHAR);
+        Name[strlenW(testAtom1)] = '*';
+        res = pRtlQueryAtomInAtomTable(AtomTable, Atom1, NULL, NULL, Name, &Len);
+        ok(!res, "Failed with exactly long enough buffer, retval: %lx\n", res);
+        ok(Name[strlenW(testAtom1)] == '*', "Writing outside buffer\n");
+        ok(0 == memcmp(Name, testAtom1, (strlenW(testAtom1) - 1) * sizeof(WCHAR)),
+           "We found wrong atom!!\n");
+
         res = pRtlPinAtomInAtomTable(AtomTable, Atom1);
         ok(!res, "Unable to pin atom in atom table, retval: %lx\n", res);
 

reactos/subsys/win32k/ntuser/class.c

@@ -207,10 +207,13 @@ IntGetClassName(struct _WINDOW_OBJECT *WindowObject, LPWSTR lpClassName,
 
    Length = 0;
    Status = RtlQueryAtomInAtomTable(WinStaObject->AtomTable,
-                                    WindowObject->Class->Atom, NULL, NULL, NULL, &Length);
+                                    WindowObject->Class->Atom, NULL, NULL,
+                                    NULL, &Length);
+   Length += sizeof(WCHAR);
    Name = ExAllocatePoolWithTag(PagedPool, Length, TAG_STRING);
    Status = RtlQueryAtomInAtomTable(WinStaObject->AtomTable,
-                                    WindowObject->Class->Atom, NULL, NULL, Name, &Length);
+                                    WindowObject->Class->Atom, NULL, NULL,
+                                    Name, &Length);
    if (!NT_SUCCESS(Status))
    {
       DPRINT("IntGetClassName: RtlQueryAtomInAtomTable failed\n");