Merge pull request #1302 from gotd/feat/client-add-on-transfer

feat(clint): add OnTransfer

Author: ernado

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)

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()
 

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 {

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
 }