Introduce a sandboxee_read_index field in AsynchronousByteTransport

The point is that the sandboxee might `execve` at some point, and lose track of where its read_index previously was. The ensures that the sandboxee will keep reading at the correct index after the `exec` system call.

PiperOrigin-RevId: 910687145
Change-Id: I7c729afc171be14f4a0932c12ae44625a77517bd

Author: paulsemel

sandboxed_api/sandbox2/util/asynchronous_byte_transport.cc

@@ -208,6 +208,7 @@ AsynchronousByteTransport::AsynchronousByteTransport(
     GetHeader()->h2s.state.store(0, std::memory_order_release);
     GetHeader()->s2h.state.store(0, std::memory_order_release);
     GetHeader()->synchronization_type = synchronization_type;
+    GetHeader()->sandboxee_read_index = 0;
   }
 }
 
@@ -278,12 +279,13 @@ absl::Status AsynchronousByteTransport::Read(absl::Span<uint8_t> data) {
       return absl::AbortedError("Write index out of bounds");
     }
 
-    if (read_index_ > write_index) {
+    uint32_t& read_index = GetReadIndex();
+    if (read_index > write_index) {
       return absl::AbortedError("Read index out of bounds");
     }
 
     absl::Span<const uint8_t> read_buffer =
-        read_data_.subspan(read_index_, write_index);
+        read_data_.subspan(read_index, write_index);
     if (read_buffer.empty()) {
       uint32_t termination_state =
           termination_state_.load(std::memory_order_relaxed);
@@ -296,18 +298,18 @@ absl::Status AsynchronousByteTransport::Read(absl::Span<uint8_t> data) {
     }
 
     size_t chunk_size =
-        std::min(kChunkSize, static_cast<size_t>(write_index - read_index_));
+        std::min(kChunkSize, static_cast<size_t>(write_index - read_index));
     absl::Span<uint8_t> chunk = data.subspan(0, chunk_size);
 
     memcpy(chunk.data(), read_buffer.data(), chunk.size());
-    read_index_ += chunk.size();
+    read_index += chunk.size();
     data.remove_prefix(chunk.size());
 
-    if (read_index_ == write_index) {
+    if (read_index == write_index) {
       // If the connection is closed, we will not read any more data at this
       // point, but we will still acknowledge the last bit of data that was sent
       // from the other side. Note that in the event of an attacker controlled
-      // client, we will read from read_index_ up to the size of the buffer
+      // client, we will read from read_index up to the size of the buffer
       // until we acknowledge the connection is closed. This is needed so that
       // we can support legitimate scenarios where the client sends size_of_buf
       // - 1 bytes and then closes the connection. In this case, the host still
@@ -348,7 +350,8 @@ bool AsynchronousByteTransport::ResetWriteIndex(ChannelHeader* channel,
   if (state & kWaitForWritingBit) {
     WakeInternal(channel);
   }
-  read_index_ = 0;
+  uint32_t& read_index = GetReadIndex();
+  read_index = 0;
   return true;
 }
 
@@ -430,7 +433,7 @@ void AsynchronousByteTransport::WakeInternal(ChannelHeader* channel,
 
 int AsynchronousByteTransport::TransferInternal(ChannelHeader* write_channel,
                                                 ChannelHeader* read_channel) {
-  uint32_t read_state = read_index_;
+  uint32_t read_state = GetReadIndex();
   // If the counterpart has already written data and is waiting for reading,
   // we just need to wake them up and get back to reading.
   if (!read_channel->state.compare_exchange_strong(

sandboxed_api/sandbox2/util/asynchronous_byte_transport.h

@@ -126,6 +126,7 @@ class AsynchronousByteTransport {
 
   struct Header {
     SynchronizationType synchronization_type;
+    uint32_t sandboxee_read_index;
     ChannelHeader h2s;  // Host to Sandboxee
     ChannelHeader s2h;  // Sandboxee to Host
     // Must be last.
@@ -164,6 +165,12 @@ class AsynchronousByteTransport {
   int TransferInternal(ChannelHeader* write_channel,
                        ChannelHeader* read_channel);
 
+  uint32_t& GetReadIndex() {
+    return client_type_ == ClientType::kHost
+               ? read_index_
+               : GetHeader()->sandboxee_read_index;
+  }
+
   std::unique_ptr<sandbox2::Buffer> buffer_;
   ClientType client_type_;
   uint32_t read_index_ = 0;

sandboxed_api/sandbox2/util/asynchronous_byte_transport_test.cc

@@ -21,6 +21,7 @@
 
 #include <cstddef>
 #include <cstdint>
+#include <cstring>
 #include <memory>
 #include <utility>
 #include <vector>
@@ -30,6 +31,7 @@
 #include "absl/log/check.h"
 #include "absl/status/status.h"
 #include "absl/synchronization/notification.h"
+#include "absl/time/clock.h"
 #include "absl/time/time.h"
 #include "absl/types/span.h"
 #include "sandboxed_api/sandbox2/buffer.h"
@@ -139,25 +141,25 @@ class TestHelper {
     while (comms.RecvInt32(reinterpret_cast<int32_t*>(&action_type))) {
       if (action_type == ActionType::kSend) {
         std::vector<uint8_t> data;
-        ASSERT_TRUE(comms.RecvBytes(&data));
-        SAPI_ASSERT_OK(transport_client->Send(data));
+        CHECK_EQ(comms.RecvBytes(&data), true);
+        CHECK_OK(transport_client->Send(data));
       } else if (action_type == ActionType::kExchange) {
         std::vector<uint8_t> data_to_send;
-        ASSERT_TRUE(comms.RecvBytes(&data_to_send));
+        CHECK_EQ(comms.RecvBytes(&data_to_send), true);
         std::vector<uint8_t> data_to_recv;
-        ASSERT_TRUE(comms.RecvBytes(&data_to_recv));
+        CHECK_EQ(comms.RecvBytes(&data_to_recv), true);
         std::vector<uint8_t> recv_data(data_to_recv.size());
-        SAPI_ASSERT_OK(transport_client->Exchange(
+        CHECK_OK(transport_client->Exchange(
             data_to_send,
             absl::Span<uint8_t>(recv_data.data(), recv_data.size())));
         ASSERT_EQ(recv_data, data_to_recv);
       } else if (action_type == ActionType::kRecv) {
         std::vector<uint8_t> data;
-        ASSERT_TRUE(comms.RecvBytes(&data));
+        CHECK_EQ(comms.RecvBytes(&data), true);
         std::vector<uint8_t> data_recv(data.size());
-        SAPI_ASSERT_OK(transport_client->Recv(
+        CHECK_OK(transport_client->Recv(
             absl::Span<uint8_t>(data_recv.data(), data_recv.size())));
-        ASSERT_EQ(data, data_recv);
+        CHECK_EQ(data, data_recv);
       } else if (action_type == ActionType::kTerminate) {
         transport_client->Terminate();
       }
@@ -185,6 +187,7 @@ class AsynchronousByteTransportTest
   void SetUp() override {
     SAPI_ASSERT_OK_AND_ASSIGN(auto buffer,
                               sandbox2::Buffer::CreateWithSize(buffer_size_));
+    memset(buffer->data(), 0x41, 4 << 10);
     int memfd = buffer->fd();
     size_t size = buffer->size();
     sandbox2::AsynchronousByteTransport::SynchronizationType
@@ -212,7 +215,7 @@ class AsynchronousByteTransportTest
   TestHelper test_helper_;
 
  private:
-  size_t buffer_size_ = 128 << 10;
+  size_t buffer_size_ = 132 << 10;
   std::unique_ptr<ScopedTimeout> timeout_;
   std::unique_ptr<sandbox2::AsynchronousByteTransport> transport_;
 };
@@ -402,6 +405,17 @@ TEST_P(AsynchronousByteTransportTest, TwoReadsAfterConnectionClosedFails) {
               StatusIs(absl::StatusCode::kAborted));
 }
 
+TEST_P(AsynchronousByteTransportTest,
+       ReadMoreThanBufferWillNotTriggerOutOfBounds) {
+  std::vector<uint8_t> data(10, 'a');
+  SAPI_ASSERT_OK(GetTransport()->Send(
+      absl::Span<const uint8_t>(data.data(), data.size() / 2)));
+  test_helper_.RequestRecv(data);
+  absl::SleepFor(absl::Milliseconds(10));
+  SAPI_ASSERT_OK(GetTransport()->Send(absl::Span<const uint8_t>(
+      data.data() + data.size() / 2, data.size() / 2)));
+}
+
 INSTANTIATE_TEST_SUITE_P(
     AsynchronousByteTransportTest, AsynchronousByteTransportTest,
     ::testing::Values(