llmgw/internal/store/keyttl.go

package store

import (
	"context"
	"database/sql"
	"errors"
	"fmt"
	"time"
)

// ScopeGateway is the whole gateway, for settings that apply everywhere.
const ScopeGateway Scope = "gateway"

// ErrKeyTTL is a key asked to live longer than a scope allows.
var ErrKeyTTL = errors.New("a key may not live that long here")

// SetMaxKeyTTL sets how long a key may live at most in a scope; zero takes
// the limit away.
func (s *DB) SetMaxKeyTTL(ctx context.Context, scope Scope, id string, max time.Duration) error {
	if max <= 0 {
		_, err := s.db.ExecContext(ctx, `DELETE FROM key_ttl WHERE scope = ? AND scope_id = ?`, string(scope), id)
		return err
	}
	_, err := s.db.ExecContext(ctx, `INSERT INTO key_ttl (scope, scope_id, max_seconds) VALUES (?, ?, ?)
		ON CONFLICT (scope, scope_id) DO UPDATE SET max_seconds = excluded.max_seconds`,
		string(scope), id, int64(max/time.Second))
	return err
}

// KeyTTLs lists every limit, "scope:id" to the longest a key may live.
func (s *DB) KeyTTLs(ctx context.Context) (map[string]time.Duration, error) {
	type row struct {
		k string
		d time.Duration
	}
	rows, err := queryAll(ctx, s.db, `SELECT scope, scope_id, max_seconds FROM key_ttl`, nil,
		func(r *sql.Rows) (row, error) {
			var sc, id string
			var n int64
			err := r.Scan(&sc, &id, &n)
			return row{sc + ":" + id, time.Duration(n) * time.Second}, err
		})
	out := map[string]time.Duration{}
	for _, r := range rows {
		out[r.k] = r.d
	}
	return out, err
}

// MaxKeyTTL is the longest a key in a project may live, and the scope that
// says so: the shortest of the gateway's, and the project's, its team's and
// its organization's. Zero when nothing limits it. A personal key (no
// project) has only the gateway's.
func (s *DB) MaxKeyTTL(ctx context.Context, projectID string) (time.Duration, string, error) {
	limits, err := s.KeyTTLs(ctx)
	if err != nil || len(limits) == 0 {
		return 0, "", err
	}
	scopes := []string{string(ScopeGateway) + ":"}
	if projectID != "" {
		parents, err := s.ScopeParents(ctx)
		if err != nil {
			return 0, "", err
		}
		scopes = append(scopes, "project:"+projectID)
		scopes = append(scopes, parents["project:"+projectID]...)
	}
	var best time.Duration
	by := ""
	for _, sc := range scopes {
		if d, ok := limits[sc]; ok && (best == 0 || d < best) {
			best, by = d, sc
		}
	}
	return best, by, nil
}

// keyTTL gives a key that asks for no expiry the longest its scope allows,
// and refuses one that asks to outlive it.
func (s *DB) keyTTL(ctx context.Context, k Key) (Key, error) {
	max, by, err := s.MaxKeyTTL(ctx, k.ProjectID)
	if err != nil || max == 0 {
		return k, err
	}
	created := k.CreatedAt
	if created.IsZero() {
		created = time.Now()
	}
	if k.ExpiresAt.IsZero() {
		k.ExpiresAt = created.Add(max)
		return k, nil
	}
	if k.ExpiresAt.Sub(created) > max+time.Minute {
		return k, fmt.Errorf("%w: %s allows %s at most", ErrKeyTTL, by, max)
	}
	return k, nil
}

// KeyByID finds a key by its id.
func (s *DB) KeyByID(ctx context.Context, id string) (Key, error) {
	row := s.db.QueryRowContext(ctx, `SELECT `+keyColumns+` FROM api_keys WHERE id = ?`, id)
	k, err := scanKey(row)
	if errors.Is(err, sql.ErrNoRows) {
		return Key{}, ErrNotFound
	}
	return k, err
}

// ExpireKeyBy has a key stop working at a moment, unless it stops sooner
// already: what rotating it does to the old one.
func (s *DB) ExpireKeyBy(ctx context.Context, id string, at time.Time) error {
	res, err := s.db.ExecContext(ctx, `UPDATE api_keys SET expires_at = ?
		WHERE id = ? AND revoked_at IS NULL AND (expires_at IS NULL OR expires_at > ?)`, millis(at), id, millis(at))
	if err != nil {
		return err
	}
	if n, _ := res.RowsAffected(); n == 0 {
		if _, err := s.KeyByID(ctx, id); err != nil {
			return err
		}
	}
	return nil
}