ai-gateway-go/internal/apikey/authenticator.go

package apikey

import (
	"context"
	"crypto/subtle"
	"encoding/json"
	"errors"
	"fmt"
	"log/slog"
	"strings"
	"sync/atomic"
	"time"

	"github.com/redis/go-redis/v9"
)

type Authenticator struct {
	repository    *Repository
	redis         *redis.Client
	bootstrap     string
	cacheTTL      time.Duration
	bootstrapUses atomic.Uint64
	logger        *slog.Logger
}

func NewAuthenticator(repository *Repository, client *redis.Client, bootstrap string) *Authenticator {
	return &Authenticator{repository: repository, redis: client, bootstrap: bootstrap, cacheTTL: 30 * time.Second}
}

// SetLogger wires an optional logger used for best-effort cache diagnostics.
func (a *Authenticator) SetLogger(logger *slog.Logger) { a.logger = logger }

func (a *Authenticator) Authenticate(ctx context.Context, secret string) error {
	_, err := a.AuthenticatePrincipal(ctx, secret)
	return err
}

// invokeScope reports whether the key carries a scope that may reach the
// gateway. Besides regular gateway keys it admits the server-side
// application runtime credential ("application:run"), which is created only
// by the portal runtime (encrypted in PostgreSQL, tenant-bound, never shown
// to browsers) so that hosted application conversations can call the gateway
// on behalf of the app owner. Both call paths in the workbench runtime and
// the main proxy use this single method, so the runtime credential works
// end-to-end without widening any other surface.
func invokeScope(scopes []string) bool {
	return HasScope(scopes, "gateway:invoke") || HasScope(scopes, "application:run")
}

func (a *Authenticator) AuthenticatePrincipal(ctx context.Context, secret string) (Principal, error) {
	secret = strings.TrimSpace(secret)
	if secret == "" {
		return Principal{}, ErrInvalid
	}
	if a.bootstrap != "" && len(secret) == len(a.bootstrap) && subtle.ConstantTimeCompare([]byte(secret), []byte(a.bootstrap)) == 1 {
		a.bootstrapUses.Add(1)
		// The bootstrap credential is a real gateway key, not an anonymous
		// pass. Return a stable, well-known identity so the admission and
		// token-quota paths run (they allow it outright: RPM/quota are 0 by
		// design for a migration key) and audit records attribute usage to
		// "bootstrap" instead of an empty principal that silently skips every
		// policy stage. An empty APIKeyID previously bypassed rate limiting,
		// quota and audit attribution entirely.
		return Principal{APIKeyID: "bootstrap", Scopes: []string{"gateway:invoke"}}, nil
	}
	hash, hexHash := Digest(secret)
	if a.redis != nil {
		payload, err := a.redis.Get(ctx, cacheKey(hexHash)).Bytes()
		if err == nil {
			var principal Principal
			if json.Unmarshal(payload, &principal) == nil && invokeScope(principal.Scopes) {
				return principal, nil
			}
			return Principal{}, ErrInvalid
		}
		if err != nil && !errors.Is(err, redis.Nil) {
			return Principal{}, fmt.Errorf("%w: %v", ErrStore, err)
		}
	}
	principal, err := a.repository.Validate(ctx, hash)
	if err != nil {
		return Principal{}, err
	}
	if !invokeScope(principal.Scopes) {
		return Principal{}, ErrInvalid
	}
	// Best-effort cache: the database is the source of truth for validation.
	// A cache write failure must not turn a key the database just accepted
	// into a 503, or a transient Redis blip would take the whole gateway down.
	if a.redis != nil {
		payload, _ := json.Marshal(principal)
		if err := a.redis.Set(ctx, cacheKey(hexHash), payload, a.cacheTTL).Err(); err != nil && a.logger != nil {
			a.logger.Warn("api key cache write failed; continuing with database result", "error", err)
		}
	}
	return principal, nil
}

func (a *Authenticator) BootstrapUses() uint64 {
	if a == nil {
		return 0
	}
	return a.bootstrapUses.Load()
}

func (a *Authenticator) Invalidate(ctx context.Context, hash []byte) error {
	if a.redis == nil {
		return nil
	}
	_, hexHash := DigestFromHash(hash)
	return a.redis.Del(ctx, cacheKey(hexHash)).Err()
}

func DigestFromHash(hash []byte) ([]byte, string) {
	const hex = "0123456789abcdef"
	encoded := make([]byte, len(hash)*2)
	for i, value := range hash {
		encoded[i*2] = hex[value>>4]
		encoded[i*2+1] = hex[value&15]
	}
	return hash, string(encoded)
}

func cacheKey(hexHash string) string { return "gateway:api-key:v1:" + hexHash }