Skip to content

fix: make buffer growth overflow-safe #3944

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
peter-jerry-ye merged 3 commits into from
Jul 28, 2026
Merged

fix: make buffer growth overflow-safe #3944

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
6f29101
fix: make buffer growth overflow-safe
peter-jerry-ye Jul 28, 2026
ca6ab3d
perf(builtin): inline StringBuilder growth sizing
peter-jerry-ye Jul 28, 2026
af2d7d8
perf(buffer): streamline fixed-size growth checks
peter-jerry-ye Jul 28, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
152 changes: 111 additions & 41 deletions buffer/buffer.mbt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,24 +41,44 @@ struct Buffer {
}

///|
/// Expand the buffer size if capacity smaller than required space.
fn Buffer::grow_if_necessary(self : Buffer, required : Int) -> Unit {
let start = if self.data.length() <= 0 { 1 } else { self.data.length() }
/// Compute the next capacity without allocating. Since appends never shrink the
/// buffer, `required < len` means the required-size calculation overflowed.
fn buffer_growth_capacity(current : Int, len : Int, required : Int) -> Int {
if required < len {
abort("Buffer capacity overflow")
}
let start = if current <= 0 { 1 } else { current }
let enough_space = for space = start {
if space >= required {
break space
}
continue space * 2
}
if enough_space != self.data.length() {
let new_data = FixedArray::make_and_blit(
self.data,
allocate_len=enough_space,
init=b'\x00',
len=self.len,
)
self.data = new_data
let next = space * 2
if next <= space {
break required
}
continue next
}
enough_space
}

///|
/// Grow the buffer to at least `required`. Callers keep the capacity check on
/// their fast path and enter here only when growth or overflow handling is
/// needed. The buffer invariant `0 <= len <= data.length()` lets fixed-size
/// appends compare against the remaining capacity without overflowing.
fn Buffer::grow(self : Buffer, required : Int) -> Unit {
let new_capacity = buffer_growth_capacity(
self.data.length(),
self.len,
required,
)
let new_data = FixedArray::make_and_blit(
self.data,
allocate_len=new_capacity,
init=b'\x00',
len=self.len,
)
self.data = new_data
}

///|
Expand Down Expand Up @@ -164,7 +184,7 @@ pub fn Buffer::Buffer(size_hint? : Int = 0) -> Buffer {
pub fn from_bytes(bytes : BytesView) -> Buffer {
let val_len = bytes.length()
let buf = Buffer(size_hint=val_len)
// inline write_bytes, skip grow_if_necessary check
// Inline write_bytes because the exact capacity is known.
// SAFETY: known bytes size
buf.data.blit_from_bytes(0, bytes.data(), bytes.start_offset(), val_len)
buf.len = val_len
Expand All @@ -176,7 +196,7 @@ pub fn from_bytes(bytes : BytesView) -> Buffer {
pub fn from_array(arr : ArrayView[Byte]) -> Buffer {
let buf = Buffer(size_hint=arr.length())
for byte in arr {
// inline write_byte, skip grow_if_necessary check
// Inline write_byte because the exact capacity is known.
// SAFETY: known array size
buf.data[buf.len] = byte
buf.len += 1
Expand All @@ -189,10 +209,9 @@ pub fn from_array(arr : ArrayView[Byte]) -> Buffer {
pub fn from_iter(iter : Iter[Byte]) -> Buffer {
let buf = Buffer()
for byte in iter; capacity = buf.data.length() {
// inline write_byte and grow_if_necessary
// only call grow_if_necessary when necessary
// Inline write_byte and keep growth off the fast path.
let capacity = if buf.len == capacity {
buf.grow_if_necessary(capacity + 1)
buf.grow(capacity + 1)
buf.data.length()
} else {
capacity
Expand Down Expand Up @@ -220,7 +239,10 @@ pub impl Logger for Buffer

///|
pub impl Logger for Buffer with fn write_string(self, value) {
self.grow_if_necessary(self.len + value.length() * 2)
let required = self.len + value.length() * 2
if required > self.data.length() || required < self.len {
self.grow(required)
}
self.data.blit_from_string(self.len, value, 0, value.length())
self.len += value.length() * 2
}
Expand Down Expand Up @@ -250,7 +272,9 @@ pub impl Logger for Buffer with fn write_string(self, value) {
/// }
/// ```
pub fn Buffer::write_uint64_be(self : Buffer, value : UInt64) -> Unit {
self.grow_if_necessary(self.len + 8)
if self.data.length() - self.len < 8 {
self.grow(self.len + 8)
}
let offset = self.len
self.data[offset] = (value >> 56).to_byte()
self.data[offset + 1] = (value >> 48).to_byte()
Expand Down Expand Up @@ -287,7 +311,9 @@ pub fn Buffer::write_uint64_be(self : Buffer, value : UInt64) -> Unit {
/// }
/// ```
pub fn Buffer::write_uint64_le(self : Buffer, value : UInt64) -> Unit {
self.grow_if_necessary(self.len + 8)
if self.data.length() - self.len < 8 {
self.grow(self.len + 8)
}
let offset = self.len
self.data[offset] = value.to_byte()
self.data[offset + 1] = (value >> 8).to_byte()
Expand Down Expand Up @@ -377,7 +403,9 @@ pub fn Buffer::write_int64_le(self : Buffer, value : Int64) -> Unit {
/// }
/// ```
pub fn Buffer::write_uint_be(self : Buffer, value : UInt) -> Unit {
self.grow_if_necessary(self.len + 4)
if self.data.length() - self.len < 4 {
self.grow(self.len + 4)
}
let offset = self.len
self.data[offset] = (value >> 24).to_byte()
self.data[offset + 1] = (value >> 16).to_byte()
Expand Down Expand Up @@ -411,7 +439,9 @@ pub fn Buffer::write_uint_be(self : Buffer, value : UInt) -> Unit {
/// }
/// ```
pub fn Buffer::write_uint_le(self : Buffer, value : UInt) -> Unit {
self.grow_if_necessary(self.len + 4)
if self.data.length() - self.len < 4 {
self.grow(self.len + 4)
}
let offset = self.len
self.data[offset] = value.to_byte()
self.data[offset + 1] = (value >> 8).to_byte()
Expand Down Expand Up @@ -494,7 +524,9 @@ pub fn Buffer::write_int_le(self : Buffer, value : Int) -> Unit {
/// }
/// ```
pub fn Buffer::write_uint16_be(self : Buffer, value : UInt16) -> Unit {
self.grow_if_necessary(self.len + 2)
if self.data.length() - self.len < 2 {
self.grow(self.len + 2)
}
let offset = self.len
self.data[offset] = (value.to_int() >> 8).to_byte()
self.data[offset + 1] = value.to_byte()
Expand Down Expand Up @@ -526,7 +558,9 @@ pub fn Buffer::write_uint16_be(self : Buffer, value : UInt16) -> Unit {
/// }
/// ```
pub fn Buffer::write_uint16_le(self : Buffer, value : UInt16) -> Unit {
self.grow_if_necessary(self.len + 2)
if self.data.length() - self.len < 2 {
self.grow(self.len + 2)
}
let offset = self.len
self.data[offset] = value.to_byte()
self.data[offset + 1] = (value.to_int() >> 8).to_byte()
Expand Down Expand Up @@ -557,7 +591,9 @@ pub fn Buffer::write_uint16_le(self : Buffer, value : UInt16) -> Unit {
/// }
/// ```
pub fn Buffer::write_int16_be(self : Buffer, value : Int16) -> Unit {
self.grow_if_necessary(self.len + 2)
if self.data.length() - self.len < 2 {
self.grow(self.len + 2)
}
let offset = self.len
self.data[offset] = (value.to_int() >> 8).to_byte()
self.data[offset + 1] = value.to_byte()
Expand Down Expand Up @@ -588,7 +624,9 @@ pub fn Buffer::write_int16_be(self : Buffer, value : Int16) -> Unit {
/// }
/// ```
pub fn Buffer::write_int16_le(self : Buffer, value : Int16) -> Unit {
self.grow_if_necessary(self.len + 2)
if self.data.length() - self.len < 2 {
self.grow(self.len + 2)
}
let offset = self.len
self.data[offset] = value.to_byte()
self.data[offset + 1] = (value.to_int() >> 8).to_byte()
Expand Down Expand Up @@ -801,7 +839,10 @@ pub fn Buffer::write_bytes(self : Buffer, value : BytesView) -> Unit {
/// ```
pub fn Buffer::write_bytesview(self : Buffer, value : BytesView) -> Unit {
let val_len = value.length()
self.grow_if_necessary(self.len + val_len)
let required = self.len + val_len
if required > self.data.length() || required < self.len {
self.grow(required)
}
self.data.blit_from_bytes(
self.len,
value.data(),
Expand All @@ -817,25 +858,33 @@ pub fn Buffer::write_char_utf8(buf : Self, value : Char) -> Unit {
let code = value.to_uint()
match code {
_..<0x80 => {
buf.grow_if_necessary(buf.len + 1)
if buf.len >= buf.data.length() {
buf.grow(buf.len + 1)
}
buf.data[buf.len] = ((code & 0x7F) | 0x00).to_byte()
buf.len += 1
}
_..<0x0800 => {
buf.grow_if_necessary(buf.len + 2)
if buf.data.length() - buf.len < 2 {
buf.grow(buf.len + 2)
}
buf.data[buf.len] = (((code >> 6) & 0x1F) | 0xC0).to_byte()
buf.data[buf.len + 1] = ((code & 0x3F) | 0x80).to_byte()
buf.len += 2
}
_..<0x010000 => {
buf.grow_if_necessary(buf.len + 3)
if buf.data.length() - buf.len < 3 {
buf.grow(buf.len + 3)
}
buf.data[buf.len] = (((code >> 12) & 0x0F) | 0xE0).to_byte()
buf.data[buf.len + 1] = (((code >> 6) & 0x3F) | 0x80).to_byte()
buf.data[buf.len + 2] = ((code & 0x3F) | 0x80).to_byte()
buf.len += 3
}
_..<0x110000 => {
buf.grow_if_necessary(buf.len + 4)
if buf.data.length() - buf.len < 4 {
buf.grow(buf.len + 4)
}
buf.data[buf.len] = (((code >> 18) & 0x07) | 0xF0).to_byte()
buf.data[buf.len + 1] = (((code >> 12) & 0x3F) | 0x80).to_byte()
buf.data[buf.len + 2] = (((code >> 6) & 0x3F) | 0x80).to_byte()
Expand All @@ -851,15 +900,19 @@ pub fn Buffer::write_char_utf8(buf : Self, value : Char) -> Unit {
pub fn Buffer::write_char_utf16le(buf : Self, value : Char) -> Unit {
let code = value.to_uint()
if code < 0x10000 {
buf.grow_if_necessary(buf.len + 2)
if buf.data.length() - buf.len < 2 {
buf.grow(buf.len + 2)
}
buf.data[buf.len + 0] = (code & 0xFF).to_byte()
buf.data[buf.len + 1] = (code >> 8).to_byte()
buf.len += 2
} else if code < 0x110000 {
let cp = code - 0x10000
let high = (cp >> 10) | 0xD800
let low = (cp & 0x3FF) | 0xDC00
buf.grow_if_necessary(buf.len + 4)
if buf.data.length() - buf.len < 4 {
buf.grow(buf.len + 4)
}
buf.data[buf.len + 0] = (high & 0xFF).to_byte()
buf.data[buf.len + 1] = (high >> 8).to_byte()
buf.data[buf.len + 2] = (low & 0xFF).to_byte()
Expand All @@ -875,12 +928,16 @@ pub fn Buffer::write_char_utf16le(buf : Self, value : Char) -> Unit {
pub fn Buffer::write_char_utf16be(buf : Self, value : Char) -> Unit {
let code = value.to_uint()
if code < 0x10000 {
buf.grow_if_necessary(buf.len + 2)
if buf.data.length() - buf.len < 2 {
buf.grow(buf.len + 2)
}
buf.data[buf.len + 0] = (code >> 8).to_byte()
buf.data[buf.len + 1] = (code & 0xFF).to_byte()
buf.len += 2
} else if code < 0x110000 {
buf.grow_if_necessary(buf.len + 4)
if buf.data.length() - buf.len < 4 {
buf.grow(buf.len + 4)
}
let cp = code - 0x10000
let high = (cp >> 10) | 0xD800
let low = (cp & 0x3FF) | 0xDC00
Expand Down Expand Up @@ -937,7 +994,10 @@ pub fn Buffer::write_string_utf8(buf : Self, string : StringView) -> Unit {
#alias(write_stringview, deprecated="use write_string_utf16le instead")
pub fn Buffer::write_string_utf16le(buf : Self, string : StringView) -> Unit {
let len = string.length()
buf.grow_if_necessary(buf.len + len * 2)
let required = buf.len + len * 2
if required > buf.data.length() || required < buf.len {
buf.grow(required)
}
for code_unit in string.code_units(); j = buf.len {
let c = code_unit.to_int().reinterpret_as_uint()
buf.data[j] = (c & 0xff).to_byte()
Expand Down Expand Up @@ -966,7 +1026,10 @@ pub fn Buffer::write_string_utf16le(buf : Self, string : StringView) -> Unit {
/// ```
pub fn Buffer::write_string_utf16be(buf : Self, string : StringView) -> Unit {
let len = string.length()
buf.grow_if_necessary(buf.len + len * 2)
let required = buf.len + len * 2
if required > buf.data.length() || required < buf.len {
buf.grow(required)
}
for code_unit in string.code_units(); j = buf.len {
let c = code_unit.to_int().reinterpret_as_uint()
buf.data[j + 1] = (c & 0xff).to_byte()
Expand All @@ -986,7 +1049,10 @@ pub fn Buffer::write_string_utf16be(buf : Self, string : StringView) -> Unit {
/// * `count` : The number of characters to write. Must be non-negative and
/// `offset + count` must not exceed the length of the source string.
pub impl Logger for Buffer with fn write_view(self : Buffer, value : StringView) -> Unit {
self.grow_if_necessary(self.len + value.length() * 2)
let required = self.len + value.length() * 2
if required > self.data.length() || required < self.len {
self.grow(required)
}
self.data.blit_from_string(
self.len,
value.data(),
Expand All @@ -1005,7 +1071,9 @@ pub impl Logger for Buffer with fn write_view(self : Buffer, value : StringView)
/// * `buffer` : The buffer to write to.
/// * `char` : The character to be written.
pub impl Logger for Buffer with fn write_char(self : Buffer, value : Char) -> Unit {
self.grow_if_necessary(self.len + 4)
if self.data.length() - self.len < 4 {
self.grow(self.len + 4)
}
let inc = self.data.set_utf16le_char(self.len, value)
self.len += inc
}
Expand Down Expand Up @@ -1034,7 +1102,9 @@ pub impl Logger for Buffer with fn write_char(self : Buffer, value : Char) -> Un
/// }
/// ```
pub fn Buffer::write_byte(self : Buffer, value : Byte) -> Unit {
self.grow_if_necessary(self.len + 1)
if self.len >= self.data.length() {
self.grow(self.len + 1)
}
self.data[self.len] = value
self.len += 1
}
Expand Down
36 changes: 36 additions & 0 deletions buffer/buffer_growth_wbtest.mbt
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,36 @@
// Copyright 2026 International Digital Economy Academy
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

///|
test "buffer growth capacity doubles while representable" {
inspect(buffer_growth_capacity(8, 8, 9), content="16")
inspect(buffer_growth_capacity(8, 8, 33), content="64")
}

///|
test "buffer growth capacity falls back to the exact requirement" {
inspect(
buffer_growth_capacity(0x40000000, 0x40000000, 0x40000001),
content="1073741825",
)
inspect(
buffer_growth_capacity(0x20000000, 0x20000000, 0x40000001),
content="1073741825",
)
}

///|
test "panic buffer growth rejects a wrapped required size" {
ignore(buffer_growth_capacity(16, 16, -1))
}
2 changes: 1 addition & 1 deletion buffer/buffer_test.mbt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -63,7 +63,7 @@ test "expect method with matching content" {
}

///|
test "grow_if_necessary method" {
test "buffer grows beyond its initial capacity" {
let buf = Buffer(size_hint=10)
buf.write_string_utf16le(
"This is a test string that is longer than the initial capacity",
Expand Down
Loading
Loading