Warp/server/internal/storage/sqlstore/federation.go

package sqlstore

import (
	"context"
	"database/sql"
	"errors"
	"fmt"

	"github.com/dariuszjaniak/warp/server/internal/storage"
)

// Federation is the shared SQL implementation of storage.FederationStore.
type Federation struct {
	db     *sql.DB
	rebind func(string) string
}

// NewFederation builds a federation store.
func NewFederation(db *sql.DB, rebind func(string) string) *Federation {
	if rebind == nil {
		rebind = func(query string) string { return query }
	}
	return &Federation{db: db, rebind: rebind}
}

// The sequence table of this component, and its one counter.
const (
	federationSequences = "federation_sequences"
	seqQueue            = "queue"
)

func (s *Federation) q(query string) string { return s.rebind(query) }

func (s *Federation) tx(ctx context.Context, fn func(*sql.Tx) error) error {
	tx, err := s.db.BeginTx(ctx, nil)
	if err != nil {
		return fmt.Errorf("begin transaction: %w", err)
	}
	defer func() { _ = tx.Rollback() }()
	if err := fn(tx); err != nil {
		return err
	}
	if err := tx.Commit(); err != nil {
		return fmt.Errorf("commit transaction: %w", err)
	}
	return nil
}

const upsertServerKeySQL = `
	INSERT INTO federation_server_keys
		(server_name, key_id, public_key, valid_until_ts, expired_ts, fetched_at_ms)
	VALUES (?, ?, ?, ?, ?, ?)
	ON CONFLICT (server_name, key_id) DO UPDATE SET
		public_key = excluded.public_key,
		valid_until_ts = excluded.valid_until_ts,
		expired_ts = excluded.expired_ts,
		fetched_at_ms = excluded.fetched_at_ms`

// StoreServerKeys implements storage.FederationStore.
func (s *Federation) StoreServerKeys(ctx context.Context, keys []storage.ServerKey) error {
	if len(keys) == 0 {
		return nil
	}
	return s.tx(ctx, func(tx *sql.Tx) error {
		for _, k := range keys {
			if _, err := tx.ExecContext(ctx, s.q(upsertServerKeySQL),
				k.ServerName, k.KeyID, k.PublicKey, k.ValidUntilTS, k.ExpiredTS, k.FetchedAtMS); err != nil {
				return fmt.Errorf("store key %s of %s: %w", k.KeyID, k.ServerName, err)
			}
		}
		return nil
	})
}

// ServerKeys implements storage.FederationStore.
func (s *Federation) ServerKeys(ctx context.Context, serverName string, keyIDs []string) ([]storage.ServerKey, error) {
	query := `
		SELECT server_name, key_id, public_key, valid_until_ts, expired_ts, fetched_at_ms
		FROM federation_server_keys WHERE server_name = ?`
	args := []any{serverName}
	if len(keyIDs) > 0 {
		query += ` AND key_id IN (` + placeholders(len(keyIDs)) + `)`
		for _, id := range keyIDs {
			args = append(args, id)
		}
	}
	query += ` ORDER BY key_id`

	rows, err := s.db.QueryContext(ctx, s.q(query), args...)
	if err != nil {
		return nil, fmt.Errorf("read the keys of %q: %w", serverName, err)
	}
	defer func() { _ = rows.Close() }()

	var out []storage.ServerKey
	for rows.Next() {
		var k storage.ServerKey
		if err := rows.Scan(&k.ServerName, &k.KeyID, &k.PublicKey,
			&k.ValidUntilTS, &k.ExpiredTS, &k.FetchedAtMS); err != nil {
			return nil, fmt.Errorf("read the keys of %q: %w", serverName, err)
		}
		out = append(out, k)
	}
	return out, rows.Err()
}

const insertQueuedSQL = `
	INSERT INTO federation_queue
		(queue_id, destination, kind, room_id, event_id, payload, queued_at_ms)
	VALUES (?, ?, ?, ?, ?, ?, ?)`

