diff --git a/telegram/client.go b/telegram/client.go index 854ad993c..7476c001b 100644 --- a/telegram/client.go +++ b/telegram/client.go @@ -128,6 +128,9 @@ type Client struct { // Tracing. tracer trace.Tracer + + // onTransfer is called in transfer. + onTransfer AuthTransferHandler } // NewClient creates new unstarted client. @@ -161,6 +164,7 @@ func NewClient(appID int, appHash string, opt Options) *Client { migrationTimeout: opt.MigrationTimeout, noUpdatesMode: opt.NoUpdates, mw: opt.Middlewares, + onTransfer: opt.OnTransfer, } if opt.TracerProvider != nil { client.tracer = opt.TracerProvider.Tracer(oteltg.Name) diff --git a/telegram/client_test.go b/telegram/client_test.go index a2890d42c..f6c8069cc 100644 --- a/telegram/client_test.go +++ b/telegram/client_test.go @@ -86,6 +86,7 @@ func newTestClient(h testHandler) *Client { ctx: context.Background(), cancel: func() {}, updateHandler: UpdateHandlerFunc(func(ctx context.Context, u tg.UpdatesClass) error { return nil }), + onTransfer: noopOnTransfer, } client.init() diff --git a/telegram/options.go b/telegram/options.go index 0ac20a521..af80cb06d 100644 --- a/telegram/options.go +++ b/telegram/options.go @@ -94,6 +94,10 @@ type Options struct { // OpenTelemetry. TracerProvider trace.TracerProvider + + // OnTransfer is called during authorization transfer. + // See [AuthTransferHandler] for details. + OnTransfer AuthTransferHandler } func (opt *Options) setDefaults() { @@ -140,6 +144,9 @@ func (opt *Options) setDefaults() { return nil }) } + if opt.OnTransfer == nil { + opt.OnTransfer = noopOnTransfer + } } func defaultBackoff(c clock.Clock) func() backoff.BackOff { diff --git a/telegram/transfer.go b/telegram/transfer.go index b6e7881d9..ec7722834 100644 --- a/telegram/transfer.go +++ b/telegram/transfer.go @@ -17,20 +17,39 @@ func (c *Client) exportAuth(ctx context.Context, dcID int) (*tg.AuthExportedAuth return export, nil } +// AuthTransferHandler is a function that is called during authorization transfer. +// +// The fn callback should be serialized by user id via external locking. +// You can call [Client.Self] to acquire current user id. +// +// The fn callback must return fn error if any. +type AuthTransferHandler func(ctx context.Context, client *Client, fn func(context.Context) error) error + +func noopOnTransfer(ctx context.Context, _ *Client, fn func(context.Context) error) error { + return fn(ctx) +} + // transfer exports current authorization and imports it to another DC. // See https://core.telegram.org/api/datacenter#authorization-transfer. func (c *Client) transfer(ctx context.Context, to *tg.Client, dc int) (tg.AuthAuthorizationClass, error) { - auth, err := c.exportAuth(ctx, dc) - if err != nil { - return nil, errors.Wrapf(err, "export to %d", dc) + var out tg.AuthAuthorizationClass + if err := c.onTransfer(ctx, c, func(ctx context.Context) error { + auth, err := c.exportAuth(ctx, dc) + if err != nil { + return errors.Wrapf(err, "export to %d", dc) + } + + req := &tg.AuthImportAuthorizationRequest{} + req.FillFrom(auth) + r, err := to.AuthImportAuthorization(ctx, req) + if err != nil { + return errors.Wrapf(err, "import from %d", dc) + } + + out = r + return nil + }); err != nil { + return nil, errors.Wrap(err, "onTransfer") } - - req := &tg.AuthImportAuthorizationRequest{} - req.FillFrom(auth) - r, err := to.AuthImportAuthorization(ctx, req) - if err != nil { - return nil, errors.Wrapf(err, "import from %d", dc) - } - - return r, nil + return out, nil }