@@ -42,6 +42,11 @@ extern HANDLE g_currentProcess;
4242extern CriticalSection g_heapMapLock;
4343extern DbgHelp g_DbgHelp;
4444
45+ // Heap handles are pointers, hence disjunct from numbers below 65536, which
46+ // hence can serve as pseudo heap handles for tracking non-memory resources.
47+ #define VLD_SOCKET_RESOURCE MAKEINTRESOURCE (1 )
48+ #define VLD_WSAEVENT_RESOURCE MAKEINTRESOURCE (2 )
49+
4550// //////////////////////////////////////////////////////////////////////////////
4651//
4752// Debug CRT and MFC IAT Replacement Functions
@@ -298,6 +303,171 @@ LPVOID VisualLeakDetector::_HeapReAlloc (HANDLE heap, DWORD flags, LPVOID mem, S
298303 return newmem;
299304}
300305
306+ // //////////////////////////////////////////////////////////////////////////////
307+ //
308+ // Winsock IAT Replacement Functions
309+ //
310+ // //////////////////////////////////////////////////////////////////////////////
311+
312+ // _socket - Calls to socket are patched through to this function.
313+ // This function invokes the real socket and then calls VLD's allocation
314+ // tracking function with a pseudo heap handle of value VLD_SOCKET_RESOURCE.
315+ //
316+ // Return Value:
317+ //
318+ // Returns the return value from socket.
319+ //
320+ SOCKET VisualLeakDetector::_socket (int af, int type, int protocol)
321+ {
322+ PRINT_HOOKED_FUNCTION2 ();
323+ // Allocate the resource.
324+ SOCKET s = socket (af, type, protocol);
325+
326+ if ((s == INVALID_SOCKET) || !g_vld.enabled ())
327+ return s;
328+
329+ if (!g_DbgHelp.IsLockedByCurrentThread ()) // skip dbghelp.dll calls
330+ {
331+ CAPTURE_CONTEXT ();
332+ CaptureContext cc (socket, context_);
333+ cc.Set (VLD_SOCKET_RESOURCE, reinterpret_cast <LPVOID>(s), NULL , 0 );
334+ }
335+
336+ return s;
337+ }
338+
339+ // _accept - Calls to accept are patched through to this function.
340+ // This function invokes the real accept and then calls VLD's allocation
341+ // tracking function with a pseudo heap handle of value VLD_SOCKET_RESOURCE.
342+ //
343+ // Return Value:
344+ //
345+ // Returns the return value from accept.
346+ //
347+ SOCKET VisualLeakDetector::_accept (SOCKET s, struct sockaddr *addr, int *addrlen)
348+ {
349+ PRINT_HOOKED_FUNCTION2 ();
350+ // Allocate the resource.
351+ s = accept (s, addr, addrlen);
352+
353+ if ((s == INVALID_SOCKET) || !g_vld.enabled ())
354+ return s;
355+
356+ if (!g_DbgHelp.IsLockedByCurrentThread ()) // skip dbghelp.dll calls
357+ {
358+ CAPTURE_CONTEXT ();
359+ CaptureContext cc (accept, context_);
360+ cc.Set (VLD_SOCKET_RESOURCE, reinterpret_cast <LPVOID>(s), NULL , 0 );
361+ }
362+
363+ return s;
364+ }
365+
366+ // _connect - Calls to connect are patched through to this function.
367+ // This function invokes the real connect and then calls VLD's allocation
368+ // tracking function with a pseudo heap handle of value VLD_SOCKET_RESOURCE.
369+ //
370+ // Return Value:
371+ //
372+ // Returns the return value from connect.
373+ //
374+ SOCKET VisualLeakDetector::_connect (SOCKET s, const struct sockaddr *name, int namelen)
375+ {
376+ PRINT_HOOKED_FUNCTION2 ();
377+ // Allocate the resource.
378+ s = connect (s, name, namelen);
379+
380+ if ((s == INVALID_SOCKET) || !g_vld.enabled ())
381+ return s;
382+
383+ if (!g_DbgHelp.IsLockedByCurrentThread ()) // skip dbghelp.dll calls
384+ {
385+ CAPTURE_CONTEXT ();
386+ CaptureContext cc (connect, context_);
387+ cc.Set (VLD_SOCKET_RESOURCE, reinterpret_cast <LPVOID>(s), NULL , 0 );
388+ }
389+
390+ return s;
391+ }
392+
393+ // _closesocket - Calls to closesocket are patched through to this function.
394+ // This function calls VLD's free tracking function with a pseudo heap handle
395+ // of value VLD_SOCKET_RESOURCE and then invokes the real closesocket.
396+ //
397+ // Return Value:
398+ //
399+ // Returns the value returned by closesocket.
400+ //
401+ int VisualLeakDetector::_closesocket (SOCKET s)
402+ {
403+ PRINT_HOOKED_FUNCTION2 ();
404+
405+ if (!g_DbgHelp.IsLockedByCurrentThread ()) // skip dbghelp.dll calls
406+ {
407+ // Record the current frame pointer.
408+ CAPTURE_CONTEXT ();
409+ context_.func = reinterpret_cast <UINT_PTR>(closesocket);
410+
411+ // Unmap the resource from the specified pseudo heap.
412+ g_vld.unmapBlock (VLD_SOCKET_RESOURCE, reinterpret_cast <LPVOID>(s), context_);
413+ }
414+
415+ return closesocket (s);
416+ }
417+
418+ // _WSACreateEvent - Calls to WSACreateEvent are patched through to this
419+ // function. This function invokes the real WSACreateEvent and then calls
420+ // VLD's allocation tracking function with a pseudo heap handle of value
421+ // VLD_HANDLE_RESOURCE.
422+ //
423+ // Return Value:
424+ //
425+ // Returns the return value from WSACreateEvent.
426+ //
427+ WSAEVENT VisualLeakDetector::_WSACreateEvent ()
428+ {
429+ PRINT_HOOKED_FUNCTION2 ();
430+ // Allocate the resource.
431+ WSAEVENT hEvent = WSACreateEvent ();
432+
433+ if ((hEvent == WSA_INVALID_EVENT) || !g_vld.enabled ())
434+ return hEvent;
435+
436+ if (!g_DbgHelp.IsLockedByCurrentThread ()) // skip dbghelp.dll calls
437+ {
438+ CAPTURE_CONTEXT ();
439+ CaptureContext cc (WSACreateEvent, context_);
440+ cc.Set (VLD_WSAEVENT_RESOURCE, hEvent, NULL , 0 );
441+ }
442+
443+ return hEvent;
444+ }
445+
446+ // _WSACloseEvent - Calls to WSACloseEvent are patched through to this function.
447+ // This function calls VLD's free tracking function with a pseudo heap handle
448+ // of value VLD_HANDLE_RESOURCE and then invokes the real WSACloseEvent.
449+ //
450+ // Return Value:
451+ //
452+ // Returns the value returned by closesocket.
453+ //
454+ BOOL VisualLeakDetector::_WSACloseEvent (WSAEVENT hEvent)
455+ {
456+ PRINT_HOOKED_FUNCTION2 ();
457+
458+ if (!g_DbgHelp.IsLockedByCurrentThread ()) // skip dbghelp.dll calls
459+ {
460+ // Record the current frame pointer.
461+ CAPTURE_CONTEXT ();
462+ context_.func = reinterpret_cast <UINT_PTR>(WSACloseEvent);
463+
464+ // Unmap the resource from the specified pseudo heap.
465+ g_vld.unmapBlock (VLD_WSAEVENT_RESOURCE, hEvent, context_);
466+ }
467+
468+ return WSACloseEvent (hEvent);
469+ }
470+
301471// //////////////////////////////////////////////////////////////////////////////
302472//
303473// COM IAT Replacement Functions
0 commit comments