Fix integer overflow vulnerability in BalanceTrajectory and ComputeSmall

[mhucka]

Internal Google security scans flagged a memory access vulnerability from a potential integer overflow (b/510431411). The vulnerability originates from the calculation (num_samples + num_threads - 1) / num_threads in both BalanceTrajectory (util_balance_trajectory.cc:78) and ComputeSmall (tfq_noisy_samples.cc:231). num_samples is an int32 provided by the user. Providing a value close to INT_MAX overflows the addition into a negative number when thread counts are added. The resulting negative p_reps trickles down to needed_samples <= 0, bypassing the core generation loop. Since tensorflow::allocate_output maps raw heap buffers without automatic initialization for numerical outputs, returning it unmodified results in a direct leakage of the uninitialized data.

This is admittedly unlikely to happen in realistic usage scenarios, but who knows, maybe someone could create a malicious scenario to force it.

The approach taken here is simply to do a bounds check on the value of num_samples.

[gemini-code-assist]

Code Review

This pull request updates the calculation of sample repetitions and trajectory balancing by switching from int to uint64_t to prevent potential overflows. Feedback indicates that using uint64_t introduces a risk of Denial of Service (DoS) if num_samples is negative, as the unsigned cast would create an extremely large loop count. It is recommended to use int64_t for these calculations and to implement validation to ensure num_samples is non-negative.

[mhucka]

/gemini review

[gemini-code-assist]

Code Review

This pull request improves input validation in tfq_noisy_samples.cc by replacing debug assertions with runtime checks for the number of inputs and adding a range check for num_samples. However, the review feedback indicates that the current upper bound check for num_samples is insufficient to prevent integer overflows in downstream calculations involving thread counts and circuit sizes. It is recommended to use int64_t for intermediate calculations, apply a safer ceiling division formula, or enforce a more conservative upper bound on num_samples to protect all downstream 32-bit calculations.

[gemini-code-assist]

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.