From c2998391eea62a4580e96cb4b34f91f93806317a Mon Sep 17 00:00:00 2001
From: Tatsuhiro Tsujikawa <tatsuhiro.t@gmail.com>
Date: Thu, 2 Jan 2025 15:38:31 +0900
Subject: [PATCH] fuzz_qpackdecoder: Fuzz memory allocator

---
 fuzz/fuzz_qpackdecoder.cc | 61 +++++++++++++++++++++++++++++++++------
 lib/nghttp3_qpack.c       |  1 +
 2 files changed, 53 insertions(+), 9 deletions(-)

diff --git a/fuzz/fuzz_qpackdecoder.cc b/fuzz/fuzz_qpackdecoder.cc
index deb32a6..2b9e02a 100644
--- a/fuzz/fuzz_qpackdecoder.cc
+++ b/fuzz/fuzz_qpackdecoder.cc
@@ -29,6 +29,8 @@ struct Request {
   Request(int64_t stream_id, const nghttp3_buf *buf);
   ~Request();
 
+  int init(const nghttp3_mem &mem);
+
   nghttp3_buf buf;
   nghttp3_qpack_stream_context *sctx;
   int64_t stream_id;
@@ -48,7 +50,7 @@ using Headers = std::vector<std::pair<std::string, std::string>>;
 
 class Decoder {
 public:
-  Decoder(size_t max_dtable_size, size_t max_blocked);
+  Decoder(size_t max_dtable_size, size_t max_blocked, const nghttp3_mem &mem);
   ~Decoder();
 
   int init();
@@ -59,7 +61,7 @@ class Decoder {
   size_t get_num_blocked() const;
 
 private:
-  const nghttp3_mem *mem_;
+  const nghttp3_mem &mem_;
   nghttp3_qpack_decoder *dec_;
   std::priority_queue<std::shared_ptr<Request>,
                       std::vector<std::shared_ptr<Request>>,
@@ -70,15 +72,22 @@ class Decoder {
 };
 
 Request::Request(int64_t stream_id, const nghttp3_buf *buf)
-  : buf(*buf), stream_id(stream_id) {
-  auto mem = nghttp3_mem_default();
-  nghttp3_qpack_stream_context_new(&sctx, stream_id, mem);
+  : buf(*buf), sctx(nullptr), stream_id(stream_id) {}
+
+int Request::init(const nghttp3_mem &mem) {
+  auto rv = nghttp3_qpack_stream_context_new(&sctx, stream_id, &mem);
+  if (rv != 0) {
+    return -1;
+  }
+
+  return 0;
 }
 
 Request::~Request() { nghttp3_qpack_stream_context_del(sctx); }
 
-Decoder::Decoder(size_t max_dtable_size, size_t max_blocked)
-  : mem_(nghttp3_mem_default()),
+Decoder::Decoder(size_t max_dtable_size, size_t max_blocked,
+                 const nghttp3_mem &mem)
+  : mem_(mem),
     dec_(nullptr),
     max_dtable_size_(max_dtable_size),
     max_blocked_(max_blocked) {}
@@ -87,7 +96,7 @@ Decoder::~Decoder() { nghttp3_qpack_decoder_del(dec_); }
 
 int Decoder::init() {
   if (auto rv =
-        nghttp3_qpack_decoder_new(&dec_, max_dtable_size_, max_blocked_, mem_);
+        nghttp3_qpack_decoder_new(&dec_, max_dtable_size_, max_blocked_, &mem_);
       rv != 0) {
     return -1;
   }
@@ -113,6 +122,10 @@ std::tuple<Headers, int> Decoder::read_request(nghttp3_buf *buf,
                                                int64_t stream_id) {
   auto req = std::make_shared<Request>(stream_id, buf);
 
+  if (req->init(mem_) != 0) {
+    return {{}, -1};
+  }
+
   auto [headers, rv] = read_request(*req);
   if (rv == -1) {
     return {Headers{}, -1};
@@ -182,15 +195,45 @@ std::tuple<int64_t, Headers, int> Decoder::process_blocked() {
   return {-1, {}, 0};
 }
 
+namespace {
+void *fuzzed_malloc(size_t size, void *user_data) {
+  auto fuzzed_data_provider = static_cast<FuzzedDataProvider *>(user_data);
+
+  return fuzzed_data_provider->ConsumeBool() ? nullptr : malloc(size);
+}
+} // namespace
+
+namespace {
+void *fuzzed_calloc(size_t nmemb, size_t size, void *user_data) {
+  auto fuzzed_data_provider = static_cast<FuzzedDataProvider *>(user_data);
+
+  return fuzzed_data_provider->ConsumeBool() ? nullptr : calloc(nmemb, size);
+}
+} // namespace
+
+namespace {
+void *fuzzed_realloc(void *ptr, size_t size, void *user_data) {
+  auto fuzzed_data_provider = static_cast<FuzzedDataProvider *>(user_data);
+
+  return fuzzed_data_provider->ConsumeBool() ? nullptr : realloc(ptr, size);
+}
+} // namespace
+
 int decode(const uint8_t *data, size_t datalen) {
   FuzzedDataProvider fuzzed_data_provider(data, datalen);
 
+  auto mem = *nghttp3_mem_default();
+  mem.user_data = &fuzzed_data_provider;
+  mem.malloc = fuzzed_malloc;
+  mem.calloc = fuzzed_calloc;
+  mem.realloc = fuzzed_realloc;
+
   auto max_dtable_size =
     fuzzed_data_provider.ConsumeIntegralInRange<size_t>(0, NGHTTP3_MAX_VARINT);
   auto max_blocked =
     fuzzed_data_provider.ConsumeIntegralInRange<size_t>(0, NGHTTP3_MAX_VARINT);
 
-  auto dec = Decoder(max_dtable_size, max_blocked);
+  auto dec = Decoder(max_dtable_size, max_blocked, mem);
   if (auto rv = dec.init(); rv != 0) {
     return rv;
   }
diff --git a/lib/nghttp3_qpack.c b/lib/nghttp3_qpack.c
index 090a28b..fc5efc5 100644
--- a/lib/nghttp3_qpack.c
+++ b/lib/nghttp3_qpack.c
@@ -4185,6 +4185,7 @@ int nghttp3_qpack_decoder_new(nghttp3_qpack_decoder **pdecoder,
   rv = nghttp3_qpack_decoder_init(p, hard_max_dtable_capacity,
                                   max_blocked_streams, mem);
   if (rv != 0) {
+    nghttp3_mem_free(mem, p);
     return rv;
   }