diff --git a/internal/app/dialogs/service.go b/internal/app/dialogs/service.go index 7cce2420..f62fede2 100644 --- a/internal/app/dialogs/service.go +++ b/internal/app/dialogs/service.go @@ -336,7 +336,12 @@ func (s *Service) appendMissingChannelPeerPreviews(ctx context.Context, userID i if !ok || view.Forbidden { continue } - if view.Self.Status != domain.ChannelMemberActive { + // Linked discussion guests need a transient peer-dialog snapshot so + // TDesktop can finish materializing the comments History after + // requestSelf. ChannelLeft keeps the snapshot out of the main chat list, + // and Guest guarantees this path never turns an ordinary public preview + // into a dialog. + if view.Self.Status != domain.ChannelMemberActive && !view.Self.Guest { continue } history, err := s.channels.ListChannelHistory(ctx, userID, domain.ChannelHistoryFilter{ diff --git a/internal/mtprotoedge/conn.go b/internal/mtprotoedge/conn.go index 2634711f..575d9afb 100644 --- a/internal/mtprotoedge/conn.go +++ b/internal/mtprotoedge/conn.go @@ -25,13 +25,31 @@ type outboundWriter interface { Send(context.Context, *bin.Buffer) error } +type connLifecycle uint32 + +const ( + // The zero value is deliberately provisional so test/embedded Conn values start + // outside every SessionManager index until they complete activation. + connLifecycleProvisional connLifecycle = iota + connLifecycleClaiming + connLifecycleActive + // retired is terminal and irreversible. A physical connection that lost an + // activation claim must never become visible again, even if its read goroutine + // was already between preflight and publish when a replacement arrived. + connLifecycleRetired +) + type Conn struct { - transport transport.Conn - writer outboundWriter - cipher crypto.Cipher - msgID *proto.MessageIDGen - writeTimeout time.Duration - metrics Metrics + transport transport.Conn + // transportLease owns exactly one generation of the physical transport. + // It is nil only for directly constructed test/embedded Conns that retain + // the legacy raw-transport close fallback. + transportLease *physicalTransportLease + writer outboundWriter + cipher crypto.Cipher + msgID *proto.MessageIDGen + writeTimeout time.Duration + metrics Metrics authKeyID [8]byte // authKeyHex 是 authKeyID 的 hex 缓存：每条 RPC 的结构化日志都会带它， @@ -66,7 +84,11 @@ type Conn struct { outboundScratchOnce sync.Once // terminal 表示该 logical Conn 已停止接受新的出站操作。写失败时由 // outbound actor 置位并只发停止信号，不能在 actor 内等待自身退出。 - terminal atomic.Bool + terminal atomic.Bool + // lifecycle is a monotonic activation state machine. In particular, retired + // never transitions back to claiming/active; this closes the stale-read-loop + // ABA where an evicted Conn observed "not active" and registered itself again. + lifecycle atomic.Uint32 transportClose sync.Once rpcScheduler *inboundRPCScheduler @@ -130,6 +152,63 @@ type Conn struct { clientLayer atomic.Int32 } +func (c *Conn) lifecycleState() connLifecycle { + if c == nil { + return connLifecycleRetired + } + return connLifecycle(c.lifecycle.Load()) +} + +func (c *Conn) beginActivationClaim() bool { + if c == nil || c.terminal.Load() || !c.isPhysicalTransportCurrentOpen() { + return false + } + if !c.lifecycle.CompareAndSwap(uint32(connLifecycleProvisional), uint32(connLifecycleClaiming)) { + return false + } + // Physical close can win after the pre-check but before the lifecycle CAS. + // Do not let a doomed claimant enter SessionManager and retire a healthy old + // owner for the same logical session. + if c.terminal.Load() || !c.isPhysicalTransportCurrentOpen() { + c.lifecycle.Store(uint32(connLifecycleRetired)) + return false + } + return true +} + +func (c *Conn) publishActivation() bool { + if c == nil || c.terminal.Load() || !c.isPhysicalTransportCurrentOpen() { + return false + } + if !c.lifecycle.CompareAndSwap(uint32(connLifecycleClaiming), uint32(connLifecycleActive)) { + return false + } + // A concurrent transport failure can retire the Conn between the first + // terminal check and the CAS. Never let that intermediate active value escape. + if c.terminal.Load() || !c.isPhysicalTransportCurrentOpen() { + c.lifecycle.Store(uint32(connLifecycleRetired)) + return false + } + return true +} + +func (c *Conn) isActive() bool { + return c != nil && !c.terminal.Load() && c.lifecycleState() == connLifecycleActive +} + +// transferTransportOwnership hands this Conn's physical socket to the next +// logical generation. The caller must have fenced and drained the old writer. +func (c *Conn) transferTransportOwnership() (*physicalTransportLease, bool) { + if c == nil || c.transportLease == nil { + return nil, false + } + return c.transportLease.Transfer() +} + +func (c *Conn) isPhysicalTransportCurrentOpen() bool { + return c != nil && (c.transportLease == nil || c.transportLease.IsCurrentOpen()) +} + // ClientLayer 返回连接协商的 TL layer；未协商时返回 canonical layer（227，不降级）。 func (c *Conn) ClientLayer() int { if l := c.clientLayer.Load(); l != 0 { diff --git a/internal/mtprotoedge/duplicate_admission_test.go b/internal/mtprotoedge/duplicate_admission_test.go new file mode 100644 index 00000000..4e4f1a7c --- /dev/null +++ b/internal/mtprotoedge/duplicate_admission_test.go @@ -0,0 +1,116 @@ +package mtprotoedge + +import ( + "context" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/gotd/td/bin" + "github.com/gotd/td/proto" + "github.com/gotd/td/tg" +) + +type blockingDuplicateRPC struct { + started chan struct{} + release chan struct{} + calls atomic.Int32 + once sync.Once +} + +func newBlockingDuplicateRPC() *blockingDuplicateRPC { + return &blockingDuplicateRPC{ + started: make(chan struct{}, 4), + release: make(chan struct{}), + } +} + +func (h *blockingDuplicateRPC) Dispatch(ctx context.Context, _ [8]byte, _ int64, _ *bin.Buffer) (bin.Encoder, error) { + h.calls.Add(1) + h.started <- struct{}{} + select { + case <-h.release: + return &tg.Config{ThisDC: 2}, nil + case <-ctx.Done(): + return nil, ctx.Err() + } +} + +func (*blockingDuplicateRPC) NegotiatedLayer([8]byte, int64) (int, bool) { return 227, true } + +func (h *blockingDuplicateRPC) unblock() { h.once.Do(func() { close(h.release) }) } + +// TestPendingSameConnectionDuplicateDoesNotBlockFreshRequest models the core +// Android startup failure: an RPC is executing, salt correction makes the client +// resend the same msg_id, then initConnection assigns a fresh msg_id. A local +// duplicate must not synchronously join the old owner on the socket read loop, +// otherwise the fresh request remains unread until the old result/replay storm +// has drained. +func TestPendingSameConnectionDuplicateDoesNotBlockFreshRequest(t *testing.T) { + const dc = 2 + handler := newBlockingDuplicateRPC() + defer handler.unblock() + addr, pub, server := startTestServer(t, Options{ + DC: dc, + RPC: handler, + RPCMaxInflight: 2, + RPCGlobalWorkers: 2, + RPCQueueSize: 8, + }) + conn, auth, cipher := dialHandshake(t, addr, dc, pub) + + ids := proto.NewMessageIDGen(time.Now) + oldID := ids.New(proto.MessageFromClient) + newID := ids.New(proto.MessageFromClient) + + sendEncryptedWithSeq(t, conn, cipher, auth, oldID, 1, &tg.HelpGetConfigRequest{}) + waitDuplicateHandlerStarts(t, handler.started, "original request") + + // Same ID/seq is a retransmission, not a second business operation. + sendEncryptedWithSeq(t, conn, cipher, auth, oldID, 1, &tg.HelpGetConfigRequest{}) + // A client-side init/layer rewrap legitimately assigns a new ID and next seq. + sendEncryptedWithSeq(t, conn, cipher, auth, newID, 3, &tg.HelpGetConfigRequest{}) + waitDuplicateHandlerStarts(t, handler.started, "fresh request behind duplicate") + if got := handler.calls.Load(); got != 2 { + t.Fatalf("handler calls before release = %d, want original + fresh only", got) + } + + handler.unblock() + frames := collectReplyFrames(t, conn, cipher, auth.AuthKey, map[uint32]int{ + proto.ResultTypeID: 2, + }) + results := make(map[int64]int) + for _, frame := range frames { + if frame.TypeID != proto.ResultTypeID { + continue + } + var result proto.Result + if err := result.Decode(frame.Plain); err != nil { + t.Fatalf("decode rpc_result: %v", err) + } + results[result.RequestMessageID]++ + } + if results[oldID] != 1 || results[newID] != 1 || len(results) != 2 { + t.Fatalf("rpc_result counts = %+v, want one for old and one for fresh id", results) + } + if got := handler.calls.Load(); got != 2 { + t.Fatalf("final handler calls = %d, want 2", got) + } + deadline := time.Now().Add(2 * time.Second) + for server.rpcResults.flightLimit.snapshot() != 0 && time.Now().Before(deadline) { + time.Sleep(time.Millisecond) + } + if got := server.rpcResults.flightLimit.snapshot(); got != 0 { + t.Fatalf("duplicate admission leaked flight slots: %d", got) + } +} + +func waitDuplicateHandlerStarts(t *testing.T, started <-chan struct{}, what string) { + t.Helper() + select { + case <-started: + case <-time.After(2 * time.Second): + t.Fatalf("%s did not reach handler; socket read loop is likely blocked on a duplicate", what) + } +} diff --git a/internal/mtprotoedge/encrypted.go b/internal/mtprotoedge/encrypted.go index fceb4244..76bbcaba 100644 --- a/internal/mtprotoedge/encrypted.go +++ b/internal/mtprotoedge/encrypted.go @@ -10,6 +10,7 @@ import ( "fmt" "io" "math" + "sync" "sync/atomic" "time" @@ -19,6 +20,7 @@ import ( "github.com/gotd/td/crypto" "github.com/gotd/td/mt" "github.com/gotd/td/proto" + "github.com/gotd/td/proto/codec" "github.com/gotd/td/tgerr" "github.com/gotd/td/transport" @@ -30,11 +32,16 @@ import ( // connState 是单连接的 MTProto 运行态。 type connState struct { - sentCreated bool - seen map[int64]clientMsgRecord // 已处理的 client msg_id，用于幂等和 msgs_state_req - order []int64 - minSeen int64 - maxSeen int64 + // createdFloor is the smallest client msg_id covered by the latest + // new_session_created notification for this server-side session generation. + // It only moves down: official clients resend every request below first_msg_id, + // so advertising an outer container id while accepting smaller inner ids would + // orphan the original rpc_result messages. + createdFloor int64 + seen map[int64]clientMsgRecord // 已处理的 client msg_id，用于幂等和 msgs_state_req + order []int64 + minSeen int64 + maxSeen int64 // maxContentMsgID/maxContentSeqNo 是已接受 content 消息的 msg_id / seq_no 高水位， // 供 validateSeq 的 O(1) 快路径使用（客户端正常发送严格递增）。二者只增不减、 // 不随 seen 淘汰回退——快路径只接受「全扫描也必然接受」的子集，其余回落全扫描。 @@ -46,6 +53,9 @@ type clientMsgRecord struct { state byte seqNo int32 content bool + // service is the constructor class admitted with this msg_id. A duplicate + // uses this committed class and never decodes/executes its replacement body. + service bool } func newConnState() *connState { @@ -76,11 +86,11 @@ const ( // MTProto service vectors operate on bounded connection tracking tables. Accepting more IDs // only burns decode/CPU and cannot improve the result. maxServiceMessageIDs = 4096 - // A decoded container descriptor is 48 bytes on 64-bit Go today. Charge 64 bytes per entry - // before allocating the exact-size slice so allocator rounding and future field growth remain - // inside the process-wide inbound budget. Message bodies stay as zero-copy views of the already - // charged plaintext frame/gzip expansion. - containerDescriptorBudgetBytes = 64 + // Charge each container entry for the decoded proto.Message view plus the staged connState, + // action and ACK descriptors retained by the single-pass inbound plan. Bodies remain zero-copy + // views of the already charged plaintext frame/gzip expansion; RPC copies have a separate batch + // admission budget. + containerDescriptorBudgetBytes = 192 msgStateUnknown byte = 1 msgStateNotReceived byte = 2 @@ -97,11 +107,14 @@ const ( badMsgContainer = 64 ) +var errActivationAuthKeyRejected = errors.New("activation auth key no longer exists") + // handleEncrypted 解密加密消息，按需注册连接，处理服务消息并分发明文 payload。 // 返回（可能新建/更新的）当前连接对象，供 serveConn 维护生命周期。 // fetchedKey 非 nil 表示本帧的 auth key 是刚从 AuthKeyStore 查出的（首帧/换 auth key/被销毁 -// 后回落）；为 nil 表示走快路径——serveConn 判定 current 仍持同一未销毁的 auth key，直接复用 -// current.key/current.salt 解密，既不回查 AuthKeyStore 也不重建 store.AuthKeyData。 +// 后回落）；为 nil 表示走连接缓存快路径——serveConn 判定 current 仍持同一未销毁的 auth key， +// 直接复用 current.key/current.salt 解密。任何 provisional 在 claim 建立后、发 required +// control 前都会最终回查 AuthKeyStore，使外部撤销与 activation 线性化。 // plain 是 serveConn 持有的复用明文缓冲，frame 的 slice 仅在下一帧解密前有效。 func (s *Server) handleEncrypted(ctx context.Context, tc transport.Conn, cs *connState, current *Conn, fetchedKey *store.AuthKeyData, b, plain *bin.Buffer) (*Conn, error) { var key crypto.AuthKey @@ -120,30 +133,27 @@ func (s *Server) handleEncrypted(ctx context.Context, tc transport.Conn, cs *con return current, fmt.Errorf("decrypt: %w", err) } - if frame.salt != serverSalt { - c := current - temp := false - if c == nil || c.sessionID != frame.sessionID || c.authKeyID != key.ID { - c = s.newConn(tc, key, frame.sessionID, serverSalt) - temp = true - } - err := s.sendBadServerSalt(ctx, c, frame.messageID, frame.seqNo, serverSalt) - if temp { - c.Close() - } - return current, err - } - - // 首个加密消息或 session 变化时（重新）注册连接到 SessionManager。 + // 首个加密消息（即使 salt 尚未修正）或 session 变化时创建并保留唯一的 + // provisional Conn。同一物理 transport 换 session 必须先不可逆地 fence/drain + // 旧 writer，再把物理 lease 原子转交给新 generation。为每个 bad_server_salt + // 临时创建 Conn 会在同一 socket 上启动多个 outbound actor，Android 的启动重试 + // 风暴随即变成并发写和重复结果放大。 if current == nil || current.sessionID != frame.sessionID || current.authKeyID != key.ID { if current != nil { cs.reset() - } - if current != nil { + current.beginTerminalShutdown() s.conns.Unregister(current) - current.Close() + if !current.waitOutboundShutdownUntil(forceCloseBatchTimeout) { + return current, errors.New("previous session outbound writer did not stop") + } + nextLease, ok := current.transferTransportOwnership() + if !ok { + return current, ErrConnClosed + } + current = s.newConnWithLease(nextLease, key, frame.sessionID, serverSalt) + } else { + current = s.newConn(tc, key, frame.sessionID, serverSalt) } - current = s.newConn(tc, key, frame.sessionID, serverSalt) // 注册即播种协商 layer：新 Conn 的 clientLayer 为 0（=canonical 227），若等到 // 首条 RPC 的 Dispatch 返回后才刷新，重连老客户端在首条 RPC handler 执行期间 // 收到的 pending flush / 并发 push 会漏降级。进程内重连时 rpc 层留有 @@ -153,68 +163,120 @@ func (s *Server) handleEncrypted(ctx context.Context, tc transport.Conn, cs *con current.SetClientLayer(layer) } } - s.conns.Register(current) } - s.maybePersistSession(ctx, current, frame.sessionID, key.ID, serverSalt) + if frame.salt != serverSalt { + // bad_server_salt 是修正后重试的物理屏障：payload 与加密 envelope 都必须携带 + // 同一个权威 salt，写失败则该 provisional/active Conn 不得继续接收状态。 + return current, s.sendBadServerSalt(ctx, current, frame.messageID, frame.seqNo, serverSalt) + } body := frame.data typeID, err := (&bin.Buffer{Buf: body}).PeekID() if err != nil { return current, fmt.Errorf("peek encrypted payload type id: %w", err) } - if code := validateClientEnvelope(s.clock.Now(), frame.messageID, frame.seqNo, typeID); code != 0 { - s.log.Debug("Sending bad_msg_notification", - zap.Int64("msg_id", frame.messageID), - zap.Int32("seq_no", frame.seqNo), - zap.Uint32("type_id", typeID), - zap.Int("code", code), - ) - return current, s.sendBadMsg(ctx, current, frame.messageID, frame.seqNo, code) + plan, err := s.preflightInbound(cs, frame.messageID, frame.seqNo, body) + if err != nil { + var bad *dispatchBadMsgError + if errors.As(err, &bad) { + s.log.Debug("Sending bad_msg_notification", + zap.Int64("msg_id", bad.msgID), + zap.Int32("seq_no", bad.seqNo), + zap.Uint32("type_id", typeID), + zap.Int("code", bad.code), + ) + return current, s.sendBadMsg(ctx, current, bad.msgID, bad.seqNo, bad.code) + } + return current, err } - if err := sendQuickAckIfRequested(ctx, tc, key, frame.plaintext, s.writeTimeout); err != nil { + defer plan.close() + if err := s.prepareInboundRPCBatch(ctx, current, plan); err != nil { + return current, err + } + if err := sendQuickAckIfRequested(ctx, current.transport, key, frame.plaintext, s.writeTimeout); err != nil { return current, err } - content := clientMessageNeedsAck(typeID) - if record, ok := cs.seenRecord(frame.messageID); ok { - s.log.Debug("Duplicate msg_id; replay cached result if available", zap.Int64("msg_id", frame.messageID)) - if err := s.replayRPCResultByRequest(ctx, current, frame.messageID); err != nil { + moveCreatedFloor := cs.createdFloor == 0 || plan.logicalMin < cs.createdFloor + claimPending := false + if current.lifecycleState() == connLifecycleProvisional { + if !moveCreatedFloor { + return current, errors.New("provisional session has no new_session_created boundary") + } + if err := s.conns.BeginActivation(current); err != nil { return current, err } - if !record.content { - return current, nil + claimPending = true + defer func() { + if claimPending { + s.conns.AbortActivation(current) + } + }() + + // BeginActivation has installed current in claimsByAuth, which is the shared + // linearization domain with auth-key revocation. A delete that completed before + // the claim is visible here as !found; a delete after this read must observe and + // fence the claim. This final check intentionally covers every activation path: + // first correct-salt frame, retained bad-salt provisional and session transfer. + fresh, found, getErr := s.authKeys.Get(ctx, current.authKeyID) + if getErr != nil { + return current, fmt.Errorf("revalidate activation auth key: %w", getErr) + } + if !found || fresh.ID != current.authKeyID || fresh.Value != [256]byte(current.key.Value) { + // Send the terminal protocol error while the claim still owns a live writer; + // the deferred abort then fences and removes it before serveConn returns. + if sendErr := s.sendProtoError(ctx, current.transport, codec.CodeAuthKeyNotFound); sendErr != nil { + return current, sendErr + } + return current, errActivationAuthKeyRejected + } + if current.terminal.Load() || !current.isPhysicalTransportCurrentOpen() { + return current, ErrConnClosed } - return current, s.sendAck(ctx, current, frame.messageID) } - if code := cs.validateSeq(frame.messageID, frame.seqNo, content); code != 0 { - s.log.Debug("Sending bad_msg_notification", - zap.Int64("msg_id", frame.messageID), + if moveCreatedFloor { + s.log.Debug("Sending new_session_created", + zap.Int64("first_msg_id", plan.logicalMin), + zap.Int64("outer_msg_id", frame.messageID), zap.Int32("seq_no", frame.seqNo), - zap.Uint32("type_id", typeID), - zap.Int("code", code), ) - return current, s.sendBadMsg(ctx, current, frame.messageID, frame.seqNo, code) - } - cs.track(frame.messageID, frame.seqNo, content, msgStateReceived) - - if !cs.sentCreated { - cs.sentCreated = true - s.log.Debug("Sending new_session_created", zap.Int64("msg_id", frame.messageID), zap.Int32("seq_no", frame.seqNo)) - if err := s.sendNewSessionCreated(ctx, current, frame.messageID); err != nil { + if err := s.sendNewSessionCreated(ctx, current, plan.logicalMin); err != nil { return current, err } } + if !current.isPhysicalTransportCurrentOpen() { + return current, ErrConnClosed + } + if claimPending { + if err := s.conns.PublishActivation(current); err != nil { + return current, err + } + claimPending = false + } + if moveCreatedFloor { + cs.createdFloor = plan.logicalMin + } + plan.commitState(cs) + s.maybePersistSession(ctx, current, frame.sessionID, key.ID, serverSalt) - var acks []int64 - if err := s.dispatch(ctx, cs, current, frame.messageID, frame.seqNo, &bin.Buffer{Buf: body}, &acks); err != nil { + if err := s.executeInboundPlan(ctx, cs, current, plan); err != nil { return current, err } - if len(acks) > 0 { - if err := s.sendAck(ctx, current, acks...); err != nil { + if err := plan.commitRPCBatch(); err != nil { + return current, err + } + if len(plan.ackIDs) > 0 { + if err := s.sendAck(ctx, current, plan.ackIDs...); err != nil { return current, err } } + // An overlapping cross-connection owner may run until RPCTimeout. ACK the + // accepted duplicate before joining it so the new client does not build a + // retransmit storm while its read loop waits for the shared result. + if err := s.executePendingRPCReplays(ctx, current, plan); err != nil { + return current, err + } return current, nil } @@ -284,236 +346,37 @@ func clientQuickAckToken(key crypto.AuthKey, plaintext []byte) uint32 { // 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 { - expanded := 0 - return s.dispatchWithBudget(ctx, cs, c, msgID, seqNo, b, acks, dispatchBudget{expanded: &expanded}) -} - -type dispatchBudget struct { - depth int - containerDepth int - expanded *int -} - -func (s *Server) dispatchWithBudget(ctx context.Context, cs *connState, c *Conn, msgID int64, seqNo int32, b *bin.Buffer, acks *[]int64, budget dispatchBudget) error { - if budget.depth > maxDispatchDepth { - return fmt.Errorf("mtproto wrapper depth %d exceeds %d", budget.depth, maxDispatchDepth) - } - id, err := b.PeekID() + plan, err := s.preflightInbound(cs, msgID, seqNo, b.Buf) if err != nil { - return fmt.Errorf("peek type id: %w", err) - } - ackContent := func() { - if clientMessageNeedsAck(id) { - *acks = append(*acks, msgID) - } + var bad *dispatchBadMsgError + if errors.As(err, &bad) && c != nil { + return s.sendBadMsg(ctx, c, bad.msgID, bad.seqNo, bad.code) + } + return err } + defer plan.close() + plan.commitState(cs) + *acks = append(*acks, plan.ackIDs...) + return s.executeInboundPlan(ctx, cs, c, plan) +} - switch id { - case proto.GZIPTypeID: - data, releaseExpansion, err := s.decodeGZIPWithGlobalBudget(b) - if err != nil { - return fmt.Errorf("decode gzip: %w", err) - } - defer releaseExpansion() - *budget.expanded += len(data) - if *budget.expanded > maxDispatchExpandedBytes { - return fmt.Errorf("cumulative gzip expansion %d exceeds %d", *budget.expanded, maxDispatchExpandedBytes) - } - budget.depth++ - return s.dispatchWithBudget(ctx, cs, c, msgID, seqNo, &bin.Buffer{Buf: data}, acks, budget) +// dispatchBadMsgError carries a protocol-level rejection discovered during the +// side-effect-free wrapper/container preflight. The caller emits the single +// bad_msg_notification only after the whole container has been inspected. +type dispatchBadMsgError struct { + msgID int64 + seqNo int32 + code int +} - case proto.MessageContainerTypeID: - if budget.containerDepth != 0 { - return s.sendBadMsg(ctx, c, msgID, seqNo, badMsgContainer) - } - count, err := containerMessageCount(b) - if err != nil { - return fmt.Errorf("decode container count: %w", err) - } - if count > maxContainerMessages { - return s.sendBadMsg(ctx, c, msgID, seqNo, badMsgContainer) - } - container, releaseContainer, err := s.decodeMessageContainerViews(b, count) - if err != nil { - return fmt.Errorf("decode container: %w", err) - } - defer releaseContainer() - if code := validateClientContainer(msgID, seqNo, container); code != 0 { - return s.sendBadMsg(ctx, c, msgID, seqNo, code) - } - budget.depth++ - budget.containerDepth++ - for i := range container.Messages { - m := container.Messages[i] - typeID, err := (&bin.Buffer{Buf: m.Body}).PeekID() - if err != nil { - return fmt.Errorf("peek container message type id: %w", err) - } - content := clientMessageNeedsAck(typeID) - if record, ok := cs.seenRecord(m.ID); ok { - if err := s.replayRPCResultByRequest(ctx, c, m.ID); err != nil { - return err - } - if record.content { - *acks = append(*acks, m.ID) - } - continue - } - if code := cs.validateSeq(m.ID, int32(m.SeqNo), content); code != 0 { - return s.sendBadMsg(ctx, c, m.ID, int32(m.SeqNo), code) - } - cs.track(m.ID, int32(m.SeqNo), content, msgStateReceived) - if err := s.dispatchWithBudget(ctx, cs, c, m.ID, int32(m.SeqNo), &bin.Buffer{Buf: m.Body}, acks, budget); err != nil { - return err - } - } - return nil - - case mt.PingRequestTypeID: - var ping mt.PingRequest - if err := ping.Decode(b); err != nil { - return fmt.Errorf("decode ping: %w", err) - } - ackContent() - return s.sendPong(ctx, c, msgID, ping.PingID) - - case mt.PingDelayDisconnectRequestTypeID: - var ping mt.PingDelayDisconnectRequest - if err := ping.Decode(b); err != nil { - return fmt.Errorf("decode ping_delay_disconnect: %w", err) - } - ackContent() - return s.sendPong(ctx, c, msgID, ping.PingID) - - case mt.GetFutureSaltsRequestTypeID: - var req mt.GetFutureSaltsRequest - if err := req.Decode(b); err != nil { - return fmt.Errorf("decode get_future_salts: %w", err) - } - ackContent() - return s.sendFutureSalts(ctx, c, msgID, req.Num) - - case mt.MsgsAckTypeID: - if err := validateFirstVectorCount(b, maxServiceMessageIDs); err != nil { - return fmt.Errorf("msgs_ack vector: %w", err) - } - var ack mt.MsgsAck - if err := ack.Decode(b); err != nil { - return fmt.Errorf("decode msgs_ack: %w", err) - } - c.AckServerMessages(ack.MsgIDs) - s.log.Debug("Received msgs_ack", zap.Int64s("msg_ids", ack.MsgIDs)) - return nil - - case mt.MsgsStateReqTypeID: - if err := validateFirstVectorCount(b, maxServiceMessageIDs); err != nil { - return fmt.Errorf("msgs_state_req vector: %w", err) - } - var req mt.MsgsStateReq - if err := req.Decode(b); err != nil { - return fmt.Errorf("decode msgs_state_req: %w", err) - } - ackContent() - outgoing, err := c.OutgoingStateInfo(ctx, req.MsgIDs) - if err != nil { - return err - } - return s.sendMsgsStateInfo(ctx, c, msgID, mergeStateInfo(outgoing, cs.stateInfo(req.MsgIDs))) - - case mt.MsgResendReqTypeID: - if err := validateFirstVectorCount(b, maxServiceMessageIDs); err != nil { - return fmt.Errorf("msg_resend_req vector: %w", err) - } - var req mt.MsgResendReq - if err := req.Decode(b); err != nil { - return fmt.Errorf("decode msg_resend_req: %w", err) - } - ackContent() - outgoing, err := c.ResendMessages(ctx, req.MsgIDs) - if err != nil { - return err - } - return s.sendMsgsStateInfo(ctx, c, msgID, mergeStateInfo(outgoing, cs.stateInfo(req.MsgIDs))) - - case mt.MsgsStateInfoTypeID: - reqMsgID, info, err := msgsStateInfoView(b) - if err != nil { - return fmt.Errorf("decode msgs_state_info: %w", err) - } - s.log.Debug("Received msgs_state_info", zap.Int64("req_msg_id", reqMsgID), zap.Int("len", len(info))) - return nil - - case mt.MsgsAllInfoTypeID: - count, info, err := msgsAllInfoView(b) - if err != nil { - return fmt.Errorf("decode msgs_all_info: %w", err) - } - if len(info) != count { - return fmt.Errorf("decode msgs_all_info: info length %d does not match msg_ids %d", len(info), count) - } - s.log.Debug("Received msgs_all_info", zap.Int("msg_ids", count), zap.Int("len", len(info))) - return nil - - case mt.DestroySessionRequestTypeID: - var req mt.DestroySessionRequest - if err := req.Decode(b); err != nil { - return fmt.Errorf("decode destroy_session: %w", err) - } - ackContent() - return s.sendDestroySession(ctx, c, req.SessionID) - - case mt.HTTPWaitRequestTypeID: - var req mt.HTTPWaitRequest - if err := req.Decode(b); err != nil { - return fmt.Errorf("decode http_wait: %w", err) - } - s.log.Debug("Received http_wait", - zap.Int("max_delay", req.MaxDelay), - zap.Int("wait_after", req.WaitAfter), - zap.Int("max_wait", req.MaxWait), - ) - return nil - - case mt.RPCDropAnswerRequestTypeID: - var req mt.RPCDropAnswerRequest - if err := req.Decode(b); err != nil { - return fmt.Errorf("decode rpc_drop_answer: %w", err) - } - ackContent() - s.log.Debug("Received rpc_drop_answer", zap.Int64("req_msg_id", req.ReqMsgID)) - return s.sendResult(ctx, c, msgID, &mt.RPCAnswerUnknown{}) - - case destroyAuthKeyRequestTypeID: - var req destroyAuthKeyRequest - if err := req.Decode(b); err != nil { - return err - } - ackContent() - s.log.Debug("Received destroy_auth_key", zap.String("auth_key_id", c.authKeyHex)) - // 真正销毁：删密钥库记录（每帧回查，删除后该 key 的入站帧立即失效）并主动 - // 断开同 key 的其他连接——出站推送用连接持有的密钥副本加密、不回查密钥库， - // 不断开的话被销毁 key 的空闲连接仍能持续收到推送。发起连接除外：响应要 - // 先送达，它的下一帧会因密钥缺失自然断开。授权（authorizations）不在此清理， - // destroy_auth_key 是 PFS 密钥轮换的清理动作，不等于登出。 - if err := s.authKeys.Delete(ctx, c.authKeyID); err != nil { - s.log.Warn("Delete auth key failed", zap.String("auth_key_id", c.authKeyHex), zap.Error(err)) - return c.SendAsync(ctx, proto.MessageServerResponse, &destroyAuthKeyFail{}) - } - // 标记密钥已销毁：发起连接被 CloseSessionsForRawAuthKeyExcept 排除（响应需先送达）， - // 它下一帧不能再走 serveConn 的密钥复用快路径，须回落到 Get→AuthKeyNotFound 自然失效。 - c.keyDestroyed.Store(true) - s.conns.CloseSessionsForRawAuthKeyExcept(c.authKeyID, c.sessionID) - return c.SendAsync(ctx, proto.MessageServerResponse, &destroyAuthKeyOk{}) - - default: - ackContent() - return s.enqueueRPC(ctx, c, msgID, id, b) - } +func (e *dispatchBadMsgError) Error() string { + return fmt.Sprintf("bad client message %d/%d: code %d", e.msgID, e.seqNo, e.code) } // decodeGZIPWithGlobalBudget reserves the maximum single-wrapper output before // decompression starts. Once the actual size is known the excess reservation is -// returned, while the actual output remains charged through recursive dispatch. +// returned, while the actual output remains charged until the inbound plan is +// executed or aborted. // This closes the gap where every connection read goroutine could otherwise hold // an unaccounted 10 MiB expansion before the shared RPC scheduler saw the body. func (s *Server) decodeGZIPWithGlobalBudget(b *bin.Buffer) ([]byte, func(), error) { @@ -571,8 +434,14 @@ func gzipPackedBytesView(b *bin.Buffer) ([]byte, error) { if binary.LittleEndian.Uint32(b.Buf[:4]) != proto.GZIPTypeID { return nil, fmt.Errorf("unexpected gzip constructor %#x", binary.LittleEndian.Uint32(b.Buf[:4])) } - payload, _, err := tlBytesView(b.Buf[4:], -1) - return payload, err + payload, consumed, err := tlBytesView(b.Buf[4:], -1) + if err != nil { + return nil, err + } + if 4+consumed != len(b.Buf) { + return nil, fmt.Errorf("gzip_packed has %d trailing bytes", len(b.Buf)-(4+consumed)) + } + return payload, nil } // tlBytesView validates one TL bytes envelope and returns a view into the caller-owned buffer. @@ -604,9 +473,9 @@ func tlBytesView(raw []byte, maxPayload int) ([]byte, int, error) { } // decodeMessageContainerViews parses the container without proto.Message.Decode's per-body -// copies. Bodies are immutable views of b and stay alive only for this synchronous dispatch; -// enqueueRPC takes its own budgeted copy before returning. Only the exact-size descriptor slice -// is new memory, and that allocation is reserved globally first. +// copies. Bodies stay as immutable views through the single-pass inbound plan; batch admission +// takes independent RPC copies before the backing frame/expansion is released. The descriptor +// reservation also covers staged state, actions and ACK metadata retained by that plan. func (s *Server) decodeMessageContainerViews(b *bin.Buffer, count int) (proto.MessageContainer, func(), error) { release := func() {} if b == nil || len(b.Buf) < 8 { @@ -663,6 +532,10 @@ func (s *Server) decodeMessageContainerViews(b *bin.Buffer, count int) (proto.Me } offset = bodyEnd } + if offset != len(b.Buf) { + release() + return proto.MessageContainer{}, func() {}, fmt.Errorf("message container has %d trailing bytes", len(b.Buf)-offset) + } return proto.MessageContainer{Messages: messages}, release, nil } @@ -673,10 +546,13 @@ func msgsStateInfoView(b *bin.Buffer) (int64, []byte, error) { if got := binary.LittleEndian.Uint32(b.Buf[:4]); got != mt.MsgsStateInfoTypeID { return 0, nil, fmt.Errorf("unexpected constructor %#x", got) } - info, _, err := tlBytesView(b.Buf[12:], maxServiceMessageIDs) + info, consumed, err := tlBytesView(b.Buf[12:], maxServiceMessageIDs) if err != nil { return 0, nil, err } + if 12+consumed != len(b.Buf) { + return 0, nil, fmt.Errorf("msgs_state_info has %d trailing bytes", len(b.Buf)-(12+consumed)) + } return int64(binary.LittleEndian.Uint64(b.Buf[4:12])), info, nil } @@ -691,10 +567,13 @@ func msgsAllInfoView(b *bin.Buffer) (int, []byte, error) { return 0, nil, io.ErrUnexpectedEOF } offset := 12 + count*8 - info, _, err := tlBytesView(b.Buf[offset:], maxServiceMessageIDs) + info, consumed, err := tlBytesView(b.Buf[offset:], maxServiceMessageIDs) if err != nil { return 0, nil, err } + if offset+consumed != len(b.Buf) { + return 0, nil, fmt.Errorf("msgs_all_info has %d trailing bytes", len(b.Buf)-(offset+consumed)) + } return count, info, nil } @@ -750,15 +629,41 @@ func mergeStateInfo(primary, fallback []byte) []byte { // （已 PeekID 过一次），method 只解析一次并随任务透传，避免同一请求三处重复 PeekID/typeName。 func (s *Server) enqueueRPC(ctx context.Context, c *Conn, msgID int64, typeID uint32, request *bin.Buffer) error { method := s.typeName(typeID) - if cached, ok := s.cachedRPCResult(c, msgID); ok { + claim, err := s.rpcResults.Acquire(c.authKeyID, c.sessionID, msgID) + if err != nil { + if errors.Is(err, ErrRPCResultFlightCapacity) { + c.metrics.InboundRPCDropped(method, "flight_capacity") + return s.handleInboundRPCAdmissionError(ctx, c, msgID, method, ErrInboundRPCQueueFull) + } + return err + } + switch claim.state { + case rpcResultAcquireCompleted: s.log.Info("RPC duplicate replay from session cache", zap.String("method", method), zap.Int64("msg_id", msgID), zap.String("auth_key_id", c.authKeyHex), zap.Int64("session_id", c.sessionID), ) - return c.SendEncoded(ctx, proto.MessageServerResponse, cached) + return s.sendCachedRPCResult(ctx, c, claim.encoded) + case rpcResultAcquirePending: + encoded, ok, waitErr := claim.waiter.Wait(ctx) + if waitErr != nil || !ok || encoded == nil { + return waitErr + } + return s.sendCachedRPCResult(ctx, c, encoded) + case rpcResultAcquireOwner: + // Ownership transfers to the queued task only after commit succeeds. + default: + return ErrRPCResultFlightInvalid } + owner := claim.owner + transferred := false + defer func() { + if !transferred { + owner.Abort() + } + }() // 两级条数/字节预算必须先于 Copy：对抗客户端不能用大量满尺寸请求在“判断队列满” // 之前制造一轮无上限的临时 body 分配。reservation 在 commit/abort 间唯一持有预算。 reservation, err := c.reserveInboundRPC(ctx, method, request.Len()) @@ -767,11 +672,21 @@ func (s *Server) enqueueRPC(ctx context.Context, c *Conn, msgID int64, typeID ui } defer reservation.abort() body := request.Copy() - responseGate := &rpcResponseGate{} + err = reservation.commit(s.newInboundRPCTask(c, msgID, method, body, owner)) + transferred = err == nil + return s.handleInboundRPCAdmissionError(ctx, c, msgID, method, err) +} + +// newInboundRPCTask builds the exactly-once timeout/result gate shared by the +// single-message and atomic container-batch admission paths. body must already +// be an independently owned, budgeted copy. +func (s *Server) newInboundRPCTask(c *Conn, msgID int64, method string, body []byte, owner *rpcResultOwnerLease) inboundRPC { + responseGate := newRPCResponseGate() timeoutResponse := func() { if !responseGate.tryTimeout() { return } + defer responseGate.finish() // 原 task context 已到期，使用有界的新 context 回显明确的可重试超时； // 500 保持 TDesktop 默认重试语义，错误名区分于容量型 FLOOD_WAIT。 writeTimeout := c.writeTimeout @@ -793,10 +708,29 @@ func (s *Server) enqueueRPC(ctx context.Context, c *Conn, msgID int64, typeID ui ) } } - err = reservation.commit(inboundRPC{ + return inboundRPC{ method: method, size: len(body), onTimeout: timeoutResponse, + release: func() { + if owner == nil { + return + } + // A running deadline callback owns the same terminal flight but executes + // on context.AfterFunc's goroutine. Let its bounded delivery complete before + // deciding that the task produced no result; otherwise release could Abort + // the owner while RPC_TIMEOUT is already inside the physical writer. + if !responseGate.wait(requiredControlMaxWait + time.Second) { + c.fenceUndeliveredRPCResult() + } + if owner.Abort() { + // connState already remembers this request. If a committed task exits + // without publishing any terminal rpc_result, a same-Conn retransmit + // would otherwise be ACKed forever. Force a fresh physical generation + // where the request can be admitted again. + c.fenceUndeliveredRPCResult() + } + }, run: func(taskCtx context.Context) error { // body 是预算成功后生成的独立副本，且每个任务只 run 一次， // 无需再 append 拷贝；直接复用，省掉一份 inbound 在途内存。 @@ -816,8 +750,7 @@ func (s *Server) enqueueRPC(ctx context.Context, c *Conn, msgID int64, typeID ui } return nil }, - }) - return s.handleInboundRPCAdmissionError(ctx, c, msgID, method, err) + } } func (s *Server) handleInboundRPCAdmissionError(ctx context.Context, c *Conn, msgID int64, method string, err error) error { @@ -839,8 +772,18 @@ func (s *Server) handleInboundRPCAdmissionError(ctx context.Context, c *Conn, ms // handleRPC 把明文 RPC 请求交给 RPC 路由，并将结果或错误包成 rpc_result 回发。 func (s *Server) handleRPC(ctx context.Context, c *Conn, msgID int64, method string, b *bin.Buffer, responseGate *rpcResponseGate) error { if s.rpc == nil { - s.log.Warn("No RPC handler configured; dropping request", zap.String("method", method)) - return nil + if err := ctx.Err(); err != nil { + return err + } + if responseGate != nil && !responseGate.tryNormal() { + return context.DeadlineExceeded + } + defer responseGate.finish() + s.log.Warn("No RPC handler configured", zap.String("method", method)) + return s.sendResult(ctx, c, msgID, &mt.RPCError{ + ErrorCode: 500, + ErrorMessage: "NOT_IMPLEMENTED", + }) } ctx = postresponse.WithCallbacks(ctx) @@ -874,10 +817,48 @@ func (s *Server) handleRPC(ctx context.Context, c *Conn, msgID int64, method str fields = dbtrace.AppendZapFields(fields, "", dbStats.Snapshot()) if ctxErr := ctx.Err(); ctxErr != nil { - // A canceled request context means neither a success nor an error can be delivered - // with this expired context. In particular, do not cache a late successful result and - // hand it to outbound: a past write deadline would correctly poison that transport and - // could prevent the scheduler's fresh-context RPC_TIMEOUT response from being sent. + // The old physical Conn may have been replaced after the business transaction + // committed. Never write with the expired context itself: a fenced generation + // publishes cache-only for its replacement, while a still-live generation uses + // a fresh bounded delivery context. A deadline callback that already won + // responseGate has published RPC_TIMEOUT and prevents this late result from + // overwriting it. + // Only a successful business result proves useful work completed. Errors + // observed after cancellation may themselves be cancellation-derived or + // transient and must remain retryable rather than poisoning the replay cache. + var terminal bin.Encoder + if err == nil && result != nil { + terminal = result + } + if terminal != nil && (responseGate == nil || responseGate.tryNormal()) { + defer responseGate.finish() + if c.terminal.Load() || !c.isPhysicalTransportCurrentOpen() { + // Replacement/shutdown already fenced this logical generation. Cache-only + // publication is safe and lets the replacement join the completed flight. + if encoded, encodeErr := s.encodeRPCResult(c, msgID, terminal); encodeErr != nil { + s.log.Warn("Encode canceled RPC result for replay failed", append(fields, zap.Error(encodeErr))...) + } else { + s.storeRPCResult(c, msgID, encoded) + postresponse.Run(context.WithoutCancel(ctx)) + } + } else { + // An individual RPC deadline can expire while the physical connection is + // still healthy. Use a fresh bounded delivery context; publishing cache-only + // here would strand same-Conn duplicates behind an ACK with no result. + writeTimeout := c.writeTimeout + if writeTimeout <= 0 || writeTimeout > 5*time.Second { + writeTimeout = 5 * time.Second + } + responseCtx, cancel := context.WithTimeout(context.Background(), writeTimeout) + sendErr := s.sendResult(responseCtx, c, msgID, terminal) + cancel() + if sendErr != nil { + s.log.Debug("Send canceled RPC result failed", append(fields, zap.Error(sendErr))...) + } else { + postresponse.Run(context.WithoutCancel(ctx)) + } + } + } cancelFields := append(fields, zap.NamedError("context_error", ctxErr)) if err != nil { cancelFields = append(cancelFields, zap.NamedError("dispatch_error", err)) @@ -891,6 +872,7 @@ func (s *Server) handleRPC(ctx context.Context, c *Conn, msgID int64, method str s.log.Info("RPC result suppressed after timeout", fields...) return context.DeadlineExceeded } + defer responseGate.finish() if err != nil { var rpcErr *tgerr.Error @@ -921,6 +903,12 @@ func (s *Server) handleRPC(ctx context.Context, c *Conn, msgID int64, method str // first owns the response, and the other path becomes a no-op. type rpcResponseGate struct { state atomic.Uint32 + done chan struct{} + once sync.Once +} + +func newRPCResponseGate() *rpcResponseGate { + return &rpcResponseGate{done: make(chan struct{})} } func (g *rpcResponseGate) tryNormal() bool { @@ -931,14 +919,84 @@ func (g *rpcResponseGate) tryTimeout() bool { return g != nil && g.state.CompareAndSwap(0, 2) } +func (g *rpcResponseGate) finish() { + if g == nil || g.done == nil { + return + } + g.once.Do(func() { close(g.done) }) +} + +// wait returns once there is no response owner, the winning response has +// completed deliver-or-fence publication, or the bounded safety deadline wins. +func (g *rpcResponseGate) wait(timeout time.Duration) bool { + if g == nil || g.state.Load() == 0 { + return true + } + if g.done == nil { + return false + } + if timeout <= 0 { + <-g.done + return true + } + timer := time.NewTimer(timeout) + defer timer.Stop() + select { + case <-g.done: + return true + case <-timer.C: + return false + } +} + // sendResult 把 RPC 结果包成 rpc_result 并加密回发。 func (s *Server) sendResult(ctx context.Context, c *Conn, reqMsgID int64, result bin.Encoder) error { + if result == nil { + result = &mt.RPCError{ErrorCode: 500, ErrorMessage: "INTERNAL"} + } encoded, err := s.encodeRPCResult(c, reqMsgID, result) if err != nil { + // The business operation has already crossed atomic admission. Convert an + // invalid result encoder into one deterministic terminal RPC error instead of + // aborting the flight and allowing a reconnect to execute the operation again. + s.log.Warn("Encode RPC result failed; sending INTERNAL", zap.Int64("req_msg_id", reqMsgID), zap.Error(err)) + encoded, err = s.encodeRPCResult(c, reqMsgID, &mt.RPCError{ + ErrorCode: 500, + ErrorMessage: "INTERNAL", + }) + if err != nil { + c.fenceUndeliveredRPCResult() + return err + } + } + if err := c.SendEncoded(ctx, proto.MessageServerResponse, encoded); err != nil { + // A completed result may be published before delivery only after this logical + // Conn is irreversibly fenced. SendEncoded has non-writing failure paths + // (queue/context/scratch deadline); without this terminal barrier a later + // same-Conn duplicate would be ACKed while no result can ever arrive. + c.fenceUndeliveredRPCResult() + s.storeRPCResult(c, reqMsgID, encoded) return err } + // On a live Conn, completed means the rpc_result has reached the reliable byte + // stream. Same-physical duplicates can therefore be ACK-only without data loss. s.storeRPCResult(c, reqMsgID, encoded) - return c.SendEncoded(ctx, proto.MessageServerResponse, encoded) + return nil +} + +// sendCachedRPCResult preserves the delivery half of the rpc_result invariant +// for completed-flight replays: either the cached result reaches this physical +// byte stream, or this logical Conn is fenced so a replacement may retry it. +func (s *Server) sendCachedRPCResult(ctx context.Context, c *Conn, encoded *encodedOutboundMessage) error { + if encoded == nil { + c.fenceUndeliveredRPCResult() + return errors.New("nil cached rpc_result") + } + if err := c.SendEncoded(ctx, proto.MessageServerResponse, encoded); err != nil { + c.fenceUndeliveredRPCResult() + return err + } + return nil } // encodeRPCResult 编码 rpc_result。内层对象与 rpc_result 头（type_id + req_msg_id） @@ -986,13 +1044,14 @@ func (s *Server) replayRPCResultByRequest(ctx context.Context, c *Conn, reqMsgID return nil } if resent, err := c.ResendByRequest(ctx, reqMsgID); err != nil { + c.fenceUndeliveredRPCResult() return err } else if resent { s.log.Debug("Resent connection cached rpc_result for duplicate msg_id", zap.Int64("msg_id", reqMsgID)) return nil } if cached, ok := s.cachedRPCResult(c, reqMsgID); ok { - if err := c.SendEncoded(ctx, proto.MessageServerResponse, cached); err != nil { + if err := s.sendCachedRPCResult(ctx, c, cached); err != nil { return err } s.log.Debug("Resent session cached rpc_result for duplicate msg_id", zap.Int64("msg_id", reqMsgID)) @@ -1044,7 +1103,10 @@ func (s *Server) sendFutureSalts(ctx context.Context, c *Conn, reqMsgID int64, n // 复用同一值会让断线重连后的 new_session_created 被吞掉，错过的差分补拉 // （Android 收到后才调 getDifference）随之丢失。 func (s *Server) sendNewSessionCreated(ctx context.Context, c *Conn, firstMsgID int64) error { - return c.SendAsync(ctx, proto.MessageFromServer, &mt.NewSessionCreated{ + // This notification changes the client's request map and update recovery + // state. Unlike best-effort ack/pong traffic, it must be written successfully + // before the corresponding RPC batch starts executing. + return c.SendRequiredControl(ctx, proto.MessageFromServer, &mt.NewSessionCreated{ FirstMsgID: firstMsgID, UniqueID: s.newServerSessionUID(), ServerSalt: c.salt, @@ -1098,7 +1160,7 @@ func (s *Server) sendBadMsg(ctx context.Context, c *Conn, badMsgID int64, badSeq // sendBadServerSalt 通知客户端修正 server_salt（error_code 48）。 func (s *Server) sendBadServerSalt(ctx context.Context, c *Conn, badMsgID int64, badSeqno int32, newSalt int64) error { - return c.SendPriority(ctx, proto.MessageFromServer, &mt.BadServerSalt{ + return c.SendRequiredControl(ctx, proto.MessageFromServer, &mt.BadServerSalt{ BadMsgID: badMsgID, BadMsgSeqno: int(badSeqno), ErrorCode: 48, @@ -1138,25 +1200,6 @@ func validateClientEnvelope(now time.Time, msgID int64, seqNo int32, typeID uint return 0 } -func validateClientContainer(containerMsgID int64, containerSeqNo int32, container proto.MessageContainer) int { - for _, m := range container.Messages { - if m.ID >= containerMsgID || int32(m.SeqNo) > containerSeqNo { - return badMsgContainer - } - typeID, err := (&bin.Buffer{Buf: m.Body}).PeekID() - if err != nil { - return badMsgContainer - } - if typeID == proto.MessageContainerTypeID { - return badMsgContainer - } - if code := validateClientContainerEnvelope(m.ID, int32(m.SeqNo), typeID); code != 0 { - return badMsgContainer - } - } - return 0 -} - func validateClientContainerEnvelope(msgID int64, seqNo int32, typeID uint32) int { if msgID == 0 || proto.MessageID(msgID).Type() != proto.MessageFromClient { return badMsgIDInvalidBits @@ -1235,10 +1278,15 @@ func (cs *connState) validateSeq(msgID int64, seqNo int32, content bool) int { } func (cs *connState) track(msgID int64, seqNo int32, content bool, state byte) { + cs.trackInbound(msgID, seqNo, content, false, state) +} + +func (cs *connState) trackInbound(msgID int64, seqNo int32, content, service bool, state byte) { cs.seen[msgID] = clientMsgRecord{ state: state, seqNo: seqNo, content: content, + service: service, } if content { if msgID > cs.maxContentMsgID { diff --git a/internal/mtprotoedge/encrypted_test.go b/internal/mtprotoedge/encrypted_test.go index caafb0eb..9c27c2d9 100644 --- a/internal/mtprotoedge/encrypted_test.go +++ b/internal/mtprotoedge/encrypted_test.go @@ -40,25 +40,72 @@ func TestEncryptedPingPong(t *testing.T) { } } -// TestDuplicateMsgIDIdempotent 验证 M4：相同 msg_id 的重复 content 请求被幂等处理， -// server 重发已缓存的 rpc_result，并重新 ack，不重复执行业务。 +// TestDuplicateMsgIDIdempotent 验证相同物理连接上的重复 content 请求只重新 ACK， +// 原 owner 仍是唯一 rpc_result 发送者。若每次重复都重放完整结果，Android 的 +// bad_server_salt 全量重试会把一个启动批次放大成 N 轮孤儿结果并饿死新 request id。 func TestDuplicateMsgIDIdempotent(t *testing.T) { const dc = 2 - addr, pub, _ := startTestServer(t, Options{DC: dc}) + handler := &admissionCountingRPC{} + addr, pub, _ := startTestServer(t, Options{DC: dc, RPC: handler}) conn, auth, cipher := dialHandshake(t, addr, dc, pub) clientMsgID := proto.NewMessageIDGen(time.Now) msgID := clientMsgID.New(proto.MessageFromClient) - sendEncrypted(t, conn, cipher, auth, msgID, &mt.RPCDropAnswerRequest{ReqMsgID: msgID - 4}) - first := collectReplies(t, conn, cipher, auth.AuthKey, mt.MsgsAckTypeID) - mustHave(t, first, proto.ResultTypeID, "first rpc_result") + sendEncrypted(t, conn, cipher, auth, msgID, &tg.HelpGetConfigRequest{}) + collectReplyFrames(t, conn, cipher, auth.AuthKey, map[uint32]int{ + proto.ResultTypeID: 1, + mt.MsgsAckTypeID: 1, + }) + waitForAtomicCalls(t, &handler.calls, 1) - // 相同 msg_id —— 幂等：重发已有 rpc_result，并重新 ack。 - sendEncrypted(t, conn, cipher, auth, msgID, &mt.RPCDropAnswerRequest{ReqMsgID: msgID - 4}) - second := collectReplies(t, conn, cipher, auth.AuthKey, mt.MsgsAckTypeID) - mustHave(t, second, proto.ResultTypeID, "resent rpc_result") - mustHave(t, second, mt.MsgsAckTypeID, "second ack") + // 用一个小型重试风暴覆盖完成后的 duplicate 路径。TCP 仍存活时原结果已在同一 + // 可靠字节流上；每个 duplicate 只需 ACK，不应产生第二个 rpc_result。 + const duplicateCount = 16 + for i := 0; i < duplicateCount; i++ { + sendEncrypted(t, conn, cipher, auth, msgID, &tg.HelpGetConfigRequest{}) + } + frames := collectReplyFrames(t, conn, cipher, auth.AuthKey, map[uint32]int{ + mt.MsgsAckTypeID: duplicateCount, + }) + for _, frame := range frames { + if frame.TypeID == proto.ResultTypeID { + t.Fatalf("same-connection duplicate emitted an extra rpc_result") + } + } + if got := handler.calls.Load(); got != 1 { + t.Fatalf("same-connection duplicate business calls = %d, want 1", got) + } +} + +func TestServiceDuplicateCannotReplaceOriginallyAdmittedPayload(t *testing.T) { + const dc = 2 + addr, pub, _ := startTestServer(t, Options{DC: dc}) + conn, auth, cipher := dialHandshake(t, addr, dc, pub) + + ids := proto.NewMessageIDGen(time.Now) + msgID := ids.New(proto.MessageFromClient) + sendEncryptedWithSeq(t, conn, cipher, auth, msgID, 1, &mt.PingRequest{PingID: 11}) + collectReplyFrames(t, conn, cipher, auth.AuthKey, map[uint32]int{ + mt.PongTypeID: 1, + mt.MsgsAckTypeID: 1, + }) + + // Same id/seq/content parity but a destructive replacement body. Duplicate + // handling must use the original committed request class and never execute it. + sendEncryptedWithSeq(t, conn, cipher, auth, msgID, 1, &destroyAuthKeyRequest{}) + collectReplyFrames(t, conn, cipher, auth.AuthKey, map[uint32]int{mt.MsgsAckTypeID: 1}) + + freshID := ids.New(proto.MessageFromClient) + sendEncryptedWithSeq(t, conn, cipher, auth, freshID, 3, &mt.PingRequest{PingID: 22}) + replies := collectReplies(t, conn, cipher, auth.AuthKey, mt.PongTypeID) + var pong mt.Pong + if err := pong.Decode(mustHave(t, replies, mt.PongTypeID, "pong after replacement attempt")); err != nil { + t.Fatalf("decode pong: %v", err) + } + if pong.MsgID != freshID || pong.PingID != 22 { + t.Fatalf("pong after replacement = %+v, want msg=%d ping=22", pong, freshID) + } } // TestGetFutureSalts 验证 MTProto service message get_future_salts 由连接层直接响应， @@ -268,6 +315,14 @@ func TestOldMessageInFreshContainerAccepted(t *testing.T) { }) replies := collectReplies(t, conn, cipher, auth.AuthKey, mt.PongTypeID) + createdBuf := mustHave(t, replies, mt.NewSessionCreatedTypeID, "new_session_created") + var created mt.NewSessionCreated + if err := created.Decode(createdBuf); err != nil { + t.Fatalf("decode new_session_created: %v", err) + } + if created.FirstMsgID != oldPingMsgID { + t.Fatalf("new_session_created.first_msg_id = %d, want accepted inner msg_id %d", created.FirstMsgID, oldPingMsgID) + } buf := mustHave(t, replies, mt.PongTypeID, "pong") var pong mt.Pong if err := pong.Decode(buf); err != nil { @@ -382,8 +437,10 @@ func TestBadServerSalt(t *testing.T) { wrongSalt := auth.ServerSalt + 1 sendEncryptedWithSalt(t, conn, cipher, auth, wrongSalt, reqMsgID, &mt.PingRequest{PingID: 1}) - replies := collectReplies(t, conn, cipher, auth.AuthKey, mt.BadServerSaltTypeID) - buf := mustHave(t, replies, mt.BadServerSaltTypeID, "bad_server_salt") + envelope, typeID, buf := readServerMessage(t, conn, cipher, auth.AuthKey) + if typeID != mt.BadServerSaltTypeID { + t.Fatalf("bad salt reply type = %#x, want %#x", typeID, mt.BadServerSaltTypeID) + } var bad mt.BadServerSalt if err := bad.Decode(buf); err != nil { @@ -398,6 +455,11 @@ func TestBadServerSalt(t *testing.T) { if bad.NewServerSalt != auth.ServerSalt { t.Fatalf("bad_server_salt.new_server_salt = %#x, want %#x", bad.NewServerSalt, auth.ServerSalt) } + // DrKLO stores the salt from the encrypted envelope, not only the TL payload. + // A mismatch makes every correction ineffective and re-enters the resend storm. + if envelope.Salt != bad.NewServerSalt { + t.Fatalf("bad_server_salt envelope salt = %#x, payload = %#x", envelope.Salt, bad.NewServerSalt) + } } func TestBadMsgSeqOddExpected(t *testing.T) { @@ -470,7 +532,7 @@ func TestBadMsgSeqTooHigh(t *testing.T) { func TestSessionChangeResetsClientSeqState(t *testing.T) { const dc = 2 - addr, pub, _ := startTestServer(t, Options{DC: dc}) + addr, pub, srv := startTestServer(t, Options{DC: dc}) conn, auth, cipher := dialHandshake(t, addr, dc, pub) clientMsgID := proto.NewMessageIDGen(time.Now) @@ -492,6 +554,19 @@ func TestSessionChangeResetsClientSeqState(t *testing.T) { } mustHave(t, replies, mt.NewSessionCreatedTypeID, "new_session_created after session change") mustHave(t, replies, mt.MsgsAckTypeID, "msgs_ack after session change") + + oldKey := sessionKey{authKeyID: auth.AuthKey.ID, sessionID: auth.SessionID} + newKey := sessionKey{authKeyID: auth.AuthKey.ID, sessionID: nextSessionID} + srv.conns.mu.RLock() + _, oldVisible := srv.conns.bySession[oldKey] + newConn := srv.conns.bySession[newKey] + claims := len(srv.conns.claims) + online := len(srv.conns.bySession) + srv.conns.mu.RUnlock() + if oldVisible || newConn == nil || !newConn.isActive() || claims != 0 || online != 1 { + t.Fatalf("same-transport switch state: old=%v new=%p active=%v claims=%d online=%d", + oldVisible, newConn, newConn != nil && newConn.isActive(), claims, online) + } } func readBadMsgNotification(t *testing.T, conn transport.Conn, cipher crypto.Cipher, key crypto.AuthKey) mt.BadMsgNotification { diff --git a/internal/mtprotoedge/helpers_test.go b/internal/mtprotoedge/helpers_test.go index 4b976f6a..241d6c12 100644 --- a/internal/mtprotoedge/helpers_test.go +++ b/internal/mtprotoedge/helpers_test.go @@ -190,12 +190,82 @@ func collectReplies(t *testing.T, conn transport.Conn, cipher crypto.Cipher, key return got } +// serverReplyFrame preserves the wire order and encrypted envelope of a server +// reply. Tests which exercise session boundaries must not collapse replies into +// a TypeID-keyed map: both duplicate response types and their order are part of +// the observable protocol behavior. +type serverReplyFrame struct { + Message *crypto.EncryptedMessageData + TypeID uint32 + Plain *bin.Buffer +} + +// collectReplyFrames reads ordered server replies until every requested TypeID +// has been observed the requested number of times. Unrequested frames are kept +// in the returned slice so callers can assert ordering around control messages. +func collectReplyFrames( + t *testing.T, + conn transport.Conn, + cipher crypto.Cipher, + key crypto.AuthKey, + wantCounts map[uint32]int, +) []serverReplyFrame { + t.Helper() + + remaining := make(map[uint32]int, len(wantCounts)) + required := 0 + for typeID, count := range wantCounts { + if count <= 0 { + continue + } + remaining[typeID] = count + required += count + } + if required == 0 { + return nil + } + + // Keep the helper bounded while allowing unrelated control replies (notably + // msgs_ack) to be interleaved with the frames under test. + // One shared deadline bounds the whole collection. A per-frame deadline would + // multiply a missing-result failure by the maximum number of unrelated frames. + ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second) + defer cancel() + frames := make([]serverReplyFrame, 0, required) + for i := 0; i < required+16; i++ { + message, typeID, plain := readServerMessageContext(t, ctx, conn, cipher, key) + frames = append(frames, serverReplyFrame{ + Message: message, + TypeID: typeID, + Plain: plain, + }) + if count, ok := remaining[typeID]; ok { + if count == 1 { + delete(remaining, typeID) + } else { + remaining[typeID] = count - 1 + } + } + if len(remaining) == 0 { + return frames + } + } + + t.Fatalf("missing reply counts after %d frames: %+v", len(frames), remaining) + return nil +} + func readServerMessage(t *testing.T, conn transport.Conn, cipher crypto.Cipher, key crypto.AuthKey) (*crypto.EncryptedMessageData, uint32, *bin.Buffer) { t.Helper() - var buf bin.Buffer ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + return readServerMessageContext(t, ctx, conn, cipher, key) +} + +func readServerMessageContext(t *testing.T, ctx context.Context, conn transport.Conn, cipher crypto.Cipher, key crypto.AuthKey) (*crypto.EncryptedMessageData, uint32, *bin.Buffer) { + t.Helper() + var buf bin.Buffer err := conn.Recv(ctx, &buf) - cancel() if err != nil { t.Fatalf("recv server message: %v", err) } diff --git a/internal/mtprotoedge/inbound_preflight.go b/internal/mtprotoedge/inbound_preflight.go new file mode 100644 index 00000000..c6aeccf5 --- /dev/null +++ b/internal/mtprotoedge/inbound_preflight.go @@ -0,0 +1,880 @@ +package mtprotoedge + +import ( + "context" + "encoding/binary" + "errors" + "fmt" + "io" + "time" + + "go.uber.org/zap" + + "github.com/gotd/td/bin" + "github.com/gotd/td/mt" + "github.com/gotd/td/proto" +) + +type inboundItemKind uint8 + +const ( + inboundItemDuplicate inboundItemKind = iota + 1 + inboundItemServiceDuplicate + inboundItemPing + inboundItemPingDelay + inboundItemFutureSalts + inboundItemMsgsAck + inboundItemStateReq + inboundItemResendReq + inboundItemStateInfo + inboundItemAllInfo + inboundItemDestroySession + inboundItemHTTPWait + inboundItemDropAnswer + inboundItemDestroyAuthKey + inboundItemRPC + inboundItemCapacityError + inboundItemPendingRPC + // inboundItemReplayRPC is a request first observed by this physical Conn whose + // terminal result already exists in the cross-connection cache. It is distinct + // from inboundItemDuplicate: a duplicate already present in this Conn's seen + // table must only be ACKed. The original owner/result is already using the same + // reliable TCP stream, so replaying it once per retransmit wave amplifies a + // client retry burst and can starve newer request IDs behind orphan results. + inboundItemReplayRPC +) + +type inboundItem struct { + kind inboundItemKind + msgID int64 + seqNo int32 + typeID uint32 + content bool + body []byte + payload any +} + +type stagedClientMessage struct { + msgID int64 + seqNo int32 + content bool + service bool +} + +type inboundPlan struct { + items []inboundItem + staged []stagedClientMessage + ackIDs []int64 + logicalMin int64 + releases []func() + + rpcPrepared bool + rpcReservation *inboundRPCBatchReservation + rpcTasks []inboundRPC + rpcOwners []*rpcResultOwnerLease +} + +func (p *inboundPlan) close() { + if p == nil { + return + } + if p.rpcReservation != nil { + p.rpcReservation.abort() + p.rpcReservation = nil + } + for _, owner := range p.rpcOwners { + owner.Abort() + } + p.rpcOwners = nil + for i := len(p.releases) - 1; i >= 0; i-- { + p.releases[i]() + } + p.releases = nil +} + +func (p *inboundPlan) commitRPCBatch() error { + if p == nil || p.rpcReservation == nil { + return nil + } + // handleEncrypted calls this only after ownership/session-control barriers, + // connState commit and every synchronous service action have completed. The + // batch is runnable immediately; using the old deferred scheduler token here + // was not a real barrier on a busy Conn because an existing ready token could + // dequeue newly appended tasks before activateRPCBatch ran. + _, err := p.rpcReservation.commit(p.rpcTasks, false) + if err != nil { + return err + } + p.rpcReservation = nil + p.rpcTasks = nil + // Every owner is now attached to exactly one queued task. Its task.release + // path aborts if no terminal rpc_result was published. + p.rpcOwners = nil + return nil +} + +func (p *inboundPlan) includeLogicalID(msgID int64) { + if p.logicalMin == 0 || msgID < p.logicalMin { + p.logicalMin = msgID + } +} + +func (p *inboundPlan) commitState(cs *connState) { + for _, m := range p.staged { + cs.trackInbound(m.msgID, m.seqNo, m.content, m.service, msgStateReceived) + } +} + +type connStateOverlay struct { + base *connState + staged []stagedClientMessage + maxContentMsgID int64 + maxContentSeqNo int32 +} + +func newConnStateOverlay(base *connState) connStateOverlay { + return connStateOverlay{ + base: base, + maxContentMsgID: base.maxContentMsgID, + maxContentSeqNo: base.maxContentSeqNo, + } +} + +func (o *connStateOverlay) seenRecord(msgID int64) (clientMsgRecord, bool) { + for i := len(o.staged) - 1; i >= 0; i-- { + m := o.staged[i] + if m.msgID == msgID { + return clientMsgRecord{state: msgStateReceived, seqNo: m.seqNo, content: m.content, service: m.service}, true + } + } + return o.base.seenRecord(msgID) +} + +func (o *connStateOverlay) validateSeq(msgID int64, seqNo int32, content bool) int { + if !content { + return 0 + } + if msgID > o.maxContentMsgID && seqNo > o.maxContentSeqNo { + return 0 + } + for seenMsgID, record := range o.base.seen { + if !record.content { + continue + } + if seenMsgID < msgID && record.seqNo >= seqNo { + return badMsgSeqTooLow + } + if seenMsgID > msgID && record.seqNo <= seqNo { + return badMsgSeqTooHigh + } + } + for _, record := range o.staged { + if !record.content { + continue + } + if record.msgID < msgID && record.seqNo >= seqNo { + return badMsgSeqTooLow + } + if record.msgID > msgID && record.seqNo <= seqNo { + return badMsgSeqTooHigh + } + } + return 0 +} + +func (o *connStateOverlay) stage(msgID int64, seqNo int32, content, service bool) { + o.staged = append(o.staged, stagedClientMessage{msgID: msgID, seqNo: seqNo, content: content, service: service}) + if content { + if msgID > o.maxContentMsgID { + o.maxContentMsgID = msgID + } + if seqNo > o.maxContentSeqNo { + o.maxContentSeqNo = seqNo + } + } +} + +type inboundScope struct { + insideContainer bool + mustBeDuplicate bool +} + +type inboundPreflightBudget struct { + depth int + containerDepth int + expanded int + now time.Time +} + +func (s *Server) preflightInbound(cs *connState, msgID int64, seqNo int32, body []byte) (*inboundPlan, error) { + plan := &inboundPlan{logicalMin: 0} + overlay := newConnStateOverlay(cs) + // All envelope checks in one transport frame use the same clock sample. Apart + // from making boundary behavior deterministic, this lets walkInbound reject + // an invalid outer msg_id before spending CPU or memory on gzip expansion. + budget := &inboundPreflightBudget{now: s.clock.Now()} + if err := s.walkInbound(plan, &overlay, msgID, seqNo, body, inboundScope{}, budget); err != nil { + plan.close() + return nil, err + } + plan.staged = overlay.staged + if plan.logicalMin == 0 { + plan.close() + return nil, fmt.Errorf("inbound plan accepted no logical message") + } + for _, item := range plan.items { + if item.kind == inboundItemDestroyAuthKey && len(plan.items) != 1 { + plan.close() + return nil, &dispatchBadMsgError{msgID: msgID, seqNo: seqNo, code: badMsgContainer} + } + } + return plan, nil +} + +func (s *Server) walkInbound( + plan *inboundPlan, + overlay *connStateOverlay, + msgID int64, + seqNo int32, + body []byte, + scope inboundScope, + budget *inboundPreflightBudget, +) error { + if budget.depth > maxDispatchDepth { + return fmt.Errorf("mtproto wrapper depth %d exceeds %d", budget.depth, maxDispatchDepth) + } + b := &bin.Buffer{Buf: body} + typeID, err := b.PeekID() + if err != nil { + return fmt.Errorf("peek type id: %w", err) + } + if code := validateInboundMessageID(budget.now, msgID, scope.insideContainer); code != 0 { + if scope.insideContainer { + code = badMsgContainer + } + return &dispatchBadMsgError{msgID: msgID, seqNo: seqNo, code: code} + } + + // A repeated outer container still has to be decoded far enough to enumerate + // its inner request ids, but none of those already-accepted inner bodies needs + // decoding again. In particular, never reinflate a gzip body merely to discover + // the content bit we already retained with its msg_id. If an outer duplicate + // introduces an unseen inner id, reject the container before touching its body. + // A seen top-level content gzip is also safe to short-circuit: a container is + // non-content, so it cannot be hidden behind that retained record. Ambiguous + // non-content top-level gzip still expands to distinguish a container wrapper. + if record, seen := overlay.seenRecord(msgID); scope.mustBeDuplicate || + (typeID == proto.GZIPTypeID && seen && (scope.insideContainer || record.content)) { + if !seen || record.seqNo != seqNo { + return &dispatchBadMsgError{msgID: msgID, seqNo: seqNo, code: badMsgContainer} + } + return appendInboundDuplicate(plan, msgID, seqNo, typeID, record) + } + if typeID == proto.GZIPTypeID { + data, release, err := s.decodeGZIPWithGlobalBudget(b) + if err != nil { + return fmt.Errorf("decode gzip: %w", err) + } + plan.releases = append(plan.releases, release) + budget.expanded += len(data) + if budget.expanded > maxDispatchExpandedBytes { + return fmt.Errorf("cumulative gzip expansion %d exceeds %d", budget.expanded, maxDispatchExpandedBytes) + } + budget.depth++ + err = s.walkInbound(plan, overlay, msgID, seqNo, data, scope, budget) + budget.depth-- + return err + } + + if typeID == proto.MessageContainerTypeID { + if scope.insideContainer || budget.containerDepth != 0 { + return &dispatchBadMsgError{msgID: msgID, seqNo: seqNo, code: badMsgContainer} + } + if code := validateClientEnvelope(budget.now, msgID, seqNo, typeID); code != 0 { + return &dispatchBadMsgError{msgID: msgID, seqNo: seqNo, code: code} + } + outerRecord, outerSeen := overlay.seenRecord(msgID) + if outerSeen { + if outerRecord.content || outerRecord.seqNo != seqNo { + return &dispatchBadMsgError{msgID: msgID, seqNo: seqNo, code: badMsgContainer} + } + } else { + if code := overlay.validateSeq(msgID, seqNo, false); code != 0 { + return &dispatchBadMsgError{msgID: msgID, seqNo: seqNo, code: code} + } + overlay.stage(msgID, seqNo, false, false) + } + + count, err := containerMessageCount(b) + if err != nil { + return fmt.Errorf("decode container count: %w", err) + } + if count > maxContainerMessages { + return &dispatchBadMsgError{msgID: msgID, seqNo: seqNo, code: badMsgContainer} + } + container, release, err := s.decodeMessageContainerViews(b, count) + if err != nil { + return fmt.Errorf("decode container: %w", err) + } + plan.releases = append(plan.releases, release) + plan.items = growInboundItems(plan.items, count) + plan.ackIDs = growInt64s(plan.ackIDs, count) + overlay.staged = growStagedMessages(overlay.staged, count+1) + if len(container.Messages) == 0 { + plan.includeLogicalID(msgID) + return nil + } + + budget.depth++ + budget.containerDepth++ + for _, m := range container.Messages { + if m.ID >= msgID || int32(m.SeqNo) > seqNo { + return &dispatchBadMsgError{msgID: msgID, seqNo: seqNo, code: badMsgContainer} + } + if err := s.walkInbound(plan, overlay, m.ID, int32(m.SeqNo), m.Body, inboundScope{ + insideContainer: true, + mustBeDuplicate: outerSeen, + }, budget); err != nil { + return err + } + } + budget.containerDepth-- + budget.depth-- + return nil + } + + if scope.insideContainer { + if code := validateClientContainerEnvelope(msgID, seqNo, typeID); code != 0 { + return &dispatchBadMsgError{msgID: msgID, seqNo: seqNo, code: badMsgContainer} + } + } else if code := validateClientEnvelope(budget.now, msgID, seqNo, typeID); code != 0 { + return &dispatchBadMsgError{msgID: msgID, seqNo: seqNo, code: code} + } + + content := clientMessageNeedsAck(typeID) + if record, seen := overlay.seenRecord(msgID); seen { + if record.seqNo != seqNo || record.content != content { + return &dispatchBadMsgError{msgID: msgID, seqNo: seqNo, code: badMsgContainer} + } + return appendInboundDuplicate(plan, msgID, seqNo, typeID, record) + } + if scope.mustBeDuplicate { + return &dispatchBadMsgError{msgID: msgID, seqNo: seqNo, code: badMsgContainer} + } + if code := overlay.validateSeq(msgID, seqNo, content); code != 0 { + return &dispatchBadMsgError{msgID: msgID, seqNo: seqNo, code: code} + } + overlay.stage(msgID, seqNo, content, inboundTypeIsService(typeID)) + plan.includeLogicalID(msgID) + + item, err := preflightInboundItem(msgID, seqNo, typeID, content, body) + if err != nil { + return err + } + plan.items = append(plan.items, item) + if content { + plan.ackIDs = append(plan.ackIDs, msgID) + } + return nil +} + +func validateInboundMessageID(now time.Time, msgID int64, insideContainer bool) int { + if msgID == 0 || proto.MessageID(msgID).Type() != proto.MessageFromClient { + return badMsgIDInvalidBits + } + // A container's outer envelope supplies the wall-clock admission boundary for + // its inner messages. Inner ids still need the client low bits checked before + // wrapper expansion, but intentionally keep the established no-time-check rule. + if insideContainer { + return 0 + } + msgTime := proto.MessageID(msgID).Time() + if msgTime.Before(now.Add(-300 * time.Second)) { + return badMsgIDTooLow + } + if msgTime.After(now.Add(30 * time.Second)) { + return badMsgIDTooHigh + } + return 0 +} + +func appendInboundDuplicate(plan *inboundPlan, msgID int64, seqNo int32, typeID uint32, record clientMsgRecord) error { + plan.includeLogicalID(msgID) + kind := inboundItemDuplicate + if record.service { + kind = inboundItemServiceDuplicate + } + plan.items = append(plan.items, inboundItem{ + kind: kind, msgID: msgID, seqNo: seqNo, typeID: typeID, content: record.content, + }) + if record.content { + plan.ackIDs = append(plan.ackIDs, msgID) + } + return nil +} + +func inboundTypeIsService(typeID uint32) bool { + switch typeID { + case mt.PingRequestTypeID, + mt.PingDelayDisconnectRequestTypeID, + mt.GetFutureSaltsRequestTypeID, + mt.MsgsAckTypeID, + mt.MsgsStateReqTypeID, + mt.MsgResendReqTypeID, + mt.MsgsStateInfoTypeID, + mt.MsgsAllInfoTypeID, + mt.DestroySessionRequestTypeID, + mt.HTTPWaitRequestTypeID, + mt.RPCDropAnswerRequestTypeID, + destroyAuthKeyRequestTypeID: + return true + default: + return false + } +} + +func growInboundItems(items []inboundItem, extra int) []inboundItem { + if extra <= cap(items)-len(items) { + return items + } + grown := make([]inboundItem, len(items), len(items)+extra) + copy(grown, items) + return grown +} + +func growStagedMessages(items []stagedClientMessage, extra int) []stagedClientMessage { + if extra <= cap(items)-len(items) { + return items + } + grown := make([]stagedClientMessage, len(items), len(items)+extra) + copy(grown, items) + return grown +} + +func growInt64s(items []int64, extra int) []int64 { + if extra <= cap(items)-len(items) { + return items + } + grown := make([]int64, len(items), len(items)+extra) + copy(grown, items) + return grown +} + +type stateInfoPayload struct { + reqMsgID int64 + info []byte +} + +type allInfoPayload struct { + count int + info []byte +} + +// int64VectorView is a bounded, immutable view of a TL Vector. Keeping the +// encoded bytes avoids one retained []int64 allocation per service message while +// a whole container is in preflight. Execution materializes the small bounded +// slice only for the existing Conn APIs that require []int64, after all wrappers +// and every sibling message have already passed structural validation. +type int64VectorView struct { + raw []byte + count int +} + +func decodeInt64VectorView(b *bin.Buffer, expectedTypeID uint32, max int) (int64VectorView, error) { + if b == nil || len(b.Buf) < 12 { + return int64VectorView{}, io.ErrUnexpectedEOF + } + if got := binary.LittleEndian.Uint32(b.Buf[:4]); got != expectedTypeID { + return int64VectorView{}, fmt.Errorf("unexpected constructor %#x", got) + } + if got := binary.LittleEndian.Uint32(b.Buf[4:8]); got != bin.TypeVector { + return int64VectorView{}, fmt.Errorf("unexpected vector constructor %#x", got) + } + count := int(int32(binary.LittleEndian.Uint32(b.Buf[8:12]))) + if count < 0 { + return int64VectorView{}, fmt.Errorf("negative vector count %d", count) + } + if count > max { + return int64VectorView{}, fmt.Errorf("vector count %d exceeds %d", count, max) + } + if count > (len(b.Buf)-12)/8 { + return int64VectorView{}, io.ErrUnexpectedEOF + } + end := 12 + count*8 + if end != len(b.Buf) { + return int64VectorView{}, fmt.Errorf("long vector has %d trailing bytes", len(b.Buf)-end) + } + return int64VectorView{raw: b.Buf[12:end:end], count: count}, nil +} + +func (v int64VectorView) materialize() []int64 { + if v.count == 0 { + return nil + } + ids := make([]int64, v.count) + for i := range ids { + offset := i * 8 + ids[i] = int64(binary.LittleEndian.Uint64(v.raw[offset : offset+8])) + } + return ids +} + +func decodeInboundServiceExact(b *bin.Buffer, decoder bin.Decoder) error { + if err := decoder.Decode(b); err != nil { + return err + } + if b.Len() != 0 { + return fmt.Errorf("service message has %d trailing bytes", b.Len()) + } + return nil +} + +func preflightInboundItem(msgID int64, seqNo int32, typeID uint32, content bool, body []byte) (inboundItem, error) { + item := inboundItem{msgID: msgID, seqNo: seqNo, typeID: typeID, content: content, body: body} + b := &bin.Buffer{Buf: body} + switch typeID { + case mt.PingRequestTypeID: + var value mt.PingRequest + if err := decodeInboundServiceExact(b, &value); err != nil { + return item, fmt.Errorf("decode ping: %w", err) + } + item.kind, item.payload = inboundItemPing, value + case mt.PingDelayDisconnectRequestTypeID: + var value mt.PingDelayDisconnectRequest + if err := decodeInboundServiceExact(b, &value); err != nil { + return item, fmt.Errorf("decode ping_delay_disconnect: %w", err) + } + item.kind, item.payload = inboundItemPingDelay, value + case mt.GetFutureSaltsRequestTypeID: + var value mt.GetFutureSaltsRequest + if err := decodeInboundServiceExact(b, &value); err != nil { + return item, fmt.Errorf("decode get_future_salts: %w", err) + } + item.kind, item.payload = inboundItemFutureSalts, value + case mt.MsgsAckTypeID: + value, err := decodeInt64VectorView(b, mt.MsgsAckTypeID, maxServiceMessageIDs) + if err != nil { + return item, fmt.Errorf("decode msgs_ack: %w", err) + } + item.kind, item.payload = inboundItemMsgsAck, value + case mt.MsgsStateReqTypeID: + value, err := decodeInt64VectorView(b, mt.MsgsStateReqTypeID, maxServiceMessageIDs) + if err != nil { + return item, fmt.Errorf("decode msgs_state_req: %w", err) + } + item.kind, item.payload = inboundItemStateReq, value + case mt.MsgResendReqTypeID: + value, err := decodeInt64VectorView(b, mt.MsgResendReqTypeID, maxServiceMessageIDs) + if err != nil { + return item, fmt.Errorf("decode msg_resend_req: %w", err) + } + item.kind, item.payload = inboundItemResendReq, value + case mt.MsgsStateInfoTypeID: + reqMsgID, info, err := msgsStateInfoView(b) + if err != nil { + return item, fmt.Errorf("decode msgs_state_info: %w", err) + } + item.kind, item.payload = inboundItemStateInfo, stateInfoPayload{reqMsgID: reqMsgID, info: info} + case mt.MsgsAllInfoTypeID: + count, info, err := msgsAllInfoView(b) + if err != nil { + return item, fmt.Errorf("decode msgs_all_info: %w", err) + } + if len(info) != count { + return item, fmt.Errorf("decode msgs_all_info: info length %d does not match msg_ids %d", len(info), count) + } + item.kind, item.payload = inboundItemAllInfo, allInfoPayload{count: count, info: info} + case mt.DestroySessionRequestTypeID: + var value mt.DestroySessionRequest + if err := decodeInboundServiceExact(b, &value); err != nil { + return item, fmt.Errorf("decode destroy_session: %w", err) + } + item.kind, item.payload = inboundItemDestroySession, value + case mt.HTTPWaitRequestTypeID: + var value mt.HTTPWaitRequest + if err := decodeInboundServiceExact(b, &value); err != nil { + return item, fmt.Errorf("decode http_wait: %w", err) + } + item.kind, item.payload = inboundItemHTTPWait, value + case mt.RPCDropAnswerRequestTypeID: + var value mt.RPCDropAnswerRequest + if err := decodeInboundServiceExact(b, &value); err != nil { + return item, fmt.Errorf("decode rpc_drop_answer: %w", err) + } + item.kind, item.payload = inboundItemDropAnswer, value + case destroyAuthKeyRequestTypeID: + var value destroyAuthKeyRequest + if err := decodeInboundServiceExact(b, &value); err != nil { + return item, err + } + item.kind, item.payload = inboundItemDestroyAuthKey, value + default: + item.kind = inboundItemRPC + } + return item, nil +} + +// prepareInboundRPCBatch performs the whole container's count/byte admission +// before copying or scheduling any API RPC. Capacity exhaustion is converted +// into one consistent terminal FLOOD_WAIT result per uncached RPC; no business +// handler from the batch is allowed to start in that case. +func (s *Server) prepareInboundRPCBatch(ctx context.Context, c *Conn, plan *inboundPlan) error { + plan.rpcPrepared = true + // Keep service-only frames (ping/ack/http_wait) allocation-free here. These + // collections are needed only after the first real API RPC acquires ownership. + var indices []int + var specs []inboundRPCSpec + var ownersInPlan map[int64]*rpcResultOwnerLease + flightCapacity := false + for i := range plan.items { + item := &plan.items[i] + if item.kind != inboundItemRPC && item.kind != inboundItemDuplicate { + continue + } + localDuplicate := item.kind == inboundItemDuplicate + if localDuplicate { + // connState already proves this msg_id was admitted on this physical + // generation. Its original owner either still holds the flight or has + // published a result after physical write; do not consume a global flight + // slot merely to ACK the duplicate. + continue + } + method := s.typeName(item.typeID) + claim, err := s.rpcResults.Acquire(c.authKeyID, c.sessionID, item.msgID) + if errors.Is(err, ErrRPCResultFlightCapacity) { + if item.kind == inboundItemRPC { + c.metrics.InboundRPCDropped(method, "flight_capacity") + flightCapacity = true + item.kind = inboundItemCapacityError + } + continue + } + if err != nil { + return err + } + switch claim.state { + case rpcResultAcquireCompleted: + s.log.Info("RPC duplicate replay from session cache", + zap.String("method", method), + zap.Int64("msg_id", item.msgID), + zap.String("auth_key_id", c.authKeyHex), + zap.Int64("session_id", c.sessionID), + ) + item.kind = inboundItemReplayRPC + item.payload = claim.encoded + case rpcResultAcquirePending: + // A malformed/replayed container may repeat the same msg_id after this + // very plan installed its owner. More generally, any request already in + // this Conn's seen table shares the owner's reliable response path. Only + // a fresh physical replacement may wait for and replay the old result. + if ownersInPlan[item.msgID] != nil { + item.kind = inboundItemDuplicate + item.payload = nil + } else { + item.kind = inboundItemPendingRPC + item.payload = claim.waiter + } + case rpcResultAcquireOwner: + if item.kind == inboundItemDuplicate { + // connState says this request was already accepted, so absence from + // both completed and in-flight tables is not authority to execute it. + claim.owner.Abort() + continue + } + if ownersInPlan == nil { + ownersInPlan = make(map[int64]*rpcResultOwnerLease) + indices = make([]int, 0, len(plan.items)) + specs = make([]inboundRPCSpec, 0, len(plan.items)) + } + ownersInPlan[item.msgID] = claim.owner + plan.rpcOwners = append(plan.rpcOwners, claim.owner) + item.payload = claim.owner + indices = append(indices, i) + specs = append(specs, inboundRPCSpec{method: method, size: len(item.body)}) + default: + return ErrRPCResultFlightInvalid + } + } + if flightCapacity { + // One container is one API admission unit. If the cross-connection + // exactly-once table cannot claim every fresh request, none of this + // batch may reach a business handler. + for _, index := range indices { + plan.items[index].kind = inboundItemCapacityError + } + return nil + } + if len(specs) == 0 { + return nil + } + + reservation, err := c.reserveInboundRPCBatch(ctx, specs) + if err != nil { + if errors.Is(err, ErrInboundRPCQueueFull) { + for _, index := range indices { + plan.items[index].kind = inboundItemCapacityError + } + return nil + } + return err + } + plan.rpcReservation = reservation + plan.rpcTasks = make([]inboundRPC, len(indices)) + for i, index := range indices { + item := &plan.items[index] + body := append([]byte(nil), item.body...) + owner, _ := item.payload.(*rpcResultOwnerLease) + plan.rpcTasks[i] = s.newInboundRPCTask(c, item.msgID, specs[i].method, body, owner) + } + return nil +} + +func (s *Server) executeInboundPlan(ctx context.Context, cs *connState, c *Conn, plan *inboundPlan) error { + for _, item := range plan.items { + switch item.kind { + case inboundItemDuplicate: + // The preflight plan already stages a content ACK for a locally seen + // duplicate. Do not wait for or replay its result on the same reliable + // stream; the original owner is solely responsible for that response. + continue + case inboundItemServiceDuplicate: + // Classify from the originally committed connState record, never from the + // retransmitted body. This prevents same-id payload replacement. Only an + // already cached answer is eligible for resend (rpc_drop_answer today); + // other best-effort service traffic uses a later fresh request. + if err := s.replayRPCResultByRequest(ctx, c, item.msgID); err != nil { + return err + } + case inboundItemReplayRPC: + if encoded, _ := item.payload.(*encodedOutboundMessage); encoded != nil { + if err := s.sendCachedRPCResult(ctx, c, encoded); err != nil { + return err + } + } else if err := s.replayRPCResultByRequest(ctx, c, item.msgID); err != nil { + return err + } + case inboundItemPendingRPC: + // Wait only after this plan's fresh RPC batch has become runnable; see + // executePendingRPCReplays. Blocking here could otherwise deadlock on + // an owner appended by the same container. + continue + case inboundItemPing: + if err := s.sendPong(ctx, c, item.msgID, item.payload.(mt.PingRequest).PingID); err != nil { + return err + } + case inboundItemPingDelay: + if err := s.sendPong(ctx, c, item.msgID, item.payload.(mt.PingDelayDisconnectRequest).PingID); err != nil { + return err + } + case inboundItemFutureSalts: + if err := s.sendFutureSalts(ctx, c, item.msgID, item.payload.(mt.GetFutureSaltsRequest).Num); err != nil { + return err + } + case inboundItemMsgsAck: + ids := item.payload.(int64VectorView).materialize() + c.AckServerMessages(ids) + s.log.Debug("Received msgs_ack", zap.Int64s("msg_ids", ids)) + case inboundItemStateReq: + ids := item.payload.(int64VectorView).materialize() + outgoing, err := c.OutgoingStateInfo(ctx, ids) + if err != nil { + return err + } + if err := s.sendMsgsStateInfo(ctx, c, item.msgID, mergeStateInfo(outgoing, cs.stateInfo(ids))); err != nil { + return err + } + case inboundItemResendReq: + ids := item.payload.(int64VectorView).materialize() + outgoing, err := c.ResendMessages(ctx, ids) + if err != nil { + return err + } + if err := s.sendMsgsStateInfo(ctx, c, item.msgID, mergeStateInfo(outgoing, cs.stateInfo(ids))); err != nil { + return err + } + case inboundItemStateInfo: + value := item.payload.(stateInfoPayload) + s.log.Debug("Received msgs_state_info", zap.Int64("req_msg_id", value.reqMsgID), zap.Int("len", len(value.info))) + case inboundItemAllInfo: + value := item.payload.(allInfoPayload) + s.log.Debug("Received msgs_all_info", zap.Int("msg_ids", value.count), zap.Int("len", len(value.info))) + case inboundItemDestroySession: + if err := s.sendDestroySession(ctx, c, item.payload.(mt.DestroySessionRequest).SessionID); err != nil { + return err + } + case inboundItemHTTPWait: + value := item.payload.(mt.HTTPWaitRequest) + s.log.Debug("Received http_wait", zap.Int("max_delay", value.MaxDelay), zap.Int("wait_after", value.WaitAfter), zap.Int("max_wait", value.MaxWait)) + case inboundItemDropAnswer: + value := item.payload.(mt.RPCDropAnswerRequest) + s.log.Debug("Received rpc_drop_answer", zap.Int64("req_msg_id", value.ReqMsgID)) + if err := s.sendResult(ctx, c, item.msgID, &mt.RPCAnswerUnknown{}); err != nil { + return err + } + case inboundItemDestroyAuthKey: + s.log.Debug("Received destroy_auth_key", zap.String("auth_key_id", c.authKeyHex)) + if err := s.authKeys.Delete(ctx, c.authKeyID); err != nil { + s.log.Warn("Delete auth key failed", zap.String("auth_key_id", c.authKeyHex), zap.Error(err)) + return c.SendAsync(ctx, proto.MessageServerResponse, &destroyAuthKeyFail{}) + } + c.keyDestroyed.Store(true) + s.conns.CloseSessionsForRawAuthKeyExceptConn(c.authKeyID, c) + return c.SendAsync(ctx, proto.MessageServerResponse, &destroyAuthKeyOk{}) + case inboundItemRPC: + if plan.rpcPrepared { + continue + } + if err := s.enqueueRPC(ctx, c, item.msgID, item.typeID, &bin.Buffer{Buf: item.body}); err != nil { + return err + } + case inboundItemCapacityError: + if err := s.sendResult(ctx, c, item.msgID, &mt.RPCError{ + ErrorCode: 420, + ErrorMessage: "FLOOD_WAIT_1", + }); err != nil { + return err + } + default: + return fmt.Errorf("unknown inbound item kind %d", item.kind) + } + } + return nil +} + +// executePendingRPCReplays joins owners already running on another physical +// connection for the same MTProto session. It never dispatches business code: +// owner Put publishes the immutable result to every waiter; owner Abort leaves +// the client free to retry after the old execution has definitely stopped. +func (s *Server) executePendingRPCReplays(ctx context.Context, c *Conn, plan *inboundPlan) error { + for _, item := range plan.items { + if item.kind != inboundItemPendingRPC { + continue + } + waiter, _ := item.payload.(*rpcResultWaiter) + if waiter == nil { + continue + } + encoded, ok, err := waiter.Wait(ctx) + if err != nil { + return err + } + if !ok || encoded == nil { + // The old owner stopped without publishing a result (normally because + // replacement cancellation reached the handler before it committed). + // A fresh-connection item still owns the decoded request body, so reacquire + // only after the prior flight is definitively gone. This is sequential + // retry, never concurrent business execution. Same-Conn seen duplicates + // have no body here and rely on the client's ordinary resend path. + if len(item.body) > 0 { + if err := s.enqueueRPC(ctx, c, item.msgID, item.typeID, &bin.Buffer{Buf: item.body}); err != nil { + return err + } + } + continue + } + if err := s.sendCachedRPCResult(ctx, c, encoded); err != nil { + return err + } + } + return nil +} diff --git a/internal/mtprotoedge/inbound_preflight_bench_test.go b/internal/mtprotoedge/inbound_preflight_bench_test.go new file mode 100644 index 00000000..5cd78686 --- /dev/null +++ b/internal/mtprotoedge/inbound_preflight_bench_test.go @@ -0,0 +1,44 @@ +package mtprotoedge + +import ( + "testing" + "time" + + "github.com/gotd/td/bin" + "github.com/gotd/td/proto" + "github.com/gotd/td/tg" +) + +func BenchmarkInboundPlan32RPCContainer(b *testing.B) { + ids := proto.NewMessageIDGen(time.Now) + messages := make([]proto.Message, 32) + var request bin.Buffer + if err := (&tg.HelpGetConfigRequest{}).Encode(&request); err != nil { + b.Fatal(err) + } + requestBody := request.Copy() + for i := range messages { + messages[i] = proto.Message{ + ID: ids.New(proto.MessageFromClient), SeqNo: 1 + i*2, Bytes: len(requestBody), Body: requestBody, + } + } + outerMsgID := ids.New(proto.MessageFromClient) + var container bin.Buffer + if err := (&proto.MessageContainer{Messages: messages}).Encode(&container); err != nil { + b.Fatal(err) + } + body := container.Copy() + s := New(Options{}) + cs := newConnState() + + b.ReportAllocs() + b.SetBytes(int64(len(body))) + b.ResetTimer() + for i := 0; i < b.N; i++ { + plan, err := s.preflightInbound(cs, outerMsgID, 64, body) + if err != nil { + b.Fatal(err) + } + plan.close() + } +} diff --git a/internal/mtprotoedge/inbound_rpc.go b/internal/mtprotoedge/inbound_rpc.go index c6bafed8..2d33f8ea 100644 --- a/internal/mtprotoedge/inbound_rpc.go +++ b/internal/mtprotoedge/inbound_rpc.go @@ -30,8 +30,13 @@ type inboundRPC struct { size int run func(context.Context) error onTimeout func() - budget *inboundRPCGlobalReservation - ticket *inboundRPCTicket + // release drops request-scoped ownership which is independent from the body + // budget (for example a cross-connection in-flight RPC claim). It runs once + // after the request has either published a terminal result or become + // impossible to run; callers should still make the callback idempotent. + release func() + budget *inboundRPCGlobalReservation + ticket *inboundRPCTicket } const ( @@ -76,7 +81,7 @@ type inboundRPCScheduler struct { type inboundRPCGlobalReservation struct { scheduler *inboundRPCScheduler size int64 - once sync.Once + released atomic.Bool } // inboundRPCReservation 同时持有全局和单连接的“Copy 前”预算。commit/abort 只能成功一次； @@ -92,6 +97,32 @@ type inboundRPCReservation struct { once sync.Once } +// inboundRPCSpec 是 container preflight 与 RPC scheduler 之间的有界 admission 描述。 +// method 仅用于 metrics，size 是在 Copy 之前必须预留的 request body 字节数。 +type inboundRPCSpec struct { + method string + size int +} + +type inboundRPCBatchEntry struct { + global *inboundRPCGlobalReservation + method string + size int +} + +// inboundRPCBatchReservation 把一个 container 内的 RPC 视为一个 admission 单元。 +// reserve 一次预留整批的全局/单连接条数和字节预算；commit 一次 append +// 全部任务；abort 一次归还全部预算。这防止 container 只执行前半批。 +type inboundRPCBatchReservation struct { + conn *Conn + ctx context.Context + entries []inboundRPCBatchEntry + totalSize int64 + once sync.Once +} + +var errInboundRPCBatchTaskCount = errors.New("inbound rpc batch task count mismatch") + func newInboundRPCScheduler(workers, maxTasks int, maxBytes int64) *inboundRPCScheduler { if workers <= 0 { workers = 1 @@ -196,17 +227,89 @@ func (s *inboundRPCScheduler) reserveGlobal(size int) (*inboundRPCGlobalReservat return &inboundRPCGlobalReservation{scheduler: s, size: size64}, "", nil } +// reserveGlobalBatch 在一次 budgetMu 临界区内检查并预留整批条数/字节。 +// 返回的每个 reservation 仍由对应 task 单独归还，避免一个慢 RPC 持有 +// 整个 container 已完成任务的预算。 +func (s *inboundRPCScheduler) reserveGlobalBatch(sizes []int) ([]*inboundRPCGlobalReservation, string, error) { + reservations := make([]*inboundRPCGlobalReservation, len(sizes)) + s.budgetMu.Lock() + defer s.budgetMu.Unlock() + + select { + case <-s.stopCh: + return nil, "scheduler_closed", ErrConnClosed + default: + } + if len(sizes) > s.maxTasks-s.tasks { + return nil, "global_task_budget", ErrInboundRPCQueueFull + } + // 逐项从剩余预算中减，避免 total 和 s.bytes+total 溢出。 + remaining := s.maxBytes - s.bytes + var total int64 + for i, size := range sizes { + if size < 0 { + size = 0 + } + size64 := int64(size) + if size64 > remaining-total { + return nil, "global_byte_budget", ErrInboundRPCQueueFull + } + total += size64 + reservations[i] = &inboundRPCGlobalReservation{scheduler: s, size: size64} + } + s.tasks += len(sizes) + s.bytes += total + return reservations, "", nil +} + func (r *inboundRPCGlobalReservation) release() { if r == nil || r.scheduler == nil { return } - r.once.Do(func() { - s := r.scheduler - s.budgetMu.Lock() - s.tasks-- - s.bytes -= r.size - s.budgetMu.Unlock() - }) + if !r.released.CompareAndSwap(false, true) { + return + } + s := r.scheduler + s.budgetMu.Lock() + s.tasks-- + s.bytes -= r.size + s.budgetMu.Unlock() +} + +// releaseInboundRPCGlobalBatch 使 batch abort/commit-failure 在一次全局锁内 +// 归还它仍持有的全部预算。每张 ticket 的 CAS 保证与任何并发释放幂等。 +func releaseInboundRPCGlobalBatch(reservations []*inboundRPCGlobalReservation) { + var ( + scheduler *inboundRPCScheduler + tasks int + bytes int64 + ) + for _, reservation := range reservations { + if reservation == nil || reservation.scheduler == nil || !reservation.released.CompareAndSwap(false, true) { + continue + } + if scheduler == nil { + scheduler = reservation.scheduler + } + // 一个 batch 只能由同一 scheduler 创建。若未来出现混合调用， + // 仍通过单张 release 正确归还，而不会扣错 scheduler。 + if reservation.scheduler != scheduler { + reservation.scheduler.budgetMu.Lock() + reservation.scheduler.tasks-- + reservation.scheduler.bytes -= reservation.size + reservation.scheduler.budgetMu.Unlock() + continue + } + tasks++ + bytes += reservation.size + } + if scheduler == nil || tasks == 0 { + return + } + scheduler.budgetMu.Lock() + scheduler.tasks -= tasks + scheduler.bytes -= bytes + scheduler.budgetMu.Unlock() } func (s *inboundRPCScheduler) budgetSnapshot() (tasks int, bytes int64) { @@ -363,6 +466,10 @@ func (c *Conn) reserveInboundRPC(ctx context.Context, method string, size int) ( return nil, ctx.Err() default: } + if c.terminal.Load() { + c.metrics.InboundRPCDropped(method, "scheduler_closed") + return nil, ErrConnClosed + } if c.rpcScheduler == nil { c.metrics.InboundRPCDropped(method, "scheduler_closed") return nil, ErrConnClosed @@ -392,7 +499,7 @@ func (c *Conn) reserveInboundRPC(ctx context.Context, method string, size int) ( c.metrics.InboundRPCDropped(method, "context_done") return nil, err } - if c.rpcClosed { + if c.rpcClosed || c.terminal.Load() { c.rpcMu.Unlock() global.release() c.metrics.InboundRPCDropped(method, "scheduler_closed") @@ -427,6 +534,102 @@ func (c *Conn) reserveInboundRPC(ctx context.Context, method string, size int) ( }, nil } +// reserveInboundRPCBatch 必须在 container 内任何 request body Copy 前调用。 +// 全局预算只锁一次，单连接预算也只锁一次；任一限制不满足时 +// 整批失败，不会留下部分 task/字节 reservation。 +func (c *Conn) reserveInboundRPCBatch(ctx context.Context, specs []inboundRPCSpec) (*inboundRPCBatchReservation, error) { + if ctx == nil { + ctx = context.Background() + } + select { + case <-ctx.Done(): + c.dropInboundRPCSpecs(specs, "context_done") + return nil, ctx.Err() + default: + } + if c.terminal.Load() { + c.dropInboundRPCSpecs(specs, "scheduler_closed") + return nil, ErrConnClosed + } + if c.rpcScheduler == nil { + c.dropInboundRPCSpecs(specs, "scheduler_closed") + return nil, ErrConnClosed + } + + normalized := make([]inboundRPCSpec, len(specs)) + sizes := make([]int, len(specs)) + for i, spec := range specs { + if spec.size < 0 { + spec.size = 0 + } + normalized[i] = spec + sizes[i] = spec.size + } + globals, reason, err := c.rpcScheduler.reserveGlobalBatch(sizes) + if err != nil { + c.dropInboundRPCSpecs(normalized, reason) + return nil, err + } + + entries := make([]inboundRPCBatchEntry, len(normalized)) + var totalSize int64 + for i, spec := range normalized { + entries[i] = inboundRPCBatchEntry{ + global: globals[i], + method: spec.method, + size: spec.size, + } + totalSize += int64(spec.size) + } + + c.rpcMu.Lock() + if err := ctx.Err(); err != nil { + c.rpcMu.Unlock() + releaseInboundRPCGlobalBatch(globals) + c.dropInboundRPCSpecs(normalized, "context_done") + return nil, err + } + if c.rpcClosed || c.terminal.Load() { + c.rpcMu.Unlock() + releaseInboundRPCGlobalBatch(globals) + c.dropInboundRPCSpecs(normalized, "scheduler_closed") + return nil, ErrConnClosed + } + // 用减法比较，避免对抗性 batch 的 len 加法溢出。 + availableSlots := c.rpcQueueSize - c.rpcReserved - len(c.rpcQueue) + if len(entries) > availableSlots { + c.rpcMu.Unlock() + releaseInboundRPCGlobalBatch(globals) + c.dropInboundRPCSpecs(normalized, "queue_full") + return nil, ErrInboundRPCQueueFull + } + if totalSize > maxInflightRPCBytes-c.inflightRPCBytes.Load() { + c.rpcMu.Unlock() + releaseInboundRPCGlobalBatch(globals) + c.dropInboundRPCSpecs(normalized, "byte_budget") + return nil, ErrInboundRPCQueueFull + } + c.rpcReserved += len(entries) + c.inflightRPCBytes.Add(totalSize) + // Add 与 close 的 Wait 由 rpcMu 排序：close 置 rpcClosed 后不会再发生 Add。 + // 整批 reservation 只需一个 waiter；commit/abort 也只会完成一次。 + c.rpcReservationWG.Add(1) + c.rpcMu.Unlock() + + return &inboundRPCBatchReservation{ + conn: c, + ctx: ctx, + entries: entries, + totalSize: totalSize, + }, nil +} + +func (c *Conn) dropInboundRPCSpecs(specs []inboundRPCSpec, reason string) { + for _, spec := range specs { + c.metrics.InboundRPCDropped(spec.method, reason) + } +} + // enqueueInboundRPC 是测试和已持有独立 body 的便捷入口。生产收包路径使用 // reserveInboundRPC -> Copy -> commit，保证真正的 Copy 前预算。 func (c *Conn) enqueueInboundRPC(ctx context.Context, task inboundRPC) error { @@ -450,7 +653,7 @@ func (r *inboundRPCReservation) commit(task inboundRPC) error { c := r.conn c.rpcMu.Lock() c.rpcReserved-- - if c.rpcClosed { + if c.rpcClosed || c.terminal.Load() { c.inflightRPCBytes.Add(-int64(r.size)) } else { // The request deadline starts when admission succeeds, not when a worker @@ -525,6 +728,143 @@ func (r *inboundRPCReservation) abort() { }) } +// commit 在一次 rpcMu 临界区内把整批 task append 到队列。 +// deferSchedule=false 保持旧的立即调度语义；true 则返回一个幂等 activate +// 函数，让调用方先完成 new_session_created 等协议 barrier 再启动 worker。 +// +// 延迟调度只能延迟本次 commit 新产生的 ready token；调用方应在连接的 +// 首个 admission batch 使用它，不得把它当作已有 worker 的全局暂停锁。 +func (r *inboundRPCBatchReservation) commit(tasks []inboundRPC, deferSchedule bool) (activate func(), result error) { + if r == nil { + return nil, ErrConnClosed + } + result = ErrConnClosed + var ( + committed bool + reschedule bool + firstQueueLen int + queueCap int + globals []*inboundRPCGlobalReservation + ) + r.once.Do(func() { + c := r.conn + // A container reservation may deliberately span protocol-critical barriers + // (session ownership claim and new_session_created's physical write). Queue + // and execution latency starts only when the fully admitted batch becomes + // runnable, not while it is waiting behind those independent barriers. + enqueuedAt := time.Now() + deadline := time.Time{} + if c.rpcTimeout > 0 { + deadline = enqueuedAt.Add(c.rpcTimeout) + } + if ctxDeadline, ok := r.ctx.Deadline(); ok && (deadline.IsZero() || ctxDeadline.Before(deadline)) { + deadline = ctxDeadline + } + globals = make([]*inboundRPCGlobalReservation, len(r.entries)) + for i := range r.entries { + globals[i] = r.entries[i].global + } + c.rpcMu.Lock() + c.rpcReserved -= len(r.entries) + if len(tasks) != len(r.entries) { + c.inflightRPCBytes.Add(-r.totalSize) + result = errInboundRPCBatchTaskCount + } else if c.rpcClosed || c.terminal.Load() { + c.inflightRPCBytes.Add(-r.totalSize) + } else { + prepared := make([]inboundRPC, len(tasks)) + // The request deadline starts when admission succeeds, not when a worker + // eventually dequeues the request. This bounds total queue + execution + // latency and lets a queued request emit its explicit timeout on time. + for i := range tasks { + task := tasks[i] + entry := r.entries[i] + if deadline.IsZero() { + task.ctx, task.cancel = context.WithCancel(r.ctx) + } else { + task.ctx, task.cancel = context.WithDeadline(r.ctx, deadline) + } + task.stopRoot = context.AfterFunc(c.rpcRootCtx, task.cancel) + task.method = entry.method + task.enqueuedAt = enqueuedAt + task.deadline = deadline + task.size = entry.size + task.budget = entry.global + ticket := &inboundRPCTicket{} + if task.onTimeout != nil { + onTimeout := task.onTimeout + var timeoutOnce sync.Once + ticket.onTimeout = func() { + timeoutOnce.Do(onTimeout) + } + task.onTimeout = ticket.onTimeout + } + task.ticket = ticket + if task.onTimeout != nil && !task.deadline.IsZero() { + taskCtx := task.ctx + task.stopTimeout = context.AfterFunc(taskCtx, func() { + if errors.Is(taskCtx.Err(), context.DeadlineExceeded) { + c.expireInboundRPCTicket(ticket) + } + }) + } + prepared[i] = task + } + firstQueueLen = len(c.rpcQueue) + 1 + c.rpcQueue = append(c.rpcQueue, prepared...) + queueCap = c.rpcQueueSize + if len(prepared) > 0 && c.rpcRunning < c.rpcMaxInflight && !c.rpcReady { + c.rpcReady = true + reschedule = true + } + committed = true + result = nil + } + c.rpcMu.Unlock() + c.rpcReservationWG.Done() + if !committed { + releaseInboundRPCGlobalBatch(globals) + } + }) + if committed { + for i, entry := range r.entries { + r.conn.metrics.InboundRPCQueued(entry.method, firstQueueLen+i, queueCap) + } + if reschedule { + var once sync.Once + activate = func() { + once.Do(func() { + r.conn.rpcScheduler.schedule(r.conn) + }) + } + if !deferSchedule { + activate() + activate = nil + } + } + } + return activate, result +} + +func (r *inboundRPCBatchReservation) abort() { + if r == nil { + return + } + r.once.Do(func() { + c := r.conn + c.rpcMu.Lock() + c.rpcReserved -= len(r.entries) + c.inflightRPCBytes.Add(-r.totalSize) + c.rpcMu.Unlock() + c.rpcReservationWG.Done() + globals := make([]*inboundRPCGlobalReservation, len(r.entries)) + for i := range r.entries { + globals[i] = r.entries[i].global + } + releaseInboundRPCGlobalBatch(globals) + }) +} + func (c *Conn) takeInboundRPC() (task inboundRPC, ok, reschedule bool) { c.rpcMu.Lock() defer c.rpcMu.Unlock() @@ -583,7 +923,7 @@ func (c *Conn) finishInboundRPC(task inboundRPC) { if task.ticket != nil { task.ticket.state.Store(inboundRPCTicketDone) } - stopInboundRPCTask(task) + timeoutHandoff := stopInboundRPCTask(task) var reschedule bool c.rpcMu.Lock() c.rpcRunning-- @@ -594,11 +934,21 @@ func (c *Conn) finishInboundRPC(task inboundRPC) { } c.rpcMu.Unlock() reservation := task.budget + release := task.release // The scheduler budget may be reused immediately after release. Clear request-owned // closures/context references first so slow metrics/rescheduling cannot overlap the old body // with a newly admitted body under the same byte accounting. task = inboundRPC{} reservation.release() + if timeoutHandoff != nil { + // stopTimeout(false) means the deadline callback may already have read + // Running but not yet entered ticket.onTimeout. Calling the sync.Once wrapper + // here either performs or joins that response before owner release/Abort. + timeoutHandoff() + } + if release != nil { + release() + } c.rpcWG.Done() if reschedule { c.rpcScheduler.schedule(c) @@ -648,7 +998,8 @@ func (c *Conn) expireInboundRPCTicket(ticket *inboundRPCTicket) { if found { method := task.method reservation := task.budget - stopInboundRPCTask(task) + release := task.release + _ = stopInboundRPCTask(task) // Drop the run/context closures before returning the byte reservation. Otherwise an // onTimeout callback that blocks or performs a slow write can keep the copied request body // reachable after the global scheduler has advertised those bytes as available again. @@ -658,6 +1009,9 @@ func (c *Conn) expireInboundRPCTicket(ticket *inboundRPCTicket) { if ticket.onTimeout != nil { ticket.onTimeout() } + if release != nil { + release() + } return } if ticket.state.Load() == inboundRPCTicketRunning && ticket.onTimeout != nil { @@ -667,11 +1021,17 @@ func (c *Conn) expireInboundRPCTicket(ticket *inboundRPCTicket) { // stopInboundRPCTask disarms callbacks before canceling the context so a normal // completion or connection close cannot manufacture an RPC_TIMEOUT response. -// A deadline callback already in flight is harmless because enqueueRPC's response -// gate makes timeout and normal rpc_result mutually exclusive. -func stopInboundRPCTask(task inboundRPC) { +// If the runtime already started a deadline callback, the returned sync.Once +// wrapper is a mandatory handoff: callers invoke it before owner release so a +// callback paused between ticket-state inspection and response-gate claim cannot +// publish into a later flight generation. +func stopInboundRPCTask(task inboundRPC) (timeoutHandoff func()) { if task.stopTimeout != nil { - task.stopTimeout() + stopped := task.stopTimeout() + if !stopped && task.ctx != nil && errors.Is(task.ctx.Err(), context.DeadlineExceeded) && + task.ticket != nil && task.ticket.onTimeout != nil { + timeoutHandoff = task.ticket.onTimeout + } } if task.stopRoot != nil { task.stopRoot() @@ -679,6 +1039,7 @@ func stopInboundRPCTask(task inboundRPC) { if task.cancel != nil { task.cancel() } + return timeoutHandoff } func (c *Conn) closeInboundRPCScheduler() { @@ -722,9 +1083,16 @@ func (c *Conn) beginCloseInboundRPCScheduler() { } method := task.method reservation := task.budget - stopInboundRPCTask(task) + release := task.release + timeoutHandoff := stopInboundRPCTask(task) task = inboundRPC{} reservation.release() + if timeoutHandoff != nil { + timeoutHandoff() + } + if release != nil { + release() + } c.metrics.InboundRPCDropped(method, "connection_closed") } }) diff --git a/internal/mtprotoedge/inbound_rpc_batch_test.go b/internal/mtprotoedge/inbound_rpc_batch_test.go new file mode 100644 index 00000000..52553152 --- /dev/null +++ b/internal/mtprotoedge/inbound_rpc_batch_test.go @@ -0,0 +1,343 @@ +package mtprotoedge + +import ( + "context" + "errors" + "testing" + "time" +) + +func TestInboundRPCBatchReservationRejectsGloballyWithoutPartialBudget(t *testing.T) { + scheduler := newInboundRPCScheduler(1, 2, 10) + c := newInboundTestConn(scheduler, 1, 8, time.Second) + defer func() { + c.closeInboundRPCScheduler() + scheduler.stop(time.Second) + }() + + _, err := c.reserveInboundRPCBatch(context.Background(), []inboundRPCSpec{ + {method: "one", size: 3}, + {method: "two", size: 3}, + {method: "three", size: 3}, + }) + if !errors.Is(err, ErrInboundRPCQueueFull) { + t.Fatalf("reserve over global task budget err = %v, want ErrInboundRPCQueueFull", err) + } + if tasks, bytes := scheduler.budgetSnapshot(); tasks != 0 || bytes != 0 { + t.Fatalf("global budget after atomic rejection = (%d, %d), want zero", tasks, bytes) + } + if got := c.inflightRPCBytes.Load(); got != 0 { + t.Fatalf("connection bytes after global rejection = %d, want zero", got) + } + c.rpcMu.Lock() + reserved := c.rpcReserved + queued := len(c.rpcQueue) + c.rpcMu.Unlock() + if reserved != 0 || queued != 0 { + t.Fatalf("connection state after global rejection = reserved %d queued %d, want zero", reserved, queued) + } +} + +func TestInboundRPCBatchReservationRejectsConnectionWithoutLeakingGlobalBudget(t *testing.T) { + scheduler := newInboundRPCScheduler(1, 16, 1<<20) + c := newInboundTestConn(scheduler, 1, 2, time.Second) + defer func() { + c.closeInboundRPCScheduler() + scheduler.stop(time.Second) + }() + + _, err := c.reserveInboundRPCBatch(context.Background(), []inboundRPCSpec{ + {method: "one", size: 3}, + {method: "two", size: 5}, + {method: "three", size: 7}, + }) + if !errors.Is(err, ErrInboundRPCQueueFull) { + t.Fatalf("reserve over connection queue budget err = %v, want ErrInboundRPCQueueFull", err) + } + if tasks, bytes := scheduler.budgetSnapshot(); tasks != 0 || bytes != 0 { + t.Fatalf("global budget after connection rejection = (%d, %d), want zero", tasks, bytes) + } + if got := c.inflightRPCBytes.Load(); got != 0 { + t.Fatalf("connection bytes after connection rejection = %d, want zero", got) + } + c.rpcMu.Lock() + reserved := c.rpcReserved + queued := len(c.rpcQueue) + c.rpcMu.Unlock() + if reserved != 0 || queued != 0 { + t.Fatalf("connection state after connection rejection = reserved %d queued %d, want zero", reserved, queued) + } +} + +func TestInboundRPCBatchReservationRejectsAggregateConnectionBytes(t *testing.T) { + scheduler := newInboundRPCScheduler(1, 8, int64(maxInflightRPCBytes)*2) + c := newInboundTestConn(scheduler, 1, 4, time.Second) + defer func() { + c.closeInboundRPCScheduler() + scheduler.stop(time.Second) + }() + + halfPlusOne := maxInflightRPCBytes/2 + 1 + _, err := c.reserveInboundRPCBatch(context.Background(), []inboundRPCSpec{ + {method: "one", size: halfPlusOne}, + {method: "two", size: halfPlusOne}, + }) + if !errors.Is(err, ErrInboundRPCQueueFull) { + t.Fatalf("reserve over aggregate connection byte budget err = %v, want ErrInboundRPCQueueFull", err) + } + if tasks, bytes := scheduler.budgetSnapshot(); tasks != 0 || bytes != 0 { + t.Fatalf("global budget after aggregate byte rejection = (%d, %d), want zero", tasks, bytes) + } + if got := c.inflightRPCBytes.Load(); got != 0 { + t.Fatalf("connection bytes after aggregate byte rejection = %d, want zero", got) + } +} + +func TestInboundRPCBatchAbortReturnsEveryReservationExactlyOnce(t *testing.T) { + scheduler := newInboundRPCScheduler(1, 8, 1<<20) + c := newInboundTestConn(scheduler, 1, 4, time.Second) + defer func() { + c.closeInboundRPCScheduler() + scheduler.stop(time.Second) + }() + + reservation, err := c.reserveInboundRPCBatch(context.Background(), []inboundRPCSpec{ + {method: "one", size: 3}, + {method: "two", size: 5}, + {method: "three", size: 7}, + }) + if err != nil { + t.Fatalf("reserve batch: %v", err) + } + if tasks, bytes := scheduler.budgetSnapshot(); tasks != 3 || bytes != 15 { + t.Fatalf("budget after reserve = (%d, %d), want (3, 15)", tasks, bytes) + } + reservation.abort() + reservation.abort() + + if tasks, bytes := scheduler.budgetSnapshot(); tasks != 0 || bytes != 0 { + t.Fatalf("budget after idempotent abort = (%d, %d), want zero", tasks, bytes) + } + if got := c.inflightRPCBytes.Load(); got != 0 { + t.Fatalf("connection bytes after abort = %d, want zero", got) + } + c.rpcMu.Lock() + reserved := c.rpcReserved + queued := len(c.rpcQueue) + c.rpcMu.Unlock() + if reserved != 0 || queued != 0 { + t.Fatalf("connection state after abort = reserved %d queued %d, want zero", reserved, queued) + } +} + +func TestInboundRPCBatchCommitAppendsAllAndDefersSchedule(t *testing.T) { + scheduler := newInboundRPCScheduler(1, 8, 1<<20) + scheduler.start() + c := newInboundTestConn(scheduler, 1, 4, time.Second) + defer func() { + c.closeInboundRPCScheduler() + scheduler.stop(time.Second) + }() + + specs := []inboundRPCSpec{ + {method: "one", size: 3}, + {method: "two", size: 5}, + {method: "three", size: 7}, + } + reservation, err := c.reserveInboundRPCBatch(context.Background(), specs) + if err != nil { + t.Fatalf("reserve batch: %v", err) + } + defer reservation.abort() + + runs := make(chan string, len(specs)) + tasks := make([]inboundRPC, len(specs)) + for i, spec := range specs { + method := spec.method + tasks[i].run = func(context.Context) error { + runs <- method + return nil + } + } + activate, err := reservation.commit(tasks, true) + if err != nil { + t.Fatalf("commit batch: %v", err) + } + if activate == nil { + t.Fatal("deferred commit did not return an activation function") + } + c.rpcMu.Lock() + queued := len(c.rpcQueue) + ready := c.rpcReady + c.rpcMu.Unlock() + if queued != len(specs) || !ready { + t.Fatalf("atomic queue state after commit = queued %d ready %v, want %d/true", queued, ready, len(specs)) + } + if got := scheduler.readyLen(); got != 0 { + t.Fatalf("scheduler ready tokens before activation = %d, want zero", got) + } + select { + case method := <-runs: + t.Fatalf("RPC %q ran before deferred activation", method) + default: + } + + activate() + activate() // activation is idempotent. + for _, want := range []string{"one", "two", "three"} { + select { + case got := <-runs: + if got != want { + t.Fatalf("execution order = %q, want %q", got, want) + } + case <-time.After(time.Second): + t.Fatalf("timed out waiting for %q", want) + } + } + waitInboundRPCBatchBudget(t, scheduler, 0, 0) +} + +func TestInboundRPCBatchCommitMismatchReleasesAllWithoutEnqueue(t *testing.T) { + scheduler := newInboundRPCScheduler(1, 8, 1<<20) + c := newInboundTestConn(scheduler, 1, 4, time.Second) + defer func() { + c.closeInboundRPCScheduler() + scheduler.stop(time.Second) + }() + + reservation, err := c.reserveInboundRPCBatch(context.Background(), []inboundRPCSpec{ + {method: "one", size: 3}, + {method: "two", size: 5}, + }) + if err != nil { + t.Fatalf("reserve batch: %v", err) + } + if _, err := reservation.commit([]inboundRPC{{}}, false); !errors.Is(err, errInboundRPCBatchTaskCount) { + t.Fatalf("commit task mismatch err = %v, want %v", err, errInboundRPCBatchTaskCount) + } + if tasks, bytes := scheduler.budgetSnapshot(); tasks != 0 || bytes != 0 { + t.Fatalf("budget after mismatched commit = (%d, %d), want zero", tasks, bytes) + } + if got := c.inflightRPCBytes.Load(); got != 0 { + t.Fatalf("connection bytes after mismatched commit = %d, want zero", got) + } + c.rpcMu.Lock() + reserved := c.rpcReserved + queued := len(c.rpcQueue) + c.rpcMu.Unlock() + if reserved != 0 || queued != 0 { + t.Fatalf("connection state after mismatched commit = reserved %d queued %d, want zero", reserved, queued) + } +} + +func TestInboundRPCBatchCommitRacingCloseNeverPartiallyEnqueues(t *testing.T) { + scheduler := newInboundRPCScheduler(1, 8, 1<<20) + c := newInboundTestConn(scheduler, 1, 4, time.Second) + defer scheduler.stop(time.Second) + + reservation, err := c.reserveInboundRPCBatch(context.Background(), []inboundRPCSpec{ + {method: "one", size: 3}, + {method: "two", size: 5}, + {method: "three", size: 7}, + }) + if err != nil { + t.Fatalf("reserve batch: %v", err) + } + closed := make(chan struct{}) + go func() { + c.closeInboundRPCScheduler() + close(closed) + }() + waitInboundRPCBatchConnClosed(t, c) + + if _, err := reservation.commit(make([]inboundRPC, 3), false); !errors.Is(err, ErrConnClosed) { + t.Fatalf("commit after close err = %v, want ErrConnClosed", err) + } + select { + case <-closed: + case <-time.After(time.Second): + t.Fatal("close did not finish after batch commit returned its reservation") + } + if tasks, bytes := scheduler.budgetSnapshot(); tasks != 0 || bytes != 0 { + t.Fatalf("budget after close/commit race = (%d, %d), want zero", tasks, bytes) + } + if got := c.inflightRPCBytes.Load(); got != 0 { + t.Fatalf("connection bytes after close/commit race = %d, want zero", got) + } + c.rpcMu.Lock() + queued := len(c.rpcQueue) + c.rpcMu.Unlock() + if queued != 0 { + t.Fatalf("queue after close/commit race = %d, want zero", queued) + } +} + +func TestInboundRPCBatchCommitAfterTerminalFenceRejectsAll(t *testing.T) { + scheduler := newInboundRPCScheduler(1, 8, 1<<20) + c := newInboundTestConn(scheduler, 1, 4, time.Second) + defer func() { + c.closeInboundRPCScheduler() + scheduler.stop(time.Second) + }() + + reservation, err := c.reserveInboundRPCBatch(context.Background(), []inboundRPCSpec{ + {method: "one", size: 3}, + {method: "two", size: 5}, + }) + if err != nil { + t.Fatalf("reserve batch: %v", err) + } + // Session replacement and revocation publish terminal before the slower + // physical-close path. A reservation held across that fence must not be able + // to append even one stale task. + c.terminal.Store(true) + if _, err := reservation.commit(make([]inboundRPC, 2), false); !errors.Is(err, ErrConnClosed) { + t.Fatalf("commit after terminal fence err = %v, want ErrConnClosed", err) + } + + if tasks, bytes := scheduler.budgetSnapshot(); tasks != 0 || bytes != 0 { + t.Fatalf("budget after terminal commit rejection = (%d, %d), want zero", tasks, bytes) + } + if got := c.inflightRPCBytes.Load(); got != 0 { + t.Fatalf("connection bytes after terminal commit rejection = %d, want zero", got) + } + c.rpcMu.Lock() + reserved := c.rpcReserved + queued := len(c.rpcQueue) + c.rpcMu.Unlock() + if reserved != 0 || queued != 0 { + t.Fatalf("connection state after terminal commit rejection = reserved %d queued %d, want zero", reserved, queued) + } +} + +func waitInboundRPCBatchBudget(t *testing.T, scheduler *inboundRPCScheduler, wantTasks int, wantBytes int64) { + t.Helper() + deadline := time.Now().Add(time.Second) + for { + tasks, bytes := scheduler.budgetSnapshot() + if tasks == wantTasks && bytes == wantBytes { + return + } + if time.Now().After(deadline) { + t.Fatalf("budget = (%d, %d), want (%d, %d)", tasks, bytes, wantTasks, wantBytes) + } + time.Sleep(time.Millisecond) + } +} + +func waitInboundRPCBatchConnClosed(t *testing.T, c *Conn) { + t.Helper() + deadline := time.Now().Add(time.Second) + for { + c.rpcMu.Lock() + closed := c.rpcClosed + c.rpcMu.Unlock() + if closed { + return + } + if time.Now().After(deadline) { + t.Fatal("connection scheduler was not marked closed") + } + time.Sleep(time.Millisecond) + } +} diff --git a/internal/mtprotoedge/inbound_rpc_test.go b/internal/mtprotoedge/inbound_rpc_test.go index d26a2922..1f16fe20 100644 --- a/internal/mtprotoedge/inbound_rpc_test.go +++ b/internal/mtprotoedge/inbound_rpc_test.go @@ -3,11 +3,58 @@ package mtprotoedge import ( "context" "errors" + "sync" "sync/atomic" "testing" "time" ) +func TestStopInboundRPCTaskJoinsStartedDeadlineCallback(t *testing.T) { + ctx, cancel := context.WithDeadline(context.Background(), time.Now().Add(-time.Second)) + defer cancel() + started := make(chan struct{}) + release := make(chan struct{}) + var once sync.Once + ticket := &inboundRPCTicket{} + ticket.onTimeout = func() { + once.Do(func() { + close(started) + <-release + }) + } + task := inboundRPC{ + ctx: ctx, + ticket: ticket, + stopTimeout: func() bool { return false }, // runtime callback already scheduled + } + handoff := stopInboundRPCTask(task) + if handoff == nil { + t.Fatal("started deadline callback did not produce handoff") + } + go ticket.onTimeout() + select { + case <-started: + case <-time.After(time.Second): + t.Fatal("deadline callback did not start") + } + joined := make(chan struct{}) + go func() { + handoff() + close(joined) + }() + select { + case <-joined: + t.Fatal("handoff returned before in-flight callback completed") + case <-time.After(20 * time.Millisecond): + } + close(release) + select { + case <-joined: + case <-time.After(time.Second): + t.Fatal("handoff did not join completed callback") + } +} + func newInboundTestConn(s *inboundRPCScheduler, maxInflight, queueSize int, timeout time.Duration) *Conn { c := &Conn{metrics: NopMetrics{}} c.startInboundRPCScheduler(s, maxInflight, queueSize, timeout) diff --git a/internal/mtprotoedge/metrics_test.go b/internal/mtprotoedge/metrics_test.go index 2e9e44ed..ef0d1b91 100644 --- a/internal/mtprotoedge/metrics_test.go +++ b/internal/mtprotoedge/metrics_test.go @@ -61,8 +61,8 @@ func TestMetricsHooks(t *testing.T) { if got := m.inbound.Load(); got != 1 { t.Errorf("InboundRPCStarted called %d times, want 1", got) } - // new_session_created / ack 走 fire-and-forget（异步），可能在 client 收到 rpc_result 后 - // 才被 outbound actor 处理；轮询等其最终发送完成。M5 验证发送计数，不约束同步时序。 + // new_session_created 已经跨过 required-control 物理写屏障；msgs_ack 仍是异步控制帧， + // 因而这里只轮询等待最终发送计数。具体 boundary→RPC 顺序由 session boundary 测试锁定。 deadline := time.Now().Add(2 * time.Second) for m.outbound.Load() < 3 && time.Now().Before(deadline) { time.Sleep(5 * time.Millisecond) diff --git a/internal/mtprotoedge/new_session_uid_test.go b/internal/mtprotoedge/new_session_uid_test.go index ca742ecb..a379a6d8 100644 --- a/internal/mtprotoedge/new_session_uid_test.go +++ b/internal/mtprotoedge/new_session_uid_test.go @@ -43,3 +43,50 @@ func TestNewSessionCreatedUniqueIDPerSession(t *testing.T) { t.Fatalf("new_session_created.unique_id reused across sessions: %d", created1.UniqueID) } } + +// TestNewSessionCreatedMovesFloorForLaterLowerMessage locks the MTProto rule +// that a server-side session must publish a fresh boundary notification when it +// later accepts a smaller logical client msg_id. Official clients use this +// boundary to decide which requests must be regenerated and resent. +func TestNewSessionCreatedMovesFloorForLaterLowerMessage(t *testing.T) { + const dc = 2 + addr, pub, _ := startTestServer(t, Options{DC: dc}) + conn, auth, cipher := dialHandshake(t, addr, dc, pub) + + freshIDs := proto.NewMessageIDGen(time.Now) + firstMsgID := freshIDs.New(proto.MessageFromClient) + sendEncryptedWithSeq(t, conn, cipher, auth, firstMsgID, 0, &mt.HTTPWaitRequest{ + MaxDelay: 0, + WaitAfter: 0, + MaxWait: 25_000, + }) + first := collectReplies(t, conn, cipher, auth.AuthKey, mt.NewSessionCreatedTypeID) + var created1 mt.NewSessionCreated + if err := created1.Decode(mustHave(t, first, mt.NewSessionCreatedTypeID, "first new_session_created")); err != nil { + t.Fatalf("decode first new_session_created: %v", err) + } + if created1.FirstMsgID != firstMsgID { + t.Fatalf("first boundary = %d, want %d", created1.FirstMsgID, firstMsgID) + } + + oldIDs := proto.NewMessageIDGen(func() time.Time { return time.Now().Add(-10 * time.Minute) }) + lowerMsgID := oldIDs.New(proto.MessageFromClient) + outerMsgID := freshIDs.New(proto.MessageFromClient) + pingBody := mustEncodeTL(t, &mt.PingRequest{PingID: 99}) + sendEncrypted(t, conn, cipher, auth, outerMsgID, &proto.MessageContainer{ + Messages: []proto.Message{{ + ID: lowerMsgID, SeqNo: 1, Bytes: len(pingBody), Body: pingBody, + }}, + }) + second := collectReplies(t, conn, cipher, auth.AuthKey, mt.PongTypeID) + var created2 mt.NewSessionCreated + if err := created2.Decode(mustHave(t, second, mt.NewSessionCreatedTypeID, "lower-boundary new_session_created")); err != nil { + t.Fatalf("decode lower-boundary new_session_created: %v", err) + } + if created2.FirstMsgID != lowerMsgID { + t.Fatalf("lower boundary = %d, want %d", created2.FirstMsgID, lowerMsgID) + } + if created2.UniqueID == created1.UniqueID { + t.Fatalf("lower-boundary notification reused unique_id %d", created2.UniqueID) + } +} diff --git a/internal/mtprotoedge/outbound.go b/internal/mtprotoedge/outbound.go index cc737833..4f8b2412 100644 --- a/internal/mtprotoedge/outbound.go +++ b/internal/mtprotoedge/outbound.go @@ -40,6 +40,11 @@ const ( defaultOutboundControlQueueSize = 32 defaultOutboundTrackedMaxBytes = int64(512 << 20) // 512 MiB / Server defaultOutboundControlMaxBytes = int64(64 << 20) // ack/state/resend vectors / Server + // requiredControlMaxWait bounds protocol barriers such as new_session_created from + // the beginning of encoding through the completed physical write. These frames gate + // subsequent session state transitions, so timing out must close the connection instead + // of degrading to the best-effort control path. + requiredControlMaxWait = 5 * time.Second maxTrackedServerMsgIDs = 4096 maxTrackedAckedMsgIDs = 1024 @@ -332,6 +337,9 @@ func (c *Conn) Close() { // transport.Close call, so a timed-out batch close cannot keep accepting memory/work. func (c *Conn) beginTerminalShutdown() { c.terminal.Store(true) + // Retirement is irreversible. SessionManager activation only uses CAS from + // provisional/claiming, so a stale goroutine cannot publish this Conn again. + c.lifecycle.Store(uint32(connLifecycleRetired)) c.signalOutboundStop() c.beginCloseInboundRPCScheduler() } @@ -342,6 +350,23 @@ func (c *Conn) waitOutboundShutdown() { } } +func (c *Conn) waitOutboundShutdownUntil(timeout time.Duration) bool { + if c == nil || c.outboundDone == nil { + return true + } + if timeout <= 0 { + return false + } + timer := time.NewTimer(timeout) + defer timer.Stop() + select { + case <-c.outboundDone: + return true + case <-timer.C: + return false + } +} + // ForceClose 停止连接并关闭底层 transport。 // 仅用于授权撤销 / destroy_auth_key 这类“必须让对端立即断线”的路径；普通生命周期仍由 // serveConn 统一关闭 transport，避免正常 push/索引清理把长连接误伤成硬断。 @@ -355,6 +380,13 @@ func (c *Conn) ForceClose() { // closeTransport 只关闭物理 transport，不等待 outbound actor。写失败路径运行在 // actor 自身 goroutine 中，若在这里调用 Close 会等待 outboundDone 而自锁。 func (c *Conn) closeTransport() { + if c == nil { + return + } + if c.transportLease != nil { + _ = c.transportLease.Close() + return + } c.transportClose.Do(func() { if c.transport != nil { _ = c.transport.Close() @@ -372,6 +404,38 @@ func (c *Conn) failTransport() { c.closeTransport() } +// fenceUndeliveredRPCResult is the no-reentry terminal path used from a task's +// release callback. That callback may itself run while rpcClose.Do is draining +// queued tasks, so calling beginCloseInboundRPCScheduler again would deadlock on +// sync.Once. Closing the socket wakes serveConn, whose ordinary defer completes +// scheduler/index cleanup; when shutdown already owns the callback, that cleanup +// is already in progress. +func (c *Conn) fenceUndeliveredRPCResult() { + if c == nil { + return + } + // A replacement/shutdown that already published terminal owns physical + // lifecycle cleanup (and may intentionally transfer the lease). Only the + // resultless task that wins false->true is allowed to close this generation. + if !c.terminal.CompareAndSwap(false, true) { + return + } + c.lifecycle.Store(uint32(connLifecycleRetired)) + c.signalOutboundStop() + if c.transportLease != nil { + c.transportLease.startCloseAlreadyFenced() + return + } + // Legacy construction-only Conns have no owner callback graph, so their + // exact transport close cannot re-enter logical lifecycle cleanup. Keep a + // pathological Close outside the shared RPC worker just like the lease path. + go c.transportClose.Do(func() { + if c.transport != nil { + _ = c.transport.Close() + } + }) +} + // dropSlowConsumer 把出站队列持续拥塞的连接降级为离线连接。它不能等待 outbound // actor：调用方位于 fan-out 热路径，等待单个慢 socket 会把同一用户的健康设备和 // transactional outbox lane 一起拖住。关闭 transport 会打断可能阻塞的写；serveConn @@ -402,6 +466,41 @@ func (c *Conn) SendPriority(ctx context.Context, t proto.MessageType, msg bin.En return c.send(ctx, t, msg, true) } +// SendRequiredControl writes a protocol-critical control message before the caller commits +// the state transition guarded by that message. One absolute deadline covers encode admission, +// body-budget reservation, control-queue admission and the physical transport write. A failure +// is terminal: continuing on the same connection could expose state whose required notification +// never reached the client. +// +// Success only confirms the physical write; it does not wait for the client's msgs_ack. +func (c *Conn) SendRequiredControl(ctx context.Context, t proto.MessageType, msg bin.Encoder) error { + if ctx == nil { + ctx = context.Background() + } + now := time.Now() + deadline := now.Add(requiredControlMaxWait) + if c.writeTimeout > 0 { + if writeDeadline := now.Add(c.writeTimeout); writeDeadline.Before(deadline) { + deadline = writeDeadline + } + } + if parentDeadline, ok := ctx.Deadline(); ok && parentDeadline.Before(deadline) { + deadline = parentDeadline + } + requiredCtx, cancel := context.WithDeadline(ctx, deadline) + defer cancel() + if err := requiredCtx.Err(); err != nil { + c.failTransport() + return err + } + + err := c.sendOutbound(requiredCtx, t, msg, nil, true) + if err != nil { + c.failTransport() + } + return err +} + // 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 { @@ -519,14 +618,27 @@ func (c *Conn) sendOutbound(ctx context.Context, t proto.MessageType, msg bin.En case res := <-op.done: return res.err case <-ctx.Done(): + // A physical write can complete at the same instant as the caller's + // deadline. Prefer the actor's terminal result when it is already + // available so required-control callers do not poison a healthy Conn. + select { + case res := <-op.done: + return res.err + default: + } return ctx.Err() case <-c.outboundStop: + select { + case res := <-op.done: + return res.err + default: + } return ErrConnClosed } } // SendAsync 入队一条 server 消息但不等待发送结果（fire-and-forget），用于读循环里的控制消息 -// （ack/pong/new_session_created/bad_msg/future_salts/state_info）：避免读循环被 outbound 写 +// （ack/pong/bad_msg/future_salts/state_info）：避免读循环被 outbound 写 // 阻塞而连带卡死。走优先(control)队列保证不被普通 push 拖后；队列满时丢弃并记 metrics——此时 // 连接多已严重拥塞，控制消息丢失由客户端重传 / 读写超时兜底。返回非 nil 仅表示连接已关闭。 func (c *Conn) SendAsync(ctx context.Context, t proto.MessageType, msg bin.Encoder) error { @@ -1038,7 +1150,10 @@ func (c *Conn) failOutboundBudget(err error) { if c.metrics != nil { c.metrics.OutboundDropped("tracked_global_byte_budget") } - c.failTransport() + // No socket bytes exist yet. If an intentional session handoff already won + // the terminal CAS, it owns close/transfer and this old producer must not close + // the still-current lease. A live connection still gets fenced and closed. + c.fenceUndeliveredRPCResult() } func (c *Conn) ensureOutboundTrackedBudget() *outboundTrackedBudget { diff --git a/internal/mtprotoedge/outbound_required_control_test.go b/internal/mtprotoedge/outbound_required_control_test.go new file mode 100644 index 00000000..5d4fca21 --- /dev/null +++ b/internal/mtprotoedge/outbound_required_control_test.go @@ -0,0 +1,262 @@ +package mtprotoedge + +import ( + "context" + "crypto/rand" + "errors" + "io" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/gotd/td/bin" + "github.com/gotd/td/crypto" + "github.com/gotd/td/mt" + "github.com/gotd/td/proto" +) + +type gatedRequiredControlTransport struct { + started chan struct{} + release chan struct{} + sendErr error + + startOnce sync.Once + closeOnce sync.Once + sends atomic.Int32 + closes atomic.Int32 +} + +func newGatedRequiredControlTransport(sendErr error) *gatedRequiredControlTransport { + return &gatedRequiredControlTransport{ + started: make(chan struct{}), + release: make(chan struct{}), + sendErr: sendErr, + } +} + +func (t *gatedRequiredControlTransport) Send(context.Context, *bin.Buffer) error { + t.sends.Add(1) + t.startOnce.Do(func() { close(t.started) }) + <-t.release + return t.sendErr +} + +func (t *gatedRequiredControlTransport) Recv(context.Context, *bin.Buffer) error { + return io.EOF +} + +func (t *gatedRequiredControlTransport) Close() error { + t.closes.Add(1) + t.closeOnce.Do(func() { close(t.release) }) + return nil +} + +func (t *gatedRequiredControlTransport) unblock() { + t.closeOnce.Do(func() { close(t.release) }) +} + +func TestSendRequiredControlWaitsForPhysicalWriteAndReturnsBudget(t *testing.T) { + tr := newGatedRequiredControlTransport(nil) + controlBudget := newOutboundTrackedBudget(1 << 20) + c := newOutboundTestConn(t, tr, newOutboundTrackedBudget(1<<20)) + c.outboundControlTrackedBudget = controlBudget + + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + done := make(chan error, 1) + go func() { + done <- c.SendRequiredControl(ctx, proto.MessageServerResponse, &mt.Pong{MsgID: 1, PingID: 2}) + }() + + select { + case <-tr.started: + case <-time.After(time.Second): + t.Fatal("required control did not reach the physical writer") + } + select { + case err := <-done: + t.Fatalf("SendRequiredControl returned before physical write completed: %v", err) + case <-time.After(20 * time.Millisecond): + } + + tr.unblock() + select { + case err := <-done: + if err != nil { + t.Fatalf("SendRequiredControl: %v", err) + } + case <-time.After(time.Second): + t.Fatal("SendRequiredControl did not return after physical write") + } + if c.terminal.Load() { + t.Fatal("successful required control terminally closed the connection") + } + if got := controlBudget.snapshot(); got != 0 { + t.Fatalf("control budget after non-pending physical write = %d, want 0", got) + } + if got := tr.sends.Load(); got != 1 { + t.Fatalf("physical sends = %d, want 1", got) + } +} + +func TestSendRequiredControlReturnsAfterWriteWithoutWaitingForAck(t *testing.T) { + tr := &failAfterTransport{} + controlBudget := newOutboundTrackedBudget(1 << 20) + c := newOutboundTestConn(t, tr, newOutboundTrackedBudget(1<<20)) + c.outboundControlTrackedBudget = controlBudget + created := &mt.NewSessionCreated{FirstMsgID: 1, UniqueID: 2, ServerSalt: 3} + encoded, err := encodeOutboundMessageWithoutSlot(created) + if err != nil { + t.Fatalf("encode new_session_created: %v", err) + } + + if err := c.SendRequiredControl(context.Background(), proto.MessageFromServer, created); err != nil { + t.Fatalf("SendRequiredControl: %v", err) + } + if got := tr.stored.Load(); got != 1 { + t.Fatalf("completed physical sends = %d, want 1", got) + } + if got := controlBudget.snapshot(); got != int64(len(encoded.body)) { + t.Fatalf("pending control budget = %d, want %d until client ACK", got, len(encoded.body)) + } + + frame, err := crypto.NewClientCipher(rand.Reader).DecryptFromBuffer(c.key, &bin.Buffer{Buf: tr.lastFrame()}) + if err != nil { + t.Fatalf("decrypt new_session_created: %v", err) + } + c.AckServerMessages([]int64{frame.MessageID}) + deadline := time.Now().Add(time.Second) + for controlBudget.snapshot() != 0 && time.Now().Before(deadline) { + time.Sleep(time.Millisecond) + } + if got := controlBudget.snapshot(); got != 0 { + t.Fatalf("control budget after ACK = %d, want 0", got) + } +} + +func TestSendRequiredControlQueueDeadlineTerminatesAndReturnsBudget(t *testing.T) { + tr := &failAfterTransport{} + controlBudget := newOutboundTrackedBudget(1 << 20) + c := &Conn{ + transport: tr, + writer: tr, + metrics: NopMetrics{}, + writeTimeout: time.Second, + outboundTrackedBudget: newOutboundTrackedBudget(1 << 20), + outboundControlTrackedBudget: controlBudget, + outbound: make(chan outboundOp, 1), + outboundControl: make(chan outboundOp, 1), + outboundStop: make(chan struct{}), + } + // No actor is running and the bounded control queue is full, so the parent + // deadline must cover queue admission and make the failure terminal. + c.outboundControl <- outboundOp{kind: outboundAck} + ctx, cancel := context.WithTimeout(context.Background(), 25*time.Millisecond) + defer cancel() + + started := time.Now() + err := c.SendRequiredControl(ctx, proto.MessageServerResponse, &mt.Pong{MsgID: 1, PingID: 2}) + if !errors.Is(err, context.DeadlineExceeded) { + t.Fatalf("full control queue error = %v, want context deadline", err) + } + if elapsed := time.Since(started); elapsed > 500*time.Millisecond { + t.Fatalf("full control queue waited %v, want parent-deadline-bounded admission", elapsed) + } + if !c.terminal.Load() { + t.Fatal("required control queue failure did not terminally close the connection") + } + if got := tr.sends.Load(); got != 0 { + t.Fatalf("physical sends = %d, want 0", got) + } + if got := tr.closes.Load(); got != 1 { + t.Fatalf("transport closes = %d, want 1", got) + } + if got := controlBudget.snapshot(); got != 0 { + t.Fatalf("control budget after queue timeout = %d, want 0", got) + } +} + +func TestSendRequiredControlBlockedWriteUsesWholeOperationDeadline(t *testing.T) { + tr := newGatedRequiredControlTransport(io.ErrClosedPipe) + controlBudget := newOutboundTrackedBudget(1 << 20) + c := newOutboundTestConn(t, tr, newOutboundTrackedBudget(1<<20)) + c.outboundControlTrackedBudget = controlBudget + c.writeTimeout = 25 * time.Millisecond + + started := time.Now() + err := c.SendRequiredControl(context.Background(), proto.MessageServerResponse, &mt.Pong{MsgID: 1, PingID: 2}) + if !errors.Is(err, context.DeadlineExceeded) { + t.Fatalf("blocked required control error = %v, want context deadline", err) + } + if elapsed := time.Since(started); elapsed > 500*time.Millisecond { + t.Fatalf("blocked required control waited %v, want write-timeout-bounded operation", elapsed) + } + select { + case <-c.outboundDone: + case <-time.After(time.Second): + t.Fatal("outbound actor did not stop after required-control write timeout") + } + if !c.terminal.Load() { + t.Fatal("blocked required control did not terminally close the connection") + } + if got := tr.closes.Load(); got != 1 { + t.Fatalf("transport closes = %d, want 1", got) + } + if got := controlBudget.snapshot(); got != 0 { + t.Fatalf("control budget after blocked write = %d, want 0", got) + } +} + +func TestSendRequiredControlWriteFailureTerminatesAndReturnsBudget(t *testing.T) { + tr := &failAfterTransport{} + tr.failAt.Store(1) + controlBudget := newOutboundTrackedBudget(1 << 20) + c := newOutboundTestConn(t, tr, newOutboundTrackedBudget(1<<20)) + c.outboundControlTrackedBudget = controlBudget + + err := c.SendRequiredControl(context.Background(), proto.MessageServerResponse, &mt.Pong{MsgID: 1, PingID: 2}) + if err == nil { + t.Fatal("write-failed required control unexpectedly succeeded") + } + select { + case <-c.outboundDone: + case <-time.After(time.Second): + t.Fatal("outbound actor did not stop after required-control write failure") + } + if !c.terminal.Load() { + t.Fatal("write-failed required control did not terminally close the connection") + } + if got := tr.closes.Load(); got != 1 { + t.Fatalf("transport closes = %d, want 1", got) + } + if got := controlBudget.snapshot(); got != 0 { + t.Fatalf("control budget after write failure = %d, want 0", got) + } +} + +func TestSendRequiredControlBudgetFailureIsTerminal(t *testing.T) { + tr := &failAfterTransport{} + controlBudget := newOutboundTrackedBudget(1) + c := newOutboundTestConn(t, tr, newOutboundTrackedBudget(1<<20)) + c.outboundControlTrackedBudget = controlBudget + + err := c.SendRequiredControl(context.Background(), proto.MessageServerResponse, &mt.Pong{MsgID: 1, PingID: 2}) + if !errors.Is(err, ErrOutboundTrackedBudget) { + t.Fatalf("required control over budget = %v, want ErrOutboundTrackedBudget", err) + } + select { + case <-c.outboundDone: + case <-time.After(time.Second): + t.Fatal("outbound actor did not stop after required-control budget failure") + } + if !c.terminal.Load() { + t.Fatal("required-control budget failure did not terminally close the connection") + } + if got := tr.sends.Load(); got != 0 { + t.Fatalf("budget-rejected required control wrote %d frames, want 0", got) + } + if got := controlBudget.snapshot(); got != 0 { + t.Fatalf("control budget after reservation failure = %d, want 0", got) + } +} diff --git a/internal/mtprotoedge/provisional_revocation_test.go b/internal/mtprotoedge/provisional_revocation_test.go new file mode 100644 index 00000000..b958a383 --- /dev/null +++ b/internal/mtprotoedge/provisional_revocation_test.go @@ -0,0 +1,202 @@ +package mtprotoedge + +import ( + "context" + "errors" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/gotd/td/bin" + "github.com/gotd/td/mt" + "github.com/gotd/td/proto" + "github.com/gotd/td/proto/codec" + "github.com/gotd/td/tg" + + "telesrv/internal/store" + "telesrv/internal/store/memory" +) + +type activationGatedAuthKeyStore struct { + store.AuthKeyStore + gets atomic.Int32 + finalStarted chan struct{} + finalRelease chan struct{} + startOnce sync.Once +} + +func (s *activationGatedAuthKeyStore) Get(ctx context.Context, id [8]byte) (store.AuthKeyData, bool, error) { + if s.gets.Add(1) == 2 { + s.startOnce.Do(func() { close(s.finalStarted) }) + select { + case <-s.finalRelease: + case <-ctx.Done(): + return store.AuthKeyData{}, false, ctx.Err() + } + } + return s.AuthKeyStore.Get(ctx, id) +} + +func waitForManagedSessionAbsent(t *testing.T, manager *SessionManager, key sessionKey) { + t.Helper() + deadline := time.Now().Add(2 * time.Second) + for { + manager.mu.RLock() + claim, active := manager.claims[key], manager.bySession[key] + manager.mu.RUnlock() + if claim == nil && active == nil { + return + } + if time.Now().After(deadline) { + t.Fatalf("managed session survived terminal rejection: claim=%p active=%p", claim, active) + } + time.Sleep(time.Millisecond) + } +} + +func TestBadSaltStormRevalidatesStoreOnlyAtActivationBoundary(t *testing.T) { + const dc = 2 + keys := &countingAuthKeyStore{AuthKeyStore: memory.NewAuthKeyStore()} + handler := &admissionCountingRPC{} + addr, pub, _ := startTestServer(t, Options{DC: dc, AuthKeys: keys, RPC: handler}) + conn, auth, cipher := dialHandshake(t, addr, dc, pub) + ids := proto.NewMessageIDGen(time.Now) + firstID := ids.New(proto.MessageFromClient) + + const wrongFrames = 16 + for i := 0; i < wrongFrames; i++ { + msgID := firstID + if i != 0 { + msgID = ids.New(proto.MessageFromClient) + } + sendEncryptedWithSalt(t, conn, cipher, auth, auth.ServerSalt+1, msgID, &tg.HelpGetConfigRequest{}) + _, typeID, _ := readServerMessage(t, conn, cipher, auth.AuthKey) + if typeID != mt.BadServerSaltTypeID { + t.Fatalf("correction %d type = %#x", i, typeID) + } + } + if got := keys.gets.Load(); got != 1 { + t.Fatalf("AuthKeyStore.Get during bad-salt storm = %d, want initial lookup only", got) + } + + sendEncrypted(t, conn, cipher, auth, firstID, &tg.HelpGetConfigRequest{}) + collectReplyFrames(t, conn, cipher, auth.AuthKey, map[uint32]int{proto.ResultTypeID: 1}) + if got := keys.gets.Load(); got != 2 { + t.Fatalf("AuthKeyStore.Get after activation boundary = %d, want 2", got) + } + waitForAtomicCalls(t, &handler.calls, 1) +} + +func TestActivationFinalAuthKeyCheckRunsAfterClaim(t *testing.T) { + const dc = 2 + base := memory.NewAuthKeyStore() + keys := &activationGatedAuthKeyStore{ + AuthKeyStore: base, + finalStarted: make(chan struct{}), + finalRelease: make(chan struct{}), + } + defer func() { + select { + case <-keys.finalRelease: + default: + close(keys.finalRelease) + } + }() + handler := &admissionCountingRPC{} + addr, pub, srv := startTestServer(t, Options{DC: dc, AuthKeys: keys, RPC: handler}) + conn, auth, cipher := dialHandshake(t, addr, dc, pub) + msgID := proto.NewMessageIDGen(time.Now).New(proto.MessageFromClient) + + // The first Get is serveConn's decrypt lookup. The second is deliberately + // blocked: it must start only after BeginActivation indexed the claim. + sendEncrypted(t, conn, cipher, auth, msgID, &tg.HelpGetConfigRequest{}) + select { + case <-keys.finalStarted: + case <-time.After(2 * time.Second): + t.Fatal("activation final auth-key check did not start") + } + key := sessionKey{authKeyID: auth.AuthKey.ID, sessionID: auth.SessionID} + srv.conns.mu.RLock() + claim, active := srv.conns.claims[key], srv.conns.bySession[key] + srv.conns.mu.RUnlock() + if claim == nil || active != nil { + t.Fatalf("final auth-key check not protected by claim: claim=%p active=%p", claim, active) + } + + // Model Delete committing after the initial decrypt lookup but before the + // activation check returns. The claimant must emit terminal -404, never publish + // or dispatch the request, even before revocation fan-out gets the manager lock. + if err := base.Delete(context.Background(), auth.AuthKey.ID); err != nil { + t.Fatalf("delete auth key during activation: %v", err) + } + close(keys.finalRelease) + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + var response bin.Buffer + err := conn.Recv(ctx, &response) + var protocolErr *codec.ProtocolErr + if !errors.As(err, &protocolErr) || protocolErr.Code != codec.CodeAuthKeyNotFound { + t.Fatalf("revoked activation recv = %T %v, want protocol -404", err, err) + } + if got := handler.calls.Load(); got != 0 { + t.Fatalf("revoked activation executed %d RPCs", got) + } + waitForManagedSessionAbsent(t, srv.conns, key) +} + +func TestBadSaltProvisionalCannotReactivateDeletedAuthKey(t *testing.T) { + const dc = 2 + handler := &admissionCountingRPC{} + addr, pub, srv := startTestServer(t, Options{DC: dc, RPC: handler}) + provisional, auth, cipher := dialHandshake(t, addr, dc, pub) + ids := proto.NewMessageIDGen(time.Now) + reqMsgID := ids.New(proto.MessageFromClient) + + // Socket A is retained as a bad-salt provisional and is intentionally absent + // from SessionManager active/claim indexes. + sendEncryptedWithSalt(t, provisional, cipher, auth, auth.ServerSalt+1, reqMsgID, &tg.HelpGetConfigRequest{}) + _, typeID, _ := readServerMessage(t, provisional, cipher, auth.AuthKey) + if typeID != mt.BadServerSaltTypeID { + t.Fatalf("provisional correction type = %#x, want bad_server_salt", typeID) + } + key := sessionKey{authKeyID: auth.AuthKey.ID, sessionID: auth.SessionID} + srv.conns.mu.RLock() + activeBefore, claimBefore := srv.conns.bySession[key], srv.conns.claims[key] + srv.conns.mu.RUnlock() + if activeBefore != nil || claimBefore != nil { + t.Fatalf("bad-salt provisional leaked into manager: active=%p claim=%p", activeBefore, claimBefore) + } + + // Socket B uses the same auth key with another session and deletes it. The + // provisional is not manager-indexed, so correctness depends on its next-frame + // AuthKeyStore recheck rather than fan-out close alone. + destroyer := dialTransportOnly(t, addr) + destroySessionID := auth.SessionID ^ 1 + destroyBody := encodeClientMessageBodyForTest(t, &destroyAuthKeyRequest{}) + sendEncryptedWithSessionSaltAndSeq( + t, destroyer, cipher, auth, destroySessionID, auth.ServerSalt, + ids.New(proto.MessageFromClient), 1, destroyBody, + ) + destroyReplies := collectReplies(t, destroyer, cipher, auth.AuthKey, destroyAuthKeyOkTypeID) + mustHave(t, destroyReplies, destroyAuthKeyOkTypeID, "destroy_auth_key_ok") + if _, found, err := srv.authKeys.Get(context.Background(), auth.AuthKey.ID); err != nil || found { + t.Fatalf("auth key after destroy: found=%v err=%v", found, err) + } + + // A corrected resend must now receive terminal -404; it must not activate or + // execute the previously rejected business request with its cached key. + sendEncrypted(t, provisional, cipher, auth, reqMsgID, &tg.HelpGetConfigRequest{}) + ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second) + defer cancel() + var response bin.Buffer + err := provisional.Recv(ctx, &response) + var protocolErr *codec.ProtocolErr + if !errors.As(err, &protocolErr) || protocolErr.Code != codec.CodeAuthKeyNotFound { + t.Fatalf("corrected revoked provisional recv = %T %v, want protocol -404", err, err) + } + if got := handler.calls.Load(); got != 0 { + t.Fatalf("revoked provisional executed %d RPCs", got) + } + waitForManagedSessionAbsent(t, srv.conns, key) +} diff --git a/internal/mtprotoedge/rpc_flight.go b/internal/mtprotoedge/rpc_flight.go new file mode 100644 index 00000000..3c790c65 --- /dev/null +++ b/internal/mtprotoedge/rpc_flight.go @@ -0,0 +1,207 @@ +package mtprotoedge + +import ( + "context" + "errors" + "sync/atomic" +) + +const rpcResultFlightDefaultMaxPending = 8192 + +var ( + // ErrRPCResultFlightCapacity is returned before installing a new owner when + // the process-wide in-flight claim table has reached its hard bound. + ErrRPCResultFlightCapacity = errors.New("mtproto rpc result in-flight capacity exhausted") + ErrRPCResultFlightInvalid = errors.New("mtproto rpc result in-flight claim is invalid") +) + +type rpcResultAcquireState uint8 + +const ( + rpcResultAcquireCompleted rpcResultAcquireState = iota + 1 + rpcResultAcquirePending + rpcResultAcquireOwner +) + +// rpcResultAcquire is the atomic outcome for one +// (auth_key_id, session_id, req_msg_id) claim. +// +// Exactly one state-specific field is non-nil: +// - completed: encoded contains the immutable completed rpc_result; +// - pending: waiter joins the already-running owner; +// - owner: owner must eventually complete through rpcResultCache.Put or Abort. +type rpcResultAcquire struct { + state rpcResultAcquireState + encoded *encodedOutboundMessage + waiter *rpcResultWaiter + owner *rpcResultOwnerLease +} + +// rpcResultFlight is not part of the completed cache LRU/TTL lifecycle. Its +// done channel is closed exactly once while holding the owning cache shard lock; +// channel close publishes encoded/ok to all waiters without a waiter goroutine. +type rpcResultFlight struct { + done chan struct{} + encoded *encodedOutboundMessage + ok bool +} + +type rpcResultWaiter struct { + flight *rpcResultFlight +} + +// Wait blocks until the owner publishes through Put, aborts, or ctx expires. +// ok=false with err=nil means the owner aborted without a result. +func (w *rpcResultWaiter) Wait(ctx context.Context) (encoded *encodedOutboundMessage, ok bool, err error) { + if w == nil || w.flight == nil || ctx == nil { + return nil, false, ErrRPCResultFlightInvalid + } + + // Prefer an already-published result over a concurrently canceled context. + select { + case <-w.flight.done: + return w.flight.encoded, w.flight.ok, nil + default: + } + + select { + case <-w.flight.done: + return w.flight.encoded, w.flight.ok, nil + case <-ctx.Done(): + // If completion raced with cancellation, prefer the terminal flight state. + select { + case <-w.flight.done: + return w.flight.encoded, w.flight.ok, nil + default: + return nil, false, ctx.Err() + } + } +} + +type rpcResultOwnerLease struct { + cache *rpcResultCache + key rpcResultCacheKey + flight *rpcResultFlight +} + +// Abort releases an unfinished owner claim and wakes every waiter with no +// result. Pointer identity prevents an old lease from deleting a later owner +// that reacquired the same key. It returns true only for the winning abort. +func (l *rpcResultOwnerLease) Abort() bool { + if l == nil || l.cache == nil || l.flight == nil { + return false + } + s := l.cache.shard(l.key) + s.mu.Lock() + defer s.mu.Unlock() + + flight, ok := s.pending[l.key] + if !ok || flight != l.flight { + return false + } + delete(s.pending, l.key) + l.cache.flightLimit.release() + close(flight.done) + return true +} + +type rpcResultFlightLimit struct { + max int64 + used atomic.Int64 +} + +func (l *rpcResultFlightLimit) reserve() bool { + if l == nil || l.max <= 0 { + return false + } + for { + used := l.used.Load() + if used >= l.max { + return false + } + if l.used.CompareAndSwap(used, used+1) { + return true + } + } +} + +func (l *rpcResultFlightLimit) release() { + if l == nil { + return + } + if remaining := l.used.Add(-1); remaining < 0 { + // Put/Abort use map removal and lease identity to make double release + // impossible. Fail fast instead of masking a capacity-accounting bug that + // could otherwise admit more owners than the configured hard limit. + panic("mtproto rpc result in-flight counter underflow") + } +} + +func (l *rpcResultFlightLimit) snapshot() int64 { + if l == nil { + return 0 + } + return l.used.Load() +} + +// Acquire atomically returns a completed result, joins the existing in-flight +// owner, or installs the unique owner lease. Pending entries have a separate +// lifecycle from completed cache trim/TTL and consume one process-wide slot. +func (c *rpcResultCache) Acquire(authKeyID [8]byte, sessionID, reqMsgID int64) (rpcResultAcquire, error) { + if c == nil || reqMsgID == 0 { + return rpcResultAcquire{}, ErrRPCResultFlightInvalid + } + key := rpcResultCacheKey{authKeyID: authKeyID, sessionID: sessionID, reqMsgID: reqMsgID} + s := c.shard(key) + now := s.now() + + s.mu.Lock() + defer s.mu.Unlock() + + if elem, ok := s.byKey[key]; ok { + entry := elem.Value.(*rpcResultCacheEntry) + if entry.expiresAt.After(now) { + return rpcResultAcquire{state: rpcResultAcquireCompleted, encoded: entry.encoded}, nil + } + s.removeElement(elem) + } + if flight, ok := s.pending[key]; ok { + return rpcResultAcquire{ + state: rpcResultAcquirePending, + waiter: &rpcResultWaiter{flight: flight}, + }, nil + } + if !c.flightLimit.reserve() { + return rpcResultAcquire{}, ErrRPCResultFlightCapacity + } + flight := &rpcResultFlight{done: make(chan struct{})} + if s.pending == nil { + s.pending = make(map[rpcResultCacheKey]*rpcResultFlight) + } + s.pending[key] = flight + return rpcResultAcquire{ + state: rpcResultAcquireOwner, + owner: &rpcResultOwnerLease{cache: c, key: key, flight: flight}, + }, nil +} + +// completeRPCResultFlightLocked publishes encoded to the current owner claim. +// The caller must hold s.mu and must publish the completed cache entry first. +func (c *rpcResultCache) completeRPCResultFlightLocked( + s *rpcResultCacheShard, + key rpcResultCacheKey, + encoded *encodedOutboundMessage, +) { + if c == nil || s == nil || encoded == nil { + return + } + flight, ok := s.pending[key] + if !ok { + return + } + delete(s.pending, key) + flight.encoded = encoded + flight.ok = true + c.flightLimit.release() + close(flight.done) +} diff --git a/internal/mtprotoedge/rpc_flight_test.go b/internal/mtprotoedge/rpc_flight_test.go new file mode 100644 index 00000000..1aac5c0d --- /dev/null +++ b/internal/mtprotoedge/rpc_flight_test.go @@ -0,0 +1,396 @@ +package mtprotoedge + +import ( + "context" + "errors" + "sync" + "testing" + "time" +) + +func rpcFlightTestAuthID(seed byte) [8]byte { + return [8]byte{seed, seed + 1, seed + 2, seed + 3} +} + +func TestRPCResultFlightConcurrentAcquireHasUniqueOwner(t *testing.T) { + const callers = 64 + cache := newRPCResultCacheWithFlightLimit(time.Now, callers) + authKeyID := rpcFlightTestAuthID(1) + start := make(chan struct{}) + results := make(chan rpcResultAcquire, callers) + errs := make(chan error, callers) + + var wg sync.WaitGroup + wg.Add(callers) + for i := 0; i < callers; i++ { + go func() { + defer wg.Done() + <-start + claim, err := cache.Acquire(authKeyID, 10, 100) + if err != nil { + errs <- err + return + } + results <- claim + }() + } + close(start) + wg.Wait() + close(results) + close(errs) + + for err := range errs { + t.Fatalf("Acquire: %v", err) + } + owners := 0 + waiters := 0 + var owner *rpcResultOwnerLease + for claim := range results { + switch claim.state { + case rpcResultAcquireOwner: + owners++ + owner = claim.owner + case rpcResultAcquirePending: + waiters++ + if claim.waiter == nil { + t.Fatal("pending claim has nil waiter") + } + default: + t.Fatalf("unexpected claim state %d", claim.state) + } + } + if owners != 1 || waiters != callers-1 { + t.Fatalf("claims = owners:%d waiters:%d, want 1/%d", owners, waiters, callers-1) + } + if got := cache.flightLimit.snapshot(); got != 1 { + t.Fatalf("pending count = %d, want 1", got) + } + if owner == nil || !owner.Abort() { + t.Fatal("unique owner did not abort its claim") + } + if got := cache.flightLimit.snapshot(); got != 0 { + t.Fatalf("pending count after abort = %d, want 0", got) + } +} + +func TestRPCResultFlightPutPublishesAndWakesAllWaiters(t *testing.T) { + const waiters = 24 + cache := newRPCResultCacheWithFlightLimit(time.Now, 32) + authKeyID := rpcFlightTestAuthID(10) + owner, err := cache.Acquire(authKeyID, 20, 200) + if err != nil || owner.state != rpcResultAcquireOwner || owner.owner == nil { + t.Fatalf("owner Acquire = state:%d err:%v", owner.state, err) + } + + waiterClaims := make([]*rpcResultWaiter, 0, waiters) + for i := 0; i < waiters; i++ { + claim, acquireErr := cache.Acquire(authKeyID, 20, 200) + if acquireErr != nil || claim.state != rpcResultAcquirePending || claim.waiter == nil { + t.Fatalf("waiter %d Acquire = state:%d err:%v", i, claim.state, acquireErr) + } + waiterClaims = append(waiterClaims, claim.waiter) + } + + type waiterResult struct { + encoded *encodedOutboundMessage + cached *encodedOutboundMessage + ok bool + err error + } + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + results := make(chan waiterResult, waiters) + for _, waiter := range waiterClaims { + go func(w *rpcResultWaiter) { + encoded, ok, waitErr := w.Wait(ctx) + cached, _ := cache.Get(authKeyID, 20, 200) + results <- waiterResult{encoded: encoded, cached: cached, ok: ok, err: waitErr} + }(waiter) + } + + want := &encodedOutboundMessage{body: []byte{1, 2, 3, 4}, typeID: 42, reqMsgID: 200} + cache.Put(authKeyID, 20, 200, want) + for i := 0; i < waiters; i++ { + got := <-results + if got.err != nil || !got.ok { + t.Fatalf("waiter %d result = ok:%v err:%v", i, got.ok, got.err) + } + if got.encoded != want || got.cached != want { + t.Fatalf("waiter %d observed direct/cache pointers %p/%p, want %p", i, got.encoded, got.cached, want) + } + } + if got := cache.flightLimit.snapshot(); got != 0 { + t.Fatalf("pending count after Put = %d, want 0", got) + } + completed, err := cache.Acquire(authKeyID, 20, 200) + if err != nil || completed.state != rpcResultAcquireCompleted || completed.encoded != want { + t.Fatalf("completed Acquire = state:%d encoded:%p err:%v", completed.state, completed.encoded, err) + } + if owner.owner.Abort() { + t.Fatal("completed owner's stale lease aborted a resolved claim") + } +} + +func TestRPCResultFlightAbortWakesAndAllowsReclaim(t *testing.T) { + cache := newRPCResultCacheWithFlightLimit(time.Now, 2) + authKeyID := rpcFlightTestAuthID(20) + first, err := cache.Acquire(authKeyID, 30, 300) + if err != nil || first.state != rpcResultAcquireOwner { + t.Fatalf("first Acquire = state:%d err:%v", first.state, err) + } + waiting, err := cache.Acquire(authKeyID, 30, 300) + if err != nil || waiting.state != rpcResultAcquirePending { + t.Fatalf("waiting Acquire = state:%d err:%v", waiting.state, err) + } + if !first.owner.Abort() { + t.Fatal("first Abort lost") + } + if encoded, ok, waitErr := waiting.waiter.Wait(context.Background()); waitErr != nil || ok || encoded != nil { + t.Fatalf("aborted Wait = encoded:%p ok:%v err:%v", encoded, ok, waitErr) + } + if got := cache.flightLimit.snapshot(); got != 0 { + t.Fatalf("pending count after Abort = %d, want 0", got) + } + + second, err := cache.Acquire(authKeyID, 30, 300) + if err != nil || second.state != rpcResultAcquireOwner || second.owner == nil { + t.Fatalf("reclaim = state:%d err:%v", second.state, err) + } + if first.owner.Abort() { + t.Fatal("stale first lease aborted the replacement owner") + } + if got := cache.flightLimit.snapshot(); got != 1 { + t.Fatalf("pending count after reclaim = %d, want 1", got) + } + if !second.owner.Abort() || cache.flightLimit.snapshot() != 0 { + t.Fatal("replacement owner did not release its claim") + } +} + +func TestRPCResultFlightCompletedCachePressureDoesNotEvictPending(t *testing.T) { + now := time.Unix(1_000, 0) + cache := newRPCResultCacheWithFlightLimit(func() time.Time { return now }, 4) + authKeyID := rpcFlightTestAuthID(30) + pending, err := cache.Acquire(authKeyID, 40, 400) + if err != nil || pending.state != rpcResultAcquireOwner { + t.Fatalf("pending Acquire = state:%d err:%v", pending.state, err) + } + + key := rpcResultCacheKey{authKeyID: authKeyID, sessionID: 40, reqMsgID: 400} + shard := cache.shard(key) + shard.mu.Lock() + shard.maxEntries = 2 + shard.maxBytes = 2 + shard.mu.Unlock() + for i := int64(0); i < 16; i++ { + cache.Put(authKeyID, 40, 500+i, &encodedOutboundMessage{body: []byte{byte(i)}}) + } + if got := cache.flightLimit.snapshot(); got != 1 { + t.Fatalf("completed trim changed pending count to %d", got) + } + joined, err := cache.Acquire(authKeyID, 40, 400) + if err != nil || joined.state != rpcResultAcquirePending { + t.Fatalf("Acquire after completed trim = state:%d err:%v", joined.state, err) + } + + // Expire the independent completed cache and prove the pending owner remains. + now = now.Add(rpcResultCacheTTL + time.Second) + _, _ = cache.Get(authKeyID, 40, 515) + joinedAfterTTL, err := cache.Acquire(authKeyID, 40, 400) + if err != nil || joinedAfterTTL.state != rpcResultAcquirePending { + t.Fatalf("Acquire after completed TTL = state:%d err:%v", joinedAfterTTL.state, err) + } + if !pending.owner.Abort() || cache.flightLimit.snapshot() != 0 { + t.Fatal("pending claim did not survive pressure through explicit Abort") + } +} + +func TestRPCResultFlightCapacityAndCountReturn(t *testing.T) { + cache := newRPCResultCacheWithFlightLimit(time.Now, 2) + authKeyID := rpcFlightTestAuthID(40) + first, err := cache.Acquire(authKeyID, 50, 501) + if err != nil || first.state != rpcResultAcquireOwner { + t.Fatalf("first Acquire = state:%d err:%v", first.state, err) + } + second, err := cache.Acquire(authKeyID, 50, 502) + if err != nil || second.state != rpcResultAcquireOwner { + t.Fatalf("second Acquire = state:%d err:%v", second.state, err) + } + joined, err := cache.Acquire(authKeyID, 50, 501) + if err != nil || joined.state != rpcResultAcquirePending { + t.Fatalf("join at capacity = state:%d err:%v", joined.state, err) + } + if _, err := cache.Acquire(authKeyID, 50, 503); !errors.Is(err, ErrRPCResultFlightCapacity) { + t.Fatalf("over-capacity Acquire err = %v, want %v", err, ErrRPCResultFlightCapacity) + } + if got := cache.flightLimit.snapshot(); got != 2 { + t.Fatalf("pending count at capacity = %d, want 2", got) + } + + want := &encodedOutboundMessage{body: []byte{9}, reqMsgID: 501} + cache.Put(authKeyID, 50, 501, want) + if got := cache.flightLimit.snapshot(); got != 1 { + t.Fatalf("pending count after Put = %d, want 1", got) + } + third, err := cache.Acquire(authKeyID, 50, 503) + if err != nil || third.state != rpcResultAcquireOwner { + t.Fatalf("Acquire after returned slot = state:%d err:%v", third.state, err) + } + completed, err := cache.Acquire(authKeyID, 50, 501) + if err != nil || completed.state != rpcResultAcquireCompleted || completed.encoded != want { + t.Fatalf("completed Acquire at capacity = state:%d err:%v", completed.state, err) + } + if encoded, ok, waitErr := joined.waiter.Wait(context.Background()); waitErr != nil || !ok || encoded != want { + t.Fatalf("joined Wait = encoded:%p ok:%v err:%v", encoded, ok, waitErr) + } + if !second.owner.Abort() || !third.owner.Abort() { + t.Fatal("owners failed to return remaining capacity") + } + if got := cache.flightLimit.snapshot(); got != 0 { + t.Fatalf("final pending count = %d, want 0", got) + } +} + +func TestRPCResultFlightOversizedPutStillResolvesWaiters(t *testing.T) { + cache := newRPCResultCacheWithFlightLimit(time.Now, 1) + authKeyID := rpcFlightTestAuthID(50) + owner, err := cache.Acquire(authKeyID, 60, 600) + if err != nil || owner.state != rpcResultAcquireOwner { + t.Fatalf("owner Acquire = state:%d err:%v", owner.state, err) + } + joined, err := cache.Acquire(authKeyID, 60, 600) + if err != nil || joined.state != rpcResultAcquirePending { + t.Fatalf("joined Acquire = state:%d err:%v", joined.state, err) + } + + key := rpcResultCacheKey{authKeyID: authKeyID, sessionID: 60, reqMsgID: 600} + shard := cache.shard(key) + shard.mu.Lock() + shard.maxBytes = 1 + shard.mu.Unlock() + want := &encodedOutboundMessage{body: []byte{1, 2}, reqMsgID: 600} + cache.Put(authKeyID, 60, 600, want) + if encoded, ok, waitErr := joined.waiter.Wait(context.Background()); waitErr != nil || !ok || encoded != want { + t.Fatalf("oversized Wait = encoded:%p ok:%v err:%v", encoded, ok, waitErr) + } + if _, ok := cache.Get(authKeyID, 60, 600); ok { + t.Fatal("oversized result changed completed-cache compatibility") + } + if got := cache.flightLimit.snapshot(); got != 0 { + t.Fatalf("oversized Put leaked pending count %d", got) + } + if owner.owner.Abort() { + t.Fatal("oversized Put left its old owner abortable") + } + retry, err := cache.Acquire(authKeyID, 60, 600) + if err != nil || retry.state != rpcResultAcquireOwner { + t.Fatalf("retry after uncacheable completion = state:%d err:%v", retry.state, err) + } + if !retry.owner.Abort() { + t.Fatal("retry owner failed to abort") + } +} + +func TestRPCResultFlightWaitContextDoesNotReleaseOwner(t *testing.T) { + cache := newRPCResultCacheWithFlightLimit(time.Now, 1) + authKeyID := rpcFlightTestAuthID(60) + owner, err := cache.Acquire(authKeyID, 70, 700) + if err != nil { + t.Fatalf("owner Acquire: %v", err) + } + joined, err := cache.Acquire(authKeyID, 70, 700) + if err != nil { + t.Fatalf("joined Acquire: %v", err) + } + ctx, cancel := context.WithCancel(context.Background()) + cancel() + if encoded, ok, waitErr := joined.waiter.Wait(ctx); !errors.Is(waitErr, context.Canceled) || ok || encoded != nil { + t.Fatalf("canceled Wait = encoded:%p ok:%v err:%v", encoded, ok, waitErr) + } + if got := cache.flightLimit.snapshot(); got != 1 { + t.Fatalf("waiter cancellation released owner count to %d", got) + } + if !owner.owner.Abort() || cache.flightLimit.snapshot() != 0 { + t.Fatal("owner did not retain and release claim after waiter cancellation") + } +} + +func TestRPCResultFlightConcurrentCapacityReturnsAllSlots(t *testing.T) { + const ( + limit = 32 + callers = 512 + ) + cache := newRPCResultCacheWithFlightLimit(time.Now, limit) + authKeyID := rpcFlightTestAuthID(70) + start := make(chan struct{}) + errs := make(chan error, callers) + var wg sync.WaitGroup + wg.Add(callers) + for i := 0; i < callers; i++ { + i := i + go func() { + defer wg.Done() + <-start + claim, err := cache.Acquire(authKeyID, 80+int64(i%4), 1_000+int64(i)) + if errors.Is(err, ErrRPCResultFlightCapacity) { + return + } + if err != nil { + errs <- err + return + } + if claim.state != rpcResultAcquireOwner || claim.owner == nil { + errs <- errors.New("unique-key claim did not become owner") + return + } + if i%2 == 0 { + cache.Put(authKeyID, 80+int64(i%4), 1_000+int64(i), &encodedOutboundMessage{body: []byte{1}}) + } else if !claim.owner.Abort() { + errs <- errors.New("owner Abort lost") + } + }() + } + close(start) + wg.Wait() + close(errs) + for err := range errs { + t.Errorf("concurrent claim: %v", err) + } + if got := cache.flightLimit.snapshot(); got != 0 { + t.Fatalf("concurrent completion leaked %d pending slots", got) + } +} + +func TestQueuedRPCConnectionCloseAbortsOwnerClaim(t *testing.T) { + s := New(Options{}) + c := newInboundTestConn(s.rpcScheduler, 1, 4, time.Second) + claim, err := s.rpcResults.Acquire([8]byte{9}, 90, 900) + if err != nil || claim.state != rpcResultAcquireOwner || claim.owner == nil { + t.Fatalf("Acquire owner = state:%d err:%v", claim.state, err) + } + reservation, err := c.reserveInboundRPC(context.Background(), "test.queuedFlight", 4) + if err != nil { + t.Fatalf("reserve queued RPC: %v", err) + } + task := s.newInboundRPCTask(c, 900, "test.queuedFlight", []byte{1, 2, 3, 4}, claim.owner) + if err := reservation.commit(task); err != nil { + t.Fatalf("commit queued RPC: %v", err) + } + + // The scheduler is intentionally not started, so close must drain the queued + // task and invoke its independent flight-release callback. + c.closeInboundRPCScheduler() + if got := s.rpcResults.flightLimit.snapshot(); got != 0 { + t.Fatalf("connection close leaked %d queued owner claims", got) + } + if tasks, bytes := s.rpcScheduler.budgetSnapshot(); tasks != 0 || bytes != 0 { + t.Fatalf("connection close leaked scheduler budget %d/%d", tasks, bytes) + } + retry, err := s.rpcResults.Acquire([8]byte{9}, 90, 900) + if err != nil || retry.state != rpcResultAcquireOwner { + t.Fatalf("reclaim after queued close = state:%d err:%v", retry.state, err) + } + if !retry.owner.Abort() { + t.Fatal("replacement owner did not abort") + } +} diff --git a/internal/mtprotoedge/rpc_result_cache.go b/internal/mtprotoedge/rpc_result_cache.go index 28b46754..98fc5827 100644 --- a/internal/mtprotoedge/rpc_result_cache.go +++ b/internal/mtprotoedge/rpc_result_cache.go @@ -30,11 +30,14 @@ type rpcResultCacheEntry struct { expiresAt time.Time } -// rpcResultCache 缓存已回发的 rpc_result（按 auth_key+session+req_msg_id），用于 -// 跨连接重放重复请求。encodedOutboundMessage 构造后不可变（push fan-out 与 pending -// resend 均依赖该契约），因此 Get/Put 直接共享指针，不做防御性拷贝。 +// rpcResultCache 缓存已有交付证明的 rpc_result（按 auth_key+session+req_msg_id）， +// 用于跨连接重放重复请求。Put 的调用方必须先证明结果已物理写出，或原 logical Conn +// 已不可逆 fenced；绝不能发布“Conn 仍 current/open 但结果尚未上 wire”的完成态。 +// encodedOutboundMessage 构造后不可变（push fan-out 与 pending resend 均依赖该契约）， +// 因此 Get/Put 直接共享指针，不做防御性拷贝。 type rpcResultCache struct { - shards [rpcResultCacheShards]rpcResultCacheShard + shards [rpcResultCacheShards]rpcResultCacheShard + flightLimit rpcResultFlightLimit } type rpcResultCacheShard struct { @@ -46,13 +49,25 @@ type rpcResultCacheShard struct { bytes int order *list.List byKey map[rpcResultCacheKey]*list.Element + // pending is deliberately independent from the completed-result order/byKey + // cache. In-flight owners and waiters must not disappear when completed + // results expire or are trimmed under entry/byte pressure. + pending map[rpcResultCacheKey]*rpcResultFlight } func newRPCResultCache(now func() time.Time) *rpcResultCache { + return newRPCResultCacheWithFlightLimit(now, rpcResultFlightDefaultMaxPending) +} + +func newRPCResultCacheWithFlightLimit(now func() time.Time, maxPending int) *rpcResultCache { if now == nil { now = time.Now } + if maxPending <= 0 { + maxPending = rpcResultFlightDefaultMaxPending + } c := &rpcResultCache{} + c.flightLimit.max = int64(maxPending) for i := range c.shards { s := &c.shards[i] s.now = now @@ -61,6 +76,7 @@ func newRPCResultCache(now func() time.Time) *rpcResultCache { s.maxBytes = rpcResultCacheMaxBytes / rpcResultCacheShards s.order = list.New() s.byKey = make(map[rpcResultCacheKey]*list.Element) + s.pending = make(map[rpcResultCacheKey]*rpcResultFlight) } return c } @@ -102,28 +118,32 @@ func (c *rpcResultCache) Put(authKeyID [8]byte, sessionID, reqMsgID int64, encod key := rpcResultCacheKey{authKeyID: authKeyID, sessionID: sessionID, reqMsgID: reqMsgID} s := c.shard(key) size := len(encoded.body) - if s.maxBytes > 0 && size > s.maxBytes { - return - } + cacheable := s.maxBytes <= 0 || size <= s.maxBytes now := s.now() s.mu.Lock() defer s.mu.Unlock() - s.expireLocked(now) - if elem, ok := s.byKey[key]; ok { - s.removeElement(elem) + if cacheable { + s.expireLocked(now) + if elem, ok := s.byKey[key]; ok { + s.removeElement(elem) + } + entry := &rpcResultCacheEntry{ + key: key, + encoded: encoded, + size: size, + expiresAt: now.Add(s.ttl), + } + elem := s.order.PushBack(entry) + s.byKey[key] = elem + s.bytes += size + s.trimLocked() } - entry := &rpcResultCacheEntry{ - key: key, - encoded: encoded, - size: size, - expiresAt: now.Add(s.ttl), - } - elem := s.order.PushBack(entry) - s.byKey[key] = elem - s.bytes += size - s.trimLocked() + // Resolve the independent in-flight entry only after the completed cache has + // been published. Waiters awakened by this close can therefore immediately + // observe either the shared encoded result or the completed Get entry. + c.completeRPCResultFlightLocked(s, key, encoded) } func (s *rpcResultCacheShard) expireLocked(now time.Time) { diff --git a/internal/mtprotoedge/rpc_result_delivery_test.go b/internal/mtprotoedge/rpc_result_delivery_test.go new file mode 100644 index 00000000..30905727 --- /dev/null +++ b/internal/mtprotoedge/rpc_result_delivery_test.go @@ -0,0 +1,199 @@ +package mtprotoedge + +import ( + "context" + "errors" + "io" + "sync" + "testing" + "time" + + "github.com/gotd/td/bin" + "github.com/gotd/td/proto" + "github.com/gotd/td/tg" +) + +type blockingCloseRPCResultTransport struct { + started chan struct{} + release chan struct{} + once sync.Once +} + +func newBlockingCloseRPCResultTransport() *blockingCloseRPCResultTransport { + return &blockingCloseRPCResultTransport{started: make(chan struct{}), release: make(chan struct{})} +} + +func (*blockingCloseRPCResultTransport) Send(context.Context, *bin.Buffer) error { return nil } +func (*blockingCloseRPCResultTransport) Recv(context.Context, *bin.Buffer) error { return io.EOF } +func (t *blockingCloseRPCResultTransport) Close() error { + t.once.Do(func() { close(t.started) }) + <-t.release + return nil +} + +func TestRPCResultCachePublishesOnlyAfterPhysicalWrite(t *testing.T) { + tr := newGatedRequiredControlTransport(nil) + s := New(Options{WriteTimeout: time.Second}) + key := newTestAuthKey(t) + c := s.newConn(tr, key, 74001, 1) + t.Cleanup(c.ForceClose) + reqMsgID := proto.NewMessageIDGen(time.Now).New(proto.MessageFromClient) + + owner, err := s.rpcResults.Acquire(key.ID, c.sessionID, reqMsgID) + if err != nil || owner.state != rpcResultAcquireOwner { + t.Fatalf("initial flight owner = %+v err=%v", owner, err) + } + done := make(chan error, 1) + go func() { + done <- s.sendResult(context.Background(), c, reqMsgID, &tg.Config{ThisDC: 2}) + }() + select { + case <-tr.started: + case <-time.After(time.Second): + t.Fatal("rpc_result did not reach physical writer") + } + if _, ok := s.rpcResults.Get(key.ID, c.sessionID, reqMsgID); ok { + t.Fatal("rpc_result became completed before physical write") + } + pending, err := s.rpcResults.Acquire(key.ID, c.sessionID, reqMsgID) + if err != nil || pending.state != rpcResultAcquirePending || pending.waiter == nil { + t.Fatalf("flight while write blocked = %+v err=%v", pending, err) + } + + tr.unblock() + select { + case err := <-done: + if err != nil { + t.Fatalf("sendResult after write: %v", err) + } + case <-time.After(time.Second): + t.Fatal("sendResult did not finish") + } + if encoded, ok, err := pending.waiter.Wait(context.Background()); err != nil || !ok || encoded == nil { + t.Fatalf("pending waiter after write = encoded:%p ok:%v err:%v", encoded, ok, err) + } + completed, err := s.rpcResults.Acquire(key.ID, c.sessionID, reqMsgID) + if err != nil || completed.state != rpcResultAcquireCompleted || completed.encoded == nil { + t.Fatalf("flight after write = %+v err=%v", completed, err) + } +} + +func TestRPCResultPrewriteFailureFencesConnBeforeCachePublication(t *testing.T) { + tr := &collectingSessionTransport{} + s := New(Options{WriteTimeout: 20 * time.Millisecond}) + // No rpc_result can reserve its conservative 3x wire scratch from one byte. + // The actor therefore fails before touching the socket, which used to leave a + // live Conn with a prematurely completed cache entry. + s.outboundScratchPool = newOutboundScratchPool(1) + key := newTestAuthKey(t) + c := s.newConn(tr, key, 74002, 1) + reqMsgID := proto.NewMessageIDGen(time.Now).New(proto.MessageFromClient) + owner, err := s.rpcResults.Acquire(key.ID, c.sessionID, reqMsgID) + if err != nil || owner.state != rpcResultAcquireOwner { + t.Fatalf("initial flight owner = %+v err=%v", owner, err) + } + + err = s.sendResult(context.Background(), c, reqMsgID, &tg.Config{ThisDC: 2}) + if err == nil || (!errors.Is(err, context.DeadlineExceeded) && + !errors.Is(err, ErrConnClosed) && !errors.Is(err, ErrOutboundTrackedBudget)) { + t.Fatalf("prewrite sendResult error = %v", err) + } + closeDeadline := time.Now().Add(time.Second) + for !tr.closed.Load() && time.Now().Before(closeDeadline) { + time.Sleep(time.Millisecond) + } + if !c.terminal.Load() || !tr.closed.Load() || c.isPhysicalTransportCurrentOpen() { + t.Fatalf("failed delivery did not fence Conn: terminal=%v closed=%v current_open=%v", c.terminal.Load(), tr.closed.Load(), c.isPhysicalTransportCurrentOpen()) + } + completed, acquireErr := s.rpcResults.Acquire(key.ID, c.sessionID, reqMsgID) + if acquireErr != nil || completed.state != rpcResultAcquireCompleted || completed.encoded == nil { + t.Fatalf("terminal replay cache = %+v err=%v", completed, acquireErr) + } + if got := s.rpcResults.flightLimit.snapshot(); got != 0 { + t.Fatalf("failed delivery leaked flight slot: %d", got) + } + c.Close() +} + +func TestRPCResultFailureAfterIntentionalTerminalDoesNotCloseTransferLease(t *testing.T) { + tr := &collectingSessionTransport{} + s := New(Options{WriteTimeout: time.Second}) + s.outboundTrackedBudget = newOutboundTrackedBudget(1) + key := newTestAuthKey(t) + oldConn := s.newConn(tr, key, 74003, 1) + reqMsgID := proto.NewMessageIDGen(time.Now).New(proto.MessageFromClient) + owner, err := s.rpcResults.Acquire(key.ID, oldConn.sessionID, reqMsgID) + if err != nil || owner.state != rpcResultAcquireOwner { + t.Fatalf("initial flight owner = %+v err=%v", owner, err) + } + + // Session replacement publishes terminal before transferring the physical + // lease. A late old-generation result sees ErrConnClosed at producer admission; + // it may publish cache-only, but must not upgrade that intentional fence into + // a physical close that makes Transfer fail. + oldConn.beginTerminalShutdown() + err = s.sendResult(context.Background(), oldConn, reqMsgID, &tg.Config{ThisDC: 2}) + if !errors.Is(err, ErrOutboundTrackedBudget) { + t.Fatalf("late result error = %v, want ErrOutboundTrackedBudget", err) + } + if tr.closed.Load() { + t.Fatal("late result closed intentionally transferable transport") + } + if !oldConn.waitOutboundShutdownUntil(time.Second) { + t.Fatal("old outbound actor did not stop") + } + nextLease, ok := oldConn.transferTransportOwnership() + if !ok || nextLease == nil { + t.Fatal("late result prevented physical transfer") + } + newConn := s.newConnWithLease(nextLease, key, 74004, 1) + oldConn.ForceClose() + if tr.closed.Load() || !newConn.isPhysicalTransportCurrentOpen() { + t.Fatalf("stale close after transfer: raw_closed=%v current_open=%v", tr.closed.Load(), newConn.isPhysicalTransportCurrentOpen()) + } + completed, acquireErr := s.rpcResults.Acquire(key.ID, oldConn.sessionID, reqMsgID) + if acquireErr != nil || completed.state != rpcResultAcquireCompleted || completed.encoded == nil { + t.Fatalf("late result cache = %+v err=%v", completed, acquireErr) + } + newConn.ForceClose() +} + +func TestRPCResultPublishesBeforePathologicalPhysicalCloseReturns(t *testing.T) { + tr := newBlockingCloseRPCResultTransport() + s := New(Options{WriteTimeout: time.Second}) + s.outboundTrackedBudget = newOutboundTrackedBudget(1) + key := newTestAuthKey(t) + c := s.newConn(tr, key, 74005, 1) + reqMsgID := proto.NewMessageIDGen(time.Now).New(proto.MessageFromClient) + owner, err := s.rpcResults.Acquire(key.ID, c.sessionID, reqMsgID) + if err != nil || owner.state != rpcResultAcquireOwner { + t.Fatalf("initial flight owner = %+v err=%v", owner, err) + } + + done := make(chan error, 1) + go func() { done <- s.sendResult(context.Background(), c, reqMsgID, &tg.Config{ThisDC: 2}) }() + select { + case <-tr.started: + case <-time.After(time.Second): + t.Fatal("physical Close did not start") + } + select { + case err := <-done: + if !errors.Is(err, ErrOutboundTrackedBudget) { + t.Fatalf("sendResult error = %v, want budget error", err) + } + case <-time.After(time.Second): + t.Fatal("pathological physical Close blocked result publication") + } + completed, acquireErr := s.rpcResults.Acquire(key.ID, c.sessionID, reqMsgID) + if acquireErr != nil || completed.state != rpcResultAcquireCompleted || completed.encoded == nil { + t.Fatalf("completed result before raw Close return = %+v err=%v", completed, acquireErr) + } + close(tr.release) + if c.transportLease != nil { + if err := c.transportLease.owner.waitClosed(); err != nil { + t.Fatalf("physical Close: %v", err) + } + } + c.Close() +} diff --git a/internal/mtprotoedge/rpc_test.go b/internal/mtprotoedge/rpc_test.go index 8663134d..8ec2e191 100644 --- a/internal/mtprotoedge/rpc_test.go +++ b/internal/mtprotoedge/rpc_test.go @@ -257,7 +257,8 @@ func TestDuplicateRPCResultAcrossReconnectUsesSessionCache(t *testing.T) { func TestCanceledRPCErrorIsNotCachedAcrossReconnect(t *testing.T) { const dc = 2 handler := &canceledInternalRPC{ - firstDone: make(chan struct{}), + firstStarted: make(chan struct{}), + firstDone: make(chan struct{}), } addr, pub, _ := startTestServer(t, Options{DC: dc, RPC: handler}) conn, auth, cipher := dialHandshake(t, addr, dc, pub) @@ -266,6 +267,11 @@ func TestCanceledRPCErrorIsNotCachedAcrossReconnect(t *testing.T) { reqMsgID := clientMsgID.New(proto.MessageFromClient) sendEncrypted(t, conn, cipher, auth, reqMsgID, &tg.HelpGetConfigRequest{}) + select { + case <-handler.firstStarted: + case <-time.After(5 * time.Second): + t.Fatal("timed out waiting for first rpc to start after session barrier") + } _ = conn.Close() select { case <-handler.firstDone: @@ -355,12 +361,14 @@ func (h *runningDeadlineRPC) Dispatch(ctx context.Context, _ [8]byte, _ int64, _ func (h *runningDeadlineRPC) NegotiatedLayer([8]byte, int64) (int, bool) { return 227, true } type canceledInternalRPC struct { - calls atomic.Int32 - firstDone chan struct{} + calls atomic.Int32 + firstStarted chan struct{} + firstDone chan struct{} } func (h *canceledInternalRPC) Dispatch(ctx context.Context, _ [8]byte, _ int64, _ *bin.Buffer) (bin.Encoder, error) { if h.calls.Add(1) == 1 { + close(h.firstStarted) <-ctx.Done() close(h.firstDone) return nil, tgerr.New(500, "INTERNAL_SERVER_ERROR") diff --git a/internal/mtprotoedge/server.go b/internal/mtprotoedge/server.go index c9d09cd5..c8d0752f 100644 --- a/internal/mtprotoedge/server.go +++ b/internal/mtprotoedge/server.go @@ -172,7 +172,7 @@ func (o *Options) setDefaults() { o.RPCGlobalWorkers = 256 } if o.RPCGlobalMaxTasks <= 0 { - o.RPCGlobalMaxTasks = 8192 + o.RPCGlobalMaxTasks = rpcResultFlightDefaultMaxPending } if o.RPCGlobalMaxBytes <= 0 { o.RPCGlobalMaxBytes = 512 << 20 @@ -295,7 +295,7 @@ func New(opts Options) *Server { clock: opts.Clock, rand: opts.Rand, types: tmap.New(tg.TypesMap(), mt.TypesMap(), proto.TypesMap()), - rpcResults: newRPCResultCache(opts.Clock.Now), + rpcResults: newRPCResultCacheWithFlightLimit(opts.Clock.Now, opts.RPCGlobalMaxTasks), admission: newAdmissionController(opts.MaxConnections, opts.MaxConnectionsPerIP, opts.MaxConcurrentHandshakes), } } @@ -307,8 +307,33 @@ func (s *Server) Conns() *SessionManager { // newConn 基于一次解密结果创建一个可发送的连接对象。 func (s *Server) newConn(tc transport.Conn, key crypto.AuthKey, sessionID, salt int64) *Conn { + if lease, ok := tc.(*physicalTransportLease); ok { + return s.newConnWithLease(lease, key, sessionID, salt) + } + if tc != nil { + _, lease := newPhysicalTransportOwner(tc) + return s.newConnWithLease(lease, key, sessionID, salt) + } + // Preserve the nil transport used by construction-only tests. + return s.buildConn(nil, nil, key, sessionID, salt) +} + +// newConnWithLease attaches a logical Conn to an explicitly owned physical +// transport generation. Production session replacement must Transfer the old +// lease first; it must never wrap the same raw transport in a second owner. +func (s *Server) newConnWithLease(lease *physicalTransportLease, key crypto.AuthKey, sessionID, salt int64) *Conn { + if lease == nil { + panic("mtprotoedge: nil physical transport lease") + } + c := s.buildConn(lease, lease, key, sessionID, salt) + lease.bindLogicalConn(c) + return c +} + +func (s *Server) buildConn(tc transport.Conn, lease *physicalTransportLease, key crypto.AuthKey, sessionID, salt int64) *Conn { c := &Conn{ transport: tc, + transportLease: lease, writer: tc, cipher: s.cipher, msgID: proto.NewMessageIDGen(s.clock.Now), @@ -558,7 +583,8 @@ func (s *Server) promoteConn(raw net.Conn, obfuscated bool) (transport.Conn, err // - auth_key_id 未注册：回 AuthKeyNotFound，促使客户端重新握手。 // // 连接建立 session 后注册到 SessionManager，结束时注销。 -func (s *Server) serveConn(ctx context.Context, conn transport.Conn) (err error) { +func (s *Server) serveConn(ctx context.Context, raw transport.Conn) (err error) { + transportOwner, conn := newPhysicalTransportOwner(raw) s.metrics.ConnOpened() s.log.Debug("Connection accepted") @@ -568,9 +594,13 @@ func (s *Server) serveConn(ctx context.Context, conn transport.Conn) (err error) // this stack has stopped using b/plain; transport.Close may have raced us earlier and must // not return that memory budget prematurely. releaseInboundFrameOwnership(conn) - // 先同步关闭物理 socket，解除可能阻塞在 writer.Send 的 outbound actor； - // 再停止 logical Conn，避免 Close 等 actor 时反过来等到 write deadline。 - _ = conn.Close() + // Publish the terminal/RPC-cancel gates before index removal or lifecycle + // observers. Physical close then releases a writer already inside Send; the + // final Close only waits for the now-fenced actors to converge. + if current != nil { + current.beginTerminalShutdown() + } + _ = transportOwner.CloseAny() if current != nil { s.conns.Unregister(current) current.Close() @@ -584,7 +614,7 @@ func (s *Server) serveConn(ctx context.Context, conn transport.Conn) (err error) defer cancel() go func() { <-ctx.Done() - _ = conn.Close() + _ = transportOwner.CloseAny() }() cs := newConnState() @@ -601,7 +631,7 @@ func (s *Server) serveConn(ctx context.Context, conn transport.Conn) (err error) // 建立 session 前（current==nil，握手 + 首个加密消息之前）用较短的 handshakeTimeout // 快速回收静默的半开 / 异常连接；建立 session 后用 readTimeout（客户端有 ping 心跳）。 timeout := s.readTimeout - if current == nil { + if current == nil || !current.isActive() { timeout = s.handshakeTimeout } if err := s.recv(ctx, conn, &b, timeout); err != nil { @@ -618,6 +648,12 @@ func (s *Server) serveConn(ctx context.Context, conn transport.Conn) (err error) } if authKeyID == emptyAuthKeyID { + // A physical socket may perform key exchange only before it owns an + // encrypted logical session. Mixing the direct exchange writer with an + // active outbound actor would bypass generation/write serialization. + if current != nil { + return errors.New("unencrypted exchange on established encrypted connection") + } releaseHandshake, admitted := s.admission.tryAcquireHandshake() if !admitted { if err := s.sendProtoError(ctx, conn, codec.CodeTransportFlood); err != nil { @@ -647,7 +683,9 @@ func (s *Server) serveConn(ctx context.Context, conn transport.Conn) (err error) // 已建立连接复用缓存密钥走快路径（fetchedKey=nil）：避开每帧回查 AuthKeyStore—— // 这是 mtprotoedge 层最热的库访问点。密钥材料创建后不可变；销毁(destroy_auth_key)/ // 撤销由 SessionManager 主动 Close 连接保证失效，不依赖此被动回查。仅 destroy_auth_key - // 的发起连接置 keyDestroyed，使其下一帧回落到 Get→AuthKeyNotFound，维持原契约。 + // 的发起连接置 keyDestroyed，使其下一帧回落到 Get→AuthKeyNotFound。尚未进入 + // SessionManager 的 bad-salt provisional 会在 handleEncrypted 建立 activation claim + // 后精确复查一次，既把撤销与激活线性化，也不把 salt storm 放大成 PG 写风暴。 var fetchedKey *store.AuthKeyData if current == nil || current.authKeyID != authKeyID || current.keyDestroyed.Load() { d, found, err := s.authKeys.Get(ctx, authKeyID) @@ -655,7 +693,11 @@ func (s *Server) serveConn(ctx context.Context, conn transport.Conn) (err error) return fmt.Errorf("lookup auth key: %w", err) } if !found { - if err := s.sendProtoError(ctx, conn, codec.CodeAuthKeyNotFound); err != nil { + writer := transport.Conn(conn) + if current != nil { + writer = current.transport + } + if err := s.sendProtoError(ctx, writer, codec.CodeAuthKeyNotFound); err != nil { return err } // -404 对 TDesktop 是 terminal key failure；继续保留 socket 只会允许 @@ -666,6 +708,11 @@ func (s *Server) serveConn(ctx context.Context, conn transport.Conn) (err error) } current, err = s.handleEncrypted(ctx, conn, cs, current, fetchedKey, &b, &plain) + if errors.Is(err, errActivationAuthKeyRejected) { + // handleEncrypted writes -404 while its activation claim still owns the + // physical writer, then its deferred abort removes/closes the claim. + return nil + } if err != nil { return err } diff --git a/internal/mtprotoedge/session_activation_test.go b/internal/mtprotoedge/session_activation_test.go new file mode 100644 index 00000000..7838e642 --- /dev/null +++ b/internal/mtprotoedge/session_activation_test.go @@ -0,0 +1,253 @@ +package mtprotoedge + +import ( + "context" + "errors" + "testing" + "time" + + "go.uber.org/zap/zaptest" + + "github.com/gotd/td/mt" + "github.com/gotd/td/proto" +) + +func newSessionActivationTestConn(t *testing.T, authKeyID [8]byte, sessionID int64) *Conn { + t.Helper() + c := &Conn{ + authKeyID: authKeyID, + sessionID: sessionID, + metrics: NopMetrics{}, + } + c.startOutbound() + t.Cleanup(c.Close) + return c +} + +func TestSessionActivationGatesReplacementBeforePublishing(t *testing.T) { + manager := NewSessionManager(zaptest.NewLogger(t)) + key := [8]byte{1, 2, 3, 4} + oldConn := newSessionActivationTestConn(t, key, 7001) + newConn := newSessionActivationTestConn(t, key, 7001) + + if err := manager.Register(oldConn); err != nil { + t.Fatalf("register initial: %v", err) + } + if !oldConn.isActive() { + t.Fatal("initial connection was not activated") + } + if newConn.lifecycleState() != connLifecycleProvisional { + t.Fatal("provisional replacement was active before registration") + } + + if err := manager.Register(newConn); err != nil { + t.Fatalf("register replacement: %v", err) + } + if !newConn.isActive() { + t.Fatal("replacement was not activated") + } + if oldConn.lifecycleState() != connLifecycleRetired || !oldConn.terminal.Load() { + t.Fatalf("old connection gates = lifecycle:%v terminal:%v", oldConn.lifecycleState(), oldConn.terminal.Load()) + } + if err := oldConn.SendAsync(context.Background(), proto.MessageFromServer, &mt.MsgsAck{}); !errors.Is(err, ErrConnClosed) { + t.Fatalf("stale old connection send error = %v, want ErrConnClosed", err) + } + + manager.mu.RLock() + got := manager.bySession[sessionKey{authKeyID: key, sessionID: 7001}] + manager.mu.RUnlock() + if got != newConn { + t.Fatalf("published session = %p, want replacement %p", got, newConn) + } +} + +func TestSessionActivationClaimPreemptionCannotReversePublish(t *testing.T) { + manager := NewSessionManager(zaptest.NewLogger(t)) + key := [8]byte{9, 8, 7, 6} + first := newSessionActivationTestConn(t, key, 8001) + second := newSessionActivationTestConn(t, key, 8001) + + if err := manager.BeginActivation(first); err != nil { + t.Fatalf("begin first activation: %v", err) + } + if first.lifecycleState() != connLifecycleClaiming { + t.Fatalf("first lifecycle = %v, want claiming", first.lifecycleState()) + } + if got := manager.Online(); got != 0 { + t.Fatalf("online during claim = %d, want 0", got) + } + + if err := manager.BeginActivation(second); err != nil { + t.Fatalf("begin superseding activation: %v", err) + } + if first.lifecycleState() != connLifecycleRetired || !first.terminal.Load() { + t.Fatalf("superseded first lifecycle=%v terminal=%v", first.lifecycleState(), first.terminal.Load()) + } + if err := manager.PublishActivation(first); !errors.Is(err, ErrSessionActivationSuperseded) { + t.Fatalf("stale publish error = %v, want superseded", err) + } + if err := manager.Register(first); !errors.Is(err, ErrSessionActivationSuperseded) { + t.Fatalf("stale register error = %v, want superseded", err) + } + if err := manager.PublishActivation(second); err != nil { + t.Fatalf("publish second activation: %v", err) + } + + manager.mu.RLock() + got := manager.bySession[sessionKey{authKeyID: key, sessionID: 8001}] + claim := manager.claims[sessionKey{authKeyID: key, sessionID: 8001}] + manager.mu.RUnlock() + if got != second || claim != nil || !second.isActive() { + t.Fatalf("activation owner=%p claim=%p second_active=%v", got, claim, second.isActive()) + } +} + +func TestBeginActivationWaitsForPreviousPhysicalWriterFence(t *testing.T) { + manager := NewSessionManager(zaptest.NewLogger(t)) + key := [8]byte{3, 3, 3, 3} + releaseClose := make(chan struct{}) + transport := newSlowCloseTransport(0, releaseClose) + oldConn := &Conn{ + authKeyID: key, + sessionID: 8501, + metrics: NopMetrics{}, + transport: transport, + writer: transport, + } + oldConn.startOutbound() + t.Cleanup(oldConn.Close) + if err := manager.Register(oldConn); err != nil { + t.Fatalf("register old: %v", err) + } + newConn := newSessionActivationTestConn(t, key, oldConn.sessionID) + + beginDone := make(chan error, 1) + go func() { beginDone <- manager.BeginActivation(newConn) }() + deadline := time.Now().Add(time.Second) + for transport.closes.Load() == 0 && time.Now().Before(deadline) { + time.Sleep(time.Millisecond) + } + if transport.closes.Load() == 0 { + t.Fatal("replacement did not start closing previous transport") + } + select { + case err := <-beginDone: + t.Fatalf("BeginActivation returned before old physical close: %v", err) + default: + } + close(releaseClose) + select { + case err := <-beginDone: + if err != nil { + t.Fatalf("BeginActivation after old close: %v", err) + } + case <-time.After(time.Second): + t.Fatal("BeginActivation did not converge after old close") + } + if err := manager.PublishActivation(newConn); err != nil { + t.Fatalf("publish replacement: %v", err) + } +} + +func TestSessionActivationClaimIndexesCleanOnPublishAndAbort(t *testing.T) { + manager := NewSessionManager(zaptest.NewLogger(t)) + key := [8]byte{4, 4, 4, 4} + published := newSessionActivationTestConn(t, key, 9001) + if err := manager.BeginActivation(published); err != nil { + t.Fatalf("begin publish claim: %v", err) + } + manager.mu.RLock() + if manager.claims[connSessionKey(published)] != published || manager.claimsByAuth[key][published.sessionID] != published { + manager.mu.RUnlock() + t.Fatal("claim indexes missing after BeginActivation") + } + manager.mu.RUnlock() + if err := manager.PublishActivation(published); err != nil { + t.Fatalf("publish claim: %v", err) + } + manager.mu.RLock() + globalClaims, authClaims := len(manager.claims), len(manager.claimsByAuth[key]) + manager.mu.RUnlock() + if globalClaims != 0 || authClaims != 0 { + t.Fatalf("claim indexes after publish = %d/%d, want 0/0", globalClaims, authClaims) + } + + aborted := newSessionActivationTestConn(t, key, 9002) + if err := manager.BeginActivation(aborted); err != nil { + t.Fatalf("begin abort claim: %v", err) + } + manager.AbortActivation(aborted) + manager.mu.RLock() + globalClaims, authClaims = len(manager.claims), len(manager.claimsByAuth[key]) + manager.mu.RUnlock() + if globalClaims != 0 || authClaims != 0 || aborted.lifecycleState() != connLifecycleRetired { + t.Fatalf("claim indexes/lifecycle after abort = %d/%d/%v", globalClaims, authClaims, aborted.lifecycleState()) + } +} + +func TestRawAuthKeyCloseExactConnDoesNotExcludeSameSessionReplacement(t *testing.T) { + manager := NewSessionManager(zaptest.NewLogger(t)) + key := [8]byte{4, 3, 2, 1} + const sessionID = 9050 + + // The destroy request may finish on a retired Conn after another physical + // connection has become the owner of the same logical session. + destroyer := newSessionActivationTestConn(t, key, sessionID) + destroyer.beginTerminalShutdown() + replacement := newSessionActivationTestConn(t, key, sessionID) + if err := manager.Register(replacement); err != nil { + t.Fatalf("register replacement: %v", err) + } + + if got := manager.CloseSessionsForRawAuthKeyExceptConn(key, destroyer); got != 1 { + t.Fatalf("closed sessions = %d, want replacement only", got) + } + manager.mu.RLock() + active, claim := manager.bySession[sessionKey{authKeyID: key, sessionID: sessionID}], manager.claims[sessionKey{authKeyID: key, sessionID: sessionID}] + manager.mu.RUnlock() + if active != nil || claim != nil || !replacement.terminal.Load() { + t.Fatalf("same-session replacement escaped exact exclusion: active=%p claim=%p terminal=%v", active, claim, replacement.terminal.Load()) + } +} + +func TestDestroyFencesOutboundAndReservedRPCBeforeRemoval(t *testing.T) { + manager := NewSessionManager(zaptest.NewLogger(t)) + key := [8]byte{5, 5, 5, 5} + c := newSessionActivationTestConn(t, key, 9101) + scheduler := newInboundRPCScheduler(1, 8, 1<<20) + c.startInboundRPCScheduler(scheduler, 1, 8, time.Second) + if err := manager.Register(c); err != nil { + t.Fatalf("register: %v", err) + } + + reservation, err := c.reserveInboundRPC(context.Background(), "test.destroyFence", 8) + if err != nil { + t.Fatalf("reserve inbound RPC: %v", err) + } + destroyed := make(chan bool, 1) + go func() { + destroyed <- manager.DestroySessionForAuthKey(key, c.sessionID) + }() + select { + case <-c.rpcRootCtx.Done(): + case <-time.After(time.Second): + t.Fatal("destroy did not synchronously cancel RPC admission") + } + if err := c.SendAsync(context.Background(), proto.MessageFromServer, &mt.MsgsAck{}); !errors.Is(err, ErrConnClosed) { + t.Fatalf("send after destroy fence = %v, want ErrConnClosed", err) + } + if err := reservation.commit(inboundRPC{}); !errors.Is(err, ErrConnClosed) { + t.Fatalf("reserved RPC commit after destroy = %v, want ErrConnClosed", err) + } + select { + case ok := <-destroyed: + if !ok { + t.Fatal("destroy returned false") + } + case <-time.After(time.Second): + t.Fatal("destroy did not converge after reservation commit") + } + if tasks, bytes := scheduler.budgetSnapshot(); tasks != 0 || bytes != 0 { + t.Fatalf("scheduler budget after destroy = %d/%d, want 0/0", tasks, bytes) + } +} diff --git a/internal/mtprotoedge/session_barrier_integration_test.go b/internal/mtprotoedge/session_barrier_integration_test.go new file mode 100644 index 00000000..e3f98e8e --- /dev/null +++ b/internal/mtprotoedge/session_barrier_integration_test.go @@ -0,0 +1,630 @@ +package mtprotoedge + +import ( + "context" + "crypto/rand" + "errors" + "io" + "sync" + "sync/atomic" + "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" + + "telesrv/internal/store" +) + +func encryptedRPCFrameForBarrierTest( + t *testing.T, + key crypto.AuthKey, + salt, sessionID, msgID int64, +) (*bin.Buffer, store.AuthKeyData) { + return encryptedRPCFrameWithAuthoritativeSaltForBarrierTest(t, key, salt, salt, sessionID, msgID, 1) +} + +func encryptedRPCFrameWithAuthoritativeSaltForBarrierTest( + t *testing.T, + key crypto.AuthKey, + frameSalt, authoritativeSalt, sessionID, msgID int64, + seqNo int32, +) (*bin.Buffer, store.AuthKeyData) { + t.Helper() + body := encodeClientMessageBodyForTest(t, &tg.HelpGetConfigRequest{}) + var frame bin.Buffer + if err := crypto.NewClientCipher(rand.Reader).Encrypt(key, crypto.EncryptedMessageData{ + Salt: frameSalt, + SessionID: sessionID, + MessageID: msgID, + SeqNo: seqNo, + MessageDataLen: int32(len(body)), + MessageDataWithPadding: body, + }, &frame); err != nil { + t.Fatalf("encrypt client RPC: %v", err) + } + return &frame, store.AuthKeyData{ + ID: key.ID, + Value: [256]byte(key.Value), + ServerSalt: authoritativeSalt, + } +} + +func TestBadServerSaltRetainsOneProvisionalConnUntilCorrected(t *testing.T) { + handler := &admissionCountingRPC{} + s := New(Options{RPC: handler, WriteTimeout: time.Second}) + s.rpcScheduler.start() + t.Cleanup(func() { s.rpcScheduler.stop(time.Second) }) + + tr := &collectingSessionTransport{} + key := newTestAuthKey(t) + const ( + serverSalt = int64(0x1020304050) + wrongSalt = serverSalt + 1 + sessionID = int64(71001) + ) + ids := proto.NewMessageIDGen(time.Now) + firstID := ids.New(proto.MessageFromClient) + firstWrong, stored := encryptedRPCFrameWithAuthoritativeSaltForBarrierTest( + t, key, wrongSalt, serverSalt, sessionID, firstID, 1, + ) + if err := s.authKeys.Save(context.Background(), stored); err != nil { + t.Fatalf("save auth key: %v", err) + } + cs := newConnState() + var plain bin.Buffer + + firstConn, err := s.handleEncrypted(context.Background(), tr, cs, nil, &stored, firstWrong, &plain) + if err != nil { + t.Fatalf("first bad salt: %v", err) + } + if firstConn == nil || firstConn.lifecycleState() != connLifecycleProvisional || firstConn.terminal.Load() { + t.Fatalf("first correction lifecycle conn=%p state=%v terminal=%v", firstConn, firstConn.lifecycleState(), firstConn.terminal.Load()) + } + if cs.createdFloor != 0 || len(cs.seen) != 0 || handler.calls.Load() != 0 { + t.Fatalf("bad salt admitted state: floor=%d seen=%d calls=%d", cs.createdFloor, len(cs.seen), handler.calls.Load()) + } + + secondID := ids.New(proto.MessageFromClient) + secondWrong, _ := encryptedRPCFrameWithAuthoritativeSaltForBarrierTest( + t, key, wrongSalt, serverSalt, sessionID, secondID, 3, + ) + secondConn, err := s.handleEncrypted(context.Background(), tr, cs, firstConn, nil, secondWrong, &plain) + if err != nil { + t.Fatalf("second bad salt: %v", err) + } + if secondConn != firstConn { + t.Fatalf("bad-salt retry replaced provisional Conn: first=%p second=%p", firstConn, secondConn) + } + if got := len(tr.snapshot()); got != 2 { + t.Fatalf("distinct bad msg ids produced %d corrections, want 2", got) + } + for i, wire := range tr.snapshot()[:2] { + data, decryptErr := crypto.NewClientCipher(rand.Reader).DecryptFromBuffer(key, &bin.Buffer{Buf: wire}) + if decryptErr != nil { + t.Fatalf("decrypt correction %d: %v", i, decryptErr) + } + if data.Salt != serverSalt { + t.Fatalf("correction %d envelope salt = %#x, want %#x", i, data.Salt, serverSalt) + } + var bad mt.BadServerSalt + body := &bin.Buffer{Buf: append([]byte(nil), data.Data()...)} + if decodeErr := bad.Decode(body); decodeErr != nil { + t.Fatalf("decode correction %d: %v", i, decodeErr) + } + if bad.NewServerSalt != serverSalt { + t.Fatalf("correction %d payload salt = %#x, want %#x", i, bad.NewServerSalt, serverSalt) + } + } + + corrected, _ := encryptedRPCFrameWithAuthoritativeSaltForBarrierTest( + t, key, serverSalt, serverSalt, sessionID, firstID, 1, + ) + activeConn, err := s.handleEncrypted(context.Background(), tr, cs, secondConn, nil, corrected, &plain) + if err != nil { + t.Fatalf("corrected retry: %v", err) + } + if activeConn != firstConn || !activeConn.isActive() { + t.Fatalf("corrected retry connection=%p first=%p active=%v", activeConn, firstConn, activeConn != nil && activeConn.isActive()) + } + if cs.createdFloor != firstID { + t.Fatalf("corrected created floor = %d, want %d", cs.createdFloor, firstID) + } + waitForAtomicCalls(t, &handler.calls, 1) + flightDeadline := time.Now().Add(2 * time.Second) + for s.rpcResults.flightLimit.snapshot() != 0 && time.Now().Before(flightDeadline) { + time.Sleep(time.Millisecond) + } + if got := s.rpcResults.flightLimit.snapshot(); got != 0 { + t.Fatalf("corrected retry leaked flight slots: %d", got) + } + activeConn.ForceClose() +} + +func TestWrongSaltSessionChangeTransfersPhysicalOwnership(t *testing.T) { + // The first handler remains pending until session replacement cancels it. Its + // owner Abort must observe the already-published terminal gate and must not + // close the physical lease that is about to transfer to the new session. + handler := &cancelThenRetryRPC{started: make(chan struct{})} + s := New(Options{RPC: handler, WriteTimeout: time.Second}) + s.rpcScheduler.start() + t.Cleanup(func() { s.rpcScheduler.stop(time.Second) }) + + tr := &collectingSessionTransport{} + key := newTestAuthKey(t) + const ( + serverSalt = int64(0x66778899) + wrongSalt = serverSalt + 7 + firstSID = int64(72001) + secondSID = int64(72002) + ) + ids := proto.NewMessageIDGen(time.Now) + firstID := ids.New(proto.MessageFromClient) + firstFrame, stored := encryptedRPCFrameWithAuthoritativeSaltForBarrierTest( + t, key, serverSalt, serverSalt, firstSID, firstID, 1, + ) + if err := s.authKeys.Save(context.Background(), stored); err != nil { + t.Fatalf("save auth key: %v", err) + } + cs := newConnState() + var plain bin.Buffer + oldConn, err := s.handleEncrypted(context.Background(), tr, cs, nil, &stored, firstFrame, &plain) + if err != nil { + t.Fatalf("activate first session: %v", err) + } + select { + case <-handler.started: + case <-time.After(2 * time.Second): + t.Fatal("first session RPC did not start") + } + + secondID := ids.New(proto.MessageFromClient) + wrongFrame, _ := encryptedRPCFrameWithAuthoritativeSaltForBarrierTest( + t, key, wrongSalt, serverSalt, secondSID, secondID, 1, + ) + newConn, err := s.handleEncrypted(context.Background(), tr, cs, oldConn, nil, wrongFrame, &plain) + if err != nil { + t.Fatalf("new session bad salt: %v", err) + } + if newConn == nil || newConn == oldConn || newConn.lifecycleState() != connLifecycleProvisional { + t.Fatalf("replacement conn old=%p new=%p state=%v", oldConn, newConn, newConn.lifecycleState()) + } + if !oldConn.terminal.Load() || tr.closed.Load() { + t.Fatalf("transfer state old_terminal=%v raw_closed=%v", oldConn.terminal.Load(), tr.closed.Load()) + } + // A delayed stale close must not tear down the generation already transferred + // to the new logical session. + oldConn.ForceClose() + if tr.closed.Load() || !newConn.isPhysicalTransportCurrentOpen() { + t.Fatalf("stale ForceClose closed replacement: raw_closed=%v current_open=%v", tr.closed.Load(), newConn.isPhysicalTransportCurrentOpen()) + } + + corrected, _ := encryptedRPCFrameWithAuthoritativeSaltForBarrierTest( + t, key, serverSalt, serverSalt, secondSID, secondID, 1, + ) + activated, err := s.handleEncrypted(context.Background(), tr, cs, newConn, nil, corrected, &plain) + if err != nil { + t.Fatalf("activate transferred session: %v", err) + } + if activated != newConn || !activated.isActive() { + t.Fatalf("transferred session active=%v conn=%p want=%p", activated != nil && activated.isActive(), activated, newConn) + } + waitForAtomicCalls(t, &handler.calls, 2) + activated.ForceClose() +} + +type collectingSessionTransport struct { + mu sync.Mutex + frames [][]byte + closed atomic.Bool +} + +func (t *collectingSessionTransport) Send(_ context.Context, b *bin.Buffer) error { + if t.closed.Load() { + return io.ErrClosedPipe + } + t.mu.Lock() + t.frames = append(t.frames, append([]byte(nil), b.Raw()...)) + t.mu.Unlock() + return nil +} + +func (*collectingSessionTransport) Recv(context.Context, *bin.Buffer) error { return io.EOF } +func (t *collectingSessionTransport) Close() error { + t.closed.Store(true) + return nil +} + +func (t *collectingSessionTransport) snapshot() [][]byte { + t.mu.Lock() + defer t.mu.Unlock() + out := make([][]byte, len(t.frames)) + for i := range t.frames { + out[i] = append([]byte(nil), t.frames[i]...) + } + return out +} + +type reconnectFlightRPC struct { + calls atomic.Int32 + started chan struct{} + release chan struct{} + once sync.Once +} + +func (h *reconnectFlightRPC) Dispatch(context.Context, [8]byte, int64, *bin.Buffer) (bin.Encoder, error) { + h.calls.Add(1) + h.once.Do(func() { close(h.started) }) + <-h.release // deliberately ignore replacement cancellation after committing work + return &tg.Config{ThisDC: 2}, nil +} + +func (*reconnectFlightRPC) NegotiatedLayer([8]byte, int64) (int, bool) { return 227, true } + +type cancelThenRetryRPC struct { + calls atomic.Int32 + active atomic.Int32 + max atomic.Int32 + started chan struct{} + once sync.Once +} + +func (h *cancelThenRetryRPC) Dispatch(ctx context.Context, _ [8]byte, _ int64, _ *bin.Buffer) (bin.Encoder, error) { + call := h.calls.Add(1) + active := h.active.Add(1) + defer h.active.Add(-1) + for { + max := h.max.Load() + if active <= max || h.max.CompareAndSwap(max, active) { + break + } + } + if call == 1 { + h.once.Do(func() { close(h.started) }) + <-ctx.Done() + return nil, ctx.Err() + } + return &tg.Config{ThisDC: 2}, nil +} + +func (*cancelThenRetryRPC) NegotiatedLayer([8]byte, int64) (int, bool) { return 227, true } + +func TestHandleEncryptedRequiredSessionBarrierPrecedesStateRegistrationAndRPC(t *testing.T) { + handler := &admissionCountingRPC{} + s := New(Options{RPC: handler, WriteTimeout: time.Second}) + s.rpcScheduler.start() + t.Cleanup(func() { s.rpcScheduler.stop(time.Second) }) + + tr := newGatedRequiredControlTransport(nil) + key := newTestAuthKey(t) + const ( + salt = int64(0x10203040) + sessionID = int64(70001) + ) + msgID := proto.NewMessageIDGen(time.Now).New(proto.MessageFromClient) + frame, stored := encryptedRPCFrameForBarrierTest(t, key, salt, sessionID, msgID) + if err := s.authKeys.Save(context.Background(), stored); err != nil { + t.Fatalf("save auth key: %v", err) + } + cs := newConnState() + var plain bin.Buffer + + type result struct { + conn *Conn + err error + } + done := make(chan result, 1) + go func() { + conn, err := s.handleEncrypted(context.Background(), tr, cs, nil, &stored, frame, &plain) + done <- result{conn: conn, err: err} + }() + + select { + case <-tr.started: + case <-time.After(time.Second): + t.Fatal("new_session_created did not reach the physical-write barrier") + } + if got := handler.calls.Load(); got != 0 { + t.Fatalf("handler calls while session barrier blocked = %d, want 0", got) + } + if cs.createdFloor != 0 || len(cs.seen) != 0 { + t.Fatalf("connState committed before session barrier: floor=%d seen=%d", cs.createdFloor, len(cs.seen)) + } + s.conns.mu.RLock() + registered := s.conns.bySession[sessionKey{authKeyID: key.ID, sessionID: sessionID}] + claim := s.conns.claims[sessionKey{authKeyID: key.ID, sessionID: sessionID}] + s.conns.mu.RUnlock() + if registered != nil || claim == nil { + t.Fatalf("blocked barrier visibility = registered:%p claim:%p, want nil/non-nil", registered, claim) + } + select { + case got := <-done: + t.Fatalf("handleEncrypted returned before physical barrier: %v", got.err) + default: + } + + tr.unblock() + var got result + select { + case got = <-done: + case <-time.After(2 * time.Second): + t.Fatal("handleEncrypted did not finish after session barrier write") + } + if got.err != nil { + t.Fatalf("handleEncrypted: %v", got.err) + } + if got.conn == nil || !got.conn.isActive() { + t.Fatalf("connection after successful barrier = %p lifecycle=%v", got.conn, got.conn.lifecycleState()) + } + if cs.createdFloor != msgID { + t.Fatalf("created floor after barrier = %d, want %d", cs.createdFloor, msgID) + } + waitForAtomicCalls(t, &handler.calls, 1) + s.conns.mu.RLock() + registered = s.conns.bySession[sessionKey{authKeyID: key.ID, sessionID: sessionID}] + s.conns.mu.RUnlock() + if registered != got.conn { + t.Fatalf("registered connection = %p, want %p", registered, got.conn) + } + got.conn.ForceClose() +} + +func TestHandleEncryptedRequiredSessionBarrierFailureIsAtomic(t *testing.T) { + handler := &admissionCountingRPC{} + s := New(Options{RPC: handler, WriteTimeout: time.Second}) + s.rpcScheduler.start() + t.Cleanup(func() { s.rpcScheduler.stop(time.Second) }) + + tr := newGatedRequiredControlTransport(io.ErrClosedPipe) + key := newTestAuthKey(t) + const ( + salt = int64(0x50607080) + sessionID = int64(70002) + ) + msgID := proto.NewMessageIDGen(time.Now).New(proto.MessageFromClient) + frame, stored := encryptedRPCFrameForBarrierTest(t, key, salt, sessionID, msgID) + if err := s.authKeys.Save(context.Background(), stored); err != nil { + t.Fatalf("save auth key: %v", err) + } + cs := newConnState() + var plain bin.Buffer + + done := make(chan error, 1) + go func() { + _, err := s.handleEncrypted(context.Background(), tr, cs, nil, &stored, frame, &plain) + done <- err + }() + select { + case <-tr.started: + case <-time.After(time.Second): + t.Fatal("failed new_session_created did not reach physical writer") + } + if handler.calls.Load() != 0 || cs.createdFloor != 0 || len(cs.seen) != 0 { + t.Fatalf("state changed while failing barrier blocked: calls=%d floor=%d seen=%d", handler.calls.Load(), cs.createdFloor, len(cs.seen)) + } + tr.unblock() + select { + case err := <-done: + if err == nil || !(errors.Is(err, io.ErrClosedPipe) || errors.Is(err, ErrConnClosed)) { + t.Fatalf("failed session barrier error = %v", err) + } + case <-time.After(2 * time.Second): + t.Fatal("failed session barrier did not return") + } + + if got := handler.calls.Load(); got != 0 { + t.Fatalf("handler ran after failed session barrier: %d", got) + } + if cs.createdFloor != 0 || len(cs.seen) != 0 || len(cs.order) != 0 { + t.Fatalf("failed barrier committed connState: floor=%d seen=%d order=%d", cs.createdFloor, len(cs.seen), len(cs.order)) + } + s.conns.mu.RLock() + registered := s.conns.bySession[sessionKey{authKeyID: key.ID, sessionID: sessionID}] + claim := s.conns.claims[sessionKey{authKeyID: key.ID, sessionID: sessionID}] + s.conns.mu.RUnlock() + if registered != nil || claim != nil { + t.Fatalf("failed barrier leaked manager ownership: registered=%p claim=%p", registered, claim) + } + if tasks, bytes := s.rpcScheduler.budgetSnapshot(); tasks != 0 || bytes != 0 { + t.Fatalf("failed barrier leaked RPC budget: tasks=%d bytes=%d", tasks, bytes) + } + if got := s.frameBudget.usedBytes(); got != 0 { + t.Fatalf("failed barrier leaked frame budget: %d", got) + } + if got := s.rpcResults.flightLimit.snapshot(); got != 0 { + t.Fatalf("failed barrier leaked RPC claim slots: %d", got) + } +} + +func TestCrossConnectionInflightRPCHasOneBusinessOwnerAndReplaysResult(t *testing.T) { + handler := &reconnectFlightRPC{started: make(chan struct{}), release: make(chan struct{})} + s := New(Options{RPC: handler, WriteTimeout: time.Second, RPCTimeout: 5 * time.Second}) + s.rpcScheduler.start() + t.Cleanup(func() { s.rpcScheduler.stop(time.Second) }) + + key := newTestAuthKey(t) + const ( + salt = int64(0x11223344) + sessionID = int64(70003) + ) + msgID := proto.NewMessageIDGen(time.Now).New(proto.MessageFromClient) + firstFrame, stored := encryptedRPCFrameForBarrierTest(t, key, salt, sessionID, msgID) + if err := s.authKeys.Save(context.Background(), stored); err != nil { + t.Fatalf("save auth key: %v", err) + } + firstTransport := &collectingSessionTransport{} + firstState := newConnState() + var firstPlain bin.Buffer + firstConn, err := s.handleEncrypted(context.Background(), firstTransport, firstState, nil, &stored, firstFrame, &firstPlain) + if err != nil { + t.Fatalf("first handleEncrypted: %v", err) + } + select { + case <-handler.started: + case <-time.After(2 * time.Second): + t.Fatal("first RPC owner did not start") + } + + secondFrame, _ := encryptedRPCFrameForBarrierTest(t, key, salt, sessionID, msgID) + secondTransport := &collectingSessionTransport{} + secondState := newConnState() + var secondPlain bin.Buffer + type handleResult struct { + conn *Conn + err error + } + secondDone := make(chan handleResult, 1) + go func() { + conn, handleErr := s.handleEncrypted(context.Background(), secondTransport, secondState, nil, &stored, secondFrame, &secondPlain) + secondDone <- handleResult{conn: conn, err: handleErr} + }() + + deadline := time.Now().Add(2 * time.Second) + for !firstConn.terminal.Load() && time.Now().Before(deadline) { + time.Sleep(time.Millisecond) + } + if !firstConn.terminal.Load() { + t.Fatal("replacement did not fence the first physical connection") + } + if got := handler.calls.Load(); got != 1 { + t.Fatalf("overlapping reconnect executed %d business handlers, want 1", got) + } + select { + case got := <-secondDone: + t.Fatalf("duplicate reconnect returned before owner completion: %v", got.err) + default: + } + + close(handler.release) + var second handleResult + select { + case second = <-secondDone: + case <-time.After(3 * time.Second): + t.Fatal("duplicate reconnect did not receive owner result") + } + if second.err != nil { + t.Fatalf("second handleEncrypted: %v", second.err) + } + if second.conn == nil || !second.conn.isActive() { + t.Fatalf("second connection was not active after replay: %p", second.conn) + } + if got := handler.calls.Load(); got != 1 { + t.Fatalf("cross-connection duplicate business calls = %d, want 1", got) + } + + resultCount := 0 + for _, wire := range secondTransport.snapshot() { + data, decryptErr := crypto.NewClientCipher(rand.Reader).DecryptFromBuffer(key, &bin.Buffer{Buf: wire}) + if decryptErr != nil { + t.Fatalf("decrypt second connection reply: %v", decryptErr) + } + plain := &bin.Buffer{Buf: append([]byte(nil), data.Data()...)} + typeID, peekErr := plain.PeekID() + if peekErr != nil { + t.Fatalf("peek second connection reply: %v", peekErr) + } + if typeID != proto.ResultTypeID { + continue + } + var result proto.Result + if decodeErr := result.Decode(plain); decodeErr != nil { + t.Fatalf("decode replayed rpc_result: %v", decodeErr) + } + if result.RequestMessageID != msgID { + t.Fatalf("replayed rpc_result request id = %d, want %d", result.RequestMessageID, msgID) + } + resultCount++ + } + if resultCount != 1 { + t.Fatalf("replayed rpc_result count on replacement = %d, want 1", resultCount) + } + waitInboundRPCBatchBudget(t, s.rpcScheduler, 0, 0) + if got := s.rpcResults.flightLimit.snapshot(); got != 0 { + t.Fatalf("completed reconnect leaked flight claims: %d", got) + } + second.conn.ForceClose() +} + +func TestCrossConnectionInflightAbortRetriesOnlyAfterOldOwnerStops(t *testing.T) { + handler := &cancelThenRetryRPC{started: make(chan struct{})} + s := New(Options{RPC: handler, WriteTimeout: time.Second, RPCTimeout: 5 * time.Second}) + s.rpcScheduler.start() + t.Cleanup(func() { s.rpcScheduler.stop(time.Second) }) + + key := newTestAuthKey(t) + const ( + salt = int64(0x55667788) + sessionID = int64(70004) + ) + msgID := proto.NewMessageIDGen(time.Now).New(proto.MessageFromClient) + firstFrame, stored := encryptedRPCFrameForBarrierTest(t, key, salt, sessionID, msgID) + if err := s.authKeys.Save(context.Background(), stored); err != nil { + t.Fatalf("save auth key: %v", err) + } + firstTransport := &collectingSessionTransport{} + firstState := newConnState() + var firstPlain bin.Buffer + firstConn, err := s.handleEncrypted(context.Background(), firstTransport, firstState, nil, &stored, firstFrame, &firstPlain) + if err != nil { + t.Fatalf("first handleEncrypted: %v", err) + } + select { + case <-handler.started: + case <-time.After(2 * time.Second): + t.Fatal("first cancellation-aware owner did not start") + } + + secondFrame, _ := encryptedRPCFrameForBarrierTest(t, key, salt, sessionID, msgID) + secondTransport := &collectingSessionTransport{} + secondState := newConnState() + var secondPlain bin.Buffer + secondConn, err := s.handleEncrypted(context.Background(), secondTransport, secondState, nil, &stored, secondFrame, &secondPlain) + if err != nil { + t.Fatalf("second handleEncrypted: %v", err) + } + waitForAtomicCalls(t, &handler.calls, 2) + waitInboundRPCBatchBudget(t, s.rpcScheduler, 0, 0) + if got := handler.max.Load(); got != 1 { + t.Fatalf("old and retry handlers overlapped: max active=%d, want 1", got) + } + if got := s.rpcResults.flightLimit.snapshot(); got != 0 { + t.Fatalf("sequential retry leaked flight claims: %d", got) + } + + resultCount := 0 + for _, wire := range secondTransport.snapshot() { + data, decryptErr := crypto.NewClientCipher(rand.Reader).DecryptFromBuffer(key, &bin.Buffer{Buf: wire}) + if decryptErr != nil { + t.Fatalf("decrypt sequential-retry reply: %v", decryptErr) + } + plain := &bin.Buffer{Buf: append([]byte(nil), data.Data()...)} + typeID, peekErr := plain.PeekID() + if peekErr != nil { + t.Fatalf("peek sequential-retry reply: %v", peekErr) + } + if typeID == proto.ResultTypeID { + resultCount++ + } + } + if resultCount != 1 { + t.Fatalf("sequential retry result count = %d, want 1", resultCount) + } + if !firstConn.terminal.Load() || secondConn == nil || !secondConn.isActive() { + t.Fatalf("replacement lifecycle = old terminal:%v new:%p active:%v", firstConn.terminal.Load(), secondConn, secondConn != nil && secondConn.isActive()) + } + secondConn.ForceClose() +} + +func waitForAtomicCalls(t *testing.T, calls interface{ Load() int32 }, want int32) { + t.Helper() + deadline := time.Now().Add(2 * time.Second) + for calls.Load() != want && time.Now().Before(deadline) { + time.Sleep(time.Millisecond) + } + if got := calls.Load(); got != want { + t.Fatalf("handler calls = %d, want %d", got, want) + } +} diff --git a/internal/mtprotoedge/session_boundary_test.go b/internal/mtprotoedge/session_boundary_test.go new file mode 100644 index 00000000..63885f3c --- /dev/null +++ b/internal/mtprotoedge/session_boundary_test.go @@ -0,0 +1,383 @@ +package mtprotoedge + +import ( + "context" + "strings" + "sync/atomic" + "testing" + "time" + + "github.com/gotd/td/bin" + "github.com/gotd/td/mt" + "github.com/gotd/td/proto" + "github.com/gotd/td/tg" +) + +// TestFirstContainerBoundaryKeepsAndroidRequestMap models DrKLO's +// new_session_created handling: it drops every running request whose msg_id is +// lower than first_msg_id. Every inner request accepted from the first +// container must therefore remain addressable when its pong arrives. +func TestFirstContainerBoundaryKeepsAndroidRequestMap(t *testing.T) { + const dc = 2 + addr, pub, _ := startTestServer(t, Options{DC: dc}) + conn, auth, cipher := dialHandshake(t, addr, dc, pub) + + msgIDs := proto.NewMessageIDGen(time.Now) + running := make(map[int64]int64, 3) + messages := make([]proto.Message, 0, 3) + for i := 0; i < 3; i++ { + msgID := msgIDs.New(proto.MessageFromClient) + pingID := int64(10_001 + i) + body := mustEncodeTL(t, &mt.PingRequest{PingID: pingID}) + messages = append(messages, proto.Message{ + ID: msgID, + SeqNo: 1 + i*2, + Bytes: len(body), + Body: body, + }) + running[msgID] = pingID + } + outerMsgID := msgIDs.New(proto.MessageFromClient) + sendEncrypted(t, conn, cipher, auth, outerMsgID, &proto.MessageContainer{Messages: messages}) + + boundaryFrames := collectReplyFrames(t, conn, cipher, auth.AuthKey, map[uint32]int{ + mt.NewSessionCreatedTypeID: 1, + }) + var created mt.NewSessionCreated + for _, frame := range boundaryFrames { + if frame.TypeID == mt.PongTypeID { + t.Fatalf("pong arrived before new_session_created boundary") + } + if frame.TypeID != mt.NewSessionCreatedTypeID { + continue + } + if err := created.Decode(frame.Plain); err != nil { + t.Fatalf("decode new_session_created: %v", err) + } + } + + // DrKLO ConnectionsManager.cpp clears running requests with + // request.messageId < first_msg_id when it receives this notification. + for msgID := range running { + if msgID < created.FirstMsgID { + delete(running, msgID) + } + } + for _, accepted := range messages { + if _, ok := running[accepted.ID]; !ok { + t.Fatalf( + "accepted inner msg_id %d was evicted by Android boundary %d (outer=%d)", + accepted.ID, + created.FirstMsgID, + outerMsgID, + ) + } + } + + pongFrames := collectReplyFrames(t, conn, cipher, auth.AuthKey, map[uint32]int{ + mt.PongTypeID: len(messages), + }) + seen := make(map[int64]struct{}, len(messages)) + for _, frame := range pongFrames { + if frame.TypeID != mt.PongTypeID { + continue + } + var pong mt.Pong + if err := pong.Decode(frame.Plain); err != nil { + t.Fatalf("decode pong: %v", err) + } + wantPingID, ok := running[pong.MsgID] + if !ok { + t.Fatalf("orphan pong for msg_id %d after Android boundary cleanup", pong.MsgID) + } + if pong.PingID != wantPingID { + t.Fatalf("pong ping_id = %d for msg_id %d, want %d", pong.PingID, pong.MsgID, wantPingID) + } + if _, duplicate := seen[pong.MsgID]; duplicate { + t.Fatalf("duplicate pong for msg_id %d", pong.MsgID) + } + seen[pong.MsgID] = struct{}{} + delete(running, pong.MsgID) + } + if len(running) != 0 { + t.Fatalf("requests without correlated pong: %+v", running) + } +} + +type admissionCountingRPC struct { + calls atomic.Int32 +} + +func (h *admissionCountingRPC) Dispatch(_ context.Context, _ [8]byte, _ int64, _ *bin.Buffer) (bin.Encoder, error) { + h.calls.Add(1) + return &tg.Config{ThisDC: 2}, nil +} + +func (*admissionCountingRPC) NegotiatedLayer([8]byte, int64) (int, bool) { return 227, true } + +func TestContainerRPCAdmissionFailureIsAtomic(t *testing.T) { + const dc = 2 + handler := &admissionCountingRPC{} + addr, pub, _ := startTestServer(t, Options{ + DC: dc, + RPC: handler, + RPCMaxInflight: 1, + RPCQueueSize: 1, + RPCGlobalWorkers: 1, + }) + conn, auth, cipher := dialHandshake(t, addr, dc, pub) + + msgIDs := proto.NewMessageIDGen(time.Now) + messages := make([]proto.Message, 0, 2) + requestIDs := make(map[int64]struct{}, 2) + for i := 0; i < 2; i++ { + msgID := msgIDs.New(proto.MessageFromClient) + body := mustEncodeTL(t, &tg.HelpGetConfigRequest{}) + messages = append(messages, proto.Message{ID: msgID, SeqNo: 1 + i*2, Bytes: len(body), Body: body}) + requestIDs[msgID] = struct{}{} + } + outerMsgID := msgIDs.New(proto.MessageFromClient) + sendEncrypted(t, conn, cipher, auth, outerMsgID, &proto.MessageContainer{Messages: messages}) + + frames := collectReplyFrames(t, conn, cipher, auth.AuthKey, map[uint32]int{proto.ResultTypeID: 2}) + results := 0 + for _, frame := range frames { + if frame.TypeID != proto.ResultTypeID { + continue + } + var result proto.Result + if err := result.Decode(frame.Plain); err != nil { + t.Fatalf("decode rpc_result: %v", err) + } + if _, ok := requestIDs[result.RequestMessageID]; !ok { + t.Fatalf("unexpected or duplicate capacity rpc_result req_msg_id %d", result.RequestMessageID) + } + var rpcErr mt.RPCError + if err := rpcErr.Decode(&bin.Buffer{Buf: result.Result}); err != nil { + t.Fatalf("decode capacity rpc_error: %v", err) + } + if rpcErr.ErrorCode != 420 || rpcErr.ErrorMessage != "FLOOD_WAIT_1" { + t.Fatalf("capacity rpc_error = %+v", rpcErr) + } + delete(requestIDs, result.RequestMessageID) + results++ + } + if results != 2 { + t.Fatalf("capacity rpc_results = %d, want 2", results) + } + if len(requestIDs) != 0 { + t.Fatalf("capacity requests without exactly one result: %+v", requestIDs) + } + if got := handler.calls.Load(); got != 0 { + t.Fatalf("partially executed handler calls = %d, want 0", got) + } +} + +func TestGZIPWrappedRPCUsesLogicalEnvelopeBoundary(t *testing.T) { + const dc = 2 + handler := &admissionCountingRPC{} + addr, pub, _ := startTestServer(t, Options{DC: dc, RPC: handler}) + conn, auth, cipher := dialHandshake(t, addr, dc, pub) + + requestBody := mustEncodeTL(t, &tg.HelpGetConfigRequest{}) + msgID := proto.NewMessageIDGen(time.Now).New(proto.MessageFromClient) + sendEncryptedWithSeq(t, conn, cipher, auth, msgID, 1, &proto.GZIP{Data: requestBody}) + frames := collectReplyFrames(t, conn, cipher, auth.AuthKey, map[uint32]int{proto.ResultTypeID: 1}) + var boundary mt.NewSessionCreated + for _, frame := range frames { + if frame.TypeID == mt.NewSessionCreatedTypeID { + if err := boundary.Decode(frame.Plain); err != nil { + t.Fatalf("decode gzip RPC boundary: %v", err) + } + } + } + if boundary.FirstMsgID != msgID { + t.Fatalf("gzip RPC boundary = %d, want envelope %d", boundary.FirstMsgID, msgID) + } + if got := handler.calls.Load(); got != 1 { + t.Fatalf("gzip RPC handler calls = %d, want 1", got) + } +} + +func TestGZIPWrappedContainerUsesLowestInnerBoundary(t *testing.T) { + const dc = 2 + addr, pub, _ := startTestServer(t, Options{DC: dc}) + conn, auth, cipher := dialHandshake(t, addr, dc, pub) + + ids := proto.NewMessageIDGen(time.Now) + innerMsgID := ids.New(proto.MessageFromClient) + outerMsgID := ids.New(proto.MessageFromClient) + pingBody := mustEncodeTL(t, &mt.PingRequest{PingID: 7001}) + containerBody := mustEncodeTL(t, &proto.MessageContainer{Messages: []proto.Message{{ + ID: innerMsgID, SeqNo: 1, Bytes: len(pingBody), Body: pingBody, + }}}) + sendEncryptedWithSeq(t, conn, cipher, auth, outerMsgID, 2, &proto.GZIP{Data: containerBody}) + frames := collectReplyFrames(t, conn, cipher, auth.AuthKey, map[uint32]int{mt.PongTypeID: 1}) + var boundary mt.NewSessionCreated + for _, frame := range frames { + if frame.TypeID == mt.NewSessionCreatedTypeID { + if err := boundary.Decode(frame.Plain); err != nil { + t.Fatalf("decode gzip container boundary: %v", err) + } + } + } + if boundary.FirstMsgID != innerMsgID { + t.Fatalf("gzip container boundary = %d, want inner %d (outer %d)", boundary.FirstMsgID, innerMsgID, outerMsgID) + } +} + +func TestEmptyContainerUsesOuterBoundary(t *testing.T) { + const dc = 2 + addr, pub, _ := startTestServer(t, Options{DC: dc}) + conn, auth, cipher := dialHandshake(t, addr, dc, pub) + + outerMsgID := proto.NewMessageIDGen(time.Now).New(proto.MessageFromClient) + sendEncryptedWithSeq(t, conn, cipher, auth, outerMsgID, 0, &proto.MessageContainer{}) + frames := collectReplyFrames(t, conn, cipher, auth.AuthKey, map[uint32]int{mt.NewSessionCreatedTypeID: 1}) + var boundary mt.NewSessionCreated + for _, frame := range frames { + if frame.TypeID != mt.NewSessionCreatedTypeID { + continue + } + if err := boundary.Decode(frame.Plain); err != nil { + t.Fatalf("decode empty container boundary: %v", err) + } + } + if boundary.FirstMsgID != outerMsgID { + t.Fatalf("empty container boundary = %d, want outer %d", boundary.FirstMsgID, outerMsgID) + } +} + +func TestDuplicateContainerAcksWithoutBusinessReexecution(t *testing.T) { + const dc = 2 + handler := &admissionCountingRPC{} + addr, pub, _ := startTestServer(t, Options{DC: dc, RPC: handler}) + conn, auth, cipher := dialHandshake(t, addr, dc, pub) + + ids := proto.NewMessageIDGen(time.Now) + messages := make([]proto.Message, 0, 2) + for i := 0; i < 2; i++ { + body := mustEncodeTL(t, &tg.HelpGetConfigRequest{}) + messages = append(messages, proto.Message{ + ID: ids.New(proto.MessageFromClient), SeqNo: 1 + i*2, Bytes: len(body), Body: body, + }) + } + outerMsgID := ids.New(proto.MessageFromClient) + container := &proto.MessageContainer{Messages: messages} + sendEncrypted(t, conn, cipher, auth, outerMsgID, container) + initialFrames := collectReplyFrames(t, conn, cipher, auth.AuthKey, map[uint32]int{proto.ResultTypeID: 2}) + answerIDs := make([]int64, 0, 2) + for _, frame := range initialFrames { + if frame.TypeID == proto.ResultTypeID { + answerIDs = append(answerIDs, frame.Message.MessageID) + } + } + if len(answerIDs) != 2 { + t.Fatalf("initial rpc_result answer ids = %d, want 2", len(answerIDs)) + } + // ACK both original server results before retransmitting the container. This is + // a wire barrier: any later rpc_result is necessarily a duplicate replay rather + // than an unconsumed response from the initial batch. + ackMsgID := ids.New(proto.MessageFromClient) + sendEncryptedWithSeq(t, conn, cipher, auth, ackMsgID, 4, &mt.MsgsAck{MsgIDs: answerIDs}) + if got := handler.calls.Load(); got != 2 { + t.Fatalf("initial handler calls = %d, want 2", got) + } + + sendEncrypted(t, conn, cipher, auth, outerMsgID, container) + duplicateFrames := collectReplyFrames(t, conn, cipher, auth.AuthKey, map[uint32]int{mt.MsgsAckTypeID: 1}) + for _, frame := range duplicateFrames { + if frame.TypeID == proto.ResultTypeID { + t.Fatal("same-connection duplicate container replayed an extra rpc_result") + } + } + if got := handler.calls.Load(); got != 2 { + t.Fatalf("duplicate container reexecuted handlers: calls=%d, want 2", got) + } +} + +type largeStartupBurstRPC struct { + calls atomic.Int32 + body string +} + +func (h *largeStartupBurstRPC) Dispatch(ctx context.Context, _ [8]byte, _ int64, _ *bin.Buffer) (bin.Encoder, error) { + call := h.calls.Add(1) + // Deliberately perturb completion order while remaining cancellation-aware. + timer := time.NewTimer(time.Duration(call%4) * time.Millisecond) + defer timer.Stop() + select { + case <-timer.C: + case <-ctx.Done(): + return nil, ctx.Err() + } + return &tg.DataJSON{Data: h.body}, nil +} + +func (*largeStartupBurstRPC) NegotiatedLayer([8]byte, int64) (int, bool) { return 227, true } + +func TestAndroidStartupBurstKeepsAllLargeRPCResultsAddressable(t *testing.T) { + const ( + dc = 2 + requests = 30 + ) + handler := &largeStartupBurstRPC{body: strings.Repeat("x", 192<<10)} + addr, pub, _ := startTestServer(t, Options{DC: dc, RPC: handler}) + conn, auth, cipher := dialHandshake(t, addr, dc, pub) + + ids := proto.NewMessageIDGen(time.Now) + running := make(map[int64]struct{}, requests) + messages := make([]proto.Message, 0, requests) + for i := 0; i < requests; i++ { + msgID := ids.New(proto.MessageFromClient) + body := mustEncodeTL(t, &tg.HelpGetConfigRequest{}) + messages = append(messages, proto.Message{ + ID: msgID, SeqNo: 1 + i*2, Bytes: len(body), Body: body, + }) + running[msgID] = struct{}{} + } + outerMsgID := ids.New(proto.MessageFromClient) + sendEncrypted(t, conn, cipher, auth, outerMsgID, &proto.MessageContainer{Messages: messages}) + + boundaryFrames := collectReplyFrames(t, conn, cipher, auth.AuthKey, map[uint32]int{mt.NewSessionCreatedTypeID: 1}) + var boundary mt.NewSessionCreated + for _, frame := range boundaryFrames { + if frame.TypeID == proto.ResultTypeID { + t.Fatal("large rpc_result arrived before session boundary") + } + if frame.TypeID == mt.NewSessionCreatedTypeID { + if err := boundary.Decode(frame.Plain); err != nil { + t.Fatalf("decode startup boundary: %v", err) + } + } + } + for msgID := range running { + if msgID < boundary.FirstMsgID { + delete(running, msgID) + } + } + if len(running) != requests { + t.Fatalf("Android boundary removed accepted startup requests: kept=%d want=%d floor=%d outer=%d", len(running), requests, boundary.FirstMsgID, outerMsgID) + } + + resultFrames := collectReplyFrames(t, conn, cipher, auth.AuthKey, map[uint32]int{proto.ResultTypeID: requests}) + for _, frame := range resultFrames { + if frame.TypeID != proto.ResultTypeID { + continue + } + var result proto.Result + if err := result.Decode(frame.Plain); err != nil { + t.Fatalf("decode startup rpc_result: %v", err) + } + if _, ok := running[result.RequestMessageID]; !ok { + t.Fatalf("orphan large rpc_result for request %d", result.RequestMessageID) + } + delete(running, result.RequestMessageID) + } + if len(running) != 0 { + t.Fatalf("startup RPCs without result: %d", len(running)) + } + if got := handler.calls.Load(); got != requests { + t.Fatalf("startup handler calls = %d, want %d", got, requests) + } +} diff --git a/internal/mtprotoedge/session_manager.go b/internal/mtprotoedge/session_manager.go index 3d172b73..fc15e676 100644 --- a/internal/mtprotoedge/session_manager.go +++ b/internal/mtprotoedge/session_manager.go @@ -22,6 +22,11 @@ var ErrSessionNotFound = errors.New("session not found") // ErrSessionAmbiguous 表示仅用 session_id 无法唯一定位连接。 var ErrSessionAmbiguous = errors.New("session id is shared by multiple auth keys") +var ( + ErrSessionActivationSuperseded = errors.New("session activation superseded") + ErrSessionActivationFence = errors.New("session activation could not fence previous writer") +) + const ( maxPendingPushesPerSession = 32 // maxFlushAttempts / flushRetryBackoff：排空暂存推送时 c.Send 失败（出站拥塞 5s 超时 @@ -117,8 +122,13 @@ type SessionLifecycleObserver interface { // 它管理运行态的在线连接，与持久化的 store.SessionStore 互补：后者记录 session 数据， // 前者持有可发送的活跃连接。所有方法并发安全。 type SessionManager struct { - mu sync.RWMutex - bySession map[sessionKey]*Conn + mu sync.RWMutex + bySession map[sessionKey]*Conn + // claims owns the provisional -> active gap. A claimant is intentionally + // absent from every push/online index until its required session control frame + // is on the wire and PublishActivation validates the same owner. + claims map[sessionKey]*Conn + claimsByAuth map[[8]byte]map[int64]*Conn // raw authKeyID -> sessionID -> provisional claim bySessionID map[int64]map[[8]byte]*Conn // sessionID → raw authKeyID → Conn，用于兼容旧 API 的唯一性检查 byAuthKey map[[8]byte]map[int64]*Conn // raw authKeyID → sessionID → Conn byBusinessAuthKey map[[8]byte]map[sessionKey]*Conn @@ -142,6 +152,8 @@ func NewSessionManager(log *zap.Logger) *SessionManager { } return &SessionManager{ bySession: make(map[sessionKey]*Conn), + claims: make(map[sessionKey]*Conn), + claimsByAuth: make(map[[8]byte]map[int64]*Conn), bySessionID: make(map[int64]map[[8]byte]*Conn), byAuthKey: make(map[[8]byte]map[int64]*Conn), byBusinessAuthKey: make(map[[8]byte]map[sessionKey]*Conn), @@ -164,32 +176,94 @@ func (m *SessionManager) SetLifecycleObserver(observer SessionLifecycleObserver) m.mu.Unlock() } -// Register 注册一个活跃连接。若同 raw auth_key_id + session_id 已存在（重连），旧连接被替换并移除索引。 -func (m *SessionManager) Register(c *Conn) { - m.mu.Lock() +// BeginActivation atomically claims auth_key_id + session_id without publishing the +// new Conn. Under the manager lock it irreversibly fences every previous owner, +// removes active indexes and closes producer/RPC admission gates. Physical close and +// outbound-actor convergence happen outside the lock; the caller may send the +// required new_session_created frame only after this method returns nil. +func (m *SessionManager) BeginActivation(c *Conn) error { + if c == nil || !c.beginActivationClaim() { + return ErrSessionActivationSuperseded + } key := connSessionKey(c) - var replaced *Conn - var evicted *Conn - if old, ok := m.bySession[key]; ok && old != c { - replaced = old - m.removeLocked(old, false) - } else if existing := m.byAuthKey[c.authKeyID]; len(existing) >= maxSessionsPerAuthKey { - // 同 raw auth_key 的 session 数达上限且本次是新 session：驱逐建连最早的 session 让位， - // 防对抗客户端用海量 session_id 撑爆索引。驱逐对象与新连接同属一个设备凭据， - // 触顶基本是该凭据自身异常；选最旧而非 map 随机，避免误杀刚建立的活跃下载/主连接。 - // 被驱逐连接的 serveConn 会在下一帧因 actor 已关而退出。O(cap) 扫描仅在触顶时发生。 - for _, ec := range existing { - if evicted == nil || ec.createdAt.Before(evicted.createdAt) { - evicted = ec - } + retired := make([]*Conn, 0, 2) + m.mu.Lock() + if c.terminal.Load() || !c.isPhysicalTransportCurrentOpen() || c.lifecycleState() != connLifecycleClaiming { + c.beginTerminalShutdown() + m.mu.Unlock() + return ErrConnClosed + } + if oldClaim := m.claims[key]; oldClaim != nil && oldClaim != c { + m.retireClaimLocked(key, oldClaim, false) + retired = append(retired, oldClaim) + } + if old := m.bySession[key]; old != nil && old != c { + m.retireConnLocked(old, false) + retired = append(retired, old) + } + + // Claims reserve a cap slot just like published sessions. Otherwise many + // concurrent handshakes could all pass the old byAuthKey-only check and publish + // beyond maxSessionsPerAuthKey. + for len(m.byAuthKey[c.authKeyID])+m.claimCountForAuthLocked(c.authKeyID) >= maxSessionsPerAuthKey { + victimKey, victim, isClaim := m.oldestAuthOwnerLocked(c.authKeyID, c) + if victim == nil { + break } - m.removeLocked(evicted, true) - m.log.Debug("Evicted oldest session for auth key at cap", + if isClaim { + m.retireClaimLocked(victimKey, victim, true) + } else { + m.retireConnLocked(victim, true) + } + retired = append(retired, victim) + m.log.Debug("Evicted oldest session activation for auth key at cap", zap.String("auth_key_id", sessionKeyLog(c.authKeyID)), zap.Int("cap", maxSessionsPerAuthKey), ) } + m.addClaimLocked(key, c) + m.mu.Unlock() + + // Do not wait for old business handlers: their root context/admission gate is + // already canceled. We only need the old physical writer and outbound actor to + // converge before the new Conn is allowed to write the session barrier. + if !closeConnBatch(retired, forceCloseBatchTimeout, false) { + m.AbortActivation(c) + return ErrSessionActivationFence + } + return nil +} + +// PublishActivation makes the current claim visible to push/online lookups. It is +// deliberately a separate operation from BeginActivation so required protocol +// control can be written while the Conn remains provisional and unindexed. +func (m *SessionManager) PublishActivation(c *Conn) error { + if c == nil { + return ErrSessionActivationSuperseded + } + key := connSessionKey(c) + m.mu.Lock() + defer m.mu.Unlock() + if m.claims[key] != c { + return ErrSessionActivationSuperseded + } + if c.terminal.Load() || c.lifecycleState() != connLifecycleClaiming { + m.removeClaimLocked(key, c) + return ErrConnClosed + } + if old := m.bySession[key]; old != nil && old != c { + // An active owner without this claim can only be a stale/unsafe publisher. + // Never reverse-replace it from here; the caller must reconnect and claim again. + m.removeClaimLocked(key, c) + c.beginTerminalShutdown() + return ErrSessionActivationSuperseded + } + if !c.publishActivation() { + m.removeClaimLocked(key, c) + return ErrConnClosed + } + m.removeClaimLocked(key, c) m.bySession[key] = c addSessionIDIndex(m.bySessionID, c.sessionID, c.authKeyID, c) addConnIndex(m.byAuthKey, c.authKeyID, c.sessionID, c) @@ -200,34 +274,88 @@ func (m *SessionManager) Register(c *Conn) { c.userIDResolved.Store(true) addUserIndex(m.byUser, uid, key, c) } - m.log.Debug("Session registered", + m.log.Debug("Session activated", zap.String("auth_key_id", sessionKeyLog(key.authKeyID)), zap.Int64("session_id", c.sessionID), zap.Int("online", len(m.bySession)), ) - m.mu.Unlock() + return nil +} - // 同 identity 的新物理连接已经原子接管索引；立即关闭旧 transport，不能只停 - // actor 后让旧 FD/read goroutine 滞留到 read timeout。replacement 与 cap eviction - // 共用一个并发关闭批次，不能把每条 Conn 的 RPC 等待上界串行相加。 - if replaced != nil || evicted != nil { - if !forceCloseConnBatch([]*Conn{replaced, evicted}, forceCloseBatchTimeout) { - m.log.Warn("Session replacement/eviction close exceeded shared deadline") - } +// AbortActivation removes a claim only when c still owns it, then terminally +// retires the provisional Conn. A superseded caller cannot delete the newer claim. +func (m *SessionManager) AbortActivation(c *Conn) { + if c == nil { + return } + key := connSessionKey(c) + owned := false + m.mu.Lock() + if m.claims[key] == c { + c.beginTerminalShutdown() + m.removeClaimLocked(key, c) + if m.bySession[key] == nil { + m.deletePendingLocked(key) + delete(m.flushing, key) + } + owned = true + } + m.mu.Unlock() + if owned { + _ = closeConnBatch([]*Conn{c}, forceCloseBatchTimeout, false) + } +} + +// Register is retained for tests and embedders that do not have a wire-level +// required-control barrier. Production encrypted traffic uses the explicit +// BeginActivation -> SendRequiredControl -> PublishActivation sequence. +func (m *SessionManager) Register(c *Conn) error { + if c == nil { + return ErrSessionActivationSuperseded + } + if c.isActive() { + m.mu.RLock() + current := m.bySession[connSessionKey(c)] + m.mu.RUnlock() + if current == c { + return nil + } + return ErrSessionActivationSuperseded + } + if err := m.BeginActivation(c); err != nil { + return err + } + if err := m.PublishActivation(c); err != nil { + m.AbortActivation(c) + return err + } + return nil } // Unregister 注销一个连接（仅当它仍是当前注册的同一对象，避免误删重连后的新连接）。 // 观察者对未登录连接（userID=0）也回调：业务层据此清理按 session 维度的缓存条目， // 否则未登录连接的元数据只能等容量上限驱逐。 func (m *SessionManager) Unregister(c *Conn) { + if c == nil { + return + } + // Close admission/outbound producer gates before removing indexes or invoking + // a lifecycle observer. An observer may block, but no old RPC/push may continue + // to write or enqueue work during that interval. + c.beginTerminalShutdown() m.mu.Lock() var ( observer SessionLifecycleObserver offlineUser int64 lastForUser bool ) - if cur, ok := m.bySession[connSessionKey(c)]; ok && cur == c { + key := connSessionKey(c) + if m.claims[key] == c { + m.removeClaimLocked(key, c) + m.deletePendingLocked(key) + delete(m.flushing, key) + } + if cur, ok := m.bySession[key]; ok && cur == c { offlineUser = m.removeLocked(c, true) if offlineUser != 0 { lastForUser = len(m.byUser[offlineUser]) == 0 @@ -256,7 +384,7 @@ func (m *SessionManager) DestroySession(sessionID int64) bool { m.mu.Unlock() return false } - offlineUser := m.removeLocked(c, true) + offlineUser := m.retireConnLocked(c, true) lastForUser := offlineUser != 0 && len(m.byUser[offlineUser]) == 0 observer := m.lifecycle m.log.Debug("Session destroyed", @@ -283,11 +411,22 @@ func (m *SessionManager) DestroySessionForAuthKey(authKeyID [8]byte, sessionID i key := sessionKey{authKeyID: authKeyID, sessionID: sessionID} c, ok := m.bySession[key] if !ok { + if claim := m.claims[key]; claim != nil { + m.retireClaimLocked(key, claim, true) + m.mu.Unlock() + if !forceCloseConnBatch([]*Conn{claim}, forceCloseBatchTimeout) { + m.log.Warn("Claimed session close exceeded shared deadline", + zap.String("auth_key_id", sessionKeyLog(authKeyID)), + zap.Int64("session_id", sessionID), + ) + } + return true + } m.deletePendingLocked(key) m.mu.Unlock() return false } - offlineUser := m.removeLocked(c, true) + offlineUser := m.retireConnLocked(c, true) lastForUser := offlineUser != 0 && len(m.byUser[offlineUser]) == 0 observer := m.lifecycle m.log.Debug("Session destroyed", @@ -505,10 +644,17 @@ func (m *SessionManager) CloseSessionsForBusinessAuthKey(authKeyID [8]byte) int if !connUsesBusinessAuthKey(c, authKeyID) { continue } - uid := m.removeLocked(c, true) + uid := m.retireConnLocked(c, true) conns = append(conns, c) events = append(events, offlineEvent{key: key, userID: uid, last: uid != 0 && len(m.byUser[uid]) == 0}) } + for key, c := range m.claims { + if !connUsesBusinessAuthKey(c, authKeyID) { + continue + } + m.retireClaimLocked(key, c, true) + conns = append(conns, c) + } observer := m.lifecycle if len(conns) > 0 { m.log.Debug("Force close sessions for revoked auth key", @@ -531,10 +677,26 @@ func (m *SessionManager) CloseSessionsForBusinessAuthKey(authKeyID [8]byte) int return len(conns) } -// CloseSessionsForRawAuthKeyExcept 强制断开指定 raw auth_key 的活跃连接，可排除 -// 一个 session（destroy_auth_key 的发起连接：响应要先送达，它的密钥已删，下一帧 -// 自然失效）。出站推送不回查密钥库，必须主动断开底层 transport 才能让销毁立即生效。 +// CloseSessionsForRawAuthKeyExcept 强制断开指定 raw auth_key 的活跃连接，可按 +// session ID 排除一个 session。该接口供业务层授权撤销使用；wire-level +// destroy_auth_key 必须改用精确 Conn 排除，避免同 session replacement 被误放过。 func (m *SessionManager) CloseSessionsForRawAuthKeyExcept(authKeyID [8]byte, exceptSessionID int64) int { + return m.closeSessionsForRawAuthKey(authKeyID, func(sessionID int64, _ *Conn) bool { + return sessionID == exceptSessionID + }) +} + +// CloseSessionsForRawAuthKeyExceptConn closes every active/claiming owner for a raw +// auth key except the exact Conn executing destroy_auth_key. A session ID is not an +// identity: a concurrent replacement may already own the same logical session while +// the retired request handler finishes deletion. +func (m *SessionManager) CloseSessionsForRawAuthKeyExceptConn(authKeyID [8]byte, except *Conn) int { + return m.closeSessionsForRawAuthKey(authKeyID, func(_ int64, c *Conn) bool { + return c == except + }) +} + +func (m *SessionManager) closeSessionsForRawAuthKey(authKeyID [8]byte, skip func(int64, *Conn) bool) int { type offlineEvent struct { key sessionKey userID int64 @@ -544,14 +706,22 @@ func (m *SessionManager) CloseSessionsForRawAuthKeyExcept(authKeyID [8]byte, exc var conns []*Conn var events []offlineEvent for sessionID, c := range m.byAuthKey[authKeyID] { - if sessionID == exceptSessionID { + if skip != nil && skip(sessionID, c) { continue } key := sessionKey{authKeyID: authKeyID, sessionID: sessionID} - uid := m.removeLocked(c, true) + uid := m.retireConnLocked(c, true) conns = append(conns, c) events = append(events, offlineEvent{key: key, userID: uid, last: uid != 0 && len(m.byUser[uid]) == 0}) } + for sessionID, c := range m.claimsByAuth[authKeyID] { + if skip != nil && skip(sessionID, c) { + continue + } + key := sessionKey{authKeyID: authKeyID, sessionID: sessionID} + m.retireClaimLocked(key, c, true) + conns = append(conns, c) + } observer := m.lifecycle m.mu.Unlock() if !forceCloseConnBatch(conns, forceCloseBatchTimeout) { @@ -574,6 +744,14 @@ func (m *SessionManager) CloseSessionsForRawAuthKeyExcept(authKeyID [8]byte, exc // deadline, but no timed-out Conn can enqueue more work in that interval. Nil/duplicate entries are // removed so Register's replacement/eviction slots cannot close the same Conn twice. func forceCloseConnBatch(conns []*Conn, timeout time.Duration) bool { + return closeConnBatch(conns, timeout, true) +} + +// closeConnBatch always converges physical writers/outbound actors. waitInbound +// is reserved for destructive control-plane operations; activation takeover sets +// it false so a canceled business handler that ignores its context cannot stall a +// healthy replacement. Its admission and response writer are already terminal. +func closeConnBatch(conns []*Conn, timeout time.Duration, waitInbound bool) bool { if len(conns) == 0 { return true } @@ -640,7 +818,7 @@ func forceCloseConnBatch(conns []*Conn, timeout time.Duration) bool { if remaining <= 0 { return false } - if c.rpcScheduler != nil && !c.waitInboundShutdown(remaining) { + if waitInbound && c.rpcScheduler != nil && !c.waitInboundShutdown(remaining) { return false } if c.outboundDone == nil { @@ -1631,6 +1809,9 @@ func (m *SessionManager) onlineChannelUsers(index map[int64]map[sessionKey]int64 func (m *SessionManager) removeLocked(c *Conn, dropPending bool) int64 { key := connSessionKey(c) + if m.bySession[key] != c { + return 0 + } delete(m.bySession, key) removeSessionIDIndex(m.bySessionID, c.sessionID, c.authKeyID) removeConnIndex(m.byAuthKey, c.authKeyID, c.sessionID) @@ -1650,6 +1831,75 @@ func (m *SessionManager) removeLocked(c *Conn, dropPending bool) int64 { return uid } +// retireConnLocked closes every admission/producer gate before the Conn leaves +// manager indexes. Callers may close the physical transport and wait outside m.mu, +// but no pointer collected by an earlier fan-out can enqueue after this returns. +func (m *SessionManager) retireConnLocked(c *Conn, dropPending bool) int64 { + if c == nil { + return 0 + } + c.beginTerminalShutdown() + return m.removeLocked(c, dropPending) +} + +func (m *SessionManager) retireClaimLocked(key sessionKey, c *Conn, dropPending bool) { + if c == nil || m.claims[key] != c { + return + } + c.beginTerminalShutdown() + m.removeClaimLocked(key, c) + if dropPending && m.bySession[key] == nil { + m.deletePendingLocked(key) + } + delete(m.flushing, key) +} + +func (m *SessionManager) claimCountForAuthLocked(authKeyID [8]byte) int { + return len(m.claimsByAuth[authKeyID]) +} + +func (m *SessionManager) oldestAuthOwnerLocked(authKeyID [8]byte, exclude *Conn) (sessionKey, *Conn, bool) { + var ( + oldestKey sessionKey + oldest *Conn + isClaim bool + ) + for sessionID, candidate := range m.byAuthKey[authKeyID] { + if candidate == nil || candidate == exclude { + continue + } + if oldest == nil || candidate.createdAt.Before(oldest.createdAt) { + oldestKey = sessionKey{authKeyID: authKeyID, sessionID: sessionID} + oldest = candidate + isClaim = false + } + } + for sessionID, candidate := range m.claimsByAuth[authKeyID] { + if candidate == nil || candidate == exclude { + continue + } + if oldest == nil || candidate.createdAt.Before(oldest.createdAt) { + oldestKey = sessionKey{authKeyID: authKeyID, sessionID: sessionID} + oldest = candidate + isClaim = true + } + } + return oldestKey, oldest, isClaim +} + +func (m *SessionManager) addClaimLocked(key sessionKey, c *Conn) { + m.claims[key] = c + addConnIndex(m.claimsByAuth, key.authKeyID, key.sessionID, c) +} + +func (m *SessionManager) removeClaimLocked(key sessionKey, c *Conn) { + if m.claims[key] != c { + return + } + delete(m.claims, key) + removeConnIndex(m.claimsByAuth, key.authKeyID, key.sessionID) +} + func (m *SessionManager) businessAuthKeyCandidatesLocked(authKeyID [8]byte) map[sessionKey]*Conn { out := make(map[sessionKey]*Conn, len(m.byBusinessAuthKey[authKeyID])+len(m.byAuthKey[authKeyID])) for key, c := range m.byBusinessAuthKey[authKeyID] { diff --git a/internal/mtprotoedge/structural_limits_test.go b/internal/mtprotoedge/structural_limits_test.go index 10904369..b8a64ce6 100644 --- a/internal/mtprotoedge/structural_limits_test.go +++ b/internal/mtprotoedge/structural_limits_test.go @@ -2,8 +2,12 @@ package mtprotoedge import ( "context" + "encoding/binary" + "errors" + "io" "strings" "testing" + "time" "github.com/gotd/td/bin" "github.com/gotd/td/mt" @@ -39,6 +43,101 @@ func TestContainerMessageCountAndServiceVectorCaps(t *testing.T) { } } +func TestServiceLongVectorViewsAreBoundedExactAndZeroCopy(t *testing.T) { + tests := []struct { + name string + typeID uint32 + value bin.Encoder + }{ + {name: "msgs_ack", typeID: mt.MsgsAckTypeID, value: &mt.MsgsAck{MsgIDs: []int64{11, 22}}}, + {name: "msgs_state_req", typeID: mt.MsgsStateReqTypeID, value: &mt.MsgsStateReq{MsgIDs: []int64{11, 22}}}, + {name: "msg_resend_req", typeID: mt.MsgResendReqTypeID, value: &mt.MsgResendReq{MsgIDs: []int64{11, 22}}}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + var encoded bin.Buffer + if err := tt.value.Encode(&encoded); err != nil { + t.Fatalf("encode service vector: %v", err) + } + view, err := decodeInt64VectorView(&encoded, tt.typeID, maxServiceMessageIDs) + if err != nil { + t.Fatalf("decode service vector view: %v", err) + } + if view.count != 2 { + t.Fatalf("view count = %d, want 2", view.count) + } + + // The retained plan payload is a view into the already-budgeted frame, + // not a generated-decoder []int64 copy. + binary.LittleEndian.PutUint64(encoded.Buf[12:20], uint64(99)) + ids := view.materialize() + if len(ids) != 2 || ids[0] != 99 || ids[1] != 22 { + t.Fatalf("materialized ids = %v, want [99 22]", ids) + } + + trailing := bin.Buffer{Buf: append(append([]byte(nil), encoded.Buf...), 0, 0, 0, 0)} + if _, err := decodeInt64VectorView(&trailing, tt.typeID, maxServiceMessageIDs); err == nil || !strings.Contains(err.Error(), "trailing") { + t.Fatalf("trailing vector err = %v, want exact-length rejection", err) + } + truncated := bin.Buffer{Buf: encoded.Buf[:len(encoded.Buf)-1]} + if _, err := decodeInt64VectorView(&truncated, tt.typeID, maxServiceMessageIDs); !errors.Is(err, io.ErrUnexpectedEOF) { + t.Fatalf("truncated vector err = %v, want unexpected EOF", err) + } + }) + } + + var oversized bin.Buffer + if err := (&mt.MsgsAck{MsgIDs: []int64{1}}).Encode(&oversized); err != nil { + t.Fatalf("encode oversized seed: %v", err) + } + binary.LittleEndian.PutUint32(oversized.Buf[8:12], uint32(maxServiceMessageIDs+1)) + if _, err := decodeInt64VectorView(&oversized, mt.MsgsAckTypeID, maxServiceMessageIDs); err == nil || !strings.Contains(err.Error(), "exceeds") { + t.Fatalf("oversized vector err = %v, want capped rejection", err) + } +} + +func TestInboundPlanContainerCapacityBoundary(t *testing.T) { + ids := proto.NewMessageIDGen(time.Now) + waitBody := encodeClientMessageBodyForTest(t, &mt.HTTPWaitRequest{MaxWait: 25_000}) + messages := make([]proto.Message, maxContainerMessages) + for i := range messages { + messages[i] = proto.Message{ + ID: ids.New(proto.MessageFromClient), SeqNo: 0, Bytes: len(waitBody), Body: waitBody, + } + } + outerMsgID := ids.New(proto.MessageFromClient) + body := encodeClientMessageBodyForTest(t, &proto.MessageContainer{Messages: messages}) + + s := New(Options{Logger: zaptest.NewLogger(t)}) + before := s.frameBudget.usedBytes() + plan, err := s.preflightInbound(newConnState(), outerMsgID, 0, body) + if err != nil { + t.Fatalf("container at cap preflight: %v", err) + } + if len(plan.items) != maxContainerMessages || len(plan.staged) != maxContainerMessages+1 { + t.Fatalf("container plan sizes = items:%d staged:%d", len(plan.items), len(plan.staged)) + } + if plan.logicalMin != messages[0].ID { + t.Fatalf("container plan floor = %d, want %d", plan.logicalMin, messages[0].ID) + } + plan.close() + if after := s.frameBudget.usedBytes(); after != before { + t.Fatalf("container plan leaked frame budget: before=%d after=%d", before, after) + } + + over := append(append([]proto.Message(nil), messages...), proto.Message{ + ID: outerMsgID, SeqNo: 0, Bytes: len(waitBody), Body: waitBody, + }) + overBody := encodeClientMessageBodyForTest(t, &proto.MessageContainer{Messages: over}) + if plan, err := s.preflightInbound(newConnState(), ids.New(proto.MessageFromClient), 0, overBody); err == nil { + plan.close() + t.Fatal("container above cap unexpectedly accepted") + } + if after := s.frameBudget.usedBytes(); after != before { + t.Fatalf("rejected container leaked frame budget: before=%d after=%d", before, after) + } +} + func TestContainerDecodeUsesBudgetedZeroCopyBodies(t *testing.T) { encoded := bin.Buffer{} wantBody := []byte{0x11, 0x22, 0x33, 0x44} @@ -83,6 +182,319 @@ func TestContainerDecodeUsesBudgetedZeroCopyBodies(t *testing.T) { } } +func TestContainerAndGZIPViewsRejectTrailingBytes(t *testing.T) { + var container bin.Buffer + if err := (&proto.MessageContainer{}).Encode(&container); err != nil { + t.Fatalf("encode empty container: %v", err) + } + container.PutInt(0x11223344) + s := New(Options{Logger: zaptest.NewLogger(t)}) + if _, release, err := s.decodeMessageContainerViews(&container, 0); err == nil { + release() + t.Fatal("container trailing bytes unexpectedly accepted") + } + + var packed bin.Buffer + if err := (proto.GZIP{Data: []byte{1, 2, 3, 4}}).Encode(&packed); err != nil { + t.Fatalf("encode gzip wrapper: %v", err) + } + packed.PutInt(0x55667788) + if _, err := gzipPackedBytesView(&packed); err == nil { + t.Fatal("gzip trailing bytes unexpectedly accepted") + } +} + +func TestContainerInvalidSequenceTailIsAtomic(t *testing.T) { + ids := proto.NewMessageIDGen(time.Now) + firstMsgID := ids.New(proto.MessageFromClient) + secondMsgID := ids.New(proto.MessageFromClient) + outerMsgID := ids.New(proto.MessageFromClient) + pingOne := encodeClientMessageBodyForTest(t, &mt.PingRequest{PingID: 1}) + pingTwo := encodeClientMessageBodyForTest(t, &mt.PingRequest{PingID: 2}) + container := proto.MessageContainer{Messages: []proto.Message{ + {ID: firstMsgID, SeqNo: 3, Bytes: len(pingOne), Body: pingOne}, + {ID: secondMsgID, SeqNo: 1, Bytes: len(pingTwo), Body: pingTwo}, + }} + encoded := encodeClientMessageBodyForTest(t, &container) + + s := New(Options{Logger: zaptest.NewLogger(t)}) + cs := newConnState() + c := &Conn{ + metrics: NopMetrics{}, + outbound: make(chan outboundOp, 4), + outboundControl: make(chan outboundOp, 4), + outboundStop: make(chan struct{}), + } + var acks []int64 + if err := s.dispatch(context.Background(), cs, c, outerMsgID, 4, &bin.Buffer{Buf: encoded}, &acks); err != nil { + t.Fatalf("dispatch invalid-tail container: %v", err) + } + + if len(cs.seen) != 0 || len(cs.order) != 0 || cs.maxContentMsgID != 0 || cs.maxContentSeqNo != 0 { + t.Fatalf("invalid container partially committed connState: %+v", cs) + } + if len(acks) != 0 { + t.Fatalf("invalid container produced ACKs: %v", acks) + } + if got := len(c.outboundControl); got != 1 { + t.Fatalf("control frames = %d, want only bad_msg_notification", got) + } + op := <-c.outboundControl + if op.encoded == nil || op.encoded.typeID != mt.BadMsgNotificationTypeID { + t.Fatalf("control frame type = %#v, want bad_msg_notification", op.encoded) + } + var bad mt.BadMsgNotification + if err := bad.Decode(&bin.Buffer{Buf: op.encoded.body}); err != nil { + t.Fatalf("decode bad_msg_notification: %v", err) + } + if bad.BadMsgID != secondMsgID || bad.BadMsgSeqno != 1 || bad.ErrorCode != badMsgSeqTooLow { + t.Fatalf("bad_msg_notification = %+v", bad) + } + op.releaseReservation(c.outboundTrackedBudget) +} + +func TestContainerMalformedGZIPTailHasNoStateEffects(t *testing.T) { + ids := proto.NewMessageIDGen(time.Now) + waitMsgID := ids.New(proto.MessageFromClient) + badMsgID := ids.New(proto.MessageFromClient) + outerMsgID := ids.New(proto.MessageFromClient) + waitBody := encodeClientMessageBodyForTest(t, &mt.HTTPWaitRequest{MaxWait: 25_000}) + var malformedGZIP bin.Buffer + malformedGZIP.PutID(proto.GZIPTypeID) + container := proto.MessageContainer{Messages: []proto.Message{ + {ID: waitMsgID, SeqNo: 0, Bytes: len(waitBody), Body: waitBody}, + {ID: badMsgID, SeqNo: 1, Bytes: malformedGZIP.Len(), Body: malformedGZIP.Buf}, + }} + encoded := encodeClientMessageBodyForTest(t, &container) + + s := New(Options{Logger: zaptest.NewLogger(t)}) + cs := newConnState() + var acks []int64 + if err := s.dispatch(context.Background(), cs, nil, outerMsgID, 2, &bin.Buffer{Buf: encoded}, &acks); err == nil { + t.Fatal("malformed gzip tail unexpectedly accepted") + } + if len(cs.seen) != 0 || len(cs.order) != 0 || cs.maxContentMsgID != 0 || cs.maxContentSeqNo != 0 { + t.Fatalf("malformed gzip tail partially committed connState: %+v", cs) + } + if len(acks) != 0 { + t.Fatalf("malformed gzip tail produced ACKs: %v", acks) + } + if got := s.frameBudget.usedBytes(); got != 0 { + t.Fatalf("malformed gzip tail leaked frame budget: %d", got) + } +} + +func TestInvalidMessageIDRejectsBeforeGZIPExpansion(t *testing.T) { + leaf := encodeClientMessageBodyForTest(t, &mt.PingRequest{PingID: 1}) + wrapped := encodeClientMessageBodyForTest(t, &proto.GZIP{Data: leaf}) + current := proto.NewMessageIDGen(time.Now).New(proto.MessageFromClient) + tests := []struct { + name string + msgID int64 + badCode int + }{ + {name: "stale", msgID: proto.NewMessageIDGen(func() time.Time { return time.Now().Add(-10 * time.Minute) }).New(proto.MessageFromClient), badCode: badMsgIDTooLow}, + {name: "future", msgID: proto.NewMessageIDGen(func() time.Time { return time.Now().Add(time.Minute) }).New(proto.MessageFromClient), badCode: badMsgIDTooHigh}, + {name: "invalid bits", msgID: current + 1, badCode: badMsgIDInvalidBits}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + s := New(Options{Logger: zaptest.NewLogger(t)}) + // Any attempted gzip expansion would fail this deliberately tiny + // budget and mask the required bad_msg_notification. + s.frameBudget = newInboundFrameBudget(1) + plan, err := s.preflightInbound(newConnState(), tt.msgID, 1, wrapped) + if plan != nil { + plan.close() + t.Fatal("invalid gzip envelope unexpectedly produced a plan") + } + var bad *dispatchBadMsgError + if !errors.As(err, &bad) || bad.code != tt.badCode { + t.Fatalf("invalid gzip envelope err = %v, want bad_msg code %d", err, tt.badCode) + } + if got := s.frameBudget.usedBytes(); got != 0 { + t.Fatalf("invalid gzip envelope consumed expansion budget: %d", got) + } + }) + } +} + +func TestInvalidContainerInnerIDRejectsBeforeGZIPExpansion(t *testing.T) { + ids := proto.NewMessageIDGen(time.Now) + validInnerID := ids.New(proto.MessageFromClient) + invalidInnerID := validInnerID + 1 + if proto.MessageID(invalidInnerID).Type() == proto.MessageFromClient { + t.Fatalf("test inner msg_id %d unexpectedly has client bits", invalidInnerID) + } + outerMsgID := ids.New(proto.MessageFromClient) + leaf := encodeClientMessageBodyForTest(t, &mt.PingRequest{PingID: 2}) + gzipped := encodeClientMessageBodyForTest(t, &proto.GZIP{Data: leaf}) + container := encodeClientMessageBodyForTest(t, &proto.MessageContainer{Messages: []proto.Message{{ + ID: invalidInnerID, SeqNo: 1, Bytes: len(gzipped), Body: gzipped, + }}}) + + s := New(Options{Logger: zaptest.NewLogger(t)}) + // Enough for the one descriptor, intentionally not enough for a gzip + // expansion. Invalid inner bits must win and every descriptor charge unwind. + s.frameBudget = newInboundFrameBudget(containerDescriptorBudgetBytes) + plan, err := s.preflightInbound(newConnState(), outerMsgID, 2, container) + if plan != nil { + plan.close() + t.Fatal("invalid inner msg_id unexpectedly produced a plan") + } + var bad *dispatchBadMsgError + if !errors.As(err, &bad) || bad.code != badMsgContainer || bad.msgID != invalidInnerID { + t.Fatalf("invalid inner gzip err = %v, want container rejection for %d", err, invalidInnerID) + } + if got := s.frameBudget.usedBytes(); got != 0 { + t.Fatalf("invalid inner gzip leaked descriptor/expansion budget: %d", got) + } +} + +func TestSeenGZIPInnerShortCircuitsBeforeExpansion(t *testing.T) { + tests := []struct { + name string + seqNo int32 + content bool + leaf bin.Encoder + wantACK bool + }{ + {name: "content", seqNo: 1, content: true, leaf: &mt.PingRequest{PingID: 3}, wantACK: true}, + {name: "non_content", seqNo: 0, content: false, leaf: &mt.HTTPWaitRequest{MaxWait: 25_000}}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + ids := proto.NewMessageIDGen(time.Now) + innerMsgID := ids.New(proto.MessageFromClient) + outerMsgID := ids.New(proto.MessageFromClient) + leaf := encodeClientMessageBodyForTest(t, tt.leaf) + gzipped := encodeClientMessageBodyForTest(t, &proto.GZIP{Data: leaf}) + container := encodeClientMessageBodyForTest(t, &proto.MessageContainer{Messages: []proto.Message{{ + ID: innerMsgID, SeqNo: int(tt.seqNo), Bytes: len(gzipped), Body: gzipped, + }}}) + + cs := newConnState() + cs.track(innerMsgID, tt.seqNo, tt.content, msgStateReceived) + s := New(Options{Logger: zaptest.NewLogger(t)}) + s.frameBudget = newInboundFrameBudget(containerDescriptorBudgetBytes) + plan, err := s.preflightInbound(cs, outerMsgID, 2, container) + if err != nil { + t.Fatalf("preflight seen gzip inner: %v", err) + } + if len(plan.items) != 1 || plan.items[0].kind != inboundItemDuplicate || plan.items[0].content != tt.content { + plan.close() + t.Fatalf("seen gzip plan items = %+v, want one duplicate", plan.items) + } + if gotACK := len(plan.ackIDs) == 1 && plan.ackIDs[0] == innerMsgID; gotACK != tt.wantACK { + plan.close() + t.Fatalf("seen gzip ACKs = %v, wantACK=%v", plan.ackIDs, tt.wantACK) + } + if got := s.frameBudget.usedBytes(); got != containerDescriptorBudgetBytes { + plan.close() + t.Fatalf("seen gzip held budget = %d, want descriptor-only %d", got, containerDescriptorBudgetBytes) + } + plan.close() + if got := s.frameBudget.usedBytes(); got != 0 { + t.Fatalf("seen gzip close leaked budget: %d", got) + } + }) + } +} + +func TestSeenTopLevelContentGZIPShortCircuitsBeforeExpansion(t *testing.T) { + msgID := proto.NewMessageIDGen(time.Now).New(proto.MessageFromClient) + leaf := encodeClientMessageBodyForTest(t, &mt.PingRequest{PingID: 31}) + wrapped := encodeClientMessageBodyForTest(t, &proto.GZIP{Data: leaf}) + cs := newConnState() + cs.track(msgID, 1, true, msgStateReceived) + + s := New(Options{Logger: zaptest.NewLogger(t)}) + s.frameBudget = newInboundFrameBudget(1) + plan, err := s.preflightInbound(cs, msgID, 1, wrapped) + if err != nil { + t.Fatalf("preflight seen top-level gzip: %v", err) + } + if len(plan.items) != 1 || plan.items[0].kind != inboundItemDuplicate || len(plan.ackIDs) != 1 || plan.ackIDs[0] != msgID { + plan.close() + t.Fatalf("seen top-level gzip plan = items:%+v acks:%v", plan.items, plan.ackIDs) + } + plan.close() + if got := s.frameBudget.usedBytes(); got != 0 { + t.Fatalf("seen top-level gzip consumed expansion budget: %d", got) + } +} + +func TestDuplicateOuterContainerOnlyInspectsInnerDescriptors(t *testing.T) { + ids := proto.NewMessageIDGen(time.Now) + innerMsgID := ids.New(proto.MessageFromClient) + outerMsgID := ids.New(proto.MessageFromClient) + leaf := encodeClientMessageBodyForTest(t, &mt.PingRequest{PingID: 4}) + gzipped := encodeClientMessageBodyForTest(t, &proto.GZIP{Data: leaf}) + container := encodeClientMessageBodyForTest(t, &proto.MessageContainer{Messages: []proto.Message{{ + ID: innerMsgID, SeqNo: 1, Bytes: len(gzipped), Body: gzipped, + }}}) + wrappedContainer := encodeClientMessageBodyForTest(t, &proto.GZIP{Data: container}) + + t.Run("seen inner replays without expansion", func(t *testing.T) { + cs := newConnState() + cs.track(outerMsgID, 2, false, msgStateReceived) + cs.track(innerMsgID, 1, true, msgStateReceived) + s := New(Options{Logger: zaptest.NewLogger(t)}) + s.frameBudget = newInboundFrameBudget(containerDescriptorBudgetBytes) + plan, err := s.preflightInbound(cs, outerMsgID, 2, container) + if err != nil { + t.Fatalf("preflight duplicate outer: %v", err) + } + if len(plan.items) != 1 || plan.items[0].kind != inboundItemDuplicate { + plan.close() + t.Fatalf("duplicate outer items = %+v", plan.items) + } + plan.close() + if got := s.frameBudget.usedBytes(); got != 0 { + t.Fatalf("duplicate outer leaked descriptor budget: %d", got) + } + }) + + t.Run("top-level gzip container still locates inner", func(t *testing.T) { + cs := newConnState() + cs.track(outerMsgID, 2, false, msgStateReceived) + cs.track(innerMsgID, 1, true, msgStateReceived) + s := New(Options{Logger: zaptest.NewLogger(t)}) + s.frameBudget = newInboundFrameBudget(maxSingleGZIPExpandedBytes + containerDescriptorBudgetBytes) + plan, err := s.preflightInbound(cs, outerMsgID, 2, wrappedContainer) + if err != nil { + t.Fatalf("preflight gzip duplicate outer: %v", err) + } + if len(plan.items) != 1 || plan.items[0].kind != inboundItemDuplicate || plan.items[0].msgID != innerMsgID { + plan.close() + t.Fatalf("gzip duplicate outer items = %+v, want inner duplicate", plan.items) + } + plan.close() + if got := s.frameBudget.usedBytes(); got != 0 { + t.Fatalf("gzip duplicate outer leaked expansion/descriptor budget: %d", got) + } + }) + + t.Run("unseen inner rejects without expansion", func(t *testing.T) { + cs := newConnState() + cs.track(outerMsgID, 2, false, msgStateReceived) + s := New(Options{Logger: zaptest.NewLogger(t)}) + s.frameBudget = newInboundFrameBudget(containerDescriptorBudgetBytes) + plan, err := s.preflightInbound(cs, outerMsgID, 2, container) + if plan != nil { + plan.close() + t.Fatal("duplicate outer with unseen inner unexpectedly produced a plan") + } + var bad *dispatchBadMsgError + if !errors.As(err, &bad) || bad.code != badMsgContainer { + t.Fatalf("duplicate outer unseen-inner err = %v, want container rejection", err) + } + if got := s.frameBudget.usedBytes(); got != 0 { + t.Fatalf("unseen duplicate inner leaked descriptor/expansion budget: %d", got) + } + }) +} + func TestServiceInfoViewsRejectOversizedBytesWithoutDecodeCopy(t *testing.T) { state := mt.MsgsStateInfo{ReqMsgID: 7, Info: make([]byte, maxServiceMessageIDs)} var encodedState bin.Buffer @@ -142,10 +554,58 @@ func TestDispatchRejectsExcessiveWrapperDepthBeforeRPC(t *testing.T) { s := New(Options{Logger: zaptest.NewLogger(t)}) var acks []int64 - err := s.dispatch(context.Background(), newConnState(), nil, 4, 0, &bin.Buffer{Buf: encoded}, &acks) + msgID := proto.NewMessageIDGen(time.Now).New(proto.MessageFromClient) + err := s.dispatch(context.Background(), newConnState(), nil, msgID, 0, &bin.Buffer{Buf: encoded}, &acks) if err == nil || !strings.Contains(err.Error(), "wrapper depth") { t.Fatalf("deep wrapper err = %v, want wrapper depth rejection", err) } + if got := s.frameBudget.usedBytes(); got != 0 { + t.Fatalf("deep wrapper rejection leaked frame budget: %d", got) + } +} + +func TestNestedWrapperTailFailuresReleaseEveryBudget(t *testing.T) { + ping := encodeClientMessageBodyForTest(t, &mt.PingRequest{PingID: 5}) + + var gzipWithTail bin.Buffer + if err := (&proto.GZIP{Data: ping}).Encode(&gzipWithTail); err != nil { + t.Fatalf("encode inner gzip: %v", err) + } + gzipWithTail.PutInt(0x11223344) + gzipTailWrapped := encodeClientMessageBodyForTest(t, &proto.GZIP{Data: gzipWithTail.Copy()}) + + var containerWithTail bin.Buffer + if err := (&proto.MessageContainer{}).Encode(&containerWithTail); err != nil { + t.Fatalf("encode inner container: %v", err) + } + containerWithTail.PutInt(0x55667788) + containerTailWrapped := encodeClientMessageBodyForTest(t, &proto.GZIP{Data: containerWithTail.Copy()}) + + tests := []struct { + name string + body []byte + seqNo int32 + }{ + {name: "gzip tail after outer expansion", body: gzipTailWrapped, seqNo: 1}, + {name: "container tail after outer expansion", body: containerTailWrapped, seqNo: 0}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + s := New(Options{Logger: zaptest.NewLogger(t)}) + msgID := proto.NewMessageIDGen(time.Now).New(proto.MessageFromClient) + plan, err := s.preflightInbound(newConnState(), msgID, tt.seqNo, tt.body) + if plan != nil { + plan.close() + t.Fatal("invalid nested wrapper unexpectedly produced a plan") + } + if err == nil || !strings.Contains(err.Error(), "trailing") { + t.Fatalf("nested wrapper err = %v, want trailing-byte rejection", err) + } + if got := s.frameBudget.usedBytes(); got != 0 { + t.Fatalf("nested wrapper failure leaked frame budget: %d", got) + } + }) + } } func TestOversizedConnectionBuffersAreReleasedAfterFrame(t *testing.T) { diff --git a/internal/mtprotoedge/transport_ownership.go b/internal/mtprotoedge/transport_ownership.go new file mode 100644 index 00000000..c7557dc5 --- /dev/null +++ b/internal/mtprotoedge/transport_ownership.go @@ -0,0 +1,339 @@ +package mtprotoedge + +import ( + "context" + "sync" + "sync/atomic" + "time" + + "github.com/gotd/td/bin" + "github.com/gotd/td/transport" +) + +const ( + physicalTransportClosedBit = uint64(1) << 63 + physicalTransportGenerationMax = physicalTransportClosedBit - 1 +) + +// physicalTransportOwner separates the lifetime of one physical socket from +// the logical Conn generations that successively use it. state is either an +// open generation in [1, physicalTransportGenerationMax], or that generation +// with physicalTransportClosedBit set. +// +// Transfer and owner-close race through one atomic state transition. This is +// the critical property: when transfer wins, a stale generation can no longer +// close the socket; when close wins, no later generation can be published. +type physicalTransportOwner struct { + raw transport.Conn + + // writeMu orders a completed send before generation transfer. CloseAny and + // owner-close deliberately do not wait for it: raw.Close must be able to + // interrupt a transport implementation blocked inside Send. + writeMu sync.Mutex + state atomic.Uint64 + + // binding identifies the logical Conn currently attached to the open + // generation. Physical close fences that exact Conn before touching raw; + // generation matching prevents a stale close from retiring a replacement. + bindingMu sync.Mutex + boundGeneration uint64 + boundLogicalConn *Conn + + closeDone chan struct{} + closeErr error +} + +// physicalTransportLease is the capability owned by one logical Conn +// generation. It implements transport.Conn so exchange/protocol-error paths +// can use the initial lease from the moment serveConn starts. +type physicalTransportLease struct { + owner *physicalTransportOwner + generation uint64 +} + +var _ transport.Conn = (*physicalTransportLease)(nil) + +func newPhysicalTransportOwner(raw transport.Conn) (*physicalTransportOwner, *physicalTransportLease) { + owner := &physicalTransportOwner{ + raw: raw, + closeDone: make(chan struct{}), + } + owner.state.Store(1) + return owner, &physicalTransportLease{owner: owner, generation: 1} +} + +// Transfer atomically hands the physical transport to the next logical Conn +// generation. Holding writeMu ensures every send admitted by the old lease has +// returned before the generation changes. CloseAny may still interrupt such a +// send and make the CAS fail, which is the required shutdown ordering. +func (l *physicalTransportLease) Transfer() (*physicalTransportLease, bool) { + if l == nil || l.owner == nil || l.generation == 0 { + return nil, false + } + owner := l.owner + owner.writeMu.Lock() + defer owner.writeMu.Unlock() + + if l.generation >= physicalTransportGenerationMax { + return nil, false + } + if !owner.state.CompareAndSwap(l.generation, l.generation+1) { + return nil, false + } + owner.bindingMu.Lock() + if owner.boundGeneration == l.generation { + owner.boundGeneration = l.generation + 1 + owner.boundLogicalConn = nil + } + owner.bindingMu.Unlock() + return &physicalTransportLease{owner: owner, generation: l.generation + 1}, true +} + +// IsCurrentOpen reports whether this lease is still the unique open owner. +// Callers use it at protocol publication barriers after a required write. +func (l *physicalTransportLease) IsCurrentOpen() bool { + return l != nil && l.owner != nil && l.generation != 0 && + l.owner.state.Load() == l.generation +} + +// bindLogicalConn publishes the logical Conn for this exact generation. If +// physical close already won, the new Conn is terminally fenced immediately +// and can never pass SessionManager publication checks. +func (l *physicalTransportLease) bindLogicalConn(c *Conn) bool { + if l == nil || l.owner == nil || c == nil { + if c != nil { + c.beginTerminalShutdown() + } + return false + } + owner := l.owner + owner.bindingMu.Lock() + open := owner.state.Load() == l.generation + if open { + owner.boundGeneration = l.generation + owner.boundLogicalConn = c + } + owner.bindingMu.Unlock() + if !open { + c.beginTerminalShutdown() + } + return open +} + +// Send admits a write only while this lease is the current open generation. +// The lock also serializes quick ACK/protocol writes with the outbound actor. +func (l *physicalTransportLease) Send(ctx context.Context, b *bin.Buffer) error { + return l.withCurrentWriter(func(raw transport.Conn) error { + return raw.Send(ctx, b) + }) +} + +// SendDeadline preserves the allocation-free fast path of compatTransportConn. +func (l *physicalTransportLease) SendDeadline(deadline time.Time, b *bin.Buffer) error { + return l.withCurrentWriter(func(raw transport.Conn) error { + 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 + } + owner := l.owner + owner.writeMu.Lock() + defer owner.writeMu.Unlock() + if owner.state.Load() != l.generation { + return ErrConnClosed + } + return send(owner.raw) +} + +// Recv is owned by serveConn rather than a logical Conn generation. It is a +// direct proxy so the read loop remains valid after ownership transfers. +func (l *physicalTransportLease) Recv(ctx context.Context, b *bin.Buffer) error { + if l == nil || l.owner == nil || l.owner.raw == nil { + return ErrConnClosed + } + return l.owner.raw.Recv(ctx, b) +} + +// RecvDeadline preserves serveConn's direct-deadline fast path. +func (l *physicalTransportLease) RecvDeadline(deadline time.Time, b *bin.Buffer) error { + if l == nil || l.owner == nil || l.owner.raw == nil { + return ErrConnClosed + } + if receiver, ok := l.owner.raw.(deadlineReceiver); ok { + return receiver.RecvDeadline(deadline, b) + } + ctx := context.Background() + cancel := func() {} + if !deadline.IsZero() { + ctx, cancel = context.WithDeadline(ctx, deadline) + } + defer cancel() + return l.owner.raw.Recv(ctx, b) +} + +// Close closes the physical transport only if this lease still owns the +// current generation. A stale logical Conn therefore cannot close a socket +// already transferred to its replacement. +func (l *physicalTransportLease) Close() error { + if l == nil || l.owner == nil || l.generation == 0 { + return nil + } + owner := l.owner + for { + state := owner.state.Load() + if state&physicalTransportClosedBit != 0 { + return owner.waitClosed() + } + if state != l.generation { + return nil + } + if owner.state.CompareAndSwap(state, state|physicalTransportClosedBit) { + owner.fenceBoundGeneration(l.generation) + return owner.closeRaw() + } + } +} + +// startCloseAlreadyFenced is the non-reentrant close capability for the logical +// Conn that has already published terminal/lifecycle gates itself. It +// synchronously publishes the physical closed bit (so Transfer cannot win), then +// runs the potentially pathological raw.Close outside the RPC worker/flight +// handoff. It neither calls fenceBoundGeneration nor waits for a CloseAny that +// already won, avoiding both lifecycle reentry and close cycles. +func (l *physicalTransportLease) startCloseAlreadyFenced() { + if l == nil || l.owner == nil || l.generation == 0 { + return + } + owner := l.owner + for { + state := owner.state.Load() + if state&physicalTransportClosedBit != 0 || state != l.generation { + return + } + if owner.state.CompareAndSwap(state, state|physicalTransportClosedBit) { + owner.bindingMu.Lock() + if owner.boundGeneration == l.generation { + owner.boundLogicalConn = nil + } + owner.bindingMu.Unlock() + go func() { _ = owner.closeRaw() }() + return + } + } +} + +// CloseAny is the unconditional physical-socket capability retained by +// serveConn. It marks the owner closed before calling raw.Close and never waits +// for writeMu, so it can break a blocked Send or Recv. +func (o *physicalTransportOwner) CloseAny() error { + if o == nil { + return nil + } + for { + state := o.state.Load() + if state&physicalTransportClosedBit != 0 { + return o.waitClosed() + } + if o.state.CompareAndSwap(state, state|physicalTransportClosedBit) { + o.fenceBoundGeneration(state) + return o.closeRaw() + } + } +} + +func (o *physicalTransportOwner) fenceBoundGeneration(generation uint64) { + o.bindingMu.Lock() + var c *Conn + if o.boundGeneration == generation { + c = o.boundLogicalConn + o.boundLogicalConn = nil + } + o.bindingMu.Unlock() + if c != nil { + c.beginTerminalShutdown() + } +} + +func (o *physicalTransportOwner) closeRaw() error { + if o.raw != nil { + o.closeErr = o.raw.Close() + } + close(o.closeDone) + return o.closeErr +} + +func (o *physicalTransportOwner) waitClosed() error { + <-o.closeDone + return o.closeErr +} + +// Forward the optional compat-transport capabilities hidden by the lease. +// These keep frame-budget ownership and quick-ack semantics unchanged while +// serveConn operates on the initial lease instead of the raw transport. +func (l *physicalTransportLease) releaseInboundFrame() { + if l == nil || l.owner == nil { + return + } + if releaser, ok := l.owner.raw.(inboundFrameOwnershipReleaser); ok { + releaser.releaseInboundFrame() + } +} + +func (l *physicalTransportLease) retainInboundFrameBytes(n int64) bool { + if l == nil || l.owner == nil { + return true + } + if retainer, ok := l.owner.raw.(inboundFrameBackingRetainer); ok { + return retainer.retainInboundFrameBytes(n) + } + return true +} + +func (l *physicalTransportLease) ConsumeQuickAckRequested() bool { + if l == nil || l.owner == nil { + return false + } + if quick, ok := l.owner.raw.(quickAckTransport); ok { + return quick.ConsumeQuickAckRequested() + } + return false +} + +func (l *physicalTransportLease) SendQuickAck(ctx context.Context, token uint32) error { + return l.withCurrentWriter(func(raw transport.Conn) error { + if quick, ok := raw.(quickAckTransport); ok { + return quick.SendQuickAck(ctx, token) + } + return nil + }) +} + +func (l *physicalTransportLease) SendQuickAckDeadline(deadline time.Time, token uint32) error { + return l.withCurrentWriter(func(raw transport.Conn) error { + if quick, ok := raw.(deadlineQuickAckTransport); ok { + return quick.SendQuickAckDeadline(deadline, token) + } + if quick, ok := raw.(quickAckTransport); ok { + ctx := context.Background() + cancel := func() {} + if !deadline.IsZero() { + ctx, cancel = context.WithDeadline(ctx, deadline) + } + defer cancel() + return quick.SendQuickAck(ctx, token) + } + return nil + }) +} diff --git a/internal/mtprotoedge/transport_ownership_test.go b/internal/mtprotoedge/transport_ownership_test.go new file mode 100644 index 00000000..18e1f90e --- /dev/null +++ b/internal/mtprotoedge/transport_ownership_test.go @@ -0,0 +1,198 @@ +package mtprotoedge + +import ( + "context" + "errors" + "io" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/gotd/td/bin" +) + +type ownershipTestTransport struct { + closeCalls atomic.Int32 + sendCalls atomic.Int32 + closed chan struct{} + closeOnce sync.Once + blockSend bool + started chan struct{} + startOnce sync.Once +} + +func newOwnershipTestTransport(blockSend bool) *ownershipTestTransport { + return &ownershipTestTransport{ + closed: make(chan struct{}), + blockSend: blockSend, + started: make(chan struct{}), + } +} + +func (t *ownershipTestTransport) Send(context.Context, *bin.Buffer) error { + t.sendCalls.Add(1) + t.startOnce.Do(func() { close(t.started) }) + if !t.blockSend { + select { + case <-t.closed: + return io.ErrClosedPipe + default: + return nil + } + } + <-t.closed + return io.ErrClosedPipe +} + +func (*ownershipTestTransport) Recv(context.Context, *bin.Buffer) error { return io.EOF } + +func (t *ownershipTestTransport) Close() error { + t.closeOnce.Do(func() { + t.closeCalls.Add(1) + close(t.closed) + }) + return nil +} + +func TestPhysicalTransportTransferFencesStaleLease(t *testing.T) { + raw := newOwnershipTestTransport(false) + owner, oldLease := newPhysicalTransportOwner(raw) + newLease, ok := oldLease.Transfer() + if !ok { + t.Fatal("transfer failed") + } + if err := oldLease.Send(context.Background(), &bin.Buffer{}); !errors.Is(err, ErrConnClosed) { + t.Fatalf("stale Send error = %v, want ErrConnClosed", err) + } + if err := oldLease.Close(); err != nil { + t.Fatalf("stale Close: %v", err) + } + if got := raw.closeCalls.Load(); got != 0 { + t.Fatalf("stale Close closed raw %d times", got) + } + if err := newLease.Send(context.Background(), &bin.Buffer{}); err != nil { + t.Fatalf("current Send: %v", err) + } + if err := owner.CloseAny(); err != nil { + t.Fatalf("owner CloseAny: %v", err) + } + if got := raw.closeCalls.Load(); got != 1 { + t.Fatalf("raw close calls = %d, want 1", got) + } +} + +func TestPhysicalTransportCloseWinningPreventsTransfer(t *testing.T) { + raw := newOwnershipTestTransport(false) + owner, lease := newPhysicalTransportOwner(raw) + if err := owner.CloseAny(); err != nil { + t.Fatalf("CloseAny: %v", err) + } + if next, ok := lease.Transfer(); ok || next != nil { + t.Fatalf("transfer after close = (%p,%v), want nil,false", next, ok) + } +} + +func TestPhysicalTransportCloseAnyInterruptsWriteAndDefeatsTransfer(t *testing.T) { + raw := newOwnershipTestTransport(true) + owner, lease := newPhysicalTransportOwner(raw) + sendDone := make(chan error, 1) + go func() { sendDone <- lease.Send(context.Background(), &bin.Buffer{}) }() + select { + case <-raw.started: + case <-time.After(time.Second): + t.Fatal("write did not block") + } + transferDone := make(chan bool, 1) + go func() { + _, ok := lease.Transfer() + transferDone <- ok + }() + select { + case ok := <-transferDone: + t.Fatalf("transfer returned before blocked write ended: %v", ok) + case <-time.After(20 * time.Millisecond): + } + if err := owner.CloseAny(); err != nil { + t.Fatalf("CloseAny: %v", err) + } + select { + case err := <-sendDone: + if !errors.Is(err, io.ErrClosedPipe) { + t.Fatalf("blocked Send error = %v", err) + } + case <-time.After(time.Second): + t.Fatal("CloseAny did not interrupt write") + } + select { + case ok := <-transferDone: + if ok { + t.Fatal("transfer won after physical close") + } + case <-time.After(time.Second): + t.Fatal("transfer did not finish") + } +} + +func TestPhysicalCloseFencesActivationPublication(t *testing.T) { + raw := newOwnershipTestTransport(false) + owner, lease := newPhysicalTransportOwner(raw) + s := New(Options{}) + c := s.newConnWithLease(lease, newTestAuthKey(t), 73001, 1) + if !c.beginActivationClaim() { + t.Fatal("activation claim failed") + } + if err := owner.CloseAny(); err != nil { + t.Fatalf("CloseAny: %v", err) + } + if c.publishActivation() { + t.Fatal("closed physical transport published active Conn") + } + if !c.terminal.Load() || c.lifecycleState() != connLifecycleRetired { + t.Fatalf("closed Conn terminal=%v lifecycle=%v", c.terminal.Load(), c.lifecycleState()) + } + c.Close() +} + +func TestPhysicalCloseBitPreventsActivationClaimBeforeLogicalFence(t *testing.T) { + raw := newOwnershipTestTransport(false) + owner, lease := newPhysicalTransportOwner(raw) + s := New(Options{}) + c := s.newConnWithLease(lease, newTestAuthKey(t), 73002, 1) + + // Hold the binding lock so CloseAny can linearize the physical closed bit but + // cannot yet publish c.terminal. beginActivationClaim must inspect the lease + // itself and refuse this otherwise-dangerous window. + owner.bindingMu.Lock() + closeDone := make(chan error, 1) + go func() { closeDone <- owner.CloseAny() }() + deadline := time.Now().Add(time.Second) + for owner.state.Load()&physicalTransportClosedBit == 0 && time.Now().Before(deadline) { + time.Sleep(time.Millisecond) + } + if owner.state.Load()&physicalTransportClosedBit == 0 { + owner.bindingMu.Unlock() + t.Fatal("CloseAny did not publish closed bit") + } + if c.terminal.Load() { + owner.bindingMu.Unlock() + t.Fatal("logical fence escaped held binding lock") + } + if c.beginActivationClaim() { + owner.bindingMu.Unlock() + t.Fatal("closed physical generation entered activation claim") + } + owner.bindingMu.Unlock() + select { + case err := <-closeDone: + if err != nil { + t.Fatalf("CloseAny: %v", err) + } + case <-time.After(time.Second): + t.Fatal("CloseAny did not finish") + } + if !c.terminal.Load() || c.lifecycleState() != connLifecycleRetired { + t.Fatalf("logical fence terminal=%v lifecycle=%v", c.terminal.Load(), c.lifecycleState()) + } + c.Close() +} diff --git a/internal/rpc/channels_dialogs_rpc_test.go b/internal/rpc/channels_dialogs_rpc_test.go index 6cfedf45..bf89991d 100644 --- a/internal/rpc/channels_dialogs_rpc_test.go +++ b/internal/rpc/channels_dialogs_rpc_test.go @@ -5,6 +5,7 @@ import ( "github.com/gotd/td/bin" "github.com/gotd/td/clock" "github.com/gotd/td/tg" + "github.com/gotd/td/tgerr" "go.uber.org/zap/zaptest" "strconv" "strings" @@ -36,7 +37,8 @@ func TestLinkedDiscussionGuestCanCommentWithoutMembership(t *testing.T) { subscriber, _ := users.Create(ctx, domain.User{AccessHash: 302, Phone: "15550003002", FirstName: "Subscriber"}) channelStore := memory.NewChannelStore() channels := appchannels.NewService(channelStore, appchannels.WithBotProfileResolver(emptyDiscussionBotProfiles{})) - r := New(Config{}, Deps{Users: appusers.NewService(users), Channels: channels}, zaptest.NewLogger(t), clock.System) + dialogs := appdialogs.NewService(memory.NewDialogStore(), channelStore) + r := New(Config{}, Deps{Users: appusers.NewService(users), Channels: channels, Dialogs: dialogs}, zaptest.NewLogger(t), clock.System) broadcast, err := channels.CreateChannel(ctx, owner.ID, domain.CreateChannelRequest{Title: "Private channel", Broadcast: true, Date: 1700003001}) if err != nil { @@ -78,14 +80,11 @@ func TestLinkedDiscussionGuestCanCommentWithoutMembership(t *testing.T) { if _, err := r.onChannelsGetFullChannel(WithUserID(ctx, subscriber.ID), inputGroup); err != nil { t.Fatalf("get linked group full as guest: %v", err) } - participant, err := r.onChannelsGetParticipant(WithUserID(ctx, subscriber.ID), &tg.ChannelsGetParticipantRequest{ + _, err = r.onChannelsGetParticipant(WithUserID(ctx, subscriber.ID), &tg.ChannelsGetParticipantRequest{ Channel: inputGroup, Participant: &tg.InputPeerSelf{}, }) - if err != nil { - t.Fatalf("get linked guest participant: %v", err) - } - if _, ok := participant.Participant.(*tg.ChannelParticipantLeft); !ok { - t.Fatalf("linked guest participant = %T, want channelParticipantLeft", participant.Participant) + if !tgerr.Is(err, "USER_NOT_PARTICIPANT") { + t.Fatalf("get linked guest self participant err = %v, want USER_NOT_PARTICIPANT", err) } if _, err := r.onChannelsGetParticipants(WithUserID(ctx, subscriber.ID), &tg.ChannelsGetParticipantsRequest{ Channel: inputGroup, Filter: &tg.ChannelParticipantsRecent{}, Limit: 20, @@ -97,11 +96,68 @@ func TestLinkedDiscussionGuestCanCommentWithoutMembership(t *testing.T) { }); err != nil { t.Fatalf("get linked group bot participants as guest: %v", err) } - if _, err := r.onMessagesGetReplies(WithUserID(ctx, subscriber.ID), &tg.MessagesGetRepliesRequest{ + directReplies, err := r.onMessagesGetReplies(WithUserID(ctx, subscriber.ID), &tg.MessagesGetRepliesRequest{ Peer: &tg.InputPeerChannel{ChannelID: group.Channel.ID, AccessHash: group.Channel.AccessHash}, MsgID: root.ID, Limit: 20, - }); err != nil { + }) + if err != nil { t.Fatalf("get linked group replies as guest: %v", err) } + directPage, ok := directReplies.(*tg.MessagesChannelMessages) + if !ok || len(directPage.Chats) != 1 { + t.Fatalf("direct linked guest replies = %T %+v, want one channel chat", directReplies, directReplies) + } + directGroup, ok := directPage.Chats[0].(*tg.Channel) + if !ok || directGroup.ID != group.Channel.ID || directGroup.Min || !directGroup.Left { + t.Fatalf("direct linked guest replies chat = %T %+v, want full left group %d", directPage.Chats[0], directPage.Chats[0], group.Channel.ID) + } + viaBroadcastReplies, err := r.onMessagesGetReplies(WithUserID(ctx, subscriber.ID), &tg.MessagesGetRepliesRequest{ + Peer: &tg.InputPeerChannel{ChannelID: broadcast.Channel.ID, AccessHash: broadcast.Channel.AccessHash}, MsgID: post.ID, Limit: 20, + }) + if err != nil { + t.Fatalf("get linked replies through broadcast as guest: %v", err) + } + viaBroadcastPage, ok := viaBroadcastReplies.(*tg.MessagesChannelMessages) + if !ok || len(viaBroadcastPage.Chats) != 2 { + t.Fatalf("broadcast linked guest replies = %T %+v, want primary group plus source channel", viaBroadcastReplies, viaBroadcastReplies) + } + viaBroadcastGroup, ok := viaBroadcastPage.Chats[0].(*tg.Channel) + if !ok || viaBroadcastGroup.ID != group.Channel.ID || viaBroadcastGroup.Min || !viaBroadcastGroup.Left { + t.Fatalf("broadcast linked guest primary chat = %T %+v, want full left group %d", viaBroadcastPage.Chats[0], viaBroadcastPage.Chats[0], group.Channel.ID) + } + viaBroadcastSource, ok := viaBroadcastPage.Chats[1].(*tg.Channel) + if !ok || viaBroadcastSource.ID != broadcast.Channel.ID || !viaBroadcastSource.Min { + t.Fatalf("broadcast linked guest companion chat = %T %+v, want min source %d", viaBroadcastPage.Chats[1], viaBroadcastPage.Chats[1], broadcast.Channel.ID) + } + peerDialogsReq := &tg.MessagesGetPeerDialogsRequest{Peers: []tg.InputDialogPeerClass{ + &tg.InputDialogPeer{Peer: &tg.InputPeerChannel{ChannelID: group.Channel.ID, AccessHash: group.Channel.AccessHash}}, + }} + var peerDialogsIn bin.Buffer + if err := peerDialogsReq.Encode(&peerDialogsIn); err != nil { + t.Fatalf("encode linked guest getPeerDialogs: %v", err) + } + peerDialogsEnc, err := r.Dispatch(WithUserID(ctx, subscriber.ID), [8]byte{}, 0, &peerDialogsIn) + if err != nil { + t.Fatalf("dispatch linked guest getPeerDialogs: %v", err) + } + peerDialogs, ok := peerDialogsEnc.(*tg.MessagesPeerDialogs) + if !ok || len(peerDialogs.Dialogs) != 1 || len(peerDialogs.Chats) != 1 || len(peerDialogs.Messages) == 0 { + t.Fatalf("linked guest peer dialogs = %T %+v, want one transient dialog/chat with top message", peerDialogsEnc, peerDialogsEnc) + } + dialog, ok := peerDialogs.Dialogs[0].(*tg.Dialog) + if !ok || dialog.TopMessage == 0 { + t.Fatalf("linked guest dialog = %T %+v, want non-zero top message", peerDialogs.Dialogs[0], peerDialogs.Dialogs[0]) + } + guestDialogView, err := channels.GetChannel(ctx, subscriber.ID, group.Channel.ID) + if err != nil { + t.Fatalf("get linked guest view for dialog pts: %v", err) + } + if pts, ok := dialog.GetPts(); !ok || pts != guestDialogView.Channel.Pts { + t.Fatalf("linked guest dialog pts = %d (set=%v), want %d", pts, ok, guestDialogView.Channel.Pts) + } + chat, ok := peerDialogs.Chats[0].(*tg.Channel) + if !ok || !chat.Left || chat.ID != group.Channel.ID { + t.Fatalf("linked guest peer dialog chat = %T %+v, want left group %d", peerDialogs.Chats[0], peerDialogs.Chats[0], group.Channel.ID) + } if changed, err := r.onMessagesReadDiscussion(WithUserID(ctx, subscriber.ID), &tg.MessagesReadDiscussionRequest{ Peer: &tg.InputPeerChannel{ChannelID: group.Channel.ID, AccessHash: group.Channel.AccessHash}, MsgID: root.ID, ReadMaxID: group.Channel.TopMessageID, }); err != nil || changed { @@ -120,6 +176,112 @@ func TestLinkedDiscussionGuestCanCommentWithoutMembership(t *testing.T) { } } +func TestPublicChannelAndMegagroupPreviewStayReadableForNonMember(t *testing.T) { + ctx := context.Background() + users := memory.NewUserStore() + owner, _ := users.Create(ctx, domain.User{AccessHash: 401, Phone: "15550004001", FirstName: "Owner"}) + viewer, _ := users.Create(ctx, domain.User{AccessHash: 402, Phone: "15550004002", FirstName: "Viewer"}) + channels := appchannels.NewService(memory.NewChannelStore()) + r := New(Config{}, Deps{Users: appusers.NewService(users), Channels: channels}, zaptest.NewLogger(t), clock.System) + + tests := []struct { + name string + broadcast bool + username string + }{ + {name: "broadcast", broadcast: true, username: "public_rpc_broadcast"}, + {name: "megagroup", username: "public_rpc_megagroup"}, + } + for i, test := range tests { + t.Run(test.name, func(t *testing.T) { + var created domain.CreateChannelResult + var err error + if test.broadcast { + created, err = channels.CreateChannel(ctx, owner.ID, domain.CreateChannelRequest{Title: test.name, Broadcast: true, Date: 1700004000 + i}) + } else { + created, err = channels.CreateMegagroupFromCreateChat(ctx, owner.ID, domain.CreateChannelRequest{Title: test.name, Date: 1700004000 + i}) + } + if err != nil { + t.Fatalf("create public peer: %v", err) + } + public, err := channels.UpdateUsername(ctx, owner.ID, domain.UpdateChannelUsernameRequest{ + UserID: owner.ID, ChannelID: created.Channel.ID, Username: test.username, + }) + if err != nil { + t.Fatalf("make public: %v", err) + } + if _, err := channels.SendMessage(ctx, owner.ID, domain.SendChannelMessageRequest{ + ChannelID: public.ID, RandomID: int64(4000 + i), Message: "public history", Date: 1700004010 + i, + }); err != nil { + t.Fatalf("seed public history: %v", err) + } + inputChannel := &tg.InputChannel{ChannelID: public.ID, AccessHash: public.AccessHash} + full, err := r.onChannelsGetFullChannel(WithUserID(ctx, viewer.ID), inputChannel) + if err != nil { + t.Fatalf("getFullChannel preview: %v", err) + } + var preview *tg.Channel + for _, chat := range full.Chats { + if channel, ok := chat.(*tg.Channel); ok && channel.ID == public.ID { + preview = channel + } + } + if preview == nil || !preview.Left { + t.Fatalf("full chats = %+v, want public peer projected left", full.Chats) + } + history := dispatchMessagesPayload(t, r, WithUserID(ctx, viewer.ID), &tg.MessagesGetHistoryRequest{ + Peer: &tg.InputPeerChannel{ChannelID: public.ID, AccessHash: public.AccessHash}, Limit: 20, + }) + messages, _, _ := searchMessagesPayload(t, history) + if len(messages) == 0 { + t.Fatal("public preview history is empty") + } + if _, err := r.onChannelsGetParticipants(WithUserID(ctx, viewer.ID), &tg.ChannelsGetParticipantsRequest{ + Channel: inputChannel, Filter: &tg.ChannelParticipantsRecent{}, Limit: 20, + }); err != nil { + t.Fatalf("getParticipants preview: %v", err) + } + if _, err := r.onChannelsGetParticipant(WithUserID(ctx, viewer.ID), &tg.ChannelsGetParticipantRequest{ + Channel: inputChannel, Participant: &tg.InputPeerSelf{}, + }); !tgerr.Is(err, "USER_NOT_PARTICIPANT") { + t.Fatalf("self participant preview err = %v, want USER_NOT_PARTICIPANT", err) + } + if _, err := channels.GetChannel(ctx, viewer.ID, public.ID); err != nil { + t.Fatalf("public preview became forbidden after participant miss: %v", err) + } + if _, err := channels.JoinChannel(ctx, viewer.ID, public.ID, 1700004020+i); err != nil { + t.Fatalf("join public peer: %v", err) + } + if _, err := channels.LeaveChannel(ctx, viewer.ID, public.ID, 1700004030+i); err != nil { + t.Fatalf("leave public peer: %v", err) + } + if _, err := r.onChannelsGetParticipant(WithUserID(ctx, viewer.ID), &tg.ChannelsGetParticipantRequest{ + Channel: inputChannel, Participant: &tg.InputPeerSelf{}, + }); !tgerr.Is(err, "USER_NOT_PARTICIPANT") { + t.Fatalf("left self participant err = %v, want USER_NOT_PARTICIPANT", err) + } + if _, err := channels.GetHistory(ctx, viewer.ID, domain.ChannelHistoryFilter{ChannelID: public.ID, Limit: 20}); err != nil { + t.Fatalf("public history after leave: %v", err) + } + }) + } + + private, err := channels.CreateMegagroupFromCreateChat(ctx, owner.ID, domain.CreateChannelRequest{Title: "private", Date: 1700004050}) + if err != nil { + t.Fatalf("create private group: %v", err) + } + privateReq := &tg.MessagesGetHistoryRequest{ + Peer: &tg.InputPeerChannel{ChannelID: private.Channel.ID, AccessHash: private.Channel.AccessHash}, Limit: 20, + } + var privateIn bin.Buffer + if err := privateReq.Encode(&privateIn); err != nil { + t.Fatalf("encode private history: %v", err) + } + if _, err := r.Dispatch(WithUserID(ctx, viewer.ID), [8]byte{}, 0, &privateIn); !tgerr.Is(err, "CHANNEL_PRIVATE") { + t.Fatalf("private non-member history err = %v, want CHANNEL_PRIVATE", err) + } +} + func (s *countingDiscussionReadChannels) ResolveDiscussionReadTarget(ctx context.Context, userID, sourceChannelID int64, sourceMessageID, readMaxID int) (domain.ChannelDiscussionReadTarget, error) { s.resolveCalls++ return s.delegate.ResolveDiscussionReadTarget(ctx, userID, sourceChannelID, sourceMessageID, readMaxID) diff --git a/internal/rpc/channels_members.go b/internal/rpc/channels_members.go index 0deb431f..f8d2d742 100644 --- a/internal/rpc/channels_members.go +++ b/internal/rpc/channels_members.go @@ -206,7 +206,7 @@ func (r *Router) onChannelsGetParticipant(ctx context.Context, req *tg.ChannelsG } member, err := r.deps.Channels.GetParticipant(ctx, userID, channelID, peer.ID) if err != nil { - return nil, channelInvalidErr(err) + return nil, channelParticipantErr(err) } participant := tgChannelParticipant(userID, member) users := r.tgUsersForIDs(ctx, userID, channelParticipantUserRefs(participant)) @@ -217,6 +217,13 @@ func (r *Router) onChannelsGetParticipant(ctx context.Context, req *tg.ChannelsG }, nil } +func channelParticipantErr(err error) error { + if errors.Is(err, domain.ErrUserNotParticipant) { + return tgerr400("USER_NOT_PARTICIPANT") + } + return channelInvalidErr(err) +} + func channelParticipantUserRefs(participant tg.ChannelParticipantClass) []int64 { ids := make([]int64, 0, 3) add := func(id int64) { diff --git a/internal/rpc/convert_channels_core.go b/internal/rpc/convert_channels_core.go index 51868c0f..23d77ca6 100644 --- a/internal/rpc/convert_channels_core.go +++ b/internal/rpc/convert_channels_core.go @@ -22,8 +22,15 @@ func tgChannelChatsWithPrimarySelf(viewerUserID int64, primary domain.Channel, e continue } seen[ch.ID] = struct{}{} + if ch.ID != primary.ID { + // Companion channels have no matching viewer membership in + // ChannelHistory. Keep them minimal so they cannot overwrite the + // client's cached left/creator/admin/banned state. + out = append(out, tgChannelChatMin(viewerUserID, ch)) + continue + } var self *domain.ChannelMember - if ch.ID == primary.ID && primarySelf.ChannelID == ch.ID && primarySelf.UserID == viewerUserID { + if primarySelf.ChannelID == ch.ID && primarySelf.UserID == viewerUserID { self = &primarySelf } out = append(out, tgChannelChat(viewerUserID, ch, self)) diff --git a/internal/store/memory/channel_linked_guest_authorization_test.go b/internal/store/memory/channel_linked_guest_authorization_test.go new file mode 100644 index 00000000..8bcffb7c --- /dev/null +++ b/internal/store/memory/channel_linked_guest_authorization_test.go @@ -0,0 +1,119 @@ +package memory + +import ( + "context" + "errors" + "testing" + + "telesrv/internal/domain" +) + +func TestLinkedDiscussionGuestAuthorizationFailsClosedMemory(t *testing.T) { + ctx := context.Background() + channels := NewChannelStore() + const ( + ownerID int64 = 9101 + outsiderID int64 = 9102 + ) + privateGroup, err := channels.CreateChannel(ctx, domain.CreateChannelRequest{ + CreatorUserID: ownerID, + Title: "unlinked private group", + Megagroup: true, + Date: 1700009100, + }) + if err != nil { + t.Fatalf("create private group: %v", err) + } + if _, err := channels.SendChannelMessage(ctx, domain.SendChannelMessageRequest{ + UserID: outsiderID, + ChannelID: privateGroup.Channel.ID, + RandomID: 91001, + Message: "must be rejected", + Date: 1700009101, + }); !errors.Is(err, domain.ErrChannelPrivate) { + t.Fatalf("unlinked outsider send err = %v, want ErrChannelPrivate", err) + } + if _, err := channels.ResolveDiscussionReadTarget(ctx, outsiderID, privateGroup.Channel.ID, privateGroup.Message.ID, privateGroup.Message.ID); !errors.Is(err, domain.ErrChannelPrivate) { + t.Fatalf("unlinked outsider discussion read err = %v, want ErrChannelPrivate", err) + } + if got := len(channels.messages[privateGroup.Channel.ID]); got != 1 { + t.Fatalf("unlinked outsider changed history length = %d, want create message only", got) + } +} + +func TestLinkedDiscussionGuestRepliesKeepViewerProjectionMemory(t *testing.T) { + ctx := context.Background() + channels := NewChannelStore() + const ( + ownerID int64 = 9201 + subscriberID int64 = 9202 + ) + broadcast, err := channels.CreateChannel(ctx, domain.CreateChannelRequest{ + CreatorUserID: ownerID, + Title: "linked source", + Broadcast: true, + Date: 1700009200, + }) + if err != nil { + t.Fatalf("create broadcast: %v", err) + } + group, err := channels.CreateChannel(ctx, domain.CreateChannelRequest{ + CreatorUserID: ownerID, + Title: "linked discussion", + Megagroup: true, + Date: 1700009201, + }) + if err != nil { + t.Fatalf("create group: %v", err) + } + if _, err := channels.SetDiscussionGroup(ctx, ownerID, broadcast.Channel.ID, group.Channel.ID); err != nil { + t.Fatalf("set discussion group: %v", err) + } + if _, err := channels.InviteToChannel(ctx, broadcast.Channel.ID, ownerID, []int64{subscriberID}, 1700009202); err != nil { + t.Fatalf("invite broadcast subscriber: %v", err) + } + post, err := channels.SendChannelMessage(ctx, domain.SendChannelMessageRequest{ + UserID: ownerID, ChannelID: broadcast.Channel.ID, RandomID: 92001, + Message: "post", Date: 1700009203, + }) + if err != nil || post.Discussion == nil { + t.Fatalf("send linked post = %+v err %v", post, err) + } + + assertGuest := func(name string, history domain.ChannelHistory, err error) { + t.Helper() + if err != nil { + t.Fatalf("%s: %v", name, err) + } + if history.Channel.ID != group.Channel.ID || history.Self.ChannelID != group.Channel.ID || history.Self.UserID != subscriberID || !history.Self.Guest || history.Self.Status != domain.ChannelMemberLeft { + t.Fatalf("%s = channel %+v self %+v, want linked left guest", name, history.Channel, history.Self) + } + } + direct, err := channels.ListChannelReplies(ctx, subscriberID, domain.ChannelRepliesFilter{ + ChannelID: group.Channel.ID, RootMessageID: post.Discussion.Message.ID, Limit: 20, + }) + assertGuest("direct replies", direct, err) + viaBroadcast, err := channels.ListChannelReplies(ctx, subscriberID, domain.ChannelRepliesFilter{ + ChannelID: broadcast.Channel.ID, RootMessageID: post.Message.ID, Limit: 20, + }) + assertGuest("broadcast replies", viaBroadcast, err) + + if _, err := channels.EditChannelBanned(ctx, domain.EditChannelBannedRequest{ + UserID: ownerID, ChannelID: group.Channel.ID, + Participant: domain.Peer{Type: domain.PeerTypeUser, ID: subscriberID}, + BannedRights: domain.ChannelBannedRights{ViewMessages: true, UntilDate: 2147483647}, + Date: 1700009204, + }); err != nil { + t.Fatalf("ban linked subscriber from target: %v", err) + } + if _, err := channels.ListChannelReplies(ctx, subscriberID, domain.ChannelRepliesFilter{ + ChannelID: group.Channel.ID, RootMessageID: post.Discussion.Message.ID, Limit: 20, + }); !errors.Is(err, domain.ErrChannelUserBanned) { + t.Fatalf("target-banned direct replies err = %v, want ErrChannelUserBanned", err) + } + if _, err := channels.ListChannelReplies(ctx, subscriberID, domain.ChannelRepliesFilter{ + ChannelID: broadcast.Channel.ID, RootMessageID: post.Message.ID, Limit: 20, + }); !errors.Is(err, domain.ErrChannelUserBanned) { + t.Fatalf("target-banned broadcast replies err = %v, want ErrChannelUserBanned", err) + } +} diff --git a/internal/store/memory/channel_members.go b/internal/store/memory/channel_members.go index e63ce9fb..f0d4cf48 100644 --- a/internal/store/memory/channel_members.go +++ b/internal/store/memory/channel_members.go @@ -12,7 +12,7 @@ import ( func (s *ChannelStore) GetParticipants(_ context.Context, viewerUserID, channelID int64, filter domain.ChannelParticipantsFilter, offset, limit int) (domain.ChannelParticipantList, error) { s.mu.RLock() defer s.mu.RUnlock() - channel, viewer, err := s.channelAndMemberOrLinkedGuestLocked(viewerUserID, channelID) + channel, viewer, _, err := s.channelForViewerLocked(viewerUserID, channelID) if err != nil { return domain.ChannelParticipantList{}, err } @@ -75,17 +75,16 @@ func (s *ChannelStore) GetParticipants(_ context.Context, viewerUserID, channelI func (s *ChannelStore) GetParticipant(_ context.Context, viewerUserID, channelID, participantUserID int64) (domain.ChannelMember, error) { s.mu.RLock() defer s.mu.RUnlock() - channel, viewer, err := s.channelAndMemberOrLinkedGuestLocked(viewerUserID, channelID) + _, viewer, _, err := s.channelForViewerLocked(viewerUserID, channelID) if err != nil { return domain.ChannelMember{}, err } if viewerUserID == participantUserID && viewer.Guest { - return viewer, nil + return domain.ChannelMember{}, domain.ErrUserNotParticipant } - _ = channel member, ok := s.members[channelID][participantUserID] - if !ok { - return domain.ChannelMember{}, domain.ErrChannelPrivate + if !ok || (participantUserID == viewerUserID && member.Status == domain.ChannelMemberLeft) { + return domain.ChannelMember{}, domain.ErrUserNotParticipant } return member, nil } @@ -868,7 +867,7 @@ func (s *ChannelStore) FilterActiveChannelMemberIDs(_ context.Context, channelID func (s *ChannelStore) ListActiveChannelMembers(_ context.Context, viewerUserID, channelID int64, limit int) (domain.Channel, domain.ChannelMember, []domain.ChannelMember, error) { s.mu.RLock() defer s.mu.RUnlock() - channel, viewer, err := s.channelAndMemberOrLinkedGuestLocked(viewerUserID, channelID) + channel, viewer, _, err := s.channelForViewerLocked(viewerUserID, channelID) if err != nil { return domain.Channel{}, domain.ChannelMember{}, nil, err } diff --git a/internal/store/memory/channel_message_history.go b/internal/store/memory/channel_message_history.go index f2199b9d..824033dd 100644 --- a/internal/store/memory/channel_message_history.go +++ b/internal/store/memory/channel_message_history.go @@ -443,7 +443,16 @@ func (s *ChannelStore) ResolveDiscussionReadTarget(_ context.Context, viewerUser if targetChannelID != sourceChannelID { _, targetMember, err = s.channelAndMemberLocked(viewerUserID, targetChannelID) if errors.Is(err, domain.ErrChannelPrivate) { - targetMember, _, err = s.linkedDiscussionGuestLocked(viewerUserID, s.channels[targetChannelID]) + guestMember, allowed, guestErr := s.linkedDiscussionGuestLocked(viewerUserID, s.channels[targetChannelID]) + switch { + case guestErr != nil: + err = guestErr + case allowed: + targetMember = guestMember + err = nil + default: + err = domain.ErrChannelPrivate + } } } if err != nil { diff --git a/internal/store/memory/channel_message_send.go b/internal/store/memory/channel_message_send.go index 263408d7..b5ebcd7c 100644 --- a/internal/store/memory/channel_message_send.go +++ b/internal/store/memory/channel_message_send.go @@ -41,10 +41,18 @@ func (s *ChannelStore) SendChannelMessage(_ context.Context, req domain.SendChan channel, member, err := s.channelAndMemberLocked(req.UserID, req.ChannelID) if errors.Is(err, domain.ErrChannelPrivate) { if candidate, ok := s.channels[req.ChannelID]; ok && !candidate.Deleted { - var guest bool - member, guest, err = s.linkedDiscussionGuestLocked(req.UserID, candidate) - if guest { + guestMember, guest, guestErr := s.linkedDiscussionGuestLocked(req.UserID, candidate) + switch { + case guestErr != nil: + err = guestErr + case guest: channel = candidate + member = guestMember + err = nil + default: + // A clean "not a linked guest" result is not authorization. + // Preserve the original private-member error and fail closed. + err = domain.ErrChannelPrivate } } } diff --git a/internal/store/memory/channel_topics.go b/internal/store/memory/channel_topics.go index d0a1baf1..99bd8eef 100644 --- a/internal/store/memory/channel_topics.go +++ b/internal/store/memory/channel_topics.go @@ -513,23 +513,19 @@ func (s *ChannelStore) ListChannelReplies(_ context.Context, viewerUserID int64, extraChannels := []domain.Channel(nil) if source.Broadcast { if root.Discussion == nil || root.Discussion.ChannelID == 0 || root.Discussion.MessageID == 0 { - return domain.ChannelHistory{Channel: source, Count: 0}, nil + return domain.ChannelHistory{Channel: source, Self: member, Count: 0}, nil } - linked, ok := s.channels[root.Discussion.ChannelID] - if !ok || linked.Deleted { - return domain.ChannelHistory{Channel: source, Count: 0}, nil + linked, linkedMember, linkedErr := s.channelAndMemberOrLinkedGuestLocked(viewerUserID, root.Discussion.ChannelID) + if linkedErr != nil { + return domain.ChannelHistory{}, linkedErr } targetChannel = linked + targetMember = linkedMember rootID = root.Discussion.MessageID - if linkedMember, ok := s.members[linked.ID][viewerUserID]; ok { - targetMember = linkedMember - } else { - targetMember = domain.ChannelMember{} - } extraChannels = append(extraChannels, source) } if targetRoot, ok := s.findMessageLocked(targetChannel.ID, rootID); !ok || targetRoot.Deleted { - return domain.ChannelHistory{Channel: targetChannel, Channels: extraChannels, Count: 0}, nil + return domain.ChannelHistory{Channel: targetChannel, Self: targetMember, Channels: extraChannels, Count: 0}, nil } limit := filter.Limit if limit <= 0 || limit > domain.MaxChannelRepliesLimit { @@ -570,7 +566,7 @@ func (s *ChannelStore) ListChannelReplies(_ context.Context, viewerUserID int64, topics = append(topics, cloneChannelForumTopic(topic)) } } - return domain.ChannelHistory{Channel: targetChannel, Channels: extraChannels, Topics: topics, Messages: out, Count: len(base)}, nil + return domain.ChannelHistory{Channel: targetChannel, Self: targetMember, Channels: extraChannels, Topics: topics, Messages: out, Count: len(base)}, nil } func (s *ChannelStore) populateChannelMessageRepliesLocked(viewerUserID, channelID int64, messages []domain.ChannelMessage) { diff --git a/internal/store/postgres/channel_core.go b/internal/store/postgres/channel_core.go index 4814fea8..c3a5363a 100644 --- a/internal/store/postgres/channel_core.go +++ b/internal/store/postgres/channel_core.go @@ -377,7 +377,20 @@ WHERE c.id = ANY($2::bigint[]) AND NOT c.deleted`, viewerUserID, ids) if err != nil { return nil, err } + linkedGuests, err := s.listLinkedDiscussionGuests(ctx, s.db, viewerUserID, remaining) + if err != nil { + return nil, err + } for _, channel := range channels { + if member, ok := linkedGuests[channel.ID]; ok { + views[channel.ID] = domain.ChannelView{ + Channel: channel, + Self: member, + Dialog: previewChannelDialog(viewerUserID, channel, member), + SelfBoostsApplied: 0, + } + continue + } if member, _, ok, err := s.monoforumAdminPreview(ctx, s.db, viewerUserID, channel); err != nil { return nil, err } else if ok { diff --git a/internal/store/postgres/channel_member_helpers.go b/internal/store/postgres/channel_member_helpers.go index 9abf479e..1092e2fe 100644 --- a/internal/store/postgres/channel_member_helpers.go +++ b/internal/store/postgres/channel_member_helpers.go @@ -96,6 +96,72 @@ func (s *ChannelStore) getLinkedDiscussionGuest(ctx context.Context, db sqlcgen. return guest, true, nil } +// listLinkedDiscussionGuests is the bounded batch equivalent of +// getLinkedDiscussionGuest. GetChannels uses it after loading the requested +// channel rows so messages.getPeerDialogs can materialize linked discussion +// histories without an N+1 query per requested peer. +// +// The target membership predicate deliberately excludes active, kicked, +// banned and view-messages-banned rows. Active members were projected by the +// primary GetChannels query; explicit target denial must always win over the +// source broadcast membership. Returned members are transient and are never +// persisted to channel_members/channel_dialogs. +func (s *ChannelStore) listLinkedDiscussionGuests(ctx context.Context, db sqlcgen.DBTX, viewerUserID int64, targetIDs []int64) (map[int64]domain.ChannelMember, error) { + out := make(map[int64]domain.ChannelMember) + if viewerUserID == 0 || len(targetIDs) == 0 { + return out, nil + } + rows, err := db.Query(ctx, ` +SELECT target.id +FROM channels target +JOIN channels source + ON source.id = target.linked_chat_id + AND NOT source.deleted + AND source.broadcast + AND source.linked_chat_id = target.id +JOIN channel_members source_member + ON source_member.channel_id = source.id + AND source_member.user_id = $1 + AND source_member.status = 'active' + AND NOT COALESCE((source_member.banned_rights->>'ViewMessages')::boolean, false) +LEFT JOIN channel_members target_member + ON target_member.channel_id = target.id + AND target_member.user_id = $1 +WHERE target.id = ANY($2::bigint[]) + AND NOT target.deleted + AND target.megagroup + AND NOT target.broadcast + AND ( + target_member.user_id IS NULL + OR ( + target_member.status NOT IN ('active', 'banned', 'kicked') + AND NOT COALESCE((target_member.banned_rights->>'ViewMessages')::boolean, false) + ) + ) +ORDER BY target.id`, viewerUserID, targetIDs) + if err != nil { + return nil, fmt.Errorf("list linked discussion guests: %w", err) + } + defer rows.Close() + for rows.Next() { + var channelID int64 + if err := rows.Scan(&channelID); err != nil { + return nil, fmt.Errorf("scan linked discussion guest: %w", err) + } + out[channelID] = domain.ChannelMember{ + ChannelID: channelID, + UserID: viewerUserID, + Status: domain.ChannelMemberLeft, + Role: domain.ChannelRoleMember, + Guest: true, + } + } + if err := rows.Err(); err != nil { + return nil, fmt.Errorf("iterate linked discussion guests: %w", err) + } + return out, nil +} + func (s *ChannelStore) getPublicPreviewMember(ctx context.Context, db sqlcgen.DBTX, viewerUserID int64, ch domain.Channel) (domain.ChannelMember, error) { member, err := s.getChannelMember(ctx, db, ch.ID, viewerUserID) if err != nil { diff --git a/internal/store/postgres/channel_member_list.go b/internal/store/postgres/channel_member_list.go index 1e0a60bc..94188cf3 100644 --- a/internal/store/postgres/channel_member_list.go +++ b/internal/store/postgres/channel_member_list.go @@ -2,6 +2,7 @@ package postgres import ( "context" + "errors" "fmt" "sort" "strings" @@ -12,7 +13,7 @@ import ( ) func (s *ChannelStore) GetParticipants(ctx context.Context, viewerUserID, channelID int64, filter domain.ChannelParticipantsFilter, offset, limit int) (domain.ChannelParticipantList, error) { - channel, viewer, err := s.getChannelForMemberOrLinkedGuest(ctx, s.db, viewerUserID, channelID) + channel, viewer, _, err := s.getChannelForViewer(ctx, s.db, viewerUserID, channelID) if err != nil { return domain.ChannelParticipantList{}, err } @@ -170,14 +171,21 @@ WHERE channel_id = $1 } func (s *ChannelStore) GetParticipant(ctx context.Context, viewerUserID, channelID, participantUserID int64) (domain.ChannelMember, error) { - _, viewer, err := s.getChannelForMemberOrLinkedGuest(ctx, s.db, viewerUserID, channelID) + _, viewer, _, err := s.getChannelForViewer(ctx, s.db, viewerUserID, channelID) if err != nil { return domain.ChannelMember{}, err } if viewerUserID == participantUserID && viewer.Guest { - return viewer, nil + return domain.ChannelMember{}, domain.ErrUserNotParticipant } - return s.getChannelMember(ctx, s.db, channelID, participantUserID) + member, err := s.getChannelMember(ctx, s.db, channelID, participantUserID) + if errors.Is(err, domain.ErrChannelPrivate) { + return domain.ChannelMember{}, domain.ErrUserNotParticipant + } + if err == nil && participantUserID == viewerUserID && member.Status == domain.ChannelMemberLeft { + return domain.ChannelMember{}, domain.ErrUserNotParticipant + } + return member, err } func (s *ChannelStore) ListActiveChannelMemberIDs(ctx context.Context, viewerUserID, channelID int64, limit int) ([]int64, error) { @@ -204,7 +212,7 @@ func (s *ChannelStore) ListActiveChannelMemberIDs(ctx context.Context, viewerUse } func (s *ChannelStore) ListActiveChannelMembers(ctx context.Context, viewerUserID, channelID int64, limit int) (domain.Channel, domain.ChannelMember, []domain.ChannelMember, error) { - channel, viewer, err := s.getChannelForMemberOrLinkedGuest(ctx, s.db, viewerUserID, channelID) + channel, viewer, _, err := s.getChannelForViewer(ctx, s.db, viewerUserID, channelID) if err != nil { return domain.Channel{}, domain.ChannelMember{}, nil, err } @@ -237,7 +245,7 @@ LIMIT $2`, channelID, limit) } func (s *ChannelStore) ListActiveChannelBotMembers(ctx context.Context, viewerUserID, channelID int64, offset, limit int) (domain.ChannelParticipantList, error) { - channel, viewer, err := s.getChannelForMemberOrLinkedGuest(ctx, s.db, viewerUserID, channelID) + channel, viewer, _, err := s.getChannelForViewer(ctx, s.db, viewerUserID, channelID) if err != nil { return domain.ChannelParticipantList{}, err } diff --git a/internal/store/postgres/channel_message_send.go b/internal/store/postgres/channel_message_send.go index 7563e2bf..615cb541 100644 --- a/internal/store/postgres/channel_message_send.go +++ b/internal/store/postgres/channel_message_send.go @@ -66,10 +66,18 @@ func (s *ChannelStore) sendChannelMessageOnce(ctx context.Context, req domain.Se channel, member, err := s.getChannelForMember(ctx, tx, req.UserID, req.ChannelID) if errors.Is(err, domain.ErrChannelPrivate) { if candidate, candidateErr := s.channelByID(ctx, tx, req.ChannelID); candidateErr == nil { - var guest bool - member, guest, err = s.getLinkedDiscussionGuest(ctx, tx, req.UserID, candidate) - if guest { + guestMember, guest, guestErr := s.getLinkedDiscussionGuest(ctx, tx, req.UserID, candidate) + switch { + case guestErr != nil: + err = guestErr + case guest: channel = candidate + member = guestMember + err = nil + default: + // A clean "not a linked guest" result is not authorization. + // Preserve the original private-member error and fail closed. + err = domain.ErrChannelPrivate } } else { err = candidateErr diff --git a/internal/store/postgres/channel_public_preview_integration_test.go b/internal/store/postgres/channel_public_preview_integration_test.go new file mode 100644 index 00000000..4a7657a5 --- /dev/null +++ b/internal/store/postgres/channel_public_preview_integration_test.go @@ -0,0 +1,263 @@ +package postgres + +import ( + "context" + "errors" + "testing" + + appdialogs "telesrv/internal/app/dialogs" + "telesrv/internal/domain" +) + +func TestPublicChannelAndMegagroupPreviewPostgres(t *testing.T) { + pool := testPool(t) + ctx := context.Background() + suffix := randomSuffix(t) + users := NewUserStore(pool) + owner, err := users.Create(ctx, domain.User{AccessHash: 941, Phone: "+1941" + suffix + "01", FirstName: "PreviewOwner"}) + if err != nil { + t.Fatalf("create owner: %v", err) + } + viewer, err := users.Create(ctx, domain.User{AccessHash: 942, Phone: "+1942" + suffix + "02", FirstName: "PreviewViewer"}) + if err != nil { + t.Fatalf("create viewer: %v", err) + } + channels := NewChannelStore(pool) + var channelIDs []int64 + t.Cleanup(func() { + if len(channelIDs) != 0 { + _, _ = pool.Exec(ctx, "DELETE FROM channels WHERE id = ANY($1::bigint[])", channelIDs) + } + _, _ = pool.Exec(ctx, "DELETE FROM users WHERE id = ANY($1::bigint[])", []int64{owner.ID, viewer.ID}) + }) + + for i, tc := range []struct { + name string + broadcast bool + }{ + {name: "broadcast", broadcast: true}, + {name: "megagroup"}, + } { + t.Run(tc.name, func(t *testing.T) { + created, err := channels.CreateChannel(ctx, domain.CreateChannelRequest{ + CreatorUserID: owner.ID, + Title: "Public Preview " + tc.name + " " + suffix, + Broadcast: tc.broadcast, + Megagroup: !tc.broadcast, + Date: 1700009400 + i, + }) + if err != nil { + t.Fatalf("create channel: %v", err) + } + channelIDs = append(channelIDs, created.Channel.ID) + public, err := channels.UpdateUsername(ctx, domain.UpdateChannelUsernameRequest{ + UserID: owner.ID, ChannelID: created.Channel.ID, Username: "pub" + tc.name + suffix, + }) + if err != nil { + t.Fatalf("make public: %v", err) + } + sent, err := channels.SendChannelMessage(ctx, domain.SendChannelMessageRequest{ + UserID: owner.ID, ChannelID: public.ID, RandomID: int64(94100 + i), Message: "public history", Date: 1700009410 + i, + }) + if err != nil { + t.Fatalf("send public message: %v", err) + } + history, err := channels.ListChannelHistory(ctx, viewer.ID, domain.ChannelHistoryFilter{ChannelID: public.ID, Limit: 20}) + if err != nil { + t.Fatalf("public preview history: %v", err) + } + found := false + for _, message := range history.Messages { + if message.ID == sent.Message.ID { + found = true + } + } + if !found || history.Self.Status != domain.ChannelMemberLeft { + t.Fatalf("preview history = %+v self=%+v", history.Messages, history.Self) + } + if _, err := channels.GetParticipants(ctx, viewer.ID, public.ID, domain.ChannelParticipantsFilter{Kind: domain.ChannelParticipantsRecent}, 0, 20); err != nil { + t.Fatalf("public preview participants: %v", err) + } + if _, err := channels.GetParticipant(ctx, viewer.ID, public.ID, viewer.ID); !errors.Is(err, domain.ErrUserNotParticipant) { + t.Fatalf("public preview self participant err = %v, want ErrUserNotParticipant", err) + } + var memberExists bool + if err := pool.QueryRow(ctx, `SELECT EXISTS ( +SELECT 1 FROM channel_members WHERE channel_id = $1 AND user_id = $2 +)`, public.ID, viewer.ID).Scan(&memberExists); err != nil { + t.Fatalf("check preview member row: %v", err) + } + if memberExists { + t.Fatal("public preview persisted a channel member row") + } + if _, err := channels.JoinChannel(ctx, public.ID, viewer.ID, 1700009420+i); err != nil { + t.Fatalf("join public peer: %v", err) + } + if _, err := channels.LeaveChannel(ctx, public.ID, viewer.ID, 1700009430+i); err != nil { + t.Fatalf("leave public peer: %v", err) + } + if _, err := channels.GetParticipant(ctx, viewer.ID, public.ID, viewer.ID); !errors.Is(err, domain.ErrUserNotParticipant) { + t.Fatalf("left self participant err = %v, want ErrUserNotParticipant", err) + } + if _, err := channels.ListChannelHistory(ctx, viewer.ID, domain.ChannelHistoryFilter{ChannelID: public.ID, Limit: 20}); err != nil { + t.Fatalf("public history after leave: %v", err) + } + }) + } + + private, err := channels.CreateChannel(ctx, domain.CreateChannelRequest{ + CreatorUserID: owner.ID, Title: "Private Preview " + suffix, Megagroup: true, Date: 1700009450, + }) + if err != nil { + t.Fatalf("create private group: %v", err) + } + channelIDs = append(channelIDs, private.Channel.ID) + if _, err := channels.ListChannelHistory(ctx, viewer.ID, domain.ChannelHistoryFilter{ChannelID: private.Channel.ID, Limit: 20}); !errors.Is(err, domain.ErrChannelPrivate) { + t.Fatalf("private preview history err = %v, want ErrChannelPrivate", err) + } +} + +func TestLinkedDiscussionGuestPeerDialogProjectionPostgres(t *testing.T) { + pool := testPool(t) + ctx := context.Background() + suffix := randomSuffix(t) + users := NewUserStore(pool) + owner, err := users.Create(ctx, domain.User{AccessHash: 951, Phone: "+1951" + suffix + "01", FirstName: "DiscussionOwner"}) + if err != nil { + t.Fatalf("create owner: %v", err) + } + subscriber, err := users.Create(ctx, domain.User{AccessHash: 952, Phone: "+1952" + suffix + "02", FirstName: "DiscussionSubscriber"}) + if err != nil { + t.Fatalf("create subscriber: %v", err) + } + outsider, err := users.Create(ctx, domain.User{AccessHash: 953, Phone: "+1953" + suffix + "03", FirstName: "DiscussionOutsider"}) + if err != nil { + t.Fatalf("create outsider: %v", err) + } + channels := NewChannelStore(pool, + WithChannelRowCache(NewChannelRowCache(32)), + WithChannelMemberCache(NewChannelMemberCache(64))) + var channelIDs []int64 + t.Cleanup(func() { + if len(channelIDs) != 0 { + _, _ = pool.Exec(ctx, "DELETE FROM channels WHERE id = ANY($1::bigint[])", channelIDs) + } + _, _ = pool.Exec(ctx, "DELETE FROM users WHERE id = ANY($1::bigint[])", []int64{owner.ID, subscriber.ID, outsider.ID}) + }) + + broadcast, err := channels.CreateChannel(ctx, domain.CreateChannelRequest{ + CreatorUserID: owner.ID, Title: "Peer Dialog Source " + suffix, Broadcast: true, Date: 1700009500, + }) + if err != nil { + t.Fatalf("create broadcast: %v", err) + } + channelIDs = append(channelIDs, broadcast.Channel.ID) + group, err := channels.CreateChannel(ctx, domain.CreateChannelRequest{ + CreatorUserID: owner.ID, Title: "Peer Dialog Group " + suffix, Megagroup: true, Date: 1700009501, + }) + if err != nil { + t.Fatalf("create discussion group: %v", err) + } + channelIDs = append(channelIDs, group.Channel.ID) + if _, err := channels.SetDiscussionGroup(ctx, owner.ID, broadcast.Channel.ID, group.Channel.ID); err != nil { + t.Fatalf("set discussion group: %v", err) + } + if _, err := channels.InviteToChannel(ctx, broadcast.Channel.ID, owner.ID, []int64{subscriber.ID}, 1700009502); err != nil { + t.Fatalf("invite broadcast subscriber: %v", err) + } + post, err := channels.SendChannelMessage(ctx, domain.SendChannelMessageRequest{ + UserID: owner.ID, ChannelID: broadcast.Channel.ID, RandomID: 95101, Message: "peer dialog root", Date: 1700009503, + }) + if err != nil || post.Discussion == nil { + t.Fatalf("send linked post = %+v err %v", post, err) + } + + views, err := channels.GetChannels(ctx, subscriber.ID, []int64{group.Channel.ID}) + if err != nil || len(views) != 1 { + t.Fatalf("batch linked guest views = %+v err %v, want one", views, err) + } + view := views[0] + if !view.Self.Guest || view.Self.Status != domain.ChannelMemberLeft || view.Dialog.TopMessageID == 0 || view.Channel.Pts == 0 { + t.Fatalf("batch linked guest view = %+v dialog=%+v channel=%+v", view.Self, view.Dialog, view.Channel) + } + dialogs := appdialogs.NewService(nil, channels) + peerDialogs, err := dialogs.GetPeerDialogs(ctx, subscriber.ID, []domain.Peer{{Type: domain.PeerTypeChannel, ID: group.Channel.ID}}) + if err != nil { + t.Fatalf("get linked guest peer dialogs: %v", err) + } + if len(peerDialogs.Dialogs) != 1 || len(peerDialogs.Channels) != 1 || len(peerDialogs.ChannelMessages) == 0 { + t.Fatalf("linked guest peer dialogs = %+v, want transient dialog/channel/top message", peerDialogs) + } + if peerDialogs.Dialogs[0].TopMessage == 0 || peerDialogs.Dialogs[0].Pts != view.Channel.Pts { + t.Fatalf("linked guest dialog = %+v, want top message and pts %d", peerDialogs.Dialogs[0], view.Channel.Pts) + } + directReplies, err := channels.ListChannelReplies(ctx, subscriber.ID, domain.ChannelRepliesFilter{ + ChannelID: group.Channel.ID, RootMessageID: post.Discussion.Message.ID, Limit: 20, + }) + if err != nil || !directReplies.Self.Guest || directReplies.Self.Status != domain.ChannelMemberLeft || directReplies.Channel.ID != group.Channel.ID { + t.Fatalf("direct linked guest replies = %+v err %v, want guest self for group", directReplies, err) + } + viaBroadcastReplies, err := channels.ListChannelReplies(ctx, subscriber.ID, domain.ChannelRepliesFilter{ + ChannelID: broadcast.Channel.ID, RootMessageID: post.Message.ID, Limit: 20, + }) + if err != nil || !viaBroadcastReplies.Self.Guest || viaBroadcastReplies.Self.Status != domain.ChannelMemberLeft || viaBroadcastReplies.Channel.ID != group.Channel.ID { + t.Fatalf("broadcast linked guest replies = %+v err %v, want guest self for target group", viaBroadcastReplies, err) + } + + outsiderViews, err := channels.GetChannels(ctx, outsider.ID, []int64{group.Channel.ID}) + if err != nil || len(outsiderViews) != 0 { + t.Fatalf("private discussion outsider views = %+v err %v, want empty", outsiderViews, err) + } + outsiderDialogs, err := dialogs.GetPeerDialogs(ctx, outsider.ID, []domain.Peer{{Type: domain.PeerTypeChannel, ID: group.Channel.ID}}) + if err != nil || len(outsiderDialogs.Dialogs) != 0 { + t.Fatalf("private discussion outsider dialogs = %+v err %v, want empty", outsiderDialogs, err) + } + if _, err := channels.InviteToChannel(ctx, broadcast.Channel.ID, owner.ID, []int64{outsider.ID}, 1700009504); err != nil { + t.Fatalf("invite second broadcast subscriber: %v", err) + } + if _, err := channels.EditChannelBanned(ctx, domain.EditChannelBannedRequest{ + UserID: owner.ID, + ChannelID: group.Channel.ID, + Participant: domain.Peer{Type: domain.PeerTypeUser, ID: outsider.ID}, + BannedRights: domain.ChannelBannedRights{ + ViewMessages: true, + UntilDate: 2147483647, + }, + Date: 1700009505, + }); err != nil { + t.Fatalf("ban linked subscriber from target group: %v", err) + } + bannedViews, err := channels.GetChannels(ctx, outsider.ID, []int64{group.Channel.ID}) + if err != nil || len(bannedViews) != 1 || !bannedViews[0].Forbidden || bannedViews[0].Self.Guest { + t.Fatalf("target-banned linked subscriber views = %+v err %v, want forbidden non-guest", bannedViews, err) + } + bannedDialogs, err := dialogs.GetPeerDialogs(ctx, outsider.ID, []domain.Peer{{Type: domain.PeerTypeChannel, ID: group.Channel.ID}}) + if !errors.Is(err, domain.ErrChannelUserBanned) || len(bannedDialogs.Dialogs) != 0 { + t.Fatalf("target-banned linked subscriber dialogs = %+v err %v, want ErrChannelUserBanned without dialog", bannedDialogs, err) + } + if _, err := channels.ListChannelReplies(ctx, outsider.ID, domain.ChannelRepliesFilter{ + ChannelID: group.Channel.ID, RootMessageID: post.Discussion.Message.ID, Limit: 20, + }); !errors.Is(err, domain.ErrChannelUserBanned) { + t.Fatalf("target-banned direct replies err = %v, want ErrChannelUserBanned", err) + } + if _, err := channels.ListChannelReplies(ctx, outsider.ID, domain.ChannelRepliesFilter{ + ChannelID: broadcast.Channel.ID, RootMessageID: post.Message.ID, Limit: 20, + }); !errors.Is(err, domain.ErrChannelUserBanned) { + t.Fatalf("target-banned broadcast replies err = %v, want ErrChannelUserBanned", err) + } + + var memberExists, dialogExists bool + if err := pool.QueryRow(ctx, `SELECT EXISTS ( +SELECT 1 FROM channel_members WHERE channel_id = $1 AND user_id = $2 +)`, group.Channel.ID, subscriber.ID).Scan(&memberExists); err != nil { + t.Fatalf("check transient guest member: %v", err) + } + if err := pool.QueryRow(ctx, `SELECT EXISTS ( +SELECT 1 FROM channel_dialogs WHERE channel_id = $1 AND user_id = $2 +)`, group.Channel.ID, subscriber.ID).Scan(&dialogExists); err != nil { + t.Fatalf("check transient guest dialog: %v", err) + } + if memberExists || dialogExists { + t.Fatalf("transient guest persisted state: member=%v dialog=%v", memberExists, dialogExists) + } +} diff --git a/internal/store/postgres/channel_topics.go b/internal/store/postgres/channel_topics.go index 181dd5e2..53426b81 100644 --- a/internal/store/postgres/channel_topics.go +++ b/internal/store/postgres/channel_topics.go @@ -595,28 +595,27 @@ func (s *ChannelStore) ListChannelReplies(ctx context.Context, viewerUserID int6 return domain.ChannelHistory{}, domain.ErrMessageIDInvalid } target := source - availableMinID := member.AvailableMinID + targetMember := member + availableMinID := targetMember.AvailableMinID extraChannels := []domain.Channel(nil) rootID := root.ID if source.Broadcast { if root.Discussion == nil || root.Discussion.ChannelID == 0 || root.Discussion.MessageID == 0 { - return domain.ChannelHistory{Channel: source}, nil + return domain.ChannelHistory{Channel: source, Self: member}, nil } - linked, err := getChannelByID(ctx, s.db, root.Discussion.ChannelID) + linked, linkedMember, err := s.getChannelForMemberOrLinkedGuest(ctx, s.db, viewerUserID, root.Discussion.ChannelID) if err != nil { - return domain.ChannelHistory{Channel: source}, nil + return domain.ChannelHistory{}, err } target = linked + targetMember = linkedMember rootID = root.Discussion.MessageID - availableMinID = 0 - if linkedMember, err := s.getChannelMember(ctx, s.db, linked.ID, viewerUserID); err == nil && validateChannelMemberVisible(linkedMember) == nil { - availableMinID = linkedMember.AvailableMinID - } + availableMinID = targetMember.AvailableMinID extraChannels = append(extraChannels, source) } targetRoot, err := s.getChannelMessage(ctx, s.db, target.ID, rootID) if err != nil || targetRoot.Deleted || targetRoot.ID <= availableMinID { - return domain.ChannelHistory{Channel: target, Channels: extraChannels}, nil + return domain.ChannelHistory{Channel: target, Self: targetMember, Channels: extraChannels}, nil } limit := filter.Limit if limit <= 0 || limit > domain.MaxChannelRepliesLimit { @@ -646,7 +645,7 @@ func (s *ChannelStore) ListChannelReplies(ctx context.Context, viewerUserID int6 return domain.ChannelHistory{}, err } } - return domain.ChannelHistory{Channel: target, Channels: extraChannels, Topics: topics, Messages: messages, Count: count}, nil + return domain.ChannelHistory{Channel: target, Self: targetMember, Channels: extraChannels, Topics: topics, Messages: messages, Count: count}, nil } func (s *ChannelStore) getForumTopic(ctx context.Context, db sqlcgen.DBTX, channelID int64, topicID int) (domain.ChannelForumTopic, error) {