# 🔐 Secure FL: Zero-Knowledge Federated Learning

A dual-verifiable framework for federated learning using zero-knowledge proofs to ensure training integrity and aggregation correctness.

## 🎯 Core Features

- **Dual ZKP Verification**: Client-side zk-STARKs + Server-side zk-SNARKs
- **FedJSCM Aggregation**: Momentum-based federated optimization
- **Dynamic Proof Rigor**: Adaptive proof complexity based on training stability
- **Parameter Quantization**: ZKP-compatible weight compression

## 🏗️ Architecture

```
Client Training + zk-STARK Proof → FL Server + zk-SNARK Proof → Verified Model
```

The system provides **dual verification**:
1. **Clients** generate zk-STARK proofs of correct local training
2. **Server** generates zk-SNARK proofs of correct aggregation

## 🚀 Quick Start

### Installation

```bash
# Install the package with uv (recommended)
uv pip install secure-fl

# Or install from source with uv
git clone https://github.com/krishantt/secure-fl
cd secure-fl
uv pip install -e .

# For development with all dependencies
uv sync --all-extras
```

### ZKP Prerequisites

Install zero-knowledge proof tools:

```bash
# Automated setup with make (recommended)
make setup-zkp

# Or manual setup:
# 1. Install Rust
curl --proto '=https' --tlsv1.2 https://sh.rustup.rs -sSf | sh

# 2. Install Circom
git clone https://github.com/iden3/circom.git
cd circom && cargo install --path circom

# 3. Install SnarkJS
npm install -g snarkjs

# Verify setup
uv run secure-fl check-zkp
```

### Basic Usage

#### Server

```python
from secure_fl import SecureFlowerServer, create_server_strategy
import torch.nn as nn

# Define model
class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = nn.Linear(784, 10)
    
    def forward(self, x):
        return self.fc(x.view(-1, 784))

# Create server with ZKP verification
strategy = create_server_strategy(
    model_fn=SimpleModel,
    enable_zkp=True,
    proof_rigor="high"
)

server = SecureFlowerServer(strategy=strategy)
server.start(num_rounds=10)
```

### Configuration

Create and use a configuration file:

```bash
# Create example config
uv run secure-fl create-config

# Edit config.yaml as needed
# Then use it:
```

#### Client

```python
from secure_fl import create_client, start_client
from torchvision import datasets, transforms

# Load data
transform = transforms.Compose([transforms.ToTensor()])
dataset = datasets.MNIST('./data', train=True, transform=transform)

# Create secure client
client = create_client(
    client_id="client_1",
    model_fn=SimpleModel,
    train_data=dataset,
    enable_zkp=True
)

# Connect to server
start_client(client, "localhost:8080")
```

#### CLI Interface

```bash
# Start server
uv run secure-fl-server --config config.yaml

# Start client  
uv run secure-fl-client --server localhost:8080 --dataset mnist --client-id client_1

# Check system status
uv run secure-fl check-zkp
```

## 🔬 Technical Details

### Zero-Knowledge Proofs

- **Client-side (zk-STARKs)**: Prove correct SGD computation using Cairo circuits
- **Server-side (zk-SNARKs)**: Prove correct FedJSCM aggregation using Circom circuits

### FedJSCM Aggregation

Momentum-based federated averaging:
```
w_{t+1} = w_t - η_g * (β * m_t + (1-β) * ∇F_t)
```
where `∇F_t` is the federated gradient and `m_t` is the momentum buffer.

### Dynamic Proof Rigor

Automatically adjusts ZKP complexity based on training stability:
- **High stability**: Reduced proof complexity for efficiency
- **Low stability**: Increased proof rigor for security

## 📊 Configuration

Create a `config.yaml`:

```yaml
server:
  host: "localhost"
  port: 8080
  num_rounds: 10

strategy:
  min_fit_clients: 2
  fraction_fit: 1.0
  momentum: 0.9

zkp:
  enable_zkp: true
  proof_rigor: "high"
  quantize_weights: true
  quantization_bits: 8
```

## 🔧 Development

### Setup Development Environment

```bash
git clone https://github.com/krishantt/secure-fl
cd secure-fl

# Complete development setup
make dev

# Or manually with uv
uv sync --all-extras
make setup-zkp
```

### Development Commands

```bash
# Run tests
make test
make test-quick      # Fast tests with early exit
make test-cov        # With coverage report

# Code quality
make lint           # Check with ruff
make format         # Format code
make type-check     # Run mypy
make check          # All quality checks

# Development workflow
make demo           # Run demonstration
make clean          # Clean artifacts
```

### Docker (Minimal)

```bash
# Build image
docker build -t secure-fl:local .

# Start local FL stack
docker compose up -d

# Stop services
docker compose down
```

## 📈 Experiments

Run benchmarks and experiments:

```bash
# Basic demo
make demo
# or: uv run python experiments/demo.py

# Reproducible benchmark suite
uv run python experiments/canonical_benchmark.py --datasets mnist synthetic_small --num-repeats 5 --require-real-proofs

# Custom training
uv run python experiments/train.py --config experiments/config.yaml

# Check environment
make env-info
```

## 🏷️ Repository Structure

```
secure-fl/
├── src/secure_fl/       # Main package
│   ├── federation/      # FL clients, server, and aggregation strategy
│   ├── zkp/             # ZKP managers and quantization
│   ├── models/          # Model definitions (MNIST/CIFAR/ResNet/MLP)
│   ├── core/            # Types, config, exceptions, versioning
│   ├── cli/             # CLI commands and setup helpers
│   └── utils/           # Logging and helper utilities
│   ├── proofs/          # ZKP circuits and proving assets
│   │   ├── client_circuits/ # zk-STARK/PySNARK client circuits
│   │   └── server/      # zk-SNARK (Circom) server circuits
├── experiments/         # Research experiments
├── tests/               # Test suite
└── docs/                # Documentation
```

## 🤝 Contributing

1. Fork the repository
2. Set up development environment: `make dev`  
3. Create a feature branch
4. Make your changes with proper type hints
5. Add tests and ensure coverage
6. Run quality checks: `make check`
7. Test your changes: `make test`
8. Submit a pull request

### Code Style

- Use type hints throughout
- Follow the established error handling patterns
- Add proper logging with context
- Write tests for new functionality
- Update documentation as needed

## 📄 License

MIT License - see [LICENSE](LICENSE) for details.

## 📚 Citation

```bibtex
@misc{timilsina2024secure,
  title={Secure FL: Dual-Verifiable Framework for Federated Learning using Zero-Knowledge Proofs},
  author={Timilsina, Krishant and Paudel, Bindu},
  year={2024},
  url={https://github.com/krishantt/secure-fl}
}
```

## 🙏 Acknowledgments

- Flower framework for federated learning infrastructure
- Circom and Cairo for zero-knowledge proof systems
- The federated learning and cryptography research communities