// Enqueue implements storage.FederationStore.
func (s *Federation) Enqueue(ctx context.Context, items []storage.QueuedItem) ([]storage.QueuedItem, error) {
	if len(items) == 0 {
		return nil, nil
	}
	out := make([]storage.QueuedItem, len(items))
	copy(out, items)
	err := s.tx(ctx, func(tx *sql.Tx) error {
		top, err := advanceSequence(ctx, tx, s.q, federationSequences, seqQueue, int64(len(out)))
		if err != nil {
			return err
		}
		first := top - int64(len(out)) + 1
		for i := range out {
			out[i].ID = first + int64(i)
			item := out[i]
			if _, err := tx.ExecContext(ctx, s.q(insertQueuedSQL),
				item.ID, item.Destination, item.Kind, item.RoomID, item.EventID,
				string(item.Payload), item.QueuedAtMS); err != nil {
				return fmt.Errorf("queue for %s: %w", item.Destination, err)
			}
		}
		return nil
	})
	if err != nil {
		return nil, fmt.Errorf("enqueue %d items: %w", len(items), err)
	}
	return out, nil
}

const selectQueuedSQL = `
	SELECT queue_id, destination, kind, room_id, event_id, payload, queued_at_ms
	FROM federation_queue
	WHERE destination = ? AND kind = ?
	ORDER BY queue_id
	LIMIT ?`

// Queued implements storage.FederationStore.
func (s *Federation) Queued(ctx context.Context, destination, kind string, limit int) ([]storage.QueuedItem, error) {
	if limit <= 0 {
		return nil, nil
	}
	rows, err := s.db.QueryContext(ctx, s.q(selectQueuedSQL), destination, kind, limit)
	if err != nil {
		return nil, fmt.Errorf("read the queue for %q: %w", destination, err)
	}
	defer func() { _ = rows.Close() }()

	var out []storage.QueuedItem
	for rows.Next() {
		var (
			item    storage.QueuedItem
			payload string
		)
		if err := rows.Scan(&item.ID, &item.Destination, &item.Kind, &item.RoomID,
			&item.EventID, &payload, &item.QueuedAtMS); err != nil {
			return nil, fmt.Errorf("read the queue for %q: %w", destination, err)
		}
		item.Payload = []byte(payload)
		out = append(out, item)
	}
	return out, rows.Err()
}

// Dequeue implements storage.FederationStore.
func (s *Federation) Dequeue(ctx context.Context, ids []int64) error {
	if len(ids) == 0 {
		return nil
	}
	args := make([]any, len(ids))
	for i, id := range ids {
		args[i] = id
	}
	query := `DELETE FROM federation_queue WHERE queue_id IN (` + placeholders(len(ids)) + `)`
	if _, err := s.db.ExecContext(ctx, s.q(query), args...); err != nil {
		return fmt.Errorf("dequeue %d items: %w", len(ids), err)
	}
	return nil
}

const selectQueuedDestinationsSQL = `
	SELECT DISTINCT destination FROM federation_queue ORDER BY destination`

// QueuedDestinations implements storage.FederationStore.
func (s *Federation) QueuedDestinations(ctx context.Context) ([]string, error) {
	rows, err := s.db.QueryContext(ctx, s.q(selectQueuedDestinationsSQL))
	if err != nil {
		return nil, fmt.Errorf("read the queued destinations: %w", err)
	}
	defer func() { _ = rows.Close() }()

	var out []string
	for rows.Next() {
		var d string
		if err := rows.Scan(&d); err != nil {
			return nil, fmt.Errorf("read the queued destinations: %w", err)
		}
		out = append(out, d)
	}
	return out, rows.Err()
}

const destinationColumns = `
	server_name, failures, retry_at_ms, last_success_ms, last_failure_ms, last_error`

func scanDestination(scan func(...any) error) (storage.Destination, error) {
	var d storage.Destination
	err := scan(&d.ServerName, &d.Failures, &d.RetryAtMS, &d.LastSuccessMS, &d.LastFailureMS, &d.LastError)
	return d, err
}

const selectDestinationSQL = `SELECT ` + destinationColumns + `
	FROM federation_destinations WHERE server_name = ?`

// Destination implements storage.FederationStore.
func (s *Federation) Destination(ctx context.Context, serverName string) (storage.Destination, error) {
	d, err := scanDestination(s.db.QueryRowContext(ctx, s.q(selectDestinationSQL), serverName).Scan)
	switch {
	case errors.Is(err, sql.ErrNoRows):
		return storage.Destination{ServerName: serverName}, nil
	case err != nil:
		return storage.Destination{}, fmt.Errorf("read destination %q: %w", serverName, err)
	}
	return d, nil
}

