Secure-FL in Rust

A type-safe port of the Secure-FL verifier and aggregation server. It speaks the same proof format and parameter representation as the Python package, so a Python client trained with PyTorch can submit to this server unchanged.

Scope: proof schema, commitments, norm enforcement, FedJSCM aggregation, and Groth16 verification. Training stays in Python — reimplementing PyTorch numerics in Rust would invalidate the existing benchmark results.

What changes vs. the Python server

Layout

rust/
├── securefl-core/          # no IO, no async
│   ├── tensor.rs           # dtype-preserving tensor + numpy-compatible math
│   ├── commit.rs           # state_hash (byte-identical to compute_hash), l2_norm
│   ├── proof.rs            # typed client proof payload
│   ├── verify.rs           # verification pipeline -> Reject
│   ├── fedjscm.rs          # momentum aggregation + trimmed mean
│   └── groth16.rs          # snarkjs vk/proof JSON -> arkworks verification
├── securefl-prover/        # native Groth16 proving from the existing .zkey
├── securefl-python/        # PyO3 extension (securefl_rust)
├── securefl-server/        # standalone axum round server
├── interop/                # fixture generator, HTTP smoke test, benchmark
└── fixtures/               # generated parity fixtures (checked in)

Wiring into the FL pipeline

There is no backend switch. secure_fl/backend.py's factories always return the Rust implementations:

make build-rust        # builds and installs src/securefl_rust.so
uv run secure-fl-server --config experiments/config.yaml

This covers the live Flower strategy (federation/server.py), the paper benchmark pipeline (benchmarks/pipeline.py, benchmarks/execution.py) and the attack benchmarks (benchmarks/attacks.py). Clients and training stay in PyTorch, so results remain comparable.

The Rust core owns the whole proof path: client delta-bound proofs, server aggregation proofs, both verifiers, aggregation, and the stability heuristic. Circom/snarkjs remain build-time tools only — make setup-zkp produces the circuit artifacts; nothing shells out at runtime.

The Python implementations of all four were removed once parity was proven.

Install as a wheel

make build-rust-wheel                            # rust/dist/securefl_core-<platform>.whl
uv pip install rust/dist/securefl_core-*.whl     # or: pip install ./rust/securefl-python

The wheel installs a top-level securefl_rust extension module, so it shares no files with the secure-fl wheel. It is published as securefl-core; installing secure-fl pulls the matching wheel, and make build-rust remains the editable-install path when you are changing Rust code.

.github/workflows/wheels.yml (manual dispatch) builds all five targets — linux x86_64/aarch64, macOS arm64/x86_64, windows x64 — and uploads them as workflow artifacts.

release.yml publishes securefl-core alongside secure-fl on every tag, pinned to the same version, then installs both from PyPI in a clean venv before the GitHub release is created. It needs a one-time trusted publisher on PyPI:

1. Open https://pypi.org/manage/account/publishing/ and add a pending publisher with:
   • PyPI Project Name: securefl-core
   • Owner: krishantt
   • Repository name: secure-fl
   • Workflow name: release.yml
   • Environment name: release
2. Nothing else. The GitHub environment release already exists (it is used by the secure-fl publish job).

The publisher is configured; securefl-core publishes first, and secure-fl waits for it so a release never ships a wrapper without its extension. The wheels.yml workflow still provides downloadable artifacts for testing.

Both packages publish through PyPI trusted publishing (OIDC), so PyPI attaches a PEP 740 attestation to every file — see the "Provenance" tab of secure-fl and securefl-core.

GitHub build attestations (actions/attest-build-provenance) are off: they need a public repo or a paid plan. Once the repo is public, add the step back to both publish jobs in release.yml along with attestations: write.

Protocol

GET  /v1/health                 -> {"status":"ok"}
PUT  /v1/model                  {tensors}       bootstrap or reset
GET  /v1/model                  -> {round, tensors, accepted_clients}
POST /v1/rounds/{r}/updates     {client_id, num_examples, tensors, proof}
POST /v1/rounds/{r}/aggregate   -> {round, tensors, server_proof?}

With --circuit-dir set, server_proof carries the aggregation Groth16 proof ({type, version, proof, public}), generated in-process and verifiable against the circuit's verification_key.json.

A tensor is {"dtype": "float32", "shape": [4,3], "data": "<base64 raw bytes>"}. Accepting updates returns {"accepted": true}; rejecting returns 422 with a typed reason in {"error": "..."}. verify order matches the Python server: proof type/version → replay guards → full state-dict hash → delta hash → norm bound → Groth16 delta-bound proof.

Run

cargo test                      # parity + crypto + unit tests

cargo run -p securefl-server -- \
  --bind 127.0.0.1:8080 \
  --vk ../src/secure_fl/proofs/client/build/delta_bound_vk.json \
  --learnable <(echo '[0,1,3,4]') \
  --max-norm 10 \
  --circuit-dir ../src/secure_fl/proofs/server

--learnable takes a JSON array of state-dict indices that are learnable parameters (BatchNorm buffers excluded), matching what the Python server gets from its BatchNorm mask. Use --learnable all when the model has no buffers.

Interop tests

# Regenerate fixtures from the Python implementation, then assert parity
../.venv/bin/python interop/gen_fixtures.py
cargo test

# Real server binary, real HTTP, real Groth16 proof
cargo build -p securefl-server && ../.venv/bin/python interop/smoke_http.py

# Run the Python suite against the Rust backend
make build-rust
uv run pytest tests/unit tests/e2e --no-cov

Historical measurements (the Python implementations have since been removed): docs/dev/RUST_END_TO_END_QUICK.md and docs/dev/RUST_BACKEND_BENCHMARK.md.

gen_fixtures.py imports the Python compute_hash, compute_parameter_norm and FedJSCMAggregator, and produces a real snarkjs Groth16 delta-bound proof plus a tampered copy. The Rust tests assert byte-identical commitments, matching aggregation output, native verification of the real proof, and rejection of both tampered proof points and tampered public signals.

The prover tests re-prove the committed input.json and assert the public outputs equal the snarkjs public.json, and that a proof built from unsatisfied inputs does not verify. They need the circuit build directory (make setup-zkp) and skip without it.

Deployment notes

• No authentication. The server binds 127.0.0.1 by default; put it behind a reverse proxy with TLS and auth before exposing it.
• Request bodies are capped at 256MB; a ResNet18 state dict is ~45MB.
• Aggregation holds the round lock while proving, so a round blocks other requests for its duration (~30ms on the 400-public circuit).

Not ported (deliberately)

• Training. No burn/candle/tch rewrite; clients keep using PyTorch and Flower. Numerics would diverge and the paper benchmarks would not transfer.
• Flower clientapi transport. The standalone server covers a Python-free deployment; reimplementing Flower's reverse-RPC would move no round cost (verify → aggregate → prove is already Rust).
• Textual TUI demo, CLI, dataset loading, benchmarks.

shortcut: aggregation crosses the PyO3 boundary by copying each tensor, so it lands within ~5–10% of the numpy implementation it replaced. Remove the copies with a borrowed-tensor API if that gap ever matters.