easyai-ai-gateway/apps/api/internal/accessruleaudit/audit.go

package accessruleaudit

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"sort"
	"strings"
	"time"

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

const (
	SchemaVersion  = "access-rule-audit/v1"
	MigrationBatch = "0102_access_rule_allow_whitelist_semantics"
)

type Snapshot struct {
	SchemaVersion string           `json:"schemaVersion"`
	GeneratedAt   time.Time        `json:"generatedAt"`
	SecretSafe    bool             `json:"secretSafe"`
	Live          RuleSet          `json:"live"`
	Archive       *ArchiveManifest `json:"archive,omitempty"`
}

type RuleSet struct {
	Total       int64       `json:"total"`
	SHA256      string      `json:"sha256"`
	AllowCount  int64       `json:"allowCount"`
	AllowSHA256 string      `json:"allowSha256"`
	DenyCount   int64       `json:"denyCount"`
	DenySHA256  string      `json:"denySha256"`
	Counts      []RuleCount `json:"counts"`
}

type RuleCount struct {
	SubjectType  string `json:"subjectType"`
	Effect       string `json:"effect"`
	ResourceType string `json:"resourceType"`
	Status       string `json:"status"`
	Count        int64  `json:"count"`
}

type ArchiveManifest struct {
	MigrationBatch     string `json:"migrationBatch"`
	ManifestAllowCount int64  `json:"manifestAllowCount"`
	ManifestAllowSHA   string `json:"manifestAllowSha256"`
	ArchivedAllowCount int64  `json:"archivedAllowCount"`
	ArchivedAllowSHA   string `json:"archivedAllowSha256"`
	ManifestDenyCount  int64  `json:"manifestDenyCount"`
	ManifestDenySHA    string `json:"manifestDenySha256"`
	Consistent         bool   `json:"consistent"`
}

type ruleDigestRow struct {
	id                 string
	subjectType        string
	subjectID          string
	resourceType       string
	resourceID         string
	effect             string
	priority           string
	minPermissionLevel string
	conditions         string
	metadata           string
	status             string
	createdEpoch       string
	updatedEpoch       string
}

func Export(ctx context.Context, pool *pgxpool.Pool, now time.Time) (Snapshot, error) {
	live, err := loadLiveRules(ctx, pool)
	if err != nil {
		return Snapshot{}, err
	}
	archive, err := loadArchiveManifest(ctx, pool)
	if err != nil {
		return Snapshot{}, err
	}
	if now.IsZero() {
		now = time.Now()
	}
	snapshot := Snapshot{
		SchemaVersion: SchemaVersion,
		GeneratedAt:   now.UTC(),
		SecretSafe:    true,
		Live:          live,
		Archive:       archive,
	}
	if err := Validate(snapshot); err != nil {
		return Snapshot{}, err
	}
	return snapshot, nil
}

func Validate(snapshot Snapshot) error {
	if snapshot.SchemaVersion != SchemaVersion {
		return fmt.Errorf("unsupported access-rule audit schema %q", snapshot.SchemaVersion)
	}
	if !snapshot.SecretSafe {
		return errors.New("access-rule audit snapshot is not marked secret-safe")
	}
	for name, value := range map[string]string{
		"live": snapshot.Live.SHA256, "allow": snapshot.Live.AllowSHA256, "deny": snapshot.Live.DenySHA256,
	} {
		if !validSHA256(value) {
			return fmt.Errorf("%s SHA-256 is invalid", name)
		}
	}
	if snapshot.Live.Total != snapshot.Live.AllowCount+snapshot.Live.DenyCount {
		return errors.New("live access-rule counts are inconsistent")
	}
	var groupedTotal int64
	for _, count := range snapshot.Live.Counts {
		if count.SubjectType == "" || count.Effect == "" || count.ResourceType == "" || count.Status == "" || count.Count < 1 {
			return errors.New("access-rule grouped count is invalid")
		}
		groupedTotal += count.Count
	}
	if groupedTotal != snapshot.Live.Total {
		return errors.New("access-rule grouped counts do not match total")
	}
	if snapshot.Archive != nil {
		archive := snapshot.Archive
		if archive.MigrationBatch != MigrationBatch || !validSHA256(archive.ManifestAllowSHA) || !validSHA256(archive.ArchivedAllowSHA) || !validSHA256(archive.ManifestDenySHA) {
			return errors.New("access-rule archive manifest is invalid")
		}
		consistent := archive.ManifestAllowCount == archive.ArchivedAllowCount && archive.ManifestAllowSHA == archive.ArchivedAllowSHA
		if archive.Consistent != consistent {
			return errors.New("access-rule archive consistency marker is incorrect")
		}
	}
	return nil
}

func VerifyMigration(before Snapshot, after Snapshot) error {
	if err := Validate(before); err != nil {
		return fmt.Errorf("before snapshot: %w", err)
	}
	if err := Validate(after); err != nil {
		return fmt.Errorf("after snapshot: %w", err)
	}
	if after.Archive == nil || !after.Archive.Consistent {
		return errors.New("verified allow archive is unavailable")
	}
	if after.Live.AllowCount != 0 {
		return fmt.Errorf("live legacy allow count=%d, want 0", after.Live.AllowCount)
	}
	if before.Live.AllowCount != after.Archive.ManifestAllowCount || before.Live.AllowSHA256 != after.Archive.ManifestAllowSHA {
		return errors.New("pre-migration allow snapshot does not match archived manifest")
	}
	if before.Live.DenyCount != after.Live.DenyCount || before.Live.DenySHA256 != after.Live.DenySHA256 {
		return errors.New("deny rules changed during allow migration")
	}
	return nil
}

func Encode(snapshot Snapshot) ([]byte, error) {
	if err := Validate(snapshot); err != nil {
		return nil, err
	}
	payload, err := json.MarshalIndent(snapshot, "", "  ")
	if err != nil {
		return nil, err
	}
	return append(payload, '\n'), nil
}

func Decode(payload []byte) (Snapshot, error) {
	decoder := json.NewDecoder(strings.NewReader(string(payload)))
	decoder.DisallowUnknownFields()
	var snapshot Snapshot
	if err := decoder.Decode(&snapshot); err != nil {
		return Snapshot{}, err
	}
	if err := Validate(snapshot); err != nil {
		return Snapshot{}, err
	}
	return snapshot, nil
}

func loadLiveRules(ctx context.Context, pool *pgxpool.Pool) (RuleSet, error) {
	rows, err := pool.Query(ctx, `
SELECT id::text, subject_type, subject_id::text, resource_type,
       resource_id::text, effect, priority::text,
       min_permission_level::text, conditions::text, metadata::text,
       status, extract(epoch FROM created_at)::text,
       extract(epoch FROM updated_at)::text
FROM gateway_access_rules
ORDER BY id`)
	if err != nil {
		return RuleSet{}, err
	}
	defer rows.Close()
	allDigests := make([]string, 0)
	allowDigests := make([]string, 0)
	denyDigests := make([]string, 0)
	counts := map[string]int64{}
	for rows.Next() {
		var row ruleDigestRow
		if err := rows.Scan(
			&row.id, &row.subjectType, &row.subjectID, &row.resourceType,
			&row.resourceID, &row.effect, &row.priority,
			&row.minPermissionLevel, &row.conditions, &row.metadata,
			&row.status, &row.createdEpoch, &row.updatedEpoch,
		); err != nil {
			return RuleSet{}, err
		}
		digest := rowSHA256(row)
		allDigests = append(allDigests, digest)
		if row.effect == "allow" {
			allowDigests = append(allowDigests, digest)
		}
		if row.effect == "deny" {
			denyDigests = append(denyDigests, digest)
		}
		counts[strings.Join([]string{row.subjectType, row.effect, row.resourceType, row.status}, "\x00")]++
	}
	if err := rows.Err(); err != nil {
		return RuleSet{}, err
	}
	grouped := make([]RuleCount, 0, len(counts))
	for key, count := range counts {
		parts := strings.Split(key, "\x00")
		grouped = append(grouped, RuleCount{SubjectType: parts[0], Effect: parts[1], ResourceType: parts[2], Status: parts[3], Count: count})
	}
	sort.Slice(grouped, func(i, j int) bool {
		left := grouped[i]
		right := grouped[j]
		return strings.Join([]string{left.SubjectType, left.Effect, left.ResourceType, left.Status}, "\x00") <
			strings.Join([]string{right.SubjectType, right.Effect, right.ResourceType, right.Status}, "\x00")
	})
	return RuleSet{
		Total:       int64(len(allDigests)),
		SHA256:      aggregateSHA256(allDigests),
		AllowCount:  int64(len(allowDigests)),
		AllowSHA256: aggregateSHA256(allowDigests),
		DenyCount:   int64(len(denyDigests)),
		DenySHA256:  aggregateSHA256(denyDigests),
		Counts:      grouped,
	}, nil
}

func loadArchiveManifest(ctx context.Context, pool *pgxpool.Pool) (*ArchiveManifest, error) {
	var available bool
	if err := pool.QueryRow(ctx, `
SELECT to_regclass('gateway_access_rule_migration_batches') IS NOT NULL
   AND to_regclass('gateway_access_rule_allow_archive') IS NOT NULL`).Scan(&available); err != nil {
		return nil, err
	}
	if !available {
		return nil, nil
	}
	manifest := &ArchiveManifest{MigrationBatch: MigrationBatch}
	if err := pool.QueryRow(ctx, `
SELECT allow_count, allow_sha256, deny_count, deny_sha256
FROM gateway_access_rule_migration_batches
WHERE migration_batch = $1`, MigrationBatch).Scan(
		&manifest.ManifestAllowCount,
		&manifest.ManifestAllowSHA,
		&manifest.ManifestDenyCount,
		&manifest.ManifestDenySHA,
	); err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return nil, nil
		}
		return nil, err
	}
	rows, err := pool.Query(ctx, `
SELECT row_sha256
FROM gateway_access_rule_allow_archive
WHERE migration_batch = $1
ORDER BY rule_id`, MigrationBatch)
	if err != nil {
		return nil, err
	}
	defer rows.Close()
	digests := make([]string, 0)
	for rows.Next() {
		var digest string
		if err := rows.Scan(&digest); err != nil {
			return nil, err
		}
		digests = append(digests, digest)
	}
	if err := rows.Err(); err != nil {
		return nil, err
	}
	manifest.ArchivedAllowCount = int64(len(digests))
	manifest.ArchivedAllowSHA = aggregateSHA256(digests)
	manifest.Consistent = manifest.ManifestAllowCount == manifest.ArchivedAllowCount && manifest.ManifestAllowSHA == manifest.ArchivedAllowSHA
	return manifest, nil
}

func rowSHA256(row ruleDigestRow) string {
	return sha256Hex(strings.Join([]string{
		row.id, row.subjectType, row.subjectID, row.resourceType,
		row.resourceID, row.effect, row.priority, row.minPermissionLevel,
		row.conditions, row.metadata, row.status, row.createdEpoch, row.updatedEpoch,
	}, "\x1f"))
}

func aggregateSHA256(rowDigests []string) string {
	return sha256Hex(strings.Join(rowDigests, "\n"))
}

func sha256Hex(value string) string {
	digest := sha256.Sum256([]byte(value))
	return hex.EncodeToString(digest[:])
}

func validSHA256(value string) bool {
	if len(value) != 64 {
		return false
	}
	_, err := hex.DecodeString(value)
	return err == nil && value == strings.ToLower(value)
}