From b55f00fa98979ed3fdfaad7f7af822e88db6c08e Mon Sep 17 00:00:00 2001
From: Neil Alexander <neilalexander@users.noreply.github.com>
Date: Fri, 18 Nov 2022 11:33:36 +0000
Subject: [PATCH] Gate registration and guest access per-virtual host

---
 clientapi/routing/register.go | 15 +++++++++++++--
 setup/config/config.go        |  2 +-
 setup/config/config_global.go | 22 ++++++++++++++++++++++
 3 files changed, 36 insertions(+), 3 deletions(-)

diff --git a/clientapi/routing/register.go b/clientapi/routing/register.go
index a92513b8b..bcde74681 100644
--- a/clientapi/routing/register.go
+++ b/clientapi/routing/register.go
@@ -650,7 +650,13 @@ func handleGuestRegistration(
 	cfg *config.ClientAPI,
 	userAPI userapi.ClientUserAPI,
 ) util.JSONResponse {
-	if cfg.RegistrationDisabled || cfg.GuestsDisabled {
+	registrationEnabled := !cfg.RegistrationDisabled
+	guestsEnabled := !cfg.GuestsDisabled
+	if r.ServerName != cfg.Matrix.ServerName {
+		registrationEnabled, guestsEnabled = cfg.Matrix.VirtualHost(r.ServerName).RegistrationAllowed()
+	}
+
+	if !registrationEnabled || !guestsEnabled {
 		return util.JSONResponse{
 			Code: http.StatusForbidden,
 			JSON: jsonerror.Forbidden("Guest registration is disabled"),
@@ -660,6 +666,7 @@ func handleGuestRegistration(
 	var res userapi.PerformAccountCreationResponse
 	err := userAPI.PerformAccountCreation(req.Context(), &userapi.PerformAccountCreationRequest{
 		AccountType: userapi.AccountTypeGuest,
+		ServerName:  r.ServerName,
 	}, &res)
 	if err != nil {
 		return util.JSONResponse{
@@ -736,7 +743,11 @@ func handleRegistrationFlow(
 		)
 	}
 
-	if cfg.RegistrationDisabled && r.Auth.Type != authtypes.LoginTypeSharedSecret {
+	registrationEnabled := !cfg.RegistrationDisabled
+	if r.ServerName != cfg.Matrix.ServerName {
+		registrationEnabled, _ = cfg.Matrix.VirtualHost(r.ServerName).RegistrationAllowed()
+	}
+	if !registrationEnabled && r.Auth.Type != authtypes.LoginTypeSharedSecret {
 		return util.JSONResponse{
 			Code: http.StatusForbidden,
 			JSON: jsonerror.Forbidden("Registration is disabled"),
diff --git a/setup/config/config.go b/setup/config/config.go
index 918bcbe3b..7e7ed1aa1 100644
--- a/setup/config/config.go
+++ b/setup/config/config.go
@@ -235,7 +235,7 @@ func loadConfig(
 		if v.KeyValidityPeriod == 0 {
 			v.KeyValidityPeriod = c.Global.KeyValidityPeriod
 		}
-		if v.PrivateKeyPath == "" {
+		if v.PrivateKeyPath == "" || v.PrivateKey == nil || v.KeyID == "" {
 			v.KeyID = c.Global.KeyID
 			v.PrivateKey = c.Global.PrivateKey
 			continue
diff --git a/setup/config/config_global.go b/setup/config/config_global.go
index 722230d9a..bb3a95972 100644
--- a/setup/config/config_global.go
+++ b/setup/config/config_global.go
@@ -151,6 +151,15 @@ func (c *Global) SplitLocalID(sigil byte, id string) (string, gomatrixserverlib.
 	return u, s, nil
 }
 
+func (c *Global) VirtualHost(serverName gomatrixserverlib.ServerName) *VirtualHost {
+	for _, v := range c.VirtualHosts {
+		if v.ServerName == serverName {
+			return v
+		}
+	}
+	return nil
+}
+
 func (c *Global) SigningIdentityFor(serverName gomatrixserverlib.ServerName) (*gomatrixserverlib.SigningIdentity, error) {
 	for _, id := range c.SigningIdentities() {
 		if id.ServerName == serverName {
@@ -202,6 +211,9 @@ type VirtualHost struct {
 
 	// Is registration enabled on this virtual host?
 	AllowRegistration bool `json:"allow_registration"`
+
+	// Is guest registration enabled on this virtual host?
+	AllowGuests bool `json:"allow_guests"`
 }
 
 func (v *VirtualHost) Verify(configErrs *ConfigErrors) {
@@ -216,6 +228,16 @@ func (v *VirtualHost) SigningIdentity() *gomatrixserverlib.SigningIdentity {
 	}
 }
 
+// RegistrationAllowed returns two bools, the first states whether registration
+// is allowed for this virtual host and the second states whether guests are
+// allowed for this virtual host.
+func (v *VirtualHost) RegistrationAllowed() (bool, bool) {
+	if v == nil {
+		return false, false
+	}
+	return v.AllowRegistration, v.AllowGuests
+}
+
 type OldVerifyKeys struct {
 	// Path to the private key.
 	PrivateKeyPath Path `yaml:"private_key"`