[ros-diffs] [gvg] 18020: Fix RtlQueryAtomInAtomTable and add regression tests

Fix RtlQueryAtomInAtomTable and add regression tests
Modified: trunk/reactos/lib/rtl/atom.c
Modified: trunk/reactos/regtests/winetests/ntdll/atom.c
Modified: trunk/reactos/subsys/win32k/ntuser/class.c
  _____  

Modified: trunk/reactos/lib/rtl/atom.c
--- trunk/reactos/lib/rtl/atom.c	2005-09-23 19:24:20 UTC (rev
18019)
+++ trunk/reactos/lib/rtl/atom.c	2005-09-23 21:08:57 UTC (rev
18020)
@@ -565,6 +565,23 @@

 
 /*
  * @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,54 +592,35 @@
                         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)
-          {
-             *RefCount = 1;
-          }
+        /* Synthesize an entry */
+        NumberEntry.AtomTableEntry.Atom = Atom;
+        NumberEntry.AtomTableEntry.NameLength =
swprintf(NumberEntry.AtomTableEntry.Name,
+                                                         L"#%lu",
+                                                         (ULONG)Atom);
+        NumberEntry.AtomTableEntry.ReferenceCount = 1;
+        NumberEntry.AtomTableEntry.Flags = RTL_ATOM_IS_PINNED;
+        Entry = &NumberEntry.AtomTableEntry;
+     }
+   else
+     {
+        RtlpLockAtomTable(AtomTable);
 
-        if (PinCount != NULL)
-          {
-             *PinCount = 1;
-          }
-
-        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;
+        Entry = RtlpGetAtomEntry(AtomTable,
+                                 (ULONG)((USHORT)Atom - 0xC000));
      }
 
-   RtlpLockAtomTable(AtomTable);
-
-   Entry = RtlpGetAtomEntry(AtomTable,
-                            (ULONG)((USHORT)Atom - 0xC000));
-
    if (Entry != NULL && Entry->Atom == (USHORT)Atom)
      {
         DPRINT("Atom name: %wZ\n", &Entry->Name);
@@ -637,27 +635,40 @@
              *PinCount = ((Entry->Flags & RTL_ATOM_IS_PINNED) != 0);
           }
 
-        if ((AtomName != NULL) && (NameLength != NULL))
+        if (NULL != NameLength)
           {
              Length = Entry->NameLength * sizeof(WCHAR);
-
-             if (*NameLength < Length + sizeof(WCHAR))
+             if (NULL != AtomName)
                {
-                  *NameLength = Length;
-                  Status = STATUS_BUFFER_TOO_SMALL;
+                  if (*NameLength < Length + sizeof(WCHAR))
+                    {
+                       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 @@
         Status = STATUS_INVALID_HANDLE;
      }
 
-   RtlpUnlockAtomTable(AtomTable);
+   if (NULL != Entry && Entry != &NumberEntry.AtomTableEntry)
+     {
+        RtlpUnlockAtomTable(AtomTable);
+     }
 
    return Status;
 }
  _____  

Modified: trunk/reactos/regtests/winetests/ntdll/atom.c
--- trunk/reactos/regtests/winetests/ntdll/atom.c	2005-09-23
19:24:20 UTC (rev 18019)
+++ trunk/reactos/regtests/winetests/ntdll/atom.c	2005-09-23
21:08:57 UTC (rev 18020)
@@ -233,6 +233,19 @@

         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);
 
  _____  

Modified: trunk/reactos/subsys/win32k/ntuser/class.c
--- trunk/reactos/subsys/win32k/ntuser/class.c	2005-09-23 19:24:20 UTC
(rev 18019)
+++ trunk/reactos/subsys/win32k/ntuser/class.c	2005-09-23 21:08:57 UTC
(rev 18020)
@@ -207,10 +207,13 @@

 
    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");