Warp/server/internal/keyserver/query.go

package keyserver

import (
	"context"
	"encoding/json"
	"errors"
	"slices"
	"sync"
	"time"

	"github.com/dariuszjaniak/warp/server/internal/federation"
	"github.com/dariuszjaniak/warp/server/internal/matrix/id"
	"github.com/dariuszjaniak/warp/server/internal/storage"
)

// Reading keys: POST /keys/query and /keys/claim, and their federation
// counterparts.

// QueryRequest is the body of POST /keys/query.
type QueryRequest struct {
	// DeviceKeys is user ID to the device IDs wanted; an empty list asks
	// for every device.
	DeviceKeys map[string][]string

	// Timeout is how long to wait for other servers.
	Timeout time.Duration
}

// QueryResponse is the answer to POST /keys/query.
type QueryResponse struct {
	DeviceKeys      map[string]map[string]json.RawMessage `json:"device_keys"`
	Failures        map[string]any                        `json:"failures"`
	MasterKeys      map[string]json.RawMessage            `json:"master_keys"`
	SelfSigningKeys map[string]json.RawMessage            `json:"self_signing_keys"`
	UserSigningKeys map[string]json.RawMessage            `json:"user_signing_keys"`
}

// ClaimRequest is the body of POST /keys/claim.
type ClaimRequest struct {
	// OneTimeKeys is user ID, then device ID, then the algorithm wanted.
	OneTimeKeys map[string]map[string]string

	Timeout time.Duration
}

// ClaimResponse is the answer to POST /keys/claim.
type ClaimResponse struct {
	OneTimeKeys map[string]map[string]map[string]json.RawMessage `json:"one_time_keys"`
	Failures    map[string]any                                   `json:"failures"`
}

// remoteFailure is what a server that could not be asked is reported as,
// in the shape other homeservers use.
func remoteFailure(err error) map[string]any {
	status := 503
	var httpErr *federation.HTTPError
	if errors.As(err, &httpErr) {
		status = httpErr.Status
	}
	return map[string]any{"status": status, "message": err.Error()}
}

// --- reading keys ---------------------------------------------------------------

// publishedDeviceKeys reads the keys of a user's devices with every stored
// signature merged in. For a local user only devices that still exist are
// returned, and names maps them to their display names when the caller
// wants those.
func (k *API) publishedDeviceKeys(ctx context.Context, userID string, deviceIDs []string) (map[string]json.RawMessage, error) {
	return k.deviceKeys(ctx, userID, deviceIDs, false)
}

func (k *API) deviceKeys(ctx context.Context, userID string, deviceIDs []string, withNames bool) (map[string]json.RawMessage, error) {
	stored, err := k.store.DeviceKeys(ctx, userID, deviceIDs)
	if err != nil {
		return nil, err
	}
	out := map[string]json.RawMessage{}
	if len(stored) == 0 {
		return out, nil
	}

	names := map[string]string{}
	if localpart, ok := k.localpartOf(userID); ok {
		devices, err := k.devices.Devices(ctx, localpart)
		if err != nil {
			return nil, err
		}
		for _, d := range devices {
			names[d.ID] = d.DisplayName
		}
	}

	sigs, err := k.store.Signatures(ctx, userID)
	if err != nil {
		return nil, err
	}
	for _, d := range stored {
		name, exists := names[d.DeviceID]
		if k.isLocal(userID) && !exists {
			continue
		}
		var mine []storage.KeySignature
		for _, sig := range sigs {
			if sig.TargetKeyID == d.DeviceID {
				mine = append(mine, sig)
			}
		}
		if !withNames {
			name = ""
		}
		merged, err := withSignatures(d.JSON, mine, name)
		if err != nil {
			return nil, err
		}
		out[d.DeviceID] = merged
	}
	return out, nil
}

// crossSigningKeys reads a user's cross-signing keys as requester is
// allowed to see them.
//
// The master key carries the signatures of the user's own devices, which
// everybody may see, and the signature of requester's user-signing key if
// they made one: who a user has verified is between them and their own
// devices. The user-signing key is only ever shown to its owner.
func (k *API) crossSigningKeys(ctx context.Context, userID, requester string) (master, selfSigning, userSigning json.RawMessage, err error) {
	keys, err := k.store.CrossSigningKeys(ctx, userID)
	if err != nil || len(keys) == 0 {
		return nil, nil, nil, err
	}
	sigs, err := k.store.Signatures(ctx, userID)
	if err != nil {
		return nil, nil, nil, err
	}

	read := func(purpose string) (json.RawMessage, error) {
		raw, ok := keys[purpose]
		if !ok {
			return nil, nil
		}
		var key crossSigningKeyObject
		if err := json.Unmarshal(raw, &key); err != nil {
			return nil, err
		}
		pub, _ := key.ed25519Key()
		var visible []storage.KeySignature
		for _, sig := range sigs {
			if sig.TargetKeyID != pub {
				continue
			}
			if sig.SignerUserID == userID || (requester != "" && sig.SignerUserID == requester) {
				visible = append(visible, sig)
			}
		}
		return withSignatures(raw, visible, "")
	}

	if master, err = read(storage.CrossSigningMaster); err != nil {
		return nil, nil, nil, err
	}
	if selfSigning, err = read(storage.CrossSigningSelfSigning); err != nil {
		return nil, nil, nil, err
	}
	if requester == userID {
		if userSigning, err = read(storage.CrossSigningUserSigning); err != nil {
			return nil, nil, nil, err
		}
	}
	return master, selfSigning, userSigning, nil
}

// --- /keys/query ----------------------------------------------------------------

// QueryKeys answers POST /keys/query for requester: local users from the
// store, remote ones from the cache or by asking their servers, all at
// once.
func (k *API) QueryKeys(ctx context.Context, requester string, req QueryRequest) (QueryResponse, error) {
	resp := QueryResponse{
		DeviceKeys:      map[string]map[string]json.RawMessage{},
		Failures:        map[string]any{},
		MasterKeys:      map[string]json.RawMessage{},
		SelfSigningKeys: map[string]json.RawMessage{},
		UserSigningKeys: map[string]json.RawMessage{},
	}

	remote := map[string]map[string][]string{}
	for userID, deviceIDs := range req.DeviceKeys {
		server, err := id.ServerNameOf(userID)
		if err != nil {
			continue
		}
		if server != k.serverName {
			if remote[server] == nil {
				remote[server] = map[string][]string{}
			}
			remote[server][userID] = deviceIDs
			continue
		}

		devices, err := k.deviceKeys(ctx, userID, deviceIDs, true)
		if err != nil {
			return QueryResponse{}, err
		}
		resp.DeviceKeys[userID] = devices

		master, selfSigning, userSigning, err := k.crossSigningKeys(ctx, userID, requester)
		if err != nil {
			return QueryResponse{}, err
		}
		setIfPresent(resp.MasterKeys, userID, master)
		setIfPresent(resp.SelfSigningKeys, userID, selfSigning)
		setIfPresent(resp.UserSigningKeys, userID, userSigning)
	}

	k.queryRemote(ctx, requester, remote, remoteTimeout(req.Timeout), &resp)
	return resp, nil
}

// queryRemote answers for other servers' users: from the cache for those
// who share a room with somebody here, by asking their server for the
// rest. A copy that is missing or stale is fetched whole and cached, so
// the next query for the same user does not leave this server.
//
// The servers are asked in parallel. One that does not answer in time is a
// failure, not an error: the keys of everybody else are still worth
// returning.
func (k *API) queryRemote(ctx context.Context, requester string, byServer map[string]map[string][]string, timeout time.Duration, resp *QueryResponse) {
	if len(byServer) == 0 {
		return
	}
	ctx, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()

	var (
		mu sync.Mutex
		wg sync.WaitGroup
	)
	fail := func(server string, err error) {
		mu.Lock()
		defer mu.Unlock()
		resp.Failures[server] = remoteFailure(err)
	}

	for server, users := range byServer {
		passOn := map[string][]string{}
		var fetch []string
		for _, userID := range sortedKeys(users) {
			if !k.cacheable(ctx, userID) {
				// A user nobody here shares a room with any more may still
				// have a copy from when somebody did.
				k.forget(ctx, userID)
				passOn[userID] = users[userID]
				continue
			}
			if keys, ok := k.freshCopy(ctx, userID); ok {
				mu.Lock()
				k.serveRemote(ctx, requester, keys, users[userID], resp)
				mu.Unlock()
				continue
			}
			fetch = append(fetch, userID)
		}

		if len(passOn) == 0 && len(fetch) == 0 {
			continue
		}
		if k.remote == nil {
			fail(server, errors.New("federation is not available on this server"))
			continue
		}

		for _, userID := range fetch {
			wg.Go(func() {
				keys, err := k.fetchRemoteKeys(ctx, server, userID)
				if err != nil {
					fail(server, err)
					return
				}
				mu.Lock()
				defer mu.Unlock()
				k.serveRemote(ctx, requester, keys, users[userID], resp)
			})
		}
		if len(passOn) == 0 {
			continue
		}
		wg.Go(func() {
			answer, err := k.remote.QueryUserKeys(ctx, server, federation.UserKeysQueryRequest{DeviceKeys: passOn})
			if err != nil {
				fail(server, err)
				return
			}

			mu.Lock()
			defer mu.Unlock()
			for userID, devices := range answer.DeviceKeys {
				// A server answers for its own users and nobody else's.
				if _, asked := passOn[userID]; asked && devices != nil {
					resp.DeviceKeys[userID] = devices
				}
			}
			for userID, master := range answer.MasterKeys {
				if _, asked := passOn[userID]; asked {
					resp.MasterKeys[userID] = k.withRequesterSignature(ctx, userID, requester, master)
				}
			}
			for userID, selfSigning := range answer.SelfSigningKeys {
				if _, asked := passOn[userID]; asked {
					resp.SelfSigningKeys[userID] = selfSigning
				}
			}
		})
	}
	wg.Wait()
}

// withRequesterSignature adds requester's own signature of a remote user's
// master key, which is stored here and nowhere else.
func (k *API) withRequesterSignature(ctx context.Context, userID, requester string, master json.RawMessage) json.RawMessage {
	if requester == "" {
		return master
	}
	sigs, err := k.store.Signatures(ctx, userID)
	if err != nil || len(sigs) == 0 {
		return master
	}
	var key crossSigningKeyObject
	if json.Unmarshal(master, &key) != nil {
		return master
	}
	pub, _ := key.ed25519Key()
	var mine []storage.KeySignature
	for _, sig := range sigs {
		if sig.TargetKeyID == pub && sig.SignerUserID == requester {
			mine = append(mine, sig)
		}
	}
	merged, err := withSignatures(master, mine, "")
	if err != nil {
		return master
	}
	return merged
}

func setIfPresent(m map[string]json.RawMessage, key string, value json.RawMessage) {
	if len(value) > 0 {
		m[key] = value
	}
}

// QueryLocalKeys answers another server's POST /user/keys/query: the keys
// of this server's users, without device display names, which Warp keeps
// to itself.
func (k *API) QueryLocalKeys(ctx context.Context, req federation.UserKeysQueryRequest) (federation.UserKeysQueryResponse, error) {
	resp := federation.UserKeysQueryResponse{
		DeviceKeys:      map[string]map[string]json.RawMessage{},
		MasterKeys:      map[string]json.RawMessage{},
		SelfSigningKeys: map[string]json.RawMessage{},
	}
	for userID, deviceIDs := range req.DeviceKeys {
		if !k.isLocal(userID) {
			continue
		}
		devices, err := k.publishedDeviceKeys(ctx, userID, deviceIDs)
		if err != nil {
			return federation.UserKeysQueryResponse{}, err
		}
		resp.DeviceKeys[userID] = devices

		master, selfSigning, _, err := k.crossSigningKeys(ctx, userID, "")
		if err != nil {
			return federation.UserKeysQueryResponse{}, err
		}
		setIfPresent(resp.MasterKeys, userID, master)
		setIfPresent(resp.SelfSigningKeys, userID, selfSigning)
	}
	return resp, nil
}

// UserDevices answers another server's GET /user/devices/{userId}: a local
// user's whole device list, and its position in the device list stream,
// which is the stream ID the server's later updates count from.
func (k *API) UserDevices(ctx context.Context, userID string) (federation.UserDevicesResponse, error) {
	resp := federation.UserDevicesResponse{UserID: userID, Devices: []federation.UserDevice{}}
	if !k.isLocal(userID) {
		return resp, invalidParamf("%s is not a user of this server", userID)
	}
	pos, err := k.store.DeviceListPosition(ctx, userID)
	if err != nil {
		return resp, err
	}
	resp.StreamID = pos

	devices, err := k.publishedDeviceKeys(ctx, userID, nil)
	if err != nil {
		return resp, err
	}
	for _, deviceID := range sortedKeys(devices) {
		resp.Devices = append(resp.Devices, federation.UserDevice{DeviceID: deviceID, Keys: devices[deviceID]})
	}
	resp.MasterKey, resp.SelfSigningKey, _, err = k.crossSigningKeys(ctx, userID, "")
	return resp, err
}

// --- /keys/claim ----------------------------------------------------------------

// ClaimKeys answers POST /keys/claim.
func (k *API) ClaimKeys(ctx context.Context, req ClaimRequest) (ClaimResponse, error) {
	resp := ClaimResponse{
		OneTimeKeys: map[string]map[string]map[string]json.RawMessage{},
		Failures:    map[string]any{},
	}

	local := map[string]map[string]string{}
	remote := map[string]map[string]map[string]string{}
	for userID, devices := range req.OneTimeKeys {
		server, err := id.ServerNameOf(userID)
		switch {
		case err != nil:
			continue
		case server == k.serverName:
			local[userID] = devices
		default:
			if remote[server] == nil {
				remote[server] = map[string]map[string]string{}
			}
			remote[server][userID] = devices
		}
	}

	claimed, err := k.ClaimLocalKeys(ctx, local)
	if err != nil {
		return ClaimResponse{}, err
	}
	for userID, devices := range claimed {
		resp.OneTimeKeys[userID] = devices
	}

	k.claimRemote(ctx, remote, remoteTimeout(req.Timeout), &resp)
	return resp, nil
}

// ClaimLocalKeys claims keys of this server's users' devices, for a local
// client or another server. A device with no key of the algorithm asked
// for is left out, which is how the spec says "none left".
func (k *API) ClaimLocalKeys(ctx context.Context, want map[string]map[string]string) (map[string]map[string]map[string]json.RawMessage, error) {
	out := map[string]map[string]map[string]json.RawMessage{}
	for userID, devices := range want {
		if !k.isLocal(userID) {
			continue
		}
		for deviceID, algorithm := range devices {
			key, ok, err := k.store.ClaimKey(ctx, userID, deviceID, algorithm)
			if err != nil {
				return nil, err
			}
			if !ok {
				continue
			}
			if out[userID] == nil {
				out[userID] = map[string]map[string]json.RawMessage{}
			}
			out[userID][deviceID] = map[string]json.RawMessage{key.KeyID: key.JSON}
		}
	}
	return out, nil
}

// claimRemote claims keys from other servers in parallel.
func (k *API) claimRemote(ctx context.Context, byServer map[string]map[string]map[string]string, timeout time.Duration, resp *ClaimResponse) {
	if len(byServer) == 0 {
		return
	}
	ctx, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()

	var (
		mu sync.Mutex
		wg sync.WaitGroup
	)
	for server, users := range byServer {
		if k.remote == nil {
			resp.Failures[server] = remoteFailure(errors.New("federation is not available on this server"))
			continue
		}
		wg.Go(func() {
			answer, err := k.remote.ClaimUserKeys(ctx, server, federation.UserKeysClaimRequest{OneTimeKeys: users})

			mu.Lock()
			defer mu.Unlock()
			if err != nil {
				resp.Failures[server] = remoteFailure(err)
				return
			}
			for userID, devices := range answer.OneTimeKeys {
				if _, asked := users[userID]; asked {
					resp.OneTimeKeys[userID] = devices
				}
			}
		})
	}
	wg.Wait()
}

// sortedKeys is a map's keys in order.
func sortedKeys[V any](m map[string]V) []string {
	out := make([]string, 0, len(m))
	for key := range m {
		out = append(out, key)
	}
	slices.Sort(out)
	return out
}