From 487d5a6de314923e8e3814b7f4e6b2e777b645b6 Mon Sep 17 00:00:00 2001
From: foodlook <16376469+foodlook@users.noreply.github.com>
Date: Sat, 13 Jun 2026 02:38:40 +0900
Subject: [PATCH] Validate dynamic_sizes bounds in LiteralProto deserialization

Piece::CopyFromProto accepted arbitrary dynamic_sizes values from
untrusted LiteralProto without validating them against the static
dimension bounds. The internal Piece::SetDynamicSize (unlike its
public counterpart MutableLiteralBase::SetDynamicSize) lacked a
CHECK_GE guard, silently breaking the invariant
  dynamic_size <= static_dimension_bound.

A subsequent call to ToStatic() would then invoke
CopyElementsWithDynamicBound, which used the poisoned dynamic size
as the copy count in std::copy_n, resulting in an out-of-bounds
heap read (confirmed via AddressSanitizer: READ of size 400 from a
16-byte buffer).

The fix adds bounds validation in CopyFromProto before calling
SetDynamicSize: for each dynamic dimension, dynamic_sizes[i] must
be in [0, shape.dimensions(i)]. Out-of-range values now return
InvalidArgument instead of silently corrupting the invariant.

Regression tests:
- InvalidProtoDynamicSizeExceedsBound: dynamic_sizes=[100] on f32[<=4] rejected
- InvalidProtoDynamicSizeNegative: dynamic_sizes=[-1] rejected
- ValidProtoDynamicSizeWithinBound: dynamic_sizes=[2] on f32[<=4] accepted, ToStatic() produces f32[2]

PiperOrigin-RevId: N/A
---
 xla/literal.cc      |  9 ++++++
 xla/literal_test.cc | 74 +++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 83 insertions(+)

diff --git a/xla/literal.cc b/xla/literal.cc
index 9638d277baec3..984c09a22bfc7 100644
--- a/xla/literal.cc
+++ b/xla/literal.cc
@@ -2558,6 +2558,15 @@ absl::Status LiteralBase::Piece::CopyFromProto(const LiteralProto& proto) {
   if (shape.is_dynamic()) {
     TF_RET_CHECK(proto.dynamic_sizes_size() == shape.dimensions().size());
     for (int64_t i = 0; i < shape.dimensions().size(); ++i) {
+      if (shape.is_dynamic_dimension(i)) {
+        const int32_t dynamic_size = proto.dynamic_sizes(i);
+        if (dynamic_size < 0 || dynamic_size > shape.dimensions(i)) {
+          return InvalidArgument(
+              "LiteralProto dynamic_sizes[%d] = %d is out of range [0, %d] "
+              "for dimension %d with static bound %d",
+              i, dynamic_size, shape.dimensions(i), i, shape.dimensions(i));
+        }
+      }
       SetDynamicSize(i, proto.dynamic_sizes(i));
     }
   }
diff --git a/xla/literal_test.cc b/xla/literal_test.cc
index b106a57250bee..e0480d584fb09 100644
--- a/xla/literal_test.cc
+++ b/xla/literal_test.cc
@@ -2728,6 +2728,80 @@ TEST_F(LiteralUtilTest, InvalidProtoMissingLayout) {
   EXPECT_THAT(status.message(), HasSubstr("LiteralProto has no layout"));
 }
 
+TEST_F(LiteralUtilTest, InvalidProtoDynamicSizeExceedsBound) {
+  // Regression test: dynamic_sizes that exceed the static dimension bound must
+  // be rejected by CreateFromProto.  Before the fix, the poisoned value was
+  // silently stored via Piece::SetDynamicSize (which lacked the CHECK_GE
+  // present in MutableLiteralBase::SetDynamicSize), leading to a heap OOB read
+  // in CopyElementsWithDynamicBound when ToStatic() was called.
+  LiteralProto proto;
+
+  // Shape: f32[<=4] — rank 1, 1 dynamic dim, static bound = 4.
+  ShapeProto* shape = proto.mutable_shape();
+  shape->set_element_type(F32);
+  shape->add_dimensions(4);
+  shape->add_is_dynamic_dimension(true);
+  LayoutProto* layout = shape->mutable_layout();
+  layout->add_minor_to_major(0);
+
+  // Provide exactly 4 float values (matches static bound).
+  proto.add_f32s(1.0f);
+  proto.add_f32s(2.0f);
+  proto.add_f32s(3.0f);
+  proto.add_f32s(4.0f);
+
+  // dynamic_sizes[0] = 100 — exceeds the static bound of 4.
+  proto.add_dynamic_sizes(100);
+
+  absl::Status status = Literal::CreateFromProto(proto).status();
+  ASSERT_FALSE(status.ok());
+  EXPECT_THAT(status.message(), HasSubstr("dynamic_sizes"));
+  EXPECT_THAT(status.message(), HasSubstr("out of range"));
+}
+
+TEST_F(LiteralUtilTest, InvalidProtoDynamicSizeNegative) {
+  // Negative dynamic_sizes must also be rejected.
+  LiteralProto proto;
+  ShapeProto* shape = proto.mutable_shape();
+  shape->set_element_type(F32);
+  shape->add_dimensions(4);
+  shape->add_is_dynamic_dimension(true);
+  LayoutProto* layout = shape->mutable_layout();
+  layout->add_minor_to_major(0);
+  proto.add_f32s(1.0f);
+  proto.add_f32s(2.0f);
+  proto.add_f32s(3.0f);
+  proto.add_f32s(4.0f);
+  proto.add_dynamic_sizes(-1);
+
+  absl::Status status = Literal::CreateFromProto(proto).status();
+  ASSERT_FALSE(status.ok());
+  EXPECT_THAT(status.message(), HasSubstr("dynamic_sizes"));
+  EXPECT_THAT(status.message(), HasSubstr("out of range"));
+}
+
+TEST_F(LiteralUtilTest, ValidProtoDynamicSizeWithinBound) {
+  // A dynamic_size within the static bound must still be accepted.
+  LiteralProto proto;
+  ShapeProto* shape = proto.mutable_shape();
+  shape->set_element_type(F32);
+  shape->add_dimensions(4);
+  shape->add_is_dynamic_dimension(true);
+  LayoutProto* layout = shape->mutable_layout();
+  layout->add_minor_to_major(0);
+  proto.add_f32s(1.0f);
+  proto.add_f32s(2.0f);
+  proto.add_f32s(3.0f);
+  proto.add_f32s(4.0f);
+  proto.add_dynamic_sizes(2);
+
+  TF_ASSERT_OK_AND_ASSIGN(Literal literal, Literal::CreateFromProto(proto));
+  EXPECT_EQ(literal.GetDynamicSize(0), 2);
+  auto static_literal = literal.ToStatic();
+  EXPECT_EQ(static_literal.shape(), ShapeUtil::MakeShape(F32, {2}));
+  EXPECT_TRUE(static_literal.shape().is_static());
+}
+
 TEST_F(LiteralUtilTest, InvalidProtoTooFewTupleElements) {
   // Proto has the too few tuple elements.
   LiteralProto proto;