rwadurian/backend/services/mpc-service/src/application/commands/participate-signing/participate-signing.handler.ts

/**
 * Participate In Signing Handler
 *
 * Handles the ParticipateInSigningCommand by:
 * 1. Joining the MPC signing session
 * 2. Loading and decrypting the party's share
 * 3. Running the TSS signing protocol
 * 4. Publishing domain events
 */

import { Injectable, Inject, Logger } from '@nestjs/common';
import { ConfigService } from '@nestjs/config';
import { ParticipateInSigningCommand } from './participate-signing.command';
import { PartyShare } from '../../../domain/entities/party-share.entity';
import { SessionState, Participant } from '../../../domain/entities/session-state.entity';
import {
  SessionId,
  PartyId,
  PublicKey,
  Threshold,
  MessageHash,
  Signature,
} from '../../../domain/value-objects';
import { SessionType, ParticipantStatus, KeyCurve } from '../../../domain/enums';
import { ShareEncryptionDomainService } from '../../../domain/services/share-encryption.domain-service';
import {
  TSS_PROTOCOL_SERVICE,
  TSSProtocolDomainService,
  TSSMessage,
  TSSParticipant,
} from '../../../domain/services/tss-protocol.domain-service';
import {
  PARTY_SHARE_REPOSITORY,
  PartyShareRepository,
} from '../../../domain/repositories/party-share.repository.interface';
import {
  SESSION_STATE_REPOSITORY,
  SessionStateRepository,
} from '../../../domain/repositories/session-state.repository.interface';
import { EventPublisherService } from '../../../infrastructure/messaging/kafka/event-publisher.service';
import { MPCCoordinatorClient, SessionInfo } from '../../../infrastructure/external/mpc-system/coordinator-client';
import { MPCMessageRouterClient } from '../../../infrastructure/external/mpc-system/message-router-client';
import { ApplicationError } from '../../../shared/exceptions/domain.exception';

export interface SigningResult {
  signature: string;
  r: string;
  s: string;
  v?: number;
  messageHash: string;
  publicKey: string;
  sessionId: string;
  partyId: string;
}

@Injectable()
export class ParticipateInSigningHandler {
  private readonly logger = new Logger(ParticipateInSigningHandler.name);

  constructor(
    @Inject(PARTY_SHARE_REPOSITORY)
    private readonly partyShareRepo: PartyShareRepository,
    @Inject(SESSION_STATE_REPOSITORY)
    private readonly sessionStateRepo: SessionStateRepository,
    @Inject(TSS_PROTOCOL_SERVICE)
    private readonly tssProtocol: TSSProtocolDomainService,
    private readonly encryptionService: ShareEncryptionDomainService,
    private readonly coordinatorClient: MPCCoordinatorClient,
    private readonly messageRouter: MPCMessageRouterClient,
    private readonly eventPublisher: EventPublisherService,
    private readonly configService: ConfigService,
  ) {}

