Auth traffic gen (#566)

Co-authored-by: Daniel Mancia <[email protected]>
Layr-Labs · Jun 8, 2024 · 44569ec · 44569ec
1 parent e02b865
commit 44569ec
Show file tree

Hide file tree

Showing 5 changed files with 70 additions and 9 deletions.
diff --git a/tools/traffic/cmd/main.go b/tools/traffic/cmd/main.go
@@ -5,6 +5,8 @@ import (
 	"log"
 	"os"
 
+	"github.com/Layr-Labs/eigenda/core"
+	"github.com/Layr-Labs/eigenda/core/auth"
 	"github.com/Layr-Labs/eigenda/tools/traffic"
 	"github.com/Layr-Labs/eigenda/tools/traffic/flags"
 	"github.com/urfave/cli"
@@ -34,7 +36,14 @@ func trafficGeneratorMain(ctx *cli.Context) error {
 	if err != nil {
 		return err
 	}
-	generator, err := traffic.NewTrafficGenerator(config)
+
+	var signer core.BlobRequestSigner
+	if config.SignerPrivateKey != "" {
+		log.Println("Using signer private key")
+		signer = auth.NewLocalBlobRequestSigner(config.SignerPrivateKey)
+	}
+
+	generator, err := traffic.NewTrafficGenerator(config, signer)
 	if err != nil {
 		panic("failed to create new traffic generator")
 	}

diff --git a/tools/traffic/config.go b/tools/traffic/config.go
@@ -19,7 +19,9 @@ type Config struct {
 	LoggingConfig          common.LoggerConfig
 	RandomizeBlobs         bool
 	InstanceLaunchInterval time.Duration
-	CustomQuorums          []uint8
+
+	SignerPrivateKey string
+	CustomQuorums    []uint8
 }
 
 func NewConfig(ctx *cli.Context) (*Config, error) {
@@ -48,6 +50,7 @@ func NewConfig(ctx *cli.Context) (*Config, error) {
 		LoggingConfig:          *loggerConfig,
 		RandomizeBlobs:         ctx.GlobalBool(flags.RandomizeBlobsFlag.Name),
 		InstanceLaunchInterval: ctx.Duration(flags.InstanceLaunchIntervalFlag.Name),
+		SignerPrivateKey:       ctx.String(flags.SignerPrivateKeyFlag.Name),
 		CustomQuorums:          customQuorumsUint8,
 	}, nil
 }
diff --git a/tools/traffic/flags/flags.go b/tools/traffic/flags/flags.go
@@ -72,6 +72,12 @@ var (
 		Required: false,
 		EnvVar:   common.PrefixEnvVar(envPrefix, "USE_SECURE_GRPC"),
 	}
+	SignerPrivateKeyFlag = cli.StringFlag{
+		Name:     common.PrefixFlag(FlagPrefix, "signer-private-key-hex"),
+		Usage:    "Private key to use for signing requests",
+		Required: false,
+		EnvVar:   common.PrefixEnvVar(envPrefix, "SIGNER_PRIVATE_KEY_HEX"),
+	}
 	CustomQuorumNumbersFlag = cli.IntSliceFlag{
 		Name:     common.PrefixFlag(FlagPrefix, "custom-quorum-numbers"),
 		Usage:    "Custom quorum numbers to use for the traffic generator",
@@ -93,6 +99,7 @@ var optionalFlags = []cli.Flag{
 	RandomizeBlobsFlag,
 	InstanceLaunchIntervalFlag,
 	UseSecureGrpcFlag,
+	SignerPrivateKeyFlag,
 	CustomQuorumNumbersFlag,
 }
 

diff --git a/tools/traffic/generator.go b/tools/traffic/generator.go
@@ -12,6 +12,7 @@ import (
 
 	"github.com/Layr-Labs/eigenda/api/clients"
 	"github.com/Layr-Labs/eigenda/common"
+	"github.com/Layr-Labs/eigenda/core"
 	"github.com/Layr-Labs/eigenda/encoding/utils/codec"
 	"github.com/Layr-Labs/eigensdk-go/logging"
 )
@@ -22,7 +23,7 @@ type TrafficGenerator struct {
 	Config          *Config
 }
 
-func NewTrafficGenerator(config *Config) (*TrafficGenerator, error) {
+func NewTrafficGenerator(config *Config, signer core.BlobRequestSigner) (*TrafficGenerator, error) {
 	loggerConfig := common.DefaultLoggerConfig()
 	logger, err := common.NewLogger(loggerConfig)
 	if err != nil {
@@ -31,7 +32,7 @@ func NewTrafficGenerator(config *Config) (*TrafficGenerator, error) {
 
 	return &TrafficGenerator{
 		Logger:          logger,
-		DisperserClient: clients.NewDisperserClient(&config.Config, nil),
+		DisperserClient: clients.NewDisperserClient(&config.Config, signer),
 		Config:          config,
 	}, nil
 }
@@ -97,11 +98,23 @@ func (g *TrafficGenerator) StartTraffic(ctx context.Context) error {
 func (g *TrafficGenerator) sendRequest(ctx context.Context, data []byte) error {
 	ctxTimeout, cancel := context.WithTimeout(ctx, g.Config.Timeout)
 	defer cancel()
-	blobStatus, key, err := g.DisperserClient.DisperseBlob(ctxTimeout, data, g.Config.CustomQuorums)
-	if err != nil {
-		return err
+
+	if g.Config.SignerPrivateKey != "" {
+		blobStatus, key, err := g.DisperserClient.DisperseBlobAuthenticated(ctxTimeout, data, g.Config.CustomQuorums)
+		if err != nil {
+			return err
+		}
+
+		g.Logger.Info("successfully dispersed new blob", "authenticated", true, "key", hex.EncodeToString(key), "status", blobStatus.String())
+		return nil
+	} else {
+		blobStatus, key, err := g.DisperserClient.DisperseBlob(ctxTimeout, data, g.Config.CustomQuorums)
+		if err != nil {
+			return err
+		}
+
+		g.Logger.Info("successfully dispersed new blob", "authenticated", false, "key", hex.EncodeToString(key), "status", blobStatus.String())
+		return nil
 	}
 
-	g.Logger.Info("successfully dispersed new blob,", "key", hex.EncodeToString(key), "status", blobStatus.String())
-	return nil
 }
diff --git a/tools/traffic/generator_test.go b/tools/traffic/generator_test.go
@@ -40,3 +40,32 @@ func TestTrafficGenerator(t *testing.T) {
 	cancel()
 	disperserClient.AssertNumberOfCalls(t, "DisperseBlob", 2)
 }
+
+func TestTrafficGeneratorAuthenticated(t *testing.T) {
+	disperserClient := clientsmock.NewMockDisperserClient()
+	logger := logging.NewNoopLogger()
+
+	trafficGenerator := &traffic.TrafficGenerator{
+		Logger: logger,
+		Config: &traffic.Config{
+			Config: clients.Config{
+				Timeout: 1 * time.Second,
+			},
+			DataSize:         1000_000,
+			RequestInterval:  2 * time.Second,
+			SignerPrivateKey: "Hi",
+		},
+		DisperserClient: disperserClient,
+	}
+
+	processing := disperser.Processing
+	disperserClient.On("DisperseBlobAuthenticated", mock.Anything, mock.Anything, mock.Anything, mock.Anything).
+		Return(&processing, []byte{1}, nil)
+	ctx, cancel := context.WithCancel(context.Background())
+	go func() {
+		_ = trafficGenerator.StartTraffic(ctx)
+	}()
+	time.Sleep(5 * time.Second)
+	cancel()
+	disperserClient.AssertNumberOfCalls(t, "DisperseBlobAuthenticated", 2)
+}