[batcher] Replace Fireblocks wallet with KMS wallet (#550)

common/kms_wallet_config.go

@@ -0,0 +1,42 @@
+package common
+
+import (
+	"github.com/urfave/cli"
+)
+
+type KMSKeyConfig struct {
+	KeyID   string
+	Region  string
+	Disable bool
+}
+
+func KMSWalletCLIFlags(envPrefix string, flagPrefix string) []cli.Flag {
+	return []cli.Flag{
+		cli.StringFlag{
+			Name:     PrefixFlag(flagPrefix, "kms-key-id"),
+			Usage:    "KMS key ID that stores the private key",
+			Required: false,
+			EnvVar:   PrefixEnvVar(envPrefix, "KMS_KEY_ID"),
+		},
+		cli.StringFlag{
+			Name:     PrefixFlag(flagPrefix, "kms-key-region"),
+			Usage:    "KMS key region",
+			Required: false,
+			EnvVar:   PrefixEnvVar(envPrefix, "KMS_KEY_REGION"),
+		},
+		cli.BoolFlag{
+			Name:     PrefixFlag(flagPrefix, "kms-key-disable"),
+			Usage:    "Disable KMS wallet",
+			Required: false,
+			EnvVar:   PrefixEnvVar(envPrefix, "KMS_KEY_DISABLE"),
+		},
+	}
+}
+
+func ReadKMSKeyConfig(ctx *cli.Context, flagPrefix string) KMSKeyConfig {
+	return KMSKeyConfig{
+		KeyID:   ctx.String(PrefixFlag(flagPrefix, "kms-key-id")),
+		Region:  ctx.String(PrefixFlag(flagPrefix, "kms-key-region")),
+		Disable: ctx.Bool(PrefixFlag(flagPrefix, "kms-key-disable")),
+	}
+}

common/logger_config.go

@@ -26,7 +26,7 @@ const (
 type LoggerConfig struct {
 	Format       LogFormat
 	OutputWriter io.Writer
-	HandlerOpts  slog.HandlerOptions
+	HandlerOpts  logging.SLoggerOptions
 }
 
 func LoggerCLIFlags(envPrefix string, flagPrefix string) []cli.Flag {
@@ -56,7 +56,7 @@ func DefaultLoggerConfig() LoggerConfig {
 	return LoggerConfig{
 		Format:       JSONLogFormat,
 		OutputWriter: os.Stdout,
-		HandlerOpts: slog.HandlerOptions{
+		HandlerOpts: logging.SLoggerOptions{
 			AddSource: true,
 			Level:     slog.LevelDebug,
 		},
@@ -95,10 +95,10 @@ func ReadLoggerCLIConfig(ctx *cli.Context, flagPrefix string) (*LoggerConfig, er
 
 func NewLogger(cfg LoggerConfig) (logging.Logger, error) {
 	if cfg.Format == JSONLogFormat {
-		return logging.NewSlogJsonLogger(cfg.OutputWriter, &cfg.HandlerOpts), nil
+		return logging.NewJsonSLogger(cfg.OutputWriter, &cfg.HandlerOpts), nil
 	}
 	if cfg.Format == TextLogFormat {
-		return logging.NewSlogTextLogger(cfg.OutputWriter, &cfg.HandlerOpts), nil
+		return logging.NewTextSLogger(cfg.OutputWriter, &cfg.HandlerOpts), nil
 	}
 	return nil, fmt.Errorf("unknown log format: %s", cfg.Format)
 }

disperser/cmd/batcher/config.go

@@ -23,7 +23,7 @@ type Config struct {
 	LoggerConfig     common.LoggerConfig
 	MetricsConfig    batcher.MetricsConfig
 	IndexerConfig    indexer.Config
-	FireblocksConfig common.FireblocksConfig
+	KMSKeyConfig     common.KMSKeyConfig
 	ChainStateConfig thegraph.Config
 	UseGraph         bool
 
@@ -39,8 +39,8 @@ func NewConfig(ctx *cli.Context) (Config, error) {
 		return Config{}, err
 	}
 	ethClientConfig := geth.ReadEthClientConfig(ctx)
-	fireblocksConfig := common.ReadFireblocksCLIConfig(ctx, flags.FlagPrefix)
-	if !fireblocksConfig.Disable {
+	kmsConfig := common.ReadKMSKeyConfig(ctx, flags.FlagPrefix)
+	if !kmsConfig.Disable {
 		ethClientConfig = geth.ReadEthClientConfigRPCOnly(ctx)
 	}
 	config := Config{
@@ -84,7 +84,7 @@ func NewConfig(ctx *cli.Context) (Config, error) {
 		EigenDAServiceManagerAddr:     ctx.GlobalString(flags.EigenDAServiceManagerFlag.Name),
 		IndexerDataDir:                ctx.GlobalString(flags.IndexerDataDirFlag.Name),
 		IndexerConfig:                 indexer.ReadIndexerConfig(ctx),
-		FireblocksConfig:              fireblocksConfig,
+		KMSKeyConfig:                  kmsConfig,
 	}
 	return config, nil
 }

disperser/cmd/batcher/flags/flags.go

@@ -236,6 +236,6 @@ func init() {
 	Flags = append(Flags, common.LoggerCLIFlags(envVarPrefix, FlagPrefix)...)
 	Flags = append(Flags, indexer.CLIFlags(envVarPrefix)...)
 	Flags = append(Flags, aws.ClientFlags(envVarPrefix, FlagPrefix)...)
-	Flags = append(Flags, common.FireblocksCLIFlags(envVarPrefix, FlagPrefix)...)
 	Flags = append(Flags, thegraph.CLIFlags(envVarPrefix)...)
+	Flags = append(Flags, common.KMSWalletCLIFlags(envVarPrefix, FlagPrefix)...)
 }

disperser/cmd/batcher/main.go

@@ -22,6 +22,7 @@ import (
 	"github.com/Layr-Labs/eigenda/disperser/cmd/batcher/flags"
 	"github.com/Layr-Labs/eigenda/disperser/common/blobstore"
 	"github.com/Layr-Labs/eigenda/disperser/encoder"
+	"github.com/Layr-Labs/eigensdk-go/aws/kms"
 	walletsdk "github.com/Layr-Labs/eigensdk-go/chainio/clients/wallet"
 	"github.com/Layr-Labs/eigensdk-go/signerv2"
 	gethcommon "github.com/ethereum/go-ethereum/common"
@@ -102,11 +103,73 @@ func RunBatcher(ctx *cli.Context) error {
 	}, logger, metrics.DispatcherMetrics)
 	asgn := &core.StdAssignmentCoordinator{}
 
-	client, err := geth.NewMultiHomingClient(config.EthClientConfig, gethcommon.HexToAddress(config.FireblocksConfig.WalletAddress), logger)
-	if err != nil {
-		logger.Error("Cannot create chain.Client", "err", err)
-		return err
+	var wallet walletsdk.Wallet
+	var client *geth.MultiHomingClient
+	if !config.KMSKeyConfig.Disable {
+		if config.KMSKeyConfig.KeyID == "" || config.KMSKeyConfig.Region == "" {
+			return errors.New("KMS key ID and region must be specified unless KMS wallet is disabled")
+		}
+		kmsClient, err := kms.NewKMSClient(context.Background(), config.KMSKeyConfig.Region)
+		if err != nil {
+			return fmt.Errorf("failed to create KMS client: %w", err)
+		}
+		pubKey, err := kms.GetECDSAPublicKey(context.Background(), kmsClient, config.KMSKeyConfig.KeyID)
+		if err != nil {
+			return fmt.Errorf("failed to get public key from KMS: %w", err)
+		}
+		addr := crypto.PubkeyToAddress(*pubKey)
+		client, err = geth.NewMultiHomingClient(config.EthClientConfig, addr, logger)
+		if err != nil {
+			logger.Error("Cannot create chain.Client", "err", err)
+			return err
+		}
+		chainID, err := client.ChainID(context.Background())
+		if err != nil {
+			return fmt.Errorf("failed to get chain ID: %w", err)
+		}
+		signer := signerv2.NewKMSSigner(context.Background(), kmsClient, pubKey, config.KMSKeyConfig.KeyID, chainID)
+		if err != nil {
+			return err
+		}
+		wallet, err = walletsdk.NewPrivateKeyWallet(client, signer, addr, logger)
+		if err != nil {
+			return err
+		}
+		logger.Info("Initialized KMS wallet", "address", addr.Hex())
+	} else if len(config.EthClientConfig.PrivateKeyString) > 0 {
+		privateKey, err := crypto.HexToECDSA(config.EthClientConfig.PrivateKeyString)
+		if err != nil {
+			return fmt.Errorf("failed to parse private key: %w", err)
+		}
+		client, err = geth.NewMultiHomingClient(config.EthClientConfig, gethcommon.Address{}, logger)
+		if err != nil {
+			logger.Error("Cannot create chain.Client", "err", err)
+			return err
+		}
+		chainID, err := client.ChainID(context.Background())
+		if err != nil {
+			return fmt.Errorf("failed to get chain ID: %w", err)
+		}
+		signerV2, address, err := signerv2.SignerFromConfig(signerv2.Config{PrivateKey: privateKey}, chainID)
+		if err != nil {
+			return err
+		}
+		wallet, err = walletsdk.NewPrivateKeyWallet(client, signerV2, address, logger.With("component", "PrivateKeyWallet"))
+		if err != nil {
+			return err
+		}
+		logger.Info("Initialized PrivateKey wallet", "address", address.Hex())
+	} else {
+		return errors.New("no wallet is configured. Either Fireblocks or PrivateKey wallet should be configured")
+	}
+
+	if wallet == nil {
+		return errors.New("wallet is not configured")
+	}
+	if client == nil {
+		return errors.New("eth client is not configured")
 	}
+
 	// used by non graph indexer
 	rpcClient, err := rpc.Dial(config.EthClientConfig.RPCURLs[0])
 	if err != nil {
@@ -166,34 +229,6 @@ func RunBatcher(ctx *cli.Context) error {
 		return err
 	}
 	finalizer := batcher.NewFinalizer(config.TimeoutConfig.ChainReadTimeout, config.BatcherConfig.FinalizerInterval, queue, client, rpcClient, config.BatcherConfig.MaxNumRetriesPerBlob, 1000, config.BatcherConfig.FinalizerPoolSize, logger, metrics.FinalizerMetrics)
-	var wallet walletsdk.Wallet
-	if !config.FireblocksConfig.Disable {
-		wallet, err = common.NewFireblocksWallet(&config.FireblocksConfig, client, logger)
-		if err != nil {
-			return err
-		}
-	} else if len(config.EthClientConfig.PrivateKeyString) > 0 {
-		privateKey, err := crypto.HexToECDSA(config.EthClientConfig.PrivateKeyString)
-		if err != nil {
-			return fmt.Errorf("failed to parse private key: %w", err)
-		}
-		chainID, err := client.ChainID(context.Background())
-		if err != nil {
-			return fmt.Errorf("failed to get chain ID: %w", err)
-		}
-		signerV2, address, err := signerv2.SignerFromConfig(signerv2.Config{PrivateKey: privateKey}, chainID)
-		if err != nil {
-			return err
-		}
-		wallet, err = walletsdk.NewPrivateKeyWallet(client, signerV2, address, logger.With("component", "PrivateKeyWallet"))
-		if err != nil {
-			return err
-		}
-		logger.Info("Initialized PrivateKey wallet", "address", address.Hex())
-	} else {
-		return errors.New("no wallet is configured. Either Fireblocks or PrivateKey wallet should be configured")
-	}
-
 	txnManager := batcher.NewTxnManager(client, wallet, config.EthClientConfig.NumConfirmations, 20, config.TimeoutConfig.TxnBroadcastTimeout, config.TimeoutConfig.ChainWriteTimeout, logger, metrics.TxnManagerMetrics)
 	batcher, err := batcher.NewBatcher(config.BatcherConfig, config.TimeoutConfig, queue, dispatcher, ics, asgn, encoderClient, agg, client, finalizer, tx, txnManager, logger, metrics, handleBatchLivenessChan)
 	if err != nil {

go.mod

@@ -6,13 +6,13 @@ toolchain go1.21.1
 
 require (
 	github.com/Layr-Labs/eigenda/api v0.0.0
-	github.com/Layr-Labs/eigensdk-go v0.1.6-0.20240414172936-84d5bc10f72f
-	github.com/aws/aws-sdk-go-v2 v1.26.0
-	github.com/aws/aws-sdk-go-v2/credentials v1.17.9
+	github.com/Layr-Labs/eigensdk-go v0.1.7-0.20240507215523-7e4891d5099a
+	github.com/aws/aws-sdk-go-v2 v1.26.1
+	github.com/aws/aws-sdk-go-v2/credentials v1.17.11
 	github.com/aws/aws-sdk-go-v2/feature/dynamodb/attributevalue v1.13.12
-	github.com/aws/aws-sdk-go-v2/service/secretsmanager v1.28.5
+	github.com/aws/aws-sdk-go-v2/service/secretsmanager v1.28.6
 	github.com/consensys/gnark-crypto v0.12.1
-	github.com/ethereum/go-ethereum v1.13.14
+	github.com/ethereum/go-ethereum v1.14.0
 	github.com/fxamacker/cbor/v2 v2.5.0
 	github.com/gin-contrib/logger v0.2.6
 	github.com/gin-gonic/gin v1.9.1
@@ -47,39 +47,40 @@ require (
 	github.com/PuerkitoBio/purell v1.1.1 // indirect
 	github.com/PuerkitoBio/urlesc v0.0.0-20170810143723-de5bf2ad4578 // indirect
 	github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.6.1 // indirect
-	github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.16.0 // indirect
-	github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.4 // indirect
-	github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.4 // indirect
+	github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.16.1 // indirect
+	github.com/aws/aws-sdk-go-v2/internal/configsources v1.3.5 // indirect
+	github.com/aws/aws-sdk-go-v2/internal/endpoints/v2 v2.6.5 // indirect
 	github.com/aws/aws-sdk-go-v2/internal/ini v1.8.0 // indirect
 	github.com/aws/aws-sdk-go-v2/internal/v4a v1.3.4 // indirect
 	github.com/aws/aws-sdk-go-v2/service/dynamodbstreams v1.20.3 // indirect
-	github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.11.1 // indirect
+	github.com/aws/aws-sdk-go-v2/service/internal/accept-encoding v1.11.2 // indirect
 	github.com/aws/aws-sdk-go-v2/service/internal/checksum v1.3.6 // indirect
 	github.com/aws/aws-sdk-go-v2/service/internal/endpoint-discovery v1.9.5 // indirect
-	github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.11.6 // indirect
+	github.com/aws/aws-sdk-go-v2/service/internal/presigned-url v1.11.7 // indirect
 	github.com/aws/aws-sdk-go-v2/service/internal/s3shared v1.17.4 // indirect
-	github.com/aws/aws-sdk-go-v2/service/sso v1.20.3 // indirect
-	github.com/aws/aws-sdk-go-v2/service/ssooidc v1.23.3 // indirect
-	github.com/aws/aws-sdk-go-v2/service/sts v1.28.5 // indirect
-	github.com/aws/smithy-go v1.20.1 // indirect
+	github.com/aws/aws-sdk-go-v2/service/kms v1.31.0 // indirect
+	github.com/aws/aws-sdk-go-v2/service/sso v1.20.5 // indirect
+	github.com/aws/aws-sdk-go-v2/service/ssooidc v1.23.4 // indirect
+	github.com/aws/aws-sdk-go-v2/service/sts v1.28.6 // indirect
+	github.com/aws/smithy-go v1.20.2 // indirect
 	github.com/bits-and-blooms/bitset v1.10.0 // indirect
 	github.com/bytedance/sonic v1.9.2 // indirect
 	github.com/cenkalti/backoff/v4 v4.2.1 // indirect
 	github.com/chenzhuoyu/base64x v0.0.0-20221115062448-fe3a3abad311 // indirect
-	github.com/cockroachdb/errors v1.9.1 // indirect
+	github.com/cockroachdb/errors v1.11.1 // indirect
 	github.com/cockroachdb/logtags v0.0.0-20230118201751-21c54148d20b // indirect
-	github.com/cockroachdb/pebble v0.0.0-20230928194634-aa077af62593 // indirect
-	github.com/cockroachdb/redact v1.1.3 // indirect
+	github.com/cockroachdb/pebble v1.1.0 // indirect
+	github.com/cockroachdb/redact v1.1.5 // indirect
 	github.com/cockroachdb/tokenbucket v0.0.0-20230807174530-cc333fc44b06 // indirect
 	github.com/containerd/continuity v0.3.0 // indirect
 	github.com/crate-crypto/go-ipa v0.0.0-20231025140028-3c0104f4b233 // indirect
-	github.com/crate-crypto/go-kzg-4844 v0.7.0 // indirect
+	github.com/crate-crypto/go-kzg-4844 v1.0.0 // indirect
 	github.com/deckarep/golang-set/v2 v2.1.0 // indirect
 	github.com/docker/cli v25.0.3+incompatible // indirect
-	github.com/docker/docker v25.0.3+incompatible // indirect
+	github.com/docker/docker v25.0.5+incompatible // indirect
 	github.com/docker/go-connections v0.5.0 // indirect
 	github.com/docker/go-units v0.5.0 // indirect
-	github.com/ethereum/c-kzg-4844 v0.4.0 // indirect
+	github.com/ethereum/c-kzg-4844 v1.0.0 // indirect
 	github.com/fjl/memsize v0.0.2 // indirect
 	github.com/gabriel-vasile/mimetype v1.4.2 // indirect
 	github.com/gammazero/deque v0.2.0 // indirect
@@ -99,7 +100,7 @@ require (
 	github.com/gogo/protobuf v1.3.2 // indirect
 	github.com/golang-jwt/jwt v3.2.2+incompatible // indirect
 	github.com/golang-jwt/jwt/v4 v4.5.0 // indirect
-	github.com/golang/protobuf v1.5.3 // indirect
+	github.com/golang/protobuf v1.5.4 // indirect
 	github.com/google/shlex v0.0.0-20191202100458-e7afc7fbc510 // indirect
 	github.com/hashicorp/errwrap v1.0.0 // indirect
 	github.com/hashicorp/go-bexpr v0.1.10 // indirect
@@ -114,6 +115,7 @@ require (
 	github.com/klauspost/cpuid/v2 v2.2.5 // indirect
 	github.com/kr/pretty v0.3.1 // indirect
 	github.com/leodido/go-urn v1.2.4 // indirect
+	github.com/lmittmann/tint v1.0.4 // indirect
 	github.com/mailru/easyjson v0.7.7 // indirect
 	github.com/mattn/go-colorable v0.1.13 // indirect
 	github.com/mitchellh/mapstructure v1.4.1 // indirect
@@ -143,10 +145,10 @@ require (
 	go.uber.org/multierr v1.11.0 // indirect
 	go.uber.org/zap v1.27.0 // indirect
 	golang.org/x/arch v0.4.0 // indirect
-	golang.org/x/mod v0.16.0 // indirect
+	golang.org/x/mod v0.17.0 // indirect
 	golang.org/x/oauth2 v0.16.0 // indirect
-	golang.org/x/sync v0.5.0 // indirect
-	golang.org/x/time v0.3.0 // indirect
+	golang.org/x/sync v0.7.0 // indirect
+	golang.org/x/time v0.5.0 // indirect
 	google.golang.org/appengine v1.6.7 // indirect
 	google.golang.org/genproto/googleapis/rpc v0.0.0-20231016165738-49dd2c1f3d0b // indirect
 	gopkg.in/natefinch/lumberjack.v2 v2.0.0 // indirect
@@ -156,7 +158,7 @@ require (
 require (
 	github.com/StackExchange/wmi v1.2.1 // indirect
 	github.com/VictoriaMetrics/fastcache v1.12.1 // indirect
-	github.com/aws/aws-sdk-go-v2/config v1.27.9
+	github.com/aws/aws-sdk-go-v2/config v1.27.11
 	github.com/aws/aws-sdk-go-v2/feature/dynamodb/expression v1.7.12
 	github.com/aws/aws-sdk-go-v2/feature/s3/manager v1.16.13
 	github.com/aws/aws-sdk-go-v2/service/dynamodb v1.31.0
@@ -171,7 +173,7 @@ require (
 	github.com/fsnotify/fsnotify v1.6.0 // indirect
 	github.com/gammazero/workerpool v1.1.3
 	github.com/gin-contrib/cors v1.4.0
-	github.com/go-logr/logr v1.2.4 // indirect
+	github.com/go-logr/logr v1.4.1 // indirect
 	github.com/go-ole/go-ole v1.3.0 // indirect
 	github.com/go-task/slim-sprig v0.0.0-20230315185526-52ccab3ef572 // indirect
 	github.com/golang/snappy v0.0.5-0.20220116011046-fa5810519dcb // indirect
@@ -200,12 +202,12 @@ require (
 	github.com/swaggo/gin-swagger v1.6.0
 	github.com/tklauser/go-sysconf v0.3.12 // indirect
 	github.com/tklauser/numcpus v0.6.1 // indirect
-	golang.org/x/crypto v0.18.0
-	golang.org/x/net v0.20.0 // indirect
-	golang.org/x/sys v0.16.0 // indirect
+	golang.org/x/crypto v0.22.0
+	golang.org/x/net v0.24.0 // indirect
+	golang.org/x/sys v0.19.0 // indirect
 	golang.org/x/text v0.14.0 // indirect
-	golang.org/x/tools v0.15.0 // indirect
-	google.golang.org/protobuf v1.32.0
+	golang.org/x/tools v0.20.0 // indirect
+	google.golang.org/protobuf v1.33.0
 	gopkg.in/yaml.v3 v3.0.1
 	rsc.io/tmplfunc v0.0.3 // indirect
 )