Extend the integration test harness to track FDs

IntegrationTests/allocation-counter-tests-framework/template/AtomicCounter/Sources/AtomicCounter/include/atomic-counter.h

@@ -29,4 +29,14 @@ void add_malloc_bytes_counter(intptr_t v);
 void reset_malloc_bytes_counter(void);
 intptr_t read_malloc_bytes_counter(void);
 
+typedef struct {
+    size_t count;
+    int *leaked;
+} LeakedFDs;
+
+void begin_tracking_fds(void);
+void track_open_fd(int fd);
+void track_closed_fd(int fd);
+LeakedFDs stop_tracking_fds(void);
+
 #endif

IntegrationTests/allocation-counter-tests-framework/template/AtomicCounter/Sources/AtomicCounter/src/atomic-counter.c

@@ -2,7 +2,7 @@
 //
 // This source file is part of the SwiftNIO open source project
 //
-// Copyright (c) 2017-2018 Apple Inc. and the SwiftNIO project authors
+// Copyright (c) 2017-2023 Apple Inc. and the SwiftNIO project authors
 // Licensed under Apache License v2.0
 //
 // See LICENSE.txt for license information
@@ -12,7 +12,13 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include <atomic-counter.h>
+#include <assert.h>
+#include <pthread.h>
 #include <stdatomic.h>
+#include <stdbool.h>
+#include <stdlib.h>
+#include <string.h>
 
 #define MAKE_COUNTER(name) /*
 */ _Atomic long g_ ## name ## _counter = ATOMIC_VAR_INIT(0); /*
@@ -34,3 +40,127 @@
 MAKE_COUNTER(free)
 MAKE_COUNTER(malloc)
 MAKE_COUNTER(malloc_bytes)
+
+// This section covers tracking leaked FDs.
+//
+// We do this by recording which FD has been set in a queue. A queue is a bad data structure here,
+// but using a better one requires writing too much code, and the performance impact here is not
+// going to be too bad.
+typedef struct {
+    size_t capacity;
+    size_t count;
+    int *allocatedFDs;
+} FDTracker;
+
+static _Bool FDTracker_search_fd(const FDTracker *tracker, int fd, size_t *foundIndex) {
+    if (tracker == NULL) { return false; }
+
+    for (size_t i = 0; i < tracker->count; i++) {
+        if (tracker->allocatedFDs[i] == fd) {
+            if (foundIndex != NULL) { 
+                *foundIndex = i;
+            }
+            return true;
+        }
+    }
+
+    return false;
+}
+
+static void FDTracker_remove_at_index(FDTracker *tracker, size_t index) {
+    assert(tracker != NULL);
+    assert(index < tracker->count);
+
+    // Shuffle everything down by 1 from index onwards.
+    const size_t lastValidTargetIndex = tracker->count - 1;
+    for (size_t i = index; i < lastValidTargetIndex; i++) {
+        tracker->allocatedFDs[i] = tracker->allocatedFDs[i + 1];
+    }
+    tracker->count--;
+}
+
+_Atomic _Bool is_tracking = ATOMIC_VAR_INIT(false);
+pthread_mutex_t tracker_lock = PTHREAD_MUTEX_INITIALIZER;
+FDTracker tracker = { 0 };
+
+void begin_tracking_fds(void) {
+    int rc = pthread_mutex_lock(&tracker_lock);
+    assert(rc == 0);
+
+    assert(tracker.capacity == 0);
+    assert(tracker.count == 0);
+    assert(tracker.allocatedFDs == NULL);
+
+    tracker.allocatedFDs = calloc(1024, sizeof(int));
+    tracker.capacity = 1024;
+
+    atomic_store_explicit(&is_tracking, true, memory_order_release);
+    rc = pthread_mutex_unlock(&tracker_lock);
+    assert(rc == 0);
+}
+
+void track_open_fd(int fd) {
+    bool should_track = atomic_load_explicit(&is_tracking, memory_order_acquire);
+    if (!should_track) { return; }
+
+    int rc = pthread_mutex_lock(&tracker_lock);
+    assert(rc == 0);
+
+    // We need to not be tracking this FD already, or there's a correctness error.
+    assert(!FDTracker_search_fd(&tracker, fd, NULL));
+
+    // We want to append to the queue.
+    if (tracker.capacity == tracker.count) {
+        // Wuh-oh, resize. We do this by doubling.
+        assert((tracker.capacity * sizeof(int)) < (SIZE_MAX / 2));
+        size_t newCapacity = tracker.capacity * 2;
+        int *new = realloc(tracker.allocatedFDs, newCapacity * sizeof(int));
+        assert(new != NULL);
+        tracker.allocatedFDs = new;
+        tracker.capacity = newCapacity;
+    }
+
+    tracker.allocatedFDs[tracker.count] = fd;
+    tracker.count++;
+
+    rc = pthread_mutex_unlock(&tracker_lock);
+    assert(rc == 0);
+}
+
+void track_closed_fd(int fd) {
+    bool should_track = atomic_load_explicit(&is_tracking, memory_order_acquire);
+    if (!should_track) { return; }
+
+    int rc = pthread_mutex_lock(&tracker_lock);
+    assert(rc == 0);
+
+    size_t index;
+    if (FDTracker_search_fd(&tracker, fd, &index)) {
+        // We're tracking this FD, let's remove it.
+        FDTracker_remove_at_index(&tracker, index);
+    }
+
+    rc = pthread_mutex_unlock(&tracker_lock);
+    assert(rc == 0);
+}
+
+LeakedFDs stop_tracking_fds(void) {
+    int rc = pthread_mutex_lock(&tracker_lock);
+    assert(rc == 0);
+
+    LeakedFDs result = {
+        .count = tracker.count,
+        .leaked = tracker.allocatedFDs
+    };
+
+    // Clear the tracker.
+    tracker.allocatedFDs = NULL;
+    tracker.capacity = 0;
+    tracker.count = 0;
+
+    atomic_store_explicit(&is_tracking, false, memory_order_release);
+    rc = pthread_mutex_unlock(&tracker_lock);
+    assert(rc == 0);
+
+    return result;
+}

IntegrationTests/allocation-counter-tests-framework/template/HookedFunctionsDoHook/Sources/HookedFunctions/include/hooked-functions.h

@@ -16,6 +16,8 @@
 #define HOOKED_FREE
 
 #include <stdlib.h>
+#include <sys/socket.h>
+#include <sys/types.h>
 #if __APPLE__
 #  include <malloc/malloc.h>
 #endif
@@ -27,6 +29,10 @@ void *replacement_realloc(void *ptr, size_t size);
 void *replacement_reallocf(void *ptr, size_t size);
 void *replacement_valloc(size_t size);
 int replacement_posix_memalign(void **memptr, size_t alignment, size_t size);
+int replacement_socket(int domain, int type, int protocol);
+int replacement_accept(int socket, struct sockaddr *restrict address, socklen_t *restrict address_len);
+int replacement_accept4(int sockfd, struct sockaddr *addr, socklen_t *addrlen, int flags);
+int replacement_close(int fildes);
 
 #if __APPLE__
 void *replacement_malloc_zone_malloc(malloc_zone_t *zone, size_t size);

IntegrationTests/allocation-counter-tests-framework/template/HookedFunctionsDoHook/Sources/HookedFunctions/src/hooked-functions-darwin.c

@@ -12,6 +12,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include <assert.h>
 #if __APPLE__
 
 #define _GNU_SOURCE
@@ -23,6 +24,7 @@
 #include <stdlib.h>
 #include <string.h>
 #include <sys/mman.h>
+#include <sys/socket.h>
 #include <sys/stat.h>
 #include <sys/types.h>
 #include <unistd.h>
@@ -148,6 +150,37 @@ int replacement_posix_memalign(void **memptr, size_t alignment, size_t size) {
     JUMP_INTO_LIBC_FUN(posix_memalign, memptr, alignment, size);
 }
 
+int replacement_socket(int domain, int type, int protocol) {
+    int fd = socket(domain, type, protocol);
+    if (fd < 0) {
+        return fd;
+    }
+
+    track_open_fd(fd);
+    return fd;
+}
+
+int replacement_accept(int socket, struct sockaddr *restrict address, socklen_t *restrict address_len) {
+    int fd = accept(socket, address, address_len);
+
+    if (fd < 0) {
+        return fd;
+    }
+
+    track_open_fd(fd);
+    return fd;
+}
+
+int replacement_accept4(int sockfd, struct sockaddr *addr, socklen_t *addrlen, int flags) {
+    // Should never be called.
+    assert(false);
+}
+
+int replacement_close(int fildes) {
+    track_closed_fd(fildes);
+    JUMP_INTO_LIBC_FUN(close, fildes);
+}
+
 DYLD_INTERPOSE(replacement_free, free)
 DYLD_INTERPOSE(replacement_malloc, malloc)
 DYLD_INTERPOSE(replacement_realloc, realloc)
@@ -161,4 +194,7 @@ DYLD_INTERPOSE(replacement_malloc_zone_valloc, malloc_zone_valloc)
 DYLD_INTERPOSE(replacement_malloc_zone_realloc, malloc_zone_realloc)
 DYLD_INTERPOSE(replacement_malloc_zone_memalign, malloc_zone_memalign)
 DYLD_INTERPOSE(replacement_malloc_zone_free, malloc_zone_free)
+DYLD_INTERPOSE(replacement_socket, socket)
+DYLD_INTERPOSE(replacement_accept, accept)
+DYLD_INTERPOSE(replacement_close, close)
 #endif

IntegrationTests/allocation-counter-tests-framework/template/HookedFunctionsDoHook/Sources/HookedFunctions/src/hooked-functions-unix.c

@@ -41,16 +41,28 @@ static _Atomic ptrdiff_t g_recursive_malloc_next_free_ptr = ATOMIC_VAR_INIT(0);
 static __thread bool g_in_malloc = false;
 static __thread bool g_in_realloc = false;
 static __thread bool g_in_free = false;
