webservices: Enforce the heap limit.

Signed-off-by: Hans Leidekker <hans@codeweavers.com>
Signed-off-by: Alexandre Julliard <julliard@winehq.org>

dlls/webservices/reader.c

@@ -170,6 +170,8 @@ static const struct prop_desc heap_props[] =
 struct heap
 {
     HANDLE      handle;
+    SIZE_T      max_size;
+    SIZE_T      allocated;
     ULONG       prop_count;
     struct prop prop[sizeof(heap_props)/sizeof(heap_props[0])];
 };
@@ -178,45 +180,75 @@ static BOOL ensure_heap( struct heap *heap )
 {
     SIZE_T size;
     if (heap->handle) return TRUE;
-    if (prop_get( heap->prop, heap->prop_count, WS_HEAP_PROPERTY_MAX_SIZE, &size, sizeof(size) ) != S_OK)
+    prop_get( heap->prop, heap->prop_count, WS_HEAP_PROPERTY_MAX_SIZE, &size, sizeof(size) );
-        return FALSE;
+    if (!(heap->handle = HeapCreate( 0, 0, 0 ))) return FALSE;
-    if (!(heap->handle = HeapCreate( 0, 0, size ))) return FALSE;
+    heap->max_size  = size;
+    heap->allocated = 0;
     return TRUE;
 }
 
 void *ws_alloc( WS_HEAP *handle, SIZE_T size )
 {
+    void *ret;
     struct heap *heap = (struct heap *)handle;
-    if (!ensure_heap( heap )) return NULL;
+    if (!ensure_heap( heap ) || size > heap->max_size - heap->allocated) return NULL;
-    return HeapAlloc( heap->handle, 0, size );
+    if ((ret = HeapAlloc( heap->handle, 0, size ))) heap->allocated += size;
+    return ret;
 }
 
 static void *ws_alloc_zero( WS_HEAP *handle, SIZE_T size )
 {
+    void *ret;
     struct heap *heap = (struct heap *)handle;
-    if (!ensure_heap( heap )) return NULL;
+    if (!ensure_heap( heap ) || size > heap->max_size - heap->allocated) return NULL;
-    return HeapAlloc( heap->handle, HEAP_ZERO_MEMORY, size );
+    if ((ret = HeapAlloc( heap->handle, HEAP_ZERO_MEMORY, size ))) heap->allocated += size;
+    return ret;
 }
 
-void *ws_realloc( WS_HEAP *handle, void *ptr, SIZE_T size )
+void *ws_realloc( WS_HEAP *handle, void *ptr, SIZE_T old_size, SIZE_T new_size )
 {
+    void *ret;
     struct heap *heap = (struct heap *)handle;
     if (!ensure_heap( heap )) return NULL;
-    return HeapReAlloc( heap->handle, 0, ptr, size );
+    if (new_size >= old_size)
+    {
+        SIZE_T size = new_size - old_size;
+        if (size > heap->max_size - heap->allocated) return NULL;
+        if ((ret = HeapReAlloc( heap->handle, 0, ptr, new_size ))) heap->allocated += size;
+    }
+    else
+    {
+        SIZE_T size = old_size - new_size;
+        if ((ret = HeapReAlloc( heap->handle, 0, ptr, new_size ))) heap->allocated -= size;
+    }
+    return ret;
 }
 
-static void *ws_realloc_zero( WS_HEAP *handle, void *ptr, SIZE_T size )
+static void *ws_realloc_zero( WS_HEAP *handle, void *ptr, SIZE_T old_size, SIZE_T new_size )
 {
+    void *ret;
     struct heap *heap = (struct heap *)handle;
     if (!ensure_heap( heap )) return NULL;
-    return HeapReAlloc( heap->handle, HEAP_ZERO_MEMORY, ptr, size );
+    if (new_size >= old_size)
+    {
+        SIZE_T size = new_size - old_size;
+        if (size > heap->max_size - heap->allocated) return NULL;
+        if ((ret = HeapReAlloc( heap->handle, HEAP_ZERO_MEMORY, ptr, new_size ))) heap->allocated += size;
+    }
+    else
+    {
+        SIZE_T size = old_size - new_size;
+        if ((ret = HeapReAlloc( heap->handle, HEAP_ZERO_MEMORY, ptr, new_size ))) heap->allocated -= size;
+    }
+    return ret;
 }
 
