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,68 +320,48 @@ 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
+        SERVER_START_REQ( get_token_groups )
         {
-            need_more_memory = FALSE;
+            TOKEN_GROUPS *groups = tokeninfo;
 
-            SERVER_START_REQ( get_token_groups )
+            req->handle = wine_server_obj_handle( token );
+            wine_server_set_reply( req, buffer, tokeninfolength );
+            status = wine_server_call( req );
+            if (status == STATUS_BUFFER_TOO_SMALL)
             {
-                TOKEN_GROUPS *groups = tokeninfo;
+                if (retlen) *retlen = reply->user_len;
-
-                req->handle = wine_server_obj_handle( token );
-                wine_server_set_reply( req, buffer, server_buf_len );
-                status = wine_server_call( req );
-                if (status == STATUS_BUFFER_TOO_SMALL)
-                {
-                    if (buffer == stack_buffer)
-                        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)
-                {
-                    struct token_groups *tg = buffer;
-                    unsigned int *attr = (unsigned int *)(tg + 1);
-                    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 (needed_bytes <= tokeninfolength)
-                    {
-                        groups->GroupCount = tg->count;
-                        memcpy( sids, (char *)buffer + non_sid_portion,
-                                reply->user_len - non_sid_portion );
-
-                        for (i = 0; i < tg->count; i++)
-                        {
-                            groups->Groups[i].Attributes = attr[i];
-                            groups->Groups[i].Sid = sids;
-                            sids = (SID *)((char *)sids + RtlLengthSid(sids));
-                        }
-                    }
-                    else status = STATUS_BUFFER_TOO_SMALL;
-                }
-                else if (retlen) *retlen = 0;
             }
-            SERVER_END_REQ;
+            else if (status == STATUS_SUCCESS)
-        } while (need_more_memory);
+            {
-        if (buffer != stack_buffer) RtlFreeHeap(GetProcessHeap(), 0, buffer);
+                struct token_groups *tg = buffer;
+                unsigned int *attr = (unsigned int *)(tg + 1);
+                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] ));
+
+                if (retlen) *retlen = reply->user_len;
+
+                groups->GroupCount = tg->count;
+                memcpy( sids, (char *)buffer + non_sid_portion,
+                        reply->user_len - non_sid_portion - FIELD_OFFSET( TOKEN_GROUPS, Groups[tg->count] ));
+
+                for (i = 0; i < tg->count; i++)
+                {
+                    groups->Groups[i].Attributes = attr[i];
+                    groups->Groups[i].Sid = sids;
+                    sids = (SID *)((char *)sids + RtlLengthSid(sids));
+                }
+             }
+             else if (retlen) *retlen = 0;
+        }
+        SERVER_END_REQ;
+
+        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)