From b1bfba680a28e6782121e0c02acbbca7e96b9efe Mon Sep 17 00:00:00 2001
From: Alok <aloknerurkar@no-reply.com>
Date: Sat, 14 Dec 2024 02:43:47 +0530
Subject: [PATCH] chore: fix stressnet with new BLS keys

---
 p2p/integrationtest/provider/client.go | 24 +++++++++++++-----------
 p2p/integrationtest/provider/main.go   | 14 +++-----------
 2 files changed, 16 insertions(+), 22 deletions(-)

diff --git a/p2p/integrationtest/provider/client.go b/p2p/integrationtest/provider/client.go
index 80b21c400..737e9cd2e 100644
--- a/p2p/integrationtest/provider/client.go
+++ b/p2p/integrationtest/provider/client.go
@@ -99,17 +99,17 @@ func (b *ProviderClient) Close() error {
 	return b.conn.Close()
 }
 
-func (b *ProviderClient) CheckAndStake(_ []string) error {
+func (b *ProviderClient) CheckAndStake() (string, error) {
 	stakeAmt, err := b.client.GetStake(context.Background(), &providerapiv1.EmptyMessage{})
 	if err != nil {
 		b.logger.Error("failed to get stake amount", "err", err)
-		return err
+		return "", err
 	}
 
 	topology, err := b.debugClient.GetTopology(context.Background(), &debugapi.EmptyMessage{})
 	if err != nil {
 		b.logger.Error("failed to get topology", "err", err)
-		return err
+		return "", err
 	}
 
 	ethAddress := topology.GetTopology().Fields["self"].GetStructValue().Fields["Ethereum Address"].GetStringValue()
@@ -120,12 +120,15 @@ func (b *ProviderClient) CheckAndStake(_ []string) error {
 	stakedAmt, set := big.NewInt(0).SetString(stakeAmt.Amount, 10)
 	if !set {
 		b.logger.Error("failed to parse stake amount")
-		return errors.New("failed to parse stake amount")
+		return "", errors.New("failed to parse stake amount")
 	}
 
 	if stakedAmt.Cmp(big.NewInt(0)) > 0 {
 		b.logger.Error("bidder already staked")
-		return nil
+		if len(stakeAmt.BlsPublicKeys) == 0 {
+			return "", errors.New("bidder already staked but no BLS public keys found")
+		}
+		return stakeAmt.BlsPublicKeys[0], nil
 	}
 
 	hashedMessage := crypto.Keccak256(common.HexToAddress(ethAddress).Bytes())
@@ -133,7 +136,7 @@ func (b *ProviderClient) CheckAndStake(_ []string) error {
 	privateKey, err := bls.KeyGen[bls.G1](ikm, nil, nil)
 	if err != nil {
 		b.logger.Error("failed to generate private key", "error", err)
-		return fmt.Errorf("failed to generate private key: %w", err)
+		return "", fmt.Errorf("failed to generate private key: %w", err)
 	}
 	publicKey := privateKey.PublicKey()
 	signature := bls.Sign(privateKey, hashedMessage)
@@ -141,13 +144,13 @@ func (b *ProviderClient) CheckAndStake(_ []string) error {
 	// Verify the signature
 	if !bls.Verify(publicKey, hashedMessage, signature) {
 		b.logger.Error("failed to verify generated BLS signature")
-		return fmt.Errorf("failed to verify generated BLS signature")
+		return "", fmt.Errorf("failed to verify generated BLS signature")
 	}
 
 	pubkeyb, err := publicKey.MarshalBinary()
 	if err != nil {
 		b.logger.Error("failed to marshal public key", "error", err)
-		return fmt.Errorf("failed to marshal public key: %w", err)
+		return "", fmt.Errorf("failed to marshal public key: %w", err)
 	}
 	b.logger.Info("generated BLS key pair",
 		"public_key", hex.EncodeToString(pubkeyb),
@@ -161,13 +164,12 @@ func (b *ProviderClient) CheckAndStake(_ []string) error {
 	})
 	if err != nil {
 		b.logger.Error("failed to register stake", "err", err)
-		return err
+		return "", err
 	}
 
 	b.logger.Info("staked 10 ETH")
 
-	return nil
-
+	return hex.EncodeToString(pubkeyb), nil
 }
 
 func (b *ProviderClient) startSender() error {
diff --git a/p2p/integrationtest/provider/main.go b/p2p/integrationtest/provider/main.go
index c29cc1a13..8cbccaeac 100644
--- a/p2p/integrationtest/provider/main.go
+++ b/p2p/integrationtest/provider/main.go
@@ -3,8 +3,6 @@ package main
 import (
 	"bytes"
 	"context"
-	crand "crypto/rand"
-	"encoding/hex"
 	"encoding/json"
 	"errors"
 	"flag"
@@ -162,19 +160,13 @@ func main() {
 	}
 	defer providerClient.Close()
 
-	blsPubKey := make([]byte, 48)
-	if _, err = crand.Read(blsPubKey); err != nil {
-		logger.Error("failed to generate BLS public key", "error", err)
-		return
-	}
-
-	payload := hex.EncodeToString(blsPubKey)
-	if err = providerClient.CheckAndStake([]string{payload}); err != nil {
+	pubKey, err := providerClient.CheckAndStake()
+	if err != nil {
 		logger.Error("failed to check and stake", "error", err)
 		return
 	}
 
-	body, err := json.Marshal([]string{payload})
+	body, err := json.Marshal([]string{pubKey})
 	if err != nil {
 		logger.Error("failed to marshal body", "error", err)
 		return