mtproto: add compat transport quick ack support

(cherry picked from commit d051bc37bd14076fdd0a83ad41cd507929b20ece)

internal/mtprotoedge/conn.go

@@ -57,6 +57,9 @@ type Conn struct {
 
 	// sentContentMessages 只由 outbound actor 访问，用于生成 MTProto seq_no。
 	sentContentMessages int32
+	// outboundPlain/outboundWire 只由 outbound actor 访问，用于复用出站加密缓冲。
+	outboundPlain bin.Buffer
+	outboundWire  bin.Buffer
 
 	identityMu              sync.RWMutex
 	businessAuthKeyID       [8]byte

internal/mtprotoedge/encrypted.go

@@ -2,6 +2,8 @@ package mtprotoedge
 
 import (
 	"context"
+	"crypto/sha256"
+	"encoding/binary"
 	"encoding/hex"
 	"errors"
 	"fmt"
@@ -125,6 +127,9 @@ func (s *Server) handleEncrypted(ctx context.Context, tc transport.Conn, cs *con
 		)
 		return current, s.sendBadMsg(ctx, current, data.MessageID, data.SeqNo, code)
 	}
+	if err := sendQuickAckIfRequested(ctx, tc, key, data); err != nil {
+		return current, err
+	}
 
 	content := clientMessageNeedsAck(typeID)
 	if record, ok := cs.seenRecord(data.MessageID); ok {
@@ -170,6 +175,30 @@ func (s *Server) handleEncrypted(ctx context.Context, tc transport.Conn, cs *con
 	return current, nil
 }
 
+func sendQuickAckIfRequested(ctx context.Context, tc transport.Conn, key crypto.AuthKey, data *crypto.EncryptedMessageData) error {
+	q, ok := tc.(quickAckTransport)
+	if !ok || !q.ConsumeQuickAckRequested() {
+		return nil
+	}
+	token, err := clientQuickAckToken(key, data)
+	if err != nil {
+		return err
+	}
+	return q.SendQuickAck(ctx, token)
+}
+
+func clientQuickAckToken(key crypto.AuthKey, data *crypto.EncryptedMessageData) (uint32, error) {
+	var plain bin.Buffer
+	if err := data.Encode(&plain); err != nil {
+		return 0, err
+	}
+	h := sha256.New()
+	_, _ = h.Write(key.Value[88:120])
+	_, _ = h.Write(plain.Raw())
+	sum := h.Sum(nil)
+	return binary.LittleEndian.Uint32(sum[:4]) &^ quickAckResponseFlag, nil
+}
+
 // dispatch 处理一条明文消息：解包 container/gzip，处理服务消息，其余转 RPC 路由。
 // content-related 消息（ping、RPC）的 msg_id 会收集到 acks 以便统一确认。
 func (s *Server) dispatch(ctx context.Context, cs *connState, c *Conn, msgID int64, seqNo int32, b *bin.Buffer, acks *[]int64) error {
@@ -560,6 +589,9 @@ func validateClientEnvelope(now time.Time, msgID int64, seqNo int32, typeID uint
 	if msgTime.After(now.Add(30 * time.Second)) {
 		return badMsgIDTooHigh
 	}
+	if clientMessageAllowsEitherSeqParity(typeID) {
+		return 0
+	}
 	if clientMessageNeedsAck(typeID) {
 		if seqNo%2 == 0 {
 			return badMsgSeqNotOdd
@@ -593,6 +625,9 @@ func validateClientContainerEnvelope(msgID int64, seqNo int32, typeID uint32) in
 	if msgID == 0 || proto.MessageID(msgID).Type() != proto.MessageFromClient {
 		return badMsgIDInvalidBits
 	}
+	if clientMessageAllowsEitherSeqParity(typeID) {
+		return 0
+	}
 	if clientMessageNeedsAck(typeID) {
 		if seqNo%2 == 0 {
 			return badMsgSeqNotOdd
@@ -603,6 +638,15 @@ func validateClientContainerEnvelope(msgID int64, seqNo int32, typeID uint32) in
 	return 0
 }
 
+func clientMessageAllowsEitherSeqParity(typeID uint32) bool {
+	switch typeID {
+	case mt.PingDelayDisconnectRequestTypeID:
+		return true
+	default:
+		return false
+	}
+}
+
 func clientMessageNeedsAck(typeID uint32) bool {
 	switch typeID {
 	case proto.MessageContainerTypeID,

internal/mtprotoedge/encrypted_test.go

@@ -301,6 +301,32 @@ func TestPingDelayDisconnectEvenSeqAccepted(t *testing.T) {
 	}
 }
 
+func TestPingDelayDisconnectOddSeqAccepted(t *testing.T) {
+	const dc = 2
+	addr, pub, _ := startTestServer(t, Options{DC: dc})
+	conn, auth, cipher := dialHandshake(t, addr, dc, pub)
+
+	clientMsgID := proto.NewMessageIDGen(time.Now)
+	reqMsgID := clientMsgID.New(proto.MessageFromClient)
+	sendEncryptedWithSeq(t, conn, cipher, auth, reqMsgID, 1, &mt.PingDelayDisconnectRequest{
+		PingID:          10,
+		DisconnectDelay: 60,
+	})
+
+	replies := collectReplies(t, conn, cipher, auth.AuthKey, mt.PongTypeID)
+	if _, ok := replies[mt.BadMsgNotificationTypeID]; ok {
+		t.Fatalf("odd ping_delay_disconnect seq_no produced bad_msg_notification")
+	}
+	buf := mustHave(t, replies, mt.PongTypeID, "pong")
+	var pong mt.Pong
+	if err := pong.Decode(buf); err != nil {
+		t.Fatalf("decode pong: %v", err)
+	}
+	if pong.MsgID != reqMsgID || pong.PingID != 10 {
+		t.Fatalf("pong = %+v, want msg_id=%d ping_id=10", pong, reqMsgID)
+	}
+}
+
 // TestDestroyAuthKey 验证 MTProto service message destroy_auth_key 由连接层直接响应，
 // 避免 TDesktop 清理旧 key 时落到业务 RPC fallback。
 func TestDestroyAuthKey(t *testing.T) {

internal/mtprotoedge/outbound.go

@@ -2,10 +2,14 @@ package mtprotoedge
 
 import (
 	"context"
+	"crypto/aes"
 	"errors"
 	"fmt"
+	"io"
 	"time"
 
+	"github.com/gotd/ige"
+
 	"github.com/gotd/td/bin"
 	"github.com/gotd/td/crypto"
 	"github.com/gotd/td/mt"
@@ -45,12 +49,19 @@ type outboundOp struct {
 	ctx        context.Context
 	msgType    proto.MessageType
 	msg        bin.Encoder
+	encoded    *encodedOutboundMessage
 	ids        []int64
 	reqMsgID   int64
 	enqueuedAt time.Time
 	done       chan outboundResult
 }
 
+type encodedOutboundMessage struct {
+	body     []byte
+	typeID   uint32
+	reqMsgID int64
+}
+
 type outboundResult struct {
 	info   []byte
 	resent bool
@@ -119,6 +130,14 @@ func (c *Conn) SendPriority(ctx context.Context, t proto.MessageType, msg bin.En
 // SendBestEffort 只等待消息进入普通 outbound 队列，不等待网络写完成。
 // 用于 updates fanout：队列拥塞时返回 ErrOutboundQueueFull，durable outbox/getDifference 负责兜底。
 func (c *Conn) SendBestEffort(ctx context.Context, t proto.MessageType, msg bin.Encoder, timeout time.Duration) error {
+	return c.sendBestEffort(ctx, t, msg, nil, timeout)
+}
+
+func (c *Conn) SendBestEffortEncoded(ctx context.Context, t proto.MessageType, encoded *encodedOutboundMessage, timeout time.Duration) error {
+	return c.sendBestEffort(ctx, t, nil, encoded, timeout)
+}
+
+func (c *Conn) sendBestEffort(ctx context.Context, t proto.MessageType, msg bin.Encoder, encoded *encodedOutboundMessage, timeout time.Duration) error {
 	if c.outbound == nil || c.outboundControl == nil {
 		return ErrConnClosed
 	}
@@ -131,6 +150,7 @@ func (c *Conn) SendBestEffort(ctx context.Context, t proto.MessageType, msg bin.
 		ctx:        writeCtx,
 		msgType:    t,
 		msg:        msg,
+		encoded:    encoded,
 		enqueuedAt: time.Now(),
 	}
 	if timeout == 0 {
@@ -164,6 +184,14 @@ func (c *Conn) SendBestEffort(ctx context.Context, t proto.MessageType, msg bin.
 }
 
 func (c *Conn) send(ctx context.Context, t proto.MessageType, msg bin.Encoder, control bool) error {
+	return c.sendOutbound(ctx, t, msg, nil, control)
+}
+
+func (c *Conn) SendEncoded(ctx context.Context, t proto.MessageType, encoded *encodedOutboundMessage) error {
+	return c.sendOutbound(ctx, t, nil, encoded, false)
+}
+
+func (c *Conn) sendOutbound(ctx context.Context, t proto.MessageType, msg bin.Encoder, encoded *encodedOutboundMessage, control bool) error {
 	if c.outbound == nil || c.outboundControl == nil {
 		return ErrConnClosed
 	}
@@ -173,6 +201,7 @@ func (c *Conn) send(ctx context.Context, t proto.MessageType, msg bin.Encoder, c
 		ctx:        ctx,
 		msgType:    t,
 		msg:        msg,
+		encoded:    encoded,
 		enqueuedAt: time.Now(),
 		done:       make(chan outboundResult, 1),
 	}
@@ -391,7 +420,7 @@ func (c *Conn) handleOutboundOp(state *outboundState, op outboundOp) {
 }
 
 func (c *Conn) handleOutboundSend(state *outboundState, op outboundOp) {
-	frame, err := c.buildFrame(op.msgType, op.msg)
+	frame, err := c.buildFrame(op.msgType, op.msg, op.encoded)
 	if err == nil {
 		err = c.writeFrame(op.ctx, frame)
 	}
@@ -465,7 +494,26 @@ func (op outboundOp) finish(res outboundResult) {
 	}
 }
 
-func (c *Conn) buildFrame(t proto.MessageType, msg bin.Encoder) (*outboundFrame, error) {
+func (c *Conn) buildFrame(t proto.MessageType, msg bin.Encoder, encoded *encodedOutboundMessage) (*outboundFrame, error) {
+	if encoded == nil {
+		var err error
+		encoded, err = encodeOutboundMessage(msg)
+		if err != nil {
+			return nil, err
+		}
+	}
+	content := frameNeedsAck(encoded.typeID)
+	msgID := c.msgID.New(t)
+	return &outboundFrame{
+		msgID:    msgID,
+		seqNo:    c.nextSeqNo(content),
+		typeID:   encoded.typeID,
+		body:     encoded.body,
+		reqMsgID: encoded.reqMsgID,
+	}, nil
+}
+
+func encodeOutboundMessage(msg bin.Encoder) (*encodedOutboundMessage, error) {
 	if msg == nil {
 		return nil, errors.New("nil outbound message")
 	}
@@ -477,13 +525,9 @@ func (c *Conn) buildFrame(t proto.MessageType, msg bin.Encoder) (*outboundFrame,
 	if err != nil {
 		return nil, fmt.Errorf("peek outbound type id: %w", err)
 	}
-	content := frameNeedsAck(typeID)
+	return &encodedOutboundMessage{
-	msgID := c.msgID.New(t)
-	return &outboundFrame{
-		msgID:    msgID,
-		seqNo:    c.nextSeqNo(content),
 		typeID:   typeID,
-		body:     body.Copy(),
+		body:     body.Raw(),
 		reqMsgID: outboundRequestMsgID(msg),
 	}, nil
 }
@@ -501,15 +545,8 @@ func (c *Conn) writeFrame(ctx context.Context, frame *outboundFrame) error {
 	if ctx == nil {
 		ctx = context.Background()
 	}
-	var out bin.Buffer
+	out, err := c.encryptOutboundFrame(frame)
-	if err := c.cipher.Encrypt(c.key, crypto.EncryptedMessageData{
+	if err != nil {
-		Salt:                   c.salt,
-		SessionID:              c.sessionID,
-		MessageID:              frame.msgID,
-		SeqNo:                  frame.seqNo,
-		MessageDataLen:         int32(len(frame.body)),
-		MessageDataWithPadding: frame.body,
-	}, &out); err != nil {
 		return fmt.Errorf("encrypt: %w", err)
 	}
 
@@ -523,7 +560,7 @@ func (c *Conn) writeFrame(ctx context.Context, frame *outboundFrame) error {
 	if writer == nil {
 		writer = c.transport
 	}
-	if err := writer.Send(sendCtx, &out); err != nil {
+	if err := writer.Send(sendCtx, out); err != nil {
 		return fmt.Errorf("send: %w", err)
 	}
 	if frame.sentAt.IsZero() {
@@ -533,6 +570,61 @@ func (c *Conn) writeFrame(ctx context.Context, frame *outboundFrame) error {
 	return nil
 }
 
+func (c *Conn) encryptOutboundFrame(frame *outboundFrame) (*bin.Buffer, error) {
+	plain := &c.outboundPlain
+	plain.Reset()
+	plain.PutLong(c.salt)
+	plain.PutLong(c.sessionID)
+	plain.PutLong(frame.msgID)
+	plain.PutInt32(frame.seqNo)
+	plain.PutInt32(int32(len(frame.body)))
+	plain.Put(frame.body)
+
+	paddingOffset := plain.Len()
+	paddingLen := encryptedPaddingLen(paddingOffset)
+	growBinBufferLen(plain, paddingOffset+paddingLen)
+	if _, err := io.ReadFull(c.cipher.Rand(), plain.Buf[paddingOffset:]); err != nil {
+		return nil, err
+	}
+
+	msgKey := crypto.MessageKey(c.key.Value, plain.Raw(), crypto.Server)
+	key, iv := crypto.Keys(c.key.Value, msgKey, crypto.Server)
+	aesBlock, err := aes.NewCipher(key[:])
+	if err != nil {
+		return nil, err
+	}
+
+	wireLen := len(c.key.ID) + len(msgKey) + plain.Len()
+	wire := &c.outboundWire
+	ensureBinBufferLen(wire, wireLen)
+	copy(wire.Buf[:len(c.key.ID)], c.key.ID[:])
+	copy(wire.Buf[len(c.key.ID):len(c.key.ID)+len(msgKey)], msgKey[:])
+	ige.EncryptBlocks(aesBlock, iv[:], wire.Buf[len(c.key.ID)+len(msgKey):], plain.Raw())
+	return wire, nil
+}
+
+func encryptedPaddingLen(l int) int {
+	return 16 + (16 - (l % 16))
+}
+
+func ensureBinBufferLen(b *bin.Buffer, n int) {
+	if cap(b.Buf) < n {
+		b.Buf = make([]byte, n)
+		return
+	}
+	b.Buf = b.Buf[:n]
+}
+
+func growBinBufferLen(b *bin.Buffer, n int) {
+	if cap(b.Buf) < n {
+		next := make([]byte, n)
+		copy(next, b.Buf)
+		b.Buf = next
+		return
+	}
+	b.Buf = b.Buf[:n]
+}
+
 func frameNeedsAck(typeID uint32) bool {
 	switch typeID {
 	case mt.MsgsAckTypeID,

internal/mtprotoedge/outbound_test.go

@@ -1,17 +1,58 @@
 package mtprotoedge
 
 import (
+	"bytes"
 	"context"
+	"crypto/rand"
 	"sync"
 	"testing"
 	"time"
 
 	"github.com/gotd/td/bin"
+	"github.com/gotd/td/crypto"
 	"github.com/gotd/td/mt"
 	"github.com/gotd/td/proto"
 	"github.com/gotd/td/tg"
 )
 
+func TestEncryptOutboundFrameDecryptsWithGotdCipher(t *testing.T) {
+	var key crypto.Key
+	if _, err := rand.Read(key[:]); err != nil {
+		t.Fatalf("rand key: %v", err)
+	}
+	authKey := key.WithID()
+	body := mustEncodeTL(t, &mt.NewSessionCreated{
+		FirstMsgID: 111,
+		UniqueID:   222,
+		ServerSalt: 333,
+	})
+	c := &Conn{
+		cipher:    crypto.NewServerCipher(rand.Reader),
+		key:       authKey,
+		salt:      12345,
+		sessionID: 67890,
+	}
+	out, err := c.encryptOutboundFrame(&outboundFrame{
+		msgID:  7649066000000000001,
+		seqNo:  1,
+		typeID: mt.NewSessionCreatedTypeID,
+		body:   body,
+	})
+	if err != nil {
+		t.Fatalf("encrypt: %v", err)
+	}
+	data, err := crypto.NewClientCipher(rand.Reader).DecryptFromBuffer(authKey, &bin.Buffer{Buf: append([]byte(nil), out.Raw()...)})
+	if err != nil {
+		t.Fatalf("decrypt: %v", err)
+	}
+	if data.Salt != c.salt || data.SessionID != c.sessionID {
+		t.Fatalf("salt/session = %d/%d, want %d/%d", data.Salt, data.SessionID, c.salt, c.sessionID)
+	}
+	if got := data.Data(); !bytes.Equal(got, body) {
+		t.Fatalf("body = %x, want %x", got, body)
+	}
+}
+
 func TestOutboundActorSerializesConcurrentSends(t *testing.T) {
 	const dc = 2
 	addr, pub, srv := startTestServer(t, Options{DC: dc})

internal/mtprotoedge/server.go

@@ -221,7 +221,7 @@ func (s *Server) Serve(ctx context.Context, ln net.Listener) error {
 	if s.obfuscated {
 		transportListener = transport.ObfuscatedListener(ln)
 	}
-	l := transport.ListenCodec(s.codec, transportListener)
+	l := newCompatTransportListener(s.codec, transportListener)
 	s.log.Info("Serving", zap.String("addr", ln.Addr().String()), zap.Int("dc", s.dc), zap.Bool("obfuscated_tcp", s.obfuscated))
 	defer s.log.Info("Stopped")
 

internal/mtprotoedge/server_test.go

@@ -163,3 +163,63 @@ func TestServerAcceptObfuscatedAbridged(t *testing.T) {
 		t.Fatal("server did not stop after ctx cancel")
 	}
 }
+
+func TestServerAcceptObfuscatedAbridgedQuickAckFrame(t *testing.T) {
+	ln, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatalf("listen: %v", err)
+	}
+
+	frames := make(chan int, 1)
+	srv := New(Options{Logger: zaptest.NewLogger(t), ObfuscatedTCP: true})
+	srv.onFrame = func(n int) {
+		select {
+		case frames <- n:
+		default:
+		}
+	}
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	serveErr := make(chan error, 1)
+	go func() { serveErr <- srv.Serve(ctx, ln) }()
+
+	raw, err := net.Dial("tcp", ln.Addr().String())
+	if err != nil {
+		t.Fatalf("dial: %v", err)
+	}
+	obfs := obfuscator.Obfuscated2(rand.Reader, raw)
+	if err := obfs.Handshake((codec.Abridged{}).ObfuscatedTag(), 2, mtproxy.Secret{}); err != nil {
+		t.Fatalf("obfuscated handshake: %v", err)
+	}
+
+	var b bin.Buffer
+	b.PutInt32(0x12345678)
+	b.PutInt32(0x0badf00d)
+	packet := append([]byte{0x80 | byte(b.Len()/4)}, b.Raw()...)
+	if _, err := obfs.Write(packet); err != nil {
+		t.Fatalf("write quick ack frame: %v", err)
+	}
+
+	select {
+	case n := <-frames:
+		if n != b.Len() {
+			t.Fatalf("received frame len = %d, want %d", n, b.Len())
+		}
+	case <-time.After(5 * time.Second):
+		t.Fatal("server did not receive quick-ack abridged frame in time")
+	}
+
+	_ = raw.Close()
+
+	cancel()
+	select {
+	case err := <-serveErr:
+		if err != nil {
+			t.Fatalf("serve returned error: %v", err)
+		}
+	case <-time.After(5 * time.Second):
+		t.Fatal("server did not stop after ctx cancel")
+	}
+}

internal/mtprotoedge/session_manager.go

@@ -488,8 +488,16 @@ func (m *SessionManager) PushToUserExceptAuthKeySession(ctx context.Context, use
 }
 
 func (m *SessionManager) pushToUser(ctx context.Context, userID int64, excludeAuthKeyID *[8]byte, excludeSessionID int64, t proto.MessageType, msg bin.Encoder) (int, error) {
+	getEncoded := onceEncodedOutbound(msg)
 	return m.pushToUserWithSender(ctx, userID, excludeAuthKeyID, excludeSessionID, t, msg, func(c *Conn) error {
-		return c.Send(ctx, t, msg)
+		if c.outbound == nil || c.outboundControl == nil {
+			return ErrConnClosed
+		}
+		encoded, err := getEncoded()
+		if err != nil {
+			return err
+		}
+		return c.SendEncoded(ctx, t, encoded)
 	})
 }
 
@@ -502,11 +510,32 @@ func (m *SessionManager) PushToUserExceptAuthKeySessionBestEffort(ctx context.Co
 }
 
 func (m *SessionManager) pushToUserBestEffort(ctx context.Context, userID int64, excludeAuthKeyID *[8]byte, excludeSessionID int64, t proto.MessageType, msg bin.Encoder, timeout time.Duration) (int, error) {
+	getEncoded := onceEncodedOutbound(msg)
 	return m.pushToUserWithSender(ctx, userID, excludeAuthKeyID, excludeSessionID, t, msg, func(c *Conn) error {
-		return c.SendBestEffort(ctx, t, msg, timeout)
+		if c.outbound == nil || c.outboundControl == nil {
+			return ErrConnClosed
+		}
+		encoded, err := getEncoded()
+		if err != nil {
+			return err
+		}
+		return c.SendBestEffortEncoded(ctx, t, encoded, timeout)
 	})
 }
 
+func onceEncodedOutbound(msg bin.Encoder) func() (*encodedOutboundMessage, error) {
+	var (
+		encoded *encodedOutboundMessage
+		err     error
+	)
+	return func() (*encodedOutboundMessage, error) {
+		if encoded == nil && err == nil {
+			encoded, err = encodeOutboundMessage(msg)
+		}
+		return encoded, err
+	}
+}
+
 func (m *SessionManager) pushToUserWithSender(ctx context.Context, userID int64, excludeAuthKeyID *[8]byte, excludeSessionID int64, t proto.MessageType, msg bin.Encoder, send func(*Conn) error) (int, error) {
 	m.mu.Lock()
 	conns := make([]*Conn, 0, len(m.byUser[userID]))

internal/mtprotoedge/session_manager_test.go

@@ -8,11 +8,21 @@ import (
 
 	"go.uber.org/zap/zaptest"
 
+	"github.com/gotd/td/bin"
 	"github.com/gotd/td/mt"
 	"github.com/gotd/td/proto"
 	"github.com/gotd/td/tg"
 )
 
+type countingOutboundEncoder struct {
+	count *int
+}
+
+func (e *countingOutboundEncoder) Encode(b *bin.Buffer) error {
+	*e.count++
+	return (&tg.UpdatesTooLong{}).Encode(b)
+}
+
 // TestSessionManagerRegistry 验证注册表的注册/注销/查找语义（不涉及网络发送）。
 func TestSessionManagerRegistry(t *testing.T) {
 	sm := NewSessionManager(zaptest.NewLogger(t))
@@ -51,6 +61,43 @@ func TestSessionManagerRegistry(t *testing.T) {
 	}
 }
 
+func TestSessionManagerBestEffortFanoutPreencodesOnce(t *testing.T) {
+	sm := NewSessionManager(zaptest.NewLogger(t))
+	const userID = int64(100)
+	for i := 0; i < 2; i++ {
+		c := &Conn{
+			sessionID:       int64(i + 1),
+			authKeyID:       [8]byte{byte(i + 1)},
+			outbound:        make(chan outboundOp, 1),
+			outboundControl: make(chan outboundOp, 1),
+			outboundStop:    make(chan struct{}),
+		}
+		c.userID.Store(userID)
+		c.userIDResolved.Store(true)
+		c.receivesUpdates.Store(true)
+		sm.Register(c)
+	}
+
+	encodes := 0
+	sent, err := sm.PushToUserExceptSessionBestEffort(
+		context.Background(),
+		userID,
+		0,
+		proto.MessageFromServer,
+		&countingOutboundEncoder{count: &encodes},
+		0,
+	)
+	if err != nil {
+		t.Fatalf("push: %v", err)
+	}
+	if sent != 2 {
+		t.Fatalf("sent = %d, want 2", sent)
+	}
+	if encodes != 1 {
+		t.Fatalf("encoded %d times, want 1", encodes)
+	}
+}
+
 func TestSessionManagerScopesSameSessionIDByAuthKey(t *testing.T) {
 	sm := NewSessionManager(zaptest.NewLogger(t))
 	raw1 := [8]byte{1}

internal/mtprotoedge/transport_compat.go

@@ -0,0 +1,456 @@
+package mtprotoedge
+
+import (
+	"bytes"
+	"context"
+	"encoding/binary"
+	"fmt"
+	"io"
+	"net"
+	"sync"
+	"time"
+
+	"github.com/go-faster/errors"
+	"go.uber.org/multierr"
+
+	"github.com/gotd/td/bin"
+	tdcrypto "github.com/gotd/td/crypto"
+	"github.com/gotd/td/proto/codec"
+	"github.com/gotd/td/transport"
+)
+
+const maxTransportMessageSize = 1 << 24
+const quickAckResponseFlag = uint32(1 << 31)
+
+type transportListener interface {
+	Accept() (transport.Conn, error)
+	Close() error
+	Addr() net.Addr
+}
+
+type quickAckTransport interface {
+	ConsumeQuickAckRequested() bool
+	SendQuickAck(ctx context.Context, token uint32) error
+}
+
+type compatTransportListener struct {
+	codec    func() transport.Codec
+	listener net.Listener
+}
+
+func newCompatTransportListener(codec func() transport.Codec, listener net.Listener) transportListener {
+	if codec != nil {
+		return transport.ListenCodec(codec, listener)
+	}
+	return &compatTransportListener{listener: listener}
+}
+
+func (l *compatTransportListener) Accept() (_ transport.Conn, rErr error) {
+	conn, err := l.listener.Accept()
+	if err != nil {
+		return nil, err
+	}
+	defer func() {
+		if rErr != nil {
+			multierr.AppendInto(&rErr, conn.Close())
+		}
+	}()
+
+	connCodec, reader, err := detectCompatCodec(conn)
+	if err != nil {
+		return nil, errors.Wrap(err, "detect codec")
+	}
+
+	return &compatTransportConn{
+		conn: wrappedCompatConn{
+			reader: reader,
+			Conn:   conn,
+		},
+		codec: connCodec,
+	}, nil
+}
+
+func (l *compatTransportListener) Close() error {
+	return l.listener.Close()
+}
+
+func (l *compatTransportListener) Addr() net.Addr {
+	return l.listener.Addr()
+}
+
+type wrappedCompatConn struct {
+	reader io.Reader
+	net.Conn
+}
+
+func (w wrappedCompatConn) Read(p []byte) (int, error) {
+	return w.reader.Read(p)
+}
+
+type compatTransportConn struct {
+	conn  net.Conn
+	codec transport.Codec
+
+	readMux  sync.Mutex
+	writeMux sync.Mutex
+}
+
+func (c *compatTransportConn) Send(ctx context.Context, b *bin.Buffer) error {
+	c.writeMux.Lock()
+	defer c.writeMux.Unlock()
+
+	if err := c.conn.SetWriteDeadline(time.Time{}); err != nil {
+		return errors.Wrap(err, "reset write deadline")
+	}
+	if deadline, ok := ctx.Deadline(); ok {
+		if err := c.conn.SetWriteDeadline(deadline); err != nil {
+			return errors.Wrap(err, "set write deadline")
+		}
+	}
+
+	if err := c.codec.Write(c.conn, b); err != nil {
+		return errors.Wrap(err, "write")
+	}
+	return nil
+}
+
+func (c *compatTransportConn) ConsumeQuickAckRequested() bool {
+	q, ok := c.codec.(quickAckCodec)
+	if !ok {
+		return false
+	}
+	return q.consumeQuickAckRequested()
+}
+
+func (c *compatTransportConn) SendQuickAck(ctx context.Context, token uint32) error {
+	q, ok := c.codec.(quickAckCodec)
+	if !ok {
+		return nil
+	}
+
+	c.writeMux.Lock()
+	defer c.writeMux.Unlock()
+
+	if err := c.conn.SetWriteDeadline(time.Time{}); err != nil {
+		return errors.Wrap(err, "reset write deadline")
+	}
+	if deadline, ok := ctx.Deadline(); ok {
+		if err := c.conn.SetWriteDeadline(deadline); err != nil {
+			return errors.Wrap(err, "set write deadline")
+		}
+	}
+
+	raw := q.quickAckResponse(token)
+	if err := writeAll(c.conn, raw[:]); err != nil {
+		return errors.Wrap(err, "write quick ack")
+	}
+	return nil
+}
+
+func (c *compatTransportConn) Recv(ctx context.Context, b *bin.Buffer) error {
+	c.readMux.Lock()
+	defer c.readMux.Unlock()
+
+	if err := c.conn.SetReadDeadline(time.Time{}); err != nil {
+		return errors.Wrap(err, "reset read deadline")
+	}
+	if deadline, ok := ctx.Deadline(); ok {
+		if err := c.conn.SetReadDeadline(deadline); err != nil {
+			return errors.Wrap(err, "set read deadline")
+		}
+	}
+
+	if err := c.codec.Read(c.conn, b); err != nil {
+		return errors.Wrap(err, "read")
+	}
+	return nil
+}
+
+func (c *compatTransportConn) Close() error {
+	return c.conn.Close()
+}
+
+func detectCompatCodec(c io.Reader) (transport.Codec, io.Reader, error) {
+	var buf [4]byte
+	if _, err := io.ReadFull(c, buf[:1]); err != nil {
+		return nil, nil, errors.Wrap(err, "read first byte")
+	}
+
+	if buf[0] == codec.AbridgedClientStart[0] {
+		return &quickAckAbridgedCodec{}, c, nil
+	}
+
+	if _, err := io.ReadFull(c, buf[1:4]); err != nil {
+		return nil, nil, errors.Wrap(err, "read header")
+	}
+	switch buf {
+	case codec.IntermediateClientStart:
+		return &quickAckIntermediateCodec{}, c, nil
+	case codec.PaddedIntermediateClientStart:
+		return &quickAckPaddedIntermediateCodec{}, c, nil
+	default:
+		return transport.Full.Codec(), io.MultiReader(bytes.NewReader(buf[:]), c), nil
+	}
+}
+
+type quickAckCodec interface {
+	transport.Codec
+	consumeQuickAckRequested() bool
+	quickAckResponse(token uint32) [4]byte
+}
+
+type quickAckAbridgedCodec struct {
+	quickAckRequested bool
+}
+
+func (*quickAckAbridgedCodec) WriteHeader(w io.Writer) error {
+	return (codec.Abridged{}).WriteHeader(w)
+}
+
+func (*quickAckAbridgedCodec) ReadHeader(r io.Reader) error {
+	return (codec.Abridged{}).ReadHeader(r)
+}
+
+func (*quickAckAbridgedCodec) Write(w io.Writer, b *bin.Buffer) error {
+	if err := validateOutgoingCompatMessage(b); err != nil {
+		return err
+	}
+
+	words := b.Len() >> 2
+	var header [4]byte
+	headerLen := 1
+	if words < 0x7f {
+		header[0] = byte(words)
+	} else {
+		header[0] = 0x7f
+		header[1] = byte(words)
+		header[2] = byte(words >> 8)
+		header[3] = byte(words >> 16)
+		headerLen = 4
+	}
+	return writeCompatPacket(w, header[:headerLen], b.Raw())
+}
+
+func (q *quickAckAbridgedCodec) Read(r io.Reader, b *bin.Buffer) error {
+	requested, err := readQuickAckAbridged(r, b)
+	if err != nil {
+		return errors.Wrap(err, "read abridged")
+	}
+	q.quickAckRequested = requested
+	return checkCompatProtocolError(b)
+}
+
+func (q *quickAckAbridgedCodec) consumeQuickAckRequested() bool {
+	v := q.quickAckRequested
+	q.quickAckRequested = false
+	return v
+}
+
+func (*quickAckAbridgedCodec) quickAckResponse(token uint32) [4]byte {
+	var raw [4]byte
+	binary.BigEndian.PutUint32(raw[:], (token&^quickAckResponseFlag)|quickAckResponseFlag)
+	return raw
+}
+
+type quickAckIntermediateCodec struct {
+	quickAckRequested bool
+}
+
+func (*quickAckIntermediateCodec) WriteHeader(w io.Writer) error {
+	return (codec.Intermediate{}).WriteHeader(w)
+}
+
+func (*quickAckIntermediateCodec) ReadHeader(r io.Reader) error {
+	return (codec.Intermediate{}).ReadHeader(r)
+}
+
+func (*quickAckIntermediateCodec) Write(w io.Writer, b *bin.Buffer) error {
+	if err := validateOutgoingCompatMessage(b); err != nil {
+		return err
+	}
+	var header [4]byte
+	binary.LittleEndian.PutUint32(header[:], uint32(b.Len()))
+	return writeCompatPacket(w, header[:], b.Raw())
+}
+
+func (q *quickAckIntermediateCodec) Read(r io.Reader, b *bin.Buffer) error {
+	requested, err := readQuickAckIntermediate(r, b, false)
+	if err != nil {
+		return errors.Wrap(err, "read intermediate")
+	}
+	q.quickAckRequested = requested
+	return checkCompatProtocolError(b)
+}
+
+func (q *quickAckIntermediateCodec) consumeQuickAckRequested() bool {
+	v := q.quickAckRequested
+	q.quickAckRequested = false
+	return v
+}
+
+func (*quickAckIntermediateCodec) quickAckResponse(token uint32) [4]byte {
+	var raw [4]byte
+	binary.LittleEndian.PutUint32(raw[:], (token&^quickAckResponseFlag)|quickAckResponseFlag)
+	return raw
+}
+
+type quickAckPaddedIntermediateCodec struct {
+	quickAckRequested bool
+}
+
+func (*quickAckPaddedIntermediateCodec) WriteHeader(w io.Writer) error {
+	return (codec.PaddedIntermediate{}).WriteHeader(w)
+}
+
+func (*quickAckPaddedIntermediateCodec) ReadHeader(r io.Reader) error {
+	return (codec.PaddedIntermediate{}).ReadHeader(r)
+}
+
+func (*quickAckPaddedIntermediateCodec) Write(w io.Writer, b *bin.Buffer) error {
+	if err := validateOutgoingCompatMessage(b); err != nil {
+		return err
+	}
+	var padding [4]byte
+	if _, err := io.ReadFull(tdcrypto.DefaultRand(), padding[:]); err != nil {
+		return err
+	}
+	n := int(padding[0] % 4)
+	payload := b.Raw()
+	if n > 0 {
+		payload = append(append([]byte(nil), payload...), padding[:n]...)
+	}
+	var header [4]byte
+	binary.LittleEndian.PutUint32(header[:], uint32(len(payload)))
+	return writeCompatPacket(w, header[:], payload)
+}
+
+func (q *quickAckPaddedIntermediateCodec) Read(r io.Reader, b *bin.Buffer) error {
+	requested, err := readQuickAckIntermediate(r, b, true)
+	if err != nil {
+		return errors.Wrap(err, "read padded intermediate")
+	}
+	q.quickAckRequested = requested
+	return checkCompatProtocolError(b)
+}
+
+func (q *quickAckPaddedIntermediateCodec) consumeQuickAckRequested() bool {
+	v := q.quickAckRequested
+	q.quickAckRequested = false
+	return v
+}
+
+func (*quickAckPaddedIntermediateCodec) quickAckResponse(token uint32) [4]byte {
+	var raw [4]byte
+	binary.LittleEndian.PutUint32(raw[:], (token&^quickAckResponseFlag)|quickAckResponseFlag)
+	return raw
+}
+
+func readQuickAckAbridged(r io.Reader, b *bin.Buffer) (bool, error) {
+	var first [1]byte
+	if _, err := io.ReadFull(r, first[:]); err != nil {
+		return false, err
+	}
+
+	requested := first[0]&0x80 != 0
+	lengthByte := first[0] & 0x7f
+	var n int
+	if lengthByte == 0x7f {
+		var tail [3]byte
+		if _, err := io.ReadFull(r, tail[:]); err != nil {
+			return false, err
+		}
+		words := uint32(tail[0]) | uint32(tail[1])<<8 | uint32(tail[2])<<16
+		n = int(words << 2)
+	} else {
+		n = int(lengthByte) << 2
+	}
+
+	if err := validateCompatTransportLength(n); err != nil {
+		return false, err
+	}
+	resetCompatBufferN(b, n)
+	if _, err := io.ReadFull(r, b.Buf); err != nil {
+		return false, errors.Wrap(err, "read payload")
+	}
+	return requested, nil
+}
+
+func readQuickAckIntermediate(r io.Reader, b *bin.Buffer, padding bool) (bool, error) {
+	var lengthBuf [4]byte
+	if _, err := io.ReadFull(r, lengthBuf[:]); err != nil {
+		return false, errors.Wrap(err, "read length")
+	}
+	rawLength := binary.LittleEndian.Uint32(lengthBuf[:])
+	requested := rawLength&quickAckResponseFlag != 0
+	n := int(rawLength &^ quickAckResponseFlag)
+	if err := validateCompatTransportLength(n); err != nil {
+		return false, err
+	}
+	resetCompatBufferN(b, n)
+	if _, err := io.ReadFull(r, b.Buf); err != nil {
+		return false, errors.Wrap(err, "read payload")
+	}
+	if padding {
+		paddingLength := n % 4
+		b.Buf = b.Buf[:n-paddingLength]
+	}
+	return requested, nil
+}
+
+func validateOutgoingCompatMessage(b *bin.Buffer) error {
+	n := b.Len()
+	if err := validateCompatTransportLength(n); err != nil {
+		return err
+	}
+	if n%bin.Word != 0 {
+		return fmt.Errorf("invalid message length %d: not aligned to %d", n, bin.Word)
+	}
+	return nil
+}
+
+func writeCompatPacket(w io.Writer, header, payload []byte) error {
+	packet := make([]byte, len(header)+len(payload))
+	copy(packet, header)
+	copy(packet[len(header):], payload)
+	return writeAll(w, packet)
+}
+
+func writeAll(w io.Writer, p []byte) error {
+	for len(p) > 0 {
+		n, err := w.Write(p)
+		if err != nil {
+			return err
+		}
+		if n == 0 {
+			return io.ErrShortWrite
+		}
+		p = p[n:]
+	}
+	return nil
+}
+
+func validateCompatTransportLength(n int) error {
+	if n <= 0 || n > maxTransportMessageSize {
+		return fmt.Errorf("invalid message length %d", n)
+	}
+	return nil
+}
+
+func resetCompatBufferN(b *bin.Buffer, n int) {
+	if cap(b.Buf) < n {
+		b.Buf = make([]byte, n)
+		return
+	}
+	b.Buf = b.Buf[:n]
+}
+
+func checkCompatProtocolError(b *bin.Buffer) error {
+	if b.Len() != bin.Word {
+		return nil
+	}
+	code, err := b.Int32()
+	if err != nil {
+		return err
+	}
+	return &codec.ProtocolErr{Code: -code}
+}

internal/mtprotoedge/transport_compat_test.go

@@ -0,0 +1,96 @@
+package mtprotoedge
+
+import (
+	"bytes"
+	"encoding/binary"
+	"testing"
+
+	"github.com/gotd/td/bin"
+)
+
+type countWriteBuffer struct {
+	bytes.Buffer
+	writes int
+}
+
+func (w *countWriteBuffer) Write(p []byte) (int, error) {
+	w.writes++
+	return w.Buffer.Write(p)
+}
+
+func TestQuickAckResponseEncoding(t *testing.T) {
+	const token = 0x01020304
+
+	abridged := (&quickAckAbridgedCodec{}).quickAckResponse(token)
+	if want := []byte{0x81, 0x02, 0x03, 0x04}; !bytes.Equal(abridged[:], want) {
+		t.Fatalf("abridged quick ack = %x, want %x", abridged, want)
+	}
+
+	intermediate := (&quickAckIntermediateCodec{}).quickAckResponse(token)
+	if want := []byte{0x04, 0x03, 0x02, 0x81}; !bytes.Equal(intermediate[:], want) {
+		t.Fatalf("intermediate quick ack = %x, want %x", intermediate, want)
+	}
+}
+
+func TestQuickAckReadFlags(t *testing.T) {
+	payload := []byte{1, 2, 3, 4, 5, 6, 7, 8}
+	abridgedPacket := append([]byte{0x80 | byte(len(payload)/bin.Word)}, payload...)
+
+	var got bin.Buffer
+	requested, err := readQuickAckAbridged(bytes.NewReader(abridgedPacket), &got)
+	if err != nil {
+		t.Fatalf("read abridged: %v", err)
+	}
+	if !requested {
+		t.Fatal("abridged quick ack flag was not detected")
+	}
+	if !bytes.Equal(got.Raw(), payload) {
+		t.Fatalf("abridged payload = %x, want %x", got.Raw(), payload)
+	}
+
+	var header [4]byte
+	binary.LittleEndian.PutUint32(header[:], uint32(len(payload))|quickAckResponseFlag)
+	intermediatePacket := append(header[:], payload...)
+	requested, err = readQuickAckIntermediate(bytes.NewReader(intermediatePacket), &got, false)
+	if err != nil {
+		t.Fatalf("read intermediate: %v", err)
+	}
+	if !requested {
+		t.Fatal("intermediate quick ack flag was not detected")
+	}
+	if !bytes.Equal(got.Raw(), payload) {
+		t.Fatalf("intermediate payload = %x, want %x", got.Raw(), payload)
+	}
+}
+
+func TestCompatTransportCodecsWriteSinglePacket(t *testing.T) {
+	var payload bin.Buffer
+	payload.PutInt32(0x01020304)
+	payload.PutInt32(0x05060708)
+
+	t.Run("abridged", func(t *testing.T) {
+		var out countWriteBuffer
+		if err := (&quickAckAbridgedCodec{}).Write(&out, &payload); err != nil {
+			t.Fatalf("write: %v", err)
+		}
+		if out.writes != 1 {
+			t.Fatalf("writes = %d, want 1", out.writes)
+		}
+		if got, want := out.Bytes()[0], byte(payload.Len()/bin.Word); got != want {
+			t.Fatalf("abridged header = %#x, want %#x", got, want)
+		}
+	})
+
+	t.Run("intermediate", func(t *testing.T) {
+		var out countWriteBuffer
+		if err := (&quickAckIntermediateCodec{}).Write(&out, &payload); err != nil {
+			t.Fatalf("write: %v", err)
+		}
+		if out.writes != 1 {
+			t.Fatalf("writes = %d, want 1", out.writes)
+		}
+		if got, want := binary.LittleEndian.Uint32(out.Bytes()[:4]), uint32(payload.Len()); got != want {
+			t.Fatalf("intermediate length = %d, want %d", got, want)
+		}
+	})
+}