Skip to content

Commit 6314da4

Browse files
falseywinchnetfalseywinchnet
authored
Merge pull request #71 from falseywinchnet/7rn4r3-main/review-core-library-for-bugs-and-vulnerabilities
Harden allocation size checks
2 parents 5b11547 + c4da00a commit 6314da4

2 files changed

Lines changed: 76 additions & 38 deletions

File tree

‎src/detail/bruun_kernel.hpp‎

Lines changed: 18 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -236,7 +236,14 @@ class heap_array {
236236

237237
bool push_back(const T& val) {
238238
if (len_ >= cap_) {
239-
if (!reserve(cap_ == 0 ? 16 : cap_ * 2)) return false;
239+
std::size_t new_cap = 16;
240+
if (cap_ != 0) {
241+
if (cap_ > static_cast<std::size_t>(-1) / 2) {
242+
return false;
243+
}
244+
new_cap = cap_ * 2;
245+
}
246+
if (!reserve(new_cap)) return false;
240247
}
241248
new (ptr_ + len_) T(val);
242249
++len_;
@@ -272,10 +279,19 @@ class heap_array {
272279

273280
private:
274281
static void* aligned_alloc_raw(std::size_t count) {
282+
if (count > static_cast<std::size_t>(-1) / sizeof(T)) {
283+
return nullptr;
284+
}
275285
const std::size_t bytes = sizeof(T) * count;
276286
std::size_t padded = bytes;
277287
const std::size_t remainder = padded % bruun_cache_alignment;
278-
if (remainder != 0) padded += bruun_cache_alignment - remainder;
288+
if (remainder != 0) {
289+
const std::size_t pad = bruun_cache_alignment - remainder;
290+
if (padded > static_cast<std::size_t>(-1) - pad) {
291+
return nullptr;
292+
}
293+
padded += pad;
294+
}
279295
void* raw = nullptr;
280296
#if defined(_MSC_VER)
281297
raw = _aligned_malloc(padded, bruun_cache_alignment);

‎src/stft.cpp‎

Lines changed: 58 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,17 @@ bool is_power_of_two(size_t n) {
4444
return n > 0 && (n & (n - 1)) == 0;
4545
}
4646

47+
bool add_overflows_size(size_t a, size_t b) {
48+
return a > std::numeric_limits<size_t>::max() - b;
49+
}
50+
51+
bool mul_overflows_size(size_t a, size_t b) {
52+
if (a == 0 || b == 0) {
53+
return false;
54+
}
55+
return a > std::numeric_limits<size_t>::max() / b;
56+
}
57+
4758
int reflect_index(long long i, size_t n) {
4859
if (n <= 1) return 0;
4960
const long long period = static_cast<long long>(2 * n - 2);
@@ -134,46 +145,57 @@ bfft_status bfft_stft_plan_create(size_t n, size_t n_fft, size_t hop_length, con
134145
if (hop_length == 0 || hop_length > n_fft || (n_fft % hop_length) != 0) return BFFT_ERROR_INVALID_ARGUMENT;
135146
if ((n % hop_length) != 0) return BFFT_ERROR_INVALID_ARGUMENT;
136147
if (transform != BFFT_STFT_RFFT && transform != BFFT_STFT_ODFT) return BFFT_ERROR_INVALID_ARGUMENT;
148+
if (n > static_cast<size_t>(std::numeric_limits<long long>::max())) return BFFT_ERROR_INVALID_ARGUMENT;
149+
if (n_fft > (static_cast<size_t>(std::numeric_limits<long long>::max()) + 2) / 2) return BFFT_ERROR_INVALID_ARGUMENT;
150+
if (n_fft > std::numeric_limits<int>::max()) return BFFT_ERROR_INVALID_ARGUMENT;
151+
if (add_overflows_size(n, n_fft - 1)) return BFFT_ERROR_INVALID_ARGUMENT;
152+
if (mul_overflows_size(n / hop_length, n_fft / 2 + 1)) return BFFT_ERROR_INVALID_ARGUMENT;
137153

138-
std::unique_ptr<bfft_stft_plan> p(new (std::nothrow) bfft_stft_plan());
139-
if (!p) return BFFT_ERROR_ALLOCATION;
140-
p->n = n;
141-
p->n_fft = n_fft;
142-
p->hop = hop_length;
143-
p->segments = n / hop_length;
144-
p->transform = transform;
145-
p->window.resize(n_fft);
146-
if (window) {
147-
for (size_t i = 0; i < n_fft; ++i) {
148-
if (!std::isfinite(window[i])) return BFFT_ERROR_INVALID_ARGUMENT;
149-
p->window[i] = window[i];
154+
try {
155+
std::unique_ptr<bfft_stft_plan> p(new (std::nothrow) bfft_stft_plan());
156+
if (!p) return BFFT_ERROR_ALLOCATION;
157+
p->n = n;
158+
p->n_fft = n_fft;
159+
p->hop = hop_length;
160+
p->segments = n / hop_length;
161+
p->transform = transform;
162+
p->window.resize(n_fft);
163+
if (window) {
164+
for (size_t i = 0; i < n_fft; ++i) {
165+
if (!std::isfinite(window[i])) return BFFT_ERROR_INVALID_ARGUMENT;
166+
p->window[i] = window[i];
167+
}
168+
} else {
169+
bfft_status st = bfft_stft_hann_window(n_fft, p->window.data());
170+
if (st != BFFT_OK) return st;
150171
}
151-
} else {
152-
bfft_status st = bfft_stft_hann_window(n_fft, p->window.data());
153-
if (st != BFFT_OK) return st;
154-
}
155-
p->analysis.resize(n_fft);
156-
ifftshift_copy(p->window, p->analysis);
157-
if (!synthesis_window(p->window, hop_length, p->synthesis)) return BFFT_ERROR_INVALID_ARGUMENT;
172+
p->analysis.resize(n_fft);
173+
ifftshift_copy(p->window, p->analysis);
174+
if (!synthesis_window(p->window, hop_length, p->synthesis)) return BFFT_ERROR_INVALID_ARGUMENT;
158175

159-
if (transform == BFFT_STFT_RFFT) {
160-
bfft_status st = bfft_plan_create(n_fft, &p->rfft);
161-
if (st != BFFT_OK) return st;
162-
p->bins = bfft_plan_bins(p->rfft);
163-
p->work.resize(bfft_plan_work_size(p->rfft));
164-
p->scratch.resize(bfft_plan_native_scratch_size(p->rfft));
165-
} else {
166-
bfft_status st = bodft_plan_create(n_fft, &p->odft);
167-
if (st != BFFT_OK) return st;
168-
p->bins = bodft_plan_bins(p->odft);
176+
if (transform == BFFT_STFT_RFFT) {
177+
bfft_status st = bfft_plan_create(n_fft, &p->rfft);
178+
if (st != BFFT_OK) return st;
179+
p->bins = bfft_plan_bins(p->rfft);
180+
p->work.resize(bfft_plan_work_size(p->rfft));
181+
p->scratch.resize(bfft_plan_native_scratch_size(p->rfft));
182+
} else {
183+
bfft_status st = bodft_plan_create(n_fft, &p->odft);
184+
if (st != BFFT_OK) return st;
185+
p->bins = bodft_plan_bins(p->odft);
186+
}
187+
p->buffer.assign(n_fft - hop_length, 0.0);
188+
p->xp.resize(n + n_fft - 1);
189+
p->segment.resize(n_fft);
190+
p->frame.resize(n_fft);
191+
p->processed.resize(n_fft);
192+
p->tmp_bins.resize(p->bins);
193+
*out_plan = p.release();
194+
} catch (const std::bad_alloc&) {
195+
return BFFT_ERROR_ALLOCATION;
196+
} catch (...) {
197+
return BFFT_ERROR_INTERNAL;
169198
}
170-
p->buffer.assign(n_fft - hop_length, 0.0);
171-
p->xp.resize(n + n_fft - 1);
172-
p->segment.resize(n_fft);
173-
p->frame.resize(n_fft);
174-
p->processed.resize(n_fft);
175-
p->tmp_bins.resize(p->bins);
176-
*out_plan = p.release();
177199
return BFFT_OK;
178200
}
179201

0 commit comments

Comments
 (0)