Skip to content

Fix integer overflow vulnerability in BalanceTrajectory and ComputeSmall #1075

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

Draft
mhucka wants to merge 5 commits into
base: master
Choose a base branch
from
+11 −1
Draft

Fix integer overflow vulnerability in BalanceTrajectory and ComputeSmall #1075

Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
ad80158
Use size_t or int64_t for trajectory size arithmetic
mhucka May 22, 2026
f598888
Format
mhucka May 22, 2026
991b01a
Change how this is done; just check input value
mhucka May 22, 2026
ec3b024
Check the number of inputs to `Compute()`
mhucka May 22, 2026
c5499e1
Merge branch 'master' into fix-noisy-samples
mhucka May 22, 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
12 changes: 11 additions & 1 deletion tensorflow_quantum/core/ops/noise/tfq_noisy_samples.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@ limitations under the License.

#include <stdlib.h>

#include <limits>
#include <string>

#include "../qsim/lib/channel.h"
Expand Down Expand Up @@ -60,7 +61,10 @@ class TfqNoisySamplesOp : public tensorflow::OpKernel {

void Compute(tensorflow::OpKernelContext* context) override {
// TODO (mbbrough): add more dimension checks for other inputs here.
DCHECK_EQ(4, context->num_inputs());
const int num_inputs = context->num_inputs();
OP_REQUIRES(context, num_inputs == 4,
tensorflow::errors::InvalidArgument(absl::StrCat(
"Expected 4 inputs, got ", num_inputs, " inputs.")));

// Parse to Program Proto and num_qubits.
std::vector<Program> programs;
Expand All @@ -79,6 +83,12 @@ class TfqNoisySamplesOp : public tensorflow::OpKernel {

int num_samples = 0;
OP_REQUIRES_OK(context, GetIndividualSample(context, &num_samples));
OP_REQUIRES(
context,
num_samples >= 0 && num_samples < std::numeric_limits<int>::max(),
Copy link
Copy Markdown
Contributor

@gemini-code-assist gemini-code-assist Bot May 22, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

security-high high

The check num_samples < std::numeric_limits<int>::max() is insufficient to prevent the integer overflow vulnerability described in the pull request. The expression (num_samples + num_threads - 1) / num_threads (used in ComputeSmall at line 241 and in BalanceTrajectory) can still overflow if num_samples is large but less than INT_MAX. For example, if num_threads is 128, any num_samples > INT_MAX - 127 will cause the addition to overflow into a negative value, leading to the memory leakage issue described.

Furthermore, other calculations such as 2 * num_samples * ncircuits.size() + 2 (line 167) and 4 * (num_samples * ncircuits.size() + num_threads) / num_threads (line 261) are also susceptible to overflow if num_samples is large.

A more robust fix would be to:

  1. Use int64_t for intermediate calculations involving num_samples in ComputeSmall and BalanceTrajectory.
  2. Use a safer ceiling division formula: num_samples / num_threads + (num_samples % num_threads != 0) which avoids the overflow-prone addition.
  3. Apply a more conservative upper bound on num_samples at this check (e.g., INT_MAX / 2) to protect all downstream 32-bit calculations.

tensorflow::errors::InvalidArgument(
absl::StrCat("num_samples must be between 0 and ",
std::numeric_limits<int>::max(), ".")));

// Construct qsim circuits.
std::vector<NoisyQsimCircuit> qsim_circuits(programs.size(),
Expand Down
Loading