Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 16 additions & 5 deletions api/clients/v2/dispersal_request_signer.go
Original file line number Diff line number Diff line change
Expand Up @@ -28,8 +28,12 @@ type DispersalRequestSignerConfig struct {
KeyID string `docs:"required"`
// PrivateKey is a hex-encoded private key for local signing (without 0x prefix). Optional if KeyID is provided.
PrivateKey string `docs:"required"`
// Region is the AWS region where the KMS key is located (e.g., "us-east-1"). Required if using KMS.
// Region is the default AWS region (e.g., "us-east-1"). Required if using KMS.
Region string `docs:"required"`
// KMSRegion is an optional AWS region override for KMS operations. When specified, this region is used
// for KMS operations instead of Region. If empty, Region is used. This allows KMS keys to be stored in
// a different region than other AWS resources.
KMSRegion string
// Endpoint is an optional custom AWS KMS endpoint URL. If empty, the standard AWS KMS endpoint is used.
// This is primarily useful for testing with LocalStack or other custom KMS implementations. Default is empty.
Endpoint string
Expand Down Expand Up @@ -96,27 +100,34 @@ func NewKMSDispersalRequestSigner(
ctx context.Context,
config DispersalRequestSignerConfig,
) (DispersalRequestSigner, error) {
// Determine which region to use for KMS operations.
// If KMSRegion is specified, use it; otherwise fall back to Region.
kmsRegion := config.Region
if config.KMSRegion != "" {
kmsRegion = config.KMSRegion
}

var kmsClient *kms.Client
if config.Endpoint != "" {
kmsClient = kms.New(kms.Options{
Region: config.Region,
Region: kmsRegion,
BaseEndpoint: aws.String(config.Endpoint),
})
} else {
// Load the AWS SDK configuration, which will automatically detect credentials
// from environment variables, IAM roles, or AWS config files
cfg, err := awsconfig.LoadDefaultConfig(ctx,
awsconfig.WithRegion(config.Region),
awsconfig.WithRegion(kmsRegion),
)
if err != nil {
return nil, fmt.Errorf("failed to load AWS config: %w", err)
return nil, fmt.Errorf("failed to load AWS config for region %s: %w", kmsRegion, err)
}
kmsClient = kms.NewFromConfig(cfg)
}

key, err := aws2.LoadPublicKeyKMS(ctx, kmsClient, config.KeyID)
if err != nil {
return nil, fmt.Errorf("failed to get ecdsa public key: %w", err)
return nil, fmt.Errorf("failed to get ecdsa public key from KMS in region %s: %w", kmsRegion, err)
}

return &kmsRequestSigner{
Expand Down
70 changes: 70 additions & 0 deletions api/clients/v2/dispersal_request_signer_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -791,3 +791,73 @@ func TestKMSSignerWithDefaultConfig(t *testing.T) {
// This will fail in test environment but we're testing the code path
require.Error(t, err, "should fail to load default AWS config in test environment")
}

func TestKMSRegionOverride(t *testing.T) {
ctx := t.Context()
_ = setupLocalStack(t)

keyManager := kms.New(kms.Options{
Region: region,
BaseEndpoint: aws.String(localstackHost),
})

keyID, _ := createTestKMSKey(t, ctx, keyManager)

// Test that KMSRegion overrides Region for KMS operations
signer, err := NewDispersalRequestSigner(ctx, DispersalRequestSignerConfig{
Region: "us-west-2", // This should be ignored for KMS
KMSRegion: region, // This should be used for KMS
Endpoint: localstackHost,
KeyID: keyID,
})
require.NoError(t, err, "should create KMS signer with region override")

kmsSigner, ok := signer.(*kmsRequestSigner)
require.True(t, ok, "should be kmsRequestSigner type")
require.NotNil(t, kmsSigner.kmsClient, "KMS client should be initialized")

// Test signing with KMS region override
rand := random.NewTestRandom()
request := auth.RandomStoreChunksRequest(rand)
request.Signature = nil

signature, err := signer.SignStoreChunksRequest(ctx, request)
require.NoError(t, err, "should sign successfully with KMS region override")
require.NotNil(t, signature, "signature should not be nil")
require.NotEmpty(t, signature, "signature should not be empty")
}

func TestKMSRegionDefault(t *testing.T) {
ctx := t.Context()
_ = setupLocalStack(t)

keyManager := kms.New(kms.Options{
Region: region,
BaseEndpoint: aws.String(localstackHost),
})

keyID, _ := createTestKMSKey(t, ctx, keyManager)

// Test that Region is used when KMSRegion is not specified
signer, err := NewDispersalRequestSigner(ctx, DispersalRequestSignerConfig{
Region: region,
Endpoint: localstackHost,
KeyID: keyID,
// KMSRegion not specified - should use Region
})
require.NoError(t, err, "should create KMS signer using default Region")

kmsSigner, ok := signer.(*kmsRequestSigner)
require.True(t, ok, "should be kmsRequestSigner type")
require.NotNil(t, kmsSigner.kmsClient, "KMS client should be initialized")

// Test signing when KMSRegion is not specified
rand := random.NewTestRandom()
request := auth.RandomStoreChunksRequest(rand)
request.Signature = nil

signature, err := signer.SignStoreChunksRequest(ctx, request)
require.NoError(t, err, "should sign successfully using default Region")
require.NotNil(t, signature, "signature should not be nil")
require.NotEmpty(t, signature, "signature should not be empty")
}
1 change: 1 addition & 0 deletions disperser/cmd/controller/config.go
Original file line number Diff line number Diff line change
Expand Up @@ -103,6 +103,7 @@ func NewConfig(ctx *cli.Context) (*controller.ControllerConfig, error) {
KeyID: ctx.GlobalString(flags.DisperserKMSKeyIDFlag.Name),
PrivateKey: ctx.GlobalString(flags.DisperserPrivateKeyFlag.Name),
Region: awsClientConfig.Region,
KMSRegion: ctx.GlobalString(flags.DisperserKMSRegionFlag.Name),
Endpoint: awsClientConfig.EndpointURL,
},
Encoder: controller.EncodingManagerConfig{
Expand Down
7 changes: 7 additions & 0 deletions disperser/cmd/controller/flags/flags.go
Original file line number Diff line number Diff line change
Expand Up @@ -240,6 +240,12 @@ var (
Required: false,
EnvVar: common.PrefixEnvVar(envVarPrefix, "DISPERSER_PRIVATE_KEY"),
}
DisperserKMSRegionFlag = cli.StringFlag{
Name: common.PrefixFlag(FlagPrefix, "disperser-kms-region"),
Usage: "AWS region for KMS key (overrides default AWS region when specified)",
Required: false,
EnvVar: common.PrefixEnvVar(envVarPrefix, "DISPERSER_KMS_REGION"),
}
ControllerReadinessProbePathFlag = cli.StringFlag{
Name: common.PrefixFlag(FlagPrefix, "controller-readiness-probe-path"),
Usage: "File path for the readiness probe; created once the controller is fully started and ready to serve traffic",
Expand Down Expand Up @@ -430,6 +436,7 @@ var optionalFlags = []cli.Flag{
DisperserStoreChunksSigningDisabledFlag,
DisperserKMSKeyIDFlag,
DisperserPrivateKeyFlag,
DisperserKMSRegionFlag,
ControllerReadinessProbePathFlag,
ControllerHealthProbePathFlag,
ControllerHeartbeatMaxStallDurationFlag,
Expand Down
Loading