login: relocate shared login finalization logic into embedded struct

All of the login methods need to do (effectively) the same thing once we
have a token, so refactor this out into something we can reuse.

pkg/connector/login.go

@@ -27,16 +27,16 @@ const LoginStepIDComplete = "fi.mau.discord.login.complete"
 
 func (d *DiscordConnector) GetLoginFlows() []bridgev2.LoginFlow {
 	return []bridgev2.LoginFlow{
-		{
-			ID:          LoginFlowIDBrowser,
-			Name:        "Browser",
-			Description: "Log in to your Discord account in a web browser.",
-		},
 		{
 			ID:          LoginFlowIDRemoteAuth,
 			Name:        "QR Code",
 			Description: "Scan a QR code with the Discord mobile app to log in.",
 		},
+		{
+			ID:          LoginFlowIDBrowser,
+			Name:        "Browser",
+			Description: "Log in to your Discord account in a web browser.",
+		},
 		{
 			ID:          LoginFlowIDToken,
 			Name:        "Token",
@@ -46,13 +46,18 @@ func (d *DiscordConnector) GetLoginFlows() []bridgev2.LoginFlow {
 }
 
 func (d *DiscordConnector) CreateLogin(ctx context.Context, user *bridgev2.User, flowID string) (bridgev2.LoginProcess, error) {
+	login := DiscordGenericLogin{
+		connector: d,
+		User:      user,
+	}
+
 	switch flowID {
 	case LoginFlowIDToken:
-		return &DiscordTokenLogin{connector: d, User: user}, nil
+		return &DiscordTokenLogin{DiscordGenericLogin: &login}, nil
 	case LoginFlowIDRemoteAuth:
-		return &DiscordRemoteAuthLogin{connector: d, User: user}, nil
+		return &DiscordRemoteAuthLogin{DiscordGenericLogin: &login}, nil
 	case LoginFlowIDBrowser:
-		return &DiscordBrowserLogin{connector: d, User: user}, nil
+		return &DiscordBrowserLogin{DiscordGenericLogin: &login}, nil
 	default:
 		return nil, fmt.Errorf("unknown discord login flow id")
 	}

pkg/connector/login_browser.go

@@ -20,35 +20,18 @@ import (
 	"context"
 	"fmt"
 
-	"github.com/bwmarrin/discordgo"
 	"github.com/rs/zerolog"
 	"maunium.net/go/mautrix/bridgev2"
-	"maunium.net/go/mautrix/bridgev2/database"
-	"maunium.net/go/mautrix/bridgev2/networkid"
 )
 
 const LoginFlowIDBrowser = "fi.mau.discord.login.browser"
 
 type DiscordBrowserLogin struct {
-	connector *DiscordConnector
-	User      *bridgev2.User
-
-	Session *discordgo.Session
+	*DiscordGenericLogin
 }
 
 var _ bridgev2.LoginProcessCookies = (*DiscordBrowserLogin)(nil)
 
-func (dl *DiscordBrowserLogin) softlyCloseSession() {
-	dl.User.Log.Debug().Msg("Closing session")
-	err := dl.Session.Close()
-	if err != nil {
-		dl.User.Log.Err(err).Msg("Couldn't close Discord session in response to login cancellation")
-	}
-}
-
-func (dl *DiscordBrowserLogin) Cancel() {
-}
-
 const ExtractDiscordTokenJS = `
 new Promise((resolve) => {
 	let mautrixDiscordTokenCheckInterval
@@ -97,57 +80,15 @@ func (dl *DiscordBrowserLogin) SubmitCookies(ctx context.Context, cookies map[st
 	}
 	log.Debug().Msg("Logging in with submitted cookie")
 
-	// FIXME FIXME: The rest of this method is basically copy and pasted from
-	// DiscordTokenLogin, so find a way to tidy this up.
-
-	session, err := NewDiscordSession(ctx, token)
+	ul, err := dl.FinalizeCreatingLogin(ctx, token)
 	if err != nil {
-		return nil, fmt.Errorf("couldn't create discord session: %w", err)
+		return nil, fmt.Errorf("couldn't log in via browser: %w", err)
 	}
 
-	client := DiscordClient{
-		connector: dl.connector,
-		Session:   session,
-	}
-	client.SetUp(ctx, nil)
-	err = client.connect(ctx)
-	if err != nil {
-		dl.softlyCloseSession()
-		return nil, err
-	}
-	// At this point we've opened a WebSocket connection to the gateway, received
-	// a READY packet, and know who we are.
-	user := session.State.User
-
-	dl.Session = session
-	ul, err := dl.User.NewLogin(ctx, &database.UserLogin{
-		ID: networkid.UserLoginID(user.ID),
-		Metadata: &UserLoginMetadata{
-			Token:            token,
-			HeartbeatSession: session.HeartbeatSession,
-		},
-	}, &bridgev2.NewLoginParams{
-		LoadUserLogin: func(ctx context.Context, login *bridgev2.UserLogin) error {
-			login.Client = &client
-			client.UserLogin = login
-
-			// Only now that we have a UserLogin can we begin syncing.
-			client.BeginSyncingIfUserLoginPresent(ctx)
-			return nil
-		},
-		DeleteOnConflict:  true,
-		DontReuseExisting: false,
-	})
-	if err != nil {
-		dl.softlyCloseSession()
-		return nil, fmt.Errorf("couldn't create login: %w", err)
-	}
-	zerolog.Ctx(ctx).Info().Str("user_id", user.ID).Str("user_username", user.Username).Msg("Connected to Discord during login")
-
 	return &bridgev2.LoginStep{
 		Type:         bridgev2.LoginStepTypeComplete,
 		StepID:       LoginStepIDComplete,
-		Instructions: fmt.Sprintf("Logged in as %s", user),
+		Instructions: dl.CompleteInstructions(),
 		CompleteParams: &bridgev2.LoginCompleteParams{
 			UserLoginID: ul.ID,
 			UserLogin:   ul,

pkg/connector/login_generic.go

@@ -0,0 +1,114 @@
+// mautrix-discord - A Matrix-Discord puppeting bridge.
+// Copyright (C) 2024 Tulir Asokan
+//
+// This program is free software: you can redistribute it and/or modify
+// it under the terms of the GNU Affero General Public License as published by
+// the Free Software Foundation, either version 3 of the License, or
+// (at your option) any later version.
+//
+// This program is distributed in the hope that it will be useful,
+// but WITHOUT ANY WARRANTY; without even the implied warranty of
+// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+// GNU Affero General Public License for more details.
+//
+// You should have received a copy of the GNU Affero General Public License
+// along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+package connector
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/bwmarrin/discordgo"
+	"github.com/rs/zerolog"
+	"maunium.net/go/mautrix/bridgev2"
+	"maunium.net/go/mautrix/bridgev2/database"
+	"maunium.net/go/mautrix/bridgev2/networkid"
+)
+
+// DiscordGenericLogin is embedded within each struct that implements
+// bridgev2.LoginProcess in order to encapsulate the common behavior that needs
+// to occur after procuring a valid user token. Namely, creating a gateway
+// connection to Discord and an associated UserLogin to wrap things up.
+//
+// It also implements a baseline Cancel method that closes the gateway
+// connection.
+type DiscordGenericLogin struct {
+	User      *bridgev2.User
+	connector *DiscordConnector
+
+	Session *discordgo.Session
+
+	// The Discord user we've authenticated as. This is only non-nil if
+	// a call to FinalizeCreatingLogin has succeeded.
+	DiscordUser *discordgo.User
+}
+
+func (dl *DiscordGenericLogin) FinalizeCreatingLogin(ctx context.Context, token string) (*bridgev2.UserLogin, error) {
+	session, err := NewDiscordSession(ctx, token)
+	if err != nil {
+		return nil, fmt.Errorf("couldn't create discord session: %w", err)
+	}
+
+	client := DiscordClient{
+		connector: dl.connector,
+		Session:   session,
+	}
+	client.SetUp(ctx, nil)
+
+	err = client.connect(ctx)
+	if err != nil {
+		dl.Cancel()
+		return nil, err
+	}
+	// At this point we've opened a WebSocket connection to the gateway, received
+	// a READY packet, and know who we are.
+	user := session.State.User
+	dl.DiscordUser = user
+
+	dl.Session = session
+	ul, err := dl.User.NewLogin(ctx, &database.UserLogin{
+		ID: networkid.UserLoginID(user.ID),
+		Metadata: &UserLoginMetadata{
+			Token:            token,
+			HeartbeatSession: session.HeartbeatSession,
+		},
+	}, &bridgev2.NewLoginParams{
+		LoadUserLogin: func(ctx context.Context, login *bridgev2.UserLogin) error {
+			login.Client = &client
+			client.UserLogin = login
+
+			// Only now that we have a UserLogin can we begin syncing.
+			client.BeginSyncingIfUserLoginPresent(ctx)
+			return nil
+		},
+		DeleteOnConflict:  true,
+		DontReuseExisting: false,
+	})
+	if err != nil {
+		dl.Cancel()
+		return nil, fmt.Errorf("couldn't create login: %w", err)
+	}
+
+	zerolog.Ctx(ctx).Info().
+		Str("user_id", user.ID).
+		Str("user_username", user.Username).
+		Msg("Logged in to Discord")
+
+	return ul, nil
+}
+
+func (dl *DiscordGenericLogin) CompleteInstructions() string {
+	return fmt.Sprintf("Logged in as %s", dl.DiscordUser.Username)
+}
+
+func (dl *DiscordGenericLogin) Cancel() {
+	if dl.Session != nil {
+		dl.User.Log.Debug().Msg("Login cancelled, closing session")
+		err := dl.Session.Close()
+		if err != nil {
+			dl.User.Log.Err(err).Msg("Couldn't close Discord session in response to login cancellation")
+		}
+	}
+}

pkg/connector/login_remoteauth.go

@@ -20,11 +20,8 @@ import (
 	"context"
 	"fmt"
 
-	"github.com/bwmarrin/discordgo"
 	"github.com/rs/zerolog"
 	"maunium.net/go/mautrix/bridgev2"
-	"maunium.net/go/mautrix/bridgev2/database"
-	"maunium.net/go/mautrix/bridgev2/networkid"
 
 	"go.mau.fi/mautrix-discord/pkg/remoteauth"
 )
@@ -32,10 +29,7 @@ import (
 const LoginFlowIDRemoteAuth = "fi.mau.discord.login.remote_auth"
 
 type DiscordRemoteAuthLogin struct {
-	connector *DiscordConnector
-	User      *bridgev2.User
-
-	Session *discordgo.Session
+	*DiscordGenericLogin
 
 	remoteAuthClient *remoteauth.Client
 	qrChan           chan string
@@ -114,66 +108,15 @@ func (dl *DiscordRemoteAuthLogin) Wait(ctx context.Context) (*bridgev2.LoginStep
 }
 
 func (dl *DiscordRemoteAuthLogin) finalizeSuccessfulLogin(ctx context.Context, user remoteauth.User) (*bridgev2.LoginStep, error) {
-	log := zerolog.Ctx(ctx)
-
-	session, err := NewDiscordSession(ctx, user.Token)
+	ul, err := dl.FinalizeCreatingLogin(ctx, user.Token)
 	if err != nil {
-		return nil, fmt.Errorf("couldn't create discord session from successful remoteauth: %w", err)
+		return nil, fmt.Errorf("couldn't log in via remoteauth: %w", err)
 	}
 
-	client := &DiscordClient{
-		connector: dl.connector,
-		Session:   session,
-	}
-	client.SetUp(ctx, nil)
-	err = client.connect(ctx)
-
-	softlyClose := func() {
-		log.Debug().Msg("Softly closing session due to error after successful remoteauth")
-		err := dl.Session.Close()
-		if err != nil {
-			log.Err(err).Msg("Couldn't softly close session due to error after successful remoteauth")
-		}
-	}
-	if err != nil {
-		softlyClose()
-		return nil, fmt.Errorf("couldn't connect to Discord: %w", err)
-	}
-	// At this point we've opened a WebSocket connection to the gateway, received
-	// a READY packet, and know who we are.
-	discordUser := session.State.User
-	dl.Session = session
-
-	ul, err := dl.User.NewLogin(ctx, &database.UserLogin{
-		ID: networkid.UserLoginID(user.UserID),
-		Metadata: &UserLoginMetadata{
-			Token:            user.Token,
-			HeartbeatSession: discordgo.NewHeartbeatSession(),
-		},
-	}, &bridgev2.NewLoginParams{
-		DeleteOnConflict: true,
-		LoadUserLogin: func(ctx context.Context, ul *bridgev2.UserLogin) error {
-			ul.Client = client
-			client.UserLogin = ul
-
-			// Only now that we have a UserLogin can we begin syncing.
-			client.BeginSyncingIfUserLoginPresent(ctx)
-			return nil
-		},
-	})
-	if err != nil {
-		softlyClose()
-		return nil, fmt.Errorf("couldn't create login after successful remoteauth: %w", err)
-	}
-	zerolog.Ctx(ctx).Info().
-		Str("user_id", discordUser.ID).
-		Str("user_username", discordUser.Username).
-		Msg("Connected to Discord during login")
-
 	return &bridgev2.LoginStep{
 		Type:         bridgev2.LoginStepTypeComplete,
 		StepID:       LoginStepIDComplete,
-		Instructions: fmt.Sprintf("Logged in as %s", user.Username),
+		Instructions: dl.CompleteInstructions(),
 		CompleteParams: &bridgev2.LoginCompleteParams{
 			UserLoginID: ul.ID,
 			UserLogin:   ul,
@@ -183,8 +126,9 @@ func (dl *DiscordRemoteAuthLogin) finalizeSuccessfulLogin(ctx context.Context, u
 
 func (dl *DiscordRemoteAuthLogin) Cancel() {
 	dl.User.Log.Debug().Msg("Discord remoteauth cancelled")
+	dl.DiscordGenericLogin.Cancel()
 
-	// remoteauth.Client doesn't seem to expose a cancellation method right now.
+	// remoteauth.Client doesn't seem to expose a cancellation method.
 	close(dl.doneChan)
 	close(dl.qrChan)
 }

pkg/connector/login_token.go

@@ -20,20 +20,13 @@ import (
 	"context"
 	"fmt"
 
-	"github.com/bwmarrin/discordgo"
-	"github.com/rs/zerolog"
 	"maunium.net/go/mautrix/bridgev2"
-	"maunium.net/go/mautrix/bridgev2/database"
-	"maunium.net/go/mautrix/bridgev2/networkid"
 )
 
 const LoginFlowIDToken = "fi.mau.discord.login.token"
 
 type DiscordTokenLogin struct {
-	connector *DiscordConnector
-	User      *bridgev2.User
-	Token     string
-	Session   *discordgo.Session
+	*DiscordGenericLogin
 }
 
 var _ bridgev2.LoginProcessUserInput = (*DiscordTokenLogin)(nil)
@@ -62,74 +55,18 @@ func (dl *DiscordTokenLogin) SubmitUserInput(ctx context.Context, input map[stri
 		return nil, fmt.Errorf("no token provided")
 	}
 
-	log := zerolog.Ctx(ctx)
-
-	log.Info().Msg("Creating session from provided token")
-	dl.Token = token
-
-	session, err := NewDiscordSession(ctx, token)
+	ul, err := dl.FinalizeCreatingLogin(ctx, token)
 	if err != nil {
-		return nil, fmt.Errorf("couldn't create discord session: %w", err)
+		return nil, fmt.Errorf("couldn't login from token: %w", err)
 	}
 
-	client := DiscordClient{
-		connector: dl.connector,
-		Session:   session,
-	}
-	client.SetUp(ctx, nil)
-	err = client.connect(ctx)
-	if err != nil {
-		dl.softlyCloseSession()
-		return nil, err
-	}
-	// At this point we've opened a WebSocket connection to the gateway, received
-	// a READY packet, and know who we are.
-	user := session.State.User
-
-	dl.Session = session
-	ul, err := dl.User.NewLogin(ctx, &database.UserLogin{
-		ID: networkid.UserLoginID(user.ID),
-		Metadata: &UserLoginMetadata{
-			Token:            token,
-			HeartbeatSession: session.HeartbeatSession,
-		},
-	}, &bridgev2.NewLoginParams{
-		LoadUserLogin: func(ctx context.Context, login *bridgev2.UserLogin) error {
-			login.Client = &client
-			client.UserLogin = login
-
-			// Only now that we have a UserLogin can we begin syncing.
-			client.BeginSyncingIfUserLoginPresent(ctx)
-			return nil
-		},
-		DeleteOnConflict:  true,
-		DontReuseExisting: false,
-	})
-	if err != nil {
-		dl.softlyCloseSession()
-		return nil, fmt.Errorf("couldn't create login: %w", err)
-	}
-	zerolog.Ctx(ctx).Info().Str("user_id", user.ID).Str("user_username", user.Username).Msg("Connected to Discord during login")
-
 	return &bridgev2.LoginStep{
 		Type:         bridgev2.LoginStepTypeComplete,
 		StepID:       LoginStepIDComplete,
-		Instructions: fmt.Sprintf("Logged in as %s", user),
+		Instructions: dl.CompleteInstructions(),
 		CompleteParams: &bridgev2.LoginCompleteParams{
 			UserLoginID: ul.ID,
 			UserLogin:   ul,
 		},
 	}, nil
 }
-
-func (dl *DiscordTokenLogin) softlyCloseSession() {
-	dl.User.Log.Debug().Msg("Closing session")
-	err := dl.Session.Close()
-	if err != nil {
-		dl.User.Log.Err(err).Msg("Couldn't close Discord session in response to login cancellation")
-	}
-}
-
-func (dl *DiscordTokenLogin) Cancel() {
-	dl.softlyCloseSession()
-}