fix(framework) Replace insecure SecAgg+ pseudo_rand_gen with SHA-256 PRF#7570
fix(framework) Replace insecure SecAgg+ pseudo_rand_gen with SHA-256 PRF#7570saurabhvmagdum wants to merge 6 commits into
Conversation
Replace weak seed folding + numpy.RandomState in pseudo_rand_gen with a SHA-256 counter-mode PRF that preserves full seed entropy and deterministically emits bytes. Enforce num_range to be a positive power-of-two, support scalar shapes, and return int64 arrays with exact shapes. Add comprehensive pytest coverage for determinism, entropy sensitivity, value range, power-of-two validation, scalar handling, and shape fidelity (framework/py/flwr/common/secure_aggregation/secaggplus_utils.py and new secaggplus_utils_test.py). Also add a brief AGENTS.md note recommending running pytest via the project's uv environment.
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 5ab6c7a510
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
There was a problem hiding this comment.
Pull request overview
This PR hardens SecAgg+ mask generation by replacing the existing pseudo_rand_gen (which compressed entropy and used NumPy’s non-cryptographic PRNG) with a deterministic SHA-256 counter-mode construction, and adds unit tests for determinism/range/shape behavior.
Changes:
- Implement SHA-256 counter-mode mask generation in
pseudo_rand_genwhile preserving the public signature/return types. - Add a new pytest module covering determinism, range bounds, shape fidelity, and invalid
num_range. - Extend
AGENTS.mdwith guidance for running Python/tests viauv.
Critical issues
- The new implementation generates masks via nested Python loops per element/per byte;
pseudo_rand_genis called on the hot path for every model tensor in SecAgg+ and this risks severe runtime regression compared to the previous vectorizedrandintapproach (see stored comment ID: 004).
Simplicity/readability suggestions
- Minor docstring cleanup (remove trailing whitespace / reflow) to avoid linter noise and match existing style (comment ID: 003).
Consistency concerns
- New test file should follow repository conventions (license header + module docstring, relative imports) as used across existing test modules (comment ID: 005).
Whether the PR should be split
- No, the changes are tightly scoped around securing and validating mask generation.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 6 comments.
| File | Description |
|---|---|
| framework/py/flwr/common/secure_aggregation/secaggplus_utils.py | Replaces insecure PRNG-based mask generation with SHA-256 counter-mode stream generation. |
| framework/py/flwr/common/secure_aggregation/secaggplus_utils_test.py | Adds tests for determinism, entropy sensitivity, range bounds, and shape/dtype fidelity. |
| AGENTS.md | Documents running Python/tests via uv from framework/. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
- Refactored `pseudo_rand_gen` to generate bytes in bulk and construct integer masks using `np.frombuffer` and vectorized bitwise shifts. This removes the per-element and per-byte Python loops, significantly improving generation performance for large tensors. - Added missing Apache license header to `secaggplus_utils_test.py` per repository style. - Maintained exact PRF cryptographic determinism and entropy behavior.
github-actions
Bot
added
the
Contributor
Switch pseudo-random number generator from naive hash concatenation to HMAC-based construction for improved cryptographic strength. Update test seed to 32 bytes to properly align with SHA-256. Improve code clarity and convert test imports to relative paths.
| import hashlib | ||
| import hmac | ||
| import struct | ||
| from collections.abc import Iterator | ||
|
|
||
| import numpy as np |
| """Seeded pseudo-random number generator for noise generation. | ||
|
|
||
| Uses SHA-256 in counter mode to generate a cryptographically strong, | ||
| deterministic byte stream from the seed, preserving full entropy. | ||
|
|
||
| Assumes `num_range` is a power of two. | ||
| """ |
| total_elements = sum(int(np.prod(shape)) for shape in dimensions_list) | ||
| total_bytes = total_elements * num_bytes | ||
|
|
||
| counter = 0 | ||
| buffer = bytearray() | ||
| while len(buffer) < total_bytes: | ||
| h = hmac.new(seed, struct.pack("<Q", counter), hashlib.sha256) | ||
| buffer.extend(h.digest()) | ||
| counter += 1 | ||
|
|
||
| buffer = buffer[:total_bytes] | ||
|
|
| if (num_range & (num_range - 1)) != 0 or num_range <= 0: | ||
| raise ValueError("num_range must be a power of two.") | ||
|
|
||
| num_bytes = (num_range.bit_length() + 6) // 8 | ||
| bitmask = num_range - 1 |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: ab0fbbd390
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| total_bytes = total_elements * num_bytes | ||
|
|
||
| counter = 0 | ||
| buffer = bytearray() | ||
| while len(buffer) < total_bytes: |
There was a problem hiding this comment.
In normal SecAgg+ runs, dimensions_list is the full model shape list used for every private and pairwise mask (secaggplus_mod.py:468/477 and server unmasking), so this allocates one raw PRF byte buffer for the entire model before the returned int64 masks are built. The fresh issue here is the new bulk implementation materializes total_bytes all at once; with the default 2**32 range this is an extra 4 bytes per parameter (about 400 MB for 100M weights) on top of the int64 masks and quantized arrays, which can push large-model secure aggregation runs into OOM. Generate and consume the PRF stream per array or in bounded chunks instead.
Useful? React with 👍 / 👎.
Optimize pseudo_rand_gen to generate random bytes per tensor instead of all upfront. This improves memory efficiency and enables streaming. Also: - Remove unused Iterator import - Clarify HMAC-SHA256 usage in docstring - Update error message to reflect >= 2 requirement - Standardize test seed value
| if num_range < 2 or (num_range & (num_range - 1)) != 0: | ||
| raise ValueError("num_range must be a power of two and >= 2.") | ||
|
|
||
| num_bytes = (num_range.bit_length() + 6) // 8 | ||
| bitmask = num_range - 1 |
| buffer = bytearray() | ||
| while len(buffer) < tensor_bytes: | ||
| h = hmac.new(seed, struct.pack("<Q", counter), hashlib.sha256) | ||
| buffer.extend(h.digest()) | ||
| counter += 1 | ||
| buffer = buffer[:tensor_bytes] |
Optimize pseudo_rand_gen by moving the buffer outside the loop to persist pseudo-random data across multiple tensor generations, reducing waste. Also add validation to ensure num_range fits within int64 masks.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: c28953a3a6
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| tensor_buffer = stream_buffer[:tensor_bytes] | ||
| stream_buffer = stream_buffer[tensor_bytes:] |
There was a problem hiding this comment.
When a model contains a large single tensor, this still materializes tensor_bytes of PRF output and copies it again before building the returned int64 mask. I checked the SecAgg+ callers in secaggplus_mod.py and server unmasking; they pass full model shapes for every private and pairwise mask, so a 100M-element tensor with the default 2**32 modulus needs hundreds of MB of extra raw bytes per mask on top of the int64 mask and quantized arrays, which can OOM large secure-aggregation runs. The fresh evidence after the prior buffering comment is that the current per-tensor refactor still allocates an entire tensor buffer here; fill the output array in bounded chunks instead.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: dd2a311ea7
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| tensor_bytes = total_elements * num_bytes | ||
|
|
||
| while len(stream_buffer) < tensor_bytes: | ||
| h = hmac.new(seed, struct.pack("<Q", counter), hashlib.sha256) |
There was a problem hiding this comment.
During a rolling upgrade, or when a server/client pair differs only by patch version, this changes the deterministic mask produced from the same SecAgg+ seed without adding any protocol/version negotiation. I checked RuntimeVersionMetadata.check_compatibility; it accepts peers with the same major.minor version and also temporarily accepts peers with missing runtime metadata, so an upgraded participant can still run a round with an old participant, but their private masks and pairwise masks will no longer cancel and the aggregate will be corrupted. Please include a SecAgg+ PRF version in the protocol/config or reject peers that may still be using the old generator.
Useful? React with 👍 / 👎.
Issue
Description
The
pseudo_rand_genhelper inframework/py/flwr/common/secure_aggregation/secaggplus_utils.pycurrently suffers from two cryptographic weaknesses that weaken the privacy guarantees of SecAgg+ masks:rd_seedis folded into at most 32 bits of entropy, regardless of its original length.np.random.RandomState(Mersenne Twister), which is designed for statistical simulation and is vulnerable to state-recovery attacks. It is not a pseudorandom function and must not be used for cryptographic masking.Related issues/PRs
Fixes #7522.
Proposal
Explanation
This PR replaces the insecure
pseudo_rand_genimplementation with a deterministic, cryptographically secure PRF based on SHA-256 in counter mode.Key changes:
seed(e.g., ECDH shared key,rd_seed) is passed directly into the PRF without any XOR compression. All input entropy contributes to the output stream.np.random.RandomStateis replaced by ahashlib.sha256counter-mode generator. The seed is concatenated with an 8-byte little-endian counter and hashed iteratively to produce an infinite deterministic byte stream.SecAggPlusWorkflowenforcesnum_range(i.e.,modulus_range) to be a power of two (e.g.,2**32,2**22). The implementation leverages this by usingbitmask = num_range - 1to extract uniformly distributed values without modulo bias or rejection sampling.List[np.ndarray],dtype=np.int64) remain unchanged. Callers insecaggplus_workflow.pyandsecaggplus_mod.pyrequire no modification.Note: Because the underlying PRNG changes from Mersenne Twister to SHA-256, the mask values generated from the same seed will differ from the previous implementation. This is expected and required for security; old persisted masks cannot be mixed with new ones.
Checklist
#contributions)Any other comments?
framework/py/flwr/common/secure_aggregation/secaggplus_utils.pyframework/py/flwr/common/secure_aggregation/secaggplus_utils_test.pyhashlib,struct) andnumpy(already a project dependency).num_rangeis a power of two, which is validated by the existingSecAggPlusWorkflowconstructor.