lux/mpc
2
0
Fork 0
mirror of https://github.com/luxfi/mpc.git synced 2026-09-24 00:26:24 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #10 from cryptoniumX/jetstream-upgrade

Browse Source 
Jetstream integration + better log format
This commit is contained in:
Thi Nguyen
2024-02-04 21:59:40 +07:00
committed by GitHub
parent 05f9337edd 21be5253f6
commit 2dd5d1a47a
10 changed files with 569 additions and 232 deletions
Download Patch File Download Diff File Expand all files Collapse all files
cmd/main.go
+12
View File
@@ -50,6 +50,15 @@ func main() {
defer natsConn.Close()
pubsub := messaging.NewNATSPubSub(natsConn)
directMessaging := messaging.NewNatsDirectMessaging(natsConn)
mqManager := messaging.NewNATsMessageQueueManager("mpc", []string{
"mpc.mpc_keygen_success.*",
"mpc.mpc_sign_success.*",
}, natsConn)
genKeySuccessQueue := mqManager.NewMessageQueue("mpc_keygen_success")
defer genKeySuccessQueue.Close()
singingSuccessQueue := mqManager.NewMessageQueue("mpc_sign_success")
defer singingSuccessQueue.Close()
logger.Info("Node is running", "peerID", nodeID, "name", *nodeName)
@@ -69,8 +78,11 @@ func main() {
eventConsumer := eventconsumer.NewEventConsumer(
mpcNode,
pubsub,
genKeySuccessQueue,
singingSuccessQueue,
)
eventConsumer.Run()
defer eventConsumer.Close()
// Create a channel to receive signals
signals := make(chan os.Signal, 1)
pkg/eventconsumer/consumer.go
+57 -10
View File
@@ -1,6 +1,7 @@
package eventconsumer
import (
"context"
"encoding/json"
"log"
"math/big"
@@ -17,17 +18,31 @@ const (
type EventConsumer interface {
Run()
Close() error
}
type eventConsumer struct {
node *mpc.Node
pubsub messaging.PubSub
genKeySucecssQueue messaging.MessageQueue
signingSuccessQueue messaging.MessageQueue
keyGenerationSub messaging.Subscription
signingSub messaging.Subscription
}
func NewEventConsumer(node *mpc.Node, pubsub messaging.PubSub) EventConsumer {
func NewEventConsumer(
node *mpc.Node,
pubsub messaging.PubSub,
genKeySucecssQueue messaging.MessageQueue,
signingSuccessQueue messaging.MessageQueue,
) EventConsumer {
return &eventConsumer{
node: node,
pubsub: pubsub,
node: node,
pubsub: pubsub,
genKeySucecssQueue: genKeySucecssQueue,
signingSuccessQueue: signingSuccessQueue,
}
}
@@ -46,36 +61,43 @@ func (ec *eventConsumer) Run() {
}
func (ec *eventConsumer) consumeKeyGenerationEvent() error {
return ec.pubsub.Subscribe(MPCGenerateEvent, func(msg []byte) {
sub, err := ec.pubsub.Subscribe(MPCGenerateEvent, func(msg []byte) {
walletID := string(msg)
// TODO: threshold is configurable
threshold := 2
session, err := ec.node.CreateKeyGenSession(walletID, threshold)
session, err := ec.node.CreateKeyGenSession(walletID, threshold, ec.genKeySucecssQueue)
if err != nil {
logger.Error("Failed to create key generation session", err)
return
}
session.Init()
ctx, done := context.WithCancel(context.Background())
go func() {
for {
select {
case <-ctx.Done():
return
case err := <-session.ErrCh:
logger.Error("Keygen session error", err)
}
}
}()
go session.GenerateKey()
go session.GenerateKey(done)
// TODO -> done and close channel
session.ListenToIncomingMessage()
})
ec.keyGenerationSub = sub
if err != nil {
return err
}
return nil
}
func (ec *eventConsumer) consumeTxSigningEvent() error {
return ec.pubsub.Subscribe(MPCSignEvent, func(raw []byte) {
sub, err := ec.pubsub.Subscribe(MPCSignEvent, func(raw []byte) {
var msg SignTxMessage
err := json.Unmarshal(raw, &msg)
if err != nil {
@@ -90,6 +112,7 @@ func (ec *eventConsumer) consumeTxSigningEvent() error {
msg.TxID,
msg.NetworkInternalCode,
threshold,
ec.signingSuccessQueue,
)
if err != nil {
logger.Error("Failed to create signing session", err)
@@ -99,9 +122,12 @@ func (ec *eventConsumer) consumeTxSigningEvent() error {
txBigInt := new(big.Int).SetBytes(msg.Tx)
session.Init(txBigInt)
ctx, done := context.WithCancel(context.Background())
go func() {
for {
select {
case <-ctx.Done():
return
case err := <-session.ErrCh:
logger.Error("Signing session error", err)
}
@@ -109,8 +135,29 @@ func (ec *eventConsumer) consumeTxSigningEvent() error {
}()
go session.Sign()
go session.Sign(done)
// TODO -> done and close channel
session.ListenToIncomingMessage()
})
ec.signingSub = sub
if err != nil {
return err
}
return nil
}
// Close and clean up
func (ec *eventConsumer) Close() error {
err := ec.keyGenerationSub.Unsubscribe()
if err != nil {
return err
}
err = ec.signingSub.Unsubscribe()
if err != nil {
return err
}
return nil
}
pkg/messaging/message_queue.go
+153
View File
@@ -0,0 +1,153 @@
package messaging
import (
"context"
"errors"
"fmt"
"github.com/cryptoniumX/mpcium/pkg/logger"
"github.com/nats-io/nats.go"
"github.com/nats-io/nats.go/jetstream"
)
var (
ErrPermament = errors.New("Permanent messaging error")
)
type MessageQueue interface {
Enqueue(topic string, message []byte, options *EnqueueOptions) error
Dequeue(topic string, handler func(message []byte) error) error
Close()
}
type EnqueueOptions struct {
IdempotententKey string
}
type msgQueue struct {
consumerName string
js jetstream.JetStream
consumer jetstream.Consumer
consumerContext jetstream.ConsumeContext
}
type NATsMessageQueueManager struct {
queueName string
js jetstream.JetStream
}
func NewNATsMessageQueueManager(queueName string, subjectWildCards []string, nc *nats.Conn) *NATsMessageQueueManager {
js, err := jetstream.New(nc)
if err != nil {
logger.Fatal("Error creating JetStream context: ", err)
}
ctx := context.Background()
stream, err := js.Stream(ctx, queueName)
if err != nil {
logger.Warn("Stream not found, creating new stream", "stream", queueName)
}
if stream != nil {
info, _ := stream.Info(ctx)
logger.Info("Stream found", "info", info)
}
_, err = js.CreateOrUpdateStream(context.Background(), jetstream.StreamConfig{
Name: queueName,
Description: "Stream for " + queueName,
Subjects: subjectWildCards,
MaxBytes: 1024,
Storage: jetstream.FileStorage,
Retention: jetstream.WorkQueuePolicy,
})
if err != nil {
logger.Fatal("Error creating JetStream stream: ", err)
}
logger.Info("Creating apex NATs Jetstream context successfully!")
return &NATsMessageQueueManager{
queueName: queueName,
js: js,
}
}
func (m *NATsMessageQueueManager) NewMessageQueue(consumerName string) MessageQueue {
mq := &msgQueue{
consumerName: consumerName,
js: m.js,
}
consumerWildCard := fmt.Sprintf("%s.%s.*", m.queueName, consumerName)
cfg := jetstream.ConsumerConfig{
Name: consumerName,
Durable: consumerName,
MaxAckPending: 4,
FilterSubjects: []string{
consumerWildCard,
},
MaxDeliver: 3,
}
logger.Info("Creating consumer for subject", "config", cfg)
consumer, err := m.js.CreateOrUpdateConsumer(context.Background(), m.queueName, cfg)
if err != nil {
logger.Fatal("Error creating JetStream consumer: ", err)
}
mq.consumer = consumer
return mq
}
func (mq *msgQueue) Enqueue(topic string, message []byte, options *EnqueueOptions) error {
header := nats.Header{}
if options != nil {
header.Add("Nats-Msg-Id", options.IdempotententKey)
}
logger.Info("Publishing message", "topic", topic)
_, err := mq.js.PublishMsg(context.Background(), &nats.Msg{
Subject: topic,
Data: message,
Header: header,
})
if err != nil {
return fmt.Errorf("Error enqueueing message: %w", err)
}
return nil
}
func (mq *msgQueue) Dequeue(topic string, handler func(message []byte) error) error {
c, err := mq.consumer.Consume(func(msg jetstream.Msg) {
meta, _ := msg.Metadata()
logger.Info("Received message", "meta", meta)
err := handler(msg.Data())
if err != nil {
if errors.Is(err, ErrPermament) {
logger.Info("Permanent error on message", "meta", meta)
msg.Term()
return
}
logger.Error("Error handling message: ", err)
msg.Nak()
return
}
logger.Info("Message Acknowledged", "meta", meta)
err = msg.Ack()
if err != nil {
logger.Error("Error acknowledging message: ", err)
}
})
mq.consumerContext = c
return err
}
func (mq *msgQueue) Close() {
mq.consumerContext.Stop()
}
func (n *msgQueue) handleReconnect(nc *nats.Conn) {
logger.Info("NATS: Reconnected to NATS")
}
pkg/messaging/point2point.go
+6 -5
View File
@@ -7,7 +7,7 @@ import (
)
type DirectMessaging interface {
Listen(target string, handler func(data []byte)) error
Listen(target string, handler func(data []byte)) (Subscription, error)
Send(target string, data []byte) error
}
@@ -29,13 +29,14 @@ func (d *natsDirectMessaging) Send(id string, message []byte) error {
return nil
}
func (d *natsDirectMessaging) Listen(id string, handler func(data []byte)) error {
_, err := d.natsConn.Subscribe(id, func(m *nats.Msg) {
func (d *natsDirectMessaging) Listen(id string, handler func(data []byte)) (Subscription, error) {
sub, err := d.natsConn.Subscribe(id, func(m *nats.Msg) {
handler(m.Data)
m.Respond([]byte("OK"))
})
if err != nil {
return err
return nil, err
}
return nil
return &natsSubscription{subscription: sub}, nil
}
pkg/messaging/pubsub.go
+19 -6
View File
@@ -1,34 +1,47 @@
package messaging
import (
"github.com/cryptoniumX/mpcium/pkg/logger"
"github.com/nats-io/nats.go"
)
type Subscription interface {
Unsubscribe() error
}
type PubSub interface {
Publish(topic string, message []byte) error
Subscribe(topic string, handler func(message []byte)) error
Subscribe(topic string, handler func(message []byte)) (Subscription, error)
}
type natsPubSub struct {
natsConn *nats.Conn
}
type natsSubscription struct {
subscription *nats.Subscription
}
func (ns *natsSubscription) Unsubscribe() error {
return ns.subscription.Unsubscribe()
}
func NewNATSPubSub(natsConn *nats.Conn) PubSub {
return &natsPubSub{natsConn}
}
func (n *natsPubSub) Publish(topic string, message []byte) error {
logger.Infof("Publishing to topic %s", topic)
return n.natsConn.Publish(topic, message)
}
func (n *natsPubSub) Subscribe(topic string, handler func(message []byte)) error {
func (n *natsPubSub) Subscribe(topic string, handler func(message []byte)) (Subscription, error) {
// TODO: Handle subscription
// handle more fields in msg
_, err := n.natsConn.Subscribe(topic, func(msg *nats.Msg) {
sub, err := n.natsConn.Subscribe(topic, func(msg *nats.Msg) {
handler(msg.Data)
})
if err != nil {
return nil, err
}
return err
return &natsSubscription{subscription: sub}, nil
}
pkg/mpc/keygen_session.go
+41 -114
View File
@@ -14,7 +14,7 @@ import (
)
const (
TypeGenerateWalletSuccess = "mpc:generate:success:%s"
TypeGenerateWalletSuccess = "mpc.mpc_keygen_success.%s"
)
type KeygenSession struct {
@@ -22,6 +22,11 @@ type KeygenSession struct {
endCh chan keygen.LocalPartySaveData
}
type KeygenSuccessEvent struct {
WalletID string `json:"wallet_id"`
PubKey []byte `json:"pub_key"`
}
func NewKeygenSession(
walletID string,
pubSub messaging.PubSub,
@@ -31,6 +36,7 @@ func NewKeygenSession(
threshold int,
preParams *keygen.LocalPreParams,
kvstore kvstore.KVStore,
successQueue messaging.MessageQueue,
) *KeygenSession {
return &KeygenSession{
Session: Session{
@@ -44,78 +50,29 @@ func NewKeygenSession(
ErrCh: make(chan error),
preParams: preParams,
kvstore: kvstore,
topicComposer: &TopicComposer{
ComposeBroadcastTopic: func() string {
return fmt.Sprintf("keygen:broadcast:%s", walletID)
},
ComposeDirectTopic: func(nodeID string) string {
return fmt.Sprintf("keygen:direct:%s:%s", walletID, nodeID)
},
},
successQueue: successQueue,
},
endCh: make(chan keygen.LocalPartySaveData),
}
}
func (s *KeygenSession) PartyID() *tss.PartyID {
return s.selfPartyID
}
func (s *KeygenSession) PartyIDs() []*tss.PartyID {
return s.partyIDs
}
func (s *KeygenSession) PartyCount() int {
return len(s.partyIDs)
}
func (s *KeygenSession) Init() {
logger.Infof("Initializing session with partyID: %s, peerIDs %s", s.selfPartyID, s.partyIDs)
ctx := tss.NewPeerContext(s.partyIDs)
params := tss.NewParameters(tss.S256(), ctx, s.selfPartyID, len(s.partyIDs), s.threshold)
s.party = keygen.NewLocalParty(params, s.outCh, s.endCh, *s.preParams)
logger.Info("Initialized session successfully")
logger.Infof("PartyID: %s is ready!", s.party.PartyID())
logger.Infof("[INITIALIZED] Initialized session successfully partyID: %s, peerIDs %s, walletID %s", s.selfPartyID, s.partyIDs, s.walletID)
}
func (s *KeygenSession) composeDirectMessageTopic(nodeID string) string {
return fmt.Sprintf("keygen:direct:%s:%s", s.walletID, nodeID)
}
func (s *KeygenSession) composeBroadcastTopic() string {
return fmt.Sprintf("keygen:broadcast:%s", s.walletID)
}
func (s *KeygenSession) handleTssMessage(keyshare tss.Message) {
data, routing, err := keyshare.WireBytes()
logger.Info("Received message", "routing", routing)
if err != nil {
s.ErrCh <- err
return
}
msg, err := MarshalTssMessage(s.walletID, data, routing.IsBroadcast, routing.From, routing.To)
if err != nil {
s.ErrCh <- fmt.Errorf("failed to marshal tss message: %w", err)
return
}
logger.Info("Preparing to send message to", "walletID", s.walletID, "to", routing.To, "isBroadcast", routing.IsBroadcast)
if routing.IsBroadcast && len(routing.To) == 0 {
logger.Info("Broadcasting message", "walletID", s.walletID, "from", s.PartyID(), "routing", routing.To)
err := s.pubSub.Publish(s.composeBroadcastTopic(), msg)
if err != nil {
s.ErrCh <- err
return
}
} else {
for _, to := range routing.To {
nodeID := PartyIDToNodeID(to)
topic := s.composeDirectMessageTopic(nodeID)
logger.Info("Direct message", "walletID", s.walletID, "from", s.PartyID(), "to", to)
logger.Info("Sending direct message to destination topic", "topic", topic)
err := s.direct.Send(topic, msg)
if err != nil {
s.ErrCh <- fmt.Errorf("Failed to send direct message to %s: %w", topic, err)
}
}
}
}
func (s *KeygenSession) GenerateKey() {
func (s *KeygenSession) GenerateKey(done func()) {
logger.Info("Starting to generate key", "walletID", s.walletID)
go func() {
if err := s.party.Start(); err != nil {
@@ -126,7 +83,6 @@ func (s *KeygenSession) GenerateKey() {
for {
select {
case msg := <-s.outCh:
logger.Info("Received message for session", "walletID", s.walletID, "msg", msg)
s.handleTssMessage(msg)
case saveData := <-s.endCh:
keyBytes, err := json.Marshal(saveData)
@@ -156,62 +112,33 @@ func (s *KeygenSession) GenerateKey() {
return
}
s.pubSub.Publish(fmt.Sprintf(TypeGenerateWalletSuccess, s.walletID), pubKeyBytes)
}
}
}
func (s *KeygenSession) ListenToIncomingMessage() {
go func() {
s.pubSub.Subscribe(s.composeBroadcastTopic(), func(msg []byte) {
logger.Info("Received tss message from broadcast channel")
s.receiveTssMessage(msg)
})
}()
nodeID := PartyIDToNodeID(s.selfPartyID)
logger.Info("Listening on target topic for incoming messages", "topic", s.composeDirectMessageTopic(nodeID))
targetID := s.composeDirectMessageTopic(nodeID)
s.direct.Listen(targetID, func(msg []byte) {
logger.Info("Received message from direct channel", "targetID", targetID)
s.receiveTssMessage(msg)
})
}
func (s *KeygenSession) receiveTssMessage(rawMsg []byte) {
msg, err := UnmarshalTssMessage(rawMsg)
if err != nil {
s.ErrCh <- fmt.Errorf("Failed to unmarshal message: %w", err)
return
}
logger.Info("Received message", "from", msg.From, "to", msg.To, "isBroadcast", msg.IsBroadcast)
isBroadcast := msg.IsBroadcast && len(msg.To) == 0
isToSelf := len(msg.To) == 1 && ComparePartyIDs(msg.To[0], s.selfPartyID)
if isBroadcast || isToSelf {
go func() {
if isBroadcast {
logger.Info("Updating broadcast message", "to", msg.To)
} else if isToSelf {
logger.Info("Updating direct message to local node", "to", msg.To)
successEvent := KeygenSuccessEvent{
WalletID: s.walletID,
PubKey: pubKeyBytes,
}
ok, err := s.party.UpdateFromBytes(msg.MsgBytes, msg.From, msg.IsBroadcast)
if !ok || err != nil {
logger.Error("Failed to update party", err, "walletID", s.walletID)
successEventBytes, err := json.Marshal(successEvent)
if err != nil {
s.ErrCh <- fmt.Errorf("failed to marshal success event: %w", err)
return
}
}()
err = s.successQueue.Enqueue(fmt.Sprintf(TypeGenerateWalletSuccess, s.walletID), successEventBytes, &messaging.EnqueueOptions{
IdempotententKey: fmt.Sprintf(TypeGenerateWalletSuccess, s.walletID),
})
if err != nil {
logger.Error("Failed to publish key generation success message", err)
s.ErrCh <- fmt.Errorf("Failed to publish key generation success message %w", err)
return
}
logger.Info("[COMPLETED KEY GEN] Key generation completed successfully", "walletID", s.walletID)
err = s.Close()
if err != nil {
logger.Error("Failed to close session", err)
}
done()
return
}
}
}
// Close and clean up
func (s *KeygenSession) Close() {
return
}
pkg/mpc/node.go
+4 -1
View File
@@ -87,7 +87,7 @@ func composeReadyTopic(nodeID string) string {
return fmt.Sprintf("%s-%s", nodeID, "ready")
}
func (p *Node) CreateKeyGenSession(walletID string, threshold int) (*KeygenSession, error) {
func (p *Node) CreateKeyGenSession(walletID string, threshold int, successQueue messaging.MessageQueue) (*KeygenSession, error) {
if !p.peerRegistry.ArePeersReady() {
return nil, fmt.Errorf("All peers are not ready!")
}
@@ -102,6 +102,7 @@ func (p *Node) CreateKeyGenSession(walletID string, threshold int) (*KeygenSessi
threshold,
p.preParams,
p.kvstore,
successQueue,
)
return session, nil
}
@@ -111,6 +112,7 @@ func (p *Node) CreateSigningSession(
txID string,
networkInternalCode string,
threshold int,
successQueue messaging.MessageQueue,
) (*SigningSession, error) {
if !p.peerRegistry.ArePeersReady() {
return nil, fmt.Errorf("All peers are not ready!")
@@ -128,6 +130,7 @@ func (p *Node) CreateSigningSession(
threshold,
p.preParams,
p.kvstore,
successQueue,
)
return session, nil
pkg/mpc/rounds.go
+126
View File
@@ -0,0 +1,126 @@
package mpc
import (
"github.com/bnb-chain/tss-lib/ecdsa/keygen"
"github.com/bnb-chain/tss-lib/ecdsa/signing"
"github.com/bnb-chain/tss-lib/tss"
"github.com/cryptoniumX/mpcium/pkg/common/errors"
)
type RoundInfo struct {
Index int
RoundMsg string
MsgIdentifier string
}
const (
KEYGEN1 = "KGRound1Message"
KEYGEN2aUnicast = "KGRound2Message1"
KEYGEN2b = "KGRound2Message2"
KEYGEN3 = "KGRound3Message"
KEYSIGN1aUnicast = "SignRound1Message1"
KEYSIGN1b = "SignRound1Message2"
KEYSIGN2Unicast = "SignRound2Message"
KEYSIGN3 = "SignRound3Message"
KEYSIGN4 = "SignRound4Message"
KEYSIGN5 = "SignRound5Message"
KEYSIGN6 = "SignRound6Message"
KEYSIGN7 = "SignRound7Message"
KEYSIGN8 = "SignRound8Message"
KEYSIGN9 = "SignRound9Message"
TSSKEYGENROUNDS = 4
TSSKEYSIGNROUNDS = 10
)
func GetMsgRound(msg []byte, partyID *tss.PartyID, isBroadcast bool) (RoundInfo, error) {
parsedMsg, err := tss.ParseWireMessage(msg, partyID, isBroadcast)
if err != nil {
return RoundInfo{}, err
}
switch parsedMsg.Content().(type) {
case *keygen.KGRound1Message:
return RoundInfo{
Index: 0,
RoundMsg: KEYGEN1,
}, nil
case *keygen.KGRound2Message1:
return RoundInfo{
Index: 1,
RoundMsg: KEYGEN2aUnicast,
}, nil
case *keygen.KGRound2Message2:
return RoundInfo{
Index: 2,
RoundMsg: KEYGEN2b,
}, nil
case *keygen.KGRound3Message:
return RoundInfo{
Index: 3,
RoundMsg: KEYGEN3,
}, nil
case *signing.SignRound1Message1:
return RoundInfo{
Index: 0,
RoundMsg: KEYSIGN1aUnicast,
}, nil
case *signing.SignRound1Message2:
return RoundInfo{
Index: 1,
RoundMsg: KEYSIGN1b,
}, nil
case *signing.SignRound2Message:
return RoundInfo{
Index: 2,
RoundMsg: KEYSIGN2Unicast,
}, nil
case *signing.SignRound3Message:
return RoundInfo{
Index: 3,
RoundMsg: KEYSIGN3,
}, nil
case *signing.SignRound4Message:
return RoundInfo{
Index: 4,
RoundMsg: KEYSIGN4,
}, nil
case *signing.SignRound5Message:
return RoundInfo{
Index: 5,
RoundMsg: KEYSIGN5,
}, nil
case *signing.SignRound6Message:
return RoundInfo{
Index: 6,
RoundMsg: KEYSIGN6,
}, nil
case *signing.SignRound7Message:
return RoundInfo{
Index: 7,
RoundMsg: KEYSIGN7,
}, nil
case *signing.SignRound8Message:
return RoundInfo{
Index: 8,
RoundMsg: KEYSIGN8,
}, nil
case *signing.SignRound9Message:
return RoundInfo{
Index: 9,
RoundMsg: KEYSIGN9,
}, nil
default:
return RoundInfo{}, errors.New("unknown round")
}
}
pkg/mpc/session.go
+130
View File
@@ -1,12 +1,22 @@
package mpc
import (
"fmt"
"strings"
"github.com/bnb-chain/tss-lib/ecdsa/keygen"
"github.com/bnb-chain/tss-lib/tss"
"github.com/cryptoniumX/mpcium/pkg/common/errors"
"github.com/cryptoniumX/mpcium/pkg/kvstore"
"github.com/cryptoniumX/mpcium/pkg/logger"
"github.com/cryptoniumX/mpcium/pkg/messaging"
)
type TopicComposer struct {
ComposeBroadcastTopic func() string
ComposeDirectTopic func(nodeID string) string
}
type Session struct {
walletID string
pubSub messaging.PubSub
@@ -20,4 +30,124 @@ type Session struct {
party tss.Party
preParams *keygen.LocalPreParams
kvstore kvstore.KVStore
broadcastSub messaging.Subscription
directSub messaging.Subscription
successQueue messaging.MessageQueue
topicComposer *TopicComposer
}
func (s *Session) PartyID() *tss.PartyID {
return s.selfPartyID
}
func (s *Session) PartyIDs() []*tss.PartyID {
return s.partyIDs
}
func (s *Session) PartyCount() int {
return len(s.partyIDs)
}
func (s *Session) handleTssMessage(keyshare tss.Message) {
data, routing, err := keyshare.WireBytes()
if err != nil {
s.ErrCh <- err
return
}
msg, err := MarshalTssMessage(s.walletID, data, routing.IsBroadcast, routing.From, routing.To)
if err != nil {
s.ErrCh <- fmt.Errorf("failed to marshal tss message: %w", err)
return
}
if routing.IsBroadcast && len(routing.To) == 0 {
err := s.pubSub.Publish(s.topicComposer.ComposeBroadcastTopic(), msg)
if err != nil {
s.ErrCh <- err
return
}
} else {
for _, to := range routing.To {
nodeID := PartyIDToNodeID(to)
topic := s.topicComposer.ComposeDirectTopic(nodeID)
err := s.direct.Send(topic, msg)
if err != nil {
s.ErrCh <- fmt.Errorf("Failed to send direct message to %s: %w", topic, err)
}
}
}
}
func (s *Session) receiveTssMessage(rawMsg []byte) {
msg, err := UnmarshalTssMessage(rawMsg)
if err != nil {
s.ErrCh <- fmt.Errorf("Failed to unmarshal message: %w", err)
return
}
toIDs := make([]string, len(msg.To))
for i, id := range msg.To {
toIDs[i] = id.String()
}
round, err := GetMsgRound(msg.MsgBytes, s.selfPartyID, msg.IsBroadcast)
if err != nil {
s.ErrCh <- errors.Wrap(err, "Broken TSS Share")
return
}
logger.Info("Received message", "from", msg.From.String(), "to", strings.Join(toIDs, ","), "isBroadcast", msg.IsBroadcast, "round", round.RoundMsg)
isBroadcast := msg.IsBroadcast && len(msg.To) == 0
isToSelf := len(msg.To) == 1 && ComparePartyIDs(msg.To[0], s.selfPartyID)
if isBroadcast || isToSelf {
ok, err := s.party.UpdateFromBytes(msg.MsgBytes, msg.From, msg.IsBroadcast)
if !ok || err != nil {
logger.Error("Failed to update party", err, "walletID", s.walletID)
return
}
}
}
func (s *Session) ListenToIncomingMessage() {
go func() {
sub, err := s.pubSub.Subscribe(s.topicComposer.ComposeBroadcastTopic(), func(msg []byte) {
s.receiveTssMessage(msg)
})
if err != nil {
s.ErrCh <- fmt.Errorf("Failed to subscribe to broadcast topic %s: %w", s.topicComposer.ComposeBroadcastTopic(), err)
return
}
s.broadcastSub = sub
}()
nodeID := PartyIDToNodeID(s.selfPartyID)
targetID := s.topicComposer.ComposeDirectTopic(nodeID)
sub, err := s.direct.Listen(targetID, func(msg []byte) {
go s.receiveTssMessage(msg) // async for avoid timeout
})
if err != nil {
s.ErrCh <- fmt.Errorf("Failed to subscribe to direct topic %s: %w", targetID, err)
}
s.directSub = sub
}
func (s *Session) Close() error {
err := s.broadcastSub.Unsubscribe()
if err != nil {
return err
}
err = s.directSub.Unsubscribe()
if err != nil {
return err
}
return nil
}
pkg/mpc/signing_session.go
+21 -96
View File
@@ -17,7 +17,7 @@ import (
)
const (
TypeSignSuccess = "mpc:sign:success"
SignSuccessTopic = "mpc.mpc_sign_success.completed"
)
type SigningSession struct {
@@ -49,6 +49,7 @@ func NewSigningSession(
threshold int,
preParams *keygen.LocalPreParams,
kvstore kvstore.KVStore,
succesQueue messaging.MessageQueue,
) *SigningSession {
return &SigningSession{
Session: Session{
@@ -62,6 +63,15 @@ func NewSigningSession(
ErrCh: make(chan error),
preParams: preParams,
kvstore: kvstore,
topicComposer: &TopicComposer{
ComposeBroadcastTopic: func() string {
return fmt.Sprintf("sign:broadcast:%s", walletID)
},
ComposeDirectTopic: func(nodeID string) string {
return fmt.Sprintf("sign:direct:%s:%s", walletID, nodeID)
},
},
successQueue: succesQueue,
},
endCh: make(chan common.SignatureData),
txID: txID,
@@ -93,7 +103,7 @@ func (s *SigningSession) Init(tx *big.Int) {
logger.Info("Initialized sigining session successfully!")
}
func (s *SigningSession) Sign() {
func (s *SigningSession) Sign(done func()) {
logger.Info("Starting signing", "walletID", s.walletID)
go func() {
if err := s.party.Start(); err != nil {
@@ -105,7 +115,6 @@ func (s *SigningSession) Sign() {
select {
case msg := <-s.outCh:
logger.Info("Generating and handle tss message", "walletID", s.walletID)
s.handleTssMessage(msg)
case sig := <-s.endCh:
publicKey := *s.data.ECDSAPub
@@ -136,107 +145,23 @@ func (s *SigningSession) Sign() {
return
}
err = s.pubSub.Publish(TypeSignSuccess, bytes)
err = s.successQueue.Enqueue(SignSuccessTopic, bytes, &messaging.EnqueueOptions{
IdempotententKey: s.txID,
})
if err != nil {
s.ErrCh <- errors.Wrap(err, "Failed to publish sign success message")
return
}
logger.Info("Sign successfully", "walletID", s.walletID)
}
logger.Info("[SIGN] Sign successfully", "walletID", s.walletID)
err = s.Close()
if err != nil {
logger.Error("Failed to close session", err)
}
}
}
func (s *SigningSession) composeDirectMessageTopic(nodeID string) string {
return fmt.Sprintf("sign:direct:%s:%s", s.walletID, nodeID)
}
func (s *SigningSession) composeBroadcastTopic() string {
return fmt.Sprintf("sign:broadcast:%s", s.walletID)
}
func (s *SigningSession) handleTssMessage(tssMsg tss.Message) {
data, routing, err := tssMsg.WireBytes()
logger.Info("Received message", "routing", routing)
if err != nil {
s.ErrCh <- err
return
}
msg, err := MarshalTssMessage(s.walletID, data, routing.IsBroadcast, routing.From, routing.To)
if err != nil {
s.ErrCh <- fmt.Errorf("failed to marshal tss message: %w", err)
return
}
logger.Info("Preparing to send message to", "walletID", s.walletID, "to", routing.To, "isBroadcast", routing.IsBroadcast)
if routing.IsBroadcast && len(routing.To) == 0 {
logger.Info("Broadcasting message", "walletID", s.walletID, "from", s.selfPartyID, "routing", routing.To)
err := s.pubSub.Publish(s.composeBroadcastTopic(), msg)
if err != nil {
s.ErrCh <- err
done()
return
}
} else {
for _, to := range routing.To {
nodeID := PartyIDToNodeID(to)
topic := s.composeDirectMessageTopic(nodeID)
logger.Info("Direct message", "walletID", s.walletID, "from", s.selfPartyID, "to", to)
logger.Info("Sending direct message to destination topic", "topic", topic)
err := s.direct.Send(topic, msg)
if err != nil {
s.ErrCh <- fmt.Errorf("Failed to send direct message to %s: %w", topic, err)
}
}
}
}
func (s *SigningSession) ListenToIncomingMessage() {
go func() {
s.pubSub.Subscribe(s.composeBroadcastTopic(), func(msg []byte) {
logger.Info("Received tss message from broadcast channel")
s.receiveTssMessage(msg)
})
}()
nodeID := PartyIDToNodeID(s.selfPartyID)
logger.Info("Listening on target topic for incoming messages", "topic", s.composeDirectMessageTopic(nodeID))
targetID := s.composeDirectMessageTopic(nodeID)
s.direct.Listen(targetID, func(msg []byte) {
logger.Info("Received message from direct channel", "targetID", targetID)
s.receiveTssMessage(msg)
})
}
func (s *SigningSession) receiveTssMessage(rawMsg []byte) {
msg, err := UnmarshalTssMessage(rawMsg)
if err != nil {
s.ErrCh <- fmt.Errorf("Failed to unmarshal message: %w", err)
return
}
logger.Info("Received message", "from", msg.From, "to", msg.To, "isBroadcast", msg.IsBroadcast)
isBroadcast := msg.IsBroadcast && len(msg.To) == 0
isToSelf := len(msg.To) == 1 && ComparePartyIDs(msg.To[0], s.selfPartyID)
if isBroadcast || isToSelf {
go func() {
if isBroadcast {
logger.Info("Updating broadcast message", "to", msg.To)
} else if isToSelf {
logger.Info("Updating direct message to local node", "to", msg.To)
}
ok, err := s.party.UpdateFromBytes(msg.MsgBytes, msg.From, msg.IsBroadcast)
if !ok || err != nil {
logger.Error("Failed to update party", err, "walletID", s.walletID)
return
}
}()
}
}