  async execute(command: ParticipateInSigningCommand): Promise<SigningResult> {
    this.logger.log(`Starting Signing participation for party: ${command.partyId}, session: ${command.sessionId}`);

    // 1. Join the signing session
    const sessionInfo = await this.joinSession(command);
    this.logger.log(`Joined signing session with ${sessionInfo.participants.length} participants`);

    // 2. Load the party's share
    const partyShare = await this.loadPartyShare(command, sessionInfo);
    this.logger.log(`Loaded share: ${partyShare.id.value}`);

    // 3. Create session state for tracking
    const sessionState = this.createSessionState(command, sessionInfo, partyShare);
    await this.sessionStateRepo.save(sessionState);

    try {
      // 4. Decrypt share data
      const masterKey = await this.getMasterKey();
      const rawShareData = this.encryptionService.decrypt(
        partyShare.shareData,
        masterKey,
      );
      this.logger.log('Share data decrypted successfully');

      // 5. Setup message channels
      const { sender, receiver } = await this.setupMessageChannels(
        command.sessionId,
        command.partyId,
      );

      // 6. Run TSS signing protocol
      this.logger.log('Starting TSS Signing protocol...');
      const messageHash = MessageHash.fromHex(command.messageHash);
      const signingResult = await this.tssProtocol.runSigning(
        command.partyId,
        this.convertParticipants(sessionInfo.participants),
        rawShareData,
        messageHash,
        Threshold.create(sessionInfo.thresholdN, sessionInfo.thresholdT),
        {
          curve: KeyCurve.SECP256K1,
          timeout: this.configService.get<number>('MPC_SIGNING_TIMEOUT', 180000),
        },
        sender,
        receiver,
      );
      this.logger.log('TSS Signing protocol completed successfully');

      // 7. Update share usage
      partyShare.markAsUsed(command.messageHash);
      await this.partyShareRepo.update(partyShare);

      // 8. Report completion to coordinator
      await this.coordinatorClient.reportCompletion({
        sessionId: command.sessionId,
        partyId: command.partyId,
        signature: signingResult.signature,
      });

      // 9. Update session state
      sessionState.completeSigning(Signature.fromHex(signingResult.signature));
      await this.sessionStateRepo.update(sessionState);

      // 10. Publish domain events
      await this.eventPublisher.publishAll(partyShare.domainEvents);
      await this.eventPublisher.publishAll(sessionState.domainEvents);
      partyShare.clearDomainEvents();
      sessionState.clearDomainEvents();

      this.logger.log(`Signing completed successfully. Signature: ${signingResult.signature.substring(0, 20)}...`);

      return {
        signature: signingResult.signature,
        r: signingResult.r,
        s: signingResult.s,
        v: signingResult.v,
        messageHash: messageHash.toHex(),
        publicKey: partyShare.publicKey.toHex(),
        sessionId: command.sessionId,
        partyId: command.partyId,
      };
    } catch (error) {
      // Handle failure
      this.logger.error(`Signing failed: ${error.message}`, error.stack);

      sessionState.fail(error.message, 'SIGNING_FAILED');
      await this.sessionStateRepo.update(sessionState);
      await this.eventPublisher.publishAll(sessionState.domainEvents);
      sessionState.clearDomainEvents();

      throw new ApplicationError(`Signing failed: ${error.message}`, 'SIGNING_FAILED');
    }
  }

  private async joinSession(command: ParticipateInSigningCommand): Promise<SessionInfo> {
    try {
      // First, create the session via coordinator to get a valid JWT token
      this.logger.log('Creating MPC signing session via coordinator...');
      const createResponse = await this.coordinatorClient.createSession({
        sessionType: 'sign',
        thresholdN: 3,  // Default 2-of-3 MPC
        thresholdT: 2,
        createdBy: command.partyId,
        messageHash: command.messageHash,
        expiresIn: 300,  // 5 minutes for signing
      });
      this.logger.log(`Signing session created: ${createResponse.sessionId}, now joining...`);

      // Now join using the valid JWT token from the coordinator
      const sessionInfo = await this.coordinatorClient.joinSession({
        sessionId: createResponse.sessionId,
        partyId: command.partyId,
        joinToken: createResponse.joinToken,
      });

      // Return session info with correct IDs and public key from command
      return {
        ...sessionInfo,
        sessionId: createResponse.sessionId,
        joinToken: createResponse.joinToken,
        publicKey: command.publicKey,  // Preserve public key from command
        messageHash: command.messageHash,
      };
    } catch (error) {
      throw new ApplicationError(
        `Failed to join signing session: ${error.message}`,
        'JOIN_SESSION_FAILED',
      );
    }
  }

  private async loadPartyShare(
    command: ParticipateInSigningCommand,
    sessionInfo: SessionInfo,
  ): Promise<PartyShare> {
    const partyId = PartyId.create(command.partyId);

    // If public key is provided in command, use it
    if (command.publicKey) {
      const publicKey = PublicKey.fromHex(command.publicKey);
      const share = await this.partyShareRepo.findByPartyIdAndPublicKey(partyId, publicKey);
      if (!share) {
        throw new ApplicationError(
          'Share not found for specified public key',
          'SHARE_NOT_FOUND',
        );
      }
      return share;
    }

    // Otherwise, get public key from session info
    if (!sessionInfo.publicKey) {
      throw new ApplicationError(
        'Public key not provided in command or session info',
        'PUBLIC_KEY_MISSING',
      );
    }

    const publicKey = PublicKey.fromHex(sessionInfo.publicKey);
    const share = await this.partyShareRepo.findByPartyIdAndPublicKey(partyId, publicKey);

    if (!share) {
      throw new ApplicationError(
        'Share not found for this party and public key',
        'SHARE_NOT_FOUND',
      );
    }

    if (!share.isActive()) {
      throw new ApplicationError(
        `Share is not active: ${share.status}`,
        'SHARE_NOT_ACTIVE',
      );
    }

    return share;
  }

  private createSessionState(
    command: ParticipateInSigningCommand,
    sessionInfo: SessionInfo,
    partyShare: PartyShare,
  ): SessionState {
    const participants: Participant[] = sessionInfo.participants.map(p => ({
      partyId: p.partyId,
      partyIndex: p.partyIndex,
      status: p.partyId === command.partyId
        ? ParticipantStatus.JOINED
        : ParticipantStatus.PENDING,
    }));

    const myParty = sessionInfo.participants.find(p => p.partyId === command.partyId);
    if (!myParty) {
      throw new ApplicationError('Party not found in session participants', 'PARTY_NOT_FOUND');
    }

    return SessionState.create({
      sessionId: SessionId.create(command.sessionId),
      partyId: PartyId.create(command.partyId),
      partyIndex: myParty.partyIndex,
      sessionType: SessionType.SIGN,
      participants,
      thresholdN: sessionInfo.thresholdN,
      thresholdT: sessionInfo.thresholdT,
      publicKey: partyShare.publicKey,
      messageHash: MessageHash.fromHex(command.messageHash),
    });
  }

  private async setupMessageChannels(
    sessionId: string,
    partyId: string,
  ): Promise<{ sender: (msg: TSSMessage) => Promise<void>; receiver: AsyncIterable<TSSMessage> }> {
    const messageStream = await this.messageRouter.subscribeMessages(sessionId, partyId);

    const sender = async (msg: TSSMessage): Promise<void> => {
      await this.messageRouter.sendMessage({
        sessionId,
        fromParty: partyId,
        toParties: msg.toParties,
        roundNumber: msg.roundNumber,
        payload: msg.payload,
      });
    };

    const receiver: AsyncIterable<TSSMessage> = {
      [Symbol.asyncIterator]: () => ({
        next: async (): Promise<IteratorResult<TSSMessage>> => {
          const message = await messageStream.next();
          if (message.done) {
            return { done: true, value: undefined };
          }
          return {
            done: false,
            value: {
              fromParty: message.value.fromParty,
              toParties: message.value.toParties,
              roundNumber: message.value.roundNumber,
              payload: message.value.payload,
            },
          };
        },
      }),
    };

    return { sender, receiver };
  }

  private convertParticipants(
    participants: Array<{ partyId: string; partyIndex: number }>,
  ): TSSParticipant[] {
    return participants.map(p => ({
      partyId: p.partyId,
      partyIndex: p.partyIndex,
    }));
  }

  private async getMasterKey(): Promise<Buffer> {
    const keyHex = this.configService.get<string>('SHARE_MASTER_KEY');
    if (!keyHex) {
      throw new ApplicationError(
        'SHARE_MASTER_KEY not configured',
        'CONFIG_ERROR',
      );
    }
    return Buffer.from(keyHex, 'hex');
  }
}