From 4fb0cdb84784d460fca905075a12e4255329c319 Mon Sep 17 00:00:00 2001 From: Skip R Date: Wed, 7 Jan 2026 20:10:18 -0800 Subject: [PATCH] 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 | 21 +++--- pkg/connector/login_browser.go | 67 ++---------------- pkg/connector/login_generic.go | 114 ++++++++++++++++++++++++++++++ pkg/connector/login_remoteauth.go | 68 ++---------------- pkg/connector/login_token.go | 71 ++----------------- 5 files changed, 141 insertions(+), 200 deletions(-) create mode 100644 pkg/connector/login_generic.go diff --git a/pkg/connector/login.go b/pkg/connector/login.go index 037851c..5f056f6 100644 --- a/pkg/connector/login.go +++ b/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") } diff --git a/pkg/connector/login_browser.go b/pkg/connector/login_browser.go index 70e9f08..d6bee8e 100644 --- a/pkg/connector/login_browser.go +++ b/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, diff --git a/pkg/connector/login_generic.go b/pkg/connector/login_generic.go new file mode 100644 index 0000000..ff094e2 --- /dev/null +++ b/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 . + +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") + } + } +} diff --git a/pkg/connector/login_remoteauth.go b/pkg/connector/login_remoteauth.go index e1ae04a..d828736 100644 --- a/pkg/connector/login_remoteauth.go +++ b/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) } diff --git a/pkg/connector/login_token.go b/pkg/connector/login_token.go index 60678b8..5b0d233 100644 --- a/pkg/connector/login_token.go +++ b/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() -}