at-container-registry/pkg/hold/labeler/cache.go

// Package labeler provides a labeler subscription client for the hold service.
//
// The hold subscribes to one labeler and mirrors active takedowns into a local
// cache (in-memory + SQLite). On takedown receipt the hold purges layer, scan,
// and image-config records for the affected manifest. The cache is consulted
// by the GC to gate blob deletion: while a takedown is within its grace window,
// the manifest's blob digests stay in the GC's referenced set so reversal can
// still restore them.
package labeler

import (
	"database/sql"
	"fmt"
	"strings"
	"sync"
	"time"
)

// Cache holds active takedown URIs and their creation timestamps. It is
// thread-safe and persists state to SQLite so the hold can answer takedown
// queries before the labeler subscription has caught up after a restart.
type Cache struct {
	mu sync.RWMutex
	// manifest URI → cts. Includes both per-record URIs (at://did/coll/rkey)
	// and per-DID URIs (at://did) for user-level takedowns.
	entries map[string]time.Time

	db *sql.DB
}

// NewCache opens (or creates) the takedown_cache and labeler_cursor tables on
// the given DB and loads any existing entries into memory.
func NewCache(db *sql.DB) (*Cache, error) {
	c := &Cache{
		entries: make(map[string]time.Time),
		db:      db,
	}

	stmts := []string{
		`CREATE TABLE IF NOT EXISTS takedown_cache (
			uri TEXT PRIMARY KEY,
			src TEXT NOT NULL,
			cts TIMESTAMP NOT NULL
		)`,
		`CREATE INDEX IF NOT EXISTS idx_takedown_cache_cts ON takedown_cache(cts)`,
		`CREATE TABLE IF NOT EXISTS labeler_cursor (
			labeler_did TEXT PRIMARY KEY,
			cursor INTEGER NOT NULL
		)`,
	}
	for _, s := range stmts {
		if _, err := db.Exec(s); err != nil {
			return nil, fmt.Errorf("init labeler schema: %w", err)
		}
	}

	rows, err := db.Query(`SELECT uri, cts FROM takedown_cache`)
	if err != nil {
		return nil, fmt.Errorf("load takedown cache: %w", err)
	}
	defer rows.Close()
	for rows.Next() {
		var uri string
		var cts time.Time
		if err := rows.Scan(&uri, &cts); err != nil {
			return nil, fmt.Errorf("scan takedown_cache row: %w", err)
		}
		c.entries[uri] = cts
	}
	return c, nil
}

// Set records a positive takedown for uri at cts. Idempotent: re-applying
// updates the timestamp (newer takedowns win).
func (c *Cache) Set(uri, src string, cts time.Time) error {
	c.mu.Lock()
	c.entries[uri] = cts
	c.mu.Unlock()

	_, err := c.db.Exec(
		`INSERT INTO takedown_cache (uri, src, cts) VALUES (?, ?, ?)
		 ON CONFLICT(uri) DO UPDATE SET src = excluded.src, cts = excluded.cts`,
		uri, src, cts,
	)
	if err != nil {
		return fmt.Errorf("persist takedown: %w", err)
	}
	return nil
}

// Negate removes a takedown entry. Idempotent.
func (c *Cache) Negate(uri string) error {
	c.mu.Lock()
	delete(c.entries, uri)
	c.mu.Unlock()

	_, err := c.db.Exec(`DELETE FROM takedown_cache WHERE uri = ?`, uri)
	if err != nil {
		return fmt.Errorf("delete takedown: %w", err)
	}
	return nil
}

// IsTakenDown reports whether a manifest URI is taken down, either directly
// (per-manifest takedown) or via a user-level takedown on its DID. The returned
// timestamp is the earliest cts that applies (i.e. the longest-standing
// takedown), which is what the grace check should compare against.
func (c *Cache) IsTakenDown(manifestURI string) (cts time.Time, ok bool) {
	c.mu.RLock()
	defer c.mu.RUnlock()

	if t, has := c.entries[manifestURI]; has {
		cts = t
		ok = true
	}

	// User-level: at://<did> with no path, applies to every record by that DID.
	did := didFromManifestURI(manifestURI)
	if did != "" {
		userURI := "at://" + did
		if t, has := c.entries[userURI]; has {
			if !ok || t.Before(cts) {
				cts = t
				ok = true
			}
		}
	}
	return cts, ok
}

// IsExpired returns true if cts is older than the grace window.
func IsExpired(cts time.Time, graceWindow time.Duration) bool {
	if graceWindow <= 0 {
		return true
	}
	return time.Since(cts) > graceWindow
}

// GetCursor returns the last persisted cursor for a labeler DID (0 if none).
func (c *Cache) GetCursor(labelerDID string) (int64, error) {
	var cursor int64
	err := c.db.QueryRow(`SELECT cursor FROM labeler_cursor WHERE labeler_did = ?`, labelerDID).Scan(&cursor)
	if err == sql.ErrNoRows {
		return 0, nil
	}
	if err != nil {
		return 0, fmt.Errorf("read labeler cursor: %w", err)
	}
	return cursor, nil
}

// SetCursor persists the cursor for a labeler DID.
func (c *Cache) SetCursor(labelerDID string, cursor int64) error {
	_, err := c.db.Exec(
		`INSERT INTO labeler_cursor (labeler_did, cursor) VALUES (?, ?)
		 ON CONFLICT(labeler_did) DO UPDATE SET cursor = excluded.cursor`,
		labelerDID, cursor,
	)
	if err != nil {
		return fmt.Errorf("persist labeler cursor: %w", err)
	}
	return nil
}

// didFromManifestURI extracts the authority (DID) from an at:// URI.
// Returns "" for malformed input.
func didFromManifestURI(uri string) string {
	const prefix = "at://"
	if !strings.HasPrefix(uri, prefix) {
		return ""
	}
	rest := uri[len(prefix):]
	if before, _, ok := strings.Cut(rest, "/"); ok {
		return before
	}
	return rest
}