Skip to content

Commit d59b0ae

Browse files
aaron-angaaron-ang
committed
std.math.big: require sufficient capacity for aliased params
1 parent d828115 commit d59b0ae

File tree

2 files changed

+103
-29
lines changed

2 files changed

+103
-29
lines changed

‎lib/std/math/big/int.zig‎

Lines changed: 94 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -2971,7 +2971,9 @@ pub const Managed = struct {
29712971
///
29722972
/// Returns an error if memory could not be allocated.
29732973
pub fn addScalar(r: *Managed, a: *const Managed, scalar: anytype) Allocator.Error!void {
2974-
try r.ensureAddScalarCapacity(a.toConst(), scalar);
2974+
const needed = @max(a.len(), calcLimbLen(scalar)) + 1;
2975+
const aliased = r.limbs.ptr == a.limbs.ptr;
2976+
try r.ensureAliasAwareCapacity(needed, aliased);
29752977
var m = r.toMutable();
29762978
m.addScalar(a.toConst(), scalar);
29772979
r.setMetadata(m.positive, m.len);
@@ -2983,7 +2985,9 @@ pub const Managed = struct {
29832985
///
29842986
/// Returns an error if memory could not be allocated.
29852987
pub fn add(r: *Managed, a: *const Managed, b: *const Managed) Allocator.Error!void {
2986-
try r.ensureAddCapacity(a.toConst(), b.toConst());
2988+
const needed = @max(a.len(), b.len()) + 1;
2989+
const aliased = r.limbs.ptr == a.limbs.ptr or r.limbs.ptr == b.limbs.ptr;
2990+
try r.ensureAliasAwareCapacity(needed, aliased);
29872991
var m = r.toMutable();
29882992
m.add(a.toConst(), b.toConst());
29892993
r.setMetadata(m.positive, m.len);
@@ -3001,7 +3005,9 @@ pub const Managed = struct {
30013005
signedness: Signedness,
30023006
bit_count: usize,
30033007
) Allocator.Error!bool {
3004-
try r.ensureTwosCompCapacity(bit_count);
3008+
const aliased = r.limbs.ptr == a.limbs.ptr or r.limbs.ptr == b.limbs.ptr;
3009+
const needed = calcTwosCompLimbCount(bit_count);
3010+
try r.ensureAliasAwareCapacity(needed, aliased);
30053011
var m = r.toMutable();
30063012
const wrapped = m.addWrap(a.toConst(), b.toConst(), signedness, bit_count);
30073013
r.setMetadata(m.positive, m.len);
@@ -3014,7 +3020,9 @@ pub const Managed = struct {
30143020
///
30153021
/// Returns an error if memory could not be allocated.
30163022
pub fn addSat(r: *Managed, a: *const Managed, b: *const Managed, signedness: Signedness, bit_count: usize) Allocator.Error!void {
3017-
try r.ensureTwosCompCapacity(bit_count);
3023+
const aliased = r.limbs.ptr == a.limbs.ptr or r.limbs.ptr == b.limbs.ptr;
3024+
const needed = calcTwosCompLimbCount(bit_count);
3025+
try r.ensureAliasAwareCapacity(needed, aliased);
30183026
var m = r.toMutable();
30193027
m.addSat(a.toConst(), b.toConst(), signedness, bit_count);
30203028
r.setMetadata(m.positive, m.len);
@@ -3026,7 +3034,9 @@ pub const Managed = struct {
30263034
///
30273035
/// Returns an error if memory could not be allocated.
30283036
pub fn sub(r: *Managed, a: *const Managed, b: *const Managed) !void {
3029-
try r.ensureCapacity(@max(a.len(), b.len()) + 1);
3037+
const aliased = r.limbs.ptr == a.limbs.ptr or r.limbs.ptr == b.limbs.ptr;
3038+
const needed = @max(a.len(), b.len()) + 1;
3039+
try r.ensureAliasAwareCapacity(needed, aliased);
30303040
var m = r.toMutable();
30313041
m.sub(a.toConst(), b.toConst());
30323042
r.setMetadata(m.positive, m.len);
@@ -3044,7 +3054,9 @@ pub const Managed = struct {
30443054
signedness: Signedness,
30453055
bit_count: usize,
30463056
) Allocator.Error!bool {
3047-
try r.ensureTwosCompCapacity(bit_count);
3057+
const aliased = r.limbs.ptr == a.limbs.ptr or r.limbs.ptr == b.limbs.ptr;
3058+
const needed = calcTwosCompLimbCount(bit_count);
3059+
try r.ensureAliasAwareCapacity(needed, aliased);
30483060
var m = r.toMutable();
30493061
const wrapped = m.subWrap(a.toConst(), b.toConst(), signedness, bit_count);
30503062
r.setMetadata(m.positive, m.len);
@@ -3063,7 +3075,9 @@ pub const Managed = struct {
30633075
signedness: Signedness,
30643076
bit_count: usize,
30653077
) Allocator.Error!void {
3066-
try r.ensureTwosCompCapacity(bit_count);
3078+
const aliased = r.limbs.ptr == a.limbs.ptr or r.limbs.ptr == b.limbs.ptr;
3079+
const needed = calcTwosCompLimbCount(bit_count);
3080+
try r.ensureAliasAwareCapacity(needed, aliased);
30673081
var m = r.toMutable();
30683082
m.subSat(a.toConst(), b.toConst(), signedness, bit_count);
30693083
r.setMetadata(m.positive, m.len);
@@ -3082,7 +3096,8 @@ pub const Managed = struct {
30823096
alias_count += 1;
30833097
if (rma.limbs.ptr == b.limbs.ptr)
30843098
alias_count += 1;
3085-
try rma.ensureMulCapacity(a.toConst(), b.toConst());
3099+
const needed = a.len() + b.len() + 1;
3100+
try rma.ensureAliasAwareCapacity(needed, alias_count > 0);
30863101
var m = rma.toMutable();
30873102
if (alias_count == 0) {
30883103
m.mulNoAlias(a.toConst(), b.toConst(), rma.allocator);
@@ -3114,8 +3129,8 @@ pub const Managed = struct {
31143129
alias_count += 1;
31153130
if (rma.limbs.ptr == b.limbs.ptr)
31163131
alias_count += 1;
3117-
3118-
try rma.ensureTwosCompCapacity(bit_count);
3132+
const needed = calcTwosCompLimbCount(bit_count);
3133+
try rma.ensureAliasAwareCapacity(needed, alias_count > 0);
31193134
var m = rma.toMutable();
31203135
if (alias_count == 0) {
31213136
m.mulWrapNoAlias(a.toConst(), b.toConst(), signedness, bit_count, rma.allocator);
@@ -3132,26 +3147,58 @@ pub const Managed = struct {
31323147
try r.ensureCapacity(calcTwosCompLimbCount(bit_count));
31333148
}
31343149

3150+
/// Ensures capacity only when parameters do not alias the result.
3151+
///
3152+
/// When aliasing is detected, this function requires the caller to have already
3153+
/// ensured sufficient capacity. This prevents use-after-free bugs that occur when
3154+
/// reallocating memory while const parameter pointers reference that same memory.
3155+
///
3156+
/// Callers using aliasing must pre-allocate capacity using the appropriate
3157+
/// `ensure*Capacity` helper before calling the operation.
3158+
///
3159+
/// See: https://github.com/ziglang/zig/issues/6167
3160+
fn ensureAliasAwareCapacity(r: *Managed, needed: usize, aliased: bool) !void {
3161+
if (aliased) {
3162+
assert(needed <= r.limbs.len);
3163+
} else {
3164+
try r.ensureCapacity(needed);
3165+
}
3166+
}
3167+
31353168
pub fn ensureAddScalarCapacity(r: *Managed, a: Const, scalar: anytype) !void {
31363169
try r.ensureCapacity(@max(a.limbs.len, calcLimbLen(scalar)) + 1);
31373170
}
31383171

3172+
pub fn ensureAddScalarCapacityManaged(r: *Managed, a: *const Managed, scalar: anytype) !void {
3173+
try r.ensureCapacity(@max(a.len(), calcLimbLen(scalar)) + 1);
3174+
}
3175+
31393176
pub fn ensureAddCapacity(r: *Managed, a: Const, b: Const) !void {
31403177
try r.ensureCapacity(@max(a.limbs.len, b.limbs.len) + 1);
31413178
}
31423179

3180+
pub fn ensureAddCapacityManaged(r: *Managed, a: *const Managed, b: *const Managed) !void {
3181+
try r.ensureCapacity(@max(a.len(), b.len()) + 1);
3182+
}
3183+
31433184
pub fn ensureMulCapacity(rma: *Managed, a: Const, b: Const) !void {
31443185
try rma.ensureCapacity(a.limbs.len + b.limbs.len + 1);
31453186
}
31463187

3188+
pub fn ensureMulCapacityManaged(rma: *Managed, a: *const Managed, b: *const Managed) !void {
3189+
try rma.ensureCapacity(a.len() + b.len() + 1);
3190+
}
3191+
31473192
/// q = a / b (rem r)
31483193
///
31493194
/// a / b are floored (rounded towards 0).
31503195
///
31513196
/// Returns an error if memory could not be allocated.
31523197
pub fn divFloor(q: *Managed, r: *Managed, a: *const Managed, b: *const Managed) !void {
3153-
try q.ensureCapacity(a.len());
3154-
try r.ensureCapacity(b.len());
3198+
const q_alias = q.limbs.ptr == a.limbs.ptr or q.limbs.ptr == b.limbs.ptr;
3199+
const r_alias = r.limbs.ptr == a.limbs.ptr or r.limbs.ptr == b.limbs.ptr;
3200+
try q.ensureAliasAwareCapacity(a.len(), q_alias);
3201+
try r.ensureAliasAwareCapacity(b.len(), r_alias);
31553202
var mq = q.toMutable();
31563203
var mr = r.toMutable();
31573204
const limbs_buffer = try q.allocator.alloc(Limb, calcDivLimbsBufferLen(a.len(), b.len()));
@@ -3167,8 +3214,10 @@ pub const Managed = struct {
31673214
///
31683215
/// Returns an error if memory could not be allocated.
31693216
pub fn divTrunc(q: *Managed, r: *Managed, a: *const Managed, b: *const Managed) !void {
3170-
try q.ensureCapacity(a.len());
3171-
try r.ensureCapacity(b.len());
3217+
const q_alias = q.limbs.ptr == a.limbs.ptr or q.limbs.ptr == b.limbs.ptr;
3218+
const r_alias = r.limbs.ptr == a.limbs.ptr or r.limbs.ptr == b.limbs.ptr;
3219+
try q.ensureAliasAwareCapacity(a.len(), q_alias);
3220+
try r.ensureAliasAwareCapacity(b.len(), r_alias);
31723221
var mq = q.toMutable();
31733222
var mr = r.toMutable();
31743223
const limbs_buffer = try q.allocator.alloc(Limb, calcDivLimbsBufferLen(a.len(), b.len()));
@@ -3181,7 +3230,9 @@ pub const Managed = struct {
31813230
/// r = a << shift, in other words, r = a * 2^shift
31823231
/// r and a may alias.
31833232
pub fn shiftLeft(r: *Managed, a: *const Managed, shift: usize) !void {
3184-
try r.ensureCapacity(a.len() + (shift / limb_bits) + 1);
3233+
const aliased = r.limbs.ptr == a.limbs.ptr;
3234+
const needed = a.len() + (shift / limb_bits) + 1;
3235+
try r.ensureAliasAwareCapacity(needed, aliased);
31853236
var m = r.toMutable();
31863237
m.shiftLeft(a.toConst(), shift);
31873238
r.setMetadata(m.positive, m.len);
@@ -3190,7 +3241,9 @@ pub const Managed = struct {
31903241
/// r = a <<| shift with 2s-complement saturating semantics.
31913242
/// r and a may alias.
31923243
pub fn shiftLeftSat(r: *Managed, a: *const Managed, shift: usize, signedness: Signedness, bit_count: usize) !void {
3193-
try r.ensureTwosCompCapacity(bit_count);
3244+
const aliased = r.limbs.ptr == a.limbs.ptr;
3245+
const needed = calcTwosCompLimbCount(bit_count);
3246+
try r.ensureAliasAwareCapacity(needed, aliased);
31943247
var m = r.toMutable();
31953248
m.shiftLeftSat(a.toConst(), shift, signedness, bit_count);
31963249
r.setMetadata(m.positive, m.len);
@@ -3212,7 +3265,9 @@ pub const Managed = struct {
32123265
return;
32133266
}
32143267

3215-
try r.ensureCapacity(a.len() - (shift / limb_bits));
3268+
const aliased = r.limbs.ptr == a.limbs.ptr;
3269+
const needed = a.len() - (shift / limb_bits);
3270+
try r.ensureAliasAwareCapacity(needed, aliased);
32163271
var m = r.toMutable();
32173272
m.shiftRight(a.toConst(), shift);
32183273
r.setMetadata(m.positive, m.len);
@@ -3221,7 +3276,9 @@ pub const Managed = struct {
32213276
/// r = ~a under 2s-complement wrapping semantics.
32223277
/// r and a may alias.
32233278
pub fn bitNotWrap(r: *Managed, a: *const Managed, signedness: Signedness, bit_count: usize) !void {
3224-
try r.ensureTwosCompCapacity(bit_count);
3279+
const aliased = r.limbs.ptr == a.limbs.ptr;
3280+
const needed = calcTwosCompLimbCount(bit_count);
3281+
try r.ensureAliasAwareCapacity(needed, aliased);
32253282
var m = r.toMutable();
32263283
m.bitNotWrap(a.toConst(), signedness, bit_count);
32273284
r.setMetadata(m.positive, m.len);
@@ -3231,7 +3288,9 @@ pub const Managed = struct {
32313288
///
32323289
/// a and b are zero-extended to the longer of a or b.
32333290
pub fn bitOr(r: *Managed, a: *const Managed, b: *const Managed) !void {
3234-
try r.ensureCapacity(@max(a.len(), b.len()));
3291+
const aliased = r.limbs.ptr == a.limbs.ptr or r.limbs.ptr == b.limbs.ptr;
3292+
const needed = @max(a.len(), b.len());
3293+
try r.ensureAliasAwareCapacity(needed, aliased);
32353294
var m = r.toMutable();
32363295
m.bitOr(a.toConst(), b.toConst());
32373296
r.setMetadata(m.positive, m.len);
@@ -3243,7 +3302,8 @@ pub const Managed = struct {
32433302
if (b.isPositive()) b.len() else if (a.isPositive()) a.len() else a.len() + 1
32443303
else if (a.isPositive()) a.len() else if (b.isPositive()) b.len() else b.len() + 1;
32453304

3246-
try r.ensureCapacity(cap);
3305+
const aliased = r.limbs.ptr == a.limbs.ptr or r.limbs.ptr == b.limbs.ptr;
3306+
try r.ensureAliasAwareCapacity(cap, aliased);
32473307
var m = r.toMutable();
32483308
m.bitAnd(a.toConst(), b.toConst());
32493309
r.setMetadata(m.positive, m.len);
@@ -3252,7 +3312,8 @@ pub const Managed = struct {
32523312
/// r = a ^ b
32533313
pub fn bitXor(r: *Managed, a: *const Managed, b: *const Managed) !void {
32543314
const cap = @max(a.len(), b.len()) + @intFromBool(a.isPositive() != b.isPositive());
3255-
try r.ensureCapacity(cap);
3315+
const aliased = r.limbs.ptr == a.limbs.ptr or r.limbs.ptr == b.limbs.ptr;
3316+
try r.ensureAliasAwareCapacity(cap, aliased);
32563317

32573318
var m = r.toMutable();
32583319
m.bitXor(a.toConst(), b.toConst());
@@ -3264,7 +3325,9 @@ pub const Managed = struct {
32643325
///
32653326
/// rma's allocator is used for temporary storage to boost multiplication performance.
32663327
pub fn gcd(rma: *Managed, x: *const Managed, y: *const Managed) !void {
3267-
try rma.ensureCapacity(@min(x.len(), y.len()));
3328+
const aliased = rma.limbs.ptr == x.limbs.ptr or rma.limbs.ptr == y.limbs.ptr;
3329+
const needed = @min(x.len(), y.len());
3330+
try rma.ensureAliasAwareCapacity(needed, aliased);
32683331
var m = rma.toMutable();
32693332
var limbs_buffer = std.array_list.Managed(Limb).init(rma.allocator);
32703333
defer limbs_buffer.deinit();
@@ -3342,15 +3405,19 @@ pub const Managed = struct {
33423405

33433406
/// r = truncate(Int(signedness, bit_count), a)
33443407
pub fn truncate(r: *Managed, a: *const Managed, signedness: Signedness, bit_count: usize) !void {
3345-
try r.ensureCapacity(calcTwosCompLimbCount(bit_count));
3408+
const aliased = r.limbs.ptr == a.limbs.ptr;
3409+
const needed = calcTwosCompLimbCount(bit_count);
3410+
try r.ensureAliasAwareCapacity(needed, aliased);
33463411
var m = r.toMutable();
33473412
m.truncate(a.toConst(), signedness, bit_count);
33483413
r.setMetadata(m.positive, m.len);
33493414
}
33503415

33513416
/// r = saturate(Int(signedness, bit_count), a)
33523417
pub fn saturate(r: *Managed, a: *const Managed, signedness: Signedness, bit_count: usize) !void {
3353-
try r.ensureCapacity(calcTwosCompLimbCount(bit_count));
3418+
const aliased = r.limbs.ptr == a.limbs.ptr;
3419+
const needed = calcTwosCompLimbCount(bit_count);
3420+
try r.ensureAliasAwareCapacity(needed, aliased);
33543421
var m = r.toMutable();
33553422
m.saturate(a.toConst(), signedness, bit_count);
33563423
r.setMetadata(m.positive, m.len);
@@ -3359,7 +3426,9 @@ pub const Managed = struct {
33593426
/// r = @popCount(a) with 2s-complement semantics.
33603427
/// r and a may be aliases.
33613428
pub fn popCount(r: *Managed, a: *const Managed, bit_count: usize) !void {
3362-
try r.ensureCapacity(calcTwosCompLimbCount(bit_count));
3429+
const aliased = r.limbs.ptr == a.limbs.ptr;
3430+
const needed = calcTwosCompLimbCount(bit_count);
3431+
try r.ensureAliasAwareCapacity(needed, aliased);
33633432
var m = r.toMutable();
33643433
m.popCount(a.toConst(), bit_count);
33653434
r.setMetadata(m.positive, m.len);

‎lib/std/math/big/int_test.zig‎

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -583,7 +583,9 @@ test "bitcount + sizeInBaseUpperBound" {
583583
try testing.expect(a.sizeInBaseUpperBound(2) >= 32);
584584
try testing.expect(a.sizeInBaseUpperBound(10) >= 10);
585585

586-
try a.shiftLeft(&a, 5000);
586+
const shift = 5000;
587+
try a.ensureCapacity(a.len() + (shift / @bitSizeOf(Limb)) + 1);
588+
try a.shiftLeft(&a, shift);
587589
try testing.expectEqual(5032, a.bitCountAbs());
588590
try testing.expect(a.sizeInBaseUpperBound(2) >= 5032);
589591
a.setSign(false);
@@ -2380,7 +2382,9 @@ test "truncate negative multi to single" {
23802382
test "truncate multi unsigned many" {
23812383
var a = try Managed.initSet(testing.allocator, 1);
23822384
defer a.deinit();
2383-
try a.shiftLeft(&a, 1023);
2385+
const shift = 1023;
2386+
try a.ensureCapacity(a.len() + (shift / @bitSizeOf(Limb)) + 1);
2387+
try a.shiftLeft(&a, shift);
23842388

23852389
var b = try Managed.init(testing.allocator);
23862390
defer b.deinit();
@@ -3263,7 +3267,7 @@ test "regression test for 1 limb overflow with alias" {
32633267
var b = try Managed.initSet(testing.allocator, 12200160415121876738);
32643268
defer b.deinit();
32653269

3266-
try a.ensureAddCapacity(a.toConst(), b.toConst());
3270+
try a.ensureAddCapacityManaged(&a, &b);
32673271
try a.add(&a, &b);
32683272

32693273
try testing.expectEqual(.eq, a.toConst().orderAgainstScalar(19740274219868223167));
@@ -3277,7 +3281,7 @@ test "regression test for realloc with alias" {
32773281
var b = try Managed.initSet(testing.allocator, 9079598147510263717870894449029933369491131786514446266146);
32783282
defer b.deinit();
32793283

3280-
try a.ensureAddCapacity(a.toConst(), b.toConst());
3284+
try a.ensureAddCapacityManaged(&a, &b);
32813285
try a.add(&a, &b);
32823286

32833287
try testing.expectEqual(.eq, a.toConst().orderAgainstScalar(14691098406862188148944207245954912110548093601382197697835));
@@ -3692,6 +3696,7 @@ test "mul multi-multi alias r with a and b" {
36923696
var a = try Managed.initSet(testing.allocator, 2 * maxInt(Limb));
36933697
defer a.deinit();
36943698

3699+
try a.ensureMulCapacityManaged(&a, &a);
36953700
try a.mul(&a, &a);
36963701

36973702
var want = try Managed.initSet(testing.allocator, 4 * maxInt(Limb) * maxInt(Limb));

0 commit comments

Comments
 (0)