diff --git a/e2e/base_test.go b/e2e/base_test.go index 194b5f1b..1e68de25 100644 --- a/e2e/base_test.go +++ b/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 == "" { diff --git a/e2e/keygen_test.go b/e2e/keygen_test.go index e4599161..387876aa 100644 --- a/e2e/keygen_test.go +++ b/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) { diff --git a/e2e/multi_client_routing_test.go b/e2e/multi_client_routing_test.go index d28d2e7d..83a23f5f 100644 --- a/e2e/multi_client_routing_test.go +++ b/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) diff --git a/e2e/setup_test_identities.sh b/e2e/setup_test_identities.sh index 00c2dfae..1cbd7861 100755 --- a/e2e/setup_test_identities.sh +++ b/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) diff --git a/pkg/eventconsumer/event_consumer.go b/pkg/eventconsumer/event_consumer.go index 4b94e2d6..f19d646e 100644 --- a/pkg/eventconsumer/event_consumer.go +++ b/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, diff --git a/pkg/mpc/node.go b/pkg/mpc/node.go index e28cf3f0..3910b6dd 100644 --- a/pkg/mpc/node.go +++ b/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) } diff --git a/setup.sh b/setup.sh index 7be28d39..37ad034e 100755 --- a/setup.sh +++ b/setup.sh @@ -1,3 +1,5 @@ +export PATH="$(go env GOPATH)/bin:$PATH" + NUM_NODES=3 make build