Implement logic for key uploads (#1197)

* begin work on storing keys

* Finish rough impl of the internal key API

* Linting

keyserver/api/api.go

@@ -14,7 +14,10 @@
 
 package api
 
-import "context"
+import (
+	"context"
+	"encoding/json"
+)
 
 type KeyInternalAPI interface {
 	PerformUploadKeys(ctx context.Context, req *PerformUploadKeysRequest, res *PerformUploadKeysResponse)
@@ -27,11 +30,62 @@ type KeyError struct {
 	Error string
 }
 
-type PerformUploadKeysRequest struct {
+// DeviceKeys represents a set of device keys for a single device
+// https://matrix.org/docs/spec/client_server/r0.6.1#post-matrix-client-r0-keys-upload
+type DeviceKeys struct {
+	// The user who owns this device
+	UserID string
+	// The device ID of this device
+	DeviceID string
+	// The raw device key JSON
+	KeyJSON []byte
 }
 
+// OneTimeKeys represents a set of one-time keys for a single device
+// https://matrix.org/docs/spec/client_server/r0.6.1#post-matrix-client-r0-keys-upload
+type OneTimeKeys struct {
+	// The user who owns this device
+	UserID string
+	// The device ID of this device
+	DeviceID string
+	// A map of algorithm:key_id => key JSON
+	KeyJSON map[string]json.RawMessage
+}
+
+// OneTimeKeysCount represents the counts of one-time keys for a single device
+type OneTimeKeysCount struct {
+	// The user who owns this device
+	UserID string
+	// The device ID of this device
+	DeviceID string
+	// algorithm to count e.g:
+	// {
+	//   "curve25519": 10,
+	//   "signed_curve25519": 20
+	// }
+	KeyCount map[string]int
+}
+
+// PerformUploadKeysRequest is the request to PerformUploadKeys
+type PerformUploadKeysRequest struct {
+	DeviceKeys  []DeviceKeys
+	OneTimeKeys []OneTimeKeys
+}
+
+// PerformUploadKeysResponse is the response to PerformUploadKeys
 type PerformUploadKeysResponse struct {
 	Error *KeyError
+	// A map of user_id -> device_id -> Error for tracking failures.
+	KeyErrors        map[string]map[string]*KeyError
+	OneTimeKeyCounts []OneTimeKeysCount
+}
+
+// KeyError sets a key error field on KeyErrors
+func (r *PerformUploadKeysResponse) KeyError(userID, deviceID string, err *KeyError) {
+	if r.KeyErrors[userID] == nil {
+		r.KeyErrors[userID] = make(map[string]*KeyError)
+	}
+	r.KeyErrors[userID][deviceID] = err
 }
 
 type PerformClaimKeysRequest struct {

keyserver/internal/internal.go

@@ -1,15 +1,37 @@
+// Copyright 2020 The Matrix.org Foundation C.I.C.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
 package internal
 
 import (
+	"bytes"
 	"context"
+	"fmt"
 
 	"github.com/matrix-org/dendrite/keyserver/api"
+	"github.com/matrix-org/dendrite/keyserver/storage"
+	"github.com/tidwall/gjson"
 )
 
-type KeyInternalAPI struct{}
+type KeyInternalAPI struct {
+	db storage.Database
+}
 
 func (a *KeyInternalAPI) PerformUploadKeys(ctx context.Context, req *api.PerformUploadKeysRequest, res *api.PerformUploadKeysResponse) {
-
+	res.KeyErrors = make(map[string]map[string]*api.KeyError)
+	a.uploadDeviceKeys(ctx, req, res)
+	a.uploadOneTimeKeys(ctx, req, res)
 }
 func (a *KeyInternalAPI) PerformClaimKeys(ctx context.Context, req *api.PerformClaimKeysRequest, res *api.PerformClaimKeysResponse) {
 
@@ -17,3 +39,84 @@ func (a *KeyInternalAPI) PerformClaimKeys(ctx context.Context, req *api.PerformC
 func (a *KeyInternalAPI) QueryKeys(ctx context.Context, req *api.QueryKeysRequest, res *api.QueryKeysResponse) {
 
 }
+
+func (a *KeyInternalAPI) uploadDeviceKeys(ctx context.Context, req *api.PerformUploadKeysRequest, res *api.PerformUploadKeysResponse) {
+	var keysToStore []api.DeviceKeys
+	// assert that the user ID / device ID are not lying for each key
+	for _, key := range req.DeviceKeys {
+		gotUserID := gjson.GetBytes(key.KeyJSON, "user_id").Str
+		gotDeviceID := gjson.GetBytes(key.KeyJSON, "device_id").Str
+		if gotUserID == key.UserID && gotDeviceID == key.DeviceID {
+			keysToStore = append(keysToStore, key)
+			continue
+		}
+
+		res.KeyError(key.UserID, key.DeviceID, &api.KeyError{
+			Error: fmt.Sprintf(
+				"user_id or device_id mismatch: users: %s - %s, devices: %s - %s",
+				gotUserID, key.UserID, gotDeviceID, key.DeviceID,
+			),
+		})
+	}
+	// get existing device keys so we can check for changes
+	existingKeys := make([]api.DeviceKeys, len(keysToStore))
+	for i := range keysToStore {
+		existingKeys[i] = api.DeviceKeys{
+			UserID:   keysToStore[i].UserID,
+			DeviceID: keysToStore[i].DeviceID,
+		}
+	}
+	if err := a.db.DeviceKeysJSON(ctx, existingKeys); err != nil {
+		res.Error = &api.KeyError{
+			Error: fmt.Sprintf("failed to query existing device keys: %s", err.Error()),
+		}
+		return
+	}
+	// store the device keys and emit changes
+	if err := a.db.StoreDeviceKeys(ctx, keysToStore); err != nil {
+		res.Error = &api.KeyError{
+			Error: fmt.Sprintf("failed to store device keys: %s", err.Error()),
+		}
+		return
+	}
+	a.emitDeviceKeyChanges(existingKeys, keysToStore)
+}
+
+func (a *KeyInternalAPI) uploadOneTimeKeys(ctx context.Context, req *api.PerformUploadKeysRequest, res *api.PerformUploadKeysResponse) {
+	for _, key := range req.OneTimeKeys {
+		// grab existing keys based on (user/device/algorithm/key ID)
+		keyIDsWithAlgorithms := make([]string, len(key.KeyJSON))
+		i := 0
+		for keyIDWithAlgo := range key.KeyJSON {
+			keyIDsWithAlgorithms[i] = keyIDWithAlgo
+			i++
+		}
+		existingKeys, err := a.db.ExistingOneTimeKeys(ctx, key.UserID, key.DeviceID, keyIDsWithAlgorithms)
+		if err != nil {
+			res.KeyError(key.UserID, key.DeviceID, &api.KeyError{
+				Error: "failed to query existing one-time keys: " + err.Error(),
+			})
+			continue
+		}
+		for keyIDWithAlgo := range existingKeys {
+			// if keys exist and the JSON doesn't match, error out as the key already exists
+			if !bytes.Equal(existingKeys[keyIDWithAlgo], key.KeyJSON[keyIDWithAlgo]) {
+				res.KeyError(key.UserID, key.DeviceID, &api.KeyError{
+					Error: fmt.Sprintf("%s device %s: algorithm / key ID %s one-time key already exists", key.UserID, key.DeviceID, keyIDWithAlgo),
+				})
+				continue
+			}
+		}
+		// store one-time keys
+		if err := a.db.StoreOneTimeKeys(ctx, key); err != nil {
+			res.KeyError(key.UserID, key.DeviceID, &api.KeyError{
+				Error: fmt.Sprintf("%s device %s : failed to store one-time keys: %s", key.UserID, key.DeviceID, err.Error()),
+			})
+		}
+	}
+
+}
+
+func (a *KeyInternalAPI) emitDeviceKeyChanges(existing, new []api.DeviceKeys) {
+	// TODO
+}

keyserver/storage/interface.go

@@ -0,0 +1,38 @@
+// Copyright 2020 The Matrix.org Foundation C.I.C.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package storage
+
+import (
+	"context"
+	"encoding/json"
+
+	"github.com/matrix-org/dendrite/keyserver/api"
+)
+
+type Database interface {
+	// ExistingOneTimeKeys returns a map of keyIDWithAlgorithm to key JSON for the given parameters. If no keys exist with this combination
+	// of user/device/key/algorithm 4-uple then it is omitted from the map. Returns an error when failing to communicate with the database.
+	ExistingOneTimeKeys(ctx context.Context, userID, deviceID string, keyIDsWithAlgorithms []string) (map[string]json.RawMessage, error)
+
+	// StoreOneTimeKeys persists the given one-time keys.
+	StoreOneTimeKeys(ctx context.Context, keys api.OneTimeKeys) error
+
+	// DeviceKeysJSON populates the KeyJSON for the given keys. If any proided `keys` have a `KeyJSON` already then it will be replaced.
+	DeviceKeysJSON(ctx context.Context, keys []api.DeviceKeys) error
+
+	// StoreDeviceKeys persists the given keys. Keys with the same user ID and device ID will be replaced.
+	// Returns an error if there was a problem storing the keys.
+	StoreDeviceKeys(ctx context.Context, keys []api.DeviceKeys) error
+}