fix(ifrt): reject dimension products that overflow int64_t in Shape::FromProto - #46403
Closed
prasanna8585 wants to merge 2 commits into
Closed
fix(ifrt): reject dimension products that overflow int64_t in Shape::FromProto#46403prasanna8585 wants to merge 2 commits into
prasanna8585 wants to merge 2 commits into
Conversation
…FromProto xla::ifrt::Shape (a class separate from xla::Shape/ShapeUtil, used by the Python/IFRT layer including ifrt_proxy) deserializes a ShapeProto and only validated that individual dimensions were non-negative -- the product of all dimensions was never checked, unlike xla::ShapeUtil::ValidateDimensions which already does this correctly for the core Shape type via OverflowSafeMultiply. Shape::num_elements() computes this product without overflow detection, so a crafted ShapeProto with individually plausible but jointly overflowing dimensions silently produces a wrapped, incorrect element count. Anything downstream that sizes a buffer from that count (e.g. byte-size computations derived from num_elements()) would be sized inconsistently with the shape's true declared dimensions. This is reachable from ifrt_proxy's deserialized RPC request protos (ArraySpec::FromProto -> Shape::FromProto), which is a genuine client-server trust boundary distinct from IFRT's general in-process API surface. Fix mirrors the existing, correct pattern in xla::ShapeUtil::ValidateDimensions: detect the overflow with __builtin_mul_overflow while walking the dimensions and reject with InvalidArgument, rather than silently returning a corrupted Shape. Two regression tests added: one confirming the overflow is rejected, one confirming ordinary non-overflowing shapes are unaffected.
…d overflow check Per golechwierowicz's review comment on this PR: replace the manual __builtin_mul_overflow-based loop with xla::OverflowSafeMultiply (xla/overflow_util.h), the same utility xla::ShapeUtil::ValidateDimensions already uses for the core XLA Shape type's equivalent check. No behavioral change -- same detection, same rejection -- just reusing the existing, already-reviewed utility instead of duplicating its logic.
golechwierowicz
approved these changes
Jul 30, 2026
copybara-service Bot
pushed a commit
that referenced
this pull request
Jul 30, 2026
… in Shape::FromProto Imported from GitHub PR #46403 `xla::ifrt::Shape::FromProto` only validated that individual dimensions were non-negative; the product of all dimensions was never checked. `xla::ShapeUtil::ValidateDimensions` already does this correctly for the core `xla::Shape` type via `OverflowSafeMultiply` -- this class is a separate type (used by the Python/IFRT layer, including `ifrt_proxy`) that didn't have the equivalent protection. `Shape::num_elements()` computes the dimension product without overflow detection, so a crafted `ShapeProto` with individually plausible but jointly overflowing dimensions silently produces a wrapped, incorrect element count rather than being rejected. This is reachable from `ifrt_proxy`'s deserialized RPC request protos (`ArraySpec::FromProto` -> `Shape::FromProto`), a client-server boundary distinct from IFRT's general in-process API. Fix mirrors the existing pattern in `ValidateDimensions`: detect overflow with `__builtin_mul_overflow` while walking dimensions, reject with `InvalidArgument` instead of returning a corrupted `Shape`. Two regression tests included. Copybara import of the project: -- 9ea3faf by Prasanna Dabi <65734642+prasanna8585@users.noreply.github.com>: fix(ifrt): reject dimension products that overflow int64_t in Shape::FromProto xla::ifrt::Shape (a class separate from xla::Shape/ShapeUtil, used by the Python/IFRT layer including ifrt_proxy) deserializes a ShapeProto and only validated that individual dimensions were non-negative -- the product of all dimensions was never checked, unlike xla::ShapeUtil::ValidateDimensions which already does this correctly for the core Shape type via OverflowSafeMultiply. Shape::num_elements() computes this product without overflow detection, so a crafted ShapeProto with individually plausible but jointly overflowing dimensions silently produces a wrapped, incorrect element count. Anything downstream that sizes a buffer from that count (e.g. byte-size computations derived from num_elements()) would be sized inconsistently with the shape's true declared dimensions. This is reachable from ifrt_proxy's deserialized RPC request protos (ArraySpec::FromProto -> Shape::FromProto), which is a genuine client-server trust boundary distinct from IFRT's general in-process API surface. Fix mirrors the existing, correct pattern in xla::ShapeUtil::ValidateDimensions: detect the overflow with __builtin_mul_overflow while walking the dimensions and reject with InvalidArgument, rather than silently returning a corrupted Shape. Two regression tests added: one confirming the overflow is rejected, one confirming ordinary non-overflowing shapes are unaffected. -- 4a1d7d9 by Prasanna Dabi <65734642+prasanna8585@users.noreply.github.com>: Address review: reuse xla::OverflowSafeMultiply instead of hand-rolled overflow check Per golechwierowicz's review comment on this PR: replace the manual __builtin_mul_overflow-based loop with xla::OverflowSafeMultiply (xla/overflow_util.h), the same utility xla::ShapeUtil::ValidateDimensions already uses for the core XLA Shape type's equivalent check. No behavioral change -- same detection, same rejection -- just reusing the existing, already-reviewed utility instead of duplicating its logic. Merging this change closes #46403 FUTURE_COPYBARA_INTEGRATE_REVIEW=#46403 from prasanna8585:fix/ifrt-shape-fromproto-overflow-check 4a1d7d9 PiperOrigin-RevId: 956346424
copybara-service Bot
pushed a commit
to tensorflow/tensorflow
that referenced
this pull request
Jul 30, 2026
… in Shape::FromProto Imported from GitHub PR openxla/xla#46403 `xla::ifrt::Shape::FromProto` only validated that individual dimensions were non-negative; the product of all dimensions was never checked. `xla::ShapeUtil::ValidateDimensions` already does this correctly for the core `xla::Shape` type via `OverflowSafeMultiply` -- this class is a separate type (used by the Python/IFRT layer, including `ifrt_proxy`) that didn't have the equivalent protection. `Shape::num_elements()` computes the dimension product without overflow detection, so a crafted `ShapeProto` with individually plausible but jointly overflowing dimensions silently produces a wrapped, incorrect element count rather than being rejected. This is reachable from `ifrt_proxy`'s deserialized RPC request protos (`ArraySpec::FromProto` -> `Shape::FromProto`), a client-server boundary distinct from IFRT's general in-process API. Fix mirrors the existing pattern in `ValidateDimensions`: detect overflow with `__builtin_mul_overflow` while walking dimensions, reject with `InvalidArgument` instead of returning a corrupted `Shape`. Two regression tests included. Copybara import of the project: -- 9ea3faf61a3396a9da49abd09199d9cb296adebf by Prasanna Dabi <65734642+prasanna8585@users.noreply.github.com>: fix(ifrt): reject dimension products that overflow int64_t in Shape::FromProto xla::ifrt::Shape (a class separate from xla::Shape/ShapeUtil, used by the Python/IFRT layer including ifrt_proxy) deserializes a ShapeProto and only validated that individual dimensions were non-negative -- the product of all dimensions was never checked, unlike xla::ShapeUtil::ValidateDimensions which already does this correctly for the core Shape type via OverflowSafeMultiply. Shape::num_elements() computes this product without overflow detection, so a crafted ShapeProto with individually plausible but jointly overflowing dimensions silently produces a wrapped, incorrect element count. Anything downstream that sizes a buffer from that count (e.g. byte-size computations derived from num_elements()) would be sized inconsistently with the shape's true declared dimensions. This is reachable from ifrt_proxy's deserialized RPC request protos (ArraySpec::FromProto -> Shape::FromProto), which is a genuine client-server trust boundary distinct from IFRT's general in-process API surface. Fix mirrors the existing, correct pattern in xla::ShapeUtil::ValidateDimensions: detect the overflow with __builtin_mul_overflow while walking the dimensions and reject with InvalidArgument, rather than silently returning a corrupted Shape. Two regression tests added: one confirming the overflow is rejected, one confirming ordinary non-overflowing shapes are unaffected. -- 4a1d7d96afe8061ea8ba8d8331654854d2c16f91 by Prasanna Dabi <65734642+prasanna8585@users.noreply.github.com>: Address review: reuse xla::OverflowSafeMultiply instead of hand-rolled overflow check Per golechwierowicz's review comment on this PR: replace the manual __builtin_mul_overflow-based loop with xla::OverflowSafeMultiply (xla/overflow_util.h), the same utility xla::ShapeUtil::ValidateDimensions already uses for the core XLA Shape type's equivalent check. No behavioral change -- same detection, same rejection -- just reusing the existing, already-reviewed utility instead of duplicating its logic. Merging this change closes #46403 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#46403 from prasanna8585:fix/ifrt-shape-fromproto-overflow-check 4a1d7d96afe8061ea8ba8d8331654854d2c16f91 PiperOrigin-RevId: 956346424
copybara-service Bot
pushed a commit
to tensorflow/tensorflow
that referenced
this pull request
Jul 30, 2026
… in Shape::FromProto Imported from GitHub PR openxla/xla#46403 `xla::ifrt::Shape::FromProto` only validated that individual dimensions were non-negative; the product of all dimensions was never checked. `xla::ShapeUtil::ValidateDimensions` already does this correctly for the core `xla::Shape` type via `OverflowSafeMultiply` -- this class is a separate type (used by the Python/IFRT layer, including `ifrt_proxy`) that didn't have the equivalent protection. `Shape::num_elements()` computes the dimension product without overflow detection, so a crafted `ShapeProto` with individually plausible but jointly overflowing dimensions silently produces a wrapped, incorrect element count rather than being rejected. This is reachable from `ifrt_proxy`'s deserialized RPC request protos (`ArraySpec::FromProto` -> `Shape::FromProto`), a client-server boundary distinct from IFRT's general in-process API. Fix mirrors the existing pattern in `ValidateDimensions`: detect overflow with `__builtin_mul_overflow` while walking dimensions, reject with `InvalidArgument` instead of returning a corrupted `Shape`. Two regression tests included. Copybara import of the project: -- 9ea3faf61a3396a9da49abd09199d9cb296adebf by Prasanna Dabi <65734642+prasanna8585@users.noreply.github.com>: fix(ifrt): reject dimension products that overflow int64_t in Shape::FromProto xla::ifrt::Shape (a class separate from xla::Shape/ShapeUtil, used by the Python/IFRT layer including ifrt_proxy) deserializes a ShapeProto and only validated that individual dimensions were non-negative -- the product of all dimensions was never checked, unlike xla::ShapeUtil::ValidateDimensions which already does this correctly for the core Shape type via OverflowSafeMultiply. Shape::num_elements() computes this product without overflow detection, so a crafted ShapeProto with individually plausible but jointly overflowing dimensions silently produces a wrapped, incorrect element count. Anything downstream that sizes a buffer from that count (e.g. byte-size computations derived from num_elements()) would be sized inconsistently with the shape's true declared dimensions. This is reachable from ifrt_proxy's deserialized RPC request protos (ArraySpec::FromProto -> Shape::FromProto), which is a genuine client-server trust boundary distinct from IFRT's general in-process API surface. Fix mirrors the existing, correct pattern in xla::ShapeUtil::ValidateDimensions: detect the overflow with __builtin_mul_overflow while walking the dimensions and reject with InvalidArgument, rather than silently returning a corrupted Shape. Two regression tests added: one confirming the overflow is rejected, one confirming ordinary non-overflowing shapes are unaffected. -- 4a1d7d96afe8061ea8ba8d8331654854d2c16f91 by Prasanna Dabi <65734642+prasanna8585@users.noreply.github.com>: Address review: reuse xla::OverflowSafeMultiply instead of hand-rolled overflow check Per golechwierowicz's review comment on this PR: replace the manual __builtin_mul_overflow-based loop with xla::OverflowSafeMultiply (xla/overflow_util.h), the same utility xla::ShapeUtil::ValidateDimensions already uses for the core XLA Shape type's equivalent check. No behavioral change -- same detection, same rejection -- just reusing the existing, already-reviewed utility instead of duplicating its logic. Merging this change closes #46403 PiperOrigin-RevId: 956425279
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
xla::ifrt::Shape::FromProtoonly validated that individual dimensions were non-negative; the product of all dimensions was never checked.xla::ShapeUtil::ValidateDimensionsalready does this correctly for the corexla::Shapetype viaOverflowSafeMultiply-- this class is a separate type (used by the Python/IFRT layer, includingifrt_proxy) that didn't have the equivalent protection.Shape::num_elements()computes the dimension product without overflow detection, so a craftedShapeProtowith individually plausible but jointly overflowing dimensions silently produces a wrapped, incorrect element count rather than being rejected.This is reachable from
ifrt_proxy's deserialized RPC request protos (ArraySpec::FromProto->Shape::FromProto), a client-server boundary distinct from IFRT's general in-process API.Fix mirrors the existing pattern in
ValidateDimensions: detect overflow with__builtin_mul_overflowwhile walking dimensions, reject withInvalidArgumentinstead of returning a corruptedShape. Two regression tests included.