From 3115998616a1af3c69b88a1c82e1b1f0b267b2f1 Mon Sep 17 00:00:00 2001
From: Tak Wai Wong <takwaiw@gmail.com>
Date: Thu, 16 Jun 2022 10:18:01 -0700
Subject: [PATCH] During registration, verify that the user ID passes the
 grammar of CAIP-10, and the Matrix ID. Also verify that the ID matches the
 authentication data. Then during login authentication, verify that the user
 ID in the auth request matches the fields in the signed message.

---
 clientapi/auth/login_publickey.go          |  1 +
 clientapi/auth/login_publickey_ethereum.go | 35 ++++++++++++++++++++++
 clientapi/routing/register.go              |  2 +-
 clientapi/routing/register_publickey.go    |  9 ++++++
 4 files changed, 46 insertions(+), 1 deletion(-)

diff --git a/clientapi/auth/login_publickey.go b/clientapi/auth/login_publickey.go
index b93420b2e..e999edeb7 100644
--- a/clientapi/auth/login_publickey.go
+++ b/clientapi/auth/login_publickey.go
@@ -30,6 +30,7 @@ import (
 
 type LoginPublicKeyHandler interface {
 	AccountExists(ctx context.Context) (string, *jsonerror.MatrixError)
+	IsValidUserIdForRegistration(userId string) bool
 	CreateLogin() *Login
 	GetSession() string
 	GetType() string
diff --git a/clientapi/auth/login_publickey_ethereum.go b/clientapi/auth/login_publickey_ethereum.go
index 592f02383..3ac367a81 100644
--- a/clientapi/auth/login_publickey_ethereum.go
+++ b/clientapi/auth/login_publickey_ethereum.go
@@ -17,6 +17,8 @@ package auth
 import (
 	"context"
 	"encoding/json"
+	"fmt"
+	"regexp"
 	"strings"
 
 	"github.com/matrix-org/dendrite/clientapi/jsonerror"
@@ -85,6 +87,24 @@ func (pk LoginPublicKeyEthereum) AccountExists(ctx context.Context) (string, *js
 	return localPart, nil
 }
 
+var validChainAgnosticIdRegex = regexp.MustCompile("^eip155=3a[0-9]+=3a0x[0-9a-fA-F]+$")
+
+func (pk LoginPublicKeyEthereum) IsValidUserIdForRegistration(userId string) bool {
+	// Verify that the user ID is a valid one according to spec.
+	// https://github.com/ChainAgnostic/CAIPs/blob/master/CAIPs/caip-10.md
+
+	// Matrix ID has additional grammar requirements for user ID.
+	// https://spec.matrix.org/v1.1/appendices/#user-identifiers
+	// Make sure disallowed characters are escaped.
+	// E.g. ":" is replaced with "=3a".
+
+	isValid := validChainAgnosticIdRegex.MatchString(userId)
+
+	// In addition, double check that the user ID for registration
+	// matches the authentication data in the request.
+	return isValid && userId == pk.UserId
+}
+
 func (pk LoginPublicKeyEthereum) ValidateLoginResponse() (bool, *jsonerror.MatrixError) {
 	// Parse the message to extract all the fields.
 	message, err := siwe.ParseMessage(pk.Message)
@@ -98,6 +118,12 @@ func (pk LoginPublicKeyEthereum) ValidateLoginResponse() (bool, *jsonerror.Matri
 		return false, jsonerror.InvalidSignature(err.Error())
 	}
 
+	// Error if the user ID does not match the signed message.
+	isVerifiedUserId := pk.verifyMessageUserId(message)
+	if !isVerifiedUserId {
+		return false, jsonerror.InvalidUsername(pk.UserId)
+	}
+
 	// Error if the chainId is not supported by the server.
 	if !contains(pk.config.PublicKeyAuthentication.Ethereum.ChainIDs, message.GetChainID()) {
 		return false, jsonerror.Forbidden("chainId")
@@ -118,6 +144,15 @@ func (pk LoginPublicKeyEthereum) CreateLogin() *Login {
 	return &login
 }
 
+func (pk LoginPublicKeyEthereum) verifyMessageUserId(message *siwe.Message) bool {
+	// Use info in the signed message to derive the expected user ID.
+	expectedUserId := fmt.Sprintf("eip155=3a%d=3a%s", message.GetChainID(), message.GetAddress())
+
+	// Case-insensitive comparison to make sure the user ID matches the expected
+	// one derived from the signed message.
+	return pk.UserId == strings.ToLower(expectedUserId)
+}
+
 func contains(list []int, element int) bool {
 	for _, i := range list {
 		if i == element {
diff --git a/clientapi/routing/register.go b/clientapi/routing/register.go
index 83c89356a..455bbf8cb 100644
--- a/clientapi/routing/register.go
+++ b/clientapi/routing/register.go
@@ -767,7 +767,7 @@ func handleRegistrationFlow(
 		sessions.addCompletedSessionStage(sessionID, authtypes.LoginTypeDummy)
 
 	case authtypes.LoginTypePublicKey:
-		isCompleted, authType, err := handlePublicKeyRegistration(cfg, reqBody, userAPI)
+		isCompleted, authType, err := handlePublicKeyRegistration(cfg, reqBody, &r, userAPI)
 		if err != nil {
 			return *err
 		}
diff --git a/clientapi/routing/register_publickey.go b/clientapi/routing/register_publickey.go
index aa0fea656..2ab2b6ca1 100644
--- a/clientapi/routing/register_publickey.go
+++ b/clientapi/routing/register_publickey.go
@@ -37,6 +37,7 @@ func newPublicKeyAuthSession(request *registerRequest, sessions *sessionsDict, s
 func handlePublicKeyRegistration(
 	cfg *config.ClientAPI,
 	reqBytes []byte,
+	r *registerRequest,
 	userAPI userapi.ClientUserAPI,
 ) (bool, authtypes.LoginType, *util.JSONResponse) {
 	if !cfg.PublicKeyAuthentication.Enabled() {
@@ -76,6 +77,14 @@ func handlePublicKeyRegistration(
 		}
 	}
 
+	isValidUserId := authHandler.IsValidUserIdForRegistration(r.Username)
+	if !isValidUserId {
+		return false, "", &util.JSONResponse{
+			Code: http.StatusUnauthorized,
+			JSON: jsonerror.InvalidUsername(r.Username),
+		}
+	}
+
 	isCompleted, jerr := authHandler.ValidateLoginResponse()
 	if jerr != nil {
 		return false, "", &util.JSONResponse{