iop/apps/control-plane/internal/credentialstore/projection.go

package credentialstore

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

	"google.golang.org/protobuf/proto"

	iop "iop/proto/gen/iop"
)

const defaultProjectionTTL = 5 * time.Minute

var (
	ErrProjectionGenerationMissing = errors.New("credentialstore: projection generation is missing")
	ErrProjectionTooLarge          = errors.New("credentialstore: projection exceeds configured bounds")
)

// ProjectionLimits bounds the amount of verifier and route metadata loaded
// into one immutable projection. MaxBytes is measured using protobuf wire
// size after the snapshot is assembled.
type ProjectionLimits struct {
	MaxTokens int
	MaxRoutes int
	MaxBytes  int
}

// DefaultProjectionLimits returns conservative hard bounds for one Edge
// authorization projection.
func DefaultProjectionLimits() ProjectionLimits {
	return ProjectionLimits{
		MaxTokens: 4096,
		MaxRoutes: 16384,
		MaxBytes:  4 << 20,
	}
}

// ProjectionBuildOptions controls snapshot expiry, bounds, and time. Clock is
// injectable so expiry fixtures do not depend on wall-clock sleeps.
type ProjectionBuildOptions struct {
	TTL    time.Duration
	Limits ProjectionLimits
	Clock  func() time.Time
}

func normalizeProjectionBuildOptions(opts ProjectionBuildOptions) (ProjectionBuildOptions, error) {
	if opts.TTL < 0 {
		return ProjectionBuildOptions{}, fmt.Errorf("credentialstore: projection ttl must be positive")
	}
	if opts.TTL == 0 {
		opts.TTL = defaultProjectionTTL
	}
	if opts.Clock == nil {
		opts.Clock = time.Now
	}
	defaults := DefaultProjectionLimits()
	if opts.Limits.MaxTokens < 0 || opts.Limits.MaxRoutes < 0 || opts.Limits.MaxBytes < 0 {
		return ProjectionBuildOptions{}, fmt.Errorf("credentialstore: projection limits must not be negative")
	}
	if opts.Limits.MaxTokens == 0 {
		opts.Limits.MaxTokens = defaults.MaxTokens
	}
	if opts.Limits.MaxRoutes == 0 {
		opts.Limits.MaxRoutes = defaults.MaxRoutes
	}
	if opts.Limits.MaxBytes == 0 {
		opts.Limits.MaxBytes = defaults.MaxBytes
	}
	return opts, nil
}

// ProjectionGeneration returns the durable generation currently committed by
// token, slot, and route mutations.
func (s *Store) ProjectionGeneration(ctx context.Context) (uint64, error) {
	if s == nil || s.db == nil {
		return 0, ErrProjectionGenerationMissing
	}
	var generation int64
	if err := s.db.QueryRowContext(ctx, `SELECT generation FROM principal_projection_state WHERE singleton=1`).Scan(&generation); err != nil {
		if errors.Is(err, sql.ErrNoRows) {
			return 0, ErrProjectionGenerationMissing
		}
		return 0, err
	}
	if generation < 0 {
		return 0, fmt.Errorf("credentialstore: invalid projection generation %d", generation)
	}
	return uint64(generation), nil
}

// bumpProjectionGenerationTx increments the singleton generation inside the
// caller's mutation transaction. A failure rolls the entire mutation back.
func (s *Store) bumpProjectionGenerationTx(ctx context.Context, tx *sql.Tx) error {
	result, err := tx.ExecContext(ctx, `UPDATE principal_projection_state SET generation=generation+1 WHERE singleton=1`)
	if err != nil {
		return fmt.Errorf("credentialstore: bump projection generation: %w", err)
	}
	affected, err := result.RowsAffected()
	if err != nil {
		return fmt.Errorf("credentialstore: inspect projection generation bump: %w", err)
	}
	if affected != 1 {
		return ErrProjectionGenerationMissing
	}
	return nil
}

// BuildPrincipalProjection reads one transactionally consistent, secret-free
// snapshot. Only active token digests and routes whose credential slot is also
// active are included. Ciphertext, envelope metadata, and raw credentials are
// deliberately absent from the protobuf contract.
func (s *Store) BuildPrincipalProjection(ctx context.Context, opts ProjectionBuildOptions) (*iop.PrincipalProjection, error) {
	if s == nil || s.db == nil {
		return nil, ErrProjectionGenerationMissing
	}
	opts, err := normalizeProjectionBuildOptions(opts)
	if err != nil {
		return nil, err
	}

	tx, err := s.db.BeginTx(ctx, &sql.TxOptions{Isolation: sql.LevelSerializable})
	if err != nil {
		return nil, fmt.Errorf("credentialstore: begin projection snapshot: %w", err)
	}
	defer func() { _ = tx.Rollback() }()

	var generation int64
	if err := tx.QueryRowContext(ctx, `SELECT generation FROM principal_projection_state WHERE singleton=1`).Scan(&generation); err != nil {
		if errors.Is(err, sql.ErrNoRows) {
			return nil, ErrProjectionGenerationMissing
		}
		return nil, fmt.Errorf("credentialstore: read projection generation: %w", err)
	}
	if generation < 0 {
		return nil, fmt.Errorf("credentialstore: invalid projection generation %d", generation)
	}

	tokens, err := s.projectedTokensTx(ctx, tx, opts.Limits.MaxTokens)
	if err != nil {
		return nil, err
	}
	routes, err := s.projectedRoutesTx(ctx, tx, opts.Limits.MaxRoutes)
	if err != nil {
		return nil, err
	}

	issuedAt := opts.Clock().UTC()
	snapshot := &iop.PrincipalProjection{
		Generation:        uint64(generation),
		IssuedAtUnixNano:  issuedAt.UnixNano(),
		ExpiresAtUnixNano: issuedAt.Add(opts.TTL).UnixNano(),
		Tokens:            tokens,
		Routes:            routes,
	}
	if size := proto.Size(snapshot); size > opts.Limits.MaxBytes {
		return nil, fmt.Errorf("%w: bytes=%d max=%d", ErrProjectionTooLarge, size, opts.Limits.MaxBytes)
	}
	if err := tx.Commit(); err != nil {
		return nil, fmt.Errorf("credentialstore: commit projection snapshot: %w", err)
	}
	return snapshot, nil
}

func (s *Store) projectedTokensTx(ctx context.Context, tx *sql.Tx, max int) ([]*iop.ProjectedPrincipalToken, error) {
	rows, err := tx.QueryContext(ctx, s.bind(`
		SELECT t.digest, p.id, p.alias, t.token_ref, t.revision
		FROM tokens t
		JOIN principals p ON p.id=t.principal_id
		WHERE t.status=?
		ORDER BY p.id, t.token_ref
		LIMIT ?
	`), StatusActive, max+1)
	if err != nil {
		return nil, fmt.Errorf("credentialstore: query projected tokens: %w", err)
	}
	defer rows.Close()

	tokens := make([]*iop.ProjectedPrincipalToken, 0)
	for rows.Next() {
		var token iop.ProjectedPrincipalToken
		var revision int64
		if err := rows.Scan(&token.TokenDigestSha256, &token.PrincipalRef, &token.PrincipalAlias, &token.TokenRef, &revision); err != nil {
			return nil, fmt.Errorf("credentialstore: scan projected token: %w", err)
		}
		if len(tokens) == max {
			return nil, fmt.Errorf("%w: tokens exceed %d", ErrProjectionTooLarge, max)
		}
		if revision < 0 {
			return nil, fmt.Errorf("credentialstore: token %s has negative revision", token.TokenRef)
		}
		token.TokenRevision = uint64(revision)
		tokens = append(tokens, &token)
	}
	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("credentialstore: iterate projected tokens: %w", err)
	}
	return tokens, nil
}

func (s *Store) projectedRoutesTx(ctx context.Context, tx *sql.Tx, max int) ([]*iop.ProjectedPrincipalRoute, error) {
	rows, err := tx.QueryContext(ctx, s.bind(`
		SELECT r.id, r.alias, r.principal_id, r.slot_id, r.profile_id,
		       r.upstream_model, r.resource_selector, r.revision, s.revision
		FROM routes r
		JOIN credential_slots s ON s.id=r.slot_id AND s.principal_id=r.principal_id
		WHERE r.status=? AND s.status=?
		ORDER BY r.principal_id, r.id
		LIMIT ?
	`), StatusActive, StatusActive, max+1)
	if err != nil {
		return nil, fmt.Errorf("credentialstore: query projected routes: %w", err)
	}
	defer rows.Close()

	routes := make([]*iop.ProjectedPrincipalRoute, 0)
	for rows.Next() {
		var route iop.ProjectedPrincipalRoute
		var alias sql.NullString
		var routeRevision, credentialRevision int64
		if err := rows.Scan(
			&route.RouteId, &alias, &route.PrincipalRef, &route.CredentialSlotRef,
			&route.ProfileId, &route.UpstreamModel, &route.ResourceSelector,
			&routeRevision, &credentialRevision,
		); err != nil {
			return nil, fmt.Errorf("credentialstore: scan projected route: %w", err)
		}
		if len(routes) == max {
			return nil, fmt.Errorf("%w: routes exceed %d", ErrProjectionTooLarge, max)
		}
		if routeRevision < 0 || credentialRevision < 0 {
			return nil, fmt.Errorf("credentialstore: route %s has a negative revision", route.RouteId)
		}
		if alias.Valid {
			route.RouteAlias = alias.String
		}
		route.RouteRevision = uint64(routeRevision)
		route.CredentialRevision = uint64(credentialRevision)
		routes = append(routes, &route)
	}
	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("credentialstore: iterate projected routes: %w", err)
	}
	return routes, nil
}