Handle inbound /keys/claim and /keys/query requests

cmd/dendrite-federation-api-server/main.go

@@ -30,10 +30,11 @@ func main() {
 	keyRing := serverKeyAPI.KeyRing()
 	fsAPI := base.FederationSenderHTTPClient()
 	rsAPI := base.RoomserverHTTPClient()
+	keyAPI := base.KeyServerHTTPClient()
 
 	federationapi.AddPublicRoutes(
 		base.PublicAPIMux, base.Cfg, userAPI, federation, keyRing,
-		rsAPI, fsAPI, base.EDUServerClient(), base.CurrentStateAPIClient(),
+		rsAPI, fsAPI, base.EDUServerClient(), base.CurrentStateAPIClient(), keyAPI,
 	)
 
 	base.SetupAndServeHTTP(string(base.Cfg.Bind.FederationAPI), string(base.Cfg.Listen.FederationAPI))

federationapi/federationapi.go

@@ -20,6 +20,7 @@ import (
 	eduserverAPI "github.com/matrix-org/dendrite/eduserver/api"
 	federationSenderAPI "github.com/matrix-org/dendrite/federationsender/api"
 	"github.com/matrix-org/dendrite/internal/config"
+	keyserverAPI "github.com/matrix-org/dendrite/keyserver/api"
 	roomserverAPI "github.com/matrix-org/dendrite/roomserver/api"
 	userapi "github.com/matrix-org/dendrite/userapi/api"
 
@@ -38,11 +39,12 @@ func AddPublicRoutes(
 	federationSenderAPI federationSenderAPI.FederationSenderInternalAPI,
 	eduAPI eduserverAPI.EDUServerInputAPI,
 	stateAPI currentstateAPI.CurrentStateInternalAPI,
+	keyAPI keyserverAPI.KeyInternalAPI,
 ) {
 
 	routing.Setup(
 		router, cfg, rsAPI,
 		eduAPI, federationSenderAPI, keyRing,
-		federation, userAPI, stateAPI,
+		federation, userAPI, stateAPI, keyAPI,
 	)
 }

federationapi/federationapi_test.go

@@ -31,7 +31,7 @@ func TestRoomsV3URLEscapeDoNot404(t *testing.T) {
 	fsAPI := base.FederationSenderHTTPClient()
 	// TODO: This is pretty fragile, as if anything calls anything on these nils this test will break.
 	// Unfortunately, it makes little sense to instantiate these dependencies when we just want to test routing.
-	federationapi.AddPublicRoutes(base.PublicAPIMux, cfg, nil, nil, keyRing, nil, fsAPI, nil, nil)
+	federationapi.AddPublicRoutes(base.PublicAPIMux, cfg, nil, nil, keyRing, nil, fsAPI, nil, nil, nil)
 	httputil.SetupHTTPAPI(
 		base.BaseMux,
 		base.PublicAPIMux,

federationapi/routing/keys.go

@@ -19,12 +19,106 @@ import (
 	"net/http"
 	"time"
 
+	"github.com/matrix-org/dendrite/clientapi/jsonerror"
 	"github.com/matrix-org/dendrite/internal/config"
+	"github.com/matrix-org/dendrite/keyserver/api"
 	"github.com/matrix-org/gomatrixserverlib"
 	"github.com/matrix-org/util"
 	"golang.org/x/crypto/ed25519"
 )
 
+type queryKeysRequest struct {
+	DeviceKeys map[string][]string `json:"device_keys"`
+}
+
+// QueryDeviceKeys returns device keys for users on this server.
+// https://matrix.org/docs/spec/server_server/latest#post-matrix-federation-v1-user-keys-query
+func QueryDeviceKeys(
+	httpReq *http.Request, request *gomatrixserverlib.FederationRequest, keyAPI api.KeyInternalAPI, thisServer gomatrixserverlib.ServerName,
+) util.JSONResponse {
+	var qkr queryKeysRequest
+	err := json.Unmarshal(request.Content(), &qkr)
+	if err != nil {
+		return util.JSONResponse{
+			Code: http.StatusBadRequest,
+			JSON: jsonerror.BadJSON("The request body could not be decoded into valid JSON. " + err.Error()),
+		}
+	}
+	// make sure we only query users on our domain
+	for userID := range qkr.DeviceKeys {
+		_, serverName, err := gomatrixserverlib.SplitID('@', userID)
+		if err != nil {
+			delete(qkr.DeviceKeys, userID)
+			continue // ignore invalid users
+		}
+		if serverName != thisServer {
+			delete(qkr.DeviceKeys, userID)
+			continue
+		}
+	}
+
+	var queryRes api.QueryKeysResponse
+	keyAPI.QueryKeys(httpReq.Context(), &api.QueryKeysRequest{
+		UserToDevices: qkr.DeviceKeys,
+	}, &queryRes)
+	if queryRes.Error != nil {
+		util.GetLogger(httpReq.Context()).WithError(queryRes.Error).Error("Failed to QueryKeys")
+		return jsonerror.InternalServerError()
+	}
+	return util.JSONResponse{
+		Code: 200,
+		JSON: struct {
+			DeviceKeys interface{} `json:"device_keys"`
+		}{queryRes.DeviceKeys},
+	}
+}
+
+type claimOTKsRequest struct {
+	OneTimeKeys map[string]map[string]string `json:"one_time_keys"`
+}
+
+// ClaimOneTimeKeys claims OTKs for users on this server.
+// https://matrix.org/docs/spec/server_server/latest#post-matrix-federation-v1-user-keys-claim
+func ClaimOneTimeKeys(
+	httpReq *http.Request, request *gomatrixserverlib.FederationRequest, keyAPI api.KeyInternalAPI, thisServer gomatrixserverlib.ServerName,
+) util.JSONResponse {
+	var cor claimOTKsRequest
+	err := json.Unmarshal(request.Content(), &cor)
+	if err != nil {
+		return util.JSONResponse{
+			Code: http.StatusBadRequest,
+			JSON: jsonerror.BadJSON("The request body could not be decoded into valid JSON. " + err.Error()),
+		}
+	}
+	// make sure we only claim users on our domain
+	for userID := range cor.OneTimeKeys {
+		_, serverName, err := gomatrixserverlib.SplitID('@', userID)
+		if err != nil {
+			delete(cor.OneTimeKeys, userID)
+			continue // ignore invalid users
+		}
+		if serverName != thisServer {
+			delete(cor.OneTimeKeys, userID)
+			continue
+		}
+	}
+
+	var claimRes api.PerformClaimKeysResponse
+	keyAPI.PerformClaimKeys(httpReq.Context(), &api.PerformClaimKeysRequest{
+		OneTimeKeys: cor.OneTimeKeys,
+	}, &claimRes)
+	if claimRes.Error != nil {
+		util.GetLogger(httpReq.Context()).WithError(claimRes.Error).Error("Failed to PerformClaimKeys")
+		return jsonerror.InternalServerError()
+	}
+	return util.JSONResponse{
+		Code: 200,
+		JSON: struct {
+			OneTimeKeys interface{} `json:"one_time_keys"`
+		}{claimRes.OneTimeKeys},
+	}
+}
+
 // LocalKeys returns the local keys for the server.
 // See https://matrix.org/docs/spec/server_server/unstable.html#publishing-keys
 func LocalKeys(cfg *config.Dendrite) util.JSONResponse {

federationapi/routing/routing.go

@@ -24,6 +24,7 @@ import (
 	federationSenderAPI "github.com/matrix-org/dendrite/federationsender/api"
 	"github.com/matrix-org/dendrite/internal/config"
 	"github.com/matrix-org/dendrite/internal/httputil"
+	keyserverAPI "github.com/matrix-org/dendrite/keyserver/api"
 	roomserverAPI "github.com/matrix-org/dendrite/roomserver/api"
 	userapi "github.com/matrix-org/dendrite/userapi/api"
 	"github.com/matrix-org/gomatrixserverlib"
@@ -54,6 +55,7 @@ func Setup(
 	federation *gomatrixserverlib.FederationClient,
 	userAPI userapi.UserInternalAPI,
 	stateAPI currentstateAPI.CurrentStateInternalAPI,
+	keyAPI keyserverAPI.KeyInternalAPI,
 ) {
 	v2keysmux := publicAPIMux.PathPrefix(pathPrefixV2Keys).Subrouter()
 	v1fedmux := publicAPIMux.PathPrefix(pathPrefixV1Federation).Subrouter()
@@ -299,4 +301,18 @@ func Setup(
 			return GetPostPublicRooms(req, rsAPI, stateAPI)
 		}),
 	).Methods(http.MethodGet)
+
+	v1fedmux.Handle("/user/keys/claim", httputil.MakeFedAPI(
+		"federation_keys_claim", cfg.Matrix.ServerName, keys, wakeup,
+		func(httpReq *http.Request, request *gomatrixserverlib.FederationRequest, vars map[string]string) util.JSONResponse {
+			return ClaimOneTimeKeys(httpReq, request, keyAPI, cfg.Matrix.ServerName)
+		},
+	)).Methods(http.MethodPost)
+
+	v1fedmux.Handle("/user/keys/query", httputil.MakeFedAPI(
+		"federation_keys_query", cfg.Matrix.ServerName, keys, wakeup,
+		func(httpReq *http.Request, request *gomatrixserverlib.FederationRequest, vars map[string]string) util.JSONResponse {
+			return QueryDeviceKeys(httpReq, request, keyAPI, cfg.Matrix.ServerName)
+		},
+	)).Methods(http.MethodPost)
 }

internal/setup/monolith.go

@@ -73,7 +73,7 @@ func (m *Monolith) AddAllPublicRoutes(publicMux *mux.Router) {
 	federationapi.AddPublicRoutes(
 		publicMux, m.Config, m.UserAPI, m.FedClient,
 		m.KeyRing, m.RoomserverAPI, m.FederationSenderAPI,
-		m.EDUInternalAPI, m.StateAPI,
+		m.EDUInternalAPI, m.StateAPI, m.KeyAPI,
 	)
 	mediaapi.AddPublicRoutes(publicMux, m.Config, m.UserAPI, m.Client)
 	syncapi.AddPublicRoutes(

keyserver/internal/internal.go

@@ -25,6 +25,7 @@ import (
 	"github.com/matrix-org/dendrite/keyserver/api"
 	"github.com/matrix-org/dendrite/keyserver/storage"
 	"github.com/matrix-org/gomatrixserverlib"
+	"github.com/matrix-org/util"
 	"github.com/tidwall/gjson"
 	"github.com/tidwall/sjson"
 )
@@ -66,11 +67,25 @@ func (a *KeyInternalAPI) PerformClaimKeys(ctx context.Context, req *api.PerformC
 				Err: fmt.Sprintf("failed to ClaimKeys locally: %s", err),
 			}
 		}
-		mergeInto(res.OneTimeKeys, keys)
+		util.GetLogger(ctx).WithField("keys_claimed", len(keys)).WithField("num_users", len(local)).Info("Claimed local keys")
+		for _, key := range keys {
+			_, ok := res.OneTimeKeys[key.UserID]
+			if !ok {
+				res.OneTimeKeys[key.UserID] = make(map[string]map[string]json.RawMessage)
+			}
+			_, ok = res.OneTimeKeys[key.UserID][key.DeviceID]
+			if !ok {
+				res.OneTimeKeys[key.UserID][key.DeviceID] = make(map[string]json.RawMessage)
+			}
+			for keyID, keyJSON := range key.KeyJSON {
+				res.OneTimeKeys[key.UserID][key.DeviceID][keyID] = keyJSON
+			}
+		}
 		delete(domainToDeviceKeys, string(a.ThisServer))
 	}
-	// claim remote keys
+	if len(domainToDeviceKeys) > 0 {
 		a.claimRemoteKeys(ctx, req.Timeout, res, domainToDeviceKeys)
+	}
 }
 
 func (a *KeyInternalAPI) claimRemoteKeys(
@@ -82,6 +97,7 @@ func (a *KeyInternalAPI) claimRemoteKeys(
 	wg.Add(len(domainToDeviceKeys))
 	// mutex for failures
 	var failMu sync.Mutex
+	util.GetLogger(ctx).WithField("num_servers", len(domainToDeviceKeys)).Info("Claiming remote keys from servers")
 
 	// fan out
 	for d, k := range domainToDeviceKeys {
@@ -108,6 +124,7 @@ func (a *KeyInternalAPI) claimRemoteKeys(
 		close(resultCh)
 	}()
 
+	keysClaimed := 0
 	for result := range resultCh {
 		for userID, nest := range result.OneTimeKeys {
 			res.OneTimeKeys[userID] = make(map[string]map[string]json.RawMessage)
@@ -119,10 +136,12 @@ func (a *KeyInternalAPI) claimRemoteKeys(
 						continue
 					}
 					res.OneTimeKeys[userID][deviceID][keyIDWithAlgo] = keyJSON
+					keysClaimed++
 				}
 			}
 		}
 	}
+	util.GetLogger(ctx).WithField("num_keys", keysClaimed).Info("Claimed remote keys")
 }
 
 func (a *KeyInternalAPI) QueryKeys(ctx context.Context, req *api.QueryKeysRequest, res *api.QueryKeysResponse) {
@@ -298,19 +317,3 @@ func (a *KeyInternalAPI) uploadOneTimeKeys(ctx context.Context, req *api.Perform
 func (a *KeyInternalAPI) emitDeviceKeyChanges(existing, new []api.DeviceKeys) {
 	// TODO
 }
-
-func mergeInto(dst map[string]map[string]map[string]json.RawMessage, src []api.OneTimeKeys) {
-	for _, key := range src {
-		_, ok := dst[key.UserID]
-		if !ok {
-			dst[key.UserID] = make(map[string]map[string]json.RawMessage)
-		}
-		_, ok = dst[key.UserID][key.DeviceID]
-		if !ok {
-			dst[key.UserID][key.DeviceID] = make(map[string]json.RawMessage)
-		}
-		for keyID, keyJSON := range key.KeyJSON {
-			dst[key.UserID][key.DeviceID][keyID] = keyJSON
-		}
-	}
-}

sytest-whitelist

@@ -125,6 +125,7 @@ Can query device keys using POST
 Can query specific device keys using POST
 query for user with no keys returns empty key dict
 Can claim one time key using POST
+Can claim remote one time key using POST
 Can add account data
 Can add account data to room
 Can get account data without syncing