From 6507b6331fc0bdc6990808054ac65fd4e42051b3 Mon Sep 17 00:00:00 2001 From: Ali Date: Sat, 16 May 2026 12:23:25 +0400 Subject: [PATCH] fix: mitigate denial of wallet attacks with rate-limits --- go.mod | 2 +- providers/awskms/README.md | 42 ++++++- providers/awskms/kms.go | 145 ++++++++++++++++++--- providers/awskms/kms_test.go | 238 +++++++++++++++++++++++++++++++++++ 4 files changed, 409 insertions(+), 18 deletions(-) diff --git a/go.mod b/go.mod index b27f19f..e9bf26b 100644 --- a/go.mod +++ b/go.mod @@ -11,6 +11,7 @@ require ( github.com/hashicorp/vault/api v1.23.0 github.com/pkg/errors v0.9.1 github.com/stretchr/testify v1.11.1 + golang.org/x/time v0.15.0 gorm.io/gorm v1.31.1 ) @@ -51,7 +52,6 @@ require ( golang.org/x/net v0.54.0 // indirect golang.org/x/sys v0.44.0 // indirect golang.org/x/text v0.37.0 // indirect - golang.org/x/time v0.15.0 // indirect gopkg.in/check.v1 v1.0.0-20201130134442-10cb98267c6c // indirect gopkg.in/yaml.v3 v3.0.1 // indirect ) diff --git a/providers/awskms/README.md b/providers/awskms/README.md index 333947d..677b1e1 100644 --- a/providers/awskms/README.md +++ b/providers/awskms/README.md @@ -13,9 +13,11 @@ package main import ( "context" + "time" "github.com/aws/aws-sdk-go-v2/config" "github.com/aws/aws-sdk-go-v2/service/kms" + "github.com/dgraph-io/ristretto" "github.com/bincyber/go-sqlcrypter" "github.com/bincyber/go-sqlcrypter/providers/awskms" @@ -29,7 +31,19 @@ func main() { client := kms.NewFromConfig(cfg) - kmsCrypter, err := awskms.New(context.Background(), client, "alias/sqlcrypter") + kmsCrypter, err := awskms.New( + context.Background(), + client, + "alias/sqlcrypter", + // Optioanlly configure request timeout, rate limits, and cache + awskms.WithRequestTimeout(2*time.Second), + awskms.WithKMSDecryptRateLimit(5, 10) + awskms.WithDEKCacheConfig(ristretto.Config{ + MaxCost: int64(100_000), + NumCounters: int64(500_000), + BufferItems: int64(64), + }), + ) if err != nil { //handle error } @@ -42,6 +56,32 @@ func main() { `KMSCrypter` uses envelope encryption. When `awskms.New()` is called, a request is made to the the KMS [GenerateDataKey](https://docs.aws.amazon.com/kms/latest/APIReference/API_GenerateDataKey.html) API to retrieve a 256-bit symmetric data encryption key (DEK). This DEK is used to encrypt data using AES GCM instead of calling the KMS [Encrypt](https://docs.aws.amazon.com/kms/latest/APIReference/API_Encrypt.html) and [Decrypt](https://docs.aws.amazon.com/kms/latest/APIReference/API_Decrypt.html) APIs every time. The encrypted DEK is stored alongside the ciphertext. To decrypt previous DEKs stored alongside ciphertext, a request is made to the KMS [Decrypt](https://docs.aws.amazon.com/kms/latest/APIReference/API_Decrypt.html) API. The decrypted DEK is then cached in memory to avoid repetitive API calls to KMS. +### Request Deadlines + +Each call to KMS (`GenerateDataKey` during `New`, and `Decrypt` when decrypting an unknown encrypted DEK) runs under a **deadline**. The default is **2 seconds**. This can be overridden with `WithRequestTimeout()` option. + +### Cache + +Previous DEKs decrypted via KMS are cached in a [Ristretto](https://github.com/dgraph-io/ristretto) in-memory cache. + +Defaults: +- `MaxCost` 1,000,000 (1MB) +- `NumCounters` 50,000 +- `BufferItems` 64. + +This can be overridden with `WithDEKCacheConfig(ristretto.Config{...})`. + +### Mitigating Denial of Wallet + +If an attacker can invoke `Decrypt` with arbitrary ciphertext, they can force **cache misses** by supplying **unique** encrypted DEK blobs. Each miss triggers a billable **KMS Decrypt** until the in-memory cache fills or entries expire. + +Mitigations in this provider: + +- **In-memory cache** of decrypted DEKs (TTL 60 minutes) so repeated blobs do not re-hit KMS. +- **Optional rate limit** on the KMS `Decrypt` path only: `WithKMSDecryptRateLimit(rps, burst)` using a token bucket. When the limit is exceeded, `Decrypt` returns `awskms.ErrKMSDecryptRateLimited` **before** calling AWS (map to HTTP 429 or similar at the application layer). + +You should still enforce **authentication**, **least-privilege IAM** on the CMK, **AWS Budgets** / billing alarms, and **CloudWatch** throttling alarms for KMS. Edge rate limiting (API gateway, WAF) is recommended for internet-facing services. + ### Testing [nsmith/local-kms](https://github.com/nsmithuk/local-kms) is used to help with testing. The seed file used is located in [testing/seed.yaml](https://github.com/bincyber/go-sqlcrypter/blob/master/providers/awskms/docs.md). diff --git a/providers/awskms/kms.go b/providers/awskms/kms.go index d28193e..12efc10 100644 --- a/providers/awskms/kms.go +++ b/providers/awskms/kms.go @@ -15,14 +15,89 @@ import ( "github.com/aws/aws-sdk-go-v2/service/kms/types" "github.com/dgraph-io/ristretto" "github.com/pkg/errors" + "golang.org/x/time/rate" "github.com/bincyber/go-sqlcrypter" ) +// defaultTimeout is the default timeout used for calls to AWS KMS. +const defaultTimeout = 2 * time.Second + +// defaults for the in-memory DEK cache (see ristretto.Config). +const ( + defaultDEKCacheMaxCost = int64(1 * 1024 * 1024) // 1 MB + defaultDEKCacheNumCounters = int64(50_000) + defaultDEKCacheBufferItems = int64(64) // 64 is the Ristretto-recommended default +) + +// ErrKMSDecryptRateLimited is returned when optional KMS Decrypt rate limiting +// rejects a cache-miss decrypt before calling AWS KMS. +var ErrKMSDecryptRateLimited = errors.New("rate limit exceeded for KMS Decrypt") + +// kmsAPI is the subset of AWS KMS client calls used by KMSCrypter (for testing with fakes). +type kmsAPI interface { + GenerateDataKey(context.Context, *kms.GenerateDataKeyInput, ...func(*kms.Options)) (*kms.GenerateDataKeyOutput, error) + Decrypt(context.Context, *kms.DecryptInput, ...func(*kms.Options)) (*kms.DecryptOutput, error) +} + +// Option configures a KMSCrypter during New. +type Option func(*KMSCrypter) error + +// WithRequestTimeout sets the per-KMS-call deadline for GenerateDataKey and Decrypt +// (encrypted DEK path). Values must be greater than zero. The default is 2 second. +func WithRequestTimeout(timeout time.Duration) Option { + return func(k *KMSCrypter) error { + if timeout <= 0 { + return errors.New("request timeout must be greater than zero") + } + k.requestTimeout = timeout + return nil + } +} + +// WithKMSDecryptRateLimit sets a token-bucket limit on KMS Decrypt calls (cache-miss path only). +// rps and burst must be greater than zero. +func WithKMSDecryptRateLimit(rps float64, burst int) Option { + return func(k *KMSCrypter) error { + if rps <= 0 { + return errors.New("KMS decrypt rate limit rps must be greater than zero") + } + if burst <= 0 { + return errors.New("KMS decrypt rate limit burst must be greater than zero") + } + k.decryptLimiter = rate.NewLimiter(rate.Limit(rps), burst) + return nil + } +} + +// WithDEKCacheConfig sets the Ristretto configuration for the decrypted-DEK cache. +// NumCounters, MaxCost, and BufferItems must be greater than zero, and NumCounters +// must be at least MaxCost (see github.com/dgraph-io/ristretto documentation). +func WithDEKCacheConfig(cfg ristretto.Config) Option { + return func(k *KMSCrypter) error { + if cfg.NumCounters <= 0 { + return errors.New("DEK cache NumCounters must be greater than zero") + } + if cfg.MaxCost <= 0 { + return errors.New("DEK cache MaxCost must be greater than zero") + } + if cfg.BufferItems <= 0 { + return errors.New("DEK cache BufferItems must be greater than zero") + } + if cfg.NumCounters < cfg.MaxCost { + return errors.New("DEK cache NumCounters must be >= MaxCost") + } + k.cacheCounters = cfg.NumCounters + k.cacheMaxCost = cfg.MaxCost + k.cacheBufferItems = cfg.BufferItems + return nil + } +} + // KMSCrypter is an implementation of the Crypterer interface // using AWS KMS with envelope encryption. type KMSCrypter struct { - client *kms.Client + client kmsAPI // keyID is the ID, ARN, or Alias for the KMS key. keyID string @@ -42,12 +117,26 @@ type KMSCrypter struct { // cache stores any previous DEKs that were stored alongside ciphertext // to avoid repetitive client.Decrypt() calls to AWS KMS. cache *ristretto.Cache + + // cacheCounters, cacheMaxCost, and cacheBufferItems map to ristretto.Config (NumCounters, MaxCost, BufferItems). + cacheCounters int64 + cacheMaxCost int64 + cacheBufferItems int64 + + // requestTimeout caps each KMS API call (GenerateDataKey, Decrypt on miss). + requestTimeout time.Duration + + // decryptLimiter, when non-nil, rate-limits KMS Decrypt on cache miss. + decryptLimiter *rate.Limiter } // New creates a new AWS KMS crypter given a KMS client and the ID/Alias/ARN of a KMS key. // A new data encryption key (DEK) is obtained from KMS which will be stored alongside the // ciphertext. 256-bit AES GCM is used to perform the encryption. -func New(ctx context.Context, client *kms.Client, keyID string) (sqlcrypter.Crypterer, error) { +// +// By default each KMS request uses a 2 second deadline. This can be overridden using WithRequestTimeout option. +// The decrypted-DEK Ristretto cache uses built-in defaults; override with WithDEKCacheConfig. +func New(ctx context.Context, client *kms.Client, keyID string, opts ...Option) (sqlcrypter.Crypterer, error) { if client == nil { return nil, errors.New("kms.Client cannot be nil") } @@ -56,13 +145,31 @@ func New(ctx context.Context, client *kms.Client, keyID string) (sqlcrypter.Cryp return nil, errors.New("keyID cannot be empty") } + k := &KMSCrypter{ + client: client, + keyID: keyID, + requestTimeout: defaultTimeout, + cacheCounters: defaultDEKCacheNumCounters, + cacheMaxCost: defaultDEKCacheMaxCost, + cacheBufferItems: defaultDEKCacheBufferItems, + } + + for _, opt := range opts { + if err := opt(k); err != nil { + return nil, errors.Wrap(err, "failed to apply KMS crypter option") + } + } + // Generate a symmetric data encryption key to encrypt new data p := &kms.GenerateDataKeyInput{ KeyId: aws.String(keyID), KeySpec: types.DataKeySpecAes256, } - resp, err := client.GenerateDataKey(ctx, p) + rCtx, cancel := context.WithTimeout(ctx, k.requestTimeout) + defer cancel() + + resp, err := k.client.GenerateDataKey(rCtx, p) if err != nil { return nil, errors.Wrap(err, "failed to retrieve data key from AWS KMS") } @@ -79,9 +186,9 @@ func New(ctx context.Context, client *kms.Client, keyID string) (sqlcrypter.Cryp // Create in-memory cache for previous DEKs cache, err := ristretto.NewCache(&ristretto.Config{ - NumCounters: 100000000, - MaxCost: 10000000, // 10MB - BufferItems: 64, + NumCounters: k.cacheCounters, + MaxCost: k.cacheMaxCost, + BufferItems: k.cacheBufferItems, }) if err != nil { return nil, errors.Wrap(err, "failed to configure in-memory cache") @@ -92,14 +199,10 @@ func New(ctx context.Context, client *kms.Client, keyID string) (sqlcrypter.Cryp return nil, errors.New("encrypted DEK length exceeds uint8 limit") } - k := &KMSCrypter{ - client: client, - keyID: keyID, - aesgcm: aesgcm, - encryptedKey: resp.CiphertextBlob, - encryptedKeyLength: uint8(dekLen), - cache: cache, - } + k.aesgcm = aesgcm + k.encryptedKey = resp.CiphertextBlob + k.encryptedKeyLength = uint8(dekLen) + k.cache = cache return k, nil } @@ -217,6 +320,10 @@ func (k *KMSCrypter) Decrypt(w io.Writer, r io.Reader) error { return nil } + if k.decryptLimiter != nil && !k.decryptLimiter.Allow() { + return ErrKMSDecryptRateLimited + } + // Since the previous DEK doesn't exist in the cache, the DEK needs to be decrypted // using KMS. Then the decrypted key can be used to decrypt the ciphertext. p := &kms.DecryptInput{ @@ -224,7 +331,10 @@ func (k *KMSCrypter) Decrypt(w io.Writer, r io.Reader) error { CiphertextBlob: encryptedKey, } - resp, err := k.client.Decrypt(context.TODO(), p) + rCtx, cancel := context.WithTimeout(context.Background(), k.requestTimeout) + defer cancel() + + resp, err := k.client.Decrypt(rCtx, p) if err != nil { return errors.Wrap(err, "failed to decrypt previous DEK using KMS") } @@ -253,4 +363,7 @@ func (k *KMSCrypter) Decrypt(w io.Writer, r io.Reader) error { return nil } -var _ sqlcrypter.Crypterer = (*KMSCrypter)(nil) +var ( + _ sqlcrypter.Crypterer = (*KMSCrypter)(nil) + _ kmsAPI = (*kms.Client)(nil) +) diff --git a/providers/awskms/kms_test.go b/providers/awskms/kms_test.go index 237e3bd..635fbe9 100644 --- a/providers/awskms/kms_test.go +++ b/providers/awskms/kms_test.go @@ -5,9 +5,13 @@ import ( "context" "crypto/aes" "crypto/cipher" + "crypto/sha256" "encoding/binary" "io" + "net/http" + "net/http/httptest" "testing" + "time" "github.com/aws/aws-sdk-go-v2/aws" "github.com/aws/aws-sdk-go-v2/config" @@ -17,6 +21,7 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "github.com/stretchr/testify/suite" + "golang.org/x/time/rate" "github.com/bincyber/go-sqlcrypter" ) @@ -41,6 +46,17 @@ func getLocalKMSClient() *kms.Client { }) } +func getTestKMSClient(baseURL string) *kms.Client { + cfg, _ := config.LoadDefaultConfig(context.TODO(), + config.WithRegion("us-west-2"), + config.WithCredentialsProvider(credentials.NewStaticCredentialsProvider("AKID", "SECRET_KEY", "TOKEN")), + ) + + return kms.NewFromConfig(cfg, func(o *kms.Options) { + o.BaseEndpoint = aws.String(baseURL) + }) +} + type KMSCrypterTestSuite struct { suite.Suite client *kms.Client @@ -228,3 +244,225 @@ func Test_KMSCrypter_Decrypt_SingleByteDoesNotPanic(t *testing.T) { require.Error(t, err) assert.Contains(t, err.Error(), "failed to read ciphertext: minimum length not met") } + +func Test_New_invalid_request_timeout(t *testing.T) { + _, err := New(context.Background(), getLocalKMSClient(), KmsKeyAlias, WithRequestTimeout(0)) + require.Error(t, err) + assert.Contains(t, err.Error(), "request timeout must be greater than zero") +} + +func Test_New_invalid_KMS_decrypt_rate_limit(t *testing.T) { + _, err := New(context.Background(), getLocalKMSClient(), KmsKeyAlias, WithKMSDecryptRateLimit(0, 1)) + require.Error(t, err) + _, err = New(context.Background(), getLocalKMSClient(), KmsKeyAlias, WithKMSDecryptRateLimit(1, 0)) + require.Error(t, err) +} + +func Test_New_invalid_DEK_cache_config(t *testing.T) { + _, err := New(context.Background(), getLocalKMSClient(), KmsKeyAlias, WithDEKCacheConfig(ristretto.Config{ + NumCounters: 100, + MaxCost: 1000, + BufferItems: 16, + })) + require.Error(t, err) + assert.Contains(t, err.Error(), "NumCounters must be >= MaxCost") + + _, err = New(context.Background(), getLocalKMSClient(), KmsKeyAlias, WithDEKCacheConfig(ristretto.Config{ + NumCounters: 10_000, + MaxCost: 0, + BufferItems: 16, + })) + require.Error(t, err) + assert.Contains(t, err.Error(), "MaxCost must be greater than zero") + + _, err = New(context.Background(), getLocalKMSClient(), KmsKeyAlias, WithDEKCacheConfig(ristretto.Config{ + NumCounters: 10_000, + MaxCost: 1000, + BufferItems: 0, + })) + require.Error(t, err) + assert.Contains(t, err.Error(), "BufferItems must be greater than zero") +} + +func Test_New_request_timeout_applies_to_GenerateDataKey(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + time.Sleep(100 * time.Millisecond) + w.WriteHeader(http.StatusOK) + })) + t.Cleanup(server.Close) + + client := getTestKMSClient(server.URL) + _, err := New(context.Background(), client, KmsKeyAlias, WithRequestTimeout(50*time.Millisecond)) + require.Error(t, err) + assert.Contains(t, err.Error(), "failed to retrieve data key from AWS KMS") + assert.ErrorIs(t, err, context.DeadlineExceeded) +} + +// dekFromEncryptedBlob returns a deterministic 32-byte AES-256 key derived from the +// KMS ciphertext blob (simulates a unique plaintext DEK per encrypted DEK). Used by +// fakeDEKStub.Decrypt and by tests that seal payloads for a given blob. +func dekFromEncryptedBlob(blob []byte) []byte { + sum := sha256.Sum256(blob) + return append([]byte(nil), sum[:]...) +} + +// fakeDEKStub implements kmsAPI for unit tests. GenerateDataKey returns a fixed current +// DEK; Decrypt returns a deterministic DEK derived from CiphertextBlob so cache key/value +// mistakes fail AES-GCM Open instead of masking bugs. +type fakeDEKStub struct { + decryptCalls int + dek []byte + currentEnc []byte +} + +func newFakeDEKStub() *fakeDEKStub { + dek := make([]byte, 32) + dek[0] = 0x11 + cur := make([]byte, 140) + for i := range cur { + cur[i] = byte(i + 1) + } + return &fakeDEKStub{dek: dek, currentEnc: cur} +} + +func (f *fakeDEKStub) GenerateDataKey(ctx context.Context, in *kms.GenerateDataKeyInput, opt ...func(*kms.Options)) (*kms.GenerateDataKeyOutput, error) { + return &kms.GenerateDataKeyOutput{ + Plaintext: append([]byte(nil), f.dek...), + CiphertextBlob: append([]byte(nil), f.currentEnc...), + }, nil +} + +func (f *fakeDEKStub) Decrypt(ctx context.Context, in *kms.DecryptInput, opt ...func(*kms.Options)) (*kms.DecryptOutput, error) { + f.decryptCalls++ + dek := dekFromEncryptedBlob(in.CiphertextBlob) + return &kms.DecryptOutput{Plaintext: dek}, nil +} + +// stallKMS implements kmsAPI; Decrypt blocks until the request context ends (for timeout tests). +type stallKMS struct { + dek []byte + currentEnc []byte +} + +func newStallKMS() *stallKMS { + dek := make([]byte, 32) + dek[0] = 0x22 + cur := make([]byte, 140) + for i := range cur { + cur[i] = byte(200 + i) + } + return &stallKMS{dek: dek, currentEnc: cur} +} + +func (s *stallKMS) GenerateDataKey(ctx context.Context, in *kms.GenerateDataKeyInput, opt ...func(*kms.Options)) (*kms.GenerateDataKeyOutput, error) { + return &kms.GenerateDataKeyOutput{ + Plaintext: append([]byte(nil), s.dek...), + CiphertextBlob: append([]byte(nil), s.currentEnc...), + }, nil +} + +func (s *stallKMS) Decrypt(ctx context.Context, in *kms.DecryptInput, opt ...func(*kms.Options)) (*kms.DecryptOutput, error) { + <-ctx.Done() + return nil, ctx.Err() +} + +func testRistretto(t *testing.T) *ristretto.Cache { + t.Helper() + c, err := ristretto.NewCache(&ristretto.Config{ + NumCounters: 10, + MaxCost: 1 << 20, + BufferItems: 64, + }) + require.NoError(t, err) + return c +} + +func aesGCMDEKFromBlob(t *testing.T, encryptedDEKBLOB []byte) cipher.AEAD { + t.Helper() + b, err := aes.NewCipher(dekFromEncryptedBlob(encryptedDEKBLOB)) + require.NoError(t, err) + gcm, err := cipher.NewGCM(b) + require.NoError(t, err) + return gcm +} + +func wireKMSCiphertext(t *testing.T, encDEK []byte, nonce []byte, plain []byte, aead cipher.AEAD) []byte { + t.Helper() + var buf bytes.Buffer + require.NoError(t, binary.Write(&buf, binary.LittleEndian, uint8(len(encDEK)))) //nolint:gosec // test blobs are length 140 + buf.Write(encDEK) + buf.Write(nonce) + buf.Write(aead.Seal(nil, nonce, plain, nil)) + return buf.Bytes() +} + +func Test_Decrypt_KMS_rate_limit_second_miss(t *testing.T) { + stub := newFakeDEKStub() + block, err := aes.NewCipher(stub.dek) + require.NoError(t, err) + currentAEAD, err := cipher.NewGCM(block) + require.NoError(t, err) + + k := &KMSCrypter{ + client: stub, + keyID: KmsKeyAlias, + encryptedKey: append([]byte(nil), stub.currentEnc...), + encryptedKeyLength: 140, + aesgcm: currentAEAD, + cache: testRistretto(t), + requestTimeout: time.Second, + decryptLimiter: rate.NewLimiter(rate.Limit(1), 1), + } + + nonce1 := make([]byte, 12) + nonce1[11] = 1 + nonce2 := make([]byte, 12) + nonce2[11] = 2 + plain := []byte("hello") + alt1 := make([]byte, 140) + alt1[0] = 0xEE + alt2 := make([]byte, 140) + alt2[0] = 0xDD + + gcm1 := aesGCMDEKFromBlob(t, alt1) + gcm2 := aesGCMDEKFromBlob(t, alt2) + + b1 := wireKMSCiphertext(t, alt1, nonce1, plain, gcm1) + out := new(bytes.Buffer) + require.NoError(t, k.Decrypt(out, bytes.NewReader(b1))) + assert.Equal(t, plain, out.Bytes()) + assert.Equal(t, 1, stub.decryptCalls) + + b2 := wireKMSCiphertext(t, alt2, nonce2, plain, gcm2) + err = k.Decrypt(new(bytes.Buffer), bytes.NewReader(b2)) + require.Error(t, err) + require.ErrorIs(t, err, ErrKMSDecryptRateLimited) + assert.Equal(t, 1, stub.decryptCalls, "second decrypt must not call KMS when rate limited") +} + +func Test_Decrypt_KMS_deadline_exceeded(t *testing.T) { + stub := newStallKMS() + block, err := aes.NewCipher(stub.dek) + require.NoError(t, err) + aeadgcm, err := cipher.NewGCM(block) + require.NoError(t, err) + + k := &KMSCrypter{ + client: stub, + keyID: KmsKeyAlias, + encryptedKey: append([]byte(nil), stub.currentEnc...), + encryptedKeyLength: 140, + aesgcm: aeadgcm, + cache: testRistretto(t), + requestTimeout: 10 * time.Millisecond, + } + + nonce := make([]byte, 12) + plain := []byte("hello") + alt := make([]byte, 140) + alt[0] = 0xCC + b := wireKMSCiphertext(t, alt, nonce, plain, aeadgcm) + err = k.Decrypt(new(bytes.Buffer), bytes.NewReader(b)) + require.Error(t, err) + require.ErrorIs(t, err, context.DeadlineExceeded) +}