const upsertDestinationSQL = `
	INSERT INTO federation_destinations (` + destinationColumns + `)
	VALUES (?, ?, ?, ?, ?, ?)
	ON CONFLICT (server_name) DO UPDATE SET
		failures = excluded.failures,
		retry_at_ms = excluded.retry_at_ms,
		last_success_ms = excluded.last_success_ms,
		last_failure_ms = excluded.last_failure_ms,
		last_error = excluded.last_error`

// SetDestination implements storage.FederationStore.
func (s *Federation) SetDestination(ctx context.Context, d storage.Destination) error {
	if _, err := s.db.ExecContext(ctx, s.q(upsertDestinationSQL),
		d.ServerName, d.Failures, d.RetryAtMS, d.LastSuccessMS, d.LastFailureMS, d.LastError); err != nil {
		return fmt.Errorf("record destination %q: %w", d.ServerName, err)
	}
	return nil
}

// selectDestinationSummariesSQL lists every server that either has a
// health row or something queued. The union is what makes a destination
// that has never failed, but has a backlog, show up at all.
const selectDestinationSummariesSQL = `
	SELECT n.server_name,
		COALESCE(d.failures, 0), COALESCE(d.retry_at_ms, 0),
		COALESCE(d.last_success_ms, 0), COALESCE(d.last_failure_ms, 0),
		COALESCE(d.last_error, ''),
		(SELECT COUNT(*) FROM federation_queue q WHERE q.destination = n.server_name AND q.kind = 'pdu'),
		(SELECT COUNT(*) FROM federation_queue q WHERE q.destination = n.server_name AND q.kind = 'edu')
	FROM (
		SELECT server_name FROM federation_destinations
		UNION
		SELECT DISTINCT destination FROM federation_queue
	) n
	LEFT JOIN federation_destinations d ON d.server_name = n.server_name
	ORDER BY n.server_name`

// Destinations implements storage.FederationStore.
func (s *Federation) Destinations(ctx context.Context) ([]storage.DestinationSummary, error) {
	rows, err := s.db.QueryContext(ctx, s.q(selectDestinationSummariesSQL))
	if err != nil {
		return nil, fmt.Errorf("read the destinations: %w", err)
	}
	defer func() { _ = rows.Close() }()

	var out []storage.DestinationSummary
	for rows.Next() {
		var d storage.DestinationSummary
		if err := rows.Scan(&d.ServerName, &d.Failures, &d.RetryAtMS, &d.LastSuccessMS,
			&d.LastFailureMS, &d.LastError, &d.QueuedPDUs, &d.QueuedEDUs); err != nil {
			return nil, fmt.Errorf("read the destinations: %w", err)
		}
		out = append(out, d)
	}
	return out, rows.Err()
}

const selectInboundTxnSQL = `
	SELECT response_json FROM federation_inbound_txns WHERE origin = ? AND txn_id = ?`

// InboundTransaction implements storage.FederationStore.
func (s *Federation) InboundTransaction(ctx context.Context, origin, txnID string) ([]byte, bool, error) {
	var response string
	err := s.db.QueryRowContext(ctx, s.q(selectInboundTxnSQL), origin, txnID).Scan(&response)
	switch {
	case errors.Is(err, sql.ErrNoRows):
		return nil, false, nil
	case err != nil:
		return nil, false, fmt.Errorf("read transaction %q from %q: %w", txnID, origin, err)
	}
	return []byte(response), true, nil
}

const insertInboundTxnSQL = `
	INSERT INTO federation_inbound_txns (origin, txn_id, response_json, received_at_ms)
	VALUES (?, ?, ?, ?)
	ON CONFLICT (origin, txn_id) DO NOTHING`

// RecordInboundTransaction implements storage.FederationStore.
func (s *Federation) RecordInboundTransaction(ctx context.Context, origin, txnID string, response []byte, atMS int64) error {
	if _, err := s.db.ExecContext(ctx, s.q(insertInboundTxnSQL),
		origin, txnID, string(response), atMS); err != nil {
		return fmt.Errorf("record transaction %q from %q: %w", txnID, origin, err)
	}
	return nil
}