From 14d3231004d06d00f332a2bfdf856696e3fd21d2 Mon Sep 17 00:00:00 2001
From: Antoine Cordelle <antoinecordelle@improbable.io>
Date: Fri, 5 Mar 2021 12:40:35 +0000
Subject: [PATCH] Add custom memory allocator support for session callbacks

---
 lib/includes/nghttp2/nghttp2.h | 25 +++++++++++++++++++++++++
 lib/nghttp2_callbacks.c        | 22 ++++++++++++++++++++--
 lib/nghttp2_callbacks.h        |  4 ++++
 3 files changed, 49 insertions(+), 2 deletions(-)

diff --git a/lib/includes/nghttp2/nghttp2.h b/lib/includes/nghttp2/nghttp2.h
index 2d7fd25e..2b87ad7f 100644
--- a/lib/includes/nghttp2/nghttp2.h
+++ b/lib/includes/nghttp2/nghttp2.h
@@ -2227,6 +2227,31 @@ typedef struct nghttp2_session_callbacks nghttp2_session_callbacks;
 NGHTTP2_EXTERN int
 nghttp2_session_callbacks_new(nghttp2_session_callbacks **callbacks_ptr);
 
+/**
+ * @function
+ *
+ * Like `nghttp2_session_callbacks_new()`, but with additional custom
+ * memory allocator specified in the |mem|.
+ *
+ * The |mem| can be ``NULL`` and the call is equivalent to
+ * `nghttp2_session_callbacks_new()`.
+ *
+ * This function does not take ownership |mem|.  The application is
+ * responsible for freeing |mem|.
+ *
+ * The library code does not refer to |mem| pointer after this
+ * function returns, so the application can safely free it.
+ *
+ * This function returns 0 if it succeeds, or one of the following
+ * negative error codes:
+ *
+ * :enum:`nghttp2_error.NGHTTP2_ERR_NOMEM`
+ *     Out of memory.
+ */
+NGHTTP2_EXTERN int
+nghttp2_session_callbacks_new2(nghttp2_session_callbacks **callbacks_ptr,
+                               nghttp2_mem *mem);
+
 /**
  * @function
  *
diff --git a/lib/nghttp2_callbacks.c b/lib/nghttp2_callbacks.c
index 3c382148..8ce7220c 100644
--- a/lib/nghttp2_callbacks.c
+++ b/lib/nghttp2_callbacks.c
@@ -23,21 +23,39 @@
  * WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
  */
 #include "nghttp2_callbacks.h"
+#include "nghttp2_mem.h"
 
 #include <stdlib.h>
 
 int nghttp2_session_callbacks_new(nghttp2_session_callbacks **callbacks_ptr) {
-  *callbacks_ptr = calloc(1, sizeof(nghttp2_session_callbacks));
+  nghttp2_session_callbacks_new2(callbacks_ptr, NULL);
+}
+
+int nghttp2_session_callbacks_new2(nghttp2_session_callbacks **callbacks_ptr,
+                                   nghttp2_mem *mem) {
+  if (mem == NULL) {
+    mem = nghttp2_mem_default();
+  }
+
+  *callbacks_ptr = nghttp2_mem_calloc(mem, 1, sizeof(nghttp2_session_callbacks));
 
   if (*callbacks_ptr == NULL) {
     return NGHTTP2_ERR_NOMEM;
   }
 
+  (*callbacks_ptr)->mem = *mem;
+
   return 0;
 }
 
 void nghttp2_session_callbacks_del(nghttp2_session_callbacks *callbacks) {
-  free(callbacks);
+  if (callbacks == NULL) {
+    return;
+  }
+  nghttp2_mem *mem;
+
+  mem = &callbacks->mem;
+  nghttp2_mem_free(mem, callbacks);
 }
 
 void nghttp2_session_callbacks_set_send_callback(
diff --git a/lib/nghttp2_callbacks.h b/lib/nghttp2_callbacks.h
index 61e51fa5..0297e223 100644
--- a/lib/nghttp2_callbacks.h
+++ b/lib/nghttp2_callbacks.h
@@ -120,6 +120,10 @@ struct nghttp2_session_callbacks {
   nghttp2_on_extension_chunk_recv_callback on_extension_chunk_recv_callback;
   nghttp2_error_callback error_callback;
   nghttp2_error_callback2 error_callback2;
+  /**
+   * Memory allocator
+   */
+  nghttp2_mem mem;
 };
 
 #endif /* NGHTTP2_CALLBACKS_H */