Merge SenderID & Per Room User Key work (#3109)

6 files changed, 47 insertions, 22 deletions

diff --git a/federationapi/federationapi_test.go b/federationapi/federationapi_test.go
index 17390843..5d167c0e 100644
--- a/federationapi/federationapi_test.go
+++ b/federationapi/federationapi_test.go
@@ -36,11 +36,11 @@ type fedRoomserverAPI struct {
 	queryRoomsForUser func(ctx context.Context, req *rsapi.QueryRoomsForUserRequest, res *rsapi.QueryRoomsForUserResponse) error
 }
 
-func (f *fedRoomserverAPI) QueryUserIDForSender(ctx context.Context, roomID string, senderID spec.SenderID) (*spec.UserID, error) {
+func (f *fedRoomserverAPI) QueryUserIDForSender(ctx context.Context, roomID spec.RoomID, senderID spec.SenderID) (*spec.UserID, error) {
 	return spec.NewUserID(string(senderID), true)
 }
 
-func (f *fedRoomserverAPI) QuerySenderIDForUser(ctx context.Context, roomID string, userID spec.UserID) (spec.SenderID, error) {
+func (f *fedRoomserverAPI) QuerySenderIDForUser(ctx context.Context, roomID spec.RoomID, userID spec.UserID) (spec.SenderID, error) {
 	return spec.SenderID(userID.String()), nil
 }
 
diff --git a/federationapi/internal/perform.go b/federationapi/internal/perform.go
index 485b79a0..7f61dba4 100644
--- a/federationapi/internal/perform.go
+++ b/federationapi/internal/perform.go
@@ -154,14 +154,9 @@ func (r *FederationInternalAPI) performJoinUsingServer(
 	if err != nil {
 		return err
 	}
-	senderID, err := r.rsAPI.QuerySenderIDForUser(ctx, roomID, *user)
-	if err != nil {
-		return err
-	}
 
 	joinInput := gomatrixserverlib.PerformJoinInput{
 		UserID:     user,
-		SenderID:   senderID,
 		RoomID:     room,
 		ServerName: serverName,
 		Content:    content,
@@ -169,12 +164,20 @@ func (r *FederationInternalAPI) performJoinUsingServer(
 		PrivateKey: r.cfg.Matrix.PrivateKey,
 		KeyID:      r.cfg.Matrix.KeyID,
 		KeyRing:    r.keyRing,
-		EventProvider: federatedEventProvider(ctx, r.federation, r.keyRing, user.Domain(), serverName, func(roomID string, senderID spec.SenderID) (*spec.UserID, error) {
+		EventProvider: federatedEventProvider(ctx, r.federation, r.keyRing, user.Domain(), serverName, func(roomID spec.RoomID, senderID spec.SenderID) (*spec.UserID, error) {
 			return r.rsAPI.QueryUserIDForSender(ctx, roomID, senderID)
 		}),
-		UserIDQuerier: func(roomID string, senderID spec.SenderID) (*spec.UserID, error) {
+		UserIDQuerier: func(roomID spec.RoomID, senderID spec.SenderID) (*spec.UserID, error) {
 			return r.rsAPI.QueryUserIDForSender(ctx, roomID, senderID)
 		},
+		SenderIDCreator: func(ctx context.Context, userID spec.UserID, roomID spec.RoomID) (spec.SenderID, error) {
+			key, keyErr := r.rsAPI.GetOrCreateUserRoomPrivateKey(ctx, userID, roomID)
+			if keyErr != nil {
+				return "", keyErr
+			}
+
+			return spec.SenderID(spec.Base64Bytes(key).Encode()), nil
+		},
 	}
 	response, joinErr := gomatrixserverlib.PerformJoin(ctx, r, joinInput)
 
@@ -368,7 +371,7 @@ func (r *FederationInternalAPI) performOutboundPeekUsingServer(
 
 	// authenticate the state returned (check its auth events etc)
 	// the equivalent of CheckSendJoinResponse()
-	userIDProvider := func(roomID string, senderID spec.SenderID) (*spec.UserID, error) {
+	userIDProvider := func(roomID spec.RoomID, senderID spec.SenderID) (*spec.UserID, error) {
 		return r.rsAPI.QueryUserIDForSender(ctx, roomID, senderID)
 	}
 	authEvents, stateEvents, err := gomatrixserverlib.CheckStateResponse(
@@ -459,7 +462,11 @@ func (r *FederationInternalAPI) PerformLeave(
 
 		// Set all the fields to be what they should be, this should be a no-op
 		// but it's possible that the remote server returned us something "odd"
-		senderID, err := r.rsAPI.QuerySenderIDForUser(ctx, request.RoomID, *userID)
+		roomID, err := spec.NewRoomID(request.RoomID)
+		if err != nil {
+			return err
+		}
+		senderID, err := r.rsAPI.QuerySenderIDForUser(ctx, *roomID, *userID)
 		if err != nil {
 			return err
 		}
@@ -527,7 +534,11 @@ func (r *FederationInternalAPI) SendInvite(
 	event gomatrixserverlib.PDU,
 	strippedState []gomatrixserverlib.InviteStrippedState,
 ) (gomatrixserverlib.PDU, error) {
-	inviter, err := r.rsAPI.QueryUserIDForSender(ctx, event.RoomID(), event.SenderID())
+	validRoomID, err := spec.NewRoomID(event.RoomID())
+	if err != nil {
+		return nil, err
+	}
+	inviter, err := r.rsAPI.QueryUserIDForSender(ctx, *validRoomID, event.SenderID())
 	if err != nil {
 		return nil, err
 	}
diff --git a/federationapi/routing/invite.go b/federationapi/routing/invite.go
index 5b15f810..e45209a2 100644
--- a/federationapi/routing/invite.go
+++ b/federationapi/routing/invite.go
@@ -95,7 +95,7 @@ func InviteV2(
 			StateQuerier:      rsAPI.StateQuerier(),
 			InviteEvent:       inviteReq.Event(),
 			StrippedState:     inviteReq.InviteRoomState(),
-			UserIDQuerier: func(roomID string, senderID spec.SenderID) (*spec.UserID, error) {
+			UserIDQuerier: func(roomID spec.RoomID, senderID spec.SenderID) (*spec.UserID, error) {
 				return rsAPI.QueryUserIDForSender(httpReq.Context(), roomID, senderID)
 			},
 		}
@@ -188,7 +188,7 @@ func InviteV1(
 		StateQuerier:      rsAPI.StateQuerier(),
 		InviteEvent:       event,
 		StrippedState:     strippedState,
-		UserIDQuerier: func(roomID string, senderID spec.SenderID) (*spec.UserID, error) {
+		UserIDQuerier: func(roomID spec.RoomID, senderID spec.SenderID) (*spec.UserID, error) {
 			return rsAPI.QueryUserIDForSender(httpReq.Context(), roomID, senderID)
 		},
 	}
diff --git a/federationapi/routing/join.go b/federationapi/routing/join.go
index d1480192..7aa50f65 100644
--- a/federationapi/routing/join.go
+++ b/federationapi/routing/join.go
@@ -98,7 +98,7 @@ func MakeJoin(
 		Roomserver: rsAPI,
 	}
 
-	senderID, err := rsAPI.QuerySenderIDForUser(httpReq.Context(), roomID.String(), userID)
+	senderID, err := rsAPI.QuerySenderIDForUser(httpReq.Context(), roomID, userID)
 	if err != nil {
 		util.GetLogger(httpReq.Context()).WithError(err).Error("rsAPI.QuerySenderIDForUser failed")
 		return util.JSONResponse{
@@ -118,7 +118,7 @@ func MakeJoin(
 		LocalServerName:   cfg.Matrix.ServerName,
 		LocalServerInRoom: res.RoomExists && res.IsInRoom,
 		RoomQuerier:       &roomQuerier,
-		UserIDQuerier: func(roomID string, senderID spec.SenderID) (*spec.UserID, error) {
+		UserIDQuerier: func(roomID spec.RoomID, senderID spec.SenderID) (*spec.UserID, error) {
 			return rsAPI.QueryUserIDForSender(httpReq.Context(), roomID, senderID)
 		},
 		BuildEventTemplate: createJoinTemplate,
@@ -215,7 +215,7 @@ func SendJoin(
 		PrivateKey:        cfg.Matrix.PrivateKey,
 		Verifier:          keys,
 		MembershipQuerier: &api.MembershipQuerier{Roomserver: rsAPI},
-		UserIDQuerier: func(roomID string, senderID spec.SenderID) (*spec.UserID, error) {
+		UserIDQuerier: func(roomID spec.RoomID, senderID spec.SenderID) (*spec.UserID, error) {
 			return rsAPI.QueryUserIDForSender(httpReq.Context(), roomID, senderID)
 		},
 	}
diff --git a/federationapi/routing/leave.go b/federationapi/routing/leave.go
index 716276be..5c8dd00f 100644
--- a/federationapi/routing/leave.go
+++ b/federationapi/routing/leave.go
@@ -87,7 +87,7 @@ func MakeLeave(
 		return event, stateEvents, nil
 	}
 
-	senderID, err := rsAPI.QuerySenderIDForUser(httpReq.Context(), roomID.String(), userID)
+	senderID, err := rsAPI.QuerySenderIDForUser(httpReq.Context(), roomID, userID)
 	if err != nil {
 		util.GetLogger(httpReq.Context()).WithError(err).Error("rsAPI.QuerySenderIDForUser failed")
 		return util.JSONResponse{
@@ -105,7 +105,7 @@ func MakeLeave(
 		LocalServerName:    cfg.Matrix.ServerName,
 		LocalServerInRoom:  res.RoomExists && res.IsInRoom,
 		BuildEventTemplate: createLeaveTemplate,
-		UserIDQuerier: func(roomID string, senderID spec.SenderID) (*spec.UserID, error) {
+		UserIDQuerier: func(roomID spec.RoomID, senderID spec.SenderID) (*spec.UserID, error) {
 			return rsAPI.QueryUserIDForSender(httpReq.Context(), roomID, senderID)
 		},
 	}
@@ -236,7 +236,14 @@ func SendLeave(
 	// Check that the sender belongs to the server that is sending us
 	// the request. By this point we've already asserted that the sender
 	// and the state key are equal so we don't need to check both.
-	sender, err := rsAPI.QueryUserIDForSender(httpReq.Context(), event.RoomID(), event.SenderID())
+	validRoomID, err := spec.NewRoomID(event.RoomID())
+	if err != nil {
+		return util.JSONResponse{
+			Code: http.StatusBadRequest,
+			JSON: spec.BadJSON("Room ID is invalid."),
+		}
+	}
+	sender, err := rsAPI.QueryUserIDForSender(httpReq.Context(), *validRoomID, event.SenderID())
 	if err != nil {
 		return util.JSONResponse{
 			Code: http.StatusForbidden,
diff --git a/federationapi/routing/threepid.go b/federationapi/routing/threepid.go
index 360802de..42ba8bfe 100644
--- a/federationapi/routing/threepid.go
+++ b/federationapi/routing/threepid.go
@@ -140,7 +140,14 @@ func ExchangeThirdPartyInvite(
 		}
 	}
 
-	userID, err := rsAPI.QueryUserIDForSender(httpReq.Context(), roomID, spec.SenderID(proto.SenderID))
+	validRoomID, err := spec.NewRoomID(roomID)
+	if err != nil {
+		return util.JSONResponse{
+			Code: http.StatusBadRequest,
+			JSON: spec.BadJSON("Invalid room ID"),
+		}
+	}
+	userID, err := rsAPI.QueryUserIDForSender(httpReq.Context(), *validRoomID, spec.SenderID(proto.SenderID))
 	if err != nil || userID == nil {
 		return util.JSONResponse{
 			Code: http.StatusBadRequest,
@@ -150,7 +157,7 @@ func ExchangeThirdPartyInvite(
 	senderDomain := userID.Domain()
 
 	// Check that the state key is correct.
-	targetUserID, err := rsAPI.QueryUserIDForSender(httpReq.Context(), roomID, spec.SenderID(*proto.StateKey))
+	targetUserID, err := rsAPI.QueryUserIDForSender(httpReq.Context(), *validRoomID, spec.SenderID(*proto.StateKey))
 	if err != nil || targetUserID == nil {
 		return util.JSONResponse{
 			Code: http.StatusBadRequest,