gsadmin/spinifex
1
0
Fork 0
mirror of https://github.com/mulgadc/spinifex.git synced 2026-10-15 09:01:31 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor: drop duplicated bluebottle logic and hand-rolled token buckets (#905)

Browse Source
This commit is contained in:
Julian Sommer
2026-08-27 10:29:06 +10:00
committed by GitHub
parent db2bf8cf0d
commit d82ee71b54
16 changed files with 99 additions and 184 deletions
Download Patch File Download Diff File Expand all files Collapse all files
cmd/spinifex/cmd/admin.go
+4 -3
View File
@@ -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)
}
cmd/spinifex/cmd/admin_ochre.go
+3 -3
View File
@@ -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)
go.mod
+2 -2
View File
@@ -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
go.sum
+2 -8
View File
@@ -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=
spinifex/daemon/daemon.go
+3 -2
View File
@@ -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)
}
spinifex/daemon/ecs_deps.go
+3 -2
View File
@@ -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
}
spinifex/daemon/eks_deps.go
+6 -4
View File
@@ -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
}
spinifex/daemon/ochre_deps.go
+2 -2
View File
@@ -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
spinifex/daemon/rds_deps.go
+2 -2
View File
@@ -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)
}
spinifex/handlers/iam/crypto.go
+1 -13
View File
@@ -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 {
spinifex/handlers/iam/crypto_test.go
+11 -9
View File
@@ -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")
}
}
spinifex/handlers/iam/instance_profiles.go
+9 -20
View File
@@ -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
spinifex/handlers/imds/dns_forwarder.go
+14 -31
View File
@@ -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
spinifex/handlers/quota/bedrock.go
+26 -74
View File
@@ -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
}
spinifex/handlers/quota/bedrock_test.go
+9 -8
View File
@@ -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
spinifex/services/awsgw/awsgw.go
+2 -1
View File
@@ -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)
}