fix: sync websocket and WebK compatibility fixes
This commit is contained in:
A
parent
ea6cc72886
commit
1f646ef024
11 changed files with 311 additions and 17 deletions
|
|
@ -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 提供客户端启动配置与国家区号目录。
|
||||
//
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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 时的最小回退。
|
||||
|
|
|
|||
|
|
@ -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) {
|
||||
|
|
|
|||
|
|
@ -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。
|
||||
|
|
|
|||
|
|
@ -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
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue