Allow "registration is idempotent, with username specified" to pass (#2488)

Co-authored-by: Neil Alexander <neilalexander@users.noreply.github.com>

2 files changed, 40 insertions, 20 deletions

diff --git a/clientapi/routing/register.go b/clientapi/routing/register.go
index eba4920c..c4ac0f2e 100644
--- a/clientapi/routing/register.go
+++ b/clientapi/routing/register.go
@@ -29,9 +29,10 @@ import (
 	"sync"
 	"time"
 
+	"github.com/tidwall/gjson"
+
 	"github.com/matrix-org/dendrite/internal/eventutil"
 	"github.com/matrix-org/dendrite/setup/config"
-	"github.com/tidwall/gjson"
 
 	"github.com/matrix-org/gomatrixserverlib"
 	"github.com/matrix-org/gomatrixserverlib/tokens"
@@ -68,9 +69,10 @@ const (
 // It shouldn't be passed by value because it contains a mutex.
 type sessionsDict struct {
 	sync.RWMutex
-	sessions map[string][]authtypes.LoginType
-	params   map[string]registerRequest
-	timer    map[string]*time.Timer
+	sessions               map[string][]authtypes.LoginType
+	sessionCompletedResult map[string]registerResponse
+	params                 map[string]registerRequest
+	timer                  map[string]*time.Timer
 	// deleteSessionToDeviceID protects requests to DELETE /devices/{deviceID} from being abused.
 	// If a UIA session is started by trying to delete device1, and then UIA is completed by deleting device2,
 	// the delete request will fail for device2 since the UIA was initiated by trying to delete device1.
@@ -115,6 +117,7 @@ func (d *sessionsDict) deleteSession(sessionID string) {
 	delete(d.params, sessionID)
 	delete(d.sessions, sessionID)
 	delete(d.deleteSessionToDeviceID, sessionID)
+	delete(d.sessionCompletedResult, sessionID)
 	// stop the timer, e.g. because the registration was completed
 	if t, ok := d.timer[sessionID]; ok {
 		if !t.Stop() {
@@ -130,6 +133,7 @@ func (d *sessionsDict) deleteSession(sessionID string) {
 func newSessionsDict() *sessionsDict {
 	return &sessionsDict{
 		sessions:                make(map[string][]authtypes.LoginType),
+		sessionCompletedResult:  make(map[string]registerResponse),
 		params:                  make(map[string]registerRequest),
 		timer:                   make(map[string]*time.Timer),
 		deleteSessionToDeviceID: make(map[string]string),
@@ -173,6 +177,19 @@ func (d *sessionsDict) addDeviceToDelete(sessionID, deviceID string) {
 	d.deleteSessionToDeviceID[sessionID] = deviceID
 }
 
+func (d *sessionsDict) addCompletedRegistration(sessionID string, response registerResponse) {
+	d.Lock()
+	defer d.Unlock()
+	d.sessionCompletedResult[sessionID] = response
+}
+
+func (d *sessionsDict) getCompletedRegistration(sessionID string) (registerResponse, bool) {
+	d.RLock()
+	defer d.RUnlock()
+	result, ok := d.sessionCompletedResult[sessionID]
+	return result, ok
+}
+
 func (d *sessionsDict) getDeviceToDelete(sessionID string) (string, bool) {
 	d.RLock()
 	defer d.RUnlock()
@@ -544,6 +561,14 @@ func Register(
 		r.DeviceID = data.DeviceID
 		r.InitialDisplayName = data.InitialDisplayName
 		r.InhibitLogin = data.InhibitLogin
+		// Check if the user already registered using this session, if so, return that result
+		if response, ok := sessions.getCompletedRegistration(sessionID); ok {
+			return util.JSONResponse{
+				Code: http.StatusOK,
+				JSON: response,
+			}
+		}
+
 	}
 	if resErr := httputil.UnmarshalJSON(reqBody, &r); resErr != nil {
 		return *resErr
@@ -839,13 +864,6 @@ func completeRegistration(
 	displayName, deviceID *string,
 	accType userapi.AccountType,
 ) util.JSONResponse {
-	var registrationOK bool
-	defer func() {
-		if registrationOK {
-			sessions.deleteSession(sessionID)
-		}
-	}()
-
 	if username == "" {
 		return util.JSONResponse{
 			Code: http.StatusBadRequest,
@@ -886,7 +904,6 @@ func completeRegistration(
 	// Check whether inhibit_login option is set. If so, don't create an access
 	// token or a device for this user
 	if inhibitLogin {
-		registrationOK = true
 		return util.JSONResponse{
 			Code: http.StatusOK,
 			JSON: registerResponse{
@@ -920,15 +937,17 @@ func completeRegistration(
 		}
 	}
 
-	registrationOK = true
+	result := registerResponse{
+		UserID:      devRes.Device.UserID,
+		AccessToken: devRes.Device.AccessToken,
+		HomeServer:  accRes.Account.ServerName,
+		DeviceID:    devRes.Device.ID,
+	}
+	sessions.addCompletedRegistration(sessionID, result)
+
 	return util.JSONResponse{
 		Code: http.StatusOK,
-		JSON: registerResponse{
-			UserID:      devRes.Device.UserID,
-			AccessToken: devRes.Device.AccessToken,
-			HomeServer:  accRes.Account.ServerName,
-			DeviceID:    devRes.Device.ID,
-		},
+		JSON: result,
 	}
 }
 
diff --git a/sytest-whitelist b/sytest-whitelist
index 5f6797a3..21bbc396 100644
--- a/sytest-whitelist
+++ b/sytest-whitelist
@@ -716,6 +716,7 @@ PUT /rooms/:room_id/redact/:event_id/:txn_id is idempotent
 Unnamed room comes with a name summary
 Named room comes with just joined member count summary
 Room summary only has 5 heroes
+registration is idempotent, with username specified
 Setting state twice is idempotent
 Joining room twice is idempotent
-Inbound federation can return missing events for shared visibility
\ No newline at end of file
+Inbound federation can return missing events for shared visibility