diff --git a/go/apps/api/routes/v2_keys_add_permissions/handler.go b/go/apps/api/routes/v2_keys_add_permissions/handler.go index 9d6deb2663..dce321c23d 100644 --- a/go/apps/api/routes/v2_keys_add_permissions/handler.go +++ b/go/apps/api/routes/v2_keys_add_permissions/handler.go @@ -30,7 +30,7 @@ type Handler struct { DB db.Database Keys keys.KeyService Auditlogs auditlogs.AuditLogService - KeyCache cache.Cache[string, db.FindKeyForVerificationRow] + KeyCache cache.Cache[string, db.CachedKeyData] } // Method returns the HTTP method this route responds to diff --git a/go/apps/api/routes/v2_keys_add_roles/handler.go b/go/apps/api/routes/v2_keys_add_roles/handler.go index f11160a75f..f6306abb5c 100644 --- a/go/apps/api/routes/v2_keys_add_roles/handler.go +++ b/go/apps/api/routes/v2_keys_add_roles/handler.go @@ -29,7 +29,7 @@ type Handler struct { DB db.Database Keys keys.KeyService Auditlogs auditlogs.AuditLogService - KeyCache cache.Cache[string, db.FindKeyForVerificationRow] + KeyCache cache.Cache[string, db.CachedKeyData] } // Method returns the HTTP method this route responds to diff --git a/go/apps/api/routes/v2_keys_delete_key/handler.go b/go/apps/api/routes/v2_keys_delete_key/handler.go index 4145db1a1d..98edf482b9 100644 --- a/go/apps/api/routes/v2_keys_delete_key/handler.go +++ b/go/apps/api/routes/v2_keys_delete_key/handler.go @@ -30,7 +30,7 @@ type Handler struct { DB db.Database Keys keys.KeyService Auditlogs auditlogs.AuditLogService - KeyCache cache.Cache[string, db.FindKeyForVerificationRow] + KeyCache cache.Cache[string, db.CachedKeyData] } // Method returns the HTTP method this route responds to diff --git a/go/apps/api/routes/v2_keys_remove_permissions/handler.go b/go/apps/api/routes/v2_keys_remove_permissions/handler.go index 4fa5b345e6..f66cba7d6c 100644 --- a/go/apps/api/routes/v2_keys_remove_permissions/handler.go +++ b/go/apps/api/routes/v2_keys_remove_permissions/handler.go @@ -28,7 +28,7 @@ type Handler struct { DB db.Database Keys keys.KeyService Auditlogs auditlogs.AuditLogService - KeyCache cache.Cache[string, db.FindKeyForVerificationRow] + KeyCache cache.Cache[string, db.CachedKeyData] } // Method returns the HTTP method this route responds to diff --git a/go/apps/api/routes/v2_keys_remove_roles/handler.go b/go/apps/api/routes/v2_keys_remove_roles/handler.go index e9efaa8d34..5d39c91b8e 100644 --- a/go/apps/api/routes/v2_keys_remove_roles/handler.go +++ b/go/apps/api/routes/v2_keys_remove_roles/handler.go @@ -29,7 +29,7 @@ type Handler struct { DB db.Database Keys keys.KeyService Auditlogs auditlogs.AuditLogService - KeyCache cache.Cache[string, db.FindKeyForVerificationRow] + KeyCache cache.Cache[string, db.CachedKeyData] } // Method returns the HTTP method this route responds to diff --git a/go/apps/api/routes/v2_keys_set_permissions/handler.go b/go/apps/api/routes/v2_keys_set_permissions/handler.go index ed0cb8c3e7..073f13af7d 100644 --- a/go/apps/api/routes/v2_keys_set_permissions/handler.go +++ b/go/apps/api/routes/v2_keys_set_permissions/handler.go @@ -31,7 +31,7 @@ type Handler struct { DB db.Database Keys keys.KeyService Auditlogs auditlogs.AuditLogService - KeyCache cache.Cache[string, db.FindKeyForVerificationRow] + KeyCache cache.Cache[string, db.CachedKeyData] } // Method returns the HTTP method this route responds to diff --git a/go/apps/api/routes/v2_keys_set_roles/handler.go b/go/apps/api/routes/v2_keys_set_roles/handler.go index b2094ebffa..95a0489a2d 100644 --- a/go/apps/api/routes/v2_keys_set_roles/handler.go +++ b/go/apps/api/routes/v2_keys_set_roles/handler.go @@ -30,7 +30,7 @@ type Handler struct { DB db.Database Keys keys.KeyService Auditlogs auditlogs.AuditLogService - KeyCache cache.Cache[string, db.FindKeyForVerificationRow] + KeyCache cache.Cache[string, db.CachedKeyData] } // Method returns the HTTP method this route responds to diff --git a/go/apps/api/routes/v2_keys_update_credits/handler.go b/go/apps/api/routes/v2_keys_update_credits/handler.go index a37b0b40d5..d3b96c173f 100644 --- a/go/apps/api/routes/v2_keys_update_credits/handler.go +++ b/go/apps/api/routes/v2_keys_update_credits/handler.go @@ -31,7 +31,7 @@ type Handler struct { DB db.Database Keys keys.KeyService Auditlogs auditlogs.AuditLogService - KeyCache cache.Cache[string, db.FindKeyForVerificationRow] + KeyCache cache.Cache[string, db.CachedKeyData] UsageLimiter usagelimiter.Service } diff --git a/go/apps/api/routes/v2_keys_update_key/handler.go b/go/apps/api/routes/v2_keys_update_key/handler.go index a7e512f863..0f3c6672dc 100644 --- a/go/apps/api/routes/v2_keys_update_key/handler.go +++ b/go/apps/api/routes/v2_keys_update_key/handler.go @@ -35,7 +35,7 @@ type Handler struct { DB db.Database Keys keys.KeyService Auditlogs auditlogs.AuditLogService - KeyCache cache.Cache[string, db.FindKeyForVerificationRow] + KeyCache cache.Cache[string, db.CachedKeyData] UsageLimiter usagelimiter.Service } diff --git a/go/apps/gw/services/caches/caches.go b/go/apps/gw/services/caches/caches.go index 4abed672f6..bc33b1b345 100644 --- a/go/apps/gw/services/caches/caches.go +++ b/go/apps/gw/services/caches/caches.go @@ -30,8 +30,8 @@ type Caches struct { // HostName -> Certificate TLSCertificate cache.Cache[string, tls.Certificate] - // KeyHash -> Key verification data (for keys service) - VerificationKeyByHash cache.Cache[string, db.FindKeyForVerificationRow] + // KeyHash -> Key verification data with pre-parsed IP whitelist (for keys service) + VerificationKeyByHash cache.Cache[string, db.CachedKeyData] } // Config defines the configuration options for initializing caches. @@ -108,7 +108,7 @@ func New(config Config) (Caches, error) { return Caches{}, fmt.Errorf("failed to create certificate cache: %w", err) } - verificationKeyByHash, err := cache.New(cache.Config[string, db.FindKeyForVerificationRow]{ + verificationKeyByHash, err := cache.New(cache.Config[string, db.CachedKeyData]{ Fresh: time.Minute, Stale: time.Minute * 10, Logger: config.Logger, diff --git a/go/internal/services/caches/caches.go b/go/internal/services/caches/caches.go index f38012b4e7..e54c646083 100644 --- a/go/internal/services/caches/caches.go +++ b/go/internal/services/caches/caches.go @@ -17,9 +17,9 @@ type Caches struct { // Keys are cache.ScopedKey and values are db.FindRatelimitNamespace. RatelimitNamespace cache.Cache[cache.ScopedKey, db.FindRatelimitNamespace] - // VerificationKeyByHash caches verification key lookups by their hash. - // Keys are string (hash) and values are db.VerificationKey. - VerificationKeyByHash cache.Cache[string, db.FindKeyForVerificationRow] + // VerificationKeyByHash caches verification key lookups by their hash with pre-parsed data. + // Keys are string (hash) and values are db.CachedKeyData (includes pre-parsed IP whitelist). + VerificationKeyByHash cache.Cache[string, db.CachedKeyData] // LiveApiByID caches live API lookups by ID. // Keys are string (ID) and values are db.FindLiveApiByIDRow. @@ -77,7 +77,7 @@ func New(config Config) (Caches, error) { return Caches{}, err } - verificationKeyByHash, err := cache.New(cache.Config[string, db.FindKeyForVerificationRow]{ + verificationKeyByHash, err := cache.New(cache.Config[string, db.CachedKeyData]{ Fresh: 10 * time.Second, Stale: 10 * time.Minute, Logger: config.Logger, diff --git a/go/internal/services/keys/get.go b/go/internal/services/keys/get.go index 4e668f7fd9..e78f5a6ebc 100644 --- a/go/internal/services/keys/get.go +++ b/go/internal/services/keys/get.go @@ -4,6 +4,7 @@ import ( "context" "encoding/json" "fmt" + "strings" "time" "github.com/unkeyed/unkey/go/internal/services/caches" @@ -67,11 +68,31 @@ func (s *service) Get(ctx context.Context, sess *zen.Session, rawKey string) (*K } h := hash.Sha256(rawKey) - key, hit, err := s.keyCache.SWR(ctx, h, func(ctx context.Context) (db.FindKeyForVerificationRow, error) { + key, hit, err := s.keyCache.SWR(ctx, h, func(ctx context.Context) (db.CachedKeyData, error) { // Use database retry with exponential backoff, skipping non-transient errors - return db.WithRetry(func() (db.FindKeyForVerificationRow, error) { + row, err := db.WithRetry(func() (db.FindKeyForVerificationRow, error) { return db.Query.FindKeyForVerification(ctx, s.db.RO(), h) }) + if err != nil { + return db.CachedKeyData{}, err + } + + // Parse IP whitelist once during cache population for performance + parsedIPWhitelist := make(map[string]struct{}) + if row.IpWhitelist.Valid && row.IpWhitelist.String != "" { + ips := strings.Split(row.IpWhitelist.String, ",") + for _, ip := range ips { + trimmed := strings.TrimSpace(ip) + if trimmed != "" { + parsedIPWhitelist[trimmed] = struct{}{} + } + } + } + + return db.CachedKeyData{ + FindKeyForVerificationRow: row, + ParsedIPWhitelist: parsedIPWhitelist, + }, nil }, caches.DefaultFindFirstOp) if err != nil { @@ -107,6 +128,7 @@ func (s *service) Get(ctx context.Context, sess *zen.Session, rawKey string) (*K }, emptyLog, nil } + // Workspace is disabled or the key is not allowed to be used for workspace operations if !key.WorkspaceEnabled || (key.ForWorkspaceEnabled.Valid && !key.ForWorkspaceEnabled.Bool) { // nolint:exhaustruct kv := &KeyVerifier{ @@ -121,7 +143,7 @@ func (s *service) Get(ctx context.Context, sess *zen.Session, rawKey string) (*K usageLimiter: s.usageLimiter, AuthorizedWorkspaceID: key.WorkspaceID, isRootKey: key.ForWorkspaceID.Valid, - Key: key, + Key: key.FindKeyForVerificationRow, } return kv, kv.log, nil @@ -180,7 +202,7 @@ func (s *service) Get(ctx context.Context, sess *zen.Session, rawKey string) (*K } kv := &KeyVerifier{ - Key: key, + Key: key.FindKeyForVerificationRow, clickhouse: s.clickhouse, rateLimiter: s.raterLimiter, usageLimiter: s.usageLimiter, @@ -193,11 +215,12 @@ func (s *service) Get(ctx context.Context, sess *zen.Session, rawKey string) (*K isRootKey: key.ForWorkspaceID.Valid, // By default we assume the key is valid unless proven otherwise - Status: StatusValid, - ratelimitConfigs: ratelimitConfigs, - Roles: roles, - Permissions: permissions, - RatelimitResults: nil, + Status: StatusValid, + ratelimitConfigs: ratelimitConfigs, + parsedIPWhitelist: key.ParsedIPWhitelist, // Use pre-parsed IPs from cache + Roles: roles, + Permissions: permissions, + RatelimitResults: nil, } if key.DeletedAtM.Valid { diff --git a/go/internal/services/keys/service.go b/go/internal/services/keys/service.go index 9afcd652b4..d743085df6 100644 --- a/go/internal/services/keys/service.go +++ b/go/internal/services/keys/service.go @@ -20,7 +20,7 @@ type Config struct { Region string // Geographic region identifier UsageLimiter usagelimiter.Service // Redis Counter for usage limiting - KeyCache cache.Cache[string, db.FindKeyForVerificationRow] // Cache for key lookups + KeyCache cache.Cache[string, db.CachedKeyData] // Cache for key lookups with pre-parsed data } type service struct { @@ -32,8 +32,8 @@ type service struct { clickhouse clickhouse.ClickHouse region string - // hash -> key - keyCache cache.Cache[string, db.FindKeyForVerificationRow] + // hash -> cached key data (includes pre-parsed IP whitelist) + keyCache cache.Cache[string, db.CachedKeyData] } // New creates a new keys service instance with the provided configuration. diff --git a/go/internal/services/keys/validation.go b/go/internal/services/keys/validation.go index 548aa93f80..85126140d8 100644 --- a/go/internal/services/keys/validation.go +++ b/go/internal/services/keys/validation.go @@ -4,11 +4,8 @@ import ( "context" "database/sql" "fmt" - "strings" "time" - "slices" - "github.com/unkeyed/unkey/go/apps/api/openapi" "github.com/unkeyed/unkey/go/internal/services/ratelimit" "github.com/unkeyed/unkey/go/internal/services/usagelimiter" @@ -58,7 +55,7 @@ func (k *KeyVerifier) withIPWhitelist() error { return nil } - if !k.Key.IpWhitelist.Valid { + if len(k.parsedIPWhitelist) == 0 { return nil } @@ -69,12 +66,7 @@ func (k *KeyVerifier) withIPWhitelist() error { return nil } - allowedIPs := strings.Split(k.Key.IpWhitelist.String, ",") - for i, ip := range allowedIPs { - allowedIPs[i] = strings.TrimSpace(ip) - } - - if !slices.Contains(allowedIPs, clientIP) { + if _, ok := k.parsedIPWhitelist[clientIP]; !ok { k.setInvalid(StatusForbidden, fmt.Sprintf("client IP %s is not in the whitelist", clientIP)) } diff --git a/go/internal/services/keys/verifier.go b/go/internal/services/keys/verifier.go index 0f001c6d75..242b09f40e 100644 --- a/go/internal/services/keys/verifier.go +++ b/go/internal/services/keys/verifier.go @@ -39,7 +39,8 @@ type KeyVerifier struct { 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) + parsedIPWhitelist map[string]struct{} // Pre-parsed IP whitelist for O(1) lookup + isRootKey bool // Whether this is a root key (special handling) message string // Internal message for validation failures tags []string // Tags associated with this verification diff --git a/go/pkg/db/cached_key_data.go b/go/pkg/db/cached_key_data.go new file mode 100644 index 0000000000..b1fa2c8389 --- /dev/null +++ b/go/pkg/db/cached_key_data.go @@ -0,0 +1,8 @@ +package db + +// CachedKeyData embeds FindKeyForVerificationRow and adds pre-processed data for caching. +// This struct is stored in the cache to avoid redundant parsing operations. +type CachedKeyData struct { + FindKeyForVerificationRow + ParsedIPWhitelist map[string]struct{} // Pre-parsed IP addresses for O(1) lookup +}