From d82ee71b54ea9338a974e7631692c0677f32dd2b Mon Sep 17 00:00:00 2001 From: Julian Sommer <115516918+juliansommer@users.noreply.github.com> Date: Thu, 27 Aug 2026 10:29:06 +1000 Subject: [PATCH] refactor: drop duplicated bluebottle logic and hand-rolled token buckets (#905) --- cmd/spinifex/cmd/admin.go | 7 +- cmd/spinifex/cmd/admin_ochre.go | 6 +- go.mod | 4 +- go.sum | 10 +-- spinifex/daemon/daemon.go | 5 +- spinifex/daemon/ecs_deps.go | 5 +- spinifex/daemon/eks_deps.go | 10 ++- spinifex/daemon/ochre_deps.go | 4 +- spinifex/daemon/rds_deps.go | 4 +- spinifex/handlers/iam/crypto.go | 14 +-- spinifex/handlers/iam/crypto_test.go | 20 +++-- spinifex/handlers/iam/instance_profiles.go | 29 ++---- spinifex/handlers/imds/dns_forwarder.go | 45 +++------- spinifex/handlers/quota/bedrock.go | 100 ++++++--------------- spinifex/handlers/quota/bedrock_test.go | 17 ++-- spinifex/services/awsgw/awsgw.go | 3 +- 16 files changed, 99 insertions(+), 184 deletions(-) diff --git a/cmd/spinifex/cmd/admin.go b/cmd/spinifex/cmd/admin.go index eb25dbf6b..527ea8303 100644 --- a/cmd/spinifex/cmd/admin.go +++ b/cmd/spinifex/cmd/admin.go @@ -27,6 +27,7 @@ import ( "time" "uuid" + "github.com/mulgadc/bluebottle/pkg/masterkey" "github.com/mulgadc/spinifex/spinifex/admin" "github.com/mulgadc/spinifex/spinifex/config" "github.com/mulgadc/spinifex/spinifex/ebsmetadata" @@ -2578,7 +2579,7 @@ func initIAMServiceFromConfig() (*handlers_iam.IAMServiceImpl, *config.ClusterCo } masterKeyPath := filepath.Join(cfg.NodeBaseDir(), "config", "master.key") - masterKey, err := handlers_iam.LoadMasterKey(masterKeyPath) + masterKey, err := masterkey.ReadShared(masterKeyPath) if err != nil { nc.Close() return nil, nil, nil, nil, fmt.Errorf("load master key: %w", err) @@ -2984,7 +2985,7 @@ type writeBootstrapResult struct { func ensureMasterKey(configDir string) (key []byte, existed bool, err error) { keyPath := filepath.Join(configDir, "master.key") if admin.FileExists(keyPath) { - key, err = handlers_iam.LoadMasterKey(keyPath) + key, err = masterkey.ReadShared(keyPath) if err != nil { return nil, false, fmt.Errorf("load master key: %w", err) } @@ -3077,7 +3078,7 @@ func ensureViperblockEncryptionKey(configDir string) ([]byte, string, error) { } keyPath := filepath.Join(viperblockDir, "encryption.key") if admin.FileExists(keyPath) { - key, err := handlers_iam.LoadMasterKey(keyPath) + key, err := masterkey.ReadShared(keyPath) if err != nil { return nil, "", fmt.Errorf("load viperblock encryption key: %w", err) } diff --git a/cmd/spinifex/cmd/admin_ochre.go b/cmd/spinifex/cmd/admin_ochre.go index 3878bd9fd..a223ded41 100644 --- a/cmd/spinifex/cmd/admin_ochre.go +++ b/cmd/spinifex/cmd/admin_ochre.go @@ -16,13 +16,13 @@ import ( "github.com/aws/aws-sdk-go/aws" "github.com/aws/aws-sdk-go/service/s3" + "github.com/mulgadc/bluebottle/pkg/masterkey" "github.com/mulgadc/spinifex/spinifex/admin" "github.com/mulgadc/spinifex/spinifex/config" "github.com/mulgadc/spinifex/spinifex/ebsprovider" gateway_bedrock "github.com/mulgadc/spinifex/spinifex/gateway/bedrock" "github.com/mulgadc/spinifex/spinifex/gateway/bedrock/hfhub" handlers_ec2_snapshot "github.com/mulgadc/spinifex/spinifex/handlers/ec2/snapshot" - handlers_iam "github.com/mulgadc/spinifex/spinifex/handlers/iam" "github.com/mulgadc/spinifex/spinifex/objectstore" "github.com/mulgadc/spinifex/spinifex/utils" "github.com/nats-io/nats.go" @@ -999,7 +999,7 @@ func resolveHFToken(ctx context.Context, cmd *cobra.Command, cfg *config.Cluster } masterKeyPath := filepath.Join(cfg.NodeBaseDir(), "config", "master.key") - masterKey, err := handlers_iam.LoadMasterKey(masterKeyPath) + masterKey, err := masterkey.ReadShared(masterKeyPath) if err != nil { return "" } @@ -1071,7 +1071,7 @@ func ochreCredentialsStore() (*gateway_bedrock.CredentialStore, func(), error) { return nil, nil, fmt.Errorf("connect to cluster: %w", err) } masterKeyPath := filepath.Join(cfg.NodeBaseDir(), "config", "master.key") - masterKey, err := handlers_iam.LoadMasterKey(masterKeyPath) + masterKey, err := masterkey.ReadShared(masterKeyPath) if err != nil { nc.Close() return nil, nil, fmt.Errorf("load master key: %w", err) diff --git a/go.mod b/go.mod index de8a1c754..c4050b4d2 100644 --- a/go.mod +++ b/go.mod @@ -21,7 +21,7 @@ require ( github.com/jackc/pgx/v5 v5.10.0 github.com/klauspost/cpuid/v2 v2.4.0 github.com/miekg/dns v1.1.73 - github.com/mulgadc/bluebottle v1.18.1-0.20260826232854-f91084812667 + github.com/mulgadc/bluebottle v1.18.1-0.20260827002656-8dfaf43f84e6 github.com/mulgadc/northstar v1.18.0 github.com/mulgadc/predastore v1.18.1-0.20260826045109-26ce1b091112 github.com/mulgadc/viperblock v1.18.0 @@ -45,6 +45,7 @@ require ( golang.org/x/crypto v0.55.0 golang.org/x/sync v0.22.0 golang.org/x/sys v0.47.0 + golang.org/x/time v0.15.0 golang.org/x/tools v0.49.0 gopkg.in/ini.v1 v1.67.3 ) @@ -200,7 +201,6 @@ require ( golang.org/x/telemetry v0.0.0-20260811182544-a038080d80e5 // indirect golang.org/x/term v0.45.0 // indirect golang.org/x/text v0.41.0 // indirect - golang.org/x/time v0.15.0 // indirect golang.org/x/vuln v1.6.0 // indirect google.golang.org/genproto/googleapis/api v0.0.0-20260803160001-6ac0973c030d // indirect google.golang.org/genproto/googleapis/rpc v0.0.0-20260803160001-6ac0973c030d // indirect diff --git a/go.sum b/go.sum index 283cd0f84..b4c6feeec 100644 --- a/go.sum +++ b/go.sum @@ -1327,16 +1327,10 @@ github.com/muesli/cancelreader v0.2.2 h1:3I4Kt4BQjOR54NavqnDogx/MIoWBFa0StPA8ELU github.com/muesli/cancelreader v0.2.2/go.mod h1:3XuTXfFS2VjM+HTLZY9Ak0l6eUKfijIfMUZ4EgX0QYo= github.com/muesli/termenv v0.16.0 h1:S5AlUN9dENB57rsbnkPyfdGuWIlkmzJjbFf0Tf5FWUc= github.com/muesli/termenv v0.16.0/go.mod h1:ZRfOIKPFDYQoDFF4Olj7/QJbW60Ol/kL1pU3VfY/Cnk= -github.com/mulgadc/bluebottle v1.18.0 h1:aNrogJBz59gJAxCwpERGBC5bGXkSoUCCCBbNvnamuRU= -github.com/mulgadc/bluebottle v1.18.0/go.mod h1:KkUyzqfKxyI8PvKR27EVqJnVm4BSbajl1Q/u+PQxAYw= -github.com/mulgadc/bluebottle v1.18.1-0.20260826044308-1a1e1242c217 h1:zRKEFadLkUQxNCK2L3MYSQxRmu4VEEHWTxMwh04Qmeo= -github.com/mulgadc/bluebottle v1.18.1-0.20260826044308-1a1e1242c217/go.mod h1:KkUyzqfKxyI8PvKR27EVqJnVm4BSbajl1Q/u+PQxAYw= -github.com/mulgadc/bluebottle v1.18.1-0.20260826232854-f91084812667 h1:3q1aWTrGgQBSicNQ68uM/wYFkmaNld76wCr/VsnLC4s= -github.com/mulgadc/bluebottle v1.18.1-0.20260826232854-f91084812667/go.mod h1:KkUyzqfKxyI8PvKR27EVqJnVm4BSbajl1Q/u+PQxAYw= +github.com/mulgadc/bluebottle v1.18.1-0.20260827002656-8dfaf43f84e6 h1:MEk7iPcOKgxWtRtJge8jiHjypiXFXEW8AfACOo3oN5U= +github.com/mulgadc/bluebottle v1.18.1-0.20260827002656-8dfaf43f84e6/go.mod h1:mvvV4d2w+kGLXxO1ymptNlxqUlNJKB2oAhk6PkpVLj8= github.com/mulgadc/northstar v1.18.0 h1:ksRjIKE6qtrw/ZoYaJh8JRBh+xdnuDw8FsmFLFCiXbg= github.com/mulgadc/northstar v1.18.0/go.mod h1:z6iol77R1vsIQyIBc1JQboUPhSJ6eG4bipzfbZZbLEM= -github.com/mulgadc/predastore v1.18.0 h1:RW1Qw1SM2iYxqi6HaHLv5Qb1h1vhZSO3x6pOhPAb9mw= -github.com/mulgadc/predastore v1.18.0/go.mod h1:+VBT3ofH/YmdTg0ug9EQwHDb8bQbHaMvWUyh5cIpBeo= github.com/mulgadc/predastore v1.18.1-0.20260826045109-26ce1b091112 h1:eCzKJDL1lmw6DNwjbDCQBoT3ZHIsYE1lMj5qitjySrY= github.com/mulgadc/predastore v1.18.1-0.20260826045109-26ce1b091112/go.mod h1:xDJMcTuHTG4aKfO2PrSlkpID/UJrvYj1yfLZ0Pi8CSU= github.com/mulgadc/viperblock v1.18.0 h1:u6yu5lPfZ+RSmfzlg1OXI3whlRBnGp8t685xUzzxlns= diff --git a/spinifex/daemon/daemon.go b/spinifex/daemon/daemon.go index dd8515d37..6b474a6e3 100644 --- a/spinifex/daemon/daemon.go +++ b/spinifex/daemon/daemon.go @@ -28,6 +28,7 @@ import ( "github.com/aws/aws-sdk-go/aws" "github.com/aws/aws-sdk-go/service/ec2" "github.com/go-chi/chi/v5" + "github.com/mulgadc/bluebottle/pkg/masterkey" "github.com/mulgadc/spinifex/internal/tlsconfig" "github.com/mulgadc/spinifex/spinifex/admin" "github.com/mulgadc/spinifex/spinifex/awserrors" @@ -1737,7 +1738,7 @@ func (d *Daemon) startCluster() error { // safe degraded mode for certificate private keys, so a missing key must // fail daemon startup rather than leave HTTPS listeners uncreatable. d.elbv2Service, err = initServiceWithRetry("ELBv2 service", func() (*handlers_elbv2.ELBv2ServiceImpl, error) { - masterKey, mkErr := handlers_iam.LoadMasterKey(filepath.Join(filepath.Dir(d.configPath), "master.key")) + masterKey, mkErr := masterkey.ReadShared(filepath.Join(filepath.Dir(d.configPath), "master.key")) if mkErr != nil { return nil, fmt.Errorf("load ELBv2 master key: %w", mkErr) } @@ -1872,7 +1873,7 @@ func (d *Daemon) startCluster() error { // concurrent boot — but a master key that never arrives fails startCluster // after the retry window instead of leaving acmService permanently nil. d.acmService, err = initServiceWithRetry("ACM service", func() (*handlers_acm.ACMServiceImpl, error) { - masterKey, mkErr := handlers_iam.LoadMasterKey(filepath.Join(filepath.Dir(d.configPath), "master.key")) + masterKey, mkErr := masterkey.ReadShared(filepath.Join(filepath.Dir(d.configPath), "master.key")) if mkErr != nil { return nil, fmt.Errorf("load ACM master key: %w", mkErr) } diff --git a/spinifex/daemon/ecs_deps.go b/spinifex/daemon/ecs_deps.go index f2bfd5e22..1542f425f 100644 --- a/spinifex/daemon/ecs_deps.go +++ b/spinifex/daemon/ecs_deps.go @@ -5,6 +5,7 @@ import ( "os" "path/filepath" + "github.com/mulgadc/bluebottle/pkg/masterkey" handlers_ecs "github.com/mulgadc/spinifex/spinifex/handlers/ecs" handlers_iam "github.com/mulgadc/spinifex/spinifex/handlers/iam" ) @@ -35,9 +36,9 @@ func (d *Daemon) buildECSServiceDeps() handlers_ecs.Deps { // find-or-create the ecsInstanceRole instance profile container instances // expose over IMDS. Only wired when the master key is present; without it // capacity provisioning is disabled. - masterKey, err := handlers_iam.LoadMasterKey(filepath.Join(filepath.Dir(d.configPath), "master.key")) + masterKey, err := masterkey.ReadShared(filepath.Join(filepath.Dir(d.configPath), "master.key")) if err != nil || masterKey == nil { - slog.Warn("ECS: LoadMasterKey failed; capacity provisioning disabled until the master key is provisioned", + slog.Warn("ECS: master key read failed; capacity provisioning disabled until the master key is provisioned", "err", err) return deps } diff --git a/spinifex/daemon/eks_deps.go b/spinifex/daemon/eks_deps.go index 2d71abde2..6c1531210 100644 --- a/spinifex/daemon/eks_deps.go +++ b/spinifex/daemon/eks_deps.go @@ -5,11 +5,13 @@ import ( "os" "path/filepath" + "log/slog" + + "github.com/mulgadc/bluebottle/pkg/masterkey" handlers_ec2_placementgroup "github.com/mulgadc/spinifex/spinifex/handlers/ec2/placementgroup" handlers_eks "github.com/mulgadc/spinifex/spinifex/handlers/eks" handlers_iam "github.com/mulgadc/spinifex/spinifex/handlers/iam" "github.com/mulgadc/spinifex/spinifex/objectstore" - "log/slog" ) // systemRoleEnsurer lazily builds — and memoizes — the KV-backed IAM service @@ -27,7 +29,7 @@ func (d *Daemon) systemRoleEnsurer() handlers_iam.SystemInstanceRoleEnsurer { if d.iamEnsurerCached != nil { return d.iamEnsurerCached } - masterKey, err := handlers_iam.LoadMasterKey(filepath.Join(filepath.Dir(d.configPath), "master.key")) + masterKey, err := masterkey.ReadShared(filepath.Join(filepath.Dir(d.configPath), "master.key")) if err != nil || masterKey == nil { slog.Warn("System role ensurer: master key unavailable; system VMs fall back to baked static creds", "err", err) @@ -50,9 +52,9 @@ func (d *Daemon) systemRoleEnsurer() handlers_iam.SystemInstanceRoleEnsurer { // buildEKSServiceDeps assembles the EKSServiceDeps. All collaborators must // already be initialised. MasterKey is loaded best-effort from the config dir. func (d *Daemon) buildEKSServiceDeps() handlers_eks.EKSServiceDeps { - masterKey, err := handlers_iam.LoadMasterKey(filepath.Join(filepath.Dir(d.configPath), "master.key")) + masterKey, err := masterkey.ReadShared(filepath.Join(filepath.Dir(d.configPath), "master.key")) if err != nil { - slog.Warn("EKS: LoadMasterKey failed; CreateCluster + DeleteCluster will reject until key is provisioned", + slog.Warn("EKS: master key read failed; CreateCluster + DeleteCluster will reject until key is provisioned", "err", err) masterKey = nil } diff --git a/spinifex/daemon/ochre_deps.go b/spinifex/daemon/ochre_deps.go index 6d5599be1..bc91d729a 100644 --- a/spinifex/daemon/ochre_deps.go +++ b/spinifex/daemon/ochre_deps.go @@ -8,10 +8,10 @@ import ( "path/filepath" "time" + "github.com/mulgadc/bluebottle/pkg/masterkey" "github.com/mulgadc/spinifex/spinifex/admin" gateway_bedrock "github.com/mulgadc/spinifex/spinifex/gateway/bedrock" handlers_bedrock "github.com/mulgadc/spinifex/spinifex/handlers/bedrock" - handlers_iam "github.com/mulgadc/spinifex/spinifex/handlers/iam" handlers_ochrevector "github.com/mulgadc/spinifex/spinifex/handlers/ochrevector" "github.com/mulgadc/spinifex/spinifex/network/host" "github.com/mulgadc/spinifex/spinifex/objectstore" @@ -60,7 +60,7 @@ func (d *Daemon) startOchreVector() { return } - masterKey, err := handlers_iam.LoadMasterKey(filepath.Join(filepath.Dir(d.configPath), "master.key")) + masterKey, err := masterkey.ReadShared(filepath.Join(filepath.Dir(d.configPath), "master.key")) if err != nil { slog.Warn("Ochre vector store disabled: master key unavailable", "err", err) return diff --git a/spinifex/daemon/rds_deps.go b/spinifex/daemon/rds_deps.go index 35aed8c6e..e9bef366b 100644 --- a/spinifex/daemon/rds_deps.go +++ b/spinifex/daemon/rds_deps.go @@ -9,8 +9,8 @@ import ( "time" "github.com/aws/aws-sdk-go/service/ec2" + "github.com/mulgadc/bluebottle/pkg/masterkey" gateway_ec2_instance "github.com/mulgadc/spinifex/spinifex/gateway/ec2/instance" - handlers_iam "github.com/mulgadc/spinifex/spinifex/handlers/iam" handlers_rds "github.com/mulgadc/spinifex/spinifex/handlers/rds" handlers_systemvpc "github.com/mulgadc/spinifex/spinifex/handlers/systemvpc" ) @@ -53,7 +53,7 @@ func (d *Daemon) buildRDSDeps() (handlers_rds.Deps, error) { // No degraded mode: a create cannot stage a master password without it, and // the fetch that replays one is answered by any node in the queue group. - masterKey, err := handlers_iam.LoadMasterKey(filepath.Join(filepath.Dir(d.configPath), "master.key")) + masterKey, err := masterkey.ReadShared(filepath.Join(filepath.Dir(d.configPath), "master.key")) if err != nil { return handlers_rds.Deps{}, fmt.Errorf("load RDS master key: %w", err) } diff --git a/spinifex/handlers/iam/crypto.go b/spinifex/handlers/iam/crypto.go index 29e42fa76..ea85f1704 100644 --- a/spinifex/handlers/iam/crypto.go +++ b/spinifex/handlers/iam/crypto.go @@ -9,7 +9,7 @@ import ( "github.com/mulgadc/bluebottle/pkg/masterkey" ) -const masterKeySize = 32 // AES-256 +const masterKeySize = masterkey.MasterKeySize // BootstrapVersion is the current bootstrap.json config version. const BootstrapVersion = "1.0" @@ -24,18 +24,6 @@ func GenerateMasterKey() ([]byte, error) { return key, nil } -// LoadMasterKey reads a master key from disk and validates it is exactly 32 bytes. -func LoadMasterKey(path string) ([]byte, error) { - key, err := os.ReadFile(path) - if err != nil { - return nil, fmt.Errorf("read master key: %w", err) - } - if len(key) != masterKeySize { - return nil, fmt.Errorf("master key must be %d bytes, got %d", masterKeySize, len(key)) - } - return key, nil -} - // SaveMasterKey writes a master key to disk with 0600 permissions. func SaveMasterKey(path string, key []byte) error { if len(key) != masterKeySize { diff --git a/spinifex/handlers/iam/crypto_test.go b/spinifex/handlers/iam/crypto_test.go index b19115954..48b9b6e8c 100644 --- a/spinifex/handlers/iam/crypto_test.go +++ b/spinifex/handlers/iam/crypto_test.go @@ -5,6 +5,8 @@ import ( "os" "path/filepath" "testing" + + "github.com/mulgadc/bluebottle/pkg/masterkey" ) func TestGenerateMasterKey(t *testing.T) { @@ -26,7 +28,7 @@ func TestGenerateMasterKey(t *testing.T) { } } -func TestSaveLoadMasterKey(t *testing.T) { +func TestSaveAndReadMasterKey(t *testing.T) { key, err := GenerateMasterKey() if err != nil { t.Fatalf("GenerateMasterKey() error: %v", err) @@ -48,16 +50,16 @@ func TestSaveLoadMasterKey(t *testing.T) { t.Fatalf("expected permissions 0600, got %04o", perm) } - loaded, err := LoadMasterKey(path) + loaded, err := masterkey.ReadShared(path) if err != nil { - t.Fatalf("LoadMasterKey() error: %v", err) + t.Fatalf("masterkey.ReadShared() error: %v", err) } if string(loaded) != string(key) { t.Fatal("loaded key does not match saved key") } } -func TestLoadMasterKeyWrongSize(t *testing.T) { +func TestReadMasterKeyWrongSize(t *testing.T) { dir := t.TempDir() path := filepath.Join(dir, "bad.key") @@ -65,16 +67,16 @@ func TestLoadMasterKeyWrongSize(t *testing.T) { t.Fatal(err) } - _, err := LoadMasterKey(path) + _, err := masterkey.ReadShared(path) if err == nil { - t.Fatal("LoadMasterKey() should fail for wrong-size key") + t.Fatal("masterkey.ReadShared() should fail for wrong-size key") } } -func TestLoadMasterKeyNotFound(t *testing.T) { - _, err := LoadMasterKey("/nonexistent/master.key") +func TestReadMasterKeyNotFound(t *testing.T) { + _, err := masterkey.ReadShared("/nonexistent/master.key") if err == nil { - t.Fatal("LoadMasterKey() should fail for missing file") + t.Fatal("masterkey.ReadShared() should fail for missing file") } } diff --git a/spinifex/handlers/iam/instance_profiles.go b/spinifex/handlers/iam/instance_profiles.go index 57feb7d28..c8834cc49 100644 --- a/spinifex/handlers/iam/instance_profiles.go +++ b/spinifex/handlers/iam/instance_profiles.go @@ -13,6 +13,8 @@ import ( "github.com/aws/aws-sdk-go/service/iam" "github.com/nats-io/nats.go/jetstream" + "github.com/mulgadc/bluebottle/pkg/auth" + "github.com/mulgadc/spinifex/spinifex/arn" "github.com/mulgadc/spinifex/spinifex/awserrors" "github.com/mulgadc/spinifex/spinifex/kvutil" @@ -268,28 +270,15 @@ func (s *IAMServiceImpl) ResolveInstanceProfile(accountID, nameOrARN string) (*I return s.getInstanceProfile(ctx, accountID, profileName) } -// parseInstanceProfileARN extracts accountID and profile name from an IAM instance-profile ARN. -func parseInstanceProfileARN(arn string) (accountID, name string, err error) { - parts := strings.SplitN(arn, ":", 6) - if len(parts) != 6 || parts[0] != "arn" || parts[1] != "aws" || parts[2] != "iam" || parts[3] != "" { +// parseInstanceProfileARN extracts accountID and profile name from an IAM +// instance-profile ARN, reporting every malformed shape as the one AWS error +// code this API returns. +func parseInstanceProfileARN(arnStr string) (accountID, name string, err error) { + accountID, name, err = auth.ParseInstanceProfileARN(arnStr) + if err != nil { return "", "", errors.New(awserrors.ErrorInvalidIamInstanceProfileArnMalformed) } - resource := parts[5] - const prefix = "instance-profile/" - if !strings.HasPrefix(resource, prefix) { - return "", "", errors.New(awserrors.ErrorInvalidIamInstanceProfileArnMalformed) - } - pathAndName := resource[len(prefix):] - slash := strings.LastIndex(pathAndName, "/") - if slash == -1 { - name = pathAndName - } else { - name = pathAndName[slash+1:] - } - if name == "" { - return "", "", errors.New(awserrors.ErrorInvalidIamInstanceProfileArnMalformed) - } - return parts[4], name, nil + return accountID, name, nil } // TagInstanceProfile upserts tags on an instance profile under CAS, like the diff --git a/spinifex/handlers/imds/dns_forwarder.go b/spinifex/handlers/imds/dns_forwarder.go index 8400010a0..1cc9dde92 100644 --- a/spinifex/handlers/imds/dns_forwarder.go +++ b/spinifex/handlers/imds/dns_forwarder.go @@ -8,6 +8,8 @@ import ( "net" "sync" "time" + + "golang.org/x/time/rate" ) // maxDNSUDPSize is the largest UDP DNS message the shim relays; EDNS0 clients @@ -283,43 +285,24 @@ func (f *dnsForwarder) proxyTCP(client net.Conn) { // — each serveUDP/serveTCP invocation owns one — so the limit is inherently // per-tap. now is injected so tests drive refill deterministically. type tokenBucket struct { - mu sync.Mutex - tokens float64 - burst float64 - rate float64 - last time.Time - now func() time.Time + limiter *rate.Limiter + now func() time.Time } -func newTokenBucket(rate, burst int, now func() time.Time) *tokenBucket { - return &tokenBucket{ - tokens: float64(burst), - burst: float64(burst), - rate: float64(rate), - last: now(), - now: now, +// newTokenBucket returns a full bucket at burst capacity refilling at rate per +// second. A non-positive rate disables limiting: rate.Inf admits everything +// regardless of burst. +func newTokenBucket(perSecond, burst int, now func() time.Time) *tokenBucket { + limit := rate.Limit(perSecond) + if perSecond <= 0 { + limit = rate.Inf } + return &tokenBucket{limiter: rate.NewLimiter(limit, burst), now: now} } -// allow refills by elapsed time and consumes one token, reporting whether the -// query may proceed. A non-positive rate disables limiting (always allow). +// allow consumes one token, reporting whether the query may proceed. func (tb *tokenBucket) allow() bool { - if tb.rate <= 0 { - return true - } - tb.mu.Lock() - defer tb.mu.Unlock() - now := tb.now() - tb.tokens += now.Sub(tb.last).Seconds() * tb.rate - if tb.tokens > tb.burst { - tb.tokens = tb.burst - } - tb.last = now - if tb.tokens >= 1 { - tb.tokens-- - return true - } - return false + return tb.limiter.AllowN(tb.now(), 1) } // bindTapDNS opens the DNS shim sockets — UDP and TCP 169.254.169.253:53 on the diff --git a/spinifex/handlers/quota/bedrock.go b/spinifex/handlers/quota/bedrock.go index 184d541a4..85ce7a7ed 100644 --- a/spinifex/handlers/quota/bedrock.go +++ b/spinifex/handlers/quota/bedrock.go @@ -12,6 +12,7 @@ import ( "github.com/mulgadc/spinifex/spinifex/kvutil" "github.com/mulgadc/spinifex/spinifex/utils" "github.com/nats-io/nats.go/jetstream" + "golang.org/x/time/rate" ) // BedrockUsageReader resolves accountID's current-period Bedrock token usage @@ -74,104 +75,54 @@ func (s *Service) CheckBedrockRPM(accountID string) error { return nil } -// tokenBucket is a classic token-bucket rate limiter: capacity tokens -// refilled continuously at capacity/60 per second, consumed one at a time. -// Refill is lazy — computed from elapsed wall-clock time on each Allow call — -// so no background goroutine is required per account. -type tokenBucket struct { - mu sync.Mutex - tokens float64 - capacity float64 - refillRate float64 // tokens per second - last time.Time - now func() time.Time -} - // newTokenBucket constructs a full bucket (a fresh account starts able to // burst its whole per-minute allowance, not throttled from zero) at -// capacityPerMinute, using now for its clock — overridable in tests, nil -// defaults to time.Now. -func newTokenBucket(capacityPerMinute int, now func() time.Time) *tokenBucket { - if now == nil { - now = time.Now - } - capacity := float64(max(capacityPerMinute, 0)) - return &tokenBucket{ - tokens: capacity, - capacity: capacity, - refillRate: capacity / 60, - last: now(), - now: now, - } +// capacityPerMinute, refilling continuously at capacity/60 per second. +func newTokenBucket(capacityPerMinute int) *rate.Limiter { + capacity := max(capacityPerMinute, 0) + return rate.NewLimiter(rate.Limit(float64(capacity)/60), capacity) } -// Allow refills for elapsed time since the last call, then attempts to -// consume one token. Returns false (throttle) when the bucket is empty. -func (b *tokenBucket) Allow() bool { - b.mu.Lock() - defer b.mu.Unlock() - b.refillLocked() - if b.tokens < 1 { - return false - } - b.tokens-- - return true -} - -// remaining reports the bucket's current occupancy (refilled up to now) -// without consuming a token, for the periodic KV sync's observability -// snapshot. -func (b *tokenBucket) remaining() int { - b.mu.Lock() - defer b.mu.Unlock() - b.refillLocked() - return int(b.tokens) -} - -// refillLocked adds tokens for the time elapsed since last, capped at -// capacity. Caller must hold mu. -func (b *tokenBucket) refillLocked() { - current := b.now() - elapsed := current.Sub(b.last).Seconds() - if elapsed <= 0 { - return - } - b.tokens = min(b.capacity, b.tokens+elapsed*b.refillRate) - b.last = current -} - -// rpmLimiter holds one tokenBucket per account, created lazily on first use -// at the configured per-minute capacity. Buckets are never evicted: an idle -// account costs one small map entry, never unbounded growth, and a restart -// resets every bucket to full — acceptable since RPM enforcement is already +// rpmLimiter holds one bucket per account, created lazily on first use at the +// configured per-minute capacity. Buckets are never evicted: an idle account +// costs one small map entry, never unbounded growth, and a restart resets +// every bucket to full — acceptable since RPM enforcement is already // documented as generous, not exact. type rpmLimiter struct { mu sync.Mutex - buckets map[string]*tokenBucket - // now overrides tokenBucket's clock for every bucket this limiter - // creates; nil (the production default) leaves each bucket on time.Now. + buckets map[string]*rate.Limiter + // now overrides the clock every bucket is measured against; nil (the + // production default) means time.Now. now func() time.Time } func newRPMLimiter() *rpmLimiter { - return &rpmLimiter{buckets: make(map[string]*tokenBucket)} + return &rpmLimiter{buckets: make(map[string]*rate.Limiter)} +} + +// clock returns the limiter's current time, honouring a test override. +func (l *rpmLimiter) clock() time.Time { + if l.now == nil { + return time.Now() + } + return l.now() } // allow returns whether accountID may proceed under capacityPerMinute, // creating its bucket on first use. func (l *rpmLimiter) allow(accountID string, capacityPerMinute int) bool { - return l.bucketFor(accountID, capacityPerMinute).Allow() + return l.bucketFor(accountID, capacityPerMinute).AllowN(l.clock(), 1) } // bucketFor returns accountID's bucket, creating it at capacityPerMinute on // first use. Capacity is fixed at creation: RPM limits are process-lifetime // config (loaded once at gateway startup), not hot-reloaded. -func (l *rpmLimiter) bucketFor(accountID string, capacityPerMinute int) *tokenBucket { +func (l *rpmLimiter) bucketFor(accountID string, capacityPerMinute int) *rate.Limiter { l.mu.Lock() defer l.mu.Unlock() b, ok := l.buckets[accountID] if !ok { - b = newTokenBucket(capacityPerMinute, l.now) + b = newTokenBucket(capacityPerMinute) l.buckets[accountID] = b } return b @@ -180,11 +131,12 @@ func (l *rpmLimiter) bucketFor(accountID string, capacityPerMinute int) *tokenBu // snapshot returns every known account's current bucket occupancy, for the // periodic KV sync below. func (l *rpmLimiter) snapshot() map[string]int { + now := l.clock() l.mu.Lock() defer l.mu.Unlock() out := make(map[string]int, len(l.buckets)) for accountID, b := range l.buckets { - out[accountID] = b.remaining() + out[accountID] = int(b.TokensAt(now)) } return out } diff --git a/spinifex/handlers/quota/bedrock_test.go b/spinifex/handlers/quota/bedrock_test.go index cf96a1844..746934c64 100644 --- a/spinifex/handlers/quota/bedrock_test.go +++ b/spinifex/handlers/quota/bedrock_test.go @@ -133,26 +133,27 @@ func TestCheckBedrockRPM_ThrottlesAfterCapacityExhausted(t *testing.T) { // refill behaviour is deterministic rather than sleep-based. func TestTokenBucket_RefillsOverTime(t *testing.T) { now := time.Now() - clock := func() time.Time { return now } - b := newTokenBucket(60, clock) // 60/min == 1/sec + l := newRPMLimiter() + l.now = func() time.Time { return now } + const capacity = 60 // 60/min == 1/sec // Drain the bucket. - for range 60 { - require.True(t, b.Allow()) + for range capacity { + require.True(t, l.allow(testAccount, capacity)) } - assert.False(t, b.Allow(), "bucket must be empty immediately after draining") + assert.False(t, l.allow(testAccount, capacity), "bucket must be empty immediately after draining") // Advance 5 seconds: 5 tokens should have refilled. now = now.Add(5 * time.Second) for i := range 5 { - assert.Truef(t, b.Allow(), "refilled token %d", i) + assert.Truef(t, l.allow(testAccount, capacity), "refilled token %d", i) } - assert.False(t, b.Allow(), "no more than the refilled amount is available") + assert.False(t, l.allow(testAccount, capacity), "no more than the refilled amount is available") // Advance well past capacity: refill caps at capacity, it never // overflows from an idle account banking unlimited tokens. now = now.Add(10 * time.Minute) - assert.Equal(t, 60, b.remaining()) + assert.Equal(t, capacity, l.snapshot()[testAccount]) } // TestRPMLimiter_PerAccountIsolation covers that one account's bucket never diff --git a/spinifex/services/awsgw/awsgw.go b/spinifex/services/awsgw/awsgw.go index c72c7c16e..3e0caa8f6 100644 --- a/spinifex/services/awsgw/awsgw.go +++ b/spinifex/services/awsgw/awsgw.go @@ -14,6 +14,7 @@ import ( "strings" "time" + "github.com/mulgadc/bluebottle/pkg/masterkey" "github.com/mulgadc/bluebottle/pkg/ratelimit" "github.com/mulgadc/spinifex/internal/tlsconfig" "github.com/mulgadc/spinifex/spinifex/admin" @@ -218,7 +219,7 @@ func launchService(config *config.ClusterConfig) error { // Load IAM master key from disk (required for all authenticated requests) masterKeyPath := filepath.Join(nodeConfig.BaseDir, "config", "master.key") - masterKey, err := handlers_iam.LoadMasterKey(masterKeyPath) + masterKey, err := masterkey.ReadShared(masterKeyPath) if err != nil { return fmt.Errorf("load IAM master key from %s: %w", masterKeyPath, err) }