diff --git a/internal/mtprotoedge/conn.go b/internal/mtprotoedge/conn.go index fe32acdf..d1180fc9 100644 --- a/internal/mtprotoedge/conn.go +++ b/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 diff --git a/internal/mtprotoedge/encrypted.go b/internal/mtprotoedge/encrypted.go index f67802ca..c3f24ff9 100644 --- a/internal/mtprotoedge/encrypted.go +++ b/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, diff --git a/internal/mtprotoedge/encrypted_test.go b/internal/mtprotoedge/encrypted_test.go index b07d8b10..35195cb6 100644 --- a/internal/mtprotoedge/encrypted_test.go +++ b/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) { diff --git a/internal/mtprotoedge/outbound.go b/internal/mtprotoedge/outbound.go index 2fe2a254..841d5c60 100644 --- a/internal/mtprotoedge/outbound.go +++ b/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) - msgID := c.msgID.New(t) - return &outboundFrame{ - msgID: msgID, - seqNo: c.nextSeqNo(content), + return &encodedOutboundMessage{ 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 - if err := c.cipher.Encrypt(c.key, crypto.EncryptedMessageData{ - Salt: c.salt, - SessionID: c.sessionID, - MessageID: frame.msgID, - SeqNo: frame.seqNo, - MessageDataLen: int32(len(frame.body)), - MessageDataWithPadding: frame.body, - }, &out); err != nil { + out, err := c.encryptOutboundFrame(frame) + if 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, diff --git a/internal/mtprotoedge/outbound_test.go b/internal/mtprotoedge/outbound_test.go index 4e2e2d75..5f62ea7a 100644 --- a/internal/mtprotoedge/outbound_test.go +++ b/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}) diff --git a/internal/mtprotoedge/server.go b/internal/mtprotoedge/server.go index 5cbcf0c0..90781a39 100644 --- a/internal/mtprotoedge/server.go +++ b/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") diff --git a/internal/mtprotoedge/server_test.go b/internal/mtprotoedge/server_test.go index 394b3635..3de7e6bf 100644 --- a/internal/mtprotoedge/server_test.go +++ b/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") + } +} diff --git a/internal/mtprotoedge/session_manager.go b/internal/mtprotoedge/session_manager.go index a397b13d..13758c2f 100644 --- a/internal/mtprotoedge/session_manager.go +++ b/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])) diff --git a/internal/mtprotoedge/session_manager_test.go b/internal/mtprotoedge/session_manager_test.go index 14070e97..5660273f 100644 --- a/internal/mtprotoedge/session_manager_test.go +++ b/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} diff --git a/internal/mtprotoedge/transport_compat.go b/internal/mtprotoedge/transport_compat.go new file mode 100644 index 00000000..9ac2fc54 --- /dev/null +++ b/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} +} diff --git a/internal/mtprotoedge/transport_compat_test.go b/internal/mtprotoedge/transport_compat_test.go new file mode 100644 index 00000000..e2ebf042 --- /dev/null +++ b/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) + } + }) +}