rwadurian/backend/mpc-system/services/server-party/application/use_cases/participate_keygen.go

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
}