From 0a7f68f382ff32d0c28d207a728c5c15611ce888 Mon Sep 17 00:00:00 2001
From: hai007 <hais@google.com>
Date: Thu, 7 Jul 2022 15:58:39 -0700
Subject: [PATCH] Cap the output stream package size to kMaxAllowedReadBytes,
 which is 1048576

PiperOrigin-RevId: 459622727
---
 connections/implementation/base_endpoint_channel.cc | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/connections/implementation/base_endpoint_channel.cc b/connections/implementation/base_endpoint_channel.cc
index c335e3e9..2328760e 100644
--- a/connections/implementation/base_endpoint_channel.cc
+++ b/connections/implementation/base_endpoint_channel.cc
@@ -223,8 +223,15 @@ Exception BaseEndpointChannel::Write(const ByteArray& data) {
       }
     }
 
+    size_t data_size = data_to_write->size();
+    if (data_size < 0 || data_size > kMaxAllowedReadBytes) {
+      NEARBY_LOGS(WARNING) << __func__ << ": Write an invalid number of bytes: "
+                           << data_size;
+      return {Exception::kIo};
+    }
+
     Exception write_exception =
-        WriteInt(writer_, static_cast<std::int32_t>(data_to_write->size()));
+        WriteInt(writer_, static_cast<std::int32_t>(data_size));
     if (write_exception.Raised()) {
       NEARBY_LOGS(WARNING) << __func__ << ": Failed to write header: "
                            << write_exception.value;