dlls/combase/malloc.c

   1 /*
   2  * Copyright 1997 Marcus Meissner
   3  *
   4  * This library is free software; you can redistribute it and/or
   5  * modify it under the terms of the GNU Lesser General Public
   6  * License as published by the Free Software Foundation; either
   7  * version 2.1 of the License, or (at your option) any later version.
   8  *
   9  * This library is distributed in the hope that it will be useful,
  10  * but WITHOUT ANY WARRANTY; without even the implied warranty of
  11  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  12  * Lesser General Public License for more details.
  13  *
  14  * You should have received a copy of the GNU Lesser General Public
  15  * License along with this library; if not, write to the Free Software
  16  * Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301, USA
  17  */
  18
  19 #define COBJMACROS
  20
  21 #include "oleauto.h"
  22
  23 #include "wine/debug.h"
  24 #include "wine/heap.h"
  25
  26 WINE_DEFAULT_DEBUG_CHANNEL(olemalloc);
  27
  28 static const IMallocVtbl allocator_vtbl;
  29
  30 struct allocator
  31 {
  32     IMalloc IMalloc_iface;
  33     IMallocSpy *spy;
  34     DWORD spyed_allocations;
  35     BOOL spy_release_pending; /* CoRevokeMallocSpy called with spyed allocations left */
  36     void **blocks;
  37     DWORD blocks_length;
  38 };
  39
  40 static struct allocator allocator = { .IMalloc_iface.lpVtbl = &allocator_vtbl };
  41
  42 static CRITICAL_SECTION allocspy_cs;
  43 static CRITICAL_SECTION_DEBUG allocspy_cs_debug =
  44 {
  45     0, 0, &allocspy_cs,
  46     { &allocspy_cs_debug.ProcessLocksList, &allocspy_cs_debug.ProcessLocksList },
  47     0, 0, { (DWORD_PTR)(__FILE__ ": allocspy_cs") }
  48 };
  49 static CRITICAL_SECTION allocspy_cs = { &allocspy_cs_debug, -1, 0, 0, 0, 0 };
  50
  51 static BOOL mallocspy_grow(DWORD length)
  52 {
  53     void **blocks;
  54
  55     if (!allocator.blocks) blocks = LocalAlloc(LMEM_ZEROINIT, length * sizeof(void *));
  56     else blocks = LocalReAlloc(allocator.blocks, length * sizeof(void *), LMEM_ZEROINIT | LMEM_MOVEABLE);
  57     if (blocks)
  58     {
  59         allocator.blocks = blocks;
  60         allocator.blocks_length = length;
  61     }
  62
  63     return blocks != NULL;
  64 }
  65
  66 static void mallocspy_add_mem(void *mem)
  67 {
  68     void **current;
  69
  70     if (!mem || (!allocator.blocks_length && !mallocspy_grow(0x1000)))
  71         return;
  72
  73     /* Find a free location */
  74     current = allocator.blocks;
  75     while (*current)
  76     {
  77         current++;
  78         if (current >= allocator.blocks + allocator.blocks_length)
  79         {
  80             DWORD old_length = allocator.blocks_length;
  81             if (!mallocspy_grow(allocator.blocks_length + 0x1000))
  82                 return;
  83             current = allocator.blocks + old_length;
  84         }
  85     }
  86
  87     *current = mem;
  88     allocator.spyed_allocations++;
  89 }
  90
  91 static void** mallocspy_is_allocation_spyed(const void *mem)
  92 {
  93     void **current = allocator.blocks;
  94
  95     while (*current != mem)
  96     {
  97         current++;
  98         if (current >= allocator.blocks + allocator.blocks_length)
  99             return NULL;
 100     }
 101
 102     return current;
 103 }
 104
 105 static BOOL mallocspy_remove_spyed_memory(const void *mem)
 106 {
 107     void **current;
 108
 109     if (!allocator.blocks_length)
 110         return FALSE;
 111
 112     if (!(current = mallocspy_is_allocation_spyed(mem)))
 113         return FALSE;
 114
 115     allocator.spyed_allocations--;
 116     *current = NULL;
 117     return TRUE;
 118 }
 119
 120 static HRESULT WINAPI allocator_QueryInterface(IMalloc *iface, REFIID riid, void **obj)
 121 {
 122     TRACE("%s, %p.\n", debugstr_guid(riid), obj);
 123
 124     if (IsEqualIID(&IID_IUnknown, riid) || IsEqualIID(&IID_IMalloc, riid))
 125     {
 126         *obj = &allocator.IMalloc_iface;
 127         return S_OK;
 128     }
 129
 130     return E_NOINTERFACE;
 131 }
 132
 133 static ULONG WINAPI allocator_AddRef(IMalloc *iface)
 134 {
 135     return 2;
 136 }
 137
 138 static ULONG WINAPI allocator_Release(IMalloc *iface)
 139 {
 140     return 1;
 141 }
 142
 143 static void * WINAPI allocator_Alloc(IMalloc *iface, SIZE_T cb)
 144 {
 145     void *addr;
 146
 147     TRACE("%ld.\n", cb);
 148
 149     if (allocator.spy)
 150     {
 151         SIZE_T preAllocResult;
 152
 153         EnterCriticalSection(&allocspy_cs);
 154         preAllocResult = IMallocSpy_PreAlloc(allocator.spy, cb);
 155         if (cb && !preAllocResult)
 156         {
 157             /* PreAlloc can force Alloc to fail, but not if cb == 0 */
 158             TRACE("returning null\n");
 159             LeaveCriticalSection(&allocspy_cs);
 160             return NULL;
 161         }
 162     }
 163
 164     addr = HeapAlloc(GetProcessHeap(), 0, cb);
 165
 166     if (allocator.spy)
 167     {
 168         addr = IMallocSpy_PostAlloc(allocator.spy, addr);
 169         mallocspy_add_mem(addr);
 170         LeaveCriticalSection(&allocspy_cs);
 171     }
 172
 173     TRACE("%p.\n",addr);
 174     return addr;
 175 }
 176
 177 static void * WINAPI allocator_Realloc(IMalloc *iface, void *pv, SIZE_T cb)
 178 {
 179     void *addr;
 180
 181     TRACE("%p, %ld.\n",pv,cb);
 182
 183     if (allocator.spy)
 184     {
 185         void *real_mem;
 186         BOOL spyed;
 187
 188         EnterCriticalSection(&allocspy_cs);
 189         spyed = mallocspy_remove_spyed_memory(pv);
 190         cb = IMallocSpy_PreRealloc(allocator.spy, pv, cb, &real_mem, spyed);
 191
 192         /* check if can release the spy */
 193         if (allocator.spy_release_pending && !allocator.spyed_allocations)
 194         {
 195             IMallocSpy_Release(allocator.spy);
 196             allocator.spy_release_pending = FALSE;
 197             allocator.spy = NULL;
 198             LeaveCriticalSection(&allocspy_cs);
 199         }
 200
 201         if (!cb)
 202         {
 203             /* PreRealloc can force Realloc to fail */
 204             if (allocator.spy)
 205                 LeaveCriticalSection(&allocspy_cs);
 206             return NULL;
 207         }
 208
 209         pv = real_mem;
 210     }
 211
 212     if (!pv) addr = HeapAlloc(GetProcessHeap(), 0, cb);
 213     else if (cb) addr = HeapReAlloc(GetProcessHeap(), 0, pv, cb);
 214     else
 215     {
 216         HeapFree(GetProcessHeap(), 0, pv);
 217         addr = NULL;
 218     }
 219
 220     if (allocator.spy)
 221     {
 222         addr = IMallocSpy_PostRealloc(allocator.spy, addr, TRUE);
 223         mallocspy_add_mem(addr);
 224         LeaveCriticalSection(&allocspy_cs);
 225     }
 226
 227     TRACE("%p.\n", addr);
 228     return addr;
 229 }
 230
 231 static void WINAPI allocator_Free(IMalloc *iface, void *mem)
 232 {
 233     BOOL spyed_block = FALSE, spy_active = FALSE;
 234
 235     TRACE("%p.\n", mem);
 236
 237     if (!mem)
 238         return;
 239
 240     if (allocator.spy)
 241     {
 242         EnterCriticalSection(&allocspy_cs);
 243         spyed_block = mallocspy_remove_spyed_memory(mem);
 244         spy_active = TRUE;
 245         mem = IMallocSpy_PreFree(allocator.spy, mem, spyed_block);
 246     }
 247
 248     HeapFree(GetProcessHeap(), 0, mem);
 249
 250     if (spy_active)
 251     {
 252         IMallocSpy_PostFree(allocator.spy, spyed_block);
 253
 254         /* check if can release the spy */
 255         if (allocator.spy_release_pending && !allocator.spyed_allocations)
 256         {
 257             IMallocSpy_Release(allocator.spy);
 258             allocator.spy_release_pending = FALSE;
 259             allocator.spy = NULL;
 260         }
 261
 262         LeaveCriticalSection(&allocspy_cs);
 263     }
 264 }
 265
 266 /******************************************************************************
 267  * NOTES
 268  *  FIXME returns:
 269  *      win95:  size allocated (4 byte boundaries)
 270  *      win2k:  size originally requested !!! (allocated on 8 byte boundaries)
 271  */
 272 static SIZE_T WINAPI allocator_GetSize(IMalloc *iface, void *mem)
 273 {
 274     BOOL spyed_block = FALSE, spy_active = FALSE;
 275     SIZE_T size;
 276
 277     TRACE("%p.\n", mem);
 278
 279     if (!mem)
 280         return (SIZE_T)-1;
 281
 282     if (allocator.spy)
 283     {
 284         EnterCriticalSection(&allocspy_cs);
 285         spyed_block = !!mallocspy_is_allocation_spyed(mem);
 286         spy_active = TRUE;
 287         mem = IMallocSpy_PreGetSize(allocator.spy, mem, spyed_block);
 288     }
 289
 290     size = HeapSize(GetProcessHeap(), 0, mem);
 291
 292     if (spy_active)
 293     {
 294         size = IMallocSpy_PostGetSize(allocator.spy, size, spyed_block);
 295         LeaveCriticalSection(&allocspy_cs);
 296     }
 297
 298     return size;
 299 }
 300
 301 static INT WINAPI allocator_DidAlloc(IMalloc *iface, void *mem)
 302 {
 303     BOOL spyed_block = FALSE, spy_active = FALSE;
 304     int did_alloc;
 305
 306     TRACE("%p.\n", mem);
 307
 308     if (!mem)
 309         return -1;
 310
 311     if (allocator.spy)
 312     {
 313         EnterCriticalSection(&allocspy_cs);
 314         spyed_block = !!mallocspy_is_allocation_spyed(mem);
 315         spy_active = TRUE;
 316         mem = IMallocSpy_PreDidAlloc(allocator.spy, mem, spyed_block);
 317     }
 318
 319     did_alloc = HeapValidate(GetProcessHeap(), 0, mem);
 320
 321     if (spy_active)
 322     {
 323         did_alloc = IMallocSpy_PostDidAlloc(allocator.spy, mem, spyed_block, did_alloc);
 324         LeaveCriticalSection(&allocspy_cs);
 325     }
 326
 327     return did_alloc;
 328 }
 329
 330 static void WINAPI allocator_HeapMinimize(IMalloc *iface)
 331 {
 332     BOOL spy_active = FALSE;
 333
 334     TRACE("\n");
 335
 336     if (allocator.spy)
 337     {
 338         EnterCriticalSection(&allocspy_cs);
 339         spy_active = TRUE;
 340         IMallocSpy_PreHeapMinimize(allocator.spy);
 341     }
 342
 343     if (spy_active)
 344     {
 345         IMallocSpy_PostHeapMinimize(allocator.spy);
 346         LeaveCriticalSection(&allocspy_cs);
 347     }
 348 }
 349
 350 static const IMallocVtbl allocator_vtbl =
 351 {
 352     allocator_QueryInterface,
 353     allocator_AddRef,
 354     allocator_Release,
 355     allocator_Alloc,
 356     allocator_Realloc,
 357     allocator_Free,
 358     allocator_GetSize,
 359     allocator_DidAlloc,
 360     allocator_HeapMinimize
 361 };
 362
 363 /******************************************************************************
 364  *                CoGetMalloc        (combase.@)
 365  */
 366 HRESULT WINAPI CoGetMalloc(DWORD context, IMalloc **imalloc)
 367 {
 368     if (context != MEMCTX_TASK)
 369     {
 370         *imalloc = NULL;
 371         return E_INVALIDARG;
 372     }
 373
 374     *imalloc = &allocator.IMalloc_iface;
 375
 376     return S_OK;
 377 }
 378
 379 /***********************************************************************
 380  *           CoTaskMemAlloc         (combase.@)
 381  */
 382 void * WINAPI CoTaskMemAlloc(SIZE_T size)
 383 {
 384     return IMalloc_Alloc(&allocator.IMalloc_iface, size);
 385 }
 386
 387 /***********************************************************************
 388  *           CoTaskMemFree          (combase.@)
 389  */
 390 void WINAPI CoTaskMemFree(void *ptr)
 391 {
 392     IMalloc_Free(&allocator.IMalloc_iface, ptr);
 393 }
 394
 395 /***********************************************************************
 396  *           CoTaskMemRealloc        (combase.@)
 397  */
 398 void * WINAPI CoTaskMemRealloc(void *ptr, SIZE_T size)
 399 {
 400     return IMalloc_Realloc(&allocator.IMalloc_iface, ptr, size);
 401 }
 402
 403 /***********************************************************************
 404  *           CoRegisterMallocSpy        (combase.@)
 405  */
 406 HRESULT WINAPI CoRegisterMallocSpy(IMallocSpy *spy)
 407 {
 408     HRESULT hr = E_INVALIDARG;
 409
 410     TRACE("%p.\n", spy);
 411
 412     if (!spy) return E_INVALIDARG;
 413
 414     EnterCriticalSection(&allocspy_cs);
 415
 416     if (allocator.spy)
 417         hr = CO_E_OBJISREG;
 418     else if (SUCCEEDED(IMallocSpy_QueryInterface(spy, &IID_IMallocSpy, (void **)&spy)))
 419     {
 420         allocator.spy = spy;
 421         hr = S_OK;
 422     }
 423
 424     LeaveCriticalSection(&allocspy_cs);
 425
 426     return hr;
 427 }
 428
 429 /***********************************************************************
 430  *           CoRevokeMallocSpy (combase.@)
 431  */
 432 HRESULT WINAPI CoRevokeMallocSpy(void)
 433 {
 434     HRESULT hr = S_OK;
 435
 436     TRACE("\n");
 437
 438     EnterCriticalSection(&allocspy_cs);
 439
 440     if (!allocator.spy)
 441         hr = CO_E_OBJNOTREG;
 442     else if (allocator.spyed_allocations)
 443     {
 444         allocator.spy_release_pending = TRUE;
 445         hr = E_ACCESSDENIED;
 446     }
 447     else
 448     {
 449         IMallocSpy_Release(allocator.spy);
 450         allocator.spy = NULL;
 451     }
 452
 453     LeaveCriticalSection(&allocspy_cs);
 454
 455     return hr;
 456 }