Code refactoring for identity store

This commit is contained in:
anhthii
2025-04-29 19:21:56 +07:00
parent 84d7fede39
commit db4785e76e
8 changed files with 174 additions and 52 deletions
+8 -1
View File
@@ -49,11 +49,17 @@ func main() {
if err != nil {
logger.Fatal("Failed to decrypt GPG file", err)
}
viper.SetConfigType("yaml")
if err := viper.MergeConfig(bytes.NewBuffer(quax)); err != nil {
logger.Fatal("failed to merge secret config: %w", err)
}
} else {
viper.SetConfigFile("./quax.yaml")
if err := viper.MergeInConfig(); err != nil {
logger.Fatal("Failed to merge quax.yaml", err)
}
logger.Info("Merge quax.yaml successfully")
}
nodeName := flag.String("name", "", "Provide node name")
@@ -132,6 +138,7 @@ func main() {
pubsub,
genKeySuccessQueue,
singingResultQueue,
identityStore,
)
eventConsumer.Run()
defer eventConsumer.Close()
-12
View File
@@ -1,12 +0,0 @@
nats:
url: tls://103.2.230.61:4223
username: apex
password: ""
consul:
address: https://apex-consul.void.exchange
username: apex
token: ""
password: ""
badger_password:
-1
View File
@@ -4,4 +4,3 @@ consul:
address: localhost:8500
mpc_threshold: 2
badger_password: 1JwFmsc9lxlLfkPl
+31 -6
View File
@@ -11,9 +11,11 @@ import (
"time"
"github.com/cryptoniumX/mpcium/pkg/event"
"github.com/cryptoniumX/mpcium/pkg/identity"
"github.com/cryptoniumX/mpcium/pkg/logger"
"github.com/cryptoniumX/mpcium/pkg/messaging"
"github.com/cryptoniumX/mpcium/pkg/mpc"
"github.com/cryptoniumX/mpcium/pkg/types"
"github.com/nats-io/nats.go"
"github.com/spf13/viper"
)
@@ -38,6 +40,7 @@ type eventConsumer struct {
keyGenerationSub messaging.Subscription
signingSub messaging.Subscription
identityStore identity.Store
// Track active sessions with timestamps for cleanup
activeSessions map[string]time.Time // Maps "walletID-txID" to creation time
@@ -52,6 +55,7 @@ func NewEventConsumer(
pubsub messaging.PubSub,
genKeySucecssQueue messaging.MessageQueue,
signingResultQueue messaging.MessageQueue,
identityStore identity.Store,
) EventConsumer {
ec := &eventConsumer{
node: node,
@@ -63,6 +67,7 @@ func NewEventConsumer(
sessionTimeout: 30 * time.Minute, // Consider sessions older than 30 minutes stale
cleanupStopChan: make(chan struct{}),
mpcThreshold: viper.GetInt("mpc_threshold"),
identityStore: identityStore,
}
// Start background cleanup goroutine
@@ -87,8 +92,22 @@ func (ec *eventConsumer) Run() {
func (ec *eventConsumer) consumeKeyGenerationEvent() error {
sub, err := ec.pubsub.Subscribe(MPCGenerateEvent, func(natMsg *nats.Msg) {
msg := natMsg.Data
walletID := string(msg)
raw := natMsg.Data
var msg types.GenerateKeyMessage
err := json.Unmarshal(raw, &msg)
if err != nil {
logger.Error("Failed to unmarshal signing message", err)
return
}
logger.Info("Received key generation event", "msg", msg)
err = ec.identityStore.VerifyInitiatorMessage(&msg)
if err != nil {
logger.Error("Failed to verify initiator message", err)
return
}
walletID := msg.WalletID
session, err := ec.node.CreateKeyGenSession(walletID, ec.mpcThreshold, ec.genKeySucecssQueue)
if err != nil {
logger.Error("Failed to create key generation session", err, "walletID", walletID)
@@ -116,7 +135,7 @@ func (ec *eventConsumer) consumeKeyGenerationEvent() error {
for {
select {
case <-ctx.Done():
successEvent.S256PubKey = session.GetPubKeyResult()
successEvent.ECDSAPubKey = session.GetPubKeyResult()
wg.Done()
return
case err := <-session.ErrCh:
@@ -177,13 +196,19 @@ func (ec *eventConsumer) consumeKeyGenerationEvent() error {
func (ec *eventConsumer) consumeTxSigningEvent() error {
sub, err := ec.pubsub.Subscribe(MPCSignEvent, func(natMsg *nats.Msg) {
raw := natMsg.Data
var msg SignTxMessage
var msg types.SignTxMessage
err := json.Unmarshal(raw, &msg)
if err != nil {
logger.Error("Failed to unmarshal signing message", err)
return
}
err = ec.identityStore.VerifyInitiatorMessage(&msg)
if err != nil {
logger.Error("Failed to verify initiator message", err)
return
}
logger.Info(
"Received signing event",
"waleltID",
@@ -204,7 +229,7 @@ func (ec *eventConsumer) consumeTxSigningEvent() error {
var session mpc.ISigningSession
switch msg.KeyType {
case KeyTypeSecp256k1:
case types.KeyTypeSecp256k1:
session, err = ec.node.CreateSigningSession(
msg.WalletID,
msg.TxID,
@@ -212,7 +237,7 @@ func (ec *eventConsumer) consumeTxSigningEvent() error {
ec.mpcThreshold,
ec.signingResultQueue,
)
case KeyTypeEd25519:
case types.KeyTypeEd25519:
session, err = ec.node.CreateEDDSASigningSession(
msg.WalletID,
msg.TxID,
+56
View File
@@ -1,5 +1,7 @@
package eventconsumer
import "encoding/json"
type KeyType string
const (
@@ -7,10 +9,64 @@ const (
KeyTypeEd25519 = "ed25519"
)
// InitiatorMessage is anything that carries a payload to verify and its signature.
type InitiatorMessage interface {
// Raw returns the canonical byte‐slice that was signed.
Raw() ([]byte, error)
// Sig returns the signature over Raw().
Sig() []byte
// InitiatorID returns the ID whose public key we have to look up.
InitiatorID() string
}
type GenerateKeyMessage struct {
WalletID string `json:"wallet_id"`
Signature []byte `json:"signature"`
}
type SignTxMessage struct {
KeyType KeyType `json:"key_type"`
WalletID string `json:"wallet_id"`
NetworkInternalCode string `json:"network_internal_code"`
TxID string `json:"tx_id"`
Tx []byte `json:"tx"`
Signature []byte `json:"signature"`
}
func (m *SignTxMessage) Raw() ([]byte, error) {
// omit the Signature field itself when computing the signed‐over data
payload := struct {
KeyType KeyType `json:"key_type"`
WalletID string `json:"wallet_id"`
NetworkInternalCode string `json:"network_internal_code"`
TxID string `json:"tx_id"`
Tx []byte `json:"tx"`
}{
KeyType: m.KeyType,
WalletID: m.WalletID,
NetworkInternalCode: m.NetworkInternalCode,
TxID: m.TxID,
Tx: m.Tx,
}
return json.Marshal(payload)
}
func (m *SignTxMessage) Sig() []byte {
return m.Signature
}
func (m *SignTxMessage) InitiatorID() string {
return m.TxID
}
func (m *GenerateKeyMessage) Raw() ([]byte, error) {
return []byte(m.WalletID), nil
}
func (m *GenerateKeyMessage) Sig() []byte {
return m.Signature
}
func (m *GenerateKeyMessage) InitiatorID() string {
return m.WalletID
}
+77 -30
View File
@@ -7,11 +7,12 @@ import (
"fmt"
"os"
"path/filepath"
"strings"
"sync"
"github.com/bnb-chain/tss-lib/v2/tss"
"github.com/cryptoniumX/mpcium/pkg/logger"
"github.com/cryptoniumX/mpcium/pkg/types"
"github.com/spf13/viper"
)
// NodeIdentity represents a node's identity information
@@ -26,7 +27,7 @@ type NodeIdentity struct {
type Store interface {
// GetPublicKey retrieves a node's public key by its ID
GetPublicKey(nodeID string) ([]byte, error)
VerifyInitiatorMessage(msg types.InitiatorMessage) error
SignMessage(msg *types.TssMessage) ([]byte, error)
VerifyMessage(msg *types.TssMessage) error
}
@@ -41,7 +42,8 @@ type fileStore struct {
mu sync.RWMutex
// Cached private key
privateKey []byte
privateKey []byte
initiatorPubKey []byte
}
// NewFileStore creates a new identity store
@@ -50,51 +52,74 @@ func NewFileStore(identityDir, nodeName string) (*fileStore, error) {
return nil, fmt.Errorf("failed to create identity directory: %w", err)
}
// Load private key immediately
// Load private key from file
privateKeyPath := filepath.Join(identityDir, fmt.Sprintf("%s_private.key", nodeName))
privateKeyData, err := os.ReadFile(privateKeyPath)
if err != nil {
return nil, fmt.Errorf("failed to read private key file: %w", err)
}
privateKey, err := hex.DecodeString(string(privateKeyData))
privateKeyHex := string(privateKeyData)
privateKey, err := hex.DecodeString(privateKeyHex)
if err != nil {
return nil, fmt.Errorf("invalid private key format: %w", err)
}
pubKeyHex := viper.GetString("event_initiator_pubkey")
if pubKeyHex == "" {
return nil, fmt.Errorf("event_initiator_pubkey not found in quax config")
}
initiatorPubKey, err := hex.DecodeString(pubKeyHex)
if err != nil {
return nil, fmt.Errorf("invalid initiator public key format: %w", err)
}
logger.Infof("Loaded initiator public key for node %s", pubKeyHex)
// Load peers.json to validate all nodes have identity files
peersData, err := os.ReadFile("peers.json")
if err != nil {
return nil, fmt.Errorf("failed to read peers.json: %w", err)
}
peers := make(map[string]string)
if err := json.Unmarshal(peersData, &peers); err != nil {
return nil, fmt.Errorf("failed to parse peers.json: %w", err)
}
store := &fileStore{
identityDir: identityDir,
currentNodeName: nodeName,
publicKeys: make(map[string][]byte),
privateKey: privateKey,
initiatorPubKey: initiatorPubKey,
}
// Load all identity files
files, err := os.ReadDir(identityDir)
if err != nil {
return nil, fmt.Errorf("failed to read identity directory: %w", err)
}
for _, file := range files {
if !file.IsDir() && strings.HasSuffix(file.Name(), "_identity.json") {
path := filepath.Join(identityDir, file.Name())
data, err := os.ReadFile(path)
if err != nil {
return nil, fmt.Errorf("failed to read identity file %s: %w", file.Name(), err)
}
var identity NodeIdentity
if err := json.Unmarshal(data, &identity); err != nil {
return nil, fmt.Errorf("failed to parse identity file %s: %w", file.Name(), err)
}
key, err := hex.DecodeString(identity.PublicKey)
if err != nil {
return nil, fmt.Errorf("invalid public key format in %s: %w", file.Name(), err)
}
store.publicKeys[identity.NodeID] = key
// Check that each node in peers.json has an identity file
for nodeName, nodeID := range peers {
identityFilePath := filepath.Join(identityDir, fmt.Sprintf("%s_identity.json", nodeName))
data, err := os.ReadFile(identityFilePath)
if err != nil {
return nil, fmt.Errorf("missing identity file for node %s (%s): %w", nodeName, nodeID, err)
}
var identity NodeIdentity
if err := json.Unmarshal(data, &identity); err != nil {
return nil, fmt.Errorf("failed to parse identity file for node %s: %w", nodeName, err)
}
// Verify that the nodeID in peers.json matches the one in the identity file
if identity.NodeID != nodeID {
return nil, fmt.Errorf("node ID mismatch for %s: %s in peers.json vs %s in identity file",
nodeName, nodeID, identity.NodeID)
}
key, err := hex.DecodeString(identity.PublicKey)
if err != nil {
return nil, fmt.Errorf("invalid public key format for node %s: %w", nodeName, err)
}
store.publicKeys[identity.NodeID] = key
}
return store, nil
@@ -152,6 +177,28 @@ func (s *fileStore) VerifyMessage(msg *types.TssMessage) error {
return nil
}
// VerifyInitiatorMessage verifies that a message was signed by the known initiator
func (s *fileStore) VerifyInitiatorMessage(msg types.InitiatorMessage) error {
// Get the raw message that was signed
msgBytes, err := msg.Raw()
if err != nil {
return fmt.Errorf("failed to get raw message data: %w", err)
}
// Get the signature
signature := msg.Sig()
if signature == nil || len(signature) == 0 {
return fmt.Errorf("message has no signature")
}
// Verify the signature using the initiator's public key
if !ed25519.Verify(s.initiatorPubKey, msgBytes, signature) {
return fmt.Errorf("invalid signature from initiator")
}
return nil
}
func partyIDToNodeID(partyID *tss.PartyID) string {
return string(partyID.KeyInt().Bytes())
}
+1 -1
View File
@@ -26,7 +26,7 @@ type KeygenSession struct {
type KeygenSuccessEvent struct {
WalletID string `json:"wallet_id"`
S256PubKey []byte `json:"s256_pub_key"`
ECDSAPubKey []byte `json:"ecdsa_pub_key"`
EDDSAPubKey []byte `json:"eddsa_pub_key"`
}
+1 -1
View File
@@ -1,3 +1,3 @@
badger_password: 1JwFmsc9lxlLfkPd
private_key: ebe2c553436606d0ed24a96eb7ca5bb70ece2f3ae11ca2fab80c4a0c8bfda62bf39da9e4afcfeca1641d4c0e95967b730740fad5fff1c77a74d94bb8572b2166
event_initiator_pubkey: 0x
event_initiator_pubkey: e2eb8eb766eb067f554f42ffbc3d6e6cf63645c6b15c90c056e139f94e1e51e6