Merge pull request #160 from johnnynanjiang/issue/140-round-1

PR for issue #140 - Missing Wallet Query Interface in SDK to Retrieve Created Wallet Information

Author: vietddude

e2e/base_test.go

@@ -80,7 +80,20 @@ func (s *E2ETestSuite) LoadConfig() error {
 		return err
 	}
 
-	return yaml.Unmarshal(data, &s.config)
+	if err := yaml.Unmarshal(data, &s.config); err != nil {
+		return err
+	}
+
+	return validateBadgerEncryptionKey(s.config.BadgerPassword)
+}
+
+func validateBadgerEncryptionKey(key string) error {
+	switch len([]byte(key)) {
+	case 16, 24, 32:
+		return nil
+	default:
+		return fmt.Errorf("badger_password must be 16, 24, or 32 bytes for Badger encryption, got %d bytes", len([]byte(key)))
+	}
 }
 
 func (s *E2ETestSuite) RunMakeClean() error {
@@ -284,6 +297,7 @@ func (s *E2ETestSuite) SeedPreParams(t *testing.T) {
 		DbPath         string `yaml:"db_path"`
 	}
 	require.NoError(t, yaml.Unmarshal(configData, &nodeConfig), "Failed to parse node config")
+	require.NoError(t, validateBadgerEncryptionKey(nodeConfig.BadgerPassword), "Invalid badger_password in %s", configPath)
 
 	dbBasePath := nodeConfig.DbPath
 	if dbBasePath == "" {

e2e/keygen_test.go

@@ -14,6 +14,7 @@ import (
 
 func TestKeyGeneration(t *testing.T) {
 	suite := NewE2ETestSuite(".")
+	logger.Init("dev", true)
 
 	// Comprehensive cleanup before starting tests
 	t.Log("Performing pre-test cleanup...")
@@ -78,10 +79,23 @@ func testKeyGeneration(t *testing.T, suite *E2ETestSuite) {
 	if suite.mpcClient == nil {
 		t.Fatal("MPC client is not initialized. Make sure Setup subtest runs first.")
 	}
+
+	walletID := uuid.New().String()
+
+	// Create a new wallet
+	createWalletAndVerify(t, suite, walletID)
+
+	// Create another new wallet with the same wallet ID
+	createWalletAndVerify(t, suite, walletID)
+
+	t.Log("Key generation test completed")
+}
+
+func createWalletAndVerify(t *testing.T, suite *E2ETestSuite, walletID string) {
 	// Generate 1 wallet ID for testing
 	walletIDs := make([]string, 0, 10)
 	for i := 0; i < 1; i++ {
-		walletIDs = append(walletIDs, uuid.New().String())
+		walletIDs = append(walletIDs, walletID)
 		suite.walletIDs = append(suite.walletIDs, walletIDs[i])
 	}
 
@@ -91,7 +105,11 @@ func testKeyGeneration(t *testing.T, suite *E2ETestSuite) {
 	err := suite.mpcClient.OnWalletCreationResult(func(result event.KeygenResultEvent) {
 		logger.Info("On wallet creation result", "event", result)
 		t.Logf("Received keygen result for wallet %s: %s", result.WalletID, result.ResultType)
-		suite.keygenResults[result.WalletID] = &result
+
+		// For testing replay properly, don't overwrite existing result for the same wallet ID to preserve the first result
+		if _, exists := suite.keygenResults[result.WalletID]; !exists {
+			suite.keygenResults[result.WalletID] = &result
+		}
 
 		if result.ResultType == event.ResultTypeError {
 			t.Logf("Keygen failed for wallet %s: %s (%s)", result.WalletID, result.ErrorReason, result.ErrorCode)
@@ -164,8 +182,6 @@ checkResults:
 			assert.NotEmpty(t, result.EDDSAPubKey, "EdDSA public key should not be empty for wallet %s", walletID)
 		}
 	}
-
-	t.Log("Key generation test completed")
 }
 
 func verifyKeyConsistency(t *testing.T, suite *E2ETestSuite) {

e2e/multi_client_routing_test.go

@@ -11,6 +11,7 @@ import (
 
 	"github.com/fystack/mpcium/pkg/client"
 	"github.com/fystack/mpcium/pkg/event"
+	"github.com/fystack/mpcium/pkg/logger"
 	"github.com/fystack/mpcium/pkg/types"
 	"github.com/google/uuid"
 	"github.com/nats-io/nats.go"
@@ -36,6 +37,7 @@ type multiClientObserver struct {
 
 func TestMultiClientResultRouting(t *testing.T) {
 	suite := NewE2ETestSuite(".")
+	logger.Init("dev", true)
 
 	t.Log("Performing pre-test cleanup...")
 	suite.CleanupTestEnvironment(t)

e2e/setup_test_identities.sh

@@ -25,7 +25,17 @@ echo "🚀 Setting up E2E Test Node Identities..."
 
 # Generate random password for badger encryption
 echo "🔐 Generating random password for badger encryption..."
-BADGER_PASSWORD=$(< /dev/urandom tr -dc 'A-Za-z0-9' | head -c 32)
+if [[ "$OSTYPE" == "darwin"* ]]; then
+    # macOS
+    BADGER_PASSWORD=$(openssl rand -hex 16)
+else
+    # Linux and others
+    BADGER_PASSWORD=$(< /dev/urandom tr -dc 'A-Za-z0-9' | head -c 32)
+fi
+if [ ${#BADGER_PASSWORD} -ne 32 ]; then
+    echo "❌ Generated Badger password must be exactly 32 bytes, got ${#BADGER_PASSWORD}"
+    exit 1
+fi
 echo "✅ Generated password: $BADGER_PASSWORD"
 
 # Generate chain_code (32-byte hex value, 64 hex characters)

pkg/eventconsumer/event_consumer.go

@@ -10,6 +10,7 @@ import (
 	"sync"
 	"time"
 
+	"github.com/dgraph-io/badger/v4"
 	"github.com/fystack/mpcium/pkg/event"
 	"github.com/fystack/mpcium/pkg/identity"
 	"github.com/fystack/mpcium/pkg/logger"
@@ -158,6 +159,31 @@ func (ec *eventConsumer) handleKeyGenEvent(natMsg *nats.Msg) {
 
 	walletID := msg.WalletID
 
+	// Attempt to get previously stored wallet creation result (if any) by the wallet ID
+	storedWalletCreationResult, storedWalletCreationResultError := ec.node.GetWalletCreationResult(walletID)
+
+	// Error when retrieving wallet creation result for the wallet ID
+	if storedWalletCreationResultError != nil && !errors.Is(storedWalletCreationResultError, badger.ErrKeyNotFound) {
+		ec.handleKeygenSessionError(walletID, storedWalletCreationResultError, "Failed to check stored wallet creation result", natMsg)
+		return
+	}
+
+	// Replay the wallet creation result if it already exists for the wallet ID
+	// TODO: to move the following duplicate logic (line 308 and line 341) into a func
+	if storedWalletCreationResult != nil {
+		key := event.KeygenResultSubject(natMsg.Header.Get(event.ClientIDHeader), walletID)
+		if err := ec.genKeyResultQueue.Enqueue(key, storedWalletCreationResult, &messaging.EnqueueOptions{
+			IdempotententKey: composeKeygenIdempotentKey(walletID, natMsg),
+		}); err != nil {
+			logger.Error("Failed to enqueue stored wallet creation result", err, "walletID", walletID)
+			ec.handleKeygenSessionError(walletID, err, "Failed to enqueue stored wallet creation result", natMsg)
+			return
+		}
+		ec.sendReplyToRemoveMsg(natMsg)
+		logger.Info("Returned stored wallet creation result for existing wallet", "walletID", walletID)
+		return
+	}
+
 	// Guard against duplicate keygen sessions for the same walletID.
 	// Under heavy load, the keygen consumer may NAK and JetStream redelivers,
 	// creating a second session on the same NATS topics which causes VSS verify failures.
@@ -272,6 +298,13 @@ func (ec *eventConsumer) handleKeyGenEvent(natMsg *nats.Msg) {
 		return
 	}
 
+	// Store wallet creation result
+	if storeErr := ec.node.StoreWalletCreationResult(walletID, payload); storeErr != nil {
+		logger.Error("Failed to store wallet creation result", storeErr, "walletID", walletID)
+		ec.handleKeygenSessionError(walletID, storeErr, "Failed to store wallet creation result", natMsg)
+		return
+	}
+
 	key := event.KeygenResultSubject(natMsg.Header.Get(event.ClientIDHeader), walletID)
 	if err := ec.genKeyResultQueue.Enqueue(
 		key,

pkg/mpc/node.go

@@ -409,6 +409,18 @@ func (p *Node) CreateReshareSession(
 	}
 }
 
+const walletCreationResultPrefix = "wallet_creation_result_prefix"
+
+func (p *Node) StoreWalletCreationResult(walletID string, result []byte) error {
+	key := fmt.Sprintf("%s:%s", walletCreationResultPrefix, walletID)
+	return p.kvstore.Put(key, result)
+}
+
+func (p *Node) GetWalletCreationResult(walletID string) ([]byte, error) {
+	key := fmt.Sprintf("%s:%s", walletCreationResultPrefix, walletID)
+	return p.kvstore.Get(key)
+}
+
 func ComposeReadyKey(nodeID string) string {
 	return fmt.Sprintf("ready/%s", nodeID)
 }

setup.sh

@@ -1,3 +1,5 @@
+export PATH="$(go env GOPATH)/bin:$PATH"
+
 NUM_NODES=3
 
 make build