Add QueryKeyChanges (#1228)

Hook some things up to call it as well.

2 files changed, 62 insertions, 22 deletions

diff --git a/syncapi/consumers/keychange.go b/syncapi/consumers/keychange.go
index 4a1c7309..78aff601 100644
--- a/syncapi/consumers/keychange.go
+++ b/syncapi/consumers/keychange.go
@@ -26,16 +26,17 @@ import (
 	"github.com/matrix-org/dendrite/syncapi/storage"
 	"github.com/matrix-org/dendrite/syncapi/types"
 	"github.com/matrix-org/gomatrixserverlib"
+	"github.com/matrix-org/util"
 	log "github.com/sirupsen/logrus"
 )
 
 // OutputKeyChangeEventConsumer consumes events that originated in the key server.
 type OutputKeyChangeEventConsumer struct {
-	keyChangeConsumer *internal.ContinualConsumer
-	db                storage.Database
-	serverName        gomatrixserverlib.ServerName // our server name
-	currentStateAPI   currentstateAPI.CurrentStateInternalAPI
-	// keyAPI            api.KeyInternalAPI
+	keyChangeConsumer   *internal.ContinualConsumer
+	db                  storage.Database
+	serverName          gomatrixserverlib.ServerName // our server name
+	currentStateAPI     currentstateAPI.CurrentStateInternalAPI
+	keyAPI              api.KeyInternalAPI
 	partitionToOffset   map[int32]int64
 	partitionToOffsetMu sync.Mutex
 }
@@ -46,6 +47,7 @@ func NewOutputKeyChangeEventConsumer(
 	serverName gomatrixserverlib.ServerName,
 	topic string,
 	kafkaConsumer sarama.Consumer,
+	keyAPI api.KeyInternalAPI,
 	currentStateAPI currentstateAPI.CurrentStateInternalAPI,
 	store storage.Database,
 ) *OutputKeyChangeEventConsumer {
@@ -60,6 +62,7 @@ func NewOutputKeyChangeEventConsumer(
 		keyChangeConsumer:   &consumer,
 		db:                  store,
 		serverName:          serverName,
+		keyAPI:              keyAPI,
 		currentStateAPI:     currentStateAPI,
 		partitionToOffset:   make(map[int32]int64),
 		partitionToOffsetMu: sync.Mutex{},
@@ -115,21 +118,44 @@ func (s *OutputKeyChangeEventConsumer) onMessage(msg *sarama.ConsumerMessage) er
 // be already filled in with join/leave information.
 func (s *OutputKeyChangeEventConsumer) Catchup(
 	ctx context.Context, userID string, res *types.Response, tok types.StreamingToken,
-) (hasNew bool, err error) {
+) (newTok *types.StreamingToken, hasNew bool, err error) {
 	// Track users who we didn't track before but now do by virtue of sharing a room with them, or not.
 	newlyJoinedRooms := joinedRooms(res, userID)
 	newlyLeftRooms := leftRooms(res)
 	if len(newlyJoinedRooms) > 0 || len(newlyLeftRooms) > 0 {
 		changed, left, err := s.trackChangedUsers(ctx, userID, newlyJoinedRooms, newlyLeftRooms)
 		if err != nil {
-			return false, err
+			return nil, false, err
 		}
 		res.DeviceLists.Changed = changed
 		res.DeviceLists.Left = left
 		hasNew = len(changed) > 0 || len(left) > 0
 	}
 
-	// TODO: now also track users who we already share rooms with but who have updated their devices between the two tokens
+	// now also track users who we already share rooms with but who have updated their devices between the two tokens
+	// TODO: Extract partition/offset from sync token
+	var partition int32
+	var offset int64
+	var queryRes api.QueryKeyChangesResponse
+	s.keyAPI.QueryKeyChanges(ctx, &api.QueryKeyChangesRequest{
+		Partition: partition,
+		Offset:    offset,
+	}, &queryRes)
+	if queryRes.Error != nil {
+		// don't fail the catchup because we may have got useful information by tracking membership
+		util.GetLogger(ctx).WithError(queryRes.Error).Error("QueryKeyChanges failed")
+	} else {
+		// TODO: Make a new streaming token using the new offset
+		userSet := make(map[string]bool)
+		for _, userID := range res.DeviceLists.Changed {
+			userSet[userID] = true
+		}
+		for _, userID := range queryRes.UserIDs {
+			if !userSet[userID] {
+				res.DeviceLists.Changed = append(res.DeviceLists.Changed, userID)
+			}
+		}
+	}
 	return
 }
 
diff --git a/syncapi/consumers/keychange_test.go b/syncapi/consumers/keychange_test.go
index 7322e208..f8e96570 100644
--- a/syncapi/consumers/keychange_test.go
+++ b/syncapi/consumers/keychange_test.go
@@ -7,6 +7,7 @@ import (
 	"testing"
 
 	"github.com/matrix-org/dendrite/currentstateserver/api"
+	keyapi "github.com/matrix-org/dendrite/keyserver/api"
 	"github.com/matrix-org/dendrite/syncapi/types"
 	"github.com/matrix-org/gomatrixserverlib"
 )
@@ -15,6 +16,19 @@ var (
 	syncingUser = "@alice:localhost"
 )
 
+type mockKeyAPI struct{}
+
+func (k *mockKeyAPI) PerformUploadKeys(ctx context.Context, req *keyapi.PerformUploadKeysRequest, res *keyapi.PerformUploadKeysResponse) {
+}
+
+// PerformClaimKeys claims one-time keys for use in pre-key messages
+func (k *mockKeyAPI) PerformClaimKeys(ctx context.Context, req *keyapi.PerformClaimKeysRequest, res *keyapi.PerformClaimKeysResponse) {
+}
+func (k *mockKeyAPI) QueryKeys(ctx context.Context, req *keyapi.QueryKeysRequest, res *keyapi.QueryKeysResponse) {
+}
+func (k *mockKeyAPI) QueryKeyChanges(ctx context.Context, req *keyapi.QueryKeyChangesRequest, res *keyapi.QueryKeyChangesResponse) {
+}
+
 type mockCurrentStateAPI struct {
 	roomIDToJoinedMembers map[string][]string
 }
@@ -144,7 +158,7 @@ func leaveResponseWithRooms(syncResponse *types.Response, userID string, roomIDs
 func TestKeyChangeCatchupOnJoinShareNewUser(t *testing.T) {
 	newShareUser := "@bill:localhost"
 	newlyJoinedRoom := "!TestKeyChangeCatchupOnJoinShareNewUser:bar"
-	consumer := NewOutputKeyChangeEventConsumer(gomatrixserverlib.ServerName("localhost"), "some_topic", nil, &mockCurrentStateAPI{
+	consumer := NewOutputKeyChangeEventConsumer(gomatrixserverlib.ServerName("localhost"), "some_topic", nil, &mockKeyAPI{}, &mockCurrentStateAPI{
 		roomIDToJoinedMembers: map[string][]string{
 			newlyJoinedRoom: {syncingUser, newShareUser},
 			"!another:room": {syncingUser},
@@ -153,7 +167,7 @@ func TestKeyChangeCatchupOnJoinShareNewUser(t *testing.T) {
 	syncResponse := types.NewResponse()
 	syncResponse = joinResponseWithRooms(syncResponse, syncingUser, []string{newlyJoinedRoom})
 
-	hasNew, err := consumer.Catchup(context.Background(), syncingUser, syncResponse, types.NewStreamToken(0, 0))
+	_, hasNew, err := consumer.Catchup(context.Background(), syncingUser, syncResponse, types.NewStreamToken(0, 0))
 	if err != nil {
 		t.Fatalf("Catchup returned an error: %s", err)
 	}
@@ -167,7 +181,7 @@ func TestKeyChangeCatchupOnJoinShareNewUser(t *testing.T) {
 func TestKeyChangeCatchupOnLeaveShareLeftUser(t *testing.T) {
 	removeUser := "@bill:localhost"
 	newlyLeftRoom := "!TestKeyChangeCatchupOnLeaveShareLeftUser:bar"
-	consumer := NewOutputKeyChangeEventConsumer(gomatrixserverlib.ServerName("localhost"), "some_topic", nil, &mockCurrentStateAPI{
+	consumer := NewOutputKeyChangeEventConsumer(gomatrixserverlib.ServerName("localhost"), "some_topic", nil, &mockKeyAPI{}, &mockCurrentStateAPI{
 		roomIDToJoinedMembers: map[string][]string{
 			newlyLeftRoom:   {removeUser},
 			"!another:room": {syncingUser},
@@ -176,7 +190,7 @@ func TestKeyChangeCatchupOnLeaveShareLeftUser(t *testing.T) {
 	syncResponse := types.NewResponse()
 	syncResponse = leaveResponseWithRooms(syncResponse, syncingUser, []string{newlyLeftRoom})
 
-	hasNew, err := consumer.Catchup(context.Background(), syncingUser, syncResponse, types.NewStreamToken(0, 0))
+	_, hasNew, err := consumer.Catchup(context.Background(), syncingUser, syncResponse, types.NewStreamToken(0, 0))
 	if err != nil {
 		t.Fatalf("Catchup returned an error: %s", err)
 	}
@@ -190,7 +204,7 @@ func TestKeyChangeCatchupOnLeaveShareLeftUser(t *testing.T) {
 func TestKeyChangeCatchupOnJoinShareNoNewUsers(t *testing.T) {
 	existingUser := "@bob:localhost"
 	newlyJoinedRoom := "!TestKeyChangeCatchupOnJoinShareNoNewUsers:bar"
-	consumer := NewOutputKeyChangeEventConsumer(gomatrixserverlib.ServerName("localhost"), "some_topic", nil, &mockCurrentStateAPI{
+	consumer := NewOutputKeyChangeEventConsumer(gomatrixserverlib.ServerName("localhost"), "some_topic", nil, &mockKeyAPI{}, &mockCurrentStateAPI{
 		roomIDToJoinedMembers: map[string][]string{
 			newlyJoinedRoom: {syncingUser, existingUser},
 			"!another:room": {syncingUser, existingUser},
@@ -199,7 +213,7 @@ func TestKeyChangeCatchupOnJoinShareNoNewUsers(t *testing.T) {
 	syncResponse := types.NewResponse()
 	syncResponse = joinResponseWithRooms(syncResponse, syncingUser, []string{newlyJoinedRoom})
 
-	hasNew, err := consumer.Catchup(context.Background(), syncingUser, syncResponse, types.NewStreamToken(0, 0))
+	_, hasNew, err := consumer.Catchup(context.Background(), syncingUser, syncResponse, types.NewStreamToken(0, 0))
 	if err != nil {
 		t.Fatalf("Catchup returned an error: %s", err)
 	}
@@ -212,7 +226,7 @@ func TestKeyChangeCatchupOnJoinShareNoNewUsers(t *testing.T) {
 func TestKeyChangeCatchupOnLeaveShareNoUsers(t *testing.T) {
 	existingUser := "@bob:localhost"
 	newlyLeftRoom := "!TestKeyChangeCatchupOnLeaveShareNoUsers:bar"
-	consumer := NewOutputKeyChangeEventConsumer(gomatrixserverlib.ServerName("localhost"), "some_topic", nil, &mockCurrentStateAPI{
+	consumer := NewOutputKeyChangeEventConsumer(gomatrixserverlib.ServerName("localhost"), "some_topic", nil, &mockKeyAPI{}, &mockCurrentStateAPI{
 		roomIDToJoinedMembers: map[string][]string{
 			newlyLeftRoom:   {existingUser},
 			"!another:room": {syncingUser, existingUser},
@@ -221,7 +235,7 @@ func TestKeyChangeCatchupOnLeaveShareNoUsers(t *testing.T) {
 	syncResponse := types.NewResponse()
 	syncResponse = leaveResponseWithRooms(syncResponse, syncingUser, []string{newlyLeftRoom})
 
-	hasNew, err := consumer.Catchup(context.Background(), syncingUser, syncResponse, types.NewStreamToken(0, 0))
+	_, hasNew, err := consumer.Catchup(context.Background(), syncingUser, syncResponse, types.NewStreamToken(0, 0))
 	if err != nil {
 		t.Fatalf("Catchup returned an error: %s", err)
 	}
@@ -234,7 +248,7 @@ func TestKeyChangeCatchupOnLeaveShareNoUsers(t *testing.T) {
 func TestKeyChangeCatchupNoNewJoinsButMessages(t *testing.T) {
 	existingUser := "@bob1:localhost"
 	roomID := "!TestKeyChangeCatchupNoNewJoinsButMessages:bar"
-	consumer := NewOutputKeyChangeEventConsumer(gomatrixserverlib.ServerName("localhost"), "some_topic", nil, &mockCurrentStateAPI{
+	consumer := NewOutputKeyChangeEventConsumer(gomatrixserverlib.ServerName("localhost"), "some_topic", nil, &mockKeyAPI{}, &mockCurrentStateAPI{
 		roomIDToJoinedMembers: map[string][]string{
 			roomID: {syncingUser, existingUser},
 		},
@@ -280,7 +294,7 @@ func TestKeyChangeCatchupNoNewJoinsButMessages(t *testing.T) {
 	jr.Timeline.Events = roomTimelineEvents
 	syncResponse.Rooms.Join[roomID] = jr
 
-	hasNew, err := consumer.Catchup(context.Background(), syncingUser, syncResponse, types.NewStreamToken(0, 0))
+	_, hasNew, err := consumer.Catchup(context.Background(), syncingUser, syncResponse, types.NewStreamToken(0, 0))
 	if err != nil {
 		t.Fatalf("Catchup returned an error: %s", err)
 	}
@@ -297,7 +311,7 @@ func TestKeyChangeCatchupChangeAndLeft(t *testing.T) {
 	newlyLeftUser2 := "@debra:localhost"
 	newlyJoinedRoom := "!join:bar"
 	newlyLeftRoom := "!left:bar"
-	consumer := NewOutputKeyChangeEventConsumer(gomatrixserverlib.ServerName("localhost"), "some_topic", nil, &mockCurrentStateAPI{
+	consumer := NewOutputKeyChangeEventConsumer(gomatrixserverlib.ServerName("localhost"), "some_topic", nil, &mockKeyAPI{}, &mockCurrentStateAPI{
 		roomIDToJoinedMembers: map[string][]string{
 			newlyJoinedRoom: {syncingUser, newShareUser, newShareUser2},
 			newlyLeftRoom:   {newlyLeftUser, newlyLeftUser2},
@@ -308,7 +322,7 @@ func TestKeyChangeCatchupChangeAndLeft(t *testing.T) {
 	syncResponse = joinResponseWithRooms(syncResponse, syncingUser, []string{newlyJoinedRoom})
 	syncResponse = leaveResponseWithRooms(syncResponse, syncingUser, []string{newlyLeftRoom})
 
-	hasNew, err := consumer.Catchup(context.Background(), syncingUser, syncResponse, types.NewStreamToken(0, 0))
+	_, hasNew, err := consumer.Catchup(context.Background(), syncingUser, syncResponse, types.NewStreamToken(0, 0))
 	if err != nil {
 		t.Fatalf("Catchup returned an error: %s", err)
 	}
@@ -333,7 +347,7 @@ func TestKeyChangeCatchupChangeAndLeftSameRoom(t *testing.T) {
 	newShareUser := "@berta:localhost"
 	newShareUser2 := "@bobby:localhost"
 	roomID := "!join:bar"
-	consumer := NewOutputKeyChangeEventConsumer(gomatrixserverlib.ServerName("localhost"), "some_topic", nil, &mockCurrentStateAPI{
+	consumer := NewOutputKeyChangeEventConsumer(gomatrixserverlib.ServerName("localhost"), "some_topic", nil, &mockKeyAPI{}, &mockCurrentStateAPI{
 		roomIDToJoinedMembers: map[string][]string{
 			roomID:          {newShareUser, newShareUser2},
 			"!another:room": {syncingUser},
@@ -393,7 +407,7 @@ func TestKeyChangeCatchupChangeAndLeftSameRoom(t *testing.T) {
 	lr.Timeline.Events = roomEvents
 	syncResponse.Rooms.Leave[roomID] = lr
 
-	hasNew, err := consumer.Catchup(context.Background(), syncingUser, syncResponse, types.NewStreamToken(0, 0))
+	_, hasNew, err := consumer.Catchup(context.Background(), syncingUser, syncResponse, types.NewStreamToken(0, 0))
 	if err != nil {
 		t.Fatalf("Catchup returned an error: %s", err)
 	}