Add QueryKeyChanges (#1228)

Hook some things up to call it as well.

4 files changed, 57 insertions, 0 deletions

diff --git a/keyserver/api/api.go b/keyserver/api/api.go
index d42fb60c..406a252d 100644
--- a/keyserver/api/api.go
+++ b/keyserver/api/api.go
@@ -26,6 +26,7 @@ type KeyInternalAPI interface {
 	// 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)
+	QueryKeyChanges(ctx context.Context, req *QueryKeyChangesRequest, res *QueryKeyChangesResponse)
 }
 
 // KeyError is returned if there was a problem performing/querying the server
@@ -131,3 +132,19 @@ type QueryKeysResponse struct {
 	// Set if there was a fatal error processing this query
 	Error *KeyError
 }
+
+type QueryKeyChangesRequest struct {
+	// The partition which had key events sent to
+	Partition int32
+	// The offset of the last received key event, or sarama.OffsetOldest if this is from the beginning
+	Offset int64
+}
+
+type QueryKeyChangesResponse struct {
+	// The set of users who have had their keys change.
+	UserIDs []string
+	// The latest offset represented in this response.
+	Offset int64
+	// Set if there was a problem handling the request.
+	Error *KeyError
+}
diff --git a/keyserver/internal/internal.go b/keyserver/internal/internal.go
index d3a6d4ba..240a5640 100644
--- a/keyserver/internal/internal.go
+++ b/keyserver/internal/internal.go
@@ -40,6 +40,17 @@ type KeyInternalAPI struct {
 	Producer   *producers.KeyChange
 }
 
+func (a *KeyInternalAPI) QueryKeyChanges(ctx context.Context, req *api.QueryKeyChangesRequest, res *api.QueryKeyChangesResponse) {
+	userIDs, latest, err := a.DB.KeyChanges(ctx, req.Partition, req.Offset)
+	if err != nil {
+		res.Error = &api.KeyError{
+			Err: err.Error(),
+		}
+	}
+	res.Offset = latest
+	res.UserIDs = userIDs
+}
+
 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)
diff --git a/keyserver/inthttp/client.go b/keyserver/inthttp/client.go
index 4c0f1e53..cd9cf70d 100644
--- a/keyserver/inthttp/client.go
+++ b/keyserver/inthttp/client.go
@@ -29,6 +29,7 @@ const (
 	PerformUploadKeysPath = "/keyserver/performUploadKeys"
 	PerformClaimKeysPath  = "/keyserver/performClaimKeys"
 	QueryKeysPath         = "/keyserver/queryKeys"
+	QueryKeyChangesPath   = "/keyserver/queryKeyChanges"
 )
 
 // NewKeyServerClient creates a KeyInternalAPI implemented by talking to a HTTP POST API.
@@ -101,3 +102,20 @@ func (h *httpKeyInternalAPI) QueryKeys(
 		}
 	}
 }
+
+func (h *httpKeyInternalAPI) QueryKeyChanges(
+	ctx context.Context,
+	request *api.QueryKeyChangesRequest,
+	response *api.QueryKeyChangesResponse,
+) {
+	span, ctx := opentracing.StartSpanFromContext(ctx, "QueryKeyChanges")
+	defer span.Finish()
+
+	apiURL := h.apiURL + QueryKeyChangesPath
+	err := httputil.PostJSON(ctx, span, h.httpClient, apiURL, request, response)
+	if err != nil {
+		response.Error = &api.KeyError{
+			Err: err.Error(),
+		}
+	}
+}
diff --git a/keyserver/inthttp/server.go b/keyserver/inthttp/server.go
index ec78b613..f3d2882c 100644
--- a/keyserver/inthttp/server.go
+++ b/keyserver/inthttp/server.go
@@ -58,4 +58,15 @@ func AddRoutes(internalAPIMux *mux.Router, s api.KeyInternalAPI) {
 			return util.JSONResponse{Code: http.StatusOK, JSON: &response}
 		}),
 	)
+	internalAPIMux.Handle(QueryKeyChangesPath,
+		httputil.MakeInternalAPI("queryKeyChanges", func(req *http.Request) util.JSONResponse {
+			request := api.QueryKeyChangesRequest{}
+			response := api.QueryKeyChangesResponse{}
+			if err := json.NewDecoder(req.Body).Decode(&request); err != nil {
+				return util.MessageResponse(http.StatusBadRequest, err.Error())
+			}
+			s.QueryKeyChanges(req.Context(), &request, &response)
+			return util.JSONResponse{Code: http.StatusOK, JSON: &response}
+		}),
+	)
 }