1 files changed, 50 insertions, 17 deletions

diff --git a/libtransport/src/utils/membuf.cc b/libtransport/src/utils/membuf.cc
index 94e5b13a1..952116bb7 100644
--- a/libtransport/src/utils/membuf.cc
+++ b/libtransport/src/utils/membuf.cc
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2017-2019 Cisco and/or its affiliates.
+ * Copyright (c) 2021 Cisco and/or its affiliates.
  * Copyright 2013-present Facebook, Inc.
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -22,6 +22,7 @@
 #include <hicn/transport/portability/win_portability.h>
 #endif
 
+#include <glog/logging.h>
 #include <hicn/transport/utils/membuf.h>
 
 #include <cassert>
@@ -145,6 +146,18 @@ void MemBuf::operator delete(void* /* ptr */, void* /* placement */) {
   // constructor.
 }
 
+bool MemBuf::operator==(const MemBuf& other) {
+  if (length() != other.length()) {
+    return false;
+  }
+
+  return (memcmp(data(), other.data(), length()) == 0);
+}
+
+bool MemBuf::operator!=(const MemBuf& other) {
+  return !this->operator==(other);
+}
+
 void MemBuf::releaseStorage(HeapStorage* storage, uint16_t freeFlags) {
   // Use relaxed memory order here.  If we are unlucky and happen to get
   // out-of-date data the compare_exchange_weak() call below will catch
@@ -197,7 +210,7 @@ MemBuf::MemBuf(CopyBufferOp /* op */, const void* buf, std::size_t size,
     : MemBuf(CREATE, headroom + size + min_tailroom) {
   advance(headroom);
   if (size > 0) {
-    assert(buf != nullptr);
+    DCHECK(buf != nullptr);
     memcpy(writableData(), buf, size);
     append(size);
   }
@@ -299,21 +312,23 @@ unique_ptr<MemBuf> MemBuf::takeOwnership(void* buf, std::size_t capacity,
   }
 }
 
-MemBuf::MemBuf(WrapBufferOp, const void* buf, std::size_t capacity) noexcept
+MemBuf::MemBuf(WrapBufferOp, const void* buf, std::size_t length,
+               std::size_t capacity) noexcept
     : MemBuf(InternalConstructor(), 0,
              // We cast away the const-ness of the buffer here.
              // This is okay since MemBuf users must use unshare() to create a
              // copy of this buffer before writing to the buffer.
              static_cast<uint8_t*>(const_cast<void*>(buf)), capacity,
-             static_cast<uint8_t*>(const_cast<void*>(buf)), capacity) {}
+             static_cast<uint8_t*>(const_cast<void*>(buf)), length) {}
 
-unique_ptr<MemBuf> MemBuf::wrapBuffer(const void* buf, std::size_t capacity) {
-  return std::make_unique<MemBuf>(WRAP_BUFFER, buf, capacity);
+unique_ptr<MemBuf> MemBuf::wrapBuffer(const void* buf, std::size_t length,
+                                      std::size_t capacity) {
+  return std::make_unique<MemBuf>(WRAP_BUFFER, buf, length, capacity);
 }
 
-MemBuf MemBuf::wrapBufferAsValue(const void* buf,
+MemBuf MemBuf::wrapBufferAsValue(const void* buf, std::size_t length,
                                  std::size_t capacity) noexcept {
-  return MemBuf(WrapBufferOp::WRAP_BUFFER, buf, capacity);
+  return MemBuf(WrapBufferOp::WRAP_BUFFER, buf, length, capacity);
 }
 
 MemBuf::MemBuf() noexcept {}
@@ -355,8 +370,8 @@ MemBuf::MemBuf(InternalConstructor, uintptr_t flagsAndSharedInfo, uint8_t* buf,
       length_(length),
       capacity_(capacity),
       flags_and_shared_info_(flagsAndSharedInfo) {
-  assert(data >= buf);
-  assert(data + length <= buf + capacity);
+  DCHECK(data >= buf);
+  DCHECK(data + length <= buf + capacity);
 }
 
 MemBuf::~MemBuf() {
@@ -545,7 +560,7 @@ void MemBuf::unshareOneSlow() {
   // minimum capacity we also maintain at least the same amount of tailroom.
   std::size_t headlen = headroom();
   if (length_ > 0) {
-    assert(data_ != nullptr);
+    DCHECK(data_ != nullptr);
     memcpy(buf + headlen, data_, length_);
   }
 
@@ -562,7 +577,7 @@ void MemBuf::unshareOneSlow() {
 void MemBuf::unshareChained() {
   // unshareChained() should only be called if we are part of a chain of
   // multiple MemBufs.  The caller should have already verified this.
-  assert(isChained());
+  DCHECK(isChained());
 
   MemBuf* current = this;
   while (true) {
@@ -592,7 +607,7 @@ void MemBuf::markExternallyShared() {
 }
 
 void MemBuf::makeManagedChained() {
-  assert(isChained());
+  DCHECK(isChained());
 
   MemBuf* current = this;
   while (true) {
@@ -663,15 +678,15 @@ void MemBuf::coalesceAndReallocate(size_t new_headroom, size_t new_length,
   size_t remaining = new_length;
   do {
     if (current->length_ > 0) {
-      assert(current->length_ <= remaining);
-      assert(current->data_ != nullptr);
+      DCHECK(current->length_ <= remaining);
+      DCHECK(current->data_ != nullptr);
       remaining -= current->length_;
       memcpy(p, current->data_, current->length_);
       p += current->length_;
     }
     current = current->next_;
   } while (current != end);
-  assert(remaining == 0);
+  DCHECK(remaining == 0);
 
   // Point at the new buffer
   decrementRefcount();
@@ -785,7 +800,7 @@ void MemBuf::reserveSlow(std::size_t min_headroom, std::size_t min_tailroom) {
     new_allocated_capacity = goodExtBufferSize(new_capacity);
     new_buffer = static_cast<uint8_t*>(malloc(new_allocated_capacity));
     if (length_ > 0) {
-      assert(data_ != nullptr);
+      DCHECK(data_ != nullptr);
       memcpy(new_buffer + min_headroom, data_, length_);
     }
     if (sharedInfo()) {
@@ -862,4 +877,22 @@ void MemBuf::initExtBuffer(uint8_t* buf, size_t mallocSize,
   *infoReturn = sharedInfo;
 }
 
+bool MemBuf::ensureCapacity(std::size_t capacity) {
+  return !isChained() && std::size_t((bufferEnd() - data())) >= capacity;
+}
+
+bool MemBuf::ensureCapacityAndFillUnused(std::size_t capacity,
+                                         uint8_t placeholder) {
+  auto ret = ensureCapacity(capacity);
+  if (!ret) {
+    return ret;
+  }
+
+  if (length() < capacity) {
+    std::memset(writableTail(), placeholder, capacity - length());
+  }
+
+  return ret;
+}
+
 }  // namespace utils
\ No newline at end of file