login: implement QR login

Signed-off-by: Sumner Evans <sumner.evans@automattic.com>

pkg/connector/client.go

@@ -317,7 +317,7 @@ func NewTelegramClient(ctx context.Context, tc *TelegramConnector, login *bridge
 // connectTelegramClient blocks until client is connected, calling Run
 // internally.
 // Technique from: https://github.com/gotd/contrib/blob/master/bg/connect.go
-func connectTelegramClient(ctx context.Context, client *telegram.Client) (context.CancelFunc, error) {
+func connectTelegramClient(ctx context.Context, client *telegram.Client) (context.Context, context.CancelFunc, error) {
 	ctx, cancel := context.WithCancel(ctx)
 
 	errC := make(chan error, 1)
@@ -337,14 +337,14 @@ func connectTelegramClient(ctx context.Context, client *telegram.Client) (contex
 	select {
 	case <-ctx.Done(): // context canceled
 		cancel()
-		return func() {}, ctx.Err()
+		return nil, func() {}, ctx.Err()
 	case err := <-errC: // startup timeout
 		cancel()
-		return func() {}, err
+		return nil, func() {}, err
 	case <-initDone: // init done
 	}
 
-	return cancel, nil
+	return ctx, cancel, nil
 }
 
 func (t *TelegramClient) onDead() {
@@ -382,7 +382,7 @@ func (t *TelegramClient) onAuthError(err error) {
 }
 
 func (t *TelegramClient) Connect(ctx context.Context) (err error) {
-	t.clientCancel, err = connectTelegramClient(ctx, t.client)
+	ctx, t.clientCancel, err = connectTelegramClient(ctx, t.client)
 	if err != nil {
 		return err
 	}

pkg/connector/login.go

@@ -18,15 +18,9 @@ package connector
 
 import (
 	"context"
-	"errors"
 	"fmt"
 
-	"github.com/gotd/td/telegram"
-	"github.com/gotd/td/telegram/auth"
 	"github.com/gotd/td/tg"
-	"github.com/rs/zerolog"
-	"go.mau.fi/zerozap"
-	"go.uber.org/zap"
 	"maunium.net/go/mautrix/bridgev2"
 	"maunium.net/go/mautrix/bridgev2/database"
 
@@ -34,159 +28,44 @@ import (
 	"go.mau.fi/mautrix-telegram/pkg/connector/util"
 )
 
-// TODO QR login support
+const (
+	LoginFlowIDPhone = "phone"
+	LoginFlowIDQR    = "qr"
 
-const LoginFlowIDPhone = "phone"
+	LoginStepIDComplete = "fi.mau.telegram.login.complete"
+)
 
 func (tg *TelegramConnector) GetLoginFlows() []bridgev2.LoginFlow {
-	return []bridgev2.LoginFlow{{
+	return []bridgev2.LoginFlow{
-		Name:        "Phone Number",
+		{
-		Description: "Login using your Telegram phone number",
+			Name:        "Phone Number",
-		ID:          LoginFlowIDPhone,
+			Description: "Login using your Telegram phone number",
-	}}
+			ID:          LoginFlowIDPhone,
+		},
+		{
+			Name:        "QR Code",
+			Description: "Login by scanning a QR code from your phone",
+			ID:          LoginFlowIDQR,
+		},
+	}
 }
 
 func (tg *TelegramConnector) CreateLogin(ctx context.Context, user *bridgev2.User, flowID string) (bridgev2.LoginProcess, error) {
-	if flowID != LoginFlowIDPhone {
+	switch flowID {
+	case LoginFlowIDPhone:
+		return &PhoneLogin{user: user, main: tg}, nil
+	case LoginFlowIDQR:
+		return &QRLogin{user: user, main: tg}, nil
+	default:
 		return nil, fmt.Errorf("unknown flow ID %s", flowID)
 	}
-	return &PhoneLogin{user: user, main: tg}, nil
 }
 
-const (
+func finalizeLogin(ctx context.Context, user *bridgev2.User, authorization *tg.AuthAuthorization, metadata UserLoginMetadata) (*bridgev2.LoginStep, error) {
-	LoginStepIDPhoneNumber = "fi.mau.telegram.phone_number"
-	LoginStepIDCode        = "fi.mau.telegram.code"
-	LoginStepIDPassword    = "fi.mau.telegram.password"
-	LoginStepIDComplete    = "fi.mau.telegram.complete"
-)
-
-type PhoneLogin struct {
-	user         *bridgev2.User
-	main         *TelegramConnector
-	authData     UserLoginSession
-	client       *telegram.Client
-	clientCancel context.CancelFunc
-
-	phone string
-	hash  string
-}
-
-var _ bridgev2.LoginProcessUserInput = (*PhoneLogin)(nil)
-
-func (p *PhoneLogin) Cancel() {
-	p.clientCancel()
-}
-
-func (p *PhoneLogin) Start(ctx context.Context) (*bridgev2.LoginStep, error) {
-	return &bridgev2.LoginStep{
-		Type:         bridgev2.LoginStepTypeUserInput,
-		StepID:       LoginStepIDPhoneNumber,
-		Instructions: "Please enter your phone number",
-		UserInputParams: &bridgev2.LoginUserInputParams{
-			Fields: []bridgev2.LoginInputDataField{
-				{
-					Type:        bridgev2.LoginInputFieldTypePhoneNumber,
-					ID:          LoginStepIDPhoneNumber,
-					Name:        "Phone Number",
-					Description: "Include the country code with +",
-				},
-			},
-		},
-	}, nil
-}
-
-func (p *PhoneLogin) SubmitUserInput(ctx context.Context, input map[string]string) (*bridgev2.LoginStep, error) {
-	if phone, ok := input[LoginStepIDPhoneNumber]; ok {
-		p.phone = phone
-		p.client = telegram.NewClient(p.main.Config.APIID, p.main.Config.APIHash, telegram.Options{
-			CustomSessionStorage: &p.authData,
-			Logger:               zap.New(zerozap.New(zerolog.Ctx(ctx).With().Str("component", "telegram_login_client").Logger())),
-		})
-		var err error
-		p.clientCancel, err = connectTelegramClient(context.Background(), p.client)
-		if err != nil {
-			return nil, err
-		}
-		sentCode, err := p.client.Auth().SendCode(ctx, p.phone, auth.SendCodeOptions{})
-		if err != nil {
-			return nil, err
-		}
-		switch s := sentCode.(type) {
-		case *tg.AuthSentCode:
-			p.hash = s.PhoneCodeHash
-			return &bridgev2.LoginStep{
-				Type:         bridgev2.LoginStepTypeUserInput,
-				StepID:       LoginStepIDCode,
-				Instructions: "Please enter the code sent to your phone",
-				UserInputParams: &bridgev2.LoginUserInputParams{
-					Fields: []bridgev2.LoginInputDataField{
-						{
-							Type: bridgev2.LoginInputFieldType2FACode,
-							ID:   LoginStepIDCode,
-							Name: "Code",
-						},
-					},
-				},
-			}, nil
-		case *tg.AuthSentCodeSuccess:
-			switch a := s.Authorization.(type) {
-			case *tg.AuthAuthorization:
-				// Looks that we are already authorized.
-				return p.handleAuthSuccess(ctx, a)
-			case *tg.AuthAuthorizationSignUpRequired:
-				return nil, fmt.Errorf("phone number does not correspond with an existing Telegram account and sign-up is not supported")
-			default:
-				return nil, fmt.Errorf("unexpected authorization type: %T", sentCode)
-			}
-		default:
-			return nil, fmt.Errorf("unexpected sent code type: %T", sentCode)
-		}
-	} else if code, ok := input[LoginStepIDCode]; ok {
-		authorization, err := p.client.Auth().SignIn(ctx, p.phone, code, p.hash)
-		if errors.Is(err, auth.ErrPasswordAuthNeeded) {
-			return &bridgev2.LoginStep{
-				Type:         bridgev2.LoginStepTypeUserInput,
-				StepID:       LoginStepIDPassword,
-				Instructions: "Please enter your password",
-				UserInputParams: &bridgev2.LoginUserInputParams{
-					Fields: []bridgev2.LoginInputDataField{
-						{
-							Type: bridgev2.LoginInputFieldTypePassword,
-							ID:   LoginStepIDPassword,
-							Name: "Password",
-						},
-					},
-				},
-			}, nil
-		} else if errors.Is(err, &auth.SignUpRequired{}) {
-			return nil, fmt.Errorf("sign-up is not supported")
-		} else if err != nil {
-			return nil, fmt.Errorf("failed to submit code: %w", err)
-		}
-		return p.handleAuthSuccess(ctx, authorization)
-	} else if password, ok := input[LoginStepIDPassword]; ok {
-		authorization, err := p.client.Auth().Password(ctx, password)
-		if err != nil {
-			return nil, fmt.Errorf("failed to submit password: %w", err)
-		}
-		return p.handleAuthSuccess(ctx, authorization)
-	}
-
-	return nil, fmt.Errorf("unexpected state during phone login")
-}
-
-func (p *PhoneLogin) handleAuthSuccess(ctx context.Context, authorization *tg.AuthAuthorization) (*bridgev2.LoginStep, error) {
-	// Now that we have the Telegram user ID, store it in the database and
-	// close the login client.
-	p.clientCancel()
-
 	userLoginID := ids.MakeUserLoginID(authorization.User.GetID())
-	ul, err := p.user.NewLogin(ctx, &database.UserLogin{
+	ul, err := user.NewLogin(ctx, &database.UserLogin{
-		ID: userLoginID,
+		ID:       userLoginID,
-		Metadata: &UserLoginMetadata{
+		Metadata: &metadata,
-			Phone:   p.phone,
-			Session: p.authData,
-		},
 	}, nil)
 	if err != nil {
 		return nil, fmt.Errorf("failed to save new login: %w", err)
@@ -196,7 +75,7 @@ func (p *PhoneLogin) handleAuthSuccess(ctx context.Context, authorization *tg.Au
 		return nil, fmt.Errorf("failed to connect after login: %w", err)
 	}
 	client := ul.Client.(*TelegramClient)
-	user, err := client.client.Self(ctx)
+	me, err := client.client.Self(ctx)
 	if err != nil {
 		return nil, err
 	}
@@ -209,7 +88,7 @@ func (p *PhoneLogin) handleAuthSuccess(ctx context.Context, authorization *tg.Au
 	return &bridgev2.LoginStep{
 		Type:         bridgev2.LoginStepTypeComplete,
 		StepID:       LoginStepIDComplete,
-		Instructions: fmt.Sprintf("Successfully logged in as %d / +%s (%s)", user.ID, user.Phone, util.FormatFullName(user.FirstName, user.LastName)),
+		Instructions: fmt.Sprintf("Successfully logged in as %d / +%s (%s)", me.ID, me.Phone, util.FormatFullName(me.FirstName, me.LastName)),
 		CompleteParams: &bridgev2.LoginCompleteParams{
 			UserLoginID: ul.ID,
 		},

pkg/connector/loginphone.go

@@ -0,0 +1,146 @@
+package connector
+
+import (
+	"context"
+	"errors"
+	"fmt"
+
+	"github.com/gotd/td/telegram"
+	"github.com/gotd/td/telegram/auth"
+	"github.com/gotd/td/tg"
+	"github.com/rs/zerolog"
+	"go.mau.fi/zerozap"
+	"go.uber.org/zap"
+	"maunium.net/go/mautrix/bridgev2"
+)
+
+const (
+	LoginStepIDPhoneNumber = "fi.mau.telegram.login.phone_number"
+	LoginStepIDCode        = "fi.mau.telegram.login.code"
+	LoginStepIDPassword    = "fi.mau.telegram.login.password"
+)
+
+type PhoneLogin struct {
+	user             *bridgev2.User
+	main             *TelegramConnector
+	authData         UserLoginSession
+	authClient       *telegram.Client
+	authClientCancel context.CancelFunc
+
+	phone string
+	hash  string
+}
+
+var _ bridgev2.LoginProcessUserInput = (*PhoneLogin)(nil)
+
+func (p *PhoneLogin) Cancel() {
+	p.authClientCancel()
+}
+
+func (p *PhoneLogin) Start(ctx context.Context) (*bridgev2.LoginStep, error) {
+	return &bridgev2.LoginStep{
+		Type:         bridgev2.LoginStepTypeUserInput,
+		StepID:       LoginStepIDPhoneNumber,
+		Instructions: "Please enter your phone number",
+		UserInputParams: &bridgev2.LoginUserInputParams{
+			Fields: []bridgev2.LoginInputDataField{
+				{
+					Type:        bridgev2.LoginInputFieldTypePhoneNumber,
+					ID:          LoginStepIDPhoneNumber,
+					Name:        "Phone Number",
+					Description: "Include the country code with +",
+				},
+			},
+		},
+	}, nil
+}
+
+func (p *PhoneLogin) SubmitUserInput(ctx context.Context, input map[string]string) (*bridgev2.LoginStep, error) {
+	if phone, ok := input[LoginStepIDPhoneNumber]; ok {
+		p.phone = phone
+		p.authClient = telegram.NewClient(p.main.Config.APIID, p.main.Config.APIHash, telegram.Options{
+			CustomSessionStorage: &p.authData,
+			Logger:               zap.New(zerozap.New(zerolog.Ctx(ctx).With().Str("component", "telegram_phone_login_client").Logger())),
+		})
+		var err error
+		_, p.authClientCancel, err = connectTelegramClient(context.Background(), p.authClient)
+		if err != nil {
+			return nil, err
+		}
+		sentCode, err := p.authClient.Auth().SendCode(ctx, p.phone, auth.SendCodeOptions{})
+		if err != nil {
+			return nil, err
+		}
+		switch s := sentCode.(type) {
+		case *tg.AuthSentCode:
+			p.hash = s.PhoneCodeHash
+			return &bridgev2.LoginStep{
+				Type:         bridgev2.LoginStepTypeUserInput,
+				StepID:       LoginStepIDCode,
+				Instructions: "Please enter the code sent to your phone",
+				UserInputParams: &bridgev2.LoginUserInputParams{
+					Fields: []bridgev2.LoginInputDataField{
+						{
+							Type: bridgev2.LoginInputFieldType2FACode,
+							ID:   LoginStepIDCode,
+							Name: "Code",
+						},
+					},
+				},
+			}, nil
+		case *tg.AuthSentCodeSuccess:
+			switch a := s.Authorization.(type) {
+			case *tg.AuthAuthorization:
+				// Looks that we are already authorized.
+				return p.handleAuthSuccess(ctx, a)
+			case *tg.AuthAuthorizationSignUpRequired:
+				return nil, fmt.Errorf("phone number does not correspond with an existing Telegram account and sign-up is not supported")
+			default:
+				return nil, fmt.Errorf("unexpected authorization type: %T", sentCode)
+			}
+		default:
+			return nil, fmt.Errorf("unexpected sent code type: %T", sentCode)
+		}
+	} else if code, ok := input[LoginStepIDCode]; ok {
+		authorization, err := p.authClient.Auth().SignIn(ctx, p.phone, code, p.hash)
+		if errors.Is(err, auth.ErrPasswordAuthNeeded) {
+			return &bridgev2.LoginStep{
+				Type:         bridgev2.LoginStepTypeUserInput,
+				StepID:       LoginStepIDPassword,
+				Instructions: "Please enter your password",
+				UserInputParams: &bridgev2.LoginUserInputParams{
+					Fields: []bridgev2.LoginInputDataField{
+						{
+							Type: bridgev2.LoginInputFieldTypePassword,
+							ID:   LoginStepIDPassword,
+							Name: "Password",
+						},
+					},
+				},
+			}, nil
+		} else if errors.Is(err, &auth.SignUpRequired{}) {
+			return nil, fmt.Errorf("sign-up is not supported")
+		} else if err != nil {
+			return nil, fmt.Errorf("failed to submit code: %w", err)
+		}
+		return p.handleAuthSuccess(ctx, authorization)
+	} else if password, ok := input[LoginStepIDPassword]; ok {
+		authorization, err := p.authClient.Auth().Password(ctx, password)
+		if err != nil {
+			return nil, fmt.Errorf("failed to submit password: %w", err)
+		}
+		return p.handleAuthSuccess(ctx, authorization)
+	}
+
+	return nil, fmt.Errorf("unexpected state during phone login")
+}
+
+func (p *PhoneLogin) handleAuthSuccess(ctx context.Context, authorization *tg.AuthAuthorization) (*bridgev2.LoginStep, error) {
+	// Stop the login client.
+	p.authClientCancel()
+
+	return finalizeLogin(ctx, p.user, authorization, UserLoginMetadata{
+		Phone:   p.phone,
+		Session: p.authData,
+	})
+}

pkg/connector/loginqr.go

@@ -0,0 +1,173 @@
+package connector
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/gotd/td/telegram"
+	"github.com/gotd/td/telegram/auth/qrlogin"
+	"github.com/gotd/td/telegram/updates"
+	"github.com/gotd/td/tg"
+	"github.com/gotd/td/tgerr"
+	"github.com/rs/zerolog"
+	"go.mau.fi/zerozap"
+	"go.uber.org/zap"
+	"maunium.net/go/mautrix/bridgev2"
+)
+
+type qrAuthResult struct {
+	PasswordNeeded bool
+	Authorization  *tg.AuthAuthorization
+	Error          error
+}
+
+type QRLogin struct {
+	user       *bridgev2.User
+	main       *TelegramConnector
+	authData   UserLoginSession
+	authClient *telegram.Client
+
+	authClientCtx    context.Context
+	authClientCancel context.CancelFunc
+
+	auth    chan qrAuthResult
+	qrToken chan qrlogin.Token
+}
+
+const LoginStepIDShowQR = "fi.mau.telegram.login.show_qr"
+
+var _ bridgev2.LoginProcessDisplayAndWait = (*QRLogin)(nil) // For showing QR code
+var _ bridgev2.LoginProcessUserInput = (*QRLogin)(nil)      // For asking for password
+
+func (q *QRLogin) Cancel() {
+	q.authClientCancel()
+}
+
+func (q *QRLogin) Start(ctx context.Context) (*bridgev2.LoginStep, error) {
+	log := zerolog.Ctx(ctx).With().Str("component", "telegram_qr_login").Logger()
+	loggedIn := make(chan struct{})
+
+	dispatcher := tg.NewUpdateDispatcher()
+	dispatcher.OnLoginToken(func(ctx context.Context, e tg.Entities, update *tg.UpdateLoginToken) error {
+		loggedIn <- struct{}{}
+		return nil
+	})
+	zaplog := zap.New(zerozap.New(log))
+	updateManager := updates.New(updates.Config{
+		Handler: dispatcher,
+		Logger:  zaplog.Named("login_update_manager"),
+	})
+	q.authClient = telegram.NewClient(q.main.Config.APIID, q.main.Config.APIHash, telegram.Options{
+		CustomSessionStorage: &q.authData,
+		UpdateHandler:        updateManager,
+		Logger:               zaplog,
+	})
+
+	var err error
+	q.authClientCtx, q.authClientCancel, err = connectTelegramClient(log.WithContext(context.Background()), q.authClient)
+	if err != nil {
+		return nil, err
+	}
+
+	qr := qrlogin.NewQR(q.authClient.API(), q.main.Config.APIID, q.main.Config.APIHash, qrlogin.Options{
+		Migrate: q.authClient.MigrateTo,
+	})
+	q.qrToken = make(chan qrlogin.Token)
+	q.auth = make(chan qrAuthResult)
+	go func() {
+		auth, err := qr.Auth(q.authClientCtx, loggedIn, func(ctx context.Context, token qrlogin.Token) error {
+			q.qrToken <- token
+			return nil
+		})
+
+		q.auth <- qrAuthResult{false, auth, err}
+	}()
+
+	// Wait for the first QR token and show it to the user.:
+	select {
+	case token := <-q.qrToken:
+		return &bridgev2.LoginStep{
+			Type:         bridgev2.LoginStepTypeDisplayAndWait,
+			StepID:       LoginStepIDShowQR,
+			Instructions: "Scan the QR code on your phone to log in",
+			DisplayAndWaitParams: &bridgev2.LoginDisplayAndWaitParams{
+				Type: bridgev2.LoginDisplayTypeQR,
+				Data: token.URL(),
+			},
+		}, nil
+	case <-ctx.Done():
+		q.Cancel()
+		return nil, ctx.Err()
+	case <-q.authClientCtx.Done():
+		return nil, q.authClientCtx.Err()
+	}
+}
+
+func (q *QRLogin) Wait(ctx context.Context) (*bridgev2.LoginStep, error) {
+	if q.qrToken == nil {
+		panic("qr token channel is nil")
+	}
+
+	select {
+	case token := <-q.qrToken:
+		// There's a new token, show it to the user.
+		return &bridgev2.LoginStep{
+			Type:         bridgev2.LoginStepTypeDisplayAndWait,
+			StepID:       LoginStepIDShowQR,
+			Instructions: "Scan the QR code on your phone to log in",
+			DisplayAndWaitParams: &bridgev2.LoginDisplayAndWaitParams{
+				Type: bridgev2.LoginDisplayTypeQR,
+				Data: token.URL(),
+			},
+		}, nil
+	case authResult := <-q.auth:
+		if tgerr.Is(authResult.Error, "SESSION_PASSWORD_NEEDED") {
+			return &bridgev2.LoginStep{
+				Type:         bridgev2.LoginStepTypeUserInput,
+				StepID:       LoginStepIDPassword,
+				Instructions: "Please enter your password",
+				UserInputParams: &bridgev2.LoginUserInputParams{
+					Fields: []bridgev2.LoginInputDataField{
+						{
+							Type: bridgev2.LoginInputFieldTypePassword,
+							ID:   LoginStepIDPassword,
+							Name: "Password",
+						},
+					},
+				},
+			}, nil
+		} else if authResult.Error != nil {
+			return nil, fmt.Errorf("failed to authenticate: %w", authResult.Error)
+		}
+
+		// Stop the login client
+		q.authClientCancel()
+
+		return finalizeLogin(ctx, q.user, authResult.Authorization, UserLoginMetadata{
+			Session: q.authData,
+		})
+	case <-ctx.Done():
+		q.Cancel()
+		return nil, ctx.Err()
+	case <-q.authClientCtx.Done():
+		return nil, q.authClientCtx.Err()
+	}
+}
+
+func (q *QRLogin) SubmitUserInput(ctx context.Context, input map[string]string) (*bridgev2.LoginStep, error) {
+	password, ok := input[LoginStepIDPassword]
+	if !ok {
+		return nil, fmt.Errorf("unexpected state during phone login")
+	}
+	authorization, err := q.authClient.Auth().Password(ctx, password)
+	if err != nil {
+		return nil, fmt.Errorf("failed to submit password: %w", err)
+	}
+
+	// Stop the login client
+	q.authClientCancel()
+
+	return finalizeLogin(ctx, q.user, authorization, UserLoginMetadata{
+		Session: q.authData,
+	})
+}