ntdll: While requesting TokenGroups calculate required user buffer size in server.

dlls/advapi32/tests/security.c

@@ -1442,6 +1442,24 @@ static void test_token_attr(void)
     ok(ret, "OpenProcessToken failed with error %d\n", GetLastError());
 
     /* groups */
+    /* insufficient buffer length */
+    SetLastError(0xdeadbeef);
+    Size2 = 0;
+    ret = GetTokenInformation(Token, TokenGroups, NULL, 0, &Size2);
+    ok(Size2 > 1, "got %d\n", Size2);
+    ok(!ret && GetLastError() == ERROR_INSUFFICIENT_BUFFER,
+        "%d with error %d\n", ret, GetLastError());
+    Size2 -= 1;
+    Groups = HeapAlloc(GetProcessHeap(), 0, Size2);
+    memset(Groups, 0xcc, Size2);
+    Size = 0;
+    ret = GetTokenInformation(Token, TokenGroups, Groups, Size2, &Size);
+    ok(Size > 1, "got %d\n", Size);
+    ok(!ret && GetLastError() == ERROR_INSUFFICIENT_BUFFER,
+        "%d with error %d\n", ret, GetLastError());
+    ok(*((BYTE*)Groups) == 0xcc, "buffer altered\n");
+    HeapFree(GetProcessHeap(), 0, Groups);
+
     SetLastError(0xdeadbeef);
     ret = GetTokenInformation(Token, TokenGroups, NULL, 0, &Size);
     ok(!ret && GetLastError() == ERROR_INSUFFICIENT_BUFFER,

dlls/ntdll/nt.c

@@ -320,35 +320,21 @@ NTSTATUS WINAPI NtQueryInformationToken(
         break;
     case TokenGroups:
     {
-        char stack_buffer[256];
+        void *buffer;
-        unsigned int server_buf_len = sizeof(stack_buffer);
-        void *buffer = stack_buffer;
-        BOOLEAN need_more_memory;
 
-        /* we cannot work out the size of the server buffer required for the
+        /* reply buffer is always shorter than output one */
-         * input size, since there are two factors affecting how much can be
+        buffer = tokeninfolength ? RtlAllocateHeap(GetProcessHeap(), 0, tokeninfolength) : NULL;
-         * stored in the buffer - number of groups and lengths of sids */
-        do
-        {
-            need_more_memory = FALSE;
 
         SERVER_START_REQ( get_token_groups )
         {
             TOKEN_GROUPS *groups = tokeninfo;
 
             req->handle = wine_server_obj_handle( token );
-                wine_server_set_reply( req, buffer, server_buf_len );
+            wine_server_set_reply( req, buffer, tokeninfolength );
             status = wine_server_call( req );
             if (status == STATUS_BUFFER_TOO_SMALL)
             {
-                    if (buffer == stack_buffer)
+                if (retlen) *retlen = reply->user_len;
-                        buffer = RtlAllocateHeap(GetProcessHeap(), 0, reply->user_len);
-                    else
-                        buffer = RtlReAllocateHeap(GetProcessHeap(), 0, buffer, reply->user_len);
-                    if (!buffer) return STATUS_NO_MEMORY;
-
-                    server_buf_len = reply->user_len;
-                    need_more_memory = TRUE;
             }
             else if (status == STATUS_SUCCESS)
             {
@@ -357,16 +343,12 @@ NTSTATUS WINAPI NtQueryInformationToken(
                 ULONG i;
                 const int non_sid_portion = (sizeof(struct token_groups) + tg->count * sizeof(unsigned int));
                 SID *sids = (SID *)((char *)tokeninfo + FIELD_OFFSET( TOKEN_GROUPS, Groups[tg->count] ));
-                    ULONG needed_bytes = FIELD_OFFSET( TOKEN_GROUPS, Groups[tg->count] ) +
-                        reply->user_len - non_sid_portion;
 
-                    if (retlen) *retlen = needed_bytes;
+                if (retlen) *retlen = reply->user_len;
 
-                    if (needed_bytes <= tokeninfolength)
-                    {
                 groups->GroupCount = tg->count;
                 memcpy( sids, (char *)buffer + non_sid_portion,
-                                reply->user_len - non_sid_portion );
+                        reply->user_len - non_sid_portion - FIELD_OFFSET( TOKEN_GROUPS, Groups[tg->count] ));
 
                 for (i = 0; i < tg->count; i++)
                 {
@@ -375,13 +357,11 @@ NTSTATUS WINAPI NtQueryInformationToken(
                     sids = (SID *)((char *)sids + RtlLengthSid(sids));
                 }
              }
-                    else status = STATUS_BUFFER_TOO_SMALL;
-                }
              else if (retlen) *retlen = 0;
         }
         SERVER_END_REQ;
-        } while (need_more_memory);
+
-        if (buffer != stack_buffer) RtlFreeHeap(GetProcessHeap(), 0, buffer);
+        RtlFreeHeap(GetProcessHeap(), 0, buffer);
         break;
     }
     case TokenPrimaryGroup:

server/token.c

@@ -1295,19 +1295,30 @@ DECL_HANDLER(get_token_groups)
                                                  &token_ops )))
     {
         size_t size_needed = sizeof(struct token_groups);
+        size_t sid_size = 0;
         unsigned int group_count = 0;
         const struct group *group;
 
         LIST_FOR_EACH_ENTRY( group, &token->groups, const struct group, entry )
         {
             group_count++;
-            size_needed += FIELD_OFFSET(SID, SubAuthority[group->sid.SubAuthorityCount]);
+            sid_size += FIELD_OFFSET(SID, SubAuthority[group->sid.SubAuthorityCount]);
         }
+        size_needed += sid_size;
+        /* attributes size */
         size_needed += sizeof(unsigned int) * group_count;
 
-        reply->user_len = size_needed;
+        /* reply buffer contains size_needed bytes formatted as:
 
-        if (size_needed <= get_reply_max_size())
+           unsigned int count;
+           unsigned int attrib[count];
+           char sid_data[];
+
+           user_len includes extra data needed for TOKEN_GROUPS representation,
+           required caller buffer size calculated here to avoid extra server call */
+        reply->user_len = FIELD_OFFSET( TOKEN_GROUPS, Groups[group_count] ) + sid_size;
+
+        if (reply->user_len <= get_reply_max_size())
         {
             struct token_groups *tg = set_reply_data_size( size_needed );
             if (tg)