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

package sqlstore

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

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

// The cache of remote users' keys: storage.KeyStore's RemoteKeys and the
// methods after it.
//
// None of these writes touch the device list stream. A remote user's
// position in it is moved by MarkDeviceListsChanged, which the keyserver
// calls for every update that concerns a local client whether or not the
// cache could take the update in: the stream says "ask again", and the
// cache is only where the answer comes from.

const (
	selectRemoteUserSQL = `
		SELECT stream_id, stale FROM keyserver_remote_users WHERE user_id = ?`

	upsertRemoteUserSQL = `
		INSERT INTO keyserver_remote_users (user_id, stream_id, stale) VALUES (?, ?, ?)
		ON CONFLICT (user_id) DO UPDATE SET stream_id = excluded.stream_id, stale = excluded.stale`

	selectRemoteDevicesSQL = `
		SELECT device_id, display_name, key_json FROM keyserver_remote_device_keys
		WHERE user_id = ? ORDER BY device_id`

	upsertRemoteDeviceSQL = `
		INSERT INTO keyserver_remote_device_keys (user_id, device_id, display_name, key_json)
		VALUES (?, ?, ?, ?)
		ON CONFLICT (user_id, device_id) DO UPDATE SET
			display_name = excluded.display_name, key_json = excluded.key_json`

	deleteRemoteDeviceSQL = `
		DELETE FROM keyserver_remote_device_keys WHERE user_id = ? AND device_id = ?`

	selectRemoteCrossSigningSQL = `
		SELECT purpose, key_json FROM keyserver_remote_cross_signing_keys WHERE user_id = ?`

	insertRemoteCrossSigningSQL = `
		INSERT INTO keyserver_remote_cross_signing_keys (user_id, purpose, key_json) VALUES (?, ?, ?)`

	deleteRemoteCrossSigningSQL = `
		DELETE FROM keyserver_remote_cross_signing_keys WHERE user_id = ?`

	selectRemoteUsersSQL = `
		SELECT user_id FROM keyserver_remote_users ORDER BY user_id`
)

// remoteKeyTables is every table a cached user has rows in.
var remoteKeyTables = []string{
	"keyserver_remote_users", "keyserver_remote_device_keys", "keyserver_remote_cross_signing_keys",
}

// remoteMarker reads a cached user's marker row.
func (s *Keys) remoteMarker(ctx context.Context, db rowQuerier, userID string) (streamID int64, stale, ok bool, err error) {
	err = db.QueryRowContext(ctx, s.q(selectRemoteUserSQL), userID).Scan(&streamID, &stale)
	switch {
	case errors.Is(err, sql.ErrNoRows):
		return 0, false, false, nil
	case err != nil:
		return 0, false, false, fmt.Errorf("read the cached keys of %q: %w", userID, err)
	}
	return streamID, stale, true, nil
}

// RemoteKeys implements storage.KeyStore.
func (s *Keys) RemoteKeys(ctx context.Context, userID string) (storage.RemoteKeys, bool, error) {
	out := storage.RemoteKeys{UserID: userID, CrossSigning: map[string][]byte{}}

	// One transaction, so that the marker and the keys are one snapshot:
	// a copy read halfway through its replacement would be served as
	// fresh with half of another one's devices.
	var found bool
	err := s.tx(ctx, func(tx *sql.Tx) error {
		var err error
		out.StreamID, out.Stale, found, err = s.remoteMarker(ctx, tx, userID)
		if err != nil || !found {
			return err
		}
		if out.Devices, err = s.remoteDevices(ctx, tx, userID); err != nil {
			return err
		}
		out.CrossSigning, err = s.remoteCrossSigning(ctx, tx, userID)
		return err
	})
	if err != nil || !found {
		return storage.RemoteKeys{}, false, err
	}
	return out, true, nil
}

func (s *Keys) remoteDevices(ctx context.Context, tx *sql.Tx, userID string) ([]storage.RemoteDevice, error) {
	rows, err := tx.QueryContext(ctx, s.q(selectRemoteDevicesSQL), userID)
	if err != nil {
		return nil, fmt.Errorf("read the cached devices of %q: %w", userID, err)
	}
	defer func() { _ = rows.Close() }()

	var out []storage.RemoteDevice
	for rows.Next() {
		var d storage.RemoteDevice
		var keyJSON string
		if err := rows.Scan(&d.DeviceID, &d.DisplayName, &keyJSON); err != nil {
			return nil, fmt.Errorf("read the cached devices of %q: %w", userID, err)
		}
		d.JSON = []byte(keyJSON)
		out = append(out, d)
	}
	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("read the cached devices of %q: %w", userID, err)
	}
	return out, nil
}

func (s *Keys) remoteCrossSigning(ctx context.Context, tx *sql.Tx, userID string) (map[string][]byte, error) {
	rows, err := tx.QueryContext(ctx, s.q(selectRemoteCrossSigningSQL), userID)
	if err != nil {
		return nil, fmt.Errorf("read the cached cross-signing keys of %q: %w", userID, err)
	}
	defer func() { _ = rows.Close() }()

	out := map[string][]byte{}
	for rows.Next() {
		var purpose, keyJSON string
		if err := rows.Scan(&purpose, &keyJSON); err != nil {
			return nil, fmt.Errorf("read the cached cross-signing keys of %q: %w", userID, err)
		}
		out[purpose] = []byte(keyJSON)
	}
	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("read the cached cross-signing keys of %q: %w", userID, err)
	}
	return out, nil
}

// SetRemoteKeys implements storage.KeyStore.
func (s *Keys) SetRemoteKeys(ctx context.Context, keys storage.RemoteKeys) (bool, error) {
	stored := false
	err := s.tx(ctx, func(tx *sql.Tx) error {
		held, _, ok, err := s.remoteMarker(ctx, tx, keys.UserID)
		if err != nil {
			return err
		}
		// A snapshot older than what is held was fetched while an update
		// was on its way in. Storing it as fresh would lose that update
		// until the next one showed the gap; leaving the copy as it is —
		// stale, as the update left it — has the next query fetch again.
		if ok && held > keys.StreamID {
			return nil
		}

		if err := s.deleteRemote(ctx, tx, keys.UserID); err != nil {
			return err
		}
		if _, err := tx.ExecContext(ctx, s.q(upsertRemoteUserSQL), keys.UserID, keys.StreamID, false); err != nil {
			return fmt.Errorf("cache the keys of %q: %w", keys.UserID, err)
		}
		for _, d := range keys.Devices {
			if _, err := tx.ExecContext(ctx, s.q(upsertRemoteDeviceSQL), keys.UserID, d.DeviceID, d.DisplayName, string(d.JSON)); err != nil {
				return fmt.Errorf("cache the keys of %s/%s: %w", keys.UserID, d.DeviceID, err)
			}
		}
		if err := s.insertRemoteCrossSigning(ctx, tx, keys.UserID, keys.CrossSigning); err != nil {
			return err
		}
		stored = true
		return nil
	})
	return stored, err
}

// UpdateRemoteDevice implements storage.KeyStore.
func (s *Keys) UpdateRemoteDevice(ctx context.Context, u storage.RemoteDeviceUpdate) (storage.RemoteUpdateResult, error) {
	var result storage.RemoteUpdateResult
	err := s.tx(ctx, func(tx *sql.Tx) error {
		held, stale, ok, err := s.remoteMarker(ctx, tx, u.UserID)
		if err != nil {
			return err
		}

		// A stale copy keeps the highest position it has heard of, so that
		// a snapshot fetched before this update cannot be stored over it.
		markStale := func(r storage.RemoteUpdateResult) error {
			result = r
			if _, err := tx.ExecContext(ctx, s.q(upsertRemoteUserSQL), u.UserID, max(held, u.StreamID), true); err != nil {
				return fmt.Errorf("mark the cached keys of %q stale: %w", u.UserID, err)
			}
			return nil
		}

		switch {
		case !ok:
			return markStale(storage.RemoteUpdateNotCached)
		case stale:
			return markStale(storage.RemoteUpdateStale)
		case u.StreamID <= held:
			result = storage.RemoteUpdateDuplicate
			return nil
		case !followsOn(held, u.PrevIDs):
			return markStale(storage.RemoteUpdateStale)
		case !u.Deleted && u.JSON == nil:
			// An update that says a device changed without saying how
			// cannot be applied; the copy has to be fetched again.
			return markStale(storage.RemoteUpdateStale)
		}

		if u.Deleted {
			_, err = tx.ExecContext(ctx, s.q(deleteRemoteDeviceSQL), u.UserID, u.DeviceID)
		} else {
			_, err = tx.ExecContext(ctx, s.q(upsertRemoteDeviceSQL), u.UserID, u.DeviceID, u.DisplayName, string(u.JSON))
		}
		if err != nil {
			return fmt.Errorf("apply an update to the cached keys of %s/%s: %w", u.UserID, u.DeviceID, err)
		}
		if _, err := tx.ExecContext(ctx, s.q(upsertRemoteUserSQL), u.UserID, u.StreamID, false); err != nil {
			return fmt.Errorf("move the cached keys of %q: %w", u.UserID, err)
		}
		result = storage.RemoteUpdateApplied
		return nil
	})
	return result, err
}

// followsOn reports whether an update with the given prev_id comes right
// after the position held. An update with no prev_id is the first its
// server ever sent for the user, which only follows on from nothing.
func followsOn(held int64, prevIDs []int64) bool {
	if len(prevIDs) == 0 {
		return held == 0
	}
	return slices.Contains(prevIDs, held)
}

// SetRemoteCrossSigningKeys implements storage.KeyStore.
func (s *Keys) SetRemoteCrossSigningKeys(ctx context.Context, userID string, keys map[string][]byte) (bool, error) {
	stored := false
	err := s.tx(ctx, func(tx *sql.Tx) error {
		_, stale, ok, err := s.remoteMarker(ctx, tx, userID)
		if err != nil || !ok || stale {
			return err
		}
		if _, err := tx.ExecContext(ctx, s.q(deleteRemoteCrossSigningSQL), userID); err != nil {
			return fmt.Errorf("replace the cached cross-signing keys of %q: %w", userID, err)
		}
		if err := s.insertRemoteCrossSigning(ctx, tx, userID, keys); err != nil {
			return err
		}
		stored = true
		return nil
	})
	return stored, err
}

func (s *Keys) insertRemoteCrossSigning(ctx context.Context, tx *sql.Tx, userID string, keys map[string][]byte) error {
	for _, purpose := range sortedKeys(keys) {
		if keys[purpose] == nil {
			continue
		}
		if _, err := tx.ExecContext(ctx, s.q(insertRemoteCrossSigningSQL), userID, purpose, string(keys[purpose])); err != nil {
			return fmt.Errorf("cache the %s key of %q: %w", purpose, userID, err)
		}
	}
	return nil
}

// RemoteKeyUsers implements storage.KeyStore.
func (s *Keys) RemoteKeyUsers(ctx context.Context) ([]string, error) {
	rows, err := s.db.QueryContext(ctx, s.q(selectRemoteUsersSQL))
	if err != nil {
		return nil, fmt.Errorf("read the users with cached keys: %w", err)
	}
	defer func() { _ = rows.Close() }()

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

// DeleteRemoteKeys implements storage.KeyStore.
func (s *Keys) DeleteRemoteKeys(ctx context.Context, userIDs []string) error {
	if len(userIDs) == 0 {
		return nil
	}
	return s.tx(ctx, func(tx *sql.Tx) error {
		for _, userID := range dedupe(userIDs) {
			if err := s.deleteRemote(ctx, tx, userID); err != nil {
				return err
			}
		}
		return nil
	})
}

// deleteRemote drops every row cached for a user.
func (s *Keys) deleteRemote(ctx context.Context, tx *sql.Tx, userID string) error {
	for _, table := range remoteKeyTables {
		if _, err := tx.ExecContext(ctx, s.q(`DELETE FROM `+table+` WHERE user_id = ?`), userID); err != nil {
			return fmt.Errorf("delete from %s for %q: %w", table, userID, err)
		}
	}
	return nil
}