E2E encryprtion test

e2e_test.go

package e2etest

import (
	"context"
	"crypto/rand"
	"crypto/sha1"
	"encoding/binary"
	"io"
	"math/big"
	"strconv"
	"testing"

	"github.com/k0kubun/pp/v3"
	"go.uber.org/zap/zapcore"
	"go.uber.org/zap/zaptest"
	"golang.org/x/xerrors"

	"github.com/gotd/td/internal/crypto"
	"github.com/gotd/td/internal/tdsync"
	"github.com/gotd/td/telegram/message"
	"github.com/gotd/td/tg"
)

var randMax = big.NewInt(0).SetBit(big.NewInt(0), crypto.RSAKeyBits, 1)

func generateBig(r io.Reader) (*big.Int, error) {
	a, err := rand.Int(r, randMax)
	if err != nil {
		return nil, err
	}
	return a, nil
}

func initDH(ctx context.Context, raw *tg.Client) (*big.Int, *tg.MessagesDhConfig, error) {
	d, err := raw.MessagesGetDhConfig(ctx, &tg.MessagesGetDhConfigRequest{
		RandomLength: crypto.RSAKeyBits,
	})
	if err != nil {
		return nil, nil, xerrors.Errorf("get DH config: %w", err)
	}

	dhCfg, ok := d.AsModified()
	if !ok {
		return nil, nil, xerrors.Errorf("unexpected type %T", d)
	}

	a, err := generateBig(rand.Reader)
	if err != nil {
		return nil, nil, xerrors.Errorf("generate random: %w", err)
	}

	return a, dhCfg, nil
}

func requestEncryptedChat(ctx context.Context, raw *tg.Client, user tg.InputUserClass) (tg.EncryptedChatClass, error) {
	a, dhCfg, err := initDH(ctx, raw)
	if err != nil {
		return nil, xerrors.Errorf("init DH: %w", err)
	}

	g := big.NewInt(int64(dhCfg.G))
	dhPrime := big.NewInt(0).SetBytes(dhCfg.P)
	gA := big.NewInt(0).Exp(g, a, dhPrime)

	c, err := raw.MessagesRequestEncryption(ctx, &tg.MessagesRequestEncryptionRequest{
		UserID:   user,
		RandomID: int(a.Int64()),
		GA:       gA.Bytes(),
	})
	if err != nil {
		return nil, xerrors.Errorf("request chat: %w", err)
	}

	return c, nil
}

func requester(ctx context.Context, suite *Suite, usernameCh <-chan string) error {
	log := suite.Logger().Named("requester")
	d := tg.NewUpdateDispatcher()

	client := suite.Client(log, d)
	raw := client.API()
	sender := message.NewSender(raw)

	accepted := make(chan tg.EncryptedChatClass)
	d.OnEncryption(func(ctx context.Context, _ tg.Entities, update *tg.UpdateEncryption) error {
		select {
		case <-ctx.Done():
			return ctx.Err()
		case accepted <- update.Chat:
			return nil
		}
	})

	return client.Run(ctx, func(ctx context.Context) error {
		if err := suite.RetryAuthenticate(ctx, client.Auth()); err != nil {
			return xerrors.Errorf("authenticate: %w", err)
		}

		var username string
		select {
		case <-ctx.Done():
			return ctx.Err()
		case u := <-usernameCh:
			username = u
		}

		user, err := sender.Resolve(username).AsInputUser(ctx)
		if err != nil {
			return xerrors.Errorf("resolve %q: %w", username, err)
		}

		c, err := requestEncryptedChat(ctx, raw, user)
		if err != nil {
			return xerrors.Errorf("request: %w", err)
		}

		log.Info("Sent request")

	dispatch:
		var chat *tg.EncryptedChat
		switch t := c.(type) {
		case *tg.EncryptedChatWaiting:
			log.Info("Waiting for response")
			select {
			case <-ctx.Done():
				return ctx.Err()
			case c = <-accepted:
				goto dispatch
			}
		case *tg.EncryptedChat:
			chat = t
		default:
			return xerrors.Errorf("unexpected type %T", t)
		}

		pp.Println("requested", chat)
		return nil
	})
}

func receiver(ctx context.Context, suite *Suite, usernameCh chan<- string) error {
	log := suite.Logger().Named("receiver")
	d := tg.NewUpdateDispatcher()

	client := suite.Client(log, d)
	raw := client.API()

	accepted := make(chan tg.EncryptedChatClass)
	d.OnEncryption(func(ctx context.Context, _ tg.Entities, update *tg.UpdateEncryption) error {
		req, ok := update.Chat.(*tg.EncryptedChatRequested)
		if !ok {
			return nil
		}

		b, dhCfg, err := initDH(ctx, raw)
		if err != nil {
			return xerrors.Errorf("init DH: %w", err)
		}

		gA := big.NewInt(0).SetBytes(req.GA)
		g := big.NewInt(int64(dhCfg.G))
		dhPrime := big.NewInt(0).SetBytes(dhCfg.P)

		// g_b := pow(g, b) mod dh_prime
		gB := big.NewInt(0).Exp(g, b, dhPrime)

		// key := pow(g_a, b) mod dh_prime
		key := make([]byte, 256)
		key = big.NewInt(0).Exp(gA, b, dhPrime).FillBytes(key[:])
		keyHash := sha1.Sum(key)

		log.Info("Accept chat")
		c, err := raw.MessagesAcceptEncryption(ctx, &tg.MessagesAcceptEncryptionRequest{
			Peer: tg.InputEncryptedChat{
				ChatID:     req.ID,
				AccessHash: req.AccessHash,
			},
			GB:             gB.Bytes(),
			KeyFingerprint: int64(binary.LittleEndian.Uint64(keyHash[12:])),
		})
		if err != nil {
			return xerrors.Errorf("accept: %w", err)
		}

		select {
		case <-ctx.Done():
			return ctx.Err()
		case accepted <- c:
			return nil
		}
	})

	return client.Run(ctx, func(ctx context.Context) error {
		if err := suite.RetryAuthenticate(ctx, client.Auth()); err != nil {
			return xerrors.Errorf("authenticate: %w", err)
		}

		var me *tg.User
		if err := retryFloodWait(ctx, func() (err error) {
			me, err = client.Self(ctx)
			return err
		}); err != nil {
			return err
		}

		expectedUsername := "echobot" + strconv.Itoa(me.ID)
		user, err := EnsureUsername(ctx, client, expectedUsername)
		if err != nil {
			return xerrors.Errorf("ensure username: %w", err)
		}
		select {
		case <-ctx.Done():
			return ctx.Err()
		case usernameCh <- user.Username:
		}

		select {
		case <-ctx.Done():
			return ctx.Err()
		case c := <-accepted:
			chat, ok := c.(*tg.EncryptedChat)
			if !ok {
				return nil
			}

			pp.Println("accepted", chat)
			return nil
		}
	})
}

func TestE2E(t *testing.T) {
	logger := zaptest.NewLogger(t, zaptest.Level(zapcore.InfoLevel))
	suite := NewSuite(t, TestOptions{
		Logger: logger,
	})

	ctx := context.Background()
	grp := tdsync.NewCancellableGroup(ctx)

	usernameCh := make(chan string)
	grp.Go(func(ctx context.Context) error {
		return requester(ctx, suite, usernameCh)
	})
	grp.Go(func(ctx context.Context) error {
		return receiver(ctx, suite, usernameCh)
	})

	if err := grp.Wait(); err != nil {
		t.Error(err)
	}
}