Skip to content

Commit 3c09a2f

Browse files
vietddudevietddude
committed
Add support for selected nodes in wallet creation and key generation
1 parent 1510cb3 commit 3c09a2f

5 files changed

Lines changed: 81 additions & 17 deletions

File tree

‎pkg/client/client.go‎

Lines changed: 12 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ const (
2121
type MPCClient interface {
2222
CreateWallet(walletID string) error
2323
CreateWalletWithAuthorizers(walletID string, authorizerSignatures []types.AuthorizerSignature) error
24+
CreateWalletWithSelectedNodes(walletID string, nodes []string) error
2425
OnWalletCreationResult(callback func(event event.KeygenResultEvent)) error
2526

2627
SignTransaction(msg *types.SignTxMessage) error
@@ -105,14 +106,19 @@ func NewMPCClient(opts Options) MPCClient {
105106

106107
// CreateWallet generates a GenerateKeyMessage, signs it, and publishes it.
107108
func (c *mpcClient) CreateWallet(walletID string) error {
108-
return c.CreateWalletWithAuthorizers(walletID, nil)
109+
return c.createWalletWithAuthorizers(walletID, nil, nil)
109110
}
110111

111-
// CreateWalletWithAuthorizers generates a GenerateKeyMessage with authorizer signatures, signs it, and publishes it.
112112
func (c *mpcClient) CreateWalletWithAuthorizers(walletID string, authorizerSignatures []types.AuthorizerSignature) error {
113+
return c.createWalletWithAuthorizers(walletID, nil, authorizerSignatures)
114+
}
115+
116+
// CreateWalletWithAuthorizers generates a GenerateKeyMessage with authorizer signatures, signs it, and publishes it.
117+
func (c *mpcClient) createWalletWithAuthorizers(walletID string, selectedNodeIDs []string, authorizerSignatures []types.AuthorizerSignature) error {
113118
// build the message
114119
msg := &types.GenerateKeyMessage{
115120
WalletID: walletID,
121+
SelectedNodeIDs: selectedNodeIDs,
116122
AuthorizerSignatures: authorizerSignatures,
117123
}
118124
// compute the canonical raw bytes
@@ -137,6 +143,10 @@ func (c *mpcClient) CreateWalletWithAuthorizers(walletID string, authorizerSigna
137143
return nil
138144
}
139145

146+
func (c *mpcClient) CreateWalletWithSelectedNodes(walletID string, nodes []string) error {
147+
return c.createWalletWithAuthorizers(walletID, nodes, nil)
148+
}
149+
140150
// The callback will be invoked whenever a wallet creation result is received.
141151
func (c *mpcClient) OnWalletCreationResult(callback func(event event.KeygenResultEvent)) error {
142152
err := c.genKeySuccessQueue.Dequeue(GenerateWalletSuccessTopic, func(msg []byte) error {

‎pkg/event/keygen.go‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,10 @@
11
package event
22

33
const (
4-
KeygenBrokerStream = "mpc-keygen"
5-
KeygenConsumerStream = "mpc-keygen-consumer"
6-
KeygenRequestTopic = "mpc.keygen_request.*"
4+
KeygenBrokerStream = "mpc-keygen"
5+
KeygenConsumerStream = "mpc-keygen-consumer"
6+
KeygenRequestTopic = "mpc.keygen_request.*"
7+
KeygenRequestTopicFmt = "mpc.keygen_request.%s"
78
)
89

910
type KeygenResultEvent struct {

‎pkg/eventconsumer/event_consumer.go‎

Lines changed: 24 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -173,13 +173,35 @@ func (ec *eventConsumer) handleKeyGenEvent(natMsg *nats.Msg) {
173173
}
174174

175175
walletID := msg.WalletID
176-
ecdsaSession, err := ec.node.CreateKeyGenSession(mpc.SessionTypeECDSA, walletID, ec.mpcThreshold, ec.genKeyResultQueue)
176+
var nodeIDs []string
177+
if len(msg.SelectedNodeIDs) > 0 {
178+
nodeIDs = msg.SelectedNodeIDs
179+
logger.Info("Using selected node IDs for keygen", "walletID", walletID, "nodeIDs", nodeIDs)
180+
}
181+
182+
ecdsaSession, err := ec.node.CreateKeyGenSession(mpc.SessionTypeECDSA, walletID, ec.mpcThreshold, ec.genKeyResultQueue, nodeIDs)
177183
if err != nil {
184+
if errors.Is(err, mpc.ErrNotInParticipantList) {
185+
logger.Info("Node is not in selected node list for this wallet, skipping keygen",
186+
"walletID", walletID,
187+
"nodeID", ec.node.ID(),
188+
)
189+
// Skip keygen instead of treating as error, same behavior as sign
190+
return
191+
}
178192
ec.handleKeygenSessionError(walletID, err, "Failed to create ECDSA key generation session", natMsg)
179193
return
180194
}
181-
eddsaSession, err := ec.node.CreateKeyGenSession(mpc.SessionTypeEDDSA, walletID, ec.mpcThreshold, ec.genKeyResultQueue)
195+
eddsaSession, err := ec.node.CreateKeyGenSession(mpc.SessionTypeEDDSA, walletID, ec.mpcThreshold, ec.genKeyResultQueue, nodeIDs)
182196
if err != nil {
197+
if errors.Is(err, mpc.ErrNotInParticipantList) {
198+
logger.Info("Node is not in selected node list for this wallet, skipping keygen",
199+
"walletID", walletID,
200+
"nodeID", ec.node.ID(),
201+
)
202+
// Skip keygen instead of treating as error, same behavior as sign
203+
return
204+
}
183205
ec.handleKeygenSessionError(walletID, err, "Failed to create EdDSA key generation session", natMsg)
184206
return
185207
}

‎pkg/mpc/node.go‎

Lines changed: 40 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -83,6 +83,7 @@ func (p *Node) CreateKeyGenSession(
8383
walletID string,
8484
threshold int,
8585
resultQueue messaging.MessageQueue,
86+
selectedNodeIDs []string,
8687
) (KeyGenSession, error) {
8788
if !p.peerRegistry.ArePeersReady() {
8889
return nil, errors.New("All nodes are not ready!")
@@ -93,24 +94,43 @@ func (p *Node) CreateKeyGenSession(
9394
return nil, fmt.Errorf("Key already exists: %s", walletID)
9495
}
9596

97+
var nodeIDs []string
98+
if len(selectedNodeIDs) > 0 {
99+
// Use selected node IDs, but filter to only include ready peers
100+
readyPeerIDs := p.peerRegistry.GetReadyPeersIncludeSelf()
101+
nodeIDs = p.filterReadyPeers(selectedNodeIDs, readyPeerIDs)
102+
103+
// Validate that we have enough ready nodes
104+
if len(nodeIDs) < threshold+1 {
105+
return nil, fmt.Errorf("not enough ready nodes in selected list: expected at least %d, got %d", threshold+1, len(nodeIDs))
106+
}
107+
108+
// Validate that self is in the selected list - return special error to skip, same as sign
109+
if !slices.Contains(nodeIDs, p.nodeID) {
110+
return nil, ErrNotInParticipantList
111+
}
112+
} else {
113+
// Use all ready peers if no selection provided
114+
nodeIDs = p.peerRegistry.GetReadyPeersIncludeSelf()
115+
}
116+
96117
switch sessionType {
97118
case SessionTypeECDSA:
98-
return p.createECDSAKeyGenSession(walletID, threshold, DefaultVersion, resultQueue)
119+
return p.createECDSAKeyGenSession(walletID, threshold, DefaultVersion, resultQueue, nodeIDs)
99120
case SessionTypeEDDSA:
100-
return p.createEDDSAKeyGenSession(walletID, threshold, DefaultVersion, resultQueue)
121+
return p.createEDDSAKeyGenSession(walletID, threshold, DefaultVersion, resultQueue, nodeIDs)
101122
default:
102123
return nil, fmt.Errorf("Unknown session type: %s", sessionType)
103124
}
104125
}
105126

106-
func (p *Node) createECDSAKeyGenSession(walletID string, threshold int, version int, resultQueue messaging.MessageQueue) (KeyGenSession, error) {
107-
readyPeerIDs := p.peerRegistry.GetReadyPeersIncludeSelf()
108-
selfPartyID, allPartyIDs := p.generatePartyIDs(PurposeKeygen, readyPeerIDs, version)
127+
func (p *Node) createECDSAKeyGenSession(walletID string, threshold int, version int, resultQueue messaging.MessageQueue, nodeIDs []string) (KeyGenSession, error) {
128+
selfPartyID, allPartyIDs := p.generatePartyIDs(PurposeKeygen, nodeIDs, version)
109129
session := newECDSAKeygenSession(
110130
walletID,
111131
p.pubSub,
112132
p.direct,
113-
readyPeerIDs,
133+
nodeIDs,
114134
selfPartyID,
115135
allPartyIDs,
116136
threshold,
@@ -123,14 +143,13 @@ func (p *Node) createECDSAKeyGenSession(walletID string, threshold int, version
123143
return session, nil
124144
}
125145

126-
func (p *Node) createEDDSAKeyGenSession(walletID string, threshold int, version int, resultQueue messaging.MessageQueue) (KeyGenSession, error) {
127-
readyPeerIDs := p.peerRegistry.GetReadyPeersIncludeSelf()
128-
selfPartyID, allPartyIDs := p.generatePartyIDs(PurposeKeygen, readyPeerIDs, version)
146+
func (p *Node) createEDDSAKeyGenSession(walletID string, threshold int, version int, resultQueue messaging.MessageQueue, nodeIDs []string) (KeyGenSession, error) {
147+
selfPartyID, allPartyIDs := p.generatePartyIDs(PurposeKeygen, nodeIDs, version)
129148
session := newEDDSAKeygenSession(
130149
walletID,
131150
p.pubSub,
132151
p.direct,
133-
readyPeerIDs,
152+
nodeIDs,
134153
selfPartyID,
135154
allPartyIDs,
136155
threshold,
@@ -250,6 +269,17 @@ func (p *Node) getReadyPeersForSession(keyInfo *keyinfo.KeyInfo, readyPeers []st
250269
return readyParticipantIDs
251270
}
252271

272+
func (p *Node) filterReadyPeers(selectedNodeIDs []string, readyPeers []string) []string {
273+
// Filter selected node IDs to only include ready peers
274+
filtered := make([]string, 0, len(selectedNodeIDs))
275+
for _, nodeID := range selectedNodeIDs {
276+
if slices.Contains(readyPeers, nodeID) {
277+
filtered = append(filtered, nodeID)
278+
}
279+
}
280+
return filtered
281+
}
282+
253283
func (p *Node) ensureNodeIsParticipant(keyInfo *keyinfo.KeyInfo) error {
254284
if !slices.Contains(keyInfo.ParticipantPeerIDs, p.nodeID) {
255285
return ErrNotInParticipantList

‎pkg/types/initiator_msg.go‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@ type InitiatorMessage interface {
3636

3737
type GenerateKeyMessage struct {
3838
WalletID string `json:"wallet_id"`
39+
SelectedNodeIDs []string `json:"selected_node_ids"`
3940
Signature []byte `json:"signature"`
4041
AuthorizerSignatures []AuthorizerSignature `json:"authorizer_signatures,omitempty"`
4142
}

0 commit comments

Comments
 (0)