diff --git a/go/apps/api/routes/v2_keys_verify_key/handler.go b/go/apps/api/routes/v2_keys_verify_key/handler.go index 617ad7fa61..a21e0c7e28 100644 --- a/go/apps/api/routes/v2_keys_verify_key/handler.go +++ b/go/apps/api/routes/v2_keys_verify_key/handler.go @@ -248,7 +248,7 @@ func (h *Handler) Handle(ctx context.Context, s *zen.Session) error { AutoApply: result.AutoApply, Duration: result.Duration.Milliseconds(), Exceeded: !result.Response.Success, - Id: result.Name, + Id: result.ID, Limit: result.Limit, Name: result.Name, Remaining: result.Response.Remaining, diff --git a/go/apps/api/run.go b/go/apps/api/run.go index e830af553a..53d2deb34a 100644 --- a/go/apps/api/run.go +++ b/go/apps/api/run.go @@ -31,7 +31,6 @@ import ( // nolint:gocognit func Run(ctx context.Context, cfg Config) error { - err := cfg.Validate() if err != nil { return fmt.Errorf("bad config: %w", err) @@ -183,6 +182,7 @@ func Run(ctx context.Context, cfg Config) error { RateLimiter: rlSvc, RBAC: rbac.New(), Clickhouse: ch, + Region: cfg.Region, }) if err != nil { return fmt.Errorf("unable to create key service: %w", err) diff --git a/go/internal/services/keys/get.go b/go/internal/services/keys/get.go index d5bb4da2e7..2d8dfc526c 100644 --- a/go/internal/services/keys/get.go +++ b/go/internal/services/keys/get.go @@ -105,6 +105,7 @@ func (s *service) Get(ctx context.Context, sess *zen.Session, rawKey string) (*K kv := &KeyVerifier{ Status: StatusWorkspaceDisabled, message: "workspace is disabled", + region: s.region, } if !key.WorkspaceEnabled || (key.ForWorkspaceEnabled.Valid && !key.ForWorkspaceEnabled.Bool) { @@ -152,6 +153,7 @@ func (s *service) Get(ctx context.Context, sess *zen.Session, rawKey string) (*K rBAC: s.rbac, session: sess, logger: s.logger, + region: s.region, message: "", isRootKey: key.ForWorkspaceID.Valid, diff --git a/go/internal/services/keys/interface.go b/go/internal/services/keys/interface.go index 88302f4041..40ea578300 100644 --- a/go/internal/services/keys/interface.go +++ b/go/internal/services/keys/interface.go @@ -11,8 +11,10 @@ import ( type KeyService interface { // Get retrieves a key and returns a KeyVerifier for validation Get(ctx context.Context, sess *zen.Session, hash string) (*KeyVerifier, func(), error) + // GetRootKey retrieves and validates a root key from the session GetRootKey(ctx context.Context, sess *zen.Session) (*KeyVerifier, func(), error) + // CreateKey generates a new secure API key CreateKey(ctx context.Context, req CreateKeyRequest) (CreateKeyResponse, error) } diff --git a/go/internal/services/keys/service.go b/go/internal/services/keys/service.go index f9084f390b..7a6c105c5f 100644 --- a/go/internal/services/keys/service.go +++ b/go/internal/services/keys/service.go @@ -19,6 +19,7 @@ type Config struct { RateLimiter ratelimit.Service // Rate limiting service RBAC *rbac.RBAC // Role-based access control Clickhouse clickhouse.ClickHouse // Clickhouse for telemetry + Region string // Geographic region identifier KeyCache cache.Cache[string, db.FindKeyForVerificationRow] // Cache for key lookups } @@ -30,6 +31,7 @@ type service struct { usageLimiter usagelimiter.Service rbac *rbac.RBAC clickhouse clickhouse.ClickHouse + region string // hash -> key keyCache cache.Cache[string, db.FindKeyForVerificationRow] @@ -52,6 +54,7 @@ func New(config Config) (*service, error) { raterLimiter: config.RateLimiter, usageLimiter: ulSvc, clickhouse: config.Clickhouse, + region: config.Region, keyCache: config.KeyCache, }, nil } diff --git a/go/internal/services/keys/validation.go b/go/internal/services/keys/validation.go index 89c6d37f91..cec8a3c2e6 100644 --- a/go/internal/services/keys/validation.go +++ b/go/internal/services/keys/validation.go @@ -143,6 +143,7 @@ func (k *KeyVerifier) withRateLimits(ctx context.Context, specifiedLimits []open } ratelimitsToCheck[name] = RatelimitConfigAndResult{ + ID: rl.ID, Cost: 1, Name: rl.Name, Duration: time.Duration(rl.Duration) * time.Millisecond, @@ -154,6 +155,7 @@ func (k *KeyVerifier) withRateLimits(ctx context.Context, specifiedLimits []open } for _, rl := range specifiedLimits { + // Custom limits are always applied on a key level if rl.Limit != nil && rl.Duration != nil { ratelimitsToCheck[rl.Name] = RatelimitConfigAndResult{ Cost: int64(ptr.SafeDeref(rl.Cost, 1)), @@ -161,8 +163,9 @@ func (k *KeyVerifier) withRateLimits(ctx context.Context, specifiedLimits []open Duration: time.Duration(*rl.Duration) * time.Millisecond, Limit: int64(*rl.Limit), AutoApply: false, - Identifier: k.Key.ID, // Specified limits use key ID + Identifier: k.Key.ID, Response: nil, + ID: "", // Doesn't exist and is custom so no ID } continue @@ -197,6 +200,7 @@ func (k *KeyVerifier) withRateLimits(ctx context.Context, specifiedLimits []open AutoApply: dbRl.AutoApply == 1, Identifier: identifier, Response: nil, + ID: dbRl.ID, } } diff --git a/go/internal/services/keys/verifier.go b/go/internal/services/keys/verifier.go index 2f9833e557..0f001c6d75 100644 --- a/go/internal/services/keys/verifier.go +++ b/go/internal/services/keys/verifier.go @@ -17,6 +17,7 @@ import ( // RatelimitConfigAndResult holds both the configuration and result for a rate limit type RatelimitConfigAndResult struct { + ID string Cost int64 Name string Duration time.Duration @@ -29,22 +30,29 @@ type RatelimitConfigAndResult struct { // KeyVerifier represents a key that has been loaded from the database and is ready for verification. // It contains all the necessary information and services to perform various validation checks. type KeyVerifier struct { - Key db.FindKeyForVerificationRow // The key data from the database - ratelimitConfigs map[string]db.KeyFindForVerificationRatelimit // Rate limits configured for this key (name -> config) - Roles []string // RBAC roles assigned to this key - Permissions []string // RBAC permissions assigned to this key - Status KeyStatus // The current validation status - AuthorizedWorkspaceID string // The workspace ID this key is authorized for - RatelimitResults map[string]RatelimitConfigAndResult // Combined config and results for rate limits (name -> config+result) - isRootKey bool // Whether this is a root key (special handling) - session *zen.Session // The current request session - rateLimiter ratelimit.Service // Rate limiting service - usageLimiter usagelimiter.Service // Usage limiting service - rBAC *rbac.RBAC // Role-based access control service - clickhouse clickhouse.ClickHouse // Clickhouse for telemetry - logger logging.Logger // Logger for verification operations - message string // Internal message for validation failures - tags []string // Tags associated with the key + Key db.FindKeyForVerificationRow // The key data from the database + Roles []string // RBAC roles assigned to this key + Permissions []string // RBAC permissions assigned to this key + Status KeyStatus // The current validation status + AuthorizedWorkspaceID string // The workspace ID this key is authorized for + + ratelimitConfigs map[string]db.KeyFindForVerificationRatelimit // Rate limits configured for this key (name -> config) + RatelimitResults map[string]RatelimitConfigAndResult // Combined config and results for rate limits (name -> config+result) + + isRootKey bool // Whether this is a root key (special handling) + + message string // Internal message for validation failures + tags []string // Tags associated with this verification + + session *zen.Session // The current request session + region string // Geographic region identifier + + // Services + rateLimiter ratelimit.Service // Rate limiting service + usageLimiter usagelimiter.Service // Usage limiting service + rBAC *rbac.RBAC // Role-based access control service + clickhouse clickhouse.ClickHouse // Clickhouse for telemetry + logger logging.Logger // Logger for verification operations } // GetRatelimitConfigs returns the rate limit configurations @@ -117,12 +125,12 @@ func (k *KeyVerifier) log() { RequestID: k.session.RequestID(), WorkspaceID: k.Key.WorkspaceID, Time: time.Now().UnixMilli(), - Region: "", Outcome: string(k.Status), KeySpaceID: k.Key.KeyAuthID, KeyID: k.Key.ID, IdentityID: k.Key.IdentityID.String, Tags: k.tags, + Region: k.region, }) keyType := "key"