sandbox2: Buffer::Expand() implementation

Adds an API for resizing a shared Buffer after it has been created.

PiperOrigin-RevId: 763920344
Change-Id: Ia45fb1cd5e89e3300f781e922b4644b9b9d6bd5e

Author: Sandboxed API Team

sandboxed_api/sandbox2/BUILD

@@ -816,8 +816,12 @@ cc_test(
         ":buffer",
         ":sandbox2",
         "//sandboxed_api:testing",
+        "//sandboxed_api/util:file_base",
         "//sandboxed_api/util:fileops",
         "//sandboxed_api/util:status_matchers",
+        "@abseil-cpp//absl/status",
+        "@abseil-cpp//absl/status:statusor",
+        "@abseil-cpp//absl/strings:string_view",
         "@googletest//:gtest_main",
     ],
 )

sandboxed_api/sandbox2/buffer.cc

@@ -14,6 +14,7 @@
 
 #include "sandboxed_api/sandbox2/buffer.h"
 
+#include <fcntl.h>
 #include <sys/mman.h>
 #include <sys/stat.h>
 #include <unistd.h>
@@ -57,6 +58,39 @@ absl::StatusOr<std::unique_ptr<Buffer>> Buffer::CreateFromFd(FDCloser fd,
   return absl::WrapUnique(new Buffer(std::move(fd), buf, size));
 }
 
+absl::StatusOr<std::unique_ptr<Buffer>> Buffer::Expand(
+    std::unique_ptr<Buffer> other, size_t size) {
+  if (!other) {
+    return absl::InvalidArgumentError("Buffer is null");
+  }
+  if (other->buf_ == nullptr) {
+    return absl::FailedPreconditionError("Buffer is not initialized");
+  }
+  if (other->fd_.get() < 0) {
+    return absl::FailedPreconditionError("Buffer is not backed by a valid fd");
+  }
+  if (size < other->size_) {
+    return absl::InvalidArgumentError("Buffer size cannot be reduced");
+  }
+  if (other->size_ == size) {
+    return other;
+  }
+  if (fallocate(other->fd_.get(), 0, 0, size) != 0) {
+    return absl::ErrnoToStatus(errno, "Could not extend buffer fd");
+  }
+  uint8_t* new_buf_ = reinterpret_cast<uint8_t*>(
+      mremap(other->buf_, other->size_, size, MREMAP_MAYMOVE));
+  if (new_buf_ == MAP_FAILED) {
+    // At this point fallocate succeeded, and remapping failed.
+    // The incoming buffer is in an undefined state. It will be destroyed
+    // when this function returns.
+    return absl::ErrnoToStatus(errno, "Could not map buffer fd");
+  }
+  other->buf_ = new_buf_;
+  other->size_ = size;
+  return other;
+}
+
 // Creates a new Buffer of the specified size, backed by a temporary file that
 // will be immediately deleted.
 absl::StatusOr<std::unique_ptr<Buffer>> Buffer::CreateWithSize(size_t size) {

sandboxed_api/sandbox2/buffer.h

@@ -21,6 +21,7 @@
 #include <utility>
 
 #include "absl/base/macros.h"
+#include "absl/status/status.h"
 #include "absl/status/statusor.h"
 #include "sandboxed_api/util/fileops.h"
 
@@ -54,6 +55,12 @@ class Buffer final {
   // will be immediately deleted.
   static absl::StatusOr<std::unique_ptr<Buffer>> CreateWithSize(size_t size);
 
+  // Expands the input buffer to the specified size.
+  // Unlike CreateWithSize, this function will pre-allocate the memory.
+  // If size is smaller than the current mapped size, the function will fail.
+  static absl::StatusOr<std::unique_ptr<Buffer>> Expand(
+      std::unique_ptr<Buffer> other, size_t size);
+
   // Returns a pointer to the buffer, which is read/write.
   uint8_t* data() const { return buf_; }
 

sandboxed_api/sandbox2/buffer_test.cc

@@ -14,7 +14,9 @@
 
 #include "sandboxed_api/sandbox2/buffer.h"
 
+#include <fcntl.h>
 #include <sys/stat.h>
+#include <sys/types.h>
 #include <unistd.h>
 
 #include <cstdint>
@@ -25,13 +27,16 @@
 
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/string_view.h"
 #include "sandboxed_api/sandbox2/executor.h"
 #include "sandboxed_api/sandbox2/ipc.h"
 #include "sandboxed_api/sandbox2/policy.h"
 #include "sandboxed_api/sandbox2/result.h"
 #include "sandboxed_api/sandbox2/sandbox2.h"
 #include "sandboxed_api/testing.h"
 #include "sandboxed_api/util/fileops.h"
+#include "sandboxed_api/util/path.h"
 #include "sandboxed_api/util/status_matchers.h"
 
 namespace sandbox2 {
@@ -94,5 +99,45 @@ TEST(BufferTest, TestWithSandboxeeMapFd) {
   struct stat stat_buf;
   EXPECT_THAT(fstat(buffer->fd(), &stat_buf), Ne(-1));
 }
+
+TEST(BufferTest, TestResize) {
+  constexpr int kSize = 1024;
+  SAPI_ASSERT_OK_AND_ASSIGN(auto buffer, Buffer::CreateWithSize(kSize));
+  EXPECT_THAT(buffer->size(), Eq(kSize));
+  uint8_t* raw_buf = buffer->data();
+  for (int i = 0; i < kSize; i++) {
+    raw_buf[i] = 'X';
+  }
+  int fd = buffer->fd();
+  SAPI_ASSERT_OK_AND_ASSIGN(buffer,
+                            Buffer::Expand(std::move(buffer), kSize * 2));
+  EXPECT_THAT(buffer->size(), Eq(kSize * 2));
+  EXPECT_THAT(buffer->data(), Ne(nullptr));
+  EXPECT_THAT(buffer->fd(), Eq(fd));  // fd should not have changed.
+  absl::string_view buf_begin_view(reinterpret_cast<char*>(buffer->data()),
+                                   kSize);
+  EXPECT_THAT(buf_begin_view, Eq(std::string(kSize, 'X')));
+}
+
+TEST(BufferTest, TestResizeDoesNotTruncateFile) {
+  const std::string tmp_file =
+      sapi::file::JoinPath(testing::TempDir(), "1mb.file");
+  FDCloser fdcloser(open(tmp_file.c_str(), O_RDWR | O_CREAT | O_TRUNC, 0600));
+  int fd = fdcloser.get();
+  ASSERT_THAT(fd, Ne(-1));
+  ASSERT_THAT(ftruncate(fd, 1 << 20), Ne(-1));
+  // Map only 4k of the 1mb file.
+  SAPI_ASSERT_OK_AND_ASSIGN(absl::StatusOr<std::unique_ptr<Buffer>> buf,
+                            Buffer::CreateFromFd(std::move(fdcloser), 4096));
+  // Expand the mapped buffer to 8k.
+  SAPI_ASSERT_OK_AND_ASSIGN(auto expanded,
+                            Buffer::Expand(std::move(*buf), 8192));
+  EXPECT_THAT(expanded->size(), Eq(8192));
+  struct stat stat_buf;
+  // File size should not have changed.
+  ASSERT_THAT(stat(tmp_file.c_str(), &stat_buf), Ne(-1));
+  EXPECT_THAT(stat_buf.st_size, Eq(1 << 20));
+}
+
 }  // namespace
 }  // namespace sandbox2