From 44569ec461c9a1dd1191e7999a72e63bd1e7aba9 Mon Sep 17 00:00:00 2001 From: Robert Raynor <35671663+mooselumph@users.noreply.github.com> Date: Fri, 7 Jun 2024 18:51:34 -0700 Subject: [PATCH] Auth traffic gen (#566) Co-authored-by: Daniel Mancia <21249320+dmanc@users.noreply.github.com> --- tools/traffic/cmd/main.go | 11 ++++++++++- tools/traffic/config.go | 5 ++++- tools/traffic/flags/flags.go | 7 +++++++ tools/traffic/generator.go | 27 ++++++++++++++++++++------- tools/traffic/generator_test.go | 29 +++++++++++++++++++++++++++++ 5 files changed, 70 insertions(+), 9 deletions(-) diff --git a/tools/traffic/cmd/main.go b/tools/traffic/cmd/main.go index 64b8b053b..62a508629 100644 --- 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 index 2e04bf65a..590c60759 100644 --- 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 index 91ffff728..aca2eb11a 100644 --- 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 index 868d736bf..2f8731ffd 100644 --- 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 index 67b4d7422..b530ef5bd 100644 --- 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) +}