From 2164664ca0d8dd298e682042d8f511a6173c080f Mon Sep 17 00:00:00 2001 From: hailin Date: Tue, 30 Dec 2025 00:43:09 -0800 Subject: [PATCH] feat(server-party): add ExecuteWithSessionInfo for co-managed keygen MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add new ExecuteWithSessionInfo method to ParticipateKeygenUseCase for server-party-co-managed to skip duplicate JoinSession call. - server-party-co-managed already calls JoinSession in session_created phase - ExecuteWithSessionInfo accepts pre-obtained SessionInfo and skips internal JoinSession - Refactor common execution logic to private executeWithSessionInfo method - Update server-party-co-managed to use ExecuteWithSessionInfo on session_started 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.5 --- .../cmd/server/main.go | 51 ++++++++++---- .../use_cases/participate_keygen.go | 69 ++++++++++++++----- 2 files changed, 90 insertions(+), 30 deletions(-) diff --git a/backend/mpc-system/services/server-party-co-managed/cmd/server/main.go b/backend/mpc-system/services/server-party-co-managed/cmd/server/main.go index 58ca8296..cef97a68 100644 --- a/backend/mpc-system/services/server-party-co-managed/cmd/server/main.go +++ b/backend/mpc-system/services/server-party-co-managed/cmd/server/main.go @@ -29,10 +29,13 @@ import ( // PendingSession stores session info between session_created and session_started events type PendingSession struct { - SessionID uuid.UUID - JoinToken string - MessageHash []byte - CreatedAt time.Time + SessionID uuid.UUID + JoinToken string + MessageHash []byte + ThresholdN int + ThresholdT int + SelectedParties []string + CreatedAt time.Time } // PendingSessionCache stores pending sessions waiting for session_started @@ -379,10 +382,13 @@ func createCoManagedSessionEventHandler( // Store pending session for later use when session_started arrives pendingSessionCache.Store(event.SessionId, &PendingSession{ - SessionID: sessionID, - JoinToken: joinToken, - MessageHash: event.MessageHash, - CreatedAt: time.Now(), + SessionID: sessionID, + JoinToken: joinToken, + MessageHash: event.MessageHash, + ThresholdN: int(event.ThresholdN), + ThresholdT: int(event.ThresholdT), + SelectedParties: event.SelectedParties, + CreatedAt: time.Now(), }) case "session_started": @@ -410,13 +416,32 @@ func createCoManagedSessionEventHandler( zap.String("session_id", event.SessionId), zap.String("party_id", partyID)) - input := use_cases.ParticipateKeygenInput{ - SessionID: pendingSession.SessionID, - PartyID: partyID, - JoinToken: pendingSession.JoinToken, + // Build SessionInfo from pending session and event data + // Note: We already called JoinSession in session_created phase, + // so we use ExecuteWithSessionInfo to skip the duplicate JoinSession call + participants := make([]use_cases.ParticipantInfo, len(pendingSession.SelectedParties)) + for i, p := range pendingSession.SelectedParties { + participants[i] = use_cases.ParticipantInfo{ + PartyID: p, + PartyIndex: i, + } } - result, err := participateKeygenUC.Execute(participateCtx, input) + sessionInfo := &use_cases.SessionInfo{ + SessionID: pendingSession.SessionID, + SessionType: "co_managed_keygen", + ThresholdN: pendingSession.ThresholdN, + ThresholdT: pendingSession.ThresholdT, + MessageHash: pendingSession.MessageHash, + Participants: participants, + } + + result, err := participateKeygenUC.ExecuteWithSessionInfo( + participateCtx, + pendingSession.SessionID, + partyID, + sessionInfo, + ) if err != nil { logger.Error("Co-managed keygen participation failed", zap.Error(err), diff --git a/backend/mpc-system/services/server-party/application/use_cases/participate_keygen.go b/backend/mpc-system/services/server-party/application/use_cases/participate_keygen.go index 88e80c3b..70b22313 100644 --- a/backend/mpc-system/services/server-party/application/use_cases/participate_keygen.go +++ b/backend/mpc-system/services/server-party/application/use_cases/participate_keygen.go @@ -134,12 +134,47 @@ func (uc *ParticipateKeygenUseCase) Execute( } } - // 2. Find self in participants and build party index map + // Delegate to the common execution logic + return uc.executeWithSessionInfo(ctx, input.SessionID, input.PartyID, sessionInfo) +} + +// ExecuteWithSessionInfo participates in a keygen session with pre-obtained SessionInfo +// This is used by server-party-co-managed which has already called JoinSession in session_created phase +// and receives session_started event when all participants have joined +func (uc *ParticipateKeygenUseCase) ExecuteWithSessionInfo( + ctx context.Context, + sessionID uuid.UUID, + partyID string, + sessionInfo *SessionInfo, +) (*ParticipateKeygenOutput, error) { + // Validate session type + if sessionInfo.SessionType != "keygen" && sessionInfo.SessionType != "co_managed_keygen" { + return nil, ErrInvalidSession + } + + logger.Info("ExecuteWithSessionInfo: starting keygen with pre-obtained session info", + zap.String("session_id", sessionID.String()), + zap.String("party_id", partyID), + zap.String("session_type", sessionInfo.SessionType), + zap.Int("participants", len(sessionInfo.Participants))) + + // Delegate to the common execution logic + return uc.executeWithSessionInfo(ctx, sessionID, partyID, sessionInfo) +} + +// executeWithSessionInfo is the common execution logic shared by Execute and ExecuteWithSessionInfo +func (uc *ParticipateKeygenUseCase) executeWithSessionInfo( + ctx context.Context, + sessionID uuid.UUID, + partyID string, + sessionInfo *SessionInfo, +) (*ParticipateKeygenOutput, error) { + // 1. Find self in participants and build party index map var selfIndex int partyIndexMap := make(map[string]int) for _, p := range sessionInfo.Participants { partyIndexMap[p.PartyID] = p.PartyIndex - if p.PartyID == input.PartyID { + if p.PartyID == partyID { selfIndex = p.PartyIndex } logger.Debug("Added participant to index map", @@ -147,13 +182,13 @@ func (uc *ParticipateKeygenUseCase) Execute( zap.Int("party_index", p.PartyIndex)) } logger.Info("Built party index map", - zap.String("session_id", input.SessionID.String()), - zap.String("self_party_id", input.PartyID), + zap.String("session_id", sessionID.String()), + zap.String("self_party_id", partyID), zap.Int("self_index", selfIndex), zap.Int("total_participants", len(sessionInfo.Participants))) // 3. Subscribe to messages - msgChan, err := uc.messageRouter.SubscribeMessages(ctx, input.SessionID, input.PartyID) + msgChan, err := uc.messageRouter.SubscribeMessages(ctx, sessionID, partyID) if err != nil { return nil, err } @@ -161,8 +196,8 @@ func (uc *ParticipateKeygenUseCase) Execute( // 4. Run TSS Keygen protocol saveData, publicKey, err := uc.runKeygenProtocol( ctx, - input.SessionID, - input.PartyID, + sessionID, + partyID, selfIndex, sessionInfo.Participants, sessionInfo.ThresholdN, @@ -175,15 +210,15 @@ func (uc *ParticipateKeygenUseCase) Execute( } // 5. Encrypt the share - encryptedShare, err := uc.cryptoService.EncryptShare(saveData, input.PartyID) + encryptedShare, err := uc.cryptoService.EncryptShare(saveData, partyID) if err != nil { return nil, err } keyShare := entities.NewPartyKeyShare( - input.PartyID, + partyID, selfIndex, - input.SessionID, + sessionID, sessionInfo.ThresholdN, sessionInfo.ThresholdT, encryptedShare, @@ -201,21 +236,21 @@ func (uc *ParticipateKeygenUseCase) Execute( return nil, ErrShareSaveFailed } logger.Info("Share saved to database (persistent party)", - zap.String("party_id", input.PartyID), - zap.String("session_id", input.SessionID.String())) + zap.String("party_id", partyID), + zap.String("session_id", sessionID.String())) case "delegate": // Delegate Party: do NOT save to database, return to user shareForUser = encryptedShare logger.Info("Share NOT saved, will be returned to user (delegate party)", - zap.String("party_id", input.PartyID), - zap.String("session_id", input.SessionID.String()), + zap.String("party_id", partyID), + zap.String("session_id", sessionID.String()), zap.Int("share_size", len(shareForUser))) case "temporary": // Temporary Party: optionally save to temp storage (not implemented yet) logger.Info("Temporary party - share not saved", - zap.String("party_id", input.PartyID)) + zap.String("party_id", partyID)) default: // Default to persistent for safety @@ -223,12 +258,12 @@ func (uc *ParticipateKeygenUseCase) Execute( return nil, ErrShareSaveFailed } logger.Warn("Unknown party role, defaulting to persistent", - zap.String("party_id", input.PartyID), + zap.String("party_id", partyID), zap.String("role", partyRole)) } // 7. Report completion to coordinator - if err := uc.sessionClient.ReportCompletion(ctx, input.SessionID, input.PartyID, publicKey); err != nil { + if err := uc.sessionClient.ReportCompletion(ctx, sessionID, partyID, publicKey); err != nil { logger.Error("failed to report completion", zap.Error(err)) // Don't fail - share is handled }