From a46bf9878b041166d9610caee43c5a199a89ced3 Mon Sep 17 00:00:00 2001
From: Neil Alexander <neilalexander@users.noreply.github.com>
Date: Mon, 7 Nov 2022 16:06:37 +0000
Subject: [PATCH] Fix some bugs

---
 clientapi/routing/threepid.go             | 14 ++++++++------
 userapi/storage/sqlite3/threepid_table.go |  4 ++--
 2 files changed, 10 insertions(+), 8 deletions(-)

diff --git a/clientapi/routing/threepid.go b/clientapi/routing/threepid.go
index 4b7989ecb..971bfcad3 100644
--- a/clientapi/routing/threepid.go
+++ b/clientapi/routing/threepid.go
@@ -136,16 +136,17 @@ func CheckAndSave3PIDAssociation(
 	}
 
 	// Save the association in the database
-	localpart, _, err := gomatrixserverlib.SplitID('@', device.UserID)
+	localpart, domain, err := gomatrixserverlib.SplitID('@', device.UserID)
 	if err != nil {
 		util.GetLogger(req.Context()).WithError(err).Error("gomatrixserverlib.SplitID failed")
 		return jsonerror.InternalServerError()
 	}
 
 	if err = threePIDAPI.PerformSaveThreePIDAssociation(req.Context(), &api.PerformSaveThreePIDAssociationRequest{
-		ThreePID:  address,
-		Localpart: localpart,
-		Medium:    medium,
+		ThreePID:   address,
+		Localpart:  localpart,
+		ServerName: domain,
+		Medium:     medium,
 	}, &struct{}{}); err != nil {
 		util.GetLogger(req.Context()).WithError(err).Error("threePIDAPI.PerformSaveThreePIDAssociation failed")
 		return jsonerror.InternalServerError()
@@ -161,7 +162,7 @@ func CheckAndSave3PIDAssociation(
 func GetAssociated3PIDs(
 	req *http.Request, threepidAPI api.ClientUserAPI, device *api.Device,
 ) util.JSONResponse {
-	localpart, _, err := gomatrixserverlib.SplitID('@', device.UserID)
+	localpart, domain, err := gomatrixserverlib.SplitID('@', device.UserID)
 	if err != nil {
 		util.GetLogger(req.Context()).WithError(err).Error("gomatrixserverlib.SplitID failed")
 		return jsonerror.InternalServerError()
@@ -169,7 +170,8 @@ func GetAssociated3PIDs(
 
 	res := &api.QueryThreePIDsForLocalpartResponse{}
 	err = threepidAPI.QueryThreePIDsForLocalpart(req.Context(), &api.QueryThreePIDsForLocalpartRequest{
-		Localpart: localpart,
+		Localpart:  localpart,
+		ServerName: domain,
 	}, res)
 	if err != nil {
 		util.GetLogger(req.Context()).WithError(err).Error("threepidAPI.QueryThreePIDsForLocalpart failed")
diff --git a/userapi/storage/sqlite3/threepid_table.go b/userapi/storage/sqlite3/threepid_table.go
index 29b0eb832..2d3c7f2ee 100644
--- a/userapi/storage/sqlite3/threepid_table.go
+++ b/userapi/storage/sqlite3/threepid_table.go
@@ -47,10 +47,10 @@ const selectLocalpartForThreePIDSQL = "" +
 	"SELECT localpart FROM userapi_threepids WHERE threepid = $1 AND medium = $2"
 
 const selectThreePIDsForLocalpartSQL = "" +
-	"SELECT threepid, medium FROM userapi_threepids WHERE localpart = $1"
+	"SELECT threepid, medium FROM userapi_threepids WHERE localpart = $1 AND server_name = $2"
 
 const insertThreePIDSQL = "" +
-	"INSERT INTO userapi_threepids (threepid, medium, localpart) VALUES ($1, $2, $3)"
+	"INSERT INTO userapi_threepids (threepid, medium, localpart, server_name) VALUES ($1, $2, $3, $4)"
 
 const deleteThreePIDSQL = "" +
 	"DELETE FROM userapi_threepids WHERE threepid = $1 AND medium = $2"