Worked. Update generate examples to test with p256

Author: tuananhk13bk

.gitignore

@@ -21,4 +21,4 @@ node0
 node1
 node2
 config.yaml
-peers.json
+.vscode

cmd/mpcium-cli/generate-initiator.go

@@ -2,11 +2,8 @@ package main
 
 import (
 	"context"
-	"crypto/ecdsa"
 	"crypto/ed25519"
-	"crypto/elliptic"
 	"crypto/rand"
-	"crypto/x509"
 	"encoding/hex"
 	"encoding/json"
 	"fmt"
@@ -18,6 +15,7 @@ import (
 
 	"filippo.io/age"
 	"github.com/fystack/mpcium/pkg/common/pathutil"
+	"github.com/fystack/mpcium/pkg/encryption"
 	"github.com/urfave/cli/v3"
 )
 
@@ -32,12 +30,6 @@ type InitiatorIdentity struct {
 	MachineName string `json:"machine_name"`
 }
 
-// KeyData holds the generated key information
-type KeyData struct {
-	PublicKeyHex  string
-	PrivateKeyHex string
-}
-
 func generateInitiatorIdentity(ctx context.Context, c *cli.Command) error {
 	nodeName := c.String("node-name")
 	outputDir := c.String("output-dir")
@@ -86,13 +78,13 @@ func generateInitiatorIdentity(ctx context.Context, c *cli.Command) error {
 	}
 
 	// Generate keys based on algorithm
-	var keyData KeyData
+	var keyData encryption.KeyData
 	var err error
 
 	if algorithm == "ed25519" {
 		keyData, err = generateEd25519Keys()
 	} else {
-		keyData, err = generateP256Keys()
+		keyData, err = encryption.GenerateP256Keys()
 	}
 
 	if err != nil {
@@ -194,40 +186,15 @@ func generateInitiatorIdentity(ctx context.Context, c *cli.Command) error {
 }
 
 // generateEd25519Keys generates Ed25519 keypair
-func generateEd25519Keys() (KeyData, error) {
+func generateEd25519Keys() (encryption.KeyData, error) {
 	pubKey, privKeyFull, err := ed25519.GenerateKey(rand.Reader)
 	if err != nil {
-		return KeyData{}, err
+		return encryption.KeyData{}, err
 	}
 
 	privKeySeed := privKeyFull.Seed()
-	return KeyData{
+	return encryption.KeyData{
 		PublicKeyHex:  hex.EncodeToString(pubKey),
 		PrivateKeyHex: hex.EncodeToString(privKeySeed),
 	}, nil
 }
-
-// generateP256Keys generates P-256 keypair
-func generateP256Keys() (KeyData, error) {
-	privateKey, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
-	if err != nil {
-		return KeyData{}, err
-	}
-
-	// Convert private key to PEM format
-	privateKeyBytes, err := x509.MarshalECPrivateKey(privateKey)
-	if err != nil {
-		return KeyData{}, err
-	}
-
-	// Convert public key to DER format
-	publicKeyBytes, err := x509.MarshalPKIXPublicKey(&privateKey.PublicKey)
-	if err != nil {
-		return KeyData{}, err
-	}
-
-	return KeyData{
-		PublicKeyHex:  hex.EncodeToString(publicKeyBytes),
-		PrivateKeyHex: hex.EncodeToString(privateKeyBytes),
-	}, nil
-}

examples/generate/main.go

@@ -23,11 +23,25 @@ import (
 func main() {
 	const environment = "development"
 	numWallets := flag.Int("n", 1, "Number of wallets to generate")
+
 	flag.Parse()
 
 	config.InitViperConfig()
 	logger.Init(environment, false)
 
+	algorithm := viper.GetString("event_initiator_algorithm")
+	if algorithm == "" {
+		algorithm = "ed25519"
+	}
+
+	// Validate algorithm
+	if algorithm != "ed25519" && algorithm != "p256" {
+		logger.Fatal(
+			"Invalid event_initiator_algorithm in config. Must be 'ed25519' or 'p256'",
+			nil,
+		)
+	}
+
 	natsURL := viper.GetString("nats.url")
 	natsConn, err := nats.Connect(natsURL)
 	if err != nil {
@@ -37,8 +51,9 @@ func main() {
 	defer natsConn.Close()
 
 	mpcClient := client.NewMPCClient(client.Options{
-		NatsConn: natsConn,
-		KeyPath:  "./event_initiator.key",
+		Algorithm: algorithm,
+		NatsConn:  natsConn,
+		KeyPath:   "./event_initiator.key",
 	})
 
 	var walletStartTimes sync.Map
@@ -89,13 +104,13 @@ func main() {
 
 	// STEP 3: Create wallets
 	for _, walletID := range walletIDs {
-		wg.Add(1)
 		if err := mpcClient.CreateWallet(walletID); err != nil {
 			logger.Error("CreateWallet failed", err)
 			walletStartTimes.Delete(walletID)
 			wg.Done()
 			continue
 		}
+		wg.Add(1)
 		logger.Info("CreateWallet sent, awaiting result...", "walletID", walletID)
 	}
 

pkg/client/client.go

@@ -2,6 +2,7 @@ package client
 
 import (
 	"context"
+	"crypto/ecdsa"
 	"crypto/ed25519"
 	"encoding/hex"
 	"encoding/json"
@@ -12,6 +13,7 @@ import (
 	"strings"
 
 	"filippo.io/age"
+	"github.com/fystack/mpcium/pkg/encryption"
 	"github.com/fystack/mpcium/pkg/event"
 	"github.com/fystack/mpcium/pkg/eventconsumer"
 	"github.com/fystack/mpcium/pkg/logger"
@@ -36,14 +38,23 @@ type MPCClient interface {
 	OnResharingResult(callback func(event event.ResharingResultEvent)) error
 }
 
+type InitiatorPrivKey struct {
+	Algorithm types.KeyType
+	Ed25519   ed25519.PrivateKey
+	P256      ecdsa.PrivateKey
+}
+
 type mpcClient struct {
 	signingBroker       messaging.MessageBroker
 	keygenBroker        messaging.MessageBroker
 	pubsub              messaging.PubSub
 	genKeySuccessQueue  messaging.MessageQueue
 	signResultQueue     messaging.MessageQueue
 	reshareSuccessQueue messaging.MessageQueue
-	privKey             ed25519.PrivateKey
+	// privKey             ed25519.PrivateKey
+	// privKeyECDSA        *ecdsa.PrivateKey
+	initiatorPrivKey *InitiatorPrivKey
+	algorithm        string
 }
 
 // Options defines configuration options for creating a new MPCClient
@@ -57,10 +68,13 @@ type Options struct {
 	// Encryption options
 	Encrypted bool   // Whether the key is encrypted
 	Password  string // Password for encrypted key
+
+	// Algorithm for key type
+	Algorithm string // Either "ed25519" or "p256" (default: "ed25519")
 }
 
 // NewMPCClient creates a new MPC client using the provided options.
-// It reads the Ed25519 private key from disk and sets up messaging connections.
+// It reads the Ed25519 or P256 private key from disk and sets up messaging connections.
 // If the key is encrypted (.age file), decryption options must be provided in the config.
 func NewMPCClient(opts Options) MPCClient {
 	// Set default paths if not provided
@@ -105,16 +119,28 @@ func NewMPCClient(opts Options) MPCClient {
 		logger.Fatal("No private key file found", nil)
 	}
 
-	privHex := string(privHexBytes)
-	// Decode private key from hex
-	privSeed, err := hex.DecodeString(privHex)
-	if err != nil {
-		fmt.Println("Failed to decode private key hex:", err)
-		os.Exit(1)
-	}
+	var priv ed25519.PrivateKey
+	var privECDSA *ecdsa.PrivateKey
+
+	if opts.Algorithm == "p256" {
+		// Parse P256 key
+		privECDSA, err = encryption.ParseP256PrivateKey(privHexBytes)
+		if err != nil {
+			logger.Fatal("Failed to parse P256 private key", err)
+		}
+	} else {
+		// Parse Ed25519 key (default behavior)
+		privHex := string(privHexBytes)
+		// Decode private key from hex
+		privSeed, err := hex.DecodeString(privHex)
+		if err != nil {
+			fmt.Println("Failed to decode private key hex:", err)
+			os.Exit(1)
+		}
 
-	// Reconstruct full Ed25519 private key from seed
-	priv := ed25519.NewKeyFromSeed(privSeed)
+		// Reconstruct full Ed25519 private key from seed
+		priv = ed25519.NewKeyFromSeed(privSeed)
+	}
 
 	// 2) Create the PubSub for both publish & subscribe
 	signingBroker, err := messaging.NewJetStreamBroker(
@@ -159,7 +185,14 @@ func NewMPCClient(opts Options) MPCClient {
 		genKeySuccessQueue:  genKeySuccessQueue,
 		signResultQueue:     signResultQueue,
 		reshareSuccessQueue: reshareSuccessQueue,
-		privKey:             priv,
+		// privKey:             priv,
+		// privKeyECDSA:        privECDSA,
+		initiatorPrivKey: &InitiatorPrivKey{
+			Algorithm: types.KeyType(opts.Algorithm),
+			Ed25519:   priv,
+			P256:      *privECDSA,
+		},
+		algorithm: opts.Algorithm,
 	}
 }
 
@@ -197,8 +230,24 @@ func (c *mpcClient) CreateWallet(walletID string) error {
 	if err != nil {
 		return fmt.Errorf("CreateWallet: raw payload error: %w", err)
 	}
-	// sign
-	msg.Signature = ed25519.Sign(c.privKey, raw)
+	logger.Info("Raw payload for signing", "raw", string(raw), "raw_bytes", raw)
+	// sign based on algorithm
+	var signature []byte
+	if c.algorithm == "p256" {
+		if c.initiatorPrivKey.P256.Curve == nil {
+			return fmt.Errorf("CreateWallet: P256 private key not initialized")
+		}
+		signature, err = encryption.SignWithP256(&c.initiatorPrivKey.P256, raw)
+		if err != nil {
+			return fmt.Errorf("CreateWallet: failed to create P256 signature: %w", err)
+		}
+		if signature == nil {
+			return fmt.Errorf("CreateWallet: failed to create P256 signature")
+		}
+	} else {
+		signature = ed25519.Sign(c.initiatorPrivKey.Ed25519, raw)
+	}
+	msg.Signature = signature
 
 	bytes, err := json.Marshal(msg)
 	if err != nil {
@@ -237,8 +286,23 @@ func (c *mpcClient) SignTransaction(msg *types.SignTxMessage) error {
 	if err != nil {
 		return fmt.Errorf("SignTransaction: raw payload error: %w", err)
 	}
-	// sign
-	msg.Signature = ed25519.Sign(c.privKey, raw)
+	// sign based on algorithm
+	var signature []byte
+	if c.algorithm == "p256" {
+		if c.initiatorPrivKey.P256.Curve == nil {
+			return fmt.Errorf("SignTransaction: P256 private key not initialized")
+		}
+		signature, err = encryption.SignWithP256(&c.initiatorPrivKey.P256, raw)
+		if err != nil {
+			return fmt.Errorf("SignTransaction: failed to create P256 signature: %w", err)
+		}
+		if signature == nil {
+			return fmt.Errorf("SignTransaction: failed to create P256 signature")
+		}
+	} else {
+		signature = ed25519.Sign(c.initiatorPrivKey.Ed25519, raw)
+	}
+	msg.Signature = signature
 
 	bytes, err := json.Marshal(msg)
 	if err != nil {
@@ -275,8 +339,23 @@ func (c *mpcClient) Resharing(msg *types.ResharingMessage) error {
 	if err != nil {
 		return fmt.Errorf("Resharing: raw payload error: %w", err)
 	}
-	// sign
-	msg.Signature = ed25519.Sign(c.privKey, raw)
+	// sign based on algorithm
+	var signature []byte
+	if c.algorithm == "p256" {
+		if c.initiatorPrivKey.P256.Curve == nil {
+			return fmt.Errorf("Resharing: P256 private key not initialized")
+		}
+		signature, err = encryption.SignWithP256(&c.initiatorPrivKey.P256, raw)
+		if err != nil {
+			return fmt.Errorf("Resharing: failed to create P256 signature: %w", err)
+		}
+		if signature == nil {
+			return fmt.Errorf("Resharing: failed to create P256 signature")
+		}
+	} else {
+		signature = ed25519.Sign(c.initiatorPrivKey.Ed25519, raw)
+	}
+	msg.Signature = signature
 
 	bytes, err := json.Marshal(msg)
 	if err != nil {