Fix heap OOB read in Literal::CreateFromProto via unchecked dynamic_sizes#44232
Open
foodlook wants to merge 1 commit into
Open
Fix heap OOB read in Literal::CreateFromProto via unchecked dynamic_sizes#44232foodlook wants to merge 1 commit into
foodlook wants to merge 1 commit into
Conversation
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
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
1 participant
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.
Summary
Piece::CopyFromProtoacceptsdynamic_sizesvalues from a serializedLiteralProtoand passes them directly toPiece::SetDynamicSizewithout validating that they fall within[0, shape.dimensions(i)]. Unlike the publicMutableLiteralBase::SetDynamicSize(which hasCHECK_GE), the internalPiece::SetDynamicSizestores the value unconditionally, breaking the invariantdynamic_size <= static_dimension_bound.A subsequent call to
ToStatic()reads the poisoned dynamic size and uses it as the copy count inCopyElementsWithDynamicBound/std::copy_n, causing an out-of-bounds heap read.Reproduction
AddressSanitizer confirms:
Fix
Add bounds validation in
CopyFromProtobefore callingSetDynamicSize: for each dynamic dimension, rejectdynamic_sizes[i]outside[0, shape.dimensions(i)]withInvalidArgument.Tests added
InvalidProtoDynamicSizeExceedsBound:dynamic_sizes=[100]onf32[<=4]-> rejectedInvalidProtoDynamicSizeNegative:dynamic_sizes=[-1]-> rejectedValidProtoDynamicSizeWithinBound:dynamic_sizes=[2]onf32[<=4]-> accepted,ToStatic()producesf32[2]