+static __thread bool g_in_socket = false;
+static __thread bool g_in_accept = false;
+static __thread bool g_in_accept4 = false;
+static __thread bool g_in_close = false;
 
 /* The types of the variables holding the libc function pointers. */
 typedef void *(*type_libc_malloc)(size_t);
 typedef void *(*type_libc_realloc)(void *, size_t);
 typedef void  (*type_libc_free)(void *);
+typedef int   (*type_libc_socket)(int, int, int);
+typedef int   (*type_libc_accept)(int, struct sockaddr*, socklen_t *);
+typedef int   (*type_libc_accept4)(int, struct sockaddr *, socklen_t *, int);
+typedef int   (*type_libc_close)(int);
 
 /* The (atomic) globals holding the pointer to the original libc implementation. */
 _Atomic type_libc_malloc g_libc_malloc;
 _Atomic type_libc_realloc g_libc_realloc;
 _Atomic type_libc_free g_libc_free;
+_Atomic type_libc_socket g_libc_socket;
+_Atomic type_libc_accept g_libc_accept;
+_Atomic type_libc_accept4 g_libc_accept4;
+_Atomic type_libc_close g_libc_close;
 
 // this is called if malloc is called whilst trying to resolve libc's realloc.
 // we just vend out pointers to a large block in the BSS (which we never free).
@@ -93,6 +105,30 @@ static void recursive_free(void *ptr) {
     abort();
 }
 
+// this is called if socket is called whilst trying to resolve libc's socket.
+static int recursive_socket(int domain, int type, int protocol) {
+    // not possible
+    abort();
+}
+
+// this is called if accept is called whilst trying to resolve libc's accept.
+static int recursive_accept(int socket, struct sockaddr *restrict address, socklen_t *restrict address_len) {
+    // not possible
+    abort();
+}
+
+// this is called if accept4 is called whilst trying to resolve libc's accept4.
+static int recursive_accept4(int sockfd, struct sockaddr *addr, socklen_t *addrlen, int flags) {
+    // not possible
+    abort();
+}
+
+// this is called if close is called whilst trying to resolve libc's close.
+static int recursive_close(int fildes) {
+    // not possible
+    abort();
+}
+
 /* On Apple platforms getting to the original libc function from a hooked
  * function is easy.  On other UNIX systems this is slightly harder because we
  * have to look up the function with the dynamic linker.  Because that isn't
@@ -192,4 +228,53 @@ int replacement_posix_memalign(void **memptr, size_t alignment, size_t size) {
     }
 }
 
+static int socket_thunk(int domain, int type, int protocol) {
+    JUMP_INTO_LIBC_FUN(socket, domain, type, protocol);
+}
+
+int replacement_socket(int domain, int type, int protocol) {
+    int fd = socket_thunk(domain, type, protocol);
+    if (fd < 0) {
+        return fd;
+    }
+
+    track_open_fd(fd);
+    return fd;
+}
+
+static int accept_thunk(int socket, struct sockaddr *restrict address, socklen_t *restrict address_len) {
+    JUMP_INTO_LIBC_FUN(accept, socket, address, address_len);
+}
+
+int replacement_accept(int socket, struct sockaddr *restrict address, socklen_t *restrict address_len) {
+    int fd = accept_thunk(socket, address, address_len);
+
+    if (fd < 0) {
+        return fd;
+    }
+
+    track_open_fd(fd);
+    return fd;
+}
+
+static int accept4_thunk(int sockfd, struct sockaddr *addr, socklen_t *addrlen, int flags) {
+    JUMP_INTO_LIBC_FUN(accept4, sockfd, addr, addrlen, flags);
+}
+
+int replacement_accept4(int sockfd, struct sockaddr *addr, socklen_t *addrlen, int flags) {
+    int fd = accept4_thunk(sockfd, addr, addrlen, flags);
+
+    if (fd < 0) {
+        return fd;
+    }
+
+    track_open_fd(fd);
+    return fd;
+}
+
+int replacement_close(int fildes) {
+    track_closed_fd(fildes);
+    JUMP_INTO_LIBC_FUN(close, fildes);
+}
+
 #endif

IntegrationTests/allocation-counter-tests-framework/template/Sources/bootstrapDoHook/main.c

@@ -41,6 +41,18 @@ void *valloc(size_t size) {
 int posix_memalign(void **memptr, size_t alignment, size_t size) {
     return replacement_posix_memalign(memptr, alignment, size);
 }
+int socket(int domain, int type, int protocol) {
+    return replacement_socket(domain, type, protocol);
+}
+int accept(int socket, struct sockaddr *restrict address, socklen_t *restrict address_len) {
+    return replacement_accept(socket, address, address_len);
+}
+int accept4(int sockfd, struct sockaddr *addr, socklen_t *addrlen, int flags) {
+    return replacement_accept4(sockfd, addr, addrlen, flags);
+}
+int close(int fildes) {
+    return replacement_close(fildes);
+}
 #endif
 
 void swift_main(void);