-void ws_free( WS_HEAP *handle, void *ptr )
+void ws_free( WS_HEAP *handle, void *ptr, SIZE_T size )
 {
     struct heap *heap = (struct heap *)handle;
     if (!heap->handle) return;
     HeapFree( heap->handle, 0, ptr );
+    heap->allocated -= size;
 }
 
 /**************************************************************************
@@ -230,8 +262,7 @@ HRESULT WINAPI WsAlloc( WS_HEAP *handle, SIZE_T size, void **ptr, WS_ERROR *erro
     if (error) FIXME( "ignoring error parameter\n" );
 
     if (!handle || !ptr) return E_INVALIDARG;
-
+    if (!(mem = ws_alloc( handle, size ))) return WS_E_QUOTA_EXCEEDED;
-    if (!(mem = ws_alloc( handle, size ))) return E_OUTOFMEMORY;
     *ptr = mem;
     return S_OK;
 }
@@ -296,7 +327,8 @@ HRESULT WINAPI WsResetHeap( WS_HEAP *handle, WS_ERROR *error )
     if (!heap) return E_INVALIDARG;
 
     HeapDestroy( heap->handle );
-    heap->handle = NULL;
+    heap->handle   = NULL;
+    heap->max_size = heap->allocated = 0;
     return S_OK;
 }
 
@@ -3858,7 +3890,8 @@ static HRESULT read_type_repeating_element( struct reader *reader, const WS_FIEL
     {
         if (nb_items >= nb_allocated)
         {
-            if (!(buf = ws_realloc_zero( heap, buf, nb_allocated * 2 * item_size )))
+            SIZE_T old_size = nb_allocated * item_size, new_size = old_size * 2;
+            if (!(buf = ws_realloc_zero( heap, buf, old_size, new_size )))
                 return WS_E_QUOTA_EXCEEDED;
             nb_allocated *= 2;
         }
@@ -3867,7 +3900,7 @@ static HRESULT read_type_repeating_element( struct reader *reader, const WS_FIEL
         if (hr == WS_E_INVALID_FORMAT) break;
         if (hr != S_OK)
         {
-            ws_free( heap, buf );
+            ws_free( heap, buf, nb_allocated * item_size );
             return hr;
         }
         offset += item_size;
@@ -3880,7 +3913,7 @@ static HRESULT read_type_repeating_element( struct reader *reader, const WS_FIEL
     {
         TRACE( "number of items %u out of range (%u-%u)\n", nb_items, desc->itemRange->minItemCount,
                desc->itemRange->maxItemCount );
-        ws_free( heap, buf );
+        ws_free( heap, buf, nb_allocated * item_size );
         return WS_E_INVALID_FORMAT;
     }
 
@@ -4035,7 +4068,7 @@ static HRESULT read_type_struct( struct reader *reader, WS_TYPE_MAPPING mapping,
     case WS_READ_REQUIRED_POINTER:
         if (hr != S_OK)
         {
-            ws_free( heap, buf );
+            ws_free( heap, buf, desc->size );
             return hr;
         }
         *(char **)ret = buf;
@@ -4045,7 +4078,7 @@ static HRESULT read_type_struct( struct reader *reader, WS_TYPE_MAPPING mapping,
     case WS_READ_NILLABLE_POINTER:
         if (is_nil_value( buf, desc->size ))
         {
-            ws_free( heap, buf );
+            ws_free( heap, buf, desc->size );
             buf = NULL;
         }
         *(char **)ret = buf;

dlls/webservices/tests/reader.c

@@ -1762,15 +1762,15 @@ static void test_WsAlloc(void)
     hr = WsCreateHeap( 256, 0, NULL, 0, &heap, NULL );
     ok( hr == S_OK, "got %08x\n", hr );
 
-    ptr = NULL;
+    ptr = (void *)0xdeadbeef;
     hr = WsAlloc( NULL, 16, &ptr, NULL );
     ok( hr == E_INVALIDARG, "got %08x\n", hr );
-    ok( ptr == NULL, "ptr set\n" );
+    ok( ptr == (void *)0xdeadbeef, "ptr set\n" );
 
-    ptr = NULL;
+    ptr = (void *)0xdeadbeef;
     hr = WsAlloc( heap, 512, &ptr, NULL );
-    todo_wine ok( hr == WS_E_QUOTA_EXCEEDED, "got %08x\n", hr );
+    ok( hr == WS_E_QUOTA_EXCEEDED, "got %08x\n", hr );
-    todo_wine ok( ptr == NULL, "ptr not set\n" );
+    ok( ptr == (void *)0xdeadbeef, "ptr set\n" );
 
     ptr = NULL;
     hr = WsAlloc( heap, 16, &ptr, NULL );

dlls/webservices/url.c

@@ -177,7 +177,7 @@ HRESULT WINAPI WsDecodeUrl( const WS_STRING *str, ULONG flags, WS_HEAP *heap, WS
     HRESULT hr = WS_E_QUOTA_EXCEEDED;
     WCHAR *p, *q, *decoded = NULL;
     WS_HTTP_URL *url = NULL;
-    ULONG len, port = 0;
+    ULONG len, len_decoded, port = 0;
 
     TRACE( "%s %08x %p %p %p\n", str ? debugstr_wn(str->chars, str->length) : "null", flags,
            heap, ret, error );
@@ -190,12 +190,13 @@ HRESULT WINAPI WsDecodeUrl( const WS_STRING *str, ULONG flags, WS_HEAP *heap, WS
         FIXME( "unimplemented flags %08x\n", flags );
         return E_NOTIMPL;
     }
-    if (!(decoded = url_decode( str->chars, str->length, heap, &len )) ||
+    if (!(decoded = url_decode( str->chars, str->length, heap, &len_decoded )) ||
         !(url = ws_alloc( heap, sizeof(*url) ))) goto error;
 
     hr = WS_E_INVALID_FORMAT;
 
     p = q = decoded;
+    len = len_decoded;
     while (len && *q != ':') { q++; len--; };
     if (*q != ':') goto error;
     if ((url->url.scheme = scheme_type( p, q - p )) == ~0u) goto error;
@@ -259,8 +260,8 @@ HRESULT WINAPI WsDecodeUrl( const WS_STRING *str, ULONG flags, WS_HEAP *heap, WS
     return S_OK;
 
 error:
-    if (decoded != str->chars) ws_free( heap, decoded );
+    if (decoded != str->chars) ws_free( heap, decoded, len_decoded );
-    ws_free( heap, url );
+    ws_free( heap, url, sizeof(*url) );
     return hr;
 }
 
@@ -418,7 +419,7 @@ HRESULT WINAPI WsEncodeUrl( const WS_URL *base, ULONG flags, WS_HEAP *heap, WS_S
                             WS_ERROR *error )
 {
     static const WCHAR fmtW[] = {':','%','u',0};
-    ULONG len = 0, len_scheme, len_enc;
+    ULONG len = 0, len_scheme, len_enc, ret_size;
     const WS_HTTP_URL *url = (const WS_HTTP_URL *)base;
     const WCHAR *scheme;
     WCHAR *str, *p, *q;
@@ -454,7 +455,8 @@ HRESULT WINAPI WsEncodeUrl( const WS_URL *base, ULONG flags, WS_HEAP *heap, WS_S
         return hr;
     len += len_enc + 1; /* '#' */
 
