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

Python	Rust
dict[str, Any] + hand-written schema check	ClientProof struct, malformed payloads fail to deserialize
verify_client_proof(...) -> bool, reason only in logs	Result<(), Reject>, one typed variant per failure
snarkjs subprocess, 300 s timeout, temp files	ark-groth16 verification in-process
json.loads → numpy, dtype inferred	Tensor { dtype, shape, bytes }, dtype carried explicitly
per-round if tree returning False	total functions returning reasons the caller can act on

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:

Role	Class
Proof verification, generation	RustProofManager
FedJSCM aggregation	RustFedJscmAdapter (exported as FedJSCMAggregator)
Stability heuristic	StabilityMonitor (Rust)

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

This covers the networked runtime (securefl-server + the gRPC clients), 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.

pyproject.toml overrides securefl-core with { path = "rust/securefl-python", editable = true } so development and CI always run the code in this checkout. That means uv sync needs a Rust toolchain; installs from PyPI are unaffected.

.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.

gRPC (primary) vs REST

--grpc-bind (default 127.0.0.1:8081) serves the same round API over protobuf. Tensors are raw bytes. Measured on a ResNet18 state dict:

Encoding	Wire size	Encode
protobuf (bytes)	44.7 MB (1.00×)	0.02 s
REST (JSON + base64)	59.7 MB (1.33×)	0.11 s

So gRPC removes a third of the traffic and most of the serialization cost per client per round. The gRPC service also carries round orchestration:

• Register(client_id) — join the roster.
• Poll(client_id) — TRAIN (with the model) if the client is in this round's sample, WAIT otherwise. The sample is a deterministic rotation over the sorted roster, with the same fraction_fit/min_fit_clients knobs.
• Submit — verified like REST; when every assigned client has submitted the server aggregates immediately and publishes the model plus its proof, so no separate driver is needed. If a round has submissions but stalls, --round-timeout-secs (default 600, 0 disables) closes it with whatever arrived on the next Poll/Submit, so one dead client cannot block the deployment.

--tls-cert/--tls-key enable TLS; --client-ca additionally requires client certificates and uses the certificate CN as the client's identity, so a client cannot submit under another id.

State is in memory only: a restart drops the model, momentum and the roster, so the operator has to PutModel again. That is deliberate for a benchmark/demo server; persistence would need a checkpoint format for momentum as well.

REST (--bind) stays as the debug/compat front-end and needs no protobuf. It checks the bearer token but not certificate identity: for authenticated clients use gRPC.

Auth

Writes (PUT /v1/model, POST .../updates, POST .../aggregate) require Authorization: Bearer <token> when --token / SECUREFL_SERVER_TOKEN is set; reads stay open. Without a token, client_id is self-declared, so anyone who can reach the port can claim unlimited identities — proofs bind identity, not authorization. A round also accepts at most --max-clients updates (429 after that), because the aggregation circuit has a fixed number of slots.

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]') \
  --buffers <(echo '[2]') \
  --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), and --buffers the complement: those tensors are FedAvg'd directly instead of going through momentum. Use --learnable all when the model has no buffers. make install-server and SecureFlServer derive both from the model, so a Python deployment never writes these by hand.

The Python side is secure_fl.client_sdk.SecureFlGrpcClient (primary) and SecureFlHttpClient (REST), plus the round loop:

from secure_fl.client_sdk import SecureFlGrpcClient
from secure_fl.round_runner import SecureFlRoundRunner

client = SecureFlGrpcClient("127.0.0.1:8081", "c1", token=TOKEN)
runner = SecureFlRoundRunner(
    client=client,
    train_fn=train,      # (parameters, round_no) -> (updated, num_examples, proof)
    rounds=10,
)
runner.run()             # register, poll until assigned, train, prove, submit

Three interop scripts exercise the real runtime: interop/smoke_runtime.py starts the server through SecureFlServer and trains two real SecureFlClients with run_client (the deployment path that replaced Flower), then checks the round closed. The other two exercise the binary over real sockets: interop/smoke_http.py (REST + gRPC round trip, a 5.6 MB model, forged proof rejected, server proof verified) and interop/smoke_mtls.py (no certificate refused, a valid certificate refused for another client_id). Both need built circuit artifacts (make setup-zkp plus one proof run), so they are a developer check rather than a CI step.

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. Numerics would diverge and the paper benchmarks would not transfer.
• A framework's clientapi transport. securefl-server covers networked rounds; reimplementing a reverse-RPC layer 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.