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

package pocketbase

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

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

// Sync is the PocketBase implementation of storage.SyncStore. Every stream
// here is a counter row advanced in the same batch as the row it stamps.
type Sync struct{ *base }

// maxToDeviceBatch bounds how many to-device messages one sync carries.
const maxToDeviceBatch = 100

// StreamPositions implements storage.SyncStore. The positions are read under
// one read lock, so that a caller gets one consistent set.
func (s *Sync) StreamPositions(ctx context.Context) (storage.StreamPositions, error) {
	var out storage.StreamPositions
	err := s.read(ctx, func(ctx context.Context) error {
		err := s.db.QueryRowContext(ctx, `SELECT COALESCE(MAX(stream_ordering), 0) FROM roomserver_events`).Scan(&out.PDU)
		if err != nil {
			return fmt.Errorf("read the top of the event stream: %w", err)
		}

		rows, err := s.db.QueryContext(ctx, `SELECT name, value FROM sync_sequences`)
		if err != nil {
			return 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 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 fmt.Errorf("read the stream positions: %w", err)
		}

		if out.DeviceList, err = s.currentSequence(ctx, keySequences, seqDeviceList); err != nil {
			return err
		}
		out.Notification, err = s.currentSequence(ctx, pushSequences, seqNotification)
		return err
	})
	if err != nil {
		return storage.StreamPositions{}, err
	}
	return out, nil
}

// SetReceipt implements storage.SyncStore. A receipt that would move
// backwards is dropped: 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.write(ctx, func(t *tx) error {
		var target int64
		err := t.queryRow(`SELECT e.stream_ordering FROM roomserver_events e
			WHERE e.room_id = ? AND e.event_id = ?`, 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 = t.queryRow(`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 = ?`,
			r.RoomID, r.UserID, r.ReceiptType, r.ThreadID).Scan(&haveEventID, &havePos, &haveStream)
		switch {
		case err == nil:
			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 = t.next(syncSequences, seqReceipt); err != nil {
			return err
		}
		t.put("sync_receipts", rec{
			"room_id": r.RoomID, "user_id": r.UserID, "receipt_type": r.ReceiptType, "thread_id": r.ThreadID,
			"event_id": r.EventID, "ts_ms": r.TimestampMS, "stream_pos": pos,
		})
		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 ` + inList + ` AND stream_pos > ?`
	args := []any{jsonList(roomIDs), from}
	if to > 0 {
		query += ` AND stream_pos <= ?`
		args = append(args, to)
	}
	query += ` ORDER BY stream_pos`

	rows, err := s.db.QueryContext(ctx, 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
	}
	rows, err := s.db.QueryContext(ctx, `SELECT `+receiptColumns+` FROM sync_receipts
		WHERE user_id = ? AND room_id IN `+inList+`
		ORDER BY stream_pos`, userID, jsonList(roomIDs))
	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
}

// SendToDevice implements storage.SyncStore. Each message takes a position
// of its own, so that one device acknowledging its copy does not drop
// another's.
func (s *Sync) SendToDevice(ctx context.Context, msgs []storage.NewToDeviceMessage) (int64, error) {
	if len(msgs) == 0 {
		top, err := s.currentSequence(ctx, syncSequences, seqToDevice)
		if err != nil {
			return 0, fmt.Errorf("send to-device messages: %w", err)
		}
		return top, nil
	}
	var top int64
	err := s.write(ctx, func(t *tx) error {
		for _, m := range msgs {
			pos, err := t.next(syncSequences, seqToDevice)
			if err != nil {
				return err
			}
			t.put("sync_to_device", rec{
				"stream_pos": pos, "localpart": m.Localpart, "device_id": m.DeviceID,
				"sender": m.Sender, "event_type": m.Type, "content": string(m.Content), "sent_at_ms": m.SentAtMS,
			})
			top = pos
		}
		return nil
	})
	if err != nil {
		return 0, fmt.Errorf("send to-device messages: %w", err)
	}
	return top, nil
}

// 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 := `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 > ?`
	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, 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()
}

// DeleteToDeviceMessages implements storage.SyncStore.
func (s *Sync) DeleteToDeviceMessages(ctx context.Context, localpart, deviceID string, pos int64) error {
	if pos <= 0 {
		return nil
	}
	err := s.write(ctx, func(t *tx) error {
		_, err := t.delWhere("sync_to_device", `localpart = ? AND device_id = ? AND stream_pos <= ?`,
			localpart, deviceID, pos)
		return err
	})
	if err != nil {
		return fmt.Errorf("delete to-device messages for %s/%s: %w", localpart, deviceID, err)
	}
	return nil
}

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

// SetPresence implements storage.SyncStore. A write that changes nothing a
// peer would notice refreshes the last-active timestamp and leaves the
// stream where it was.
func (s *Sync) SetPresence(ctx context.Context, p storage.PresenceState) (int64, bool, error) {
	var (
		pos   int64
		moved bool
	)
	err := s.write(ctx, func(t *tx) error {
		var existing storage.PresenceState
		err := t.queryRow(`SELECT `+presenceColumns+` FROM sync_presence WHERE user_id = ?`, 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
				t.put("sync_presence", rec{"user_id": p.UserID, "last_active_ms": p.LastActiveMS})
				return nil
			}
		case !errors.Is(err, sql.ErrNoRows):
			return fmt.Errorf("read the current presence: %w", err)
		}

		if pos, err = t.next(syncSequences, seqPresence); err != nil {
			return err
		}
		t.put("sync_presence", rec{
			"user_id": p.UserID, "presence": p.Presence, "status_msg": p.StatusMsg,
			"last_active_ms": p.LastActiveMS, "stream_pos": pos,
		})
		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
}

// 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, `SELECT `+presenceColumns+` FROM sync_presence WHERE user_id = ?`, 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 ` + inList + ` AND stream_pos > ?`
	args := []any{jsonList(userIDs), from}
	if to > 0 {
		query += ` AND stream_pos <= ?`
		args = append(args, to)
	}
	query += ` ORDER BY stream_pos`

	rows, err := s.db.QueryContext(ctx, 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()
}

// AccountDataChanges implements storage.SyncStore.
func (s *Sync) AccountDataChanges(ctx context.Context, localpart string, from, to int64) ([]storage.AccountDataEntry, error) {
	query := `SELECT room_id, data_type, stream_pos FROM user_account_data
		WHERE localpart = ? AND stream_pos > ?`
	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)
}

// AccountDataAt implements storage.SyncStore.
func (s *Sync) AccountDataAt(ctx context.Context, localpart string, to int64) ([]storage.AccountDataEntry, error) {
	query := `SELECT room_id, data_type, stream_pos FROM user_account_data WHERE localpart = ?`
	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, 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()
}

// SaveFilter implements storage.SyncStore. The ID is random rather than a
// counter, so that it says nothing about whose filter it is.
func (s *Sync) SaveFilter(ctx context.Context, localpart string, filter []byte, createdAtMS int64) (string, error) {
	var filterID string
	err := s.write(ctx, func(t *tx) error {
		for range 3 {
			id, err := newFilterID()
			if err != nil {
				return err
			}
			taken, err := t.exists("sync_filters", `localpart = ? AND filter_id = ?`, localpart, id)
			if err != nil {
				return err
			}
			if taken {
				continue
			}
			t.put("sync_filters", rec{
				"localpart": localpart, "filter_id": id,
				"filter_json": string(filter), "created_at_ms": createdAtMS,
			})
			filterID = id
			return nil
		}
		return errors.New("could not find an unused filter ID")
	})
	if err != nil {
		return "", fmt.Errorf("save a filter for %q: %w", localpart, err)
	}
	return filterID, nil
}

// 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
}

// Filter implements storage.SyncStore.
func (s *Sync) Filter(ctx context.Context, localpart, filterID string) ([]byte, error) {
	var body string
	err := s.db.QueryRowContext(ctx, `SELECT filter_json FROM sync_filters
		WHERE localpart = ? AND filter_id = ?`, 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
}

var _ storage.SyncStore = (*Sync)(nil)