-    if (!(str = ws_alloc( heap, len * sizeof(WCHAR) ))) return WS_E_QUOTA_EXCEEDED;
+    ret_size = len * sizeof(WCHAR);
+    if (!(str = ws_alloc( heap, ret_size ))) return WS_E_QUOTA_EXCEEDED;
 
     memcpy( str, scheme, len_scheme * sizeof(WCHAR) );
     p = str + len_scheme;
@@ -520,6 +522,6 @@ HRESULT WINAPI WsEncodeUrl( const WS_URL *base, ULONG flags, WS_HEAP *heap, WS_S
     return S_OK;
 
 error:
-    ws_free( heap, str );
+    ws_free( heap, str, ret_size );
     return hr;
 }

dlls/webservices/webservices_private.h

@@ -27,8 +27,8 @@ struct xmlbuf
 };
 
 void *ws_alloc( WS_HEAP *, SIZE_T ) DECLSPEC_HIDDEN;
-void *ws_realloc( WS_HEAP *, void *, SIZE_T ) DECLSPEC_HIDDEN;
+void *ws_realloc( WS_HEAP *, void *, SIZE_T, SIZE_T ) DECLSPEC_HIDDEN;
-void ws_free( WS_HEAP *, void * ) DECLSPEC_HIDDEN;
+void ws_free( WS_HEAP *, void *, SIZE_T ) DECLSPEC_HIDDEN;
 const char *debugstr_xmlstr( const WS_XML_STRING * ) DECLSPEC_HIDDEN;
 WS_XML_STRING *alloc_xml_string( const unsigned char *, ULONG ) DECLSPEC_HIDDEN;
 WS_XML_UTF8_TEXT *alloc_utf8_text( const unsigned char *, ULONG ) DECLSPEC_HIDDEN;

