ProofedMesh/services/signing-service/internal/registry/registry.go

// Package registry resolves ProofedMesh agent identities against the
// trust-api `agents` table (db/trust/migrations/0002_agents.up.sql).
//
// The signing-service calls this package twice per request — once by
// pocket_id_client (from the Pocket-ID JWT `sub` claim) and once by
// (k8s_namespace, k8s_service_account) (from a Kubernetes TokenReview).
// The results are handed to OPA which denies unless both lookups land on
// the same, internally consistent agent row.
package registry

import (
	"context"
	"errors"
	"fmt"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgxpool"
)

// ErrNotFound means no agent row matched the given identity. Callers
// should treat this as a hard identity failure (403), never a 500.
var ErrNotFound = errors.New("registry: agent not found")

// Agent is the subset of the trust-api `agents` row that the signing
// service needs for identity reconciliation and policy evaluation.
type Agent struct {
	AgentID           string
	PocketIDClient    string
	K8sNamespace      string
	K8sServiceAccount string
	DisplayName       string
	AllowedActions    []string
}

// Resolver abstracts the backing store so tests can swap in an in-memory
// implementation.
type Resolver interface {
	ByPocketID(ctx context.Context, client string) (Agent, error)
	ByK8sIdentity(ctx context.Context, namespace, serviceAccount string) (Agent, error)
}

// Static is an in-memory resolver suitable for unit tests and the
// never-happens case where we want to pin the signing-service to a
// deterministic set of identities.
type Static struct {
	ByPocketIDMap    map[string]Agent
	ByK8sIdentityMap map[string]Agent // key: ns + "/" + sa
}

// NewStatic builds a Static resolver from a slice of agents.
func NewStatic(agents []Agent) *Static {
	byP := make(map[string]Agent, len(agents))
	byK := make(map[string]Agent, len(agents))
	for _, a := range agents {
		byP[a.PocketIDClient] = a
		byK[a.K8sNamespace+"/"+a.K8sServiceAccount] = a
	}
	return &Static{ByPocketIDMap: byP, ByK8sIdentityMap: byK}
}

func (s *Static) ByPocketID(_ context.Context, client string) (Agent, error) {
	a, ok := s.ByPocketIDMap[client]
	if !ok {
		return Agent{}, ErrNotFound
	}
	return a, nil
}

func (s *Static) ByK8sIdentity(_ context.Context, namespace, serviceAccount string) (Agent, error) {
	a, ok := s.ByK8sIdentityMap[namespace+"/"+serviceAccount]
	if !ok {
		return Agent{}, ErrNotFound
	}
	return a, nil
}

// Postgres is the production resolver, backed by the trust-api database.
type Postgres struct {
	pool *pgxpool.Pool
}

// NewPostgres opens a connection pool at the given DSN. Caller owns Close.
func NewPostgres(ctx context.Context, dsn string) (*Postgres, error) {
	pool, err := pgxpool.New(ctx, dsn)
	if err != nil {
		return nil, fmt.Errorf("registry: pgxpool.New: %w", err)
	}
	return &Postgres{pool: pool}, nil
}

// Close releases the pool.
func (p *Postgres) Close() {
	if p.pool != nil {
		p.pool.Close()
	}
}

const selectCols = `agent_id, pocket_id_client, k8s_namespace, k8s_service_account, display_name, allowed_actions`

func (p *Postgres) ByPocketID(ctx context.Context, client string) (Agent, error) {
	row := p.pool.QueryRow(ctx, `SELECT `+selectCols+` FROM agents WHERE pocket_id_client = $1`, client)
	return scanAgent(row)
}

func (p *Postgres) ByK8sIdentity(ctx context.Context, namespace, serviceAccount string) (Agent, error) {
	row := p.pool.QueryRow(ctx,
		`SELECT `+selectCols+` FROM agents WHERE k8s_namespace = $1 AND k8s_service_account = $2`,
		namespace, serviceAccount)
	return scanAgent(row)
}

type scannable interface {
	Scan(dest ...any) error
}

func scanAgent(row scannable) (Agent, error) {
	var a Agent
	if err := row.Scan(&a.AgentID, &a.PocketIDClient, &a.K8sNamespace, &a.K8sServiceAccount, &a.DisplayName, &a.AllowedActions); err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return Agent{}, ErrNotFound
		}
		return Agent{}, fmt.Errorf("registry: scan: %w", err)
	}
	return a, nil
}