Implement key uploads (#1202)

* Add storage layer for postgres/sqlite

* Return OTK counts when inserting new keys

* Hook up the key DB and make a test pass

* Convert postgres queries to be sqlite queries

* Blacklist test due to requiring rejected events

* Unbreak tests

* Update blacklist

3 files changed, 64 insertions, 4 deletions

diff --git a/clientapi/clientapi.go b/clientapi/clientapi.go
index 029a73da..9ed285a8 100644
--- a/clientapi/clientapi.go
+++ b/clientapi/clientapi.go
@@ -26,6 +26,7 @@ import (
 	federationSenderAPI "github.com/matrix-org/dendrite/federationsender/api"
 	"github.com/matrix-org/dendrite/internal/config"
 	"github.com/matrix-org/dendrite/internal/transactions"
+	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/dendrite/userapi/storage/accounts"
@@ -48,6 +49,7 @@ func AddPublicRoutes(
 	transactionsCache *transactions.Cache,
 	fsAPI federationSenderAPI.FederationSenderInternalAPI,
 	userAPI userapi.UserInternalAPI,
+	keyAPI keyserverAPI.KeyInternalAPI,
 	extRoomsProvider api.ExtraPublicRoomsProvider,
 ) {
 	syncProducer := &producers.SyncAPIProducer{
@@ -58,6 +60,6 @@ func AddPublicRoutes(
 	routing.Setup(
 		router, cfg, eduInputAPI, rsAPI, asAPI,
 		accountsDB, deviceDB, userAPI, federation,
-		syncProducer, transactionsCache, fsAPI, stateAPI, extRoomsProvider,
+		syncProducer, transactionsCache, fsAPI, stateAPI, keyAPI, extRoomsProvider,
 	)
 }
diff --git a/clientapi/routing/keys.go b/clientapi/routing/keys.go
index 5c1a657f..8c5c5bbe 100644
--- a/clientapi/routing/keys.go
+++ b/clientapi/routing/keys.go
@@ -15,8 +15,13 @@
 package routing
 
 import (
+	"encoding/json"
 	"net/http"
 
+	"github.com/matrix-org/dendrite/clientapi/httputil"
+	"github.com/matrix-org/dendrite/clientapi/jsonerror"
+	"github.com/matrix-org/dendrite/keyserver/api"
+	userapi "github.com/matrix-org/dendrite/userapi/api"
 	"github.com/matrix-org/util"
 )
 
@@ -32,9 +37,55 @@ func QueryKeys(
 	}
 }
 
-func UploadKeys(req *http.Request) util.JSONResponse {
+type uploadKeysRequest struct {
+	DeviceKeys  json.RawMessage            `json:"device_keys"`
+	OneTimeKeys map[string]json.RawMessage `json:"one_time_keys"`
+}
+
+func UploadKeys(req *http.Request, keyAPI api.KeyInternalAPI, device *userapi.Device) util.JSONResponse {
+	var r uploadKeysRequest
+	resErr := httputil.UnmarshalJSONRequest(req, &r)
+	if resErr != nil {
+		return *resErr
+	}
+
+	uploadReq := &api.PerformUploadKeysRequest{}
+	if r.DeviceKeys != nil {
+		uploadReq.DeviceKeys = []api.DeviceKeys{
+			{
+				DeviceID: device.ID,
+				UserID:   device.UserID,
+				KeyJSON:  r.DeviceKeys,
+			},
+		}
+	}
+	if r.OneTimeKeys != nil {
+		uploadReq.OneTimeKeys = []api.OneTimeKeys{
+			{
+				DeviceID: device.ID,
+				UserID:   device.UserID,
+				KeyJSON:  r.OneTimeKeys,
+			},
+		}
+	}
+
+	var uploadRes api.PerformUploadKeysResponse
+	keyAPI.PerformUploadKeys(req.Context(), uploadReq, &uploadRes)
+	if uploadRes.Error != nil {
+		util.GetLogger(req.Context()).WithError(uploadRes.Error).Error("Failed to PerformUploadKeys")
+		return jsonerror.InternalServerError()
+	}
+	if len(uploadRes.KeyErrors) > 0 {
+		util.GetLogger(req.Context()).WithField("key_errors", uploadRes.KeyErrors).Error("Failed to upload one or more keys")
+		return util.JSONResponse{
+			Code: 400,
+			JSON: uploadRes.KeyErrors,
+		}
+	}
 	return util.JSONResponse{
 		Code: 200,
-		JSON: struct{}{},
+		JSON: struct {
+			OTKCounts interface{} `json:"one_time_key_counts"`
+		}{uploadRes.OneTimeKeyCounts[0].KeyCount},
 	}
 }
diff --git a/clientapi/routing/routing.go b/clientapi/routing/routing.go
index 965a46d2..4879ddaa 100644
--- a/clientapi/routing/routing.go
+++ b/clientapi/routing/routing.go
@@ -31,6 +31,7 @@ import (
 	"github.com/matrix-org/dendrite/internal/config"
 	"github.com/matrix-org/dendrite/internal/httputil"
 	"github.com/matrix-org/dendrite/internal/transactions"
+	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/dendrite/userapi/storage/accounts"
@@ -62,6 +63,7 @@ func Setup(
 	transactionsCache *transactions.Cache,
 	federationSender federationSenderAPI.FederationSenderInternalAPI,
 	stateAPI currentstateAPI.CurrentStateInternalAPI,
+	keyAPI keyserverAPI.KeyInternalAPI,
 	extRoomsProvider api.ExtraPublicRoomsProvider,
 ) {
 	userInteractiveAuth := auth.NewUserInteractive(accountDB.GetAccountByPassword, cfg)
@@ -705,7 +707,12 @@ func Setup(
 	// Supplying a device ID is deprecated.
 	r0mux.Handle("/keys/upload/{deviceID}",
 		httputil.MakeAuthAPI("keys_upload", userAPI, func(req *http.Request, device *userapi.Device) util.JSONResponse {
-			return UploadKeys(req)
+			return UploadKeys(req, keyAPI, device)
+		}),
+	).Methods(http.MethodPost, http.MethodOptions)
+	r0mux.Handle("/keys/upload",
+		httputil.MakeAuthAPI("keys_upload", userAPI, func(req *http.Request, device *userapi.Device) util.JSONResponse {
+			return UploadKeys(req, keyAPI, device)
 		}),
 	).Methods(http.MethodPost, http.MethodOptions)
 }