restic-manager/internal/store/enrollment.go

package store

import (
	"context"
	"database/sql"
	"encoding/json"
	"errors"
	"fmt"
	"time"
)

// CreateEnrollmentToken persists a fresh one-time token. The caller
// has already hashed the raw token; the raw form is returned to the
// operator (printed in the install snippet) and never persisted.
//
// encRepoCreds is the AEAD-encrypted blob of {repo_url, repo_username,
// repo_password} that ConsumeEnrollmentToken will promote to a
// host_credentials row. Empty string = operator chose to set creds
// later via PUT /api/hosts/{id}/repo-credentials; the agent will
// refuse backup jobs until that lands.
//
// initialPaths is the JSON-encoded path list seeded into the host's
// initial manual schedule on consume. Empty string is treated as
// "[]". Not encrypted — paths aren't secret.
func (s *Store) CreateEnrollmentToken(ctx context.Context, tokenHash string, ttl time.Duration, encRepoCreds, initialPaths string) error {
	now := time.Now().UTC()
	var enc any = nil
	if encRepoCreds != "" {
		enc = encRepoCreds
	}
	if initialPaths == "" {
		initialPaths = "[]"
	}
	_, err := s.db.ExecContext(ctx,
		`INSERT INTO enrollment_tokens (token_hash, created_at, expires_at, enc_repo_creds, initial_paths)
		 VALUES (?, ?, ?, ?, ?)`,
		tokenHash,
		now.Format(time.RFC3339Nano),
		now.Add(ttl).Format(time.RFC3339Nano),
		enc, initialPaths)
	if err != nil {
		return fmt.Errorf("store: create enrollment token: %w", err)
	}
	return nil
}

// ConsumeEnrollmentToken atomically validates a token (must exist,
// not be consumed, not be expired) and marks it consumed by hostID.
// Returns ErrNotFound on any failure.
//
// The associated repo creds (if any) are promoted into
// host_credentials by the caller via SetHostCredentials *after* the
// host row exists — host_credentials has a FK to hosts that would
// otherwise fire here, since the host is created by a separate
// statement immediately after this returns.
func (s *Store) ConsumeEnrollmentToken(ctx context.Context, tokenHash, hostID string) error {
	now := time.Now().UTC().Format(time.RFC3339Nano)
	res, err := s.db.ExecContext(ctx,
		`UPDATE enrollment_tokens
		 SET consumed_at = ?, consumed_host = ?
		 WHERE token_hash = ? AND consumed_at IS NULL AND expires_at > ?`,
		now, hostID, tokenHash, now)
	if err != nil {
		return fmt.Errorf("store: consume enrollment token: %w", err)
	}
	n, _ := res.RowsAffected()
	if n == 0 {
		return ErrNotFound
	}
	return nil
}

// EnrollmentTokenAttachments is everything the enrolment handler
// needs from a token row at consume time, fetched in one round-trip.
type EnrollmentTokenAttachments struct {
	// EncRepoCreds is the AEAD ciphertext bound (additional-data) to
	// "token:" + token_hash. Empty if no creds were stashed.
	EncRepoCreds string
	// InitialPaths is the operator-supplied path list seeded into
	// the host's initial manual schedule. Always non-nil (empty
	// slice if none were set).
	InitialPaths []string
}

// GetEnrollmentTokenAttachments returns the operator-supplied
// attachments on a still-valid enrolment token: the encrypted repo
// creds and the default-paths list. Returns ErrNotFound if the
// token is gone / consumed / expired.
//
// The caller decrypts EncRepoCreds using token_hash as AEAD
// additional data, then re-encrypts using host_id as additional
// data before passing to ConsumeEnrollmentToken.
func (s *Store) GetEnrollmentTokenAttachments(ctx context.Context, tokenHash string) (EnrollmentTokenAttachments, error) {
	now := time.Now().UTC().Format(time.RFC3339Nano)
	row := s.db.QueryRowContext(ctx,
		`SELECT enc_repo_creds, initial_paths FROM enrollment_tokens
		 WHERE token_hash = ? AND consumed_at IS NULL AND expires_at > ?`,
		tokenHash, now)
	var (
		enc          sql.NullString
		initialPaths string
	)
	if err := row.Scan(&enc, &initialPaths); err != nil {
		if errors.Is(err, sql.ErrNoRows) {
			return EnrollmentTokenAttachments{}, ErrNotFound
		}
		return EnrollmentTokenAttachments{}, fmt.Errorf("store: get enrollment token attachments: %w", err)
	}
	out := EnrollmentTokenAttachments{InitialPaths: []string{}}
	if enc.Valid {
		out.EncRepoCreds = enc.String
	}
	if initialPaths != "" {
		_ = json.Unmarshal([]byte(initialPaths), &out.InitialPaths)
	}
	return out, nil
}

// EnrollmentTokenStatus is what the awaiting-agent panel polls for
// after Add-host. Returned by GetEnrollmentTokenStatus; the
// consuming code branches on Consumed + the (optional) ConsumedHost.
type EnrollmentTokenStatus struct {
	ExpiresAt    time.Time
	ConsumedAt   *time.Time
	ConsumedHost *string
}

// GetEnrollmentTokenStatus reports whether a token has been
// consumed yet (the agent has called /api/agents/enroll). Returns
// ErrNotFound if the token is unknown — the polling endpoint maps
// that to "token expired or invalid; stop polling".
func (s *Store) GetEnrollmentTokenStatus(ctx context.Context, tokenHash string) (EnrollmentTokenStatus, error) {
	row := s.db.QueryRowContext(ctx,
		`SELECT expires_at, consumed_at, consumed_host
		 FROM enrollment_tokens WHERE token_hash = ?`,
		tokenHash)
	var (
		expiresAt        string
		consumedAt, host sql.NullString
	)
	if err := row.Scan(&expiresAt, &consumedAt, &host); err != nil {
		if errors.Is(err, sql.ErrNoRows) {
			return EnrollmentTokenStatus{}, ErrNotFound
		}
		return EnrollmentTokenStatus{}, fmt.Errorf("store: get enrollment token status: %w", err)
	}
	out := EnrollmentTokenStatus{}
	if t, err := time.Parse(time.RFC3339Nano, expiresAt); err == nil {
		out.ExpiresAt = t
	}
	if consumedAt.Valid {
		if t, err := time.Parse(time.RFC3339Nano, consumedAt.String); err == nil {
			out.ConsumedAt = &t
		}
	}
	if host.Valid {
		s := host.String
		out.ConsumedHost = &s
	}
	return out, nil
}

// PurgeExpiredEnrollmentTokens deletes long-expired token rows. Tokens
// retained for ~24h after expiry so audit traces still resolve them.
func (s *Store) PurgeExpiredEnrollmentTokens(ctx context.Context) (int64, error) {
	cutoff := time.Now().Add(-24 * time.Hour).UTC().Format(time.RFC3339Nano)
	res, err := s.db.ExecContext(ctx,
		`DELETE FROM enrollment_tokens WHERE expires_at <= ?`, cutoff)
	if err != nil {
		return 0, fmt.Errorf("store: purge enrollment tokens: %w", err)
	}
	n, _ := res.RowsAffected()
	return n, nil
}