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

package sqlstore

import (
	"context"
	"crypto/rand"
	"database/sql"
	"encoding/base64"
	"errors"
	"fmt"

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

// Sync is the shared SQL implementation of storage.SyncStore: the streams a
// sync response carries that are not room events.
//
// Like the other component stores it is written once with `?` placeholders
// and rebound per driver. Every stream here follows the same shape — a
// counter row handed out inside the writing transaction, and a stream_pos
// column on the row it stamps — so that a reader can ask any of them "what
// changed after position N" with one index lookup.
type Sync struct {
	db       *sql.DB
	rebind   func(string) string
	isUnique func(error) bool
}

// NewSync builds a sync store.
func NewSync(db *sql.DB, rebind func(string) string, isUnique func(error) bool) *Sync {
	if rebind == nil {
		rebind = func(query string) string { return query }
	}
	if isUnique == nil {
		isUnique = func(error) bool { return false }
	}
	return &Sync{db: db, rebind: rebind, isUnique: isUnique}
}

// maxToDeviceBatch bounds how many to-device messages one sync carries. A
// device that has been offline for a week can have a long queue, and the
// rest of it is still there on the next sync.
const maxToDeviceBatch = 100

// q rebinds a statement for this driver.
func (s *Sync) q(query string) string { return s.rebind(query) }

// tx runs fn in a transaction, rolling back on any error.
func (s *Sync) 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 selectStreamPositionsSQL = `
	SELECT name, value FROM sync_sequences`

// StreamPositions implements storage.SyncStore.
//
// The PDU position comes from the events table and the rest from the
// counter rows, and both are read here so that a caller gets one consistent
// set. The counters are read rather than the rows they stamp because a
// stream whose latest row was deleted — a drained to-device queue — must
// not appear to have gone backwards.
func (s *Sync) StreamPositions(ctx context.Context) (storage.StreamPositions, error) {
	var out storage.StreamPositions

	err := s.db.QueryRowContext(ctx, s.q(selectMaxStreamOrderingSQL)).Scan(&out.PDU)
	if err != nil {
		return storage.StreamPositions{}, fmt.Errorf("read the top of the event stream: %w", err)
	}

	rows, err := s.db.QueryContext(ctx, s.q(selectStreamPositionsSQL))
	if err != nil {
		return storage.StreamPositions{}, fmt.Errorf("read the stream positions: %w", err)
	}
	defer func() { _ = rows.Close() }()

	for rows.Next() {
		var (
			name  string
			value int64
		)
		if err := rows.Scan(&name, &value); err != nil {
			return storage.StreamPositions{}, fmt.Errorf("read the stream positions: %w", err)
		}
		switch name {
		case seqReceipt:
			out.Receipt = value
		case seqToDevice:
			out.ToDevice = value
		case seqAccountData:
			out.AccountData = value
		case seqPresence:
			out.Presence = value
		}
	}
	if err := rows.Err(); err != nil {
		return storage.StreamPositions{}, fmt.Errorf("read the stream positions: %w", err)
	}

	// The device list stream is the keyserver's, and is read here all the
	// same, for the same reason the event stream is: a token has to name
	// one moment in every stream it carries.
	out.DeviceList, err = currentSequence(ctx, s.db, s.q, keySequences, seqDeviceList)
	if err != nil {
		return storage.StreamPositions{}, err
	}
	// And the push store's notification counts, for the same reason.
	out.Notification, err = currentSequence(ctx, s.db, s.q, pushSequences, seqNotification)
	if err != nil {
		return storage.StreamPositions{}, err
	}
	return out, nil
}

const (
	selectReceiptPositionSQL = `
		SELECT e.stream_ordering FROM roomserver_events e
		WHERE e.room_id = ? AND e.event_id = ?`

	selectCurrentReceiptSQL = `
		SELECT r.event_id, e.stream_ordering, r.stream_pos
		FROM sync_receipts r
		LEFT JOIN roomserver_events e ON e.event_id = r.event_id
		WHERE r.room_id = ? AND r.user_id = ? AND r.receipt_type = ? AND r.thread_id = ?`

	upsertReceiptSQL = `
		INSERT INTO sync_receipts
			(room_id, user_id, receipt_type, thread_id, event_id, ts_ms, stream_pos)
		VALUES (?, ?, ?, ?, ?, ?, ?)
		ON CONFLICT (room_id, user_id, receipt_type, thread_id) DO UPDATE SET
			event_id = excluded.event_id,
			ts_ms = excluded.ts_ms,
			stream_pos = excluded.stream_pos`
)

// SetReceipt implements storage.SyncStore.
//
// A receipt that would move backwards is dropped. Two clients of one user
// read a room at different speeds, and letting the slower one win would
// make the unread marker jump backwards on the other — so the comparison is
// on the position of the event being acknowledged, not on when the receipt
// arrived.
func (s *Sync) SetReceipt(ctx context.Context, r storage.Receipt) (int64, bool, error) {
	var (
		pos   int64
		moved bool
	)
	err := s.tx(ctx, func(tx *sql.Tx) error {
		var target int64
		err := tx.QueryRowContext(ctx, s.q(selectReceiptPositionSQL), r.RoomID, r.EventID).Scan(&target)
		switch {
		case errors.Is(err, sql.ErrNoRows):
			return fmt.Errorf("event %q: %w", r.EventID, storage.ErrNoSuchEvent)
		case err != nil:
			return fmt.Errorf("locate the event a receipt names: %w", err)
		}

		var (
			haveEventID string
			havePos     sql.NullInt64
			haveStream  int64
		)
		err = tx.QueryRowContext(ctx, s.q(selectCurrentReceiptSQL),
			r.RoomID, r.UserID, r.ReceiptType, r.ThreadID).Scan(&haveEventID, &havePos, &haveStream)
		switch {
		case err == nil:
			// A receipt whose event has since been purged has no position
			// to compare against, and the new one supersedes it.
			if havePos.Valid && havePos.Int64 >= target {
				pos, moved = haveStream, false
				return nil
			}
		case !errors.Is(err, sql.ErrNoRows):
			return fmt.Errorf("read the current receipt: %w", err)
		}

		if pos, err = nextSequence(ctx, tx, s.q, syncSequences, seqReceipt); err != nil {
			return err
		}
		if _, err := tx.ExecContext(ctx, s.q(upsertReceiptSQL),
			r.RoomID, r.UserID, r.ReceiptType, r.ThreadID, r.EventID, r.TimestampMS, pos); err != nil {
			return fmt.Errorf("write the receipt: %w", err)
		}
		moved = true
		return nil
	})
	if err != nil {
		return 0, false, fmt.Errorf("set receipt in %q: %w", r.RoomID, err)
	}
	return pos, moved, nil
}

const receiptColumns = `
	room_id, user_id, receipt_type, thread_id, event_id, ts_ms, stream_pos`

func scanReceipts(rows *sql.Rows) ([]storage.Receipt, error) {
	defer func() { _ = rows.Close() }()

	var out []storage.Receipt
	for rows.Next() {
		var r storage.Receipt
		if err := rows.Scan(&r.RoomID, &r.UserID, &r.ReceiptType, &r.ThreadID,
			&r.EventID, &r.TimestampMS, &r.StreamPos); err != nil {
			return nil, err
		}
		out = append(out, r)
	}
	return out, rows.Err()
}

// Receipts implements storage.SyncStore.
func (s *Sync) Receipts(ctx context.Context, roomIDs []string, from, to int64) ([]storage.Receipt, error) {
	if len(roomIDs) == 0 {
		return nil, nil
	}

	query := `SELECT ` + receiptColumns + ` FROM sync_receipts
		WHERE room_id IN (` + placeholders(len(roomIDs)) + `) AND stream_pos > ?`
	args := make([]any, 0, len(roomIDs)+2)
	for _, id := range roomIDs {
		args = append(args, id)
	}
	args = append(args, from)
	if to > 0 {
		query += ` AND stream_pos <= ?`
		args = append(args, to)
	}
	query += ` ORDER BY stream_pos`

	rows, err := s.db.QueryContext(ctx, s.q(query), args...)
	if err != nil {
		return nil, fmt.Errorf("read receipts: %w", err)
	}
	receipts, err := scanReceipts(rows)
	if err != nil {
		return nil, fmt.Errorf("read receipts: %w", err)
	}
	return receipts, nil
}

// ReceiptsForUser implements storage.SyncStore.
func (s *Sync) ReceiptsForUser(ctx context.Context, userID string, roomIDs []string) ([]storage.Receipt, error) {
	if len(roomIDs) == 0 {
		return nil, nil
	}

	query := `SELECT ` + receiptColumns + ` FROM sync_receipts
		WHERE user_id = ? AND room_id IN (` + placeholders(len(roomIDs)) + `)
		ORDER BY stream_pos`
	args := make([]any, 0, len(roomIDs)+1)
	args = append(args, userID)
	for _, id := range roomIDs {
		args = append(args, id)
	}

	rows, err := s.db.QueryContext(ctx, s.q(query), args...)
	if err != nil {
		return nil, fmt.Errorf("read the receipts of %q: %w", userID, err)
	}
	receipts, err := scanReceipts(rows)
	if err != nil {
		return nil, fmt.Errorf("read the receipts of %q: %w", userID, err)
	}
	return receipts, nil
}

const insertToDeviceSQL = `
	INSERT INTO sync_to_device
		(stream_pos, localpart, device_id, sender, event_type, content, sent_at_ms)
	VALUES (?, ?, ?, ?, ?, ?, ?)`

// SendToDevice implements storage.SyncStore.
//
// The whole batch shares one transaction but takes a position each: a
// message to two of a user's devices is two independent deliveries, and one
// device acknowledging its copy must not drop the other's.
func (s *Sync) SendToDevice(ctx context.Context, msgs []storage.NewToDeviceMessage) (int64, error) {
	var top int64
	err := s.tx(ctx, func(tx *sql.Tx) error {
		if len(msgs) == 0 {
			var err error
			top, err = currentSequence(ctx, tx, s.q, syncSequences, seqToDevice)
			return err
		}
		for _, m := range msgs {
			pos, err := nextSequence(ctx, tx, s.q, syncSequences, seqToDevice)
			if err != nil {
				return err
			}
			if _, err := tx.ExecContext(ctx, s.q(insertToDeviceSQL),
				pos, m.Localpart, m.DeviceID, m.Sender, m.Type, string(m.Content), m.SentAtMS); err != nil {
				return fmt.Errorf("queue a to-device message for %s/%s: %w", m.Localpart, m.DeviceID, err)
			}
			top = pos
		}
		return nil
	})
	if err != nil {
		return 0, fmt.Errorf("send to-device messages: %w", err)
	}
	return top, nil
}

const selectToDeviceSQL = `
	SELECT stream_pos, localpart, device_id, sender, event_type, content, sent_at_ms
	FROM sync_to_device
	WHERE localpart = ? AND device_id = ? AND stream_pos > ?`

// ToDeviceMessages implements storage.SyncStore.
func (s *Sync) ToDeviceMessages(ctx context.Context, localpart, deviceID string, from, to int64, limit int) ([]storage.ToDeviceMessage, error) {
	if limit <= 0 || limit > maxToDeviceBatch {
		limit = maxToDeviceBatch
	}

	query := selectToDeviceSQL
	args := []any{localpart, deviceID, from}
	if to > 0 {
		query += ` AND stream_pos <= ?`
		args = append(args, to)
	}
	query += ` ORDER BY stream_pos LIMIT ?`
	args = append(args, limit)

	rows, err := s.db.QueryContext(ctx, s.q(query), args...)
	if err != nil {
		return nil, fmt.Errorf("read to-device messages for %s/%s: %w", localpart, deviceID, err)
	}
	defer func() { _ = rows.Close() }()

	var out []storage.ToDeviceMessage
	for rows.Next() {
		var (
			m       storage.ToDeviceMessage
			content string
		)
		if err := rows.Scan(&m.StreamPos, &m.Localpart, &m.DeviceID,
			&m.Sender, &m.Type, &content, &m.SentAtMS); err != nil {
			return nil, fmt.Errorf("read to-device messages for %s/%s: %w", localpart, deviceID, err)
		}
		m.Content = []byte(content)
		out = append(out, m)
	}
	return out, rows.Err()
}

const deleteToDeviceSQL = `
	DELETE FROM sync_to_device
	WHERE localpart = ? AND device_id = ? AND stream_pos <= ?`

// DeleteToDeviceMessages implements storage.SyncStore.
func (s *Sync) DeleteToDeviceMessages(ctx context.Context, localpart, deviceID string, pos int64) error {
	if pos <= 0 {
		return nil
	}
	if _, err := s.db.ExecContext(ctx, s.q(deleteToDeviceSQL), localpart, deviceID, pos); err != nil {
		return fmt.Errorf("delete to-device messages for %s/%s: %w", localpart, deviceID, err)
	}
	return nil
}

const (
	upsertPresenceSQL = `
		INSERT INTO sync_presence (user_id, presence, status_msg, last_active_ms, stream_pos)
		VALUES (?, ?, ?, ?, ?)
		ON CONFLICT (user_id) DO UPDATE SET
			presence = excluded.presence,
			status_msg = excluded.status_msg,
			last_active_ms = excluded.last_active_ms,
			stream_pos = excluded.stream_pos`

	touchPresenceSQL = `
		UPDATE sync_presence SET last_active_ms = ? WHERE user_id = ?`
)

// SetPresence implements storage.SyncStore.
//
// A write that changes nothing a peer would notice refreshes the
// last-active timestamp in place and leaves the stream where it was. Every
// /sync is such a write — a client says it is still online by syncing — and
// moving the stream each time would wake every peer's long poll, which
// would sync, which would move it again.
func (s *Sync) SetPresence(ctx context.Context, p storage.PresenceState) (int64, bool, error) {
	var (
		pos   int64
		moved bool
	)
	err := s.tx(ctx, func(tx *sql.Tx) error {
		var existing storage.PresenceState
		err := tx.QueryRowContext(ctx, s.q(selectPresenceSQL), p.UserID).Scan(
			&existing.UserID, &existing.Presence, &existing.StatusMsg,
			&existing.LastActiveMS, &existing.StreamPos)
		switch {
		case err == nil:
			if existing.Presence == p.Presence && existing.StatusMsg == p.StatusMsg {
				pos = existing.StreamPos
				if _, err := tx.ExecContext(ctx, s.q(touchPresenceSQL),
					p.LastActiveMS, p.UserID); err != nil {
					return fmt.Errorf("refresh the last-active timestamp: %w", err)
				}
				return nil
			}
		case !errors.Is(err, sql.ErrNoRows):
			return fmt.Errorf("read the current presence: %w", err)
		}

		if pos, err = nextSequence(ctx, tx, s.q, syncSequences, seqPresence); err != nil {
			return err
		}
		if _, err := tx.ExecContext(ctx, s.q(upsertPresenceSQL),
			p.UserID, p.Presence, p.StatusMsg, p.LastActiveMS, pos); err != nil {
			return fmt.Errorf("write presence: %w", err)
		}
		moved = true
		return nil
	})
	if err != nil {
		return 0, false, fmt.Errorf("set the presence of %q: %w", p.UserID, err)
	}
	return pos, moved, nil
}

const presenceColumns = `user_id, presence, status_msg, last_active_ms, stream_pos`

const selectPresenceSQL = `
	SELECT ` + presenceColumns + ` FROM sync_presence WHERE user_id = ?`

// Presence implements storage.SyncStore.
func (s *Sync) Presence(ctx context.Context, userID string) (storage.PresenceState, error) {
	var p storage.PresenceState
	err := s.db.QueryRowContext(ctx, s.q(selectPresenceSQL), userID).
		Scan(&p.UserID, &p.Presence, &p.StatusMsg, &p.LastActiveMS, &p.StreamPos)
	switch {
	case errors.Is(err, sql.ErrNoRows):
		return storage.PresenceState{}, fmt.Errorf("presence of %q: %w", userID, storage.ErrNoSuchUser)
	case err != nil:
		return storage.PresenceState{}, fmt.Errorf("read the presence of %q: %w", userID, err)
	}
	return p, nil
}

// PresenceForUsers implements storage.SyncStore.
func (s *Sync) PresenceForUsers(ctx context.Context, userIDs []string, from, to int64) ([]storage.PresenceState, error) {
	if len(userIDs) == 0 {
		return nil, nil
	}

	query := `SELECT ` + presenceColumns + ` FROM sync_presence
		WHERE user_id IN (` + placeholders(len(userIDs)) + `) AND stream_pos > ?`
	args := make([]any, 0, len(userIDs)+2)
	for _, id := range userIDs {
		args = append(args, id)
	}
	args = append(args, from)
	if to > 0 {
		query += ` AND stream_pos <= ?`
		args = append(args, to)
	}
	query += ` ORDER BY stream_pos`

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

	var out []storage.PresenceState
	for rows.Next() {
		var p storage.PresenceState
		if err := rows.Scan(&p.UserID, &p.Presence, &p.StatusMsg, &p.LastActiveMS, &p.StreamPos); err != nil {
			return nil, fmt.Errorf("read presence: %w", err)
		}
		out = append(out, p)
	}
	return out, rows.Err()
}

const selectAccountDataChangesSQL = `
	SELECT room_id, data_type, stream_pos FROM user_account_data
	WHERE localpart = ? AND stream_pos > ?`

// AccountDataChanges implements storage.SyncStore.
func (s *Sync) AccountDataChanges(ctx context.Context, localpart string, from, to int64) ([]storage.AccountDataEntry, error) {
	query := selectAccountDataChangesSQL
	args := []any{localpart, from}
	if to > 0 {
		query += ` AND stream_pos <= ?`
		args = append(args, to)
	}
	query += ` ORDER BY stream_pos`

	return s.accountData(ctx, localpart, query, args)
}

const selectAllAccountDataSQL = `
	SELECT room_id, data_type, stream_pos FROM user_account_data
	WHERE localpart = ?`

// AccountDataAt implements storage.SyncStore.
func (s *Sync) AccountDataAt(ctx context.Context, localpart string, to int64) ([]storage.AccountDataEntry, error) {
	query := selectAllAccountDataSQL
	args := []any{localpart}
	if to > 0 {
		query += ` AND stream_pos <= ?`
		args = append(args, to)
	}
	query += ` ORDER BY room_id, data_type`

	return s.accountData(ctx, localpart, query, args)
}

func (s *Sync) accountData(ctx context.Context, localpart, query string, args []any) ([]storage.AccountDataEntry, error) {
	rows, err := s.db.QueryContext(ctx, s.q(query), args...)
	if err != nil {
		return nil, fmt.Errorf("read the account data of %q: %w", localpart, err)
	}
	defer func() { _ = rows.Close() }()

	var out []storage.AccountDataEntry
	for rows.Next() {
		var e storage.AccountDataEntry
		if err := rows.Scan(&e.RoomID, &e.DataType, &e.StreamPos); err != nil {
			return nil, fmt.Errorf("read the account data of %q: %w", localpart, err)
		}
		out = append(out, e)
	}
	return out, rows.Err()
}

const insertFilterSQL = `
	INSERT INTO sync_filters (localpart, filter_id, filter_json, created_at_ms)
	VALUES (?, ?, ?, ?)`

// SaveFilter implements storage.SyncStore.
//
// The ID is random rather than a counter. The spec says only that it is
// opaque, and a counter would be a per-user one: two users would both have
// a filter "1", so a filter ID leaking into a URL or a log would say which
// user it belonged to only by accident of context.
func (s *Sync) SaveFilter(ctx context.Context, localpart string, filter []byte, createdAtMS int64) (string, error) {
	for attempt := range 3 {
		filterID, err := newFilterID()
		if err != nil {
			return "", err
		}
		_, err = s.db.ExecContext(ctx, s.q(insertFilterSQL),
			localpart, filterID, string(filter), createdAtMS)
		switch {
		case err == nil:
			return filterID, nil
		case s.isUnique(err) && attempt < 2:
			// 96 bits of randomness collided, or the caller's own ID did.
			// Either way the answer is another ID.
			continue
		default:
			return "", fmt.Errorf("save a filter for %q: %w", localpart, err)
		}
	}
	return "", fmt.Errorf("save a filter for %q: could not find an unused filter ID", localpart)
}

// newFilterID mints a filter ID: 12 random bytes, URL-safe base64.
func newFilterID() (string, error) {
	var raw [12]byte
	if _, err := rand.Read(raw[:]); err != nil {
		return "", fmt.Errorf("generate a filter ID: %w", err)
	}
	return base64.RawURLEncoding.EncodeToString(raw[:]), nil
}

const selectFilterSQL = `
	SELECT filter_json FROM sync_filters WHERE localpart = ? AND filter_id = ?`

// Filter implements storage.SyncStore.
func (s *Sync) Filter(ctx context.Context, localpart, filterID string) ([]byte, error) {
	var body string
	err := s.db.QueryRowContext(ctx, s.q(selectFilterSQL), localpart, filterID).Scan(&body)
	switch {
	case errors.Is(err, sql.ErrNoRows):
		return nil, fmt.Errorf("filter %q: %w", filterID, storage.ErrNoSuchFilter)
	case err != nil:
		return nil, fmt.Errorf("read filter %q: %w", filterID, err)
	}
	return []byte(body), nil
}

// Compile-time proof that the shared implementation satisfies the interface.
var _ storage.SyncStore = (*Sync)(nil)