dlls/webservices/writer.c

@@ -233,7 +233,7 @@ static struct xmlbuf *alloc_xmlbuf( WS_HEAP *heap )
     if (!(ret = ws_alloc( heap, sizeof(*ret) ))) return NULL;
     if (!(ret->ptr = ws_alloc( heap, XML_BUFFER_INITIAL_ALLOCATED_SIZE )))
     {
-        ws_free( heap, ret );
+        ws_free( heap, ret, sizeof(*ret) );
         return NULL;
     }
     ret->heap           = heap;
@@ -245,8 +245,8 @@ static struct xmlbuf *alloc_xmlbuf( WS_HEAP *heap )
 static void free_xmlbuf( struct xmlbuf *xmlbuf )
 {
     if (!xmlbuf) return;
-    ws_free( xmlbuf->heap, xmlbuf->ptr );
+    ws_free( xmlbuf->heap, xmlbuf->ptr, xmlbuf->size_allocated );
-    ws_free( xmlbuf->heap, xmlbuf );
+    ws_free( xmlbuf->heap, xmlbuf, sizeof(*xmlbuf) );
 }
 
 /**************************************************************************
@@ -260,7 +260,7 @@ HRESULT WINAPI WsCreateXmlBuffer( WS_HEAP *heap, const WS_XML_BUFFER_PROPERTY *p
     if (!heap || !handle) return E_INVALIDARG;
     if (count) FIXME( "properties not implemented\n" );
 
-    if (!(xmlbuf = alloc_xmlbuf( heap ))) return E_OUTOFMEMORY;
+    if (!(xmlbuf = alloc_xmlbuf( heap ))) return WS_E_QUOTA_EXCEEDED;
 
     *handle = (WS_XML_BUFFER *)xmlbuf;
     return S_OK;
@@ -355,7 +355,7 @@ HRESULT WINAPI WsSetOutput( WS_XML_WRITER *handle, const WS_XML_WRITER_ENCODING
     {
         struct xmlbuf *xmlbuf;
 
-        if (!(xmlbuf = alloc_xmlbuf( writer->output_heap ))) return E_OUTOFMEMORY;
+        if (!(xmlbuf = alloc_xmlbuf( writer->output_heap ))) return WS_E_QUOTA_EXCEEDED;
         set_output_buffer( writer, xmlbuf );
         break;
     }
@@ -414,7 +414,7 @@ static HRESULT write_grow_buffer( struct writer *writer, ULONG size )
         return S_OK;
     }
     new_size = max( buf->size_allocated * 2, writer->write_pos + size );
-    if (!(tmp = ws_realloc( buf->heap, buf->ptr, new_size ))) return E_OUTOFMEMORY;
+    if (!(tmp = ws_realloc( buf->heap, buf->ptr, buf->size_allocated, new_size ))) return WS_E_QUOTA_EXCEEDED;
     writer->write_bufptr = buf->ptr = tmp;
     buf->size_allocated = new_size;
     buf->size = writer->write_pos + size;