From 1f646ef02474fb2b8fb4b21c0e3be9411dfc07cb Mon Sep 17 00:00:00 2001 From: A Date: Sun, 12 Jul 2026 16:20:23 +0800 Subject: [PATCH] fix: sync websocket and WebK compatibility fixes --- internal/app/help/service.go | 7 +- internal/app/help/service_premium_test.go | 8 ++ internal/compat/tdesktop/defaults.go | 5 +- .../compat/tdesktop/startup_stubs_test.go | 15 +++ internal/mtprotoedge/outbound.go | 8 +- internal/mtprotoedge/outbound_scratch.go | 17 ++- internal/mtprotoedge/outbound_test.go | 28 +++- internal/mtprotoedge/server.go | 3 +- internal/mtprotoedge/transport_compat.go | 120 +++++++++++++++++- internal/mtprotoedge/transport_compat_test.go | 97 ++++++++++++++ internal/mtprotoedge/transport_ownership.go | 20 +++ 11 files changed, 311 insertions(+), 17 deletions(-) diff --git a/internal/app/help/service.go b/internal/app/help/service.go index ce33c53c..4d201356 100644 --- a/internal/app/help/service.go +++ b/internal/app/help/service.go @@ -50,9 +50,12 @@ const tdesktopClient = "tdesktop" // 隐身模式本地 UI/乐观状态用的时间常量，与当前 bounded stealth update stub 保持一致。 // - aicompose_tone_* 与 domain/app/ai 默认值一致：TDesktop/DrKLO 创建/预览 tone 时 // 直接读取这些 key 做本地输入限制和示例数量。 -const tdesktopDefaultAppConfigBase = `{"chat_read_mark_expire_period":604800,"chat_read_mark_size_threshold":50,"pm_read_date_expire_period":604800,"quote_length_max":1024,"telegram_antispam_group_size_min":200,"telegram_antispam_user_id":"5434988373","forum_upgrade_participants_min":2,"reactions_default":{"_":"reactionEmoji","emoticon":"👍"},"reactions_uniq_max":11,"reactions_user_max_default":1,"reactions_user_max_premium":3,"reactions_in_chat_max":3,"boosts_channel_level_max":100,"rich_message_posting":"enabled","upload_markup_video":true,"emojies_send_dice":["🎲","🎯","🏀","⚽","⚽️","🎳","🎰"],"premium_purchase_blocked":false,"stars_purchase_blocked":false,"stargifts_blocked":false,"stories_stealth_future_period":1500,"stories_stealth_past_period":300,"stories_stealth_cooldown_period":10800,"quick_replies_limit":100,"quick_reply_messages_limit":20,"business_chat_links_limit":100,"dialog_filters_enabled":true,"chatlist_update_period":3600,"chatlist_invites_limit_default":3,"chatlist_invites_limit_premium":20,"chatlists_joined_limit_default":2,"chatlists_joined_limit_premium":20,"about_length_limit_default":70,"about_length_limit_premium":140,"caption_length_limit_default":1024,"caption_length_limit_premium":4096,"channels_limit_default":500,"channels_limit_premium":1000,"channels_public_limit_default":10,"channels_public_limit_premium":20,"dialog_filters_limit_default":10,"dialog_filters_limit_premium":20,"dialog_filters_chats_limit_default":100,"dialog_filters_chats_limit_premium":200,"dialogs_pinned_limit_default":5,"dialogs_pinned_limit_premium":10,"dialogs_folder_pinned_limit_default":100,"dialogs_folder_pinned_limit_premium":200,"saved_dialogs_pinned_limit_default":5,"saved_dialogs_pinned_limit_premium":100,"saved_gifs_limit_default":200,"saved_gifs_limit_premium":400,"stickers_faved_limit_default":5,"stickers_faved_limit_premium":10,"recommended_channels_limit_default":10,"recommended_channels_limit_premium":100,"aicompose_tone_examples_num":3,"aicompose_tone_title_length_max":12,"aicompose_tone_prompt_length_max":1024,"aicompose_tone_saved_limit_default":5,"aicompose_tone_saved_limit_premium":20,"upload_max_fileparts_default":4000,"upload_max_fileparts_premium":8000` +// +// WebK directly calls Array.some on fragment_prefixes while rendering user profiles, +// so this compatibility key must always remain an array, even when it is empty. +const tdesktopDefaultAppConfigBase = `{"chat_read_mark_expire_period":604800,"chat_read_mark_size_threshold":50,"pm_read_date_expire_period":604800,"quote_length_max":1024,"telegram_antispam_group_size_min":200,"telegram_antispam_user_id":"5434988373","fragment_prefixes":["888"],"forum_upgrade_participants_min":2,"reactions_default":{"_":"reactionEmoji","emoticon":"👍"},"reactions_uniq_max":11,"reactions_user_max_default":1,"reactions_user_max_premium":3,"reactions_in_chat_max":3,"boosts_channel_level_max":100,"rich_message_posting":"enabled","upload_markup_video":true,"emojies_send_dice":["🎲","🎯","🏀","⚽","⚽️","🎳","🎰"],"premium_purchase_blocked":false,"stars_purchase_blocked":false,"stargifts_blocked":false,"stories_stealth_future_period":1500,"stories_stealth_past_period":300,"stories_stealth_cooldown_period":10800,"quick_replies_limit":100,"quick_reply_messages_limit":20,"business_chat_links_limit":100,"dialog_filters_enabled":true,"chatlist_update_period":3600,"chatlist_invites_limit_default":3,"chatlist_invites_limit_premium":20,"chatlists_joined_limit_default":2,"chatlists_joined_limit_premium":20,"about_length_limit_default":70,"about_length_limit_premium":140,"caption_length_limit_default":1024,"caption_length_limit_premium":4096,"channels_limit_default":500,"channels_limit_premium":1000,"channels_public_limit_default":10,"channels_public_limit_premium":20,"dialog_filters_limit_default":10,"dialog_filters_limit_premium":20,"dialog_filters_chats_limit_default":100,"dialog_filters_chats_limit_premium":200,"dialogs_pinned_limit_default":5,"dialogs_pinned_limit_premium":10,"dialogs_folder_pinned_limit_default":100,"dialogs_folder_pinned_limit_premium":200,"saved_dialogs_pinned_limit_default":5,"saved_dialogs_pinned_limit_premium":100,"saved_gifs_limit_default":200,"saved_gifs_limit_premium":400,"stickers_faved_limit_default":5,"stickers_faved_limit_premium":10,"recommended_channels_limit_default":10,"recommended_channels_limit_premium":100,"aicompose_tone_examples_num":3,"aicompose_tone_title_length_max":12,"aicompose_tone_prompt_length_max":1024,"aicompose_tone_saved_limit_default":5,"aicompose_tone_saved_limit_premium":20,"upload_max_fileparts_default":4000,"upload_max_fileparts_premium":8000` -const defaultAppConfigHash = 22 // 默认 app config 内容变更时必须递增，否则缓存端只会收到 notModified。 +const defaultAppConfigHash = 23 // 默认 app config 内容变更时必须递增，否则缓存端只会收到 notModified。 // Service 提供客户端启动配置与国家区号目录。 // diff --git a/internal/app/help/service_premium_test.go b/internal/app/help/service_premium_test.go index db9312d3..a150ae67 100644 --- a/internal/app/help/service_premium_test.go +++ b/internal/app/help/service_premium_test.go @@ -18,6 +18,10 @@ func TestAppConfigPremiumKeys(t *testing.T) { if cfg.Hash != defaultAppConfigHash || cfg.Hash < 10 { t.Fatalf("hash = %d, want defaultAppConfigHash(≥10)", cfg.Hash) } + oldCfg, oldNotModified, err := (*Service)(nil).GetAppConfig(context.Background(), defaultAppConfigHash-1) + if err != nil || oldNotModified || oldCfg.Hash != defaultAppConfigHash { + t.Fatalf("GetAppConfig(old hash) = hash %d notModified %v err %v, want refreshed config", oldCfg.Hash, oldNotModified, err) + } var decoded map[string]any if err := json.Unmarshal(cfg.JSON, &decoded); err != nil { t.Fatalf("app config json invalid: %v", err) @@ -36,6 +40,10 @@ func TestAppConfigPremiumKeys(t *testing.T) { if posting, ok := decoded["rich_message_posting"].(string); !ok || posting != "enabled" { t.Fatalf("rich_message_posting = %v, want enabled (TDesktop 富文本编辑入口默认打开)", decoded["rich_message_posting"]) } + fragmentPrefixes, ok := decoded["fragment_prefixes"].([]any) + if !ok || len(fragmentPrefixes) != 1 || fragmentPrefixes[0] != "888" { + t.Fatalf("fragment_prefixes = %#v, want [\"888\"]", decoded["fragment_prefixes"]) + } wantNumbers := map[string]float64{ "reactions_user_max_default": 1, "reactions_user_max_premium": 3, diff --git a/internal/compat/tdesktop/defaults.go b/internal/compat/tdesktop/defaults.go index fc610cc8..fb1aa145 100644 --- a/internal/compat/tdesktop/defaults.go +++ b/internal/compat/tdesktop/defaults.go @@ -7,7 +7,7 @@ import ( ) const ( - appConfigHash = 16 // app config 内容变更时必须递增，否则缓存端只会收到 notModified。 + appConfigHash = 17 // app config 内容变更时必须递增，否则缓存端只会收到 notModified。 countriesListHash = 1 timezonesListHash = 1 ) @@ -31,6 +31,9 @@ func readMarkAppConfig(mapboxToken string) *tg.JSONObject { {Key: "quote_length_max", Value: &tg.JSONNumber{Value: 1024}}, {Key: "telegram_antispam_group_size_min", Value: &tg.JSONNumber{Value: 200}}, {Key: "telegram_antispam_user_id", Value: &tg.JSONString{Value: "5434988373"}}, + {Key: "fragment_prefixes", Value: &tg.JSONArray{Value: []tg.JSONValueClass{ + &tg.JSONString{Value: "888"}, + }}}, // premium_purchase_blocked=false：客户端把 star gift「Send a Gift」入口与 // premiumCanBuy()=!premium_purchase_blocked 耦合，置 true 会同时隐藏送礼入口 // （详见 app/help/service.go 主配置注释）。这里是无 HelpService 时的最小回退。 diff --git a/internal/compat/tdesktop/startup_stubs_test.go b/internal/compat/tdesktop/startup_stubs_test.go index 613e1358..d7ad3e60 100644 --- a/internal/compat/tdesktop/startup_stubs_test.go +++ b/internal/compat/tdesktop/startup_stubs_test.go @@ -40,6 +40,7 @@ func TestAppConfigIncludesStoryStealthPeriods(t *testing.T) { } values := make(map[string]float64) strings := make(map[string]string) + arrays := make(map[string]*tg.JSONArray) if object, ok := got.Config.(*tg.JSONObject); ok && object != nil { for _, entry := range object.Value { if number, ok := entry.Value.(*tg.JSONNumber); ok { @@ -48,6 +49,9 @@ func TestAppConfigIncludesStoryStealthPeriods(t *testing.T) { if str, ok := entry.Value.(*tg.JSONString); ok { strings[entry.Key] = str.Value } + if array, ok := entry.Value.(*tg.JSONArray); ok { + arrays[entry.Key] = array + } } } want := map[string]float64{ @@ -63,9 +67,20 @@ func TestAppConfigIncludesStoryStealthPeriods(t *testing.T) { if strings["rich_message_posting"] != "enabled" { t.Fatalf("AppConfig[rich_message_posting] = %q, want enabled", strings["rich_message_posting"]) } + fragmentPrefixes := arrays["fragment_prefixes"] + if fragmentPrefixes == nil || len(fragmentPrefixes.Value) != 1 { + t.Fatalf("AppConfig[fragment_prefixes] = %#v, want one-element array", fragmentPrefixes) + } + prefix, ok := fragmentPrefixes.Value[0].(*tg.JSONString) + if !ok || prefix.Value != "888" { + t.Fatalf("AppConfig[fragment_prefixes][0] = %#v, want \"888\"", fragmentPrefixes.Value[0]) + } if _, ok := AppConfig(got.Hash).(*tg.HelpAppConfigNotModified); !ok { t.Fatalf("AppConfig(hash) = %#v, want notModified", AppConfig(got.Hash)) } + if _, ok := AppConfig(got.Hash - 1).(*tg.HelpAppConfig); !ok { + t.Fatalf("AppConfig(old hash) = %#v, want refreshed config", AppConfig(got.Hash-1)) + } } func TestFallbackAppConfigOmitsMapboxToken(t *testing.T) { diff --git a/internal/mtprotoedge/outbound.go b/internal/mtprotoedge/outbound.go index 1bd21b19..95778775 100644 --- a/internal/mtprotoedge/outbound.go +++ b/internal/mtprotoedge/outbound.go @@ -1279,6 +1279,10 @@ type deadlineOutboundWriter interface { SendDeadline(deadline time.Time, b *bin.Buffer) error } +type deadlineOutboundScratchWriter interface { + SendDeadlineWithScratch(deadline time.Time, b *bin.Buffer, scratch *[]byte) error +} + func (c *Conn) writeFrame(ctx context.Context, frame *outboundFrame) error { if ctx == nil { ctx = context.Background() @@ -1309,7 +1313,9 @@ func (c *Conn) writeFrame(ctx context.Context, frame *outboundFrame) error { if writer == nil { writer = c.transport } - if dw, ok := writer.(deadlineOutboundWriter); ok { + if sw, ok := writer.(deadlineOutboundScratchWriter); ok { + err = sw.SendDeadlineWithScratch(deadline, out, &scratch.codec) + } else if dw, ok := writer.(deadlineOutboundWriter); ok { err = dw.SendDeadline(deadline, out) } else { // 回落路径：gotd full codec / 测试注入 codec 仍走 ctx deadline。 diff --git a/internal/mtprotoedge/outbound_scratch.go b/internal/mtprotoedge/outbound_scratch.go index a840bd38..390d2d7d 100644 --- a/internal/mtprotoedge/outbound_scratch.go +++ b/internal/mtprotoedge/outbound_scratch.go @@ -13,10 +13,10 @@ const ( ) // outboundScratchPool bounds and reuses the encrypted wire buffer across connections. A lease -// reserves a conservative 3x wire size while writing (wire + codec/obfuscation copies), then -// shrinks to the actual retained capacity while idle in the bounded pool. Large one-off frames are -// dropped on return. This removes attacker-warmable per-Conn MiB buffers without returning to an -// unbounded allocation-per-message design. +// reserves wire + codec/obfuscation copies plus their bounded transport overhead, then shrinks to +// the actual retained capacity while idle in the bounded pool. Large one-off frames are dropped on +// return. This removes attacker-warmable per-Conn MiB buffers without returning to an unbounded +// allocation-per-message design. type outboundScratchPool struct { budget *outboundTrackedBudget idle chan *outboundScratch @@ -24,6 +24,7 @@ type outboundScratchPool struct { type outboundScratch struct { wire bin.Buffer + codec []byte reserved int } @@ -45,7 +46,7 @@ func (p *outboundScratchPool) acquireUntil(ctx context.Context, stop <-chan stru if p == nil || wireBytes <= 0 { return nil, ErrOutboundMessageTooLarge } - peak := wireBytes * 3 + peak := wireBytes*3 + 2*maxCompatPacketOverhead if peak < wireBytes { // int overflow return nil, ErrOutboundMessageTooLarge } @@ -79,6 +80,7 @@ func (p *outboundScratchPool) acquireUntil(ctx context.Context, stop <-chan stru // that would fit after replacement. old := scratch.reserved scratch.wire.Buf = nil + scratch.codec = nil scratch.reserved = 0 p.budget.release(old) if err := p.budget.waitReserveUntil(ctx, stop, peak, deadline); err != nil { @@ -93,10 +95,11 @@ func (p *outboundScratchPool) release(scratch *outboundScratch) { if p == nil || scratch == nil { return } - retained := cap(scratch.wire.Buf) + retained := cap(scratch.wire.Buf) + cap(scratch.codec) if retained > maxRetainedConnBuffer { p.budget.release(scratch.reserved) scratch.wire.Buf = nil + scratch.codec = nil scratch.reserved = 0 return } @@ -105,6 +108,7 @@ func (p *outboundScratchPool) release(scratch *outboundScratch) { scratch.reserved = retained } scratch.wire.Buf = scratch.wire.Buf[:0] + scratch.codec = scratch.codec[:0] p.putIdle(scratch) } @@ -114,6 +118,7 @@ func (p *outboundScratchPool) putIdle(scratch *outboundScratch) { default: p.budget.release(scratch.reserved) scratch.wire.Buf = nil + scratch.codec = nil scratch.reserved = 0 } } diff --git a/internal/mtprotoedge/outbound_test.go b/internal/mtprotoedge/outbound_test.go index 82b3cf12..f54de3fc 100644 --- a/internal/mtprotoedge/outbound_test.go +++ b/internal/mtprotoedge/outbound_test.go @@ -237,8 +237,8 @@ func TestEncodedControlFramesUseIndependentBudgetForQueuedAndPendingLifetime(t * } func TestOutboundScratchPoolBoundsConcurrentWireCopies(t *testing.T) { - pool := newOutboundScratchPool(300) - first, err := pool.acquire(context.Background(), nil, 100) // 3x peak = full budget. + pool := newOutboundScratchPool(300 + 2*maxCompatPacketOverhead) + first, err := pool.acquire(context.Background(), nil, 100) // Full wire+codec+obfuscation budget. if err != nil { t.Fatalf("acquire first scratch: %v", err) } @@ -261,9 +261,31 @@ func TestOutboundScratchPoolBoundsConcurrentWireCopies(t *testing.T) { } } +func TestOutboundScratchPoolAccountsRetainedCodecScratch(t *testing.T) { + pool := newOutboundScratchPool(300 + 2*maxCompatPacketOverhead) + scratch, err := pool.acquire(context.Background(), nil, 100) + if err != nil { + t.Fatalf("acquire scratch: %v", err) + } + scratch.codec = make([]byte, 0, 80) + pool.release(scratch) + if got := pool.snapshot(); got != 180 { + t.Fatalf("retained wire+codec scratch = %d, want 180", got) + } + + reused, err := pool.acquire(context.Background(), nil, 100) + if err != nil { + t.Fatalf("reuse scratch: %v", err) + } + if cap(reused.codec) != 80 { + t.Fatalf("reused codec scratch capacity = %d, want 80", cap(reused.codec)) + } + pool.release(reused) +} + func TestOutboundScratchAdmissionUsesWriteTimeoutWithoutClosingHealthyConnection(t *testing.T) { wireBytes := encryptedOutboundWireLen(4) - pool := newOutboundScratchPool(int64(wireBytes * 3)) + pool := newOutboundScratchPool(int64(wireBytes*3 + 2*maxCompatPacketOverhead)) blocker, err := pool.acquire(context.Background(), nil, wireBytes) if err != nil { t.Fatalf("occupy shared scratch budget: %v", err) diff --git a/internal/mtprotoedge/server.go b/internal/mtprotoedge/server.go index b141f4dd..718b3c08 100644 --- a/internal/mtprotoedge/server.go +++ b/internal/mtprotoedge/server.go @@ -383,7 +383,8 @@ func (s *Server) serveMixed(ctx context.Context, ln net.Listener) error { // 触发客户端 6s 重连风暴并误判「后端不健康」回退到外部 DNS。per-conn goroutine 模型已消解 // slow-loris 接入饥饿，故嗅探用满 handshakeTimeout 是安全的。 mux := newSamePortMux(ln, s.handshakeTimeout) - wsLn, wsHandler := transport.WebsocketListener(ln.Addr()) + wsRawLn, wsHandler := transport.WebsocketListener(ln.Addr()) + wsLn := newTransportPacketMessageListener(wsRawLn) httpServer := &http.Server{ Handler: websocketRouteHandler(wsHandler, s.websocketOrigins), diff --git a/internal/mtprotoedge/transport_compat.go b/internal/mtprotoedge/transport_compat.go index 5cfa87e0..122aa329 100644 --- a/internal/mtprotoedge/transport_compat.go +++ b/internal/mtprotoedge/transport_compat.go @@ -23,6 +23,11 @@ import ( const maxTransportMessageSize = 1 << 24 const quickAckResponseFlag = uint32(1 << 31) +const ( + maxCompatPacketOverhead = 7 // 4-byte header + up to 3 bytes padded-intermediate padding. + maxRetainedDirectMessageScratch = 64 << 10 +) + type transportListener interface { Accept() (transport.Conn, error) Close() error @@ -38,6 +43,35 @@ type deadlineQuickAckTransport interface { SendQuickAckDeadline(deadline time.Time, token uint32) error } +// transportPacketMessageConn marks transports where one Write is one message instead of an +// arbitrary byte-stream segment. coder/websocket.NetConn has exactly this contract, so a complete +// MTProto transport packet must be encoded before the single underlying Write. +type transportPacketMessageConn struct { + net.Conn +} + +func (*transportPacketMessageConn) transportPacketsAreMessages() {} + +type transportPacketMessageMarker interface { + transportPacketsAreMessages() +} + +type transportPacketMessageListener struct { + net.Listener +} + +func newTransportPacketMessageListener(listener net.Listener) net.Listener { + return &transportPacketMessageListener{Listener: listener} +} + +func (l *transportPacketMessageListener) Accept() (net.Conn, error) { + conn, err := l.Listener.Accept() + if err != nil { + return nil, err + } + return &transportPacketMessageConn{Conn: conn}, nil +} + type compatTransportListener struct { codec func() transport.Codec listener net.Listener @@ -122,11 +156,17 @@ func (l *compatTransportListener) Accept() (_ transport.Conn, rErr error) { reader: reader, Conn: conn, }, - codec: connCodec, - budget: l.budget, + codec: connCodec, + budget: l.budget, + transportPacketMessages: isTransportPacketMessageConn(conn), }, nil } +func isTransportPacketMessageConn(conn net.Conn) bool { + _, ok := conn.(transportPacketMessageMarker) + return ok +} + func (l *compatTransportListener) Close() error { return l.listener.Close() } @@ -149,6 +189,9 @@ type compatTransportConn struct { codec transport.Codec budget *inboundFrameBudget + transportPacketMessages bool + directMessageScratch []byte + readMux sync.Mutex writeMux sync.Mutex @@ -166,18 +209,72 @@ func (c *compatTransportConn) Send(ctx context.Context, b *bin.Buffer) error { // SendDeadline 按显式写超时发送一帧（deadline 为零值表示不设超时）。 // 出站热路径（Conn.writeFrame）走这里，免去 per-frame context timer 分配。 func (c *compatTransportConn) SendDeadline(deadline time.Time, b *bin.Buffer) error { + return c.sendDeadline(deadline, b, nil) +} + +// SendDeadlineWithScratch lets the authenticated outbound path lend its globally budgeted scratch +// to message-oriented transports. Handshake/control writes that do not own such a lease use the +// small bounded per-connection fallback instead. +func (c *compatTransportConn) SendDeadlineWithScratch(deadline time.Time, b *bin.Buffer, scratch *[]byte) error { + return c.sendDeadline(deadline, b, scratch) +} + +func (c *compatTransportConn) sendDeadline(deadline time.Time, b *bin.Buffer, scratch *[]byte) error { c.writeMux.Lock() defer c.writeMux.Unlock() if err := c.conn.SetWriteDeadline(deadline); err != nil { return errors.Wrap(err, "set write deadline") } + if c.transportPacketMessages { + direct := scratch == nil + if direct { + scratch = &c.directMessageScratch + defer c.releaseDirectMessageScratch() + } + if err := c.writeTransportPacketMessage(b, scratch); err != nil { + return errors.Wrap(err, "write message") + } + return nil + } if err := c.codec.Write(c.conn, b); err != nil { return errors.Wrap(err, "write") } return nil } +func (c *compatTransportConn) writeTransportPacketMessage(b *bin.Buffer, scratch *[]byte) error { + required := b.Len() + maxCompatPacketOverhead + if cap(*scratch) < required { + *scratch = make([]byte, 0, required) + } else { + *scratch = (*scratch)[:0] + } + + writer := appendPacketWriter{buf: scratch} + if err := c.codec.Write(&writer, b); err != nil { + return errors.Wrap(err, "encode packet") + } + return writeSingle(c.conn, *scratch) +} + +func (c *compatTransportConn) releaseDirectMessageScratch() { + if cap(c.directMessageScratch) > maxRetainedDirectMessageScratch { + c.directMessageScratch = nil + return + } + c.directMessageScratch = c.directMessageScratch[:0] +} + +type appendPacketWriter struct { + buf *[]byte +} + +func (w *appendPacketWriter) Write(p []byte) (int, error) { + *w.buf = append(*w.buf, p...) + return len(p), nil +} + func (c *compatTransportConn) ConsumeQuickAckRequested() bool { q, ok := c.codec.(quickAckCodec) if !ok { @@ -205,7 +302,13 @@ func (c *compatTransportConn) SendQuickAckDeadline(deadline time.Time, token uin } raw := q.quickAckResponse(token) - if err := writeAll(c.conn, raw[:]); err != nil { + var err error + if c.transportPacketMessages { + err = writeSingle(c.conn, raw[:]) + } else { + err = writeAll(c.conn, raw[:]) + } + if err != nil { return errors.Wrap(err, "write quick ack") } return nil @@ -667,6 +770,17 @@ func writeAll(w io.Writer, p []byte) error { return nil } +func writeSingle(w io.Writer, p []byte) error { + n, err := w.Write(p) + if err != nil { + return err + } + if n != len(p) { + return io.ErrShortWrite + } + return nil +} + func validateCompatTransportLength(n int) error { if n <= 0 || n > maxTransportMessageSize { return fmt.Errorf("invalid message length %d", n) diff --git a/internal/mtprotoedge/transport_compat_test.go b/internal/mtprotoedge/transport_compat_test.go index e8ca2cf8..34e9944c 100644 --- a/internal/mtprotoedge/transport_compat_test.go +++ b/internal/mtprotoedge/transport_compat_test.go @@ -3,11 +3,40 @@ package mtprotoedge import ( "bytes" "encoding/binary" + "io" + "net" "testing" + "time" "github.com/gotd/td/bin" ) +type messageWriteTestConn struct { + bytes.Buffer + writes int + maxWrite int +} + +func (c *messageWriteTestConn) Read([]byte) (int, error) { return 0, io.EOF } +func (c *messageWriteTestConn) Write(p []byte) (int, error) { + c.writes++ + if c.maxWrite > 0 && len(p) > c.maxWrite { + p = p[:c.maxWrite] + } + return c.Buffer.Write(p) +} +func (*messageWriteTestConn) Close() error { return nil } +func (*messageWriteTestConn) LocalAddr() net.Addr { return messageWriteTestAddr("local") } +func (*messageWriteTestConn) RemoteAddr() net.Addr { return messageWriteTestAddr("remote") } +func (*messageWriteTestConn) SetDeadline(time.Time) error { return nil } +func (*messageWriteTestConn) SetReadDeadline(time.Time) error { return nil } +func (*messageWriteTestConn) SetWriteDeadline(time.Time) error { return nil } + +type messageWriteTestAddr string + +func (a messageWriteTestAddr) Network() string { return "message-write-test" } +func (a messageWriteTestAddr) String() string { return string(a) } + type countWriteBuffer struct { bytes.Buffer writes int @@ -128,3 +157,71 @@ func TestCompatTransportCodecsWriteSegmentedPacketWithoutFullCopy(t *testing.T) } }) } + +func TestCompatTransportWebSocketWritesOneMessagePerPacket(t *testing.T) { + var payload bin.Buffer + payload.PutInt32(0x01020304) + payload.PutInt32(0x05060708) + + raw := &messageWriteTestConn{} + conn := &compatTransportConn{ + conn: &transportPacketMessageConn{Conn: raw}, + codec: &quickAckAbridgedCodec{}, + transportPacketMessages: true, + } + var scratch []byte + if err := conn.SendDeadlineWithScratch(time.Time{}, &payload, &scratch); err != nil { + t.Fatalf("send: %v", err) + } + if raw.writes != 1 { + t.Fatalf("websocket writes = %d, want one complete message", raw.writes) + } + want := append([]byte{byte(payload.Len() / bin.Word)}, payload.Raw()...) + if !bytes.Equal(raw.Bytes(), want) { + t.Fatalf("websocket message = %x, want %x", raw.Bytes(), want) + } + if len(scratch) != len(want) { + t.Fatalf("shared scratch length = %d, want %d", len(scratch), len(want)) + } +} + +func TestCompatTransportWebSocketDirectWriteUsesBoundedFallback(t *testing.T) { + var payload bin.Buffer + payload.PutInt32(0x11223344) + payload.PutInt32(0x55667788) + + raw := &messageWriteTestConn{} + conn := &compatTransportConn{ + conn: &transportPacketMessageConn{Conn: raw}, + codec: &quickAckAbridgedCodec{}, + transportPacketMessages: true, + } + if err := conn.SendDeadline(time.Time{}, &payload); err != nil { + t.Fatalf("direct send: %v", err) + } + if raw.writes != 1 { + t.Fatalf("direct websocket writes = %d, want one complete message", raw.writes) + } + if cap(conn.directMessageScratch) == 0 || cap(conn.directMessageScratch) > maxRetainedDirectMessageScratch { + t.Fatalf("direct scratch capacity = %d, want bounded retained buffer", cap(conn.directMessageScratch)) + } +} + +func TestCompatTransportWebSocketDoesNotRetryShortMessageWrite(t *testing.T) { + var payload bin.Buffer + payload.PutInt32(0x01020304) + payload.PutInt32(0x05060708) + + raw := &messageWriteTestConn{maxWrite: 1} + conn := &compatTransportConn{ + conn: &transportPacketMessageConn{Conn: raw}, + codec: &quickAckAbridgedCodec{}, + transportPacketMessages: true, + } + if err := conn.SendDeadline(time.Time{}, &payload); err == nil { + t.Fatal("short websocket message write unexpectedly succeeded") + } + if raw.writes != 1 { + t.Fatalf("short websocket writes = %d, want no multi-message retry", raw.writes) + } +} diff --git a/internal/mtprotoedge/transport_ownership.go b/internal/mtprotoedge/transport_ownership.go index c7557dc5..a99ebabc 100644 --- a/internal/mtprotoedge/transport_ownership.go +++ b/internal/mtprotoedge/transport_ownership.go @@ -144,6 +144,26 @@ func (l *physicalTransportLease) SendDeadline(deadline time.Time, b *bin.Buffer) }) } +// SendDeadlineWithScratch forwards the globally budgeted codec scratch through the generation +// lease while preserving the same write-ownership barrier as SendDeadline. +func (l *physicalTransportLease) SendDeadlineWithScratch(deadline time.Time, b *bin.Buffer, scratch *[]byte) error { + return l.withCurrentWriter(func(raw transport.Conn) error { + if writer, ok := raw.(deadlineOutboundScratchWriter); ok { + return writer.SendDeadlineWithScratch(deadline, b, scratch) + } + if writer, ok := raw.(deadlineOutboundWriter); ok { + return writer.SendDeadline(deadline, b) + } + ctx := context.Background() + cancel := func() {} + if !deadline.IsZero() { + ctx, cancel = context.WithDeadline(ctx, deadline) + } + defer cancel() + return raw.Send(ctx, b) + }) +} + func (l *physicalTransportLease) withCurrentWriter(send func(transport.Conn) error) error { if l == nil || l.owner == nil || l.owner.raw == nil { return ErrConnClosed