use crate::error::{OptimError, Result};
use scirs2_core::ndarray::Array1;
use scirs2_core::numeric::Float;
use scirs2_core::random::{Random, Rng};
use sha2::{Digest, Sha256};
use std::collections::HashMap;
use std::fmt::Debug;
use std::marker::PhantomData;
use super::coordinator::{
HomomorphicCiphertext, MaliciousTolerance, PrivacyLevel, SecureAggregationResult,
};
use super::helpers::{
add_mod, ct_eq, evaluate_polynomial, field_to_value, hash_values, inv_mod, mul_mod, neg_mod,
os_seeded_rng, random_bytes, require_supported_security, sub_mod, unimplemented_homomorphic,
value_to_field, SecureRng, AGGREGATE_DOMAIN, COMMITMENT_DOMAIN, COMMITMENT_NONCE_LEN,
SHAMIR_PRIME, VALUE_DIGEST_DOMAIN, VERIFICATION_DOMAIN,
};
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CommunicationSecurity {
SemiHonest,
MaliciousAbort,
MaliciousGuaranteed,
}
#[derive(Debug, Clone)]
pub enum SMPCProtocolState {
Initialization,
Setup,
InputSharing,
Computation,
OutputReconstruction,
Completed,
Aborted(String),
}
#[derive(Debug, Clone)]
pub struct SMPCSecurityGuarantees {
pub protocol_variant: SMPCProtocol,
pub communication_security: CommunicationSecurity,
pub malicious_tolerance: usize,
pub privacy_level: PrivacyLevel,
pub completeness: bool,
pub soundness: bool,
pub limitations: Vec<String>,
}
pub struct ShamirSecretSharing<T: Float + Debug + Send + Sync + 'static> {
pub(super) threshold: usize,
pub(super) num_shares: usize,
pub(super) prime_field: u128,
pub(super) rng: SecureRng,
pub(super) _phantom: PhantomData<T>,
}
impl<T: Float + Debug + Send + Sync + 'static> ShamirSecretSharing<T> {
pub fn new(threshold: usize, num_shares: usize) -> Result<Self> {
Self::validate_parameters(threshold, num_shares)?;
Ok(Self {
threshold,
num_shares,
prime_field: SHAMIR_PRIME,
rng: os_seeded_rng(),
_phantom: PhantomData,
})
}
pub fn with_seed(threshold: usize, num_shares: usize, seed: u64) -> Result<Self> {
Self::validate_parameters(threshold, num_shares)?;
Ok(Self {
threshold,
num_shares,
prime_field: SHAMIR_PRIME,
rng: Random::seed(seed),
_phantom: PhantomData,
})
}
pub(super) fn validate_parameters(threshold: usize, num_shares: usize) -> Result<()> {
if threshold == 0 {
return Err(OptimError::InvalidConfig(
"Shamir threshold must be at least 1".to_string(),
));
}
if num_shares == 0 {
return Err(OptimError::InvalidConfig(
"Shamir requires at least one share".to_string(),
));
}
if threshold > num_shares {
return Err(OptimError::InvalidConfig(format!(
"Shamir threshold {threshold} exceeds the number of shares {num_shares}"
)));
}
Ok(())
}
pub fn threshold(&self) -> usize {
self.threshold
}
pub fn num_shares(&self) -> usize {
self.num_shares
}
pub fn prime_field(&self) -> u128 {
self.prime_field
}
pub fn share_secret(&mut self, secret: T) -> Result<Vec<Share>> {
let element = value_to_field(secret)?;
self.share_field_element(element)
}
pub fn share_field_element(&mut self, secret: u128) -> Result<Vec<Share>> {
let mut coefficients = Vec::with_capacity(self.threshold);
coefficients.push(secret % self.prime_field);
for _ in 1..self.threshold {
coefficients.push(self.random_field_element());
}
let mut shares = Vec::with_capacity(self.num_shares);
for index in 1..=self.num_shares {
let x = index as u128;
if x >= self.prime_field {
return Err(OptimError::InvalidConfig(
"number of shares exceeds the size of the field".to_string(),
));
}
shares.push(Share {
x: index,
y: evaluate_polynomial(&coefficients, x),
});
}
Ok(shares)
}
pub fn reconstruct_secret(&self, shares: &[Share]) -> Result<T> {
let element = self.reconstruct_field_element(shares)?;
field_to_value(element)
}
pub fn reconstruct_field_element(&self, shares: &[Share]) -> Result<u128> {
if shares.len() < self.threshold {
return Err(OptimError::InvalidConfig(format!(
"insufficient shares for reconstruction: got {}, need {}",
shares.len(),
self.threshold
)));
}
let used = &shares[..self.threshold];
for (i, share) in used.iter().enumerate() {
if share.x == 0 {
return Err(OptimError::InvalidConfig(
"share x-coordinate must not be zero".to_string(),
));
}
if used.iter().skip(i + 1).any(|other| other.x == share.x) {
return Err(OptimError::InvalidConfig(format!(
"duplicate share x-coordinate {}",
share.x
)));
}
}
let mut result = 0u128;
for (i, share) in used.iter().enumerate() {
let xi = (share.x as u128) % self.prime_field;
let mut numerator = 1u128;
let mut denominator = 1u128;
for (j, other) in used.iter().enumerate() {
if i == j {
continue;
}
let xj = (other.x as u128) % self.prime_field;
numerator = mul_mod(numerator, neg_mod(xj));
denominator = mul_mod(denominator, sub_mod(xi, xj));
}
let lagrange = mul_mod(numerator, inv_mod(denominator)?);
result = add_mod(result, mul_mod(share.y % self.prime_field, lagrange));
}
Ok(result)
}
pub(super) fn random_field_element(&mut self) -> u128 {
loop {
let high = self.rng.next_u64() as u128;
let low = self.rng.next_u64() as u128;
let candidate = ((high << 64) | low) & ((1u128 << 127) - 1);
if candidate < self.prime_field {
return candidate;
}
}
}
}
pub struct HomomorphicEngine<T: Float + Debug + Send + Sync + 'static> {
pub(super) params: HomomorphicParameters<T>,
pub(super) digest_key: Vec<u8>,
}
impl<T: Float + Debug + Send + Sync + 'static> HomomorphicEngine<T> {
pub fn new() -> Self {
let mut rng = os_seeded_rng();
Self {
params: HomomorphicParameters::new(),
digest_key: random_bytes(&mut rng, 32),
}
}
pub fn with_seed(seed: u64) -> Self {
let mut rng = Random::seed(seed);
Self {
params: HomomorphicParameters::new(),
digest_key: random_bytes(&mut rng, 32),
}
}
pub fn params(&self) -> &HomomorphicParameters<T> {
&self.params
}
pub fn digest_values(&self, data: &Array1<T>) -> Result<HomomorphicCiphertext<T>> {
let mut digests = Vec::with_capacity(data.len());
for &value in data.iter() {
digests.push(self.digest_value(value)?);
}
Ok(HomomorphicCiphertext {
data: digests,
params: self.params.clone(),
})
}
pub fn encrypt(&self, data: &Array1<T>) -> Result<HomomorphicCiphertext<T>> {
self.digest_values(data)
}
pub fn decrypt(&self, ciphertext: &HomomorphicCiphertext<T>) -> Result<Array1<T>> {
ciphertext.validate()?;
Err(unimplemented_homomorphic("decryption"))
}
pub fn add_encrypted(
&self,
a: &HomomorphicCiphertext<T>,
b: &HomomorphicCiphertext<T>,
) -> Result<HomomorphicCiphertext<T>> {
a.validate()?;
b.validate()?;
if a.data.len() != b.data.len() {
return Err(OptimError::DimensionMismatch(
"digest vectors have different lengths".to_string(),
));
}
Err(unimplemented_homomorphic("addition"))
}
pub(super) fn digest_value(&self, value: T) -> Result<Vec<u8>> {
let as_f64 = value.to_f64().ok_or_else(|| {
OptimError::InvalidConfig("value cannot be converted to f64 for digesting".to_string())
})?;
let mut hasher = Sha256::new();
hasher.update(VALUE_DIGEST_DOMAIN);
hasher.update((self.digest_key.len() as u64).to_le_bytes());
hasher.update(&self.digest_key);
hasher.update(as_f64.to_le_bytes());
Ok(hasher.finalize().to_vec())
}
}
pub struct CryptographicAggregator<T: Float + Debug + Send + Sync + 'static> {
pub(super) config: SMPCConfig,
pub(super) commitment_scheme: CommitmentScheme<T>,
pub(super) verification_params: VerificationParameters<T>,
pub(super) aggregation_proofs: Vec<AggregationProof<T>>,
}
impl<T: Float + Debug + Send + Sync + 'static + scirs2_core::ndarray::ScalarOperand>
CryptographicAggregator<T>
{
pub fn new(config: SMPCConfig) -> Self {
Self {
config,
commitment_scheme: CommitmentScheme::new(),
verification_params: VerificationParameters::new(),
aggregation_proofs: Vec::new(),
}
}
pub fn with_seed(config: SMPCConfig, seed: u64) -> Self {
Self {
config,
commitment_scheme: CommitmentScheme::with_seed(seed),
verification_params: VerificationParameters::with_seed(seed),
aggregation_proofs: Vec::new(),
}
}
pub fn aggregation_proofs(&self) -> &[AggregationProof<T>] {
&self.aggregation_proofs
}
pub fn verification_params(&self) -> &VerificationParameters<T> {
&self.verification_params
}
pub fn secure_aggregate(
&mut self,
participant_inputs: &HashMap<String, Array1<T>>,
participants: &HashMap<String, Participant>,
) -> Result<SecureAggregationResult<T>> {
require_supported_security(self.config.communication_security)?;
let honest_participants =
self.select_honest_participants(participant_inputs, participants)?;
let aggregate = self.aggregate_honest_inputs(participant_inputs, &honest_participants)?;
let mut commitments = HashMap::new();
for id in &honest_participants {
if let Some(commitment) = participants.get(id).and_then(|p| p.commitment.clone()) {
commitments.insert(id.clone(), commitment);
}
}
let proof = self.generate_aggregation_proof(&aggregate, &commitments)?;
Ok(SecureAggregationResult {
aggregate,
honest_participants,
proof,
security_level: self.config.communication_security,
})
}
pub(super) fn select_honest_participants(
&self,
inputs: &HashMap<String, Array1<T>>,
participants: &HashMap<String, Participant>,
) -> Result<Vec<String>> {
let mut ids: Vec<&String> = participants.keys().collect();
ids.sort();
let mut honest_participants = Vec::new();
for id in ids {
let participant = match participants.get(id) {
Some(participant) => participant,
None => continue,
};
let input = match inputs.get(id) {
Some(input) => input,
None => continue,
};
if self.verify_participant_honesty(participant, input)? {
honest_participants.push(id.clone());
}
}
if honest_participants.len() < self.config.threshold {
return Err(OptimError::InvalidConfig(format!(
"insufficient honest participants for secure aggregation: {} verified, {} required",
honest_participants.len(),
self.config.threshold
)));
}
Ok(honest_participants)
}
pub(super) fn aggregate_honest_inputs(
&self,
inputs: &HashMap<String, Array1<T>>,
honest_participants: &[String],
) -> Result<Array1<T>> {
let first_participant = honest_participants.first().ok_or_else(|| {
OptimError::InvalidConfig("no honest participants for aggregation".to_string())
})?;
let first_input = inputs.get(first_participant).ok_or_else(|| {
OptimError::InvalidConfig(format!("missing input for participant {first_participant}"))
})?;
let dimension = first_input.len();
let mut aggregate = Array1::zeros(dimension);
let mut count = 0usize;
for participant_id in honest_participants {
let input = inputs.get(participant_id).ok_or_else(|| {
OptimError::InvalidConfig(format!("missing input for participant {participant_id}"))
})?;
if input.len() != dimension {
return Err(OptimError::DimensionMismatch(format!(
"participant {} submitted {} values, expected {}",
participant_id,
input.len(),
dimension
)));
}
aggregate = aggregate + input;
count += 1;
}
let divisor = T::from(count).ok_or_else(|| {
OptimError::InvalidConfig("participant count is not representable in T".to_string())
})?;
if divisor == T::zero() {
return Err(OptimError::InvalidConfig(
"no honest participants for aggregation".to_string(),
));
}
Ok(aggregate / divisor)
}
pub(super) fn generate_aggregation_proof(
&mut self,
aggregate: &Array1<T>,
commitments: &HashMap<String, Vec<u8>>,
) -> Result<AggregationProof<T>> {
let proof = AggregationProof {
aggregate_digest: hash_values(AGGREGATE_DOMAIN, &[], aggregate)?,
participant_commitments: commitments.clone(),
verification_data: self
.verification_params
.generate_verification_data(aggregate)?,
timestamp: std::time::SystemTime::now(),
_phantom: PhantomData,
};
self.aggregation_proofs.push(proof.clone());
Ok(proof)
}
pub fn verify_participant_honesty(
&self,
participant: &Participant,
submitted_input: &Array1<T>,
) -> Result<bool> {
if participant.status != ParticipantStatus::Active {
return Ok(false);
}
let commitment = match &participant.commitment {
Some(commitment) => commitment,
None => return Ok(false),
};
let nonce = match &participant.commitment_nonce {
Some(nonce) => nonce,
None => return Ok(false),
};
if !self
.commitment_scheme
.open(commitment, submitted_input, nonce)?
{
return Ok(false);
}
Ok(participant.trust_score >= self.config.malicious_tolerance.verification_threshold)
}
}
#[derive(Debug, Clone)]
pub struct SMPCConfig {
pub num_participants: usize,
pub threshold: usize,
pub security_parameter: usize,
pub enable_homomorphic: bool,
pub enable_zk_proofs: bool,
pub protocol_variant: SMPCProtocol,
pub communication_security: CommunicationSecurity,
pub malicious_tolerance: MaliciousTolerance,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Share {
pub x: usize,
pub y: u128,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SMPCProtocol {
BGW,
GMW,
SPDZ,
ABY,
FederatedSMPC,
}
pub struct CommitmentScheme<T: Float + Debug + Send + Sync + 'static> {
pub(super) rng: SecureRng,
pub(super) _phantom: PhantomData<T>,
}
impl<T: Float + Debug + Send + Sync + 'static> CommitmentScheme<T> {
pub fn new() -> Self {
Self {
rng: os_seeded_rng(),
_phantom: PhantomData,
}
}
pub fn with_seed(seed: u64) -> Self {
Self {
rng: Random::seed(seed),
_phantom: PhantomData,
}
}
pub fn commit(&mut self, value: &Array1<T>) -> Result<(Vec<u8>, CommitmentNonce)> {
let mut nonce_bytes = [0u8; COMMITMENT_NONCE_LEN];
let random = random_bytes(&mut self.rng, COMMITMENT_NONCE_LEN);
nonce_bytes.copy_from_slice(&random);
let nonce = CommitmentNonce(nonce_bytes);
let commitment = hash_values(COMMITMENT_DOMAIN, nonce.as_bytes(), value)?;
Ok((commitment, nonce))
}
pub fn open(
&self,
commitment: &[u8],
value: &Array1<T>,
nonce: &CommitmentNonce,
) -> Result<bool> {
let expected = hash_values(COMMITMENT_DOMAIN, nonce.as_bytes(), value)?;
Ok(ct_eq(commitment, &expected))
}
}
#[derive(Debug, Clone)]
pub struct Participant {
pub id: String,
pub public_key: Vec<u8>,
pub status: ParticipantStatus,
pub trust_score: f64,
pub commitment: Option<Vec<u8>>,
pub commitment_nonce: Option<CommitmentNonce>,
}
impl Participant {
pub fn new(id: impl Into<String>, public_key: Vec<u8>, trust_score: f64) -> Self {
Self {
id: id.into(),
public_key,
status: ParticipantStatus::Active,
trust_score,
commitment: None,
commitment_nonce: None,
}
}
pub fn with_commitment(mut self, commitment: Vec<u8>, nonce: CommitmentNonce) -> Self {
self.commitment = Some(commitment);
self.commitment_nonce = Some(nonce);
self
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ParticipantStatus {
Active,
Unavailable,
Suspicious,
Malicious,
}
pub struct VerificationParameters<T: Float + Debug + Send + Sync + 'static> {
pub(super) verification_key: Vec<u8>,
pub(super) _phantom: PhantomData<T>,
}
impl<T: Float + Debug + Send + Sync + 'static> VerificationParameters<T> {
pub fn new() -> Self {
let mut rng = os_seeded_rng();
Self {
verification_key: random_bytes(&mut rng, 64),
_phantom: PhantomData,
}
}
pub fn with_seed(seed: u64) -> Self {
let mut rng = Random::seed(seed);
Self {
verification_key: random_bytes(&mut rng, 64),
_phantom: PhantomData,
}
}
pub fn generate_verification_data(&self, aggregate: &Array1<T>) -> Result<Vec<u8>> {
hash_values(VERIFICATION_DOMAIN, &self.verification_key, aggregate)
}
pub fn verify_verification_data(&self, aggregate: &Array1<T>, tag: &[u8]) -> Result<bool> {
let expected = self.generate_verification_data(aggregate)?;
Ok(ct_eq(tag, &expected))
}
}
#[derive(Debug, Clone)]
pub struct HomomorphicParameters<T: Float + Debug + Send + Sync + 'static> {
pub security_level: usize,
pub modulus: u128,
pub(super) _phantom: PhantomData<T>,
}
impl<T: Float + Debug + Send + Sync + 'static> HomomorphicParameters<T> {
pub fn new() -> Self {
Self {
security_level: 128,
modulus: u64::MAX as u128,
_phantom: PhantomData,
}
}
}
#[derive(Clone, PartialEq, Eq)]
pub struct CommitmentNonce(pub(super) [u8; COMMITMENT_NONCE_LEN]);
impl CommitmentNonce {
pub fn from_bytes(bytes: [u8; COMMITMENT_NONCE_LEN]) -> Self {
Self(bytes)
}
pub fn as_bytes(&self) -> &[u8; COMMITMENT_NONCE_LEN] {
&self.0
}
}
#[derive(Debug, Clone)]
pub struct AggregationProof<T: Float + Debug + Send + Sync + 'static> {
pub aggregate_digest: Vec<u8>,
pub participant_commitments: HashMap<String, Vec<u8>>,
pub verification_data: Vec<u8>,
pub timestamp: std::time::SystemTime,
pub(super) _phantom: PhantomData<T>,
}
impl Debug for CommitmentNonce {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str("CommitmentNonce(<redacted>)")
}
}
impl<T: Float + Debug + Send + Sync + 'static> Default for CommitmentScheme<T> {
fn default() -> Self {
Self::new()
}
}
impl<T: Float + Debug + Send + Sync + 'static> Default for HomomorphicEngine<T> {
fn default() -> Self {
Self::new()
}
}
impl<T: Float + Debug + Send + Sync + 'static> Default for HomomorphicParameters<T> {
fn default() -> Self {
Self::new()
}
}
impl<T: Float + Debug + Send + Sync + 'static> Default for VerificationParameters<T> {
fn default() -> Self {
Self::new()
}
}