Skip to content

Commit deb84da

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

File tree

2 files changed

+113
-29
lines changed

2 files changed

+113
-29
lines changed

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

Lines changed: 104 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -2979,7 +2979,9 @@ pub const Managed = struct {
29792979
///
29802980
/// Returns an error if memory could not be allocated.
29812981
pub fn addScalar(r: *Managed, a: *const Managed, scalar: anytype) Allocator.Error!void {
2982-
try r.ensureAddScalarCapacity(a.toConst(), scalar);
2982+
const needed = @max(a.len(), calcLimbLen(scalar)) + 1;
2983+
const aliased = limbsAliasDistinct(r, a);
2984+
try r.ensureAliasAwareCapacity(needed, aliased);
29832985
var m = r.toMutable();
29842986
m.addScalar(a.toConst(), scalar);
29852987
r.setMetadata(m.positive, m.len);
@@ -2991,7 +2993,9 @@ pub const Managed = struct {
29912993
///
29922994
/// Returns an error if memory could not be allocated.
29932995
pub fn add(r: *Managed, a: *const Managed, b: *const Managed) Allocator.Error!void {
2994-
try r.ensureAddCapacity(a.toConst(), b.toConst());
2996+
const needed = @max(a.len(), b.len()) + 1;
2997+
const aliased = limbsAliasDistinct(r, a) or limbsAliasDistinct(r, b);
2998+
try r.ensureAliasAwareCapacity(needed, aliased);
29952999
var m = r.toMutable();
29963000
m.add(a.toConst(), b.toConst());
29973001
r.setMetadata(m.positive, m.len);
@@ -3009,7 +3013,9 @@ pub const Managed = struct {
30093013
signedness: Signedness,
30103014
bit_count: usize,
30113015
) Allocator.Error!bool {
3012-
try r.ensureTwosCompCapacity(bit_count);
3016+
const aliased = limbsAliasDistinct(r, a) or limbsAliasDistinct(r, b);
3017+
const needed = calcTwosCompLimbCount(bit_count);
3018+
try r.ensureAliasAwareCapacity(needed, aliased);
30133019
var m = r.toMutable();
30143020
const wrapped = m.addWrap(a.toConst(), b.toConst(), signedness, bit_count);
30153021
r.setMetadata(m.positive, m.len);
@@ -3022,7 +3028,9 @@ pub const Managed = struct {
30223028
///
30233029
/// Returns an error if memory could not be allocated.
30243030
pub fn addSat(r: *Managed, a: *const Managed, b: *const Managed, signedness: Signedness, bit_count: usize) Allocator.Error!void {
3025-
try r.ensureTwosCompCapacity(bit_count);
3031+
const aliased = limbsAliasDistinct(r, a) or limbsAliasDistinct(r, b);
3032+
const needed = calcTwosCompLimbCount(bit_count);
3033+
try r.ensureAliasAwareCapacity(needed, aliased);
30263034
var m = r.toMutable();
30273035
m.addSat(a.toConst(), b.toConst(), signedness, bit_count);
30283036
r.setMetadata(m.positive, m.len);
@@ -3034,7 +3042,9 @@ pub const Managed = struct {
30343042
///
30353043
/// Returns an error if memory could not be allocated.
30363044
pub fn sub(r: *Managed, a: *const Managed, b: *const Managed) !void {
3037-
try r.ensureCapacity(@max(a.len(), b.len()) + 1);
3045+
const aliased = limbsAliasDistinct(r, a) or limbsAliasDistinct(r, b);
3046+
const needed = @max(a.len(), b.len()) + 1;
3047+
try r.ensureAliasAwareCapacity(needed, aliased);
30383048
var m = r.toMutable();
30393049
m.sub(a.toConst(), b.toConst());
30403050
r.setMetadata(m.positive, m.len);
@@ -3052,7 +3062,9 @@ pub const Managed = struct {
30523062
signedness: Signedness,
30533063
bit_count: usize,
30543064
) Allocator.Error!bool {
3055-
try r.ensureTwosCompCapacity(bit_count);
3065+
const aliased = limbsAliasDistinct(r, a) or limbsAliasDistinct(r, b);
3066+
const needed = calcTwosCompLimbCount(bit_count);
3067+
try r.ensureAliasAwareCapacity(needed, aliased);
30563068
var m = r.toMutable();
30573069
const wrapped = m.subWrap(a.toConst(), b.toConst(), signedness, bit_count);
30583070
r.setMetadata(m.positive, m.len);
@@ -3071,7 +3083,9 @@ pub const Managed = struct {
30713083
signedness: Signedness,
30723084
bit_count: usize,
30733085
) Allocator.Error!void {
3074-
try r.ensureTwosCompCapacity(bit_count);
3086+
const aliased = limbsAliasDistinct(r, a) or limbsAliasDistinct(r, b);
3087+
const needed = calcTwosCompLimbCount(bit_count);
3088+
try r.ensureAliasAwareCapacity(needed, aliased);
30753089
var m = r.toMutable();
30763090
m.subSat(a.toConst(), b.toConst(), signedness, bit_count);
30773091
r.setMetadata(m.positive, m.len);
@@ -3090,7 +3104,9 @@ pub const Managed = struct {
30903104
alias_count += 1;
30913105
if (rma.limbs.ptr == b.limbs.ptr)
30923106
alias_count += 1;
3093-
try rma.ensureMulCapacity(a.toConst(), b.toConst());
3107+
const needed = a.len() + b.len() + 1;
3108+
const capacity_alias = limbsAliasDistinct(rma, a) or limbsAliasDistinct(rma, b);
3109+
try rma.ensureAliasAwareCapacity(needed, capacity_alias);
30943110
var m = rma.toMutable();
30953111
if (alias_count == 0) {
30963112
m.mulNoAlias(a.toConst(), b.toConst(), rma.allocator);
@@ -3122,8 +3138,9 @@ pub const Managed = struct {
31223138
alias_count += 1;
31233139
if (rma.limbs.ptr == b.limbs.ptr)
31243140
alias_count += 1;
3125-
3126-
try rma.ensureTwosCompCapacity(bit_count);
3141+
const needed = calcTwosCompLimbCount(bit_count);
3142+
const capacity_alias = limbsAliasDistinct(rma, a) or limbsAliasDistinct(rma, b);
3143+
try rma.ensureAliasAwareCapacity(needed, capacity_alias);
31273144
var m = rma.toMutable();
31283145
if (alias_count == 0) {
31293146
m.mulWrapNoAlias(a.toConst(), b.toConst(), signedness, bit_count, rma.allocator);
@@ -3140,26 +3157,66 @@ pub const Managed = struct {
31403157
try r.ensureCapacity(calcTwosCompLimbCount(bit_count));
31413158
}
31423159

3160+
/// Returns true when two distinct `Managed` instances share the same limbs buffer.
3161+
/// Reallocating one will not update the other.
3162+
fn limbsAliasDistinct(a: *const Managed, b: *const Managed) bool {
3163+
return @intFromPtr(a) != @intFromPtr(b) and a.limbs.ptr == b.limbs.ptr;
3164+
}
3165+
3166+
/// Ensures capacity only when parameters do not alias the result.
3167+
///
3168+
/// When aliasing is detected, this function requires the caller to have already
3169+
/// ensured sufficient capacity. This prevents use-after-free bugs that occur when
3170+
/// reallocating memory while const parameter pointers reference that same memory.
3171+
/// `aliased` should only be set when the limbs buffer is shared with a different
3172+
/// `Managed` instance that would not observe a reallocation.
3173+
///
3174+
/// Callers using aliasing must pre-allocate capacity using the appropriate
3175+
/// `ensure*Capacity` helper before calling the arithmetic operation.
3176+
///
3177+
/// See: https://github.com/ziglang/zig/issues/6167
3178+
fn ensureAliasAwareCapacity(r: *Managed, needed: usize, aliased: bool) !void {
3179+
if (aliased) {
3180+
assert(needed <= r.limbs.len);
3181+
} else {
3182+
try r.ensureCapacity(needed);
3183+
}
3184+
}
3185+
31433186
pub fn ensureAddScalarCapacity(r: *Managed, a: Const, scalar: anytype) !void {
31443187
try r.ensureCapacity(@max(a.limbs.len, calcLimbLen(scalar)) + 1);
31453188
}
31463189

3190+
pub fn ensureAddScalarCapacityManaged(r: *Managed, a: *const Managed, scalar: anytype) !void {
3191+
try r.ensureCapacity(@max(a.len(), calcLimbLen(scalar)) + 1);
3192+
}
3193+
31473194
pub fn ensureAddCapacity(r: *Managed, a: Const, b: Const) !void {
31483195
try r.ensureCapacity(@max(a.limbs.len, b.limbs.len) + 1);
31493196
}
31503197

3198+
pub fn ensureAddCapacityManaged(r: *Managed, a: *const Managed, b: *const Managed) !void {
3199+
try r.ensureCapacity(@max(a.len(), b.len()) + 1);
3200+
}
3201+
31513202
pub fn ensureMulCapacity(rma: *Managed, a: Const, b: Const) !void {
31523203
try rma.ensureCapacity(a.limbs.len + b.limbs.len + 1);
31533204
}
31543205

3206+
pub fn ensureMulCapacityManaged(rma: *Managed, a: *const Managed, b: *const Managed) !void {
3207+
try rma.ensureCapacity(a.len() + b.len() + 1);
3208+
}
3209+
31553210
/// q = a / b (rem r)
31563211
///
31573212
/// a / b are floored (rounded towards 0).
31583213
///
31593214
/// Returns an error if memory could not be allocated.
31603215
pub fn divFloor(q: *Managed, r: *Managed, a: *const Managed, b: *const Managed) !void {
3161-
try q.ensureCapacity(a.len());
3162-
try r.ensureCapacity(b.len());
3216+
const q_alias = limbsAliasDistinct(q, a) or limbsAliasDistinct(q, b);
3217+
const r_alias = limbsAliasDistinct(r, a) or limbsAliasDistinct(r, b);
3218+
try q.ensureAliasAwareCapacity(a.len(), q_alias);
3219+
try r.ensureAliasAwareCapacity(b.len(), r_alias);
31633220
var mq = q.toMutable();
31643221
var mr = r.toMutable();
31653222
const limbs_buffer = try q.allocator.alloc(Limb, calcDivLimbsBufferLen(a.len(), b.len()));
@@ -3175,8 +3232,10 @@ pub const Managed = struct {
31753232
///
31763233
/// Returns an error if memory could not be allocated.
31773234
pub fn divTrunc(q: *Managed, r: *Managed, a: *const Managed, b: *const Managed) !void {
3178-
try q.ensureCapacity(a.len());
3179-
try r.ensureCapacity(b.len());
3235+
const q_alias = limbsAliasDistinct(q, a) or limbsAliasDistinct(q, b);
3236+
const r_alias = limbsAliasDistinct(r, a) or limbsAliasDistinct(r, b);
3237+
try q.ensureAliasAwareCapacity(a.len(), q_alias);
3238+
try r.ensureAliasAwareCapacity(b.len(), r_alias);
31803239
var mq = q.toMutable();
31813240
var mr = r.toMutable();
31823241
const limbs_buffer = try q.allocator.alloc(Limb, calcDivLimbsBufferLen(a.len(), b.len()));
@@ -3189,7 +3248,9 @@ pub const Managed = struct {
31893248
/// r = a << shift, in other words, r = a * 2^shift
31903249
/// r and a may alias.
31913250
pub fn shiftLeft(r: *Managed, a: *const Managed, shift: usize) !void {
3192-
try r.ensureCapacity(a.len() + (shift / limb_bits) + 1);
3251+
const aliased = limbsAliasDistinct(r, a);
3252+
const needed = a.len() + (shift / limb_bits) + 1;
3253+
try r.ensureAliasAwareCapacity(needed, aliased);
31933254
var m = r.toMutable();
31943255
m.shiftLeft(a.toConst(), shift);
31953256
r.setMetadata(m.positive, m.len);
@@ -3198,7 +3259,9 @@ pub const Managed = struct {
31983259
/// r = a <<| shift with 2s-complement saturating semantics.
31993260
/// r and a may alias.
32003261
pub fn shiftLeftSat(r: *Managed, a: *const Managed, shift: usize, signedness: Signedness, bit_count: usize) !void {
3201-
try r.ensureTwosCompCapacity(bit_count);
3262+
const aliased = limbsAliasDistinct(r, a);
3263+
const needed = calcTwosCompLimbCount(bit_count);
3264+
try r.ensureAliasAwareCapacity(needed, aliased);
32023265
var m = r.toMutable();
32033266
m.shiftLeftSat(a.toConst(), shift, signedness, bit_count);
32043267
r.setMetadata(m.positive, m.len);
@@ -3220,7 +3283,9 @@ pub const Managed = struct {
32203283
return;
32213284
}
32223285

3223-
try r.ensureCapacity(a.len() - (shift / limb_bits));
3286+
const aliased = limbsAliasDistinct(r, a);
3287+
const needed = a.len() - (shift / limb_bits);
3288+
try r.ensureAliasAwareCapacity(needed, aliased);
32243289
var m = r.toMutable();
32253290
m.shiftRight(a.toConst(), shift);
32263291
r.setMetadata(m.positive, m.len);
@@ -3229,7 +3294,9 @@ pub const Managed = struct {
32293294
/// r = ~a under 2s-complement wrapping semantics.
32303295
/// r and a may alias.
32313296
pub fn bitNotWrap(r: *Managed, a: *const Managed, signedness: Signedness, bit_count: usize) !void {
3232-
try r.ensureTwosCompCapacity(bit_count);
3297+
const aliased = limbsAliasDistinct(r, a);
3298+
const needed = calcTwosCompLimbCount(bit_count);
3299+
try r.ensureAliasAwareCapacity(needed, aliased);
32333300
var m = r.toMutable();
32343301
m.bitNotWrap(a.toConst(), signedness, bit_count);
32353302
r.setMetadata(m.positive, m.len);
@@ -3239,7 +3306,9 @@ pub const Managed = struct {
32393306
///
32403307
/// a and b are zero-extended to the longer of a or b.
32413308
pub fn bitOr(r: *Managed, a: *const Managed, b: *const Managed) !void {
3242-
try r.ensureCapacity(@max(a.len(), b.len()));
3309+
const aliased = limbsAliasDistinct(r, a) or limbsAliasDistinct(r, b);
3310+
const needed = @max(a.len(), b.len());
3311+
try r.ensureAliasAwareCapacity(needed, aliased);
32433312
var m = r.toMutable();
32443313
m.bitOr(a.toConst(), b.toConst());
32453314
r.setMetadata(m.positive, m.len);
@@ -3251,7 +3320,8 @@ pub const Managed = struct {
32513320
if (b.isPositive()) b.len() else if (a.isPositive()) a.len() else a.len() + 1
32523321
else if (a.isPositive()) a.len() else if (b.isPositive()) b.len() else b.len() + 1;
32533322

3254-
try r.ensureCapacity(cap);
3323+
const aliased = limbsAliasDistinct(r, a) or limbsAliasDistinct(r, b);
3324+
try r.ensureAliasAwareCapacity(cap, aliased);
32553325
var m = r.toMutable();
32563326
m.bitAnd(a.toConst(), b.toConst());
32573327
r.setMetadata(m.positive, m.len);
@@ -3260,7 +3330,8 @@ pub const Managed = struct {
32603330
/// r = a ^ b
32613331
pub fn bitXor(r: *Managed, a: *const Managed, b: *const Managed) !void {
32623332
const cap = @max(a.len(), b.len()) + @intFromBool(a.isPositive() != b.isPositive());
3263-
try r.ensureCapacity(cap);
3333+
const aliased = limbsAliasDistinct(r, a) or limbsAliasDistinct(r, b);
3334+
try r.ensureAliasAwareCapacity(cap, aliased);
32643335

32653336
var m = r.toMutable();
32663337
m.bitXor(a.toConst(), b.toConst());
@@ -3272,7 +3343,9 @@ pub const Managed = struct {
32723343
///
32733344
/// rma's allocator is used for temporary storage to boost multiplication performance.
32743345
pub fn gcd(rma: *Managed, x: *const Managed, y: *const Managed) !void {
3275-
try rma.ensureCapacity(@min(x.len(), y.len()));
3346+
const aliased = limbsAliasDistinct(rma, x) or limbsAliasDistinct(rma, y);
3347+
const needed = @min(x.len(), y.len());
3348+
try rma.ensureAliasAwareCapacity(needed, aliased);
32763349
var m = rma.toMutable();
32773350
var limbs_buffer = std.array_list.Managed(Limb).init(rma.allocator);
32783351
defer limbs_buffer.deinit();
@@ -3350,15 +3423,19 @@ pub const Managed = struct {
33503423

33513424
/// r = truncate(Int(signedness, bit_count), a)
33523425
pub fn truncate(r: *Managed, a: *const Managed, signedness: Signedness, bit_count: usize) !void {
3353-
try r.ensureCapacity(calcTwosCompLimbCount(bit_count));
3426+
const aliased = limbsAliasDistinct(r, a);
3427+
const needed = calcTwosCompLimbCount(bit_count);
3428+
try r.ensureAliasAwareCapacity(needed, aliased);
33543429
var m = r.toMutable();
33553430
m.truncate(a.toConst(), signedness, bit_count);
33563431
r.setMetadata(m.positive, m.len);
33573432
}
33583433

33593434
/// r = saturate(Int(signedness, bit_count), a)
33603435
pub fn saturate(r: *Managed, a: *const Managed, signedness: Signedness, bit_count: usize) !void {
3361-
try r.ensureCapacity(calcTwosCompLimbCount(bit_count));
3436+
const aliased = limbsAliasDistinct(r, a);
3437+
const needed = calcTwosCompLimbCount(bit_count);
3438+
try r.ensureAliasAwareCapacity(needed, aliased);
33623439
var m = r.toMutable();
33633440
m.saturate(a.toConst(), signedness, bit_count);
33643441
r.setMetadata(m.positive, m.len);
@@ -3367,7 +3444,9 @@ pub const Managed = struct {
33673444
/// r = @popCount(a) with 2s-complement semantics.
33683445
/// r and a may be aliases.
33693446
pub fn popCount(r: *Managed, a: *const Managed, bit_count: usize) !void {
3370-
try r.ensureCapacity(calcTwosCompLimbCount(bit_count));
3447+
const aliased = limbsAliasDistinct(r, a);
3448+
const needed = calcTwosCompLimbCount(bit_count);
3449+
try r.ensureAliasAwareCapacity(needed, aliased);
33713450
var m = r.toMutable();
33723451
m.popCount(a.toConst(), bit_count);
33733452
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)