fix(e2e): self-hosted Telegram server (#757)

.github/workflows/master.yml

@@ -37,7 +37,21 @@ jobs:
         with:
           version: v1.54
           working-directory: ${{ matrix.directory }}
-  build-and-test:
+  unit-test:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout
+        uses: actions/checkout@v4
+      - name: Setup Golang env
+        uses: actions/setup-go@v5
+        with:
+          go-version-file: go.mod
+          cache: true
+      - name: Build
+        run: go build
+      - name: Unit Test
+        run: go test -v $(go list ./... | grep -v /test) # skip e2e test
+  e2e-test:
     runs-on: ubuntu-latest
     steps:
       - name: Checkout
@@ -49,11 +63,14 @@ jobs:
           cache: true
       - name: Install Ginkgo
         run: go install github.com/onsi/ginkgo/v2/ginkgo
+      - name: Setup Teamgram Env
+        run: |
+          git clone https://github.com/iyear/teamgram-server.git
+          cd teamgram-server
+          git checkout 3cc9864cda9a4eb45b61542494dfe517bf643372
+          sudo docker compose -f ./docker-compose-env.yaml up -d --quiet-pull
+          sudo docker compose up -d --quiet-pull
       - name: Build
         run: go build
-      - name: Unit Test
-        run: go test -v $(go list ./... | grep -v /test) # skip e2e test
       - name: E2E Test
         run: ginkgo -v -r ./test
-        env:
-          TDL_TEST: ${{ secrets.TEST_SESSION }}

app/login/code.go

@@ -32,7 +32,6 @@ func Code(ctx context.Context) error {
 		Proxy:            viper.GetString(consts.FlagProxy),
 		NTP:              viper.GetString(consts.FlagNTP),
 		ReconnectTimeout: viper.GetDuration(consts.FlagReconnectTimeout),
-		Test:             viper.GetString(consts.FlagTest),
 		UpdateHandler:    nil,
 	}, true)
 	if err != nil {

app/login/qr.go

@@ -38,7 +38,6 @@ func QR(ctx context.Context) error {
 		Proxy:            viper.GetString(consts.FlagProxy),
 		NTP:              viper.GetString(consts.FlagNTP),
 		ReconnectTimeout: viper.GetDuration(consts.FlagReconnectTimeout),
-		Test:             viper.GetString(consts.FlagTest),
 		UpdateHandler:    d,
 	}, true)
 	if err != nil {

cmd/root.go

@@ -133,9 +133,6 @@ func New() *cobra.Command {
 	cmd.PersistentFlags().String(consts.FlagNTP, "", "ntp server host, if not set, use system time")
 	cmd.PersistentFlags().Duration(consts.FlagReconnectTimeout, 5*time.Minute, "Telegram client reconnection backoff timeout, infinite if set to 0") // #158
 
-	cmd.PersistentFlags().String(consts.FlagTest, "", "use test Telegram client, only for developer")
-	_ = cmd.PersistentFlags().MarkHidden(consts.FlagTest)
-
 	// completion
 	_ = cmd.RegisterFlagCompletionFunc(consts.FlagNamespace, func(cmd *cobra.Command, args []string, toComplete string) ([]string, cobra.ShellCompDirective) {
 		engine := kv.From(cmd.Context())
@@ -183,7 +180,6 @@ func tRun(ctx context.Context, f func(ctx context.Context, c *telegram.Client, k
 		Proxy:            viper.GetString(consts.FlagProxy),
 		NTP:              viper.GetString(consts.FlagNTP),
 		ReconnectTimeout: viper.GetDuration(consts.FlagReconnectTimeout),
-		Test:             viper.GetString(consts.FlagTest),
 		UpdateHandler:    nil,
 	}
 

core/dcpool/dcpool.go

@@ -13,6 +13,13 @@ import (
 	"github.com/iyear/tdl/core/middlewares/takeout"
 )
 
+var testMode = false
+
+// EnableTestMode enables test mode, which disables takeout and pooling and directly returns original client.
+func EnableTestMode() {
+	testMode = true
+}
+
 type Pool interface {
 	Client(ctx context.Context, dc int) *tg.Client
 	Takeout(ctx context.Context, dc int) *tg.Client
@@ -55,6 +62,12 @@ func (p *pool) Client(ctx context.Context, dc int) *tg.Client {
 }
 
 func (p *pool) invoker(ctx context.Context, dc int) tg.Invoker {
+	// self-hosted Telegram server can't properly handle pooling connections,
+	// so directly return original client
+	if testMode {
+		return p.api
+	}
+
 	if i, ok := p.invokers[dc]; ok {
 		return i
 	}

core/tclient/tclient.go

@@ -9,13 +9,11 @@ import (
 	"github.com/go-faster/errors"
 	"github.com/gotd/contrib/clock"
 	"github.com/gotd/contrib/middleware/floodwait"
-	"github.com/gotd/contrib/middleware/ratelimit"
 	tdclock "github.com/gotd/td/clock"
-	"github.com/gotd/td/session"
+	"github.com/gotd/td/exchange"
 	"github.com/gotd/td/telegram"
 	"github.com/gotd/td/telegram/dcs"
 	"golang.org/x/net/proxy"
-	"golang.org/x/time/rate"
 
 	"github.com/iyear/tdl/core/logctx"
 	"github.com/iyear/tdl/core/middlewares/recovery"
@@ -24,6 +22,13 @@ import (
 	"github.com/iyear/tdl/core/util/tutil"
 )
 
+// dc values can be overridden globally
+var (
+	DCList     dcs.List
+	DC         int
+	PublicKeys []exchange.PublicKey
+)
+
 type Options struct {
 	AppID            int
 	AppHash          string
@@ -32,7 +37,6 @@ type Options struct {
 	Proxy            string
 	NTP              string
 	ReconnectTimeout time.Duration
-	Test             string
 	UpdateHandler    telegram.UpdateHandler
 }
 
@@ -66,6 +70,9 @@ func New(ctx context.Context, o Options) (*telegram.Client, error) {
 		ReconnectionBackoff: func() backoff.BackOff {
 			return newBackoff(o.ReconnectTimeout)
 		},
+		DC:             DC,
+		DCList:         DCList,
+		PublicKeys:     PublicKeys,
 		UpdateHandler:  o.UpdateHandler,
 		Device:         tutil.Device,
 		SessionStorage: o.Session,
@@ -77,18 +84,6 @@ func New(ctx context.Context, o Options) (*telegram.Client, error) {
 		Logger:         logctx.From(ctx).Named("td"),
 	}
 
-	// test account session
-	if o.Test != "" {
-		storage := &session.StorageMemory{}
-		if err := storage.StoreSession(ctx, []byte(o.Test)); err != nil {
-			return nil, errors.Wrap(err, "store test session")
-		}
-		opts.SessionStorage = storage // hook original session storage
-
-		// add rate limit to avoid frequent flood wait
-		opts.Middlewares = append(opts.Middlewares, ratelimit.New(rate.Every(100*time.Millisecond), 5))
-	}
-
 	return telegram.NewClient(o.AppID, o.AppHash, opts), nil
 }
 

pkg/consts/flag.go

@@ -13,5 +13,4 @@ const (
 	FlagNTP              = "ntp"
 	FlagReconnectTimeout = "reconnect-timeout"
 	FlagDlTemplate       = "template"
-	FlagTest             = "test"
 )

pkg/tclient/tclient.go

@@ -18,7 +18,6 @@ type Options struct {
 	Proxy            string
 	NTP              string
 	ReconnectTimeout time.Duration
-	Test             string
 	UpdateHandler    telegram.UpdateHandler
 }
 
@@ -40,7 +39,6 @@ func New(ctx context.Context, o Options, login bool, middlewares ...telegram.Mid
 		Proxy:            o.Proxy,
 		NTP:              o.NTP,
 		ReconnectTimeout: o.ReconnectTimeout,
-		Test:             o.Test,
 		UpdateHandler:    o.UpdateHandler,
 	})
 }

test/suite_test.go

@@ -1,19 +1,20 @@
 package test
 
 import (
+	"context"
+	_ "embed"
 	"fmt"
 	"io"
 	"log"
+	"math/rand"
 	"os"
-	"path/filepath"
-	"strconv"
 	"testing"
-	"time"
 
 	"github.com/fatih/color"
 	"github.com/spf13/cobra"
 
 	tcmd "github.com/iyear/tdl/cmd"
+	"github.com/iyear/tdl/test/testserver"
 
 	. "github.com/onsi/ginkgo/v2"
 	. "github.com/onsi/gomega"
@@ -25,25 +26,23 @@ func TestCommand(t *testing.T) {
 }
 
 var (
-	cmd     *cobra.Command
-	args    []string
-	output  string
-	storage string
+	cmd         *cobra.Command
+	args        []string
+	output      string
+	testAccount string
+	sessionFile string
 )
 
-var _ = BeforeSuite(func() {
-	// used to avoid "open db: timeout" conflict
-	storage = fmt.Sprintf("type=file,path=%s",
-		filepath.Join(os.TempDir(), "tdl", strconv.FormatInt(time.Now().UnixNano(), 10)))
+var _ = BeforeSuite(func(ctx context.Context) {
+	var err error
+	testAccount, sessionFile, err = testserver.Setup(ctx, rand.NewSource(GinkgoRandomSeed()))
+	Expect(err).To(Succeed())
 
 	log.SetOutput(GinkgoWriter)
 })
 
 var _ = BeforeEach(func() {
 	cmd = tcmd.New()
-
-	// wait before each test to avoid rate limit
-	time.Sleep(10 * time.Second)
 })
 
 func exec(cmd *cobra.Command, args []string, success bool) {
@@ -54,7 +53,9 @@ func exec(cmd *cobra.Command, args []string, success bool) {
 
 	log.Printf("args: %s\n", args)
 	cmd.SetArgs(append([]string{
-		"--storage", storage,
+		"-s", "131072", // self-hosted Telegram server don't support 1MiB
+		"-n", testAccount,
+		"--storage", fmt.Sprintf("type=file,path=%s", sessionFile),
 	}, args...))
 	if err = cmd.Execute(); success {
 		Expect(err).To(Succeed())

test/testserver/public_key.pem

@@ -0,0 +1,8 @@
+-----BEGIN RSA PUBLIC KEY-----
+MIIBCgKCAQEAvKLEOWTzt9Hn3/9Kdp/RdHcEhzmd8xXeLSpHIIzaXTLJDw8BhJy1
+jR/iqeG8Je5yrtVabqMSkA6ltIpgylH///FojMsX1BHu4EPYOXQgB0qOi6kr08iX
+ZIH9/iOPQOWDsL+Lt8gDG0xBy+sPe/2ZHdzKMjX6O9B4sOsxjFrk5qDoWDrioJor
+AJ7eFAfPpOBf2w73ohXudSrJE0lbQ8pCWNpMY8cB9i8r+WBitcvouLDAvmtnTX7a
+khoDzmKgpJBYliAY4qA73v7u5UIepE8QgV0jCOhxJCPubP8dg+/PlLLVKyxU5Cdi
+QtZj2EMy4s9xlNKzX8XezE0MHEa6bQpnFwIDAQAB
+-----END RSA PUBLIC KEY-----

test/testserver/testserver.go

@@ -0,0 +1,141 @@
+package testserver
+
+import (
+	"context"
+	_ "embed"
+	"log"
+	"math/rand"
+	"os"
+	"path/filepath"
+	"strconv"
+
+	"github.com/go-faster/errors"
+	"github.com/gotd/td/crypto"
+	"github.com/gotd/td/exchange"
+	"github.com/gotd/td/telegram"
+	"github.com/gotd/td/telegram/auth"
+	"github.com/gotd/td/telegram/dcs"
+	"github.com/gotd/td/tg"
+
+	"github.com/iyear/tdl/core/dcpool"
+	tclientcore "github.com/iyear/tdl/core/tclient"
+	"github.com/iyear/tdl/pkg/kv"
+	"github.com/iyear/tdl/pkg/storage"
+	"github.com/iyear/tdl/pkg/tclient"
+)
+
+//go:embed public_key.pem
+var publicKeyData []byte
+
+var (
+	dc     = 1
+	dcList = dcs.List{
+		Options: []tg.DCOption{
+			{
+				ID:        1,
+				IPAddress: "127.0.0.1",
+				Port:      10443,
+			},
+		},
+		Domains: nil,
+		Test:    false,
+	}
+	publicKeys []exchange.PublicKey
+	phone      = "+86 13858528382"
+)
+
+func init() {
+	keys, _ := crypto.ParseRSAPublicKeys(publicKeyData)
+	for _, k := range keys {
+		publicKeys = append(publicKeys, exchange.PublicKey{RSA: k})
+	}
+}
+
+// Setup creates test user and returns account and session file path. Namespace is the value of account.
+func Setup(ctx context.Context, rnd rand.Source) (account string, sessionFile string, _ error) {
+	tclientcore.DC = dc
+	tclientcore.DCList = dcList
+	tclientcore.PublicKeys = publicKeys
+
+	dcpool.EnableTestMode()
+
+	account = strconv.FormatInt(rand.Int63(), 10)
+	sessionFile = filepath.Join(os.TempDir(), "tdl", account)
+
+	return account, sessionFile, setupTestUser(ctx, rand.New(rnd), account, sessionFile)
+}
+
+func setupTestUser(ctx context.Context, rnd *rand.Rand, account, sessionFile string) error {
+	kvd, err := kv.New(kv.DriverFile, map[string]any{
+		"path": sessionFile,
+	})
+	if err != nil {
+		return errors.Wrapf(err, "create kv storage: %s", sessionFile)
+	}
+	log.Printf("session file: %s", sessionFile)
+
+	stg, err := kvd.Open(account)
+	if err != nil {
+		return errors.Wrap(err, "open test namespace")
+	}
+
+	sess := storage.NewSession(stg, true)
+
+	opts := telegram.Options{
+		DC:             dc,
+		DCList:         dcList,
+		PublicKeys:     publicKeys,
+		SessionStorage: sess,
+	}
+
+	app := tclient.Apps[tclient.AppDesktop]
+	c := telegram.NewClient(app.AppID, app.AppHash, opts)
+
+	if err = c.Run(ctx, func(ctx context.Context) error {
+		if err = c.Ping(ctx); err != nil {
+			return err
+		}
+
+		authClient := auth.NewClient(c.API(), rnd, app.AppID, app.AppHash)
+
+		if err = auth.NewFlow(
+			testAuth{phone: phone},
+			auth.SendCodeOptions{},
+		).Run(ctx, authClient); err != nil {
+			return errors.Wrap(err, "register test user")
+		}
+
+		user, err := c.Self(ctx)
+		if err != nil {
+			return errors.Wrap(err, "get self")
+		}
+
+		log.Printf("user: %v, %v, %v", user.ID, user.FirstName, user.LastName)
+		return nil
+	}); err != nil {
+		return errors.Wrap(err, "run auth")
+	}
+
+	return nil
+}
+
+type testAuth struct {
+	phone string
+}
+
+func (t testAuth) Phone(_ context.Context) (string, error)    { return t.phone, nil }
+func (t testAuth) Password(_ context.Context) (string, error) { return "", auth.ErrPasswordNotProvided }
+func (t testAuth) Code(_ context.Context, _ *tg.AuthSentCode) (string, error) {
+	return "12345", nil
+}
+
+func (t testAuth) AcceptTermsOfService(_ context.Context, _ tg.HelpTermsOfService) error {
+	return nil
+}
+
+func (t testAuth) SignUp(_ context.Context) (auth.UserInfo, error) {
+	return auth.UserInfo{
+		FirstName: "Test",
+		LastName:  "User",
+	}, nil
+}