From c2715b89b4ef792de0fb94414c0ad69df18a3487 Mon Sep 17 00:00:00 2001
From: Darshan Sen <darshan.sen@postman.com>
Date: Wed, 22 Sep 2021 21:10:57 +0530
Subject: [PATCH] src: add functions to manage allocated buffers to the
 Environment class

Signed-off-by: Darshan Sen <darshan.sen@postman.com>
---
 src/env-inl.h   | 25 +++++++++++++++++++++++++
 src/env.h       |  3 +++
 src/udp_wrap.cc | 28 +++++-----------------------
 3 files changed, 33 insertions(+), 23 deletions(-)

diff --git a/src/env-inl.h b/src/env-inl.h
index 156534bdeee8a1..7e2b25644739b6 100644
--- a/src/env-inl.h
+++ b/src/env-inl.h
@@ -25,6 +25,7 @@
 #if defined(NODE_WANT_INTERNALS) && NODE_WANT_INTERNALS
 
 #include "aliased_buffer.h"
+#include "allocated_buffer-inl.h"
 #include "callback_queue-inl.h"
 #include "env.h"
 #include "node.h"
@@ -966,6 +967,30 @@ inline IsolateData* Environment::isolate_data() const {
   return isolate_data_;
 }
 
+inline uv_buf_t Environment::allocate_managed_buffer(
+    const size_t suggested_size) {
+  const NoArrayBufferZeroFillScope no_zero_fill_scope(isolate_data());
+  std::unique_ptr<v8::BackingStore> bs =
+      v8::ArrayBuffer::NewBackingStore(isolate(), suggested_size);
+  const uv_buf_t buf =
+      uv_buf_init(static_cast<char*>(bs->Data()), bs->ByteLength());
+  released_allocated_buffers()->emplace(buf.base, std::move(bs));
+  return buf;
+}
+
+inline std::unique_ptr<v8::BackingStore> Environment::release_managed_buffer(
+    const uv_buf_t& buf) {
+  std::unique_ptr<v8::BackingStore> bs;
+  if (buf.base != nullptr) {
+    auto map = released_allocated_buffers();
+    auto it = map->find(buf.base);
+    CHECK_NE(it, map->end());
+    bs = std::move(it->second);
+    map->erase(it);
+  }
+  return bs;
+}
+
 std::unordered_map<char*, std::unique_ptr<v8::BackingStore>>*
     Environment::released_allocated_buffers() {
   return &released_allocated_buffers_;
diff --git a/src/env.h b/src/env.h
index f7d1dace414963..00874dd099c8ec 100644
--- a/src/env.h
+++ b/src/env.h
@@ -1437,6 +1437,9 @@ class Environment : public MemoryRetainer {
   void RunAndClearNativeImmediates(bool only_refed = false);
   void RunAndClearInterrupts();
 
+  inline uv_buf_t allocate_managed_buffer(const size_t suggested_size);
+  inline std::unique_ptr<v8::BackingStore> release_managed_buffer(
+      const uv_buf_t& buf);
   inline std::unordered_map<char*, std::unique_ptr<v8::BackingStore>>*
       released_allocated_buffers();
 
diff --git a/src/udp_wrap.cc b/src/udp_wrap.cc
index bb3464543786f2..2b96df23daca07 100644
--- a/src/udp_wrap.cc
+++ b/src/udp_wrap.cc
@@ -683,16 +683,7 @@ void UDPWrap::OnAlloc(uv_handle_t* handle,
 }
 
 uv_buf_t UDPWrap::OnAlloc(size_t suggested_size) {
-  Environment* env = this->env();
-  uv_buf_t buf;
-  {
-    NoArrayBufferZeroFillScope no_zero_fill_scope(env->isolate_data());
-    std::unique_ptr<BackingStore> bs =
-        ArrayBuffer::NewBackingStore(env->isolate(), suggested_size);
-    buf = uv_buf_init(static_cast<char*>(bs->Data()), bs->ByteLength());
-    env->released_allocated_buffers()->emplace(buf.base, std::move(bs));
-  }
-  return buf;
+  return this->env()->allocate_managed_buffer(suggested_size);
 }
 
 void UDPWrap::OnRecv(uv_udp_t* handle,
@@ -710,22 +701,13 @@ void UDPWrap::OnRecv(ssize_t nread,
                      unsigned int flags) {
   Environment* env = this->env();
   Isolate* isolate = env->isolate();
-
-  std::unique_ptr<BackingStore> bs;
-  if (buf_.base != nullptr) {
-    auto map = env->released_allocated_buffers();
-    auto it = map->find(buf_.base);
-    CHECK_NE(it, map->end());
-    bs = std::move(it->second);
-    map->erase(it);
-  }
-
+  std::unique_ptr<BackingStore> bs = env->release_managed_buffer(buf_);
   if (nread == 0 && addr == nullptr) {
     return;
   }
 
-  HandleScope handle_scope(isolate);
-  Context::Scope context_scope(env->context());
+  const HandleScope handle_scope(isolate);
+  const Context::Scope context_scope(env->context());
 
   Local<Value> argv[] = {
       Integer::New(isolate, static_cast<int32_t>(nread)),
@@ -739,7 +721,7 @@ void UDPWrap::OnRecv(ssize_t nread,
   } else if (nread == 0) {
     bs = ArrayBuffer::NewBackingStore(isolate, 0);
   } else {
-    NoArrayBufferZeroFillScope no_zero_fill_scope(env->isolate_data());
+    const NoArrayBufferZeroFillScope no_zero_fill_scope(env->isolate_data());
     bs = BackingStore::Reallocate(isolate, std::move(bs), nread);
   }