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

package credentialstore

import (
	"context"
	"crypto/rand"
	"crypto/sha256"
	"database/sql"
	"encoding/hex"
	"errors"
	"fmt"
	"strings"
	"time"

	"github.com/google/uuid"
)

// Status values for principal tokens. These are the only values accepted by
// the CHECK constraint on tokens.status.
const (
	StatusActive   = "active"
	StatusDisabled = "disabled"
	StatusRevoked  = "revoked"
)

// CreatePrincipalInput is the caller-provided payload for CreatePrincipalWithToken.
type CreatePrincipalInput struct {
	// Alias is a human-readable name for the principal. It must be unique
	// across all principals and is used for display and management.
	Alias string
}

// IssuedPrincipal is the one-time response from CreatePrincipalWithToken.
// RawToken is returned exactly once to the issuer and must be delivered to
// the corresponding Edge configuration before the reference is lost. The
// TokenRecord embedded in the response contains only the digest; the raw
// value is never persisted.
type IssuedPrincipal struct {
	Principal PrincipalRecord
	Token     TokenRecord
	RawToken  string
}

// IssuedToken is the one-time response when generating an additional token for an existing principal.
type IssuedToken struct {
	Token    TokenRecord
	RawToken string
}

// PrincipalRecord is the persisted principal row returned by queries.
type PrincipalRecord struct {
	ID        string
	Alias     string
	CreatedAt time.Time
	UpdatedAt time.Time
}

// TokenRecord is the persisted token row returned by queries. It never
// contains the raw token value; only the SHA-256 digest hex string.
type TokenRecord struct {
	ID          string
	PrincipalID string
	TokenRef    string
	Digest      string
	Status      string
	Revision    int64
	CreatedAt   time.Time
	UpdatedAt   time.Time
	RevokedAt   *time.Time
}

// ErrPrincipalAlreadyExists is returned when CreatePrincipalWithToken is
// called with an alias that is already in use.
var ErrPrincipalAlreadyExists = errors.New("credentialstore: principal alias already exists")

// ErrPrincipalNotFound is returned when a principal lookup finds no matching row.
var ErrPrincipalNotFound = errors.New("credentialstore: principal not found")

// ErrTokenNotFound is returned when a token lookup finds no matching row.
var ErrTokenNotFound = errors.New("credentialstore: token not found")

// ErrTokenNotActive is returned when an operation requires an active token.
var ErrTokenNotActive = errors.New("credentialstore: token is not active")

// ErrRevisionMismatch is returned by CAS operations when the provided
// revision does not match the current row revision.
var ErrRevisionMismatch = errors.New("credentialstore: revision mismatch")

// ErrTokenRevoked is returned by re-enable attempts on a revoked token.
var ErrTokenRevoked = errors.New("credentialstore: revoked token cannot be re-enabled")

// generateRawToken produces 32 random bytes encoded as a hex string.
// This is the only place the plaintext token exists; it is returned to the
// caller and never persisted.
func generateRawToken() (string, error) {
	b := make([]byte, 32)
	if _, err := rand.Read(b); err != nil {
		return "", fmt.Errorf("credentialstore: generate random: %w", err)
	}
	return hex.EncodeToString(b), nil
}

// digestOf returns the SHA-256 hex digest of raw.
func digestOf(raw string) string {
	sum := sha256.Sum256([]byte(raw))
	return hex.EncodeToString(sum[:])
}

// CreatePrincipalWithToken atomically creates a principal and an active
// token. The raw token is returned exactly once; the persisted TokenRecord
// contains only the SHA-256 digest.
//
// If the alias already exists, ErrPrincipalAlreadyExists is returned and no
// row is modified.
func (s *Store) CreatePrincipalWithToken(ctx context.Context, in CreatePrincipalInput) (*IssuedPrincipal, error) {
	if in.Alias == "" {
		return nil, fmt.Errorf("credentialstore: CreatePrincipalWithToken: alias is required")
	}
	now := time.Now().UTC()
	raw, err := generateRawToken()
	if err != nil {
		return nil, err
	}
	dig := digestOf(raw)
	principalID := uuid.New().String()
	tokenID := uuid.New().String()

	var issued IssuedPrincipal
	err = s.RunInTransaction(ctx, func(ctx context.Context, tx *sql.Tx) error {
		// Check for duplicate alias inside the transaction.
		var count int
		if err := tx.QueryRowContext(ctx, s.bind(`SELECT COUNT(*) FROM principals WHERE alias=?`), in.Alias).Scan(&count); err != nil {
			return fmt.Errorf("credentialstore: check alias: %w", err)
		}
		if count > 0 {
			return ErrPrincipalAlreadyExists
		}

		// Insert principal.
		if _, err := tx.ExecContext(ctx, s.bind(
			`INSERT INTO principals (id, alias, created_at, updated_at) VALUES (?,?,?,?)`),
			principalID, in.Alias, now, now,
		); err != nil {
			return fmt.Errorf("credentialstore: insert principal: %w", err)
		}

		// Insert token with digest only.
		if _, err := tx.ExecContext(ctx, s.bind(
			`INSERT INTO tokens (id, principal_id, token_ref, digest, status, revision, created_at, updated_at) VALUES (?,?,?,?,?,?,?,?)`),
			tokenID, principalID, tokenRefFor(dig), dig, StatusActive, 0, now, now,
		); err != nil {
			return fmt.Errorf("credentialstore: insert token: %w", err)
		}
		if err := s.bumpProjectionGenerationTx(ctx, tx); err != nil {
			return err
		}

		issued = IssuedPrincipal{
			Principal: PrincipalRecord{ID: principalID, Alias: in.Alias, CreatedAt: now, UpdatedAt: now},
			Token: TokenRecord{
				ID: tokenID, PrincipalID: principalID, TokenRef: tokenRefFor(dig),
				Digest: dig, Status: StatusActive, Revision: 0, CreatedAt: now, UpdatedAt: now,
			},
			RawToken: raw,
		}
		return nil
	})
	if err != nil {
		return nil, err
	}
	return &issued, nil
}

// CreateFirstPrincipalWithToken atomically creates the first principal and its
// active token. It locks the projection state singleton row to serialize
// concurrent callers, verifies the principal store is currently empty, and
// rejects any non-empty store with ErrPrincipalAlreadyExists.
func (s *Store) CreateFirstPrincipalWithToken(ctx context.Context, in CreatePrincipalInput) (*IssuedPrincipal, error) {
	if in.Alias == "" {
		return nil, fmt.Errorf("credentialstore: CreateFirstPrincipalWithToken: alias is required")
	}
	now := time.Now().UTC()
	raw, err := generateRawToken()
	if err != nil {
		return nil, err
	}
	dig := digestOf(raw)
	principalID := uuid.New().String()
	tokenID := uuid.New().String()

	var issued IssuedPrincipal
	err = s.RunInTransaction(ctx, func(ctx context.Context, tx *sql.Tx) error {
		res, err := tx.ExecContext(ctx, s.bind(`UPDATE principal_projection_state SET generation=generation WHERE singleton=1`))
		if err != nil {
			return fmt.Errorf("credentialstore: acquire first principal lock: %w", err)
		}
		affected, err := res.RowsAffected()
		if err != nil {
			return fmt.Errorf("credentialstore: inspect first principal lock: %w", err)
		}
		if affected != 1 {
			return ErrProjectionGenerationMissing
		}

		var count int
		if err := tx.QueryRowContext(ctx, s.bind(`SELECT COUNT(*) FROM principals`)).Scan(&count); err != nil {
			return fmt.Errorf("credentialstore: count principals: %w", err)
		}
		if count > 0 {
			return ErrPrincipalAlreadyExists
		}

		if _, err := tx.ExecContext(ctx, s.bind(
			`INSERT INTO principals (id, alias, created_at, updated_at) VALUES (?,?,?,?)`),
			principalID, in.Alias, now, now,
		); err != nil {
			return fmt.Errorf("credentialstore: insert principal: %w", err)
		}

		if _, err := tx.ExecContext(ctx, s.bind(
			`INSERT INTO tokens (id, principal_id, token_ref, digest, status, revision, created_at, updated_at) VALUES (?,?,?,?,?,?,?,?)`),
			tokenID, principalID, tokenRefFor(dig), dig, StatusActive, 0, now, now,
		); err != nil {
			return fmt.Errorf("credentialstore: insert token: %w", err)
		}
		if err := s.bumpProjectionGenerationTx(ctx, tx); err != nil {
			return err
		}

		issued = IssuedPrincipal{
			Principal: PrincipalRecord{ID: principalID, Alias: in.Alias, CreatedAt: now, UpdatedAt: now},
			Token: TokenRecord{
				ID: tokenID, PrincipalID: principalID, TokenRef: tokenRefFor(dig),
				Digest: dig, Status: StatusActive, Revision: 0, CreatedAt: now, UpdatedAt: now,
			},
			RawToken: raw,
		}
		return nil
	})
	if err != nil {
		return nil, err
	}
	return &issued, nil
}

// CreateToken atomically generates and inserts an active token for an existing principal.
// The raw token is returned exactly once in IssuedToken; the persisted TokenRecord contains only the SHA-256 digest.
func (s *Store) CreateToken(ctx context.Context, principalID string) (*IssuedToken, error) {
	if principalID == "" {
		return nil, fmt.Errorf("credentialstore: CreateToken: principalID is required")
	}
	now := time.Now().UTC()
	raw, err := generateRawToken()
	if err != nil {
		return nil, err
	}
	dig := digestOf(raw)
	tokenID := uuid.New().String()

	var issued IssuedToken
	err = s.RunInTransaction(ctx, func(ctx context.Context, tx *sql.Tx) error {
		var count int
		if err := tx.QueryRowContext(ctx, s.bind(`SELECT COUNT(*) FROM principals WHERE id=?`), principalID).Scan(&count); err != nil {
			return fmt.Errorf("credentialstore: check principal: %w", err)
		}
		if count == 0 {
			return ErrPrincipalNotFound
		}

		if _, err := tx.ExecContext(ctx, s.bind(
			`INSERT INTO tokens (id, principal_id, token_ref, digest, status, revision, created_at, updated_at) VALUES (?,?,?,?,?,?,?,?)`),
			tokenID, principalID, tokenRefFor(dig), dig, StatusActive, 0, now, now,
		); err != nil {
			return fmt.Errorf("credentialstore: insert token: %w", err)
		}
		if err := s.bumpProjectionGenerationTx(ctx, tx); err != nil {
			return err
		}

		issued = IssuedToken{
			Token: TokenRecord{
				ID:          tokenID,
				PrincipalID: principalID,
				TokenRef:    tokenRefFor(dig),
				Digest:      dig,
				Status:      StatusActive,
				Revision:    0,
				CreatedAt:   now,
				UpdatedAt:   now,
			},
			RawToken: raw,
		}
		return nil
	})
	if err != nil {
		return nil, err
	}
	return &issued, nil
}

// ListTokens returns all tokens for a given principal.
func (s *Store) ListTokens(ctx context.Context, principalID string) ([]TokenRecord, error) {
	rows, err := s.db.QueryContext(ctx, s.bind(`
		SELECT id, principal_id, token_ref, digest, status, revision, created_at, updated_at, revoked_at
		FROM tokens
		WHERE principal_id=?
		ORDER BY created_at ASC, id ASC
	`), principalID)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var out []TokenRecord
	for rows.Next() {
		var t TokenRecord
		var createdAt, updatedAt string
		var revokedAt sql.NullString
		if err := rows.Scan(&t.ID, &t.PrincipalID, &t.TokenRef, &t.Digest, &t.Status, &t.Revision, &createdAt, &updatedAt, &revokedAt); err != nil {
			return nil, err
		}
		t.CreatedAt, _ = parseTime(createdAt)
		t.UpdatedAt, _ = parseTime(updatedAt)
		if revokedAt.Valid {
			parsed, _ := parseTime(revokedAt.String)
			t.RevokedAt = &parsed
		}
		out = append(out, t)
	}
	return out, rows.Err()
}

// tokenRefFor derives a stable, opaque token reference from a digest. It is a
// safe identifier for metadata and projections; callers authenticate with the
// one-time raw token, never with this reference or the digest itself.
func tokenRefFor(digest string) string {
	return "tok_" + digest[:16]
}

// GetPrincipal returns the principal with the given alias, or (nil, nil) if
// no such principal exists.
func (s *Store) GetPrincipal(ctx context.Context, alias string) (*PrincipalRecord, error) {
	var r PrincipalRecord
	var createdAt, updatedAt string
	err := s.db.QueryRowContext(ctx, s.bind(
		`SELECT id, alias, created_at, updated_at FROM principals WHERE alias=?`), alias,
	).Scan(&r.ID, &r.Alias, &createdAt, &updatedAt)
	if err != nil {
		if err == sql.ErrNoRows {
			return nil, nil
		}
		return nil, err
	}
	r.CreatedAt, _ = time.Parse(time.RFC3339, createdAt)
	r.UpdatedAt, _ = time.Parse(time.RFC3339, updatedAt)
	return &r, nil
}

// ListPrincipals returns all principals with their active token metadata.
func (s *Store) ListPrincipals(ctx context.Context) ([]PrincipalWithToken, error) {
	rows, err := s.db.QueryContext(ctx, `
		SELECT p.id, p.alias, p.created_at, p.updated_at,
		       t.id AS token_id, t.token_ref, t.digest, t.status, t.revision, t.created_at AS token_created_at, t.updated_at AS token_updated_at, t.revoked_at
		FROM principals p
		LEFT JOIN tokens t ON t.principal_id = p.id AND t.status = 'active'
		ORDER BY p.alias
	`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	var out []PrincipalWithToken
	for rows.Next() {
		var p PrincipalRecord
		var t TokenRecord
		var principalCreatedAt, principalUpdatedAt string
		var tokenID, tokenRef, tokenDigest, tokenStatus sql.NullString
		var tokenRevision sql.NullInt64
		var tokenCreatedAt, tokenUpdatedAt, tokenRevokedAt sql.NullString
		if err := rows.Scan(&p.ID, &p.Alias, &principalCreatedAt, &principalUpdatedAt,
			&tokenID, &tokenRef, &tokenDigest, &tokenStatus, &tokenRevision,
			&tokenCreatedAt, &tokenUpdatedAt, &tokenRevokedAt); err != nil {
			return nil, err
		}
		p.CreatedAt, _ = parseTime(principalCreatedAt)
		p.UpdatedAt, _ = parseTime(principalUpdatedAt)
		if tokenID.Valid {
			t.ID = tokenID.String
			t.PrincipalID = p.ID
			t.TokenRef = tokenRef.String
			t.Digest = tokenDigest.String
			t.Status = tokenStatus.String
			t.Revision = tokenRevision.Int64
			t.CreatedAt, _ = time.Parse(time.RFC3339, tokenCreatedAt.String)
			t.UpdatedAt, _ = time.Parse(time.RFC3339, tokenUpdatedAt.String)
			if tokenRevokedAt.Valid {
				timestamp := tokenRevokedAt.String
				parsed, err := time.Parse(time.RFC3339, timestamp)
				if err != nil {
					parsed, _ = time.Parse("2006-01-02T15:04:05.999999999Z", timestamp)
				}
				t.RevokedAt = &parsed
			}
		}
		out = append(out, PrincipalWithToken{Principal: p, Token: t})
	}
	return out, rows.Err()
}

// PrincipalWithToken pairs a principal record with its active token metadata.
type PrincipalWithToken struct {
	Principal PrincipalRecord
	Token     TokenRecord
}

// LookupTokenByDigest verifies a token digest against the store and returns
// the associated principal and token metadata. It is the authentication
// lookup path used by Edge projection validators.
//
// Only active tokens are returned; disabled and revoked tokens are hidden and
// yield ErrTokenNotFound just like an unknown digest.
func (s *Store) LookupTokenByDigest(ctx context.Context, digest string) (*PrincipalRecord, *TokenRecord, error) {
	var p PrincipalRecord
	var t TokenRecord
	var principalCreatedAt, principalUpdatedAt, tokenCreatedAt, tokenUpdatedAt string
	var revokedAt sql.NullString
	err := s.db.QueryRowContext(ctx, s.bind(`
		SELECT p.id, p.alias, p.created_at, p.updated_at,
		       t.id, t.token_ref, t.digest, t.status, t.revision, t.created_at, t.updated_at, t.revoked_at
		FROM principals p
		JOIN tokens t ON t.principal_id = p.id
		WHERE t.digest = ? AND t.status = ?
	`), digest, StatusActive).Scan(
		&p.ID, &p.Alias, &principalCreatedAt, &principalUpdatedAt,
		&t.ID, &t.TokenRef, &t.Digest, &t.Status, &t.Revision, &tokenCreatedAt, &tokenUpdatedAt, &revokedAt,
	)
	if err != nil {
		if err == sql.ErrNoRows {
			return nil, nil, ErrTokenNotFound
		}
		return nil, nil, err
	}
	p.CreatedAt, _ = parseTime(principalCreatedAt)
	p.UpdatedAt, _ = parseTime(principalUpdatedAt)
	t.CreatedAt, _ = parseTime(tokenCreatedAt)
	t.UpdatedAt, _ = parseTime(tokenUpdatedAt)
	if revokedAt.Valid {
		parsed, _ := parseTime(revokedAt.String)
		t.RevokedAt = &parsed
	}
	return &p, &t, nil
}

// DisableToken transitions an active token to disabled using CAS on revision.
// Disabled tokens cannot authenticate. Revoked tokens are terminal.
func (s *Store) DisableToken(ctx context.Context, principalID, tokenRef string, currentRevision int64) (*TokenRecord, error) {
	return s.casStatus(ctx, principalID, tokenRef, currentRevision, []string{StatusActive}, StatusDisabled)
}

// RevokeToken transitions an active or disabled token to revoked. Revocation
// is irreversible: no method can return a revoked token to any other state.
func (s *Store) RevokeToken(ctx context.Context, principalID, tokenRef string, currentRevision int64) (*TokenRecord, error) {
	return s.casStatus(ctx, principalID, tokenRef, currentRevision, []string{StatusActive, StatusDisabled}, StatusRevoked)
}

// casStatus performs a compare-and-swap on the token status. It requires
// currentRevision to match the row's current revision and transitions from
// fromStatus to toStatus. The revision is incremented atomically.
func (s *Store) casStatus(ctx context.Context, principalID, tokenRef string, currentRevision int64, fromStatuses []string, toStatus string) (*TokenRecord, error) {
	var updated TokenRecord
	err := s.RunInTransaction(ctx, func(ctx context.Context, tx *sql.Tx) error {
		now := time.Now().UTC()
		placeholders := strings.TrimRight(strings.Repeat("?,", len(fromStatuses)), ",")
		query := `UPDATE tokens SET status=?, revision=revision+1, updated_at=?, revoked_at=?
			WHERE principal_id=? AND token_ref=? AND revision=? AND status IN (` + placeholders + `)`
		args := []any{toStatus, now, nil, principalID, tokenRef, currentRevision}
		if toStatus == StatusRevoked {
			args[2] = now
		}
		for _, status := range fromStatuses {
			args = append(args, status)
		}
		result, err := tx.ExecContext(ctx, s.bind(query), args...)
		if err != nil {
			return err
		}
		affected, err := result.RowsAffected()
		if err != nil {
			return err
		}
		if affected != 1 {
			return s.classifyTransitionMiss(ctx, tx, principalID, tokenRef, currentRevision, fromStatuses)
		}
		if err := s.loadToken(ctx, tx, principalID, tokenRef, &updated); err != nil {
			return err
		}
		return s.bumpProjectionGenerationTx(ctx, tx)
	})
	if err != nil {
		return nil, err
	}
	return &updated, nil
}

func (s *Store) classifyTransitionMiss(ctx context.Context, tx *sql.Tx, principalID, tokenRef string, currentRevision int64, fromStatuses []string) error {
	var status string
	var revision int64
	err := tx.QueryRowContext(ctx, s.bind(
		`SELECT status, revision FROM tokens WHERE principal_id=? AND token_ref=?`), principalID, tokenRef,
	).Scan(&status, &revision)
	if err == sql.ErrNoRows {
		return ErrTokenNotFound
	}
	if err != nil {
		return err
	}
	if revision != currentRevision {
		return fmt.Errorf("%w: have %d want %d", ErrRevisionMismatch, revision, currentRevision)
	}
	if status == StatusRevoked {
		return ErrTokenRevoked
	}
	return fmt.Errorf("%w: have %s want one of %s", ErrTokenNotActive, status, strings.Join(fromStatuses, ","))
}

func (s *Store) loadToken(ctx context.Context, tx *sql.Tx, principalID, tokenRef string, target *TokenRecord) error {
	var createdAt, updatedAt string
	var revokedAt sql.NullString
	err := tx.QueryRowContext(ctx, s.bind(
		`SELECT id, principal_id, token_ref, digest, status, revision, created_at, updated_at, revoked_at FROM tokens WHERE principal_id=? AND token_ref=?`),
		principalID, tokenRef,
	).Scan(&target.ID, &target.PrincipalID, &target.TokenRef, &target.Digest, &target.Status, &target.Revision, &createdAt, &updatedAt, &revokedAt)
	if err != nil {
		return err
	}
	target.CreatedAt, _ = parseTime(createdAt)
	target.UpdatedAt, _ = parseTime(updatedAt)
	if revokedAt.Valid {
		parsed, _ := parseTime(revokedAt.String)
		target.RevokedAt = &parsed
	}
	return nil
}

// parseTime attempts RFC3339 first, then falls back to microsecond-precision
// SQLite datetime format.
func parseTime(s string) (time.Time, error) {
	t, err := time.Parse(time.RFC3339, s)
	if err != nil {
		t, err = time.Parse("2006-01-02T15:04:05.999999999Z", s)
	}
	return t, err
}