package use_cases import ( "context" "encoding/json" "errors" "math/big" "time" "github.com/google/uuid" "github.com/rwadurian/mpc-system/pkg/crypto" "github.com/rwadurian/mpc-system/pkg/logger" "github.com/rwadurian/mpc-system/services/server-party/domain/entities" "github.com/rwadurian/mpc-system/services/server-party/domain/repositories" "go.uber.org/zap" ) var ( ErrKeygenFailed = errors.New("keygen failed") ErrKeygenTimeout = errors.New("keygen timeout") ErrInvalidSession = errors.New("invalid session") ErrShareSaveFailed = errors.New("failed to save share") ) // ParticipateKeygenInput contains input for keygen participation type ParticipateKeygenInput struct { SessionID uuid.UUID PartyID string JoinToken string } // ParticipateKeygenOutput contains output from keygen participation type ParticipateKeygenOutput struct { Success bool KeyShare *entities.PartyKeyShare PublicKey []byte } // SessionCoordinatorClient defines the interface for session coordinator communication type SessionCoordinatorClient interface { JoinSession(ctx context.Context, sessionID uuid.UUID, partyID, joinToken string) (*SessionInfo, error) ReportCompletion(ctx context.Context, sessionID uuid.UUID, partyID string, publicKey []byte) error } // MessageRouterClient defines the interface for message router communication type MessageRouterClient interface { RouteMessage(ctx context.Context, sessionID uuid.UUID, fromParty string, toParties []string, roundNumber int, payload []byte) error SubscribeMessages(ctx context.Context, sessionID uuid.UUID, partyID string) (<-chan *MPCMessage, error) } // SessionInfo contains session information from coordinator type SessionInfo struct { SessionID uuid.UUID SessionType string ThresholdN int ThresholdT int MessageHash []byte Participants []ParticipantInfo } // ParticipantInfo contains participant information type ParticipantInfo struct { PartyID string PartyIndex int } // MPCMessage represents an MPC message from the router type MPCMessage struct { FromParty string IsBroadcast bool RoundNumber int Payload []byte } // ParticipateKeygenUseCase handles keygen participation type ParticipateKeygenUseCase struct { keyShareRepo repositories.KeyShareRepository sessionClient SessionCoordinatorClient messageRouter MessageRouterClient cryptoService *crypto.CryptoService } // NewParticipateKeygenUseCase creates a new participate keygen use case func NewParticipateKeygenUseCase( keyShareRepo repositories.KeyShareRepository, sessionClient SessionCoordinatorClient, messageRouter MessageRouterClient, cryptoService *crypto.CryptoService, ) *ParticipateKeygenUseCase { return &ParticipateKeygenUseCase{ keyShareRepo: keyShareRepo, sessionClient: sessionClient, messageRouter: messageRouter, cryptoService: cryptoService, } } // Execute participates in a keygen session // Note: This is a simplified implementation. Real implementation would use tss-lib func (uc *ParticipateKeygenUseCase) Execute( ctx context.Context, input ParticipateKeygenInput, ) (*ParticipateKeygenOutput, error) { // 1. Join session via coordinator sessionInfo, err := uc.sessionClient.JoinSession(ctx, input.SessionID, input.PartyID, input.JoinToken) if err != nil { return nil, err } if sessionInfo.SessionType != "keygen" { return nil, ErrInvalidSession } // 2. Find self in participants var selfIndex int for _, p := range sessionInfo.Participants { if p.PartyID == input.PartyID { selfIndex = p.PartyIndex break } } // 3. Subscribe to messages msgChan, err := uc.messageRouter.SubscribeMessages(ctx, input.SessionID, input.PartyID) if err != nil { return nil, err } // 4. Run TSS Keygen protocol // This is a placeholder - real implementation would use tss-lib saveData, publicKey, err := uc.runKeygenProtocol( ctx, input.SessionID, input.PartyID, selfIndex, sessionInfo.Participants, sessionInfo.ThresholdN, sessionInfo.ThresholdT, msgChan, ) if err != nil { return nil, err } // 5. Encrypt and save the share encryptedShare, err := uc.cryptoService.EncryptShare(saveData, input.PartyID) if err != nil { return nil, err } keyShare := entities.NewPartyKeyShare( input.PartyID, selfIndex, input.SessionID, sessionInfo.ThresholdN, sessionInfo.ThresholdT, encryptedShare, publicKey, ) if err := uc.keyShareRepo.Save(ctx, keyShare); err != nil { return nil, ErrShareSaveFailed } // 6. Report completion to coordinator if err := uc.sessionClient.ReportCompletion(ctx, input.SessionID, input.PartyID, publicKey); err != nil { logger.Error("failed to report completion", zap.Error(err)) // Don't fail - share is saved } return &ParticipateKeygenOutput{ Success: true, KeyShare: keyShare, PublicKey: publicKey, }, nil } // runKeygenProtocol runs the TSS keygen protocol // This is a placeholder implementation func (uc *ParticipateKeygenUseCase) runKeygenProtocol( ctx context.Context, sessionID uuid.UUID, partyID string, selfIndex int, participants []ParticipantInfo, n, t int, msgChan <-chan *MPCMessage, ) ([]byte, []byte, error) { /* Real implementation would: 1. Create tss.PartyID list 2. Create tss.Parameters 3. Create keygen.LocalParty 4. Handle outgoing messages via messageRouter 5. Handle incoming messages from msgChan 6. Wait for keygen completion 7. Return LocalPartySaveData and ECDSAPub Example with tss-lib: parties := make([]*tss.PartyID, len(participants)) for i, p := range participants { parties[i] = tss.NewPartyID(p.PartyID, p.PartyID, big.NewInt(int64(p.PartyIndex))) } selfPartyID := parties[selfIndex] tssCtx := tss.NewPeerContext(parties) params := tss.NewParameters(tss.S256(), tssCtx, selfPartyID, n, t) outCh := make(chan tss.Message, n*10) endCh := make(chan keygen.LocalPartySaveData, 1) party := keygen.NewLocalParty(params, outCh, endCh) go handleOutgoingMessages(ctx, sessionID, partyID, outCh) go handleIncomingMessages(ctx, party, msgChan) party.Start() select { case saveData := <-endCh: return saveData.Bytes(), saveData.ECDSAPub.Bytes(), nil case <-time.After(10*time.Minute): return nil, nil, ErrKeygenTimeout } */ // Placeholder: Generate mock data for demonstration // In production, this would be real TSS keygen logger.Info("Running keygen protocol (placeholder)", zap.String("session_id", sessionID.String()), zap.String("party_id", partyID), zap.Int("self_index", selfIndex), zap.Int("n", n), zap.Int("t", t)) // Simulate keygen delay select { case <-ctx.Done(): return nil, nil, ctx.Err() case <-time.After(2 * time.Second): } // Generate placeholder data mockSaveData := map[string]interface{}{ "party_id": partyID, "party_index": selfIndex, "threshold_n": n, "threshold_t": t, "created_at": time.Now().Unix(), } saveDataBytes, _ := json.Marshal(mockSaveData) // Generate a placeholder public key (32 bytes) mockPublicKey := make([]byte, 65) // Uncompressed secp256k1 public key mockPublicKey[0] = 0x04 // Uncompressed prefix copy(mockPublicKey[1:], big.NewInt(int64(selfIndex+1)).Bytes()) return saveDataBytes, mockPublicKey, nil }