From f5a39a12da6720418e353163ce34b3691774de17 Mon Sep 17 00:00:00 2001
From: Kegan Dougal <kegan@matrix.org>
Date: Tue, 21 Jul 2020 10:55:57 +0100
Subject: [PATCH] Add API shape for claiming keys

---
 clientapi/routing/keys.go    | 37 ++++++++++++++++++++++++++++++++++++
 clientapi/routing/routing.go |  5 +++++
 keyserver/api/api.go         |  9 +++++++++
 3 files changed, 51 insertions(+)

diff --git a/clientapi/routing/keys.go b/clientapi/routing/keys.go
index 5f7bfb187..ba03a352f 100644
--- a/clientapi/routing/keys.go
+++ b/clientapi/routing/keys.go
@@ -117,3 +117,40 @@ func QueryKeys(req *http.Request, keyAPI api.KeyInternalAPI) util.JSONResponse {
 		},
 	}
 }
+
+type claimKeysRequest struct {
+	TimeoutMS int `json:"timeout"`
+	//  The keys to be claimed. A map from user ID, to a map from device ID to algorithm name.
+	OneTimeKeys map[string]map[string]string `json:"one_time_keys"`
+}
+
+func (r *claimKeysRequest) GetTimeout() time.Duration {
+	if r.TimeoutMS == 0 {
+		return 10 * time.Second
+	}
+	return time.Duration(r.TimeoutMS) * time.Millisecond
+}
+
+func ClaimKeys(req *http.Request, keyAPI api.KeyInternalAPI) util.JSONResponse {
+	var r claimKeysRequest
+	resErr := httputil.UnmarshalJSONRequest(req, &r)
+	if resErr != nil {
+		return *resErr
+	}
+	claimRes := api.PerformClaimKeysResponse{}
+	keyAPI.PerformClaimKeys(req.Context(), &api.PerformClaimKeysRequest{
+		OneTimeKeys: r.OneTimeKeys,
+		Timeout:     r.GetTimeout(),
+	}, &claimRes)
+	if claimRes.Error != nil {
+		util.GetLogger(req.Context()).WithError(claimRes.Error).Error("failed to PerformClaimKeys")
+		return jsonerror.InternalServerError()
+	}
+	return util.JSONResponse{
+		Code: 200,
+		JSON: map[string]interface{}{
+			"one_time_keys": claimRes.OneTimeKeys,
+			"failures":      claimRes.Failures,
+		},
+	}
+}
diff --git a/clientapi/routing/routing.go b/clientapi/routing/routing.go
index 492b7e253..c9ed5ea5c 100644
--- a/clientapi/routing/routing.go
+++ b/clientapi/routing/routing.go
@@ -714,4 +714,9 @@ func Setup(
 			return QueryKeys(req, keyAPI)
 		}),
 	).Methods(http.MethodPost, http.MethodOptions)
+	r0mux.Handle("/keys/claim",
+		httputil.MakeAuthAPI("keys_claim", userAPI, func(req *http.Request, device *userapi.Device) util.JSONResponse {
+			return ClaimKeys(req, keyAPI)
+		}),
+	).Methods(http.MethodPost, http.MethodOptions)
 }
diff --git a/keyserver/api/api.go b/keyserver/api/api.go
index 0f6cb7979..d42fb60cf 100644
--- a/keyserver/api/api.go
+++ b/keyserver/api/api.go
@@ -23,6 +23,7 @@ import (
 
 type KeyInternalAPI interface {
 	PerformUploadKeys(ctx context.Context, req *PerformUploadKeysRequest, res *PerformUploadKeysResponse)
+	// PerformClaimKeys claims one-time keys for use in pre-key messages
 	PerformClaimKeys(ctx context.Context, req *PerformClaimKeysRequest, res *PerformClaimKeysResponse)
 	QueryKeys(ctx context.Context, req *QueryKeysRequest, res *QueryKeysResponse)
 }
@@ -102,9 +103,17 @@ func (r *PerformUploadKeysResponse) KeyError(userID, deviceID string, err *KeyEr
 }
 
 type PerformClaimKeysRequest struct {
+	// Map of user_id to device_id to algorithm name
+	OneTimeKeys map[string]map[string]string
+	Timeout     time.Duration
 }
 
 type PerformClaimKeysResponse struct {
+	// Map of user_id to device_id to algorithm:key_id to key JSON
+	OneTimeKeys map[string]map[string]map[string]json.RawMessage
+	// Map of remote server domain to error JSON
+	Failures map[string]interface{}
+	// Set if there was a fatal error processing this action
 	Error *KeyError
 }