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

package pocketbase

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

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

// The cache of other servers' users' keys, as the SQL store keeps it.

var remoteKeyTables = []string{
	"keyserver_remote_users", "keyserver_remote_device_keys", "keyserver_remote_cross_signing_keys",
}

// remoteMarker reads a cached user's marker row.
func remoteMarker(row rowQuerier, userID string) (streamID int64, stale, ok bool, err error) {
	err = row(`SELECT stream_id, stale FROM keyserver_remote_users WHERE user_id = ?`, 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. The marker and the keys are read
// under one read lock, so that a copy halfway through its replacement is
// never served.
func (s *Keys) RemoteKeys(ctx context.Context, userID string) (storage.RemoteKeys, bool, error) {
	out := storage.RemoteKeys{UserID: userID, CrossSigning: map[string][]byte{}}
	var found bool
	err := s.read(ctx, func(ctx context.Context) error {
		var err error
		out.StreamID, out.Stale, found, err = remoteMarker(s.poolRow(ctx), userID)
		if err != nil || !found {
			return err
		}
		if out.Devices, err = s.remoteDevices(ctx, userID); err != nil {
			return err
		}
		out.CrossSigning, err = s.remoteCrossSigning(ctx, userID)
		return err
	})
	if err != nil || !found {
		return storage.RemoteKeys{}, false, err
	}
	return out, true, nil
}

func (s *Keys) remoteDevices(ctx context.Context, userID string) ([]storage.RemoteDevice, error) {
	rows, err := s.db.QueryContext(ctx, `SELECT device_id, display_name, key_json FROM keyserver_remote_device_keys
		WHERE user_id = ? ORDER BY device_id`, 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, userID string) (map[string][]byte, error) {
	rows, err := s.db.QueryContext(ctx, `SELECT purpose, key_json FROM keyserver_remote_cross_signing_keys
		WHERE user_id = ?`, 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. A snapshot older than what is
// held is not stored: it was fetched while an update was on its way in.
func (s *Keys) SetRemoteKeys(ctx context.Context, keys storage.RemoteKeys) (bool, error) {
	stored := false
	err := s.write(ctx, func(t *tx) error {
		held, _, ok, err := remoteMarker(t.queryRow, keys.UserID)
		if err != nil {
			return err
		}
		if ok && held > keys.StreamID {
			return nil
		}

		if err := deleteRemote(t, keys.UserID); err != nil {
			return err
		}
		t.put("keyserver_remote_users", rec{"user_id": keys.UserID, "stream_id": keys.StreamID, "stale": false})
		for _, d := range keys.Devices {
			t.put("keyserver_remote_device_keys", rec{
				"user_id": keys.UserID, "device_id": d.DeviceID,
				"display_name": d.DisplayName, "key_json": string(d.JSON),
			})
		}
		putRemoteCrossSigning(t, keys.UserID, keys.CrossSigning)
		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.write(ctx, func(t *tx) error {
		held, stale, ok, err := remoteMarker(t.queryRow, 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
			t.put("keyserver_remote_users", rec{"user_id": u.UserID, "stream_id": max(held, u.StreamID), "stale": true})
			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:
			return markStale(storage.RemoteUpdateStale)
		}

		if u.Deleted {
			if _, err := t.delKey("keyserver_remote_device_keys", u.UserID, u.DeviceID); err != nil {
				return fmt.Errorf("apply an update to the cached keys of %s/%s: %w", u.UserID, u.DeviceID, err)
			}
		} else {
			t.put("keyserver_remote_device_keys", rec{
				"user_id": u.UserID, "device_id": u.DeviceID,
				"display_name": u.DisplayName, "key_json": string(u.JSON),
			})
		}
		t.put("keyserver_remote_users", rec{"user_id": u.UserID, "stream_id": u.StreamID, "stale": false})
		result = storage.RemoteUpdateApplied
		return nil
	})
	return result, err
}

// followsOn reports whether an update with the given prev_id comes right
// after the position held.
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.write(ctx, func(t *tx) error {
		_, stale, ok, err := remoteMarker(t.queryRow, userID)
		if err != nil || !ok || stale {
			return err
		}
		if _, err := t.delWhere("keyserver_remote_cross_signing_keys", `user_id = ?`, userID); err != nil {
			return fmt.Errorf("replace the cached cross-signing keys of %q: %w", userID, err)
		}
		putRemoteCrossSigning(t, userID, keys)
		stored = true
		return nil
	})
	return stored, err
}

func putRemoteCrossSigning(t *tx, userID string, keys map[string][]byte) {
	for _, purpose := range sortedKeys(keys) {
		if keys[purpose] == nil {
			continue
		}
		t.put("keyserver_remote_cross_signing_keys", rec{"user_id": userID, "purpose": purpose, "key_json": string(keys[purpose])})
	}
}

// RemoteKeyUsers implements storage.KeyStore.
func (s *Keys) RemoteKeyUsers(ctx context.Context) ([]string, error) {
	out, err := s.queryStrings(ctx, `SELECT user_id FROM keyserver_remote_users ORDER BY user_id`)
	if err != nil {
		return nil, fmt.Errorf("read the users with cached keys: %w", err)
	}
	return out, nil
}

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

// deleteRemote drops every row cached for a user.
func deleteRemote(t *tx, userID string) error {
	for _, table := range remoteKeyTables {
		if _, err := t.delWhere(table, `user_id = ?`, userID); err != nil {
			return fmt.Errorf("delete from %s for %q: %w", table, userID, err)
		}
	}
	return nil
}