saorsa_core/threshold/

frost.rs

// Copyright 2024 Saorsa Labs Limited
//
// This software is dual-licensed under:
// - GNU Affero General Public License v3.0 or later (AGPL-3.0-or-later)
// - Commercial License
//
// For AGPL-3.0 license, see LICENSE-AGPL-3.0
// For commercial licensing, contact: saorsalabs@gmail.com
//
// Unless required by applicable law or agreed to in writing, software
// distributed under these licenses is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

//! FROST (Flexible Round-Optimized Schnorr Threshold) signatures implementation
//!
//! Provides threshold signatures where t-of-n participants can create valid signatures

use super::{Result, ThresholdError};
use crate::quantum_crypto::types::*;
// use frost_ed25519 as frost; // Temporarily disabled
// Removed unused OsRng import
use serde::{Deserialize, Serialize};
use std::collections::HashMap;

/// FROST signing session
pub struct FrostSession {
    /// Session identifier
    pub session_id: [u8; 32],

    /// Message to be signed
    pub message: Vec<u8>,

    /// Threshold value
    pub threshold: u16,

    /// Signing commitments from participants
    pub commitments: HashMap<ParticipantId, SigningCommitments>,

    /// Signing shares from participants
    pub shares: HashMap<ParticipantId, SigningShare>,

    /// Group public key
    pub group_public_key: FrostGroupPublicKey,

    /// Session state
    pub state: SessionState,
}

/// Signing commitments from a participant
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SigningCommitments {
    pub hiding: Vec<u8>,
    pub binding: Vec<u8>,
}

/// Signing share from a participant
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct SigningShare {
    pub share: Vec<u8>,
}

/// FROST session state
#[derive(Debug, Clone, PartialEq)]
pub enum SessionState {
    /// Collecting commitments
    CollectingCommitments,

    /// Collecting shares
    CollectingShares,

    /// Ready to aggregate
    ReadyToAggregate,

    /// Completed
    Completed,

    /// Failed
    Failed(String),
}

/// FROST key generation result
pub struct KeyGenerationResult {
    /// Group public key
    pub group_public_key: FrostGroupPublicKey,

    /// Participant shares
    pub shares: HashMap<ParticipantId, ParticipantShare>,

    /// Public commitments for verification
    pub commitments: Vec<Vec<u8>>,
}

/// Participant's share in the group
#[derive(Clone)]
pub struct ParticipantShare {
    pub participant_id: ParticipantId,
    pub signing_share: Vec<u8>, // Placeholder for frost::keys::SigningShare
    pub verifying_share: Vec<u8>, // Placeholder for frost::keys::VerifyingShare
}

impl FrostSession {
    /// Create new FROST signing session
    pub fn new(message: Vec<u8>, threshold: u16, group_public_key: FrostGroupPublicKey) -> Self {
        Self {
            session_id: rand::random(),
            message,
            threshold,
            commitments: HashMap::new(),
            shares: HashMap::new(),
            group_public_key,
            state: SessionState::CollectingCommitments,
        }
    }

    /// Add signing commitments from a participant
    pub fn add_commitments(
        &mut self,
        participant_id: ParticipantId,
        commitments: SigningCommitments,
    ) -> Result<()> {
        if self.state != SessionState::CollectingCommitments {
            return Err(ThresholdError::InvalidShare(
                "Not in commitment collection phase".to_string(),
            ));
        }

        self.commitments.insert(participant_id, commitments);

        // Check if we have enough commitments
        if self.commitments.len() >= self.threshold as usize {
            self.state = SessionState::CollectingShares;
        }

        Ok(())
    }

    /// Add signing share from a participant
    pub fn add_share(&mut self, participant_id: ParticipantId, share: SigningShare) -> Result<()> {
        if self.state != SessionState::CollectingShares {
            return Err(ThresholdError::InvalidShare(
                "Not in share collection phase".to_string(),
            ));
        }

        // Verify participant provided commitments
        if !self.commitments.contains_key(&participant_id) {
            return Err(ThresholdError::InvalidShare(
                "Participant did not provide commitments".to_string(),
            ));
        }

        self.shares.insert(participant_id, share);

        // Check if we have enough shares
        if self.shares.len() >= self.threshold as usize {
            self.state = SessionState::ReadyToAggregate;
        }

        Ok(())
    }

    /// Aggregate shares into final signature
    pub fn aggregate(&mut self) -> Result<FrostSignature> {
        if self.state != SessionState::ReadyToAggregate {
            return Err(ThresholdError::AggregationFailed(
                "Not ready to aggregate".to_string(),
            ));
        }

        // Convert to FROST types and aggregate
        // This is a simplified version - actual implementation would use frost crate

        // For now, concatenate all shares as a simple aggregation
        let mut aggregated = Vec::new();
        for (participant_id, share) in &self.shares {
            aggregated.extend_from_slice(&participant_id.0.to_be_bytes());
            aggregated.extend_from_slice(&share.share);
        }

        self.state = SessionState::Completed;

        Ok(FrostSignature(aggregated))
    }

    /// Check if session is complete
    pub fn is_complete(&self) -> bool {
        matches!(self.state, SessionState::Completed)
    }

    /// Get session progress
    pub fn get_progress(&self) -> SessionProgress {
        SessionProgress {
            session_id: self.session_id,
            state: self.state.clone(),
            commitments_received: self.commitments.len() as u16,
            shares_received: self.shares.len() as u16,
            threshold: self.threshold,
        }
    }
}

/// Session progress information
#[derive(Debug, Clone)]
pub struct SessionProgress {
    pub session_id: [u8; 32],
    pub state: SessionState,
    pub commitments_received: u16,
    pub shares_received: u16,
    pub threshold: u16,
}

/// FROST coordinator for managing signing sessions
pub struct FrostCoordinator {
    /// Active signing sessions
    pub sessions: HashMap<[u8; 32], FrostSession>,

    /// Group information
    pub groups: HashMap<GroupId, GroupInfo>,
}

/// Group information for FROST
pub struct GroupInfo {
    pub group_public_key: FrostGroupPublicKey,
    pub threshold: u16,
    pub participants: Vec<ParticipantId>,
}

impl Default for FrostCoordinator {
    fn default() -> Self {
        Self::new()
    }
}

impl FrostCoordinator {
    /// Create new coordinator
    pub fn new() -> Self {
        Self {
            sessions: HashMap::new(),
            groups: HashMap::new(),
        }
    }

    /// Register a group
    pub fn register_group(
        &mut self,
        group_id: GroupId,
        group_public_key: FrostGroupPublicKey,
        threshold: u16,
        participants: Vec<ParticipantId>,
    ) {
        self.groups.insert(
            group_id,
            GroupInfo {
                group_public_key,
                threshold,
                participants,
            },
        );
    }

    /// Initiate signing session
    pub fn initiate_signing(&mut self, group_id: &GroupId, message: Vec<u8>) -> Result<[u8; 32]> {
        let group_info = self.groups.get(group_id).ok_or_else(|| {
            ThresholdError::GroupOperationFailed("Group not registered".to_string())
        })?;

        let session = FrostSession::new(
            message,
            group_info.threshold,
            group_info.group_public_key.clone(),
        );

        let session_id = session.session_id;
        self.sessions.insert(session_id, session);

        Ok(session_id)
    }

    /// Process signing commitment
    pub fn process_commitment(
        &mut self,
        session_id: &[u8; 32],
        participant_id: ParticipantId,
        commitments: SigningCommitments,
    ) -> Result<()> {
        let session = self
            .sessions
            .get_mut(session_id)
            .ok_or_else(|| ThresholdError::InvalidShare("Session not found".to_string()))?;

        session.add_commitments(participant_id, commitments)
    }

    /// Process signing share
    pub fn process_share(
        &mut self,
        session_id: &[u8; 32],
        participant_id: ParticipantId,
        share: SigningShare,
    ) -> Result<()> {
        let session = self
            .sessions
            .get_mut(session_id)
            .ok_or_else(|| ThresholdError::InvalidShare("Session not found".to_string()))?;

        session.add_share(participant_id, share)
    }

    /// Complete signing session
    pub fn complete_signing(&mut self, session_id: &[u8; 32]) -> Result<FrostSignature> {
        let session = self
            .sessions
            .get_mut(session_id)
            .ok_or_else(|| ThresholdError::AggregationFailed("Session not found".to_string()))?;

        let signature = session.aggregate()?;

        // Clean up completed session after a delay
        // In practice, would schedule cleanup

        Ok(signature)
    }

    /// Get session status
    pub fn get_session_status(&self, session_id: &[u8; 32]) -> Option<SessionProgress> {
        self.sessions.get(session_id).map(|s| s.get_progress())
    }

    /// Clean up old sessions
    pub fn cleanup_old_sessions(&mut self, _max_age: std::time::Duration) {
        let _now = std::time::SystemTime::now();

        self.sessions.retain(|_, session| {
            // In practice, would check session creation time
            !matches!(
                session.state,
                SessionState::Completed | SessionState::Failed(_)
            )
        });
    }
}

/// Generate FROST key shares for a group
pub async fn generate_key_shares(threshold: u16, participants: u16) -> Result<KeyGenerationResult> {
    if threshold > participants {
        return Err(ThresholdError::InvalidParameters(
            "Threshold cannot exceed participants".to_string(),
        ));
    }

    if threshold == 0 {
        return Err(ThresholdError::InvalidParameters(
            "Threshold must be at least 1".to_string(),
        ));
    }

    // Placeholder implementation (frost crate disabled)
    let mut participant_shares = HashMap::new();

    for i in 0..participants {
        let participant_id = ParticipantId(i);
        participant_shares.insert(
            participant_id.clone(),
            ParticipantShare {
                participant_id: participant_id.clone(),
                signing_share: vec![i as u8; 32],   // Placeholder
                verifying_share: vec![i as u8; 32], // Placeholder
            },
        );
    }

    // Placeholder group public key
    let group_public_key = FrostGroupPublicKey(vec![0; 32]);

    Ok(KeyGenerationResult {
        group_public_key,
        shares: participant_shares,
        commitments: vec![], // Would include actual commitments
    })
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_frost_session_lifecycle() {
        let message = b"Test message".to_vec();
        let group_key = FrostGroupPublicKey(vec![0; 32]);
        let mut session = FrostSession::new(message, 2, group_key);

        // Add commitments
        assert_eq!(session.state, SessionState::CollectingCommitments);

        session
            .add_commitments(
                ParticipantId(1),
                SigningCommitments {
                    hiding: vec![1; 32],
                    binding: vec![2; 32],
                },
            )
            .unwrap();

        session
            .add_commitments(
                ParticipantId(2),
                SigningCommitments {
                    hiding: vec![3; 32],
                    binding: vec![4; 32],
                },
            )
            .unwrap();

        // Should move to collecting shares
        assert_eq!(session.state, SessionState::CollectingShares);

        // Add shares
        session
            .add_share(ParticipantId(1), SigningShare { share: vec![5; 32] })
            .unwrap();

        session
            .add_share(ParticipantId(2), SigningShare { share: vec![6; 32] })
            .unwrap();

        // Should be ready to aggregate
        assert_eq!(session.state, SessionState::ReadyToAggregate);

        // Aggregate
        let _signature = session.aggregate().unwrap();
        assert!(session.is_complete());
    }

    #[tokio::test]
    async fn test_key_generation() {
        let result = generate_key_shares(2, 3).await.unwrap();

        assert_eq!(result.shares.len(), 3);
        assert!(!result.group_public_key.0.is_empty());
    }

    #[test]
    fn test_coordinator() {
        let mut coordinator = FrostCoordinator::new();
        let group_id = GroupId([1; 32]);
        let group_key = FrostGroupPublicKey(vec![0; 32]);

        // Register group
        coordinator.register_group(
            group_id.clone(),
            group_key,
            2,
            vec![ParticipantId(1), ParticipantId(2), ParticipantId(3)],
        );

        // Initiate signing
        let message = b"Test message".to_vec();
        let session_id = coordinator.initiate_signing(&group_id, message).unwrap();

        // Check session status
        let status = coordinator.get_session_status(&session_id).unwrap();
        assert_eq!(status.threshold, 2);
        assert_eq!(status.commitments_received, 0);
    }
}