E2E call DH initialization test

call_test.go

package e2etest

import (
	"context"
	"crypto/sha1"
	"crypto/sha256"
	"encoding/binary"
	"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/tdsync"
	"github.com/gotd/td/telegram/message"
	"github.com/gotd/td/tg"
)

var protocol = tg.PhoneCallProtocol{
	UDPReflector:    true,
	MinLayer:        65,
	MaxLayer:        92,
	LibraryVersions: []string{"2.4.4"},
}

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

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

	accepted := make(chan tg.PhoneCallClass)
	d.OnPhoneCall(func(ctx context.Context, _ tg.Entities, update *tg.UpdatePhoneCall) error {
		select {
		case <-ctx.Done():
			return ctx.Err()
		case accepted <- update.PhoneCall:
			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)
		}

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

		h := sha256.New()
		g := big.NewInt(int64(dhCfg.G))
		dhPrime := big.NewInt(0).SetBytes(dhCfg.P)
		gA := big.NewInt(0).Exp(g, a, dhPrime)
		h.Write(gA.Bytes())

		req, err := raw.PhoneRequestCall(ctx, &tg.PhoneRequestCallRequest{
			Video:    false,
			UserID:   user,
			RandomID: int(a.Int64()),
			GAHash:   h.Sum(nil),
			Protocol: protocol,
		})
		if err != nil {
			return xerrors.Errorf("request call: %w", err)
		}
		c := req.PhoneCall

		log.Info("Sent request")

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

		gB := big.NewInt(0).SetBytes(call.GB)
		key := make([]byte, 256)
		key = big.NewInt(0).Exp(gB, a, dhPrime).FillBytes(key[:])
		keyHash := sha1.Sum(key)

		req, err = raw.PhoneConfirmCall(ctx, &tg.PhoneConfirmCallRequest{
			Peer: tg.InputPhoneCall{
				ID:         call.ID,
				AccessHash: call.AccessHash,
			},
			GA:             gA.Bytes(),
			KeyFingerprint: int64(binary.LittleEndian.Uint64(keyHash[12:])),
			Protocol:       protocol,
		})
		if err != nil {
			return xerrors.Errorf("confirm call: %w", err)
		}

		pp.Println("request", req.PhoneCall)
		return nil
	})
}

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

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

	accepted := make(chan *tg.PhoneCall)
	d.OnPhoneCall(func(ctx context.Context, _ tg.Entities, update *tg.UpdatePhoneCall) error {

		switch req := update.PhoneCall.(type) {
		case *tg.PhoneCall:
			select {
			case <-ctx.Done():
				return ctx.Err()
			case accepted <- req:
				return nil
			}
		case *tg.PhoneCallRequested:
			b, dhCfg, err := initDH(ctx, raw)
			if err != nil {
				return xerrors.Errorf("init DH: %w", err)
			}

			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)

			_, err = raw.PhoneAcceptCall(ctx, &tg.PhoneAcceptCallRequest{
				Peer: tg.InputPhoneCall{
					ID:         req.ID,
					AccessHash: req.AccessHash,
				},
				GB:       gB.Bytes(),
				Protocol: protocol,
			})
			if err != nil {
				return xerrors.Errorf("accept: %w", err)
			}
		}

		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:
			pp.Println("accepted", c)
			return nil
		}
	})
}

func TestCall(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 caller(ctx, suite, usernameCh)
	})
	grp.Go(func(ctx context.Context) error {
		return accepter(ctx, suite, usernameCh)
	})

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