astra/owpengram-server
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

feat: sync multilayer td integration

Browse source
This commit is contained in:
A 2026-07-15 13:32:06 +08:00
parent 20a310f6ca
commit 766c5db992
491 changed files with 26235 additions and 35340 deletions
Download patch file Download diff file Expand all files Collapse all files

4
internal/mtprotoedge/admission_test.go
View file

@ -9,8 +9,8 @@ import (
"testing"
"time"
"github.com/gotd/td/bin"
"github.com/gotd/td/proto/codec"
"github.com/iamxvbaba/td/bin"
"github.com/iamxvbaba/td/proto/codec"
"telesrv/internal/store"
"telesrv/internal/store/memory"

40
internal/mtprotoedge/auth_key_expiry_test.go
View file

@ -12,15 +12,15 @@ import (
"go.uber.org/zap/zaptest"
"github.com/gotd/log/logzap"
"github.com/gotd/td/bin"
"github.com/gotd/td/clock"
"github.com/gotd/td/crypto"
"github.com/gotd/td/exchange"
"github.com/gotd/td/mt"
"github.com/gotd/td/proto"
"github.com/gotd/td/proto/codec"
"github.com/gotd/td/tg"
"github.com/gotd/td/transport"
"github.com/iamxvbaba/td/bin"
"github.com/iamxvbaba/td/clock"
"github.com/iamxvbaba/td/crypto"
"github.com/iamxvbaba/td/exchange"
"github.com/iamxvbaba/td/mt"
"github.com/iamxvbaba/td/proto"
"github.com/iamxvbaba/td/proto/codec"
"github.com/iamxvbaba/td/tg"
"github.com/iamxvbaba/td/transport"
)
func TestAuthKeyProtocolUnavailable(t *testing.T) {
@ -118,9 +118,9 @@ func TestActiveTemporaryAuthKeyExpiresBeforeNextRPCDispatch(t *testing.T) {
testClock := newExpiryTestClock(now)
handler := &admissionCountingRPC{}
addr, pub, srv := startTestServer(t, Options{
DC: dc,
Clock: testClock,
RPC: handler,
DC: dc,
Clock: testClock,
legacyRPC: handler,
})
conn, auth, cipher := dialTemporaryHandshakeForExpiryTest(t, addr, dc, expiresIn, pub)
@ -179,7 +179,7 @@ func TestExpiredTemporaryAuthKeyRejectsServerPushWithoutWireWrite(t *testing.T)
c.authKeyExpiresAt = int(now.Unix())
err := c.SendBestEffortEncoded(context.Background(), proto.MessageFromServer,
&encodedOutboundMessage{typeID: tg.UpdatesTooLongTypeID, body: []byte{0x0b, 0xa1, 0x01, 0xe3}}, 0)
exactTestUpdatesTooLong(t, c), 0)
if !errors.Is(err, ErrConnClosed) {
t.Fatalf("push on expired temp key = %v, want ErrConnClosed", err)
}
@ -198,7 +198,7 @@ func TestQueuedPushCannotCrossTemporaryAuthKeyExpiry(t *testing.T) {
c := newOutboundTestConn(t, tr, nil)
c.now = clock.Now
c.authKeyExpiresAt = int(now.Add(time.Minute).Unix())
encoded := &encodedOutboundMessage{typeID: tg.UpdatesTooLongTypeID, body: []byte{0x0b, 0xa1, 0x01, 0xe3}}
encoded := exactTestUpdatesTooLong(t, c)
if err := c.SendBestEffortEncoded(context.Background(), proto.MessageFromServer, encoded, 0); err != nil {
t.Fatalf("enqueue first push: %v", err)
@ -247,7 +247,7 @@ func TestTemporaryAuthKeyExpiryWhileWaitingForPhysicalWriterSkipsRawSend(t *test
t.Fatal("direct protocol write did not acquire physical writer")
}
encoded := &encodedOutboundMessage{typeID: tg.UpdatesTooLongTypeID, body: []byte{0x0b, 0xa1, 0x01, 0xe3}}
encoded := exactTestUpdatesTooLong(t, c)
actorDone := make(chan error, 1)
go func() {
actorDone <- c.SendEncoded(context.Background(), proto.MessageFromServer, encoded)
@ -295,9 +295,9 @@ func TestRetiredActorWaitingForPhysicalWriterDoesNotDefeatLeaseTransfer(t *testi
}
actorDone := make(chan error, 1)
encoded := exactTestUpdatesTooLong(t, c)
go func() {
actorDone <- c.SendEncoded(context.Background(), proto.MessageFromServer,
&encodedOutboundMessage{typeID: tg.UpdatesTooLongTypeID, body: []byte{0x0b, 0xa1, 0x01, 0xe3}})
actorDone <- c.SendEncoded(context.Background(), proto.MessageFromServer, encoded)
}()
select {
case <-signaling.entered:
@ -350,9 +350,9 @@ func TestTerminalAuthKeyNotFoundSurvivesActorWaitingForPhysicalWriter(t *testing
case <-time.After(time.Second):
t.Fatal("direct protocol write did not acquire physical writer")
}
encoded := exactTestUpdatesTooLong(t, c)
go func() {
_ = c.SendEncoded(context.Background(), proto.MessageFromServer,
&encodedOutboundMessage{typeID: tg.UpdatesTooLongTypeID, body: []byte{0x0b, 0xa1, 0x01, 0xe3}})
_ = c.SendEncoded(context.Background(), proto.MessageFromServer, encoded)
}()
select {
case <-signaling.entered:
@ -402,7 +402,7 @@ func TestTerminalAuthKeyNotFoundWaitsForOutboundAndIsLastFrame(t *testing.T) {
c.transportLease = lease
c.now = clock.Now
c.authKeyExpiresAt = int(now.Add(time.Minute).Unix())
encoded := &encodedOutboundMessage{typeID: tg.UpdatesTooLongTypeID, body: []byte{0x0b, 0xa1, 0x01, 0xe3}}
encoded := exactTestUpdatesTooLong(t, c)
if err := c.SendBestEffortEncoded(context.Background(), proto.MessageFromServer, encoded, 0); err != nil {
t.Fatalf("enqueue blocked push: %v", err)
}

4
internal/mtprotoedge/auth_key_switch_test.go
View file

@ -4,8 +4,8 @@ import (
"testing"
"time"
"github.com/gotd/td/mt"
"github.com/gotd/td/proto"
"github.com/iamxvbaba/td/mt"
"github.com/iamxvbaba/td/proto"
)
func TestEncryptedConnectionSwitchesAuthKeyEvenWhenSessionIDIsReused(t *testing.T) {

29
internal/mtprotoedge/bot_callback_e2e_test.go
View file

@ -12,13 +12,13 @@ import (
"go.uber.org/zap/zaptest"
"github.com/gotd/log/logzap"
"github.com/gotd/td/clock"
"github.com/gotd/td/exchange"
"github.com/gotd/td/session"
"github.com/gotd/td/telegram"
"github.com/gotd/td/telegram/dcs"
"github.com/gotd/td/tg"
"github.com/gotd/td/transport"
"github.com/iamxvbaba/td/clock"
"github.com/iamxvbaba/td/exchange"
"github.com/iamxvbaba/td/session"
"github.com/iamxvbaba/td/telegram"
"github.com/iamxvbaba/td/telegram/dcs"
"github.com/iamxvbaba/td/tg"
"github.com/iamxvbaba/td/transport"
"telesrv/internal/app/account"
"telesrv/internal/app/auth"
@ -84,8 +84,19 @@ func newBotCallbackEnv(t *testing.T, ctx context.Context) *botCallbackEnv {
router := rpc.New(rpc.Config{DC: dc, IP: tcpAddr.IP.String(), Port: tcpAddr.Port}, deps, zaptest.NewLogger(t), clock.System)
botsService.SetRouterHooks(router)
botsService.SetTextDraftPusher(router)
srv := New(Options{Logger: zaptest.NewLogger(t), DC: dc, RSAKey: rsaKey, AuthKeys: authKeyStore, RPC: router, ActiveSessions: activeSessions})
go func() { _ = srv.Serve(ctx, ln) }()
srv := New(Options{Logger: zaptest.NewLogger(t), DC: dc, RSAKey: rsaKey, AuthKeys: authKeyStore, LayerRPC: router, ActiveSessions: activeSessions})
serveErr := make(chan error, 1)
go func() { serveErr <- srv.Serve(ctx, ln) }()
t.Cleanup(func() {
select {
case err := <-serveErr:
if err != nil {
t.Errorf("serve: %v", err)
}
case <-time.After(5 * time.Second):
t.Error("server did not stop after callback test context cancellation")
}
})
newCli := func(storage *session.StorageMemory, handler telegram.UpdateHandler) *telegram.Client {
if handler == nil {

20
internal/mtprotoedge/bot_e2e_test.go
View file

@ -13,14 +13,14 @@ import (
"go.uber.org/zap/zaptest"
"github.com/gotd/log/logzap"
"github.com/gotd/td/clock"
"github.com/gotd/td/exchange"
"github.com/gotd/td/session"
"github.com/gotd/td/telegram"
"github.com/gotd/td/telegram/dcs"
"github.com/gotd/td/tg"
"github.com/gotd/td/tgerr"
"github.com/gotd/td/transport"
"github.com/iamxvbaba/td/clock"
"github.com/iamxvbaba/td/exchange"
"github.com/iamxvbaba/td/session"
"github.com/iamxvbaba/td/telegram"
"github.com/iamxvbaba/td/telegram/dcs"
"github.com/iamxvbaba/td/tg"
"github.com/iamxvbaba/td/tgerr"
"github.com/iamxvbaba/td/transport"
"telesrv/internal/app/account"
"telesrv/internal/app/auth"
@ -87,7 +87,7 @@ func TestBotManagementRPCFlow(t *testing.T) {
router := rpc.New(rpc.Config{DC: dc, IP: tcpAddr.IP.String(), Port: tcpAddr.Port}, deps, zaptest.NewLogger(t), clock.System)
botsService.SetRouterHooks(router)
botsService.SetTextDraftPusher(router)
srv := New(Options{Logger: zaptest.NewLogger(t), DC: dc, RSAKey: rsaKey, AuthKeys: authKeyStore, RPC: router, ActiveSessions: activeSessions})
srv := New(Options{Logger: zaptest.NewLogger(t), DC: dc, RSAKey: rsaKey, AuthKeys: authKeyStore, LayerRPC: router, ActiveSessions: activeSessions})
ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
defer cancel()
@ -323,7 +323,7 @@ func TestBotFatherCreateAndBotLoginFlow(t *testing.T) {
router := rpc.New(rpc.Config{DC: dc, IP: tcpAddr.IP.String(), Port: tcpAddr.Port}, deps, zaptest.NewLogger(t), clock.System)
botsService.SetRouterHooks(router)
botsService.SetTextDraftPusher(router)
srv := New(Options{Logger: zaptest.NewLogger(t), DC: dc, RSAKey: rsaKey, AuthKeys: authKeyStore, RPC: router, ActiveSessions: activeSessions})
srv := New(Options{Logger: zaptest.NewLogger(t), DC: dc, RSAKey: rsaKey, AuthKeys: authKeyStore, LayerRPC: router, ActiveSessions: activeSessions})
ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
defer cancel()

117
internal/mtprotoedge/conn.go
View file

@ -4,16 +4,16 @@ import (
"bufio"
"context"
"encoding/hex"
"errors"
"sync"
"sync/atomic"
"time"
"github.com/gotd/td/bin"
"github.com/gotd/td/crypto"
"github.com/gotd/td/proto"
"github.com/gotd/td/transport"
"telesrv/internal/compat/layerwire"
"github.com/iamxvbaba/td/bin"
"github.com/iamxvbaba/td/crypto"
"github.com/iamxvbaba/td/proto"
"github.com/iamxvbaba/td/tg"
"github.com/iamxvbaba/td/transport"
)
// Conn 是一个已识别 session 的客户端连接，持有向其加密发送消息所需的全部上下文。
@ -39,6 +39,21 @@ const (
connLifecycleRetired
)
var (
// ErrLayerProfileUnsupported means the requested profile has no generated
// exact codec. Profiles are never clamped to the nearest supported layer.
ErrLayerProfileUnsupported = errors.New("unsupported exact layer profile")
// ErrLayerProfileConflict means one admitted request carried contradictory
// profile evidence. A later, well-formed invokeWithLayer is allowed to correct
// the connection profile and therefore does not use this error.
ErrLayerProfileConflict = errors.New("connection layer profile conflict")
// ErrLayerProfileEpochExhausted is a defensive terminal guard. Reaching it
// would require more than four billion effective layer corrections on one
// physical connection, so wrapping the epoch and making stale pushes current
// again is never safe.
ErrLayerProfileEpochExhausted = errors.New("connection layer profile epoch exhausted")
)
type Conn struct {
transport transport.Conn
// transportLease owns exactly one generation of the physical transport.
@ -114,14 +129,22 @@ type Conn struct {
rpcRunning int
rpcReady bool
rpcClosed bool
// rpcReplayRestores is a per-physical-connection ordering barrier. An exact
// cached/rewrapped init request has already executed its business handler,
// but its wrapper/client/readiness state becomes authoritative only after the
// replacement rpc_result is physically written. Queued naked RPCs remain
// admitted and budgeted, but are not scheduler-runnable until every such
// restore finishes or the connection is fenced.
rpcReplayRestores int
// Rewrap aliasing never delays execution. initialized stops collecting
// candidates after the first valid init wrapper on this physical generation.
rpcRewrapInitialized atomic.Bool
// rpcResultAcked is invoked by the sole outbound actor after it resolves an
// acknowledged server frame back to the rpc_result request msg_id.
rpcResultAcked func(*Conn, int64)
// inflightRPCBytes 跟踪已入队未完成的 inbound RPC body 总字节，配合 maxInflightRPCBytes
// 给 RPC 队列设字节预算（不止限条数），防对抗客户端发大请求撑内存。
// inflightRPCBytes 跟踪已预留/入队/执行中 inbound RPC 的 memory charge；legacy
// 等于 copied body，exact 是 typed materialization 的保守放大值。它配合
// maxInflightRPCBytes 给 RPC 队列设内存预算（不止限条数）。
inflightRPCBytes atomic.Int64
// 单连接只保留并发配额；实际 worker 来自 Server 共享池，避免每连接预留 goroutine。
rpcRootCtx context.Context
@ -152,9 +175,26 @@ type Conn struct {
membershipGen atomic.Int64
// createdAt 是连接建立时刻，供同 auth_key session 数触顶时驱逐真正最旧的连接。
createdAt time.Time
// clientLayer 是本连接协商的 TL layer（invokeWithLayer/initConnection），由 handleRPC
// 在每次 Dispatch 后从 RPC 注册表刷新。出站(rpc_result/push)按此把 227 对象降级给老客户端；
// 0 表示尚未协商，按 canonical(227) 处理=不降级。
// layerProfileState atomically packs profile, provenance and epoch. A profile
// inherited from auth-key metadata is only a default; a later well-formed
// invokeWithLayer may correct it. The epoch fences proactive updates prepared
// before that correction without invalidating request-bound RPC results.
layerProfileMu sync.RWMutex
// layerProfileEvidenceMsgID is protected by layerProfileMu. Zero means the
// selected profile came from inherited/legacy recovery and therefore has no
// ordered client-message cursor yet. Positive values are the newest accepted
// invokeWithLayer message for this exact MTProto session.
layerProfileEvidenceMsgID int64
// layerProfileEvidenceLayer retains the raw negotiated Layer even when this
// binary has no generated codec for it. In that case the packed profile stays
// Unknown, but msg_id ordering can still admit a newer supported correction
// and reject an older/same-id rollback.
layerProfileEvidenceLayer int
layerProfileState atomic.Uint64
// clientLayer is the package-internal mirror used only by legacy state-machine
// regression tests. Production application RPC/result/update encoding uses
// the structured profile state. Zero is unknown and must fail closed if a legacy
// application value reaches an outbound boundary.
clientLayer atomic.Int32
}
@ -236,16 +276,55 @@ 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 {
return int(l)
}
return layerwire.CanonicalLayer
// LayerProfile returns the exact TL profile currently selected for this
// connection. ok is false until admission or an inherited auth-key default
// supplies a supported generated profile.
func (c *Conn) LayerProfile() (profile tg.LayerProfile, ok bool) {
state := c.LayerProfileState()
return state.Profile, state.Origin != LayerProfileUnknown
}
// SetClientLayer 记录连接协商的 TL layer。
func (c *Conn) SetClientLayer(layer int) { c.clientLayer.Store(int32(layer)) }
// FreezeLayerProfile records explicit protocol evidence observed during ordered
// admission. Repeating the same value is idempotent. A later well-formed
// invokeWithLayer may replace either an inherited default or older explicit
// evidence; already-admitted requests retain their own immutable profile.
func (c *Conn) FreezeLayerProfile(profile tg.LayerProfile) error {
_, err := c.setLayerProfile(profile, LayerProfileExplicit, true)
return err
}
// FreezeLayerProfileAt applies explicit protocol evidence in client msg_id
// order. A duplicate older than the last accepted evidence is inert; the same
// msg_id carrying another Layer is a protocol conflict. Advancing the evidence
// cursor at an unchanged Layer does not rotate the outbound epoch because the
// wire profile itself did not change.
func (c *Conn) FreezeLayerProfileAt(profile tg.LayerProfile, msgID int64) (bool, error) {
return c.freezeLayerProfileAt(profile, msgID)
}
// SeedLayerProfile restores explicit evidence previously proven for this exact
// logical session. It is kept as the compatible same-session restore API;
// auth-key-wide metadata must use SeedInheritedLayerProfile instead.
func (c *Conn) SeedLayerProfile(profile tg.LayerProfile) error {
_, err := c.setLayerProfile(profile, LayerProfileExplicit, true)
return err
}
// SeedInheritedLayerProfile installs an auth-key-wide default only while the
// connection is still unknown. It never overwrites explicit evidence or an
// already selected inherited default; client protocol evidence owns correction.
func (c *Conn) SeedInheritedLayerProfile(profile tg.LayerProfile) error {
_, err := c.setLayerProfile(profile, LayerProfileInherited, false)
return err
}
// legacyClientLayer returns the test-only canonical-transcoder profile. Zero
// deliberately remains unknown; it must never be converted into an implicit
// canonical application profile.
func (c *Conn) legacyClientLayer() int { return int(c.clientLayer.Load()) }
// setLegacyClientLayer records the package-internal legacy mirror.
func (c *Conn) setLegacyClientLayer(layer int) { c.clientLayer.Store(int32(layer)) }
// AuthKeyID 返回连接的 auth_key_id。
func (c *Conn) AuthKeyID() [8]byte { return c.authKeyID }

506
internal/mtprotoedge/conn_layer_profile.go Normal file
View file

@ -0,0 +1,506 @@
package mtprotoedge
import (
"context"
"fmt"
"math"
"github.com/iamxvbaba/td/tg"
"go.uber.org/zap"
)
// LayerProfileOrigin records why a Conn currently uses a wire profile. The
// distinction is protocol-significant: inherited auth-key metadata is a useful
// availability default, while explicit invokeWithLayer evidence may correct it.
type LayerProfileOrigin uint8
const (
LayerProfileUnknown LayerProfileOrigin = iota
LayerProfileInherited
LayerProfileExplicit
)
// LayerProfileSnapshot is one atomic observation of a connection's current
// profile. Epoch advances on every effective correction, including promotion
// from inherited to explicit evidence at the same numeric layer.
type LayerProfileSnapshot struct {
Profile tg.LayerProfile
Origin LayerProfileOrigin
Epoch uint32
}
const (
layerProfileValueBits = 16
layerProfileOriginBits = 8
layerProfileOriginShift = layerProfileValueBits
layerProfileEpochShift = 32
layerProfileValueMask = uint64(1<<layerProfileValueBits - 1)
layerProfileOriginMask = uint64(1<<layerProfileOriginBits - 1)
)
func packLayerProfileState(state LayerProfileSnapshot) uint64 {
return uint64(state.Epoch)<<layerProfileEpochShift |
(uint64(state.Origin)&layerProfileOriginMask)<<layerProfileOriginShift |
(uint64(state.Profile) & layerProfileValueMask)
}
func unpackLayerProfileState(raw uint64) LayerProfileSnapshot {
if raw == 0 {
return LayerProfileSnapshot{}
}
return LayerProfileSnapshot{
Profile: tg.LayerProfile(raw & layerProfileValueMask),
Origin: LayerProfileOrigin((raw >> layerProfileOriginShift) & layerProfileOriginMask),
Epoch: uint32(raw >> layerProfileEpochShift),
}
}
// LayerProfileState returns profile, provenance and epoch from one atomic load.
func (c *Conn) LayerProfileState() LayerProfileSnapshot {
if c == nil {
return LayerProfileSnapshot{}
}
return unpackLayerProfileState(c.layerProfileState.Load())
}
func (c *Conn) setLayerProfile(profile tg.LayerProfile, origin LayerProfileOrigin, replace bool) (bool, error) {
if err := validateLayerProfile(profile); err != nil {
return false, err
}
if origin != LayerProfileInherited && origin != LayerProfileExplicit {
return false, fmt.Errorf("invalid layer profile origin %d", origin)
}
c.layerProfileMu.Lock()
defer c.layerProfileMu.Unlock()
for {
raw := c.layerProfileState.Load()
current := unpackLayerProfileState(raw)
if origin == LayerProfileInherited && c.layerProfileEvidenceMsgID > 0 {
return false, nil
}
if current.Origin != LayerProfileUnknown {
if !replace {
return false, nil
}
if current.Profile == profile && current.Origin == origin {
// The force-style compatibility APIs carry no ordered msg_id.
// Clearing the cursor keeps a subsequent production observation
// eligible instead of comparing it with unrelated test/recovery state.
c.layerProfileEvidenceMsgID = 0
c.layerProfileEvidenceLayer = 0
return false, nil
}
}
if current.Epoch == math.MaxUint32 {
return false, ErrLayerProfileEpochExhausted
}
next := LayerProfileSnapshot{Profile: profile, Origin: origin, Epoch: current.Epoch + 1}
if c.layerProfileState.CompareAndSwap(raw, packLayerProfileState(next)) {
c.layerProfileEvidenceMsgID = 0
c.layerProfileEvidenceLayer = 0
c.setLegacyClientLayer(int(profile))
return true, nil
}
}
}
func validateLayerProfile(profile tg.LayerProfile) error {
resolved, ok := tg.ResolveLayerProfile(int(profile))
if !ok || resolved != profile || uint64(profile) > layerProfileValueMask {
return fmt.Errorf("%w: %d", ErrLayerProfileUnsupported, profile)
}
return nil
}
// layerProfileEvidenceState observes the packed wire state and its ordering
// cursor under one read lock. LayerProfileState remains the allocation-free
// atomic hot-path accessor used by outbound encoding.
func (c *Conn) layerProfileEvidenceState() (LayerProfileSnapshot, int64) {
state, _, msgID := c.layerProfileRawEvidenceState()
return state, msgID
}
func (c *Conn) layerProfileRawEvidenceState() (LayerProfileSnapshot, int, int64) {
if c == nil {
return LayerProfileSnapshot{}, 0, 0
}
c.layerProfileMu.RLock()
state := unpackLayerProfileState(c.layerProfileState.Load())
layer := c.layerProfileEvidenceLayer
msgID := c.layerProfileEvidenceMsgID
c.layerProfileMu.RUnlock()
if layer == 0 && state.Origin == LayerProfileExplicit {
layer = int(state.Profile)
}
return state, layer, msgID
}
// freezeLayerProfileAt is the production explicit-evidence transition. The
// positive client msg_id is the protocol ordering authority across TCP
// reconnects and cached request replays.
func (c *Conn) freezeLayerProfileAt(profile tg.LayerProfile, msgID int64) (bool, error) {
if c == nil {
return false, fmt.Errorf("nil connection layer profile")
}
if msgID <= 0 {
return false, fmt.Errorf("invalid layer evidence msg_id %d", msgID)
}
if err := validateLayerProfile(profile); err != nil {
return false, err
}
return c.freezeRawLayerProfileAt(int(profile), msgID)
}
func (c *Conn) freezeRawLayerProfileAt(layer int, msgID int64) (bool, error) {
if c == nil {
return false, fmt.Errorf("nil connection layer profile")
}
if layer <= 0 || msgID <= 0 {
return false, fmt.Errorf("invalid raw layer evidence layer=%d msg_id=%d", layer, msgID)
}
profile, supported := tg.ResolveLayerProfile(layer)
c.layerProfileMu.Lock()
defer c.layerProfileMu.Unlock()
current := unpackLayerProfileState(c.layerProfileState.Load())
if c.layerProfileEvidenceMsgID > 0 {
switch {
case msgID < c.layerProfileEvidenceMsgID:
return false, nil
case msgID == c.layerProfileEvidenceMsgID:
if c.layerProfileEvidenceLayer != layer {
return false, fmt.Errorf("%w: msg_id %d selected both layer %d and %d", ErrLayerProfileConflict, msgID, c.layerProfileEvidenceLayer, layer)
}
return false, nil
}
}
desired := LayerProfileSnapshot{Epoch: current.Epoch}
if supported {
desired.Profile = profile
desired.Origin = LayerProfileExplicit
}
// A newer proof at the same Layer advances only the ordering cursor. No wire
// bytes prepared under this profile become stale, so rotating epoch would be
// unnecessary push churn.
if current.Profile == desired.Profile && current.Origin == desired.Origin {
c.layerProfileEvidenceLayer = layer
c.layerProfileEvidenceMsgID = msgID
return true, nil
}
if current.Epoch == math.MaxUint32 {
return false, ErrLayerProfileEpochExhausted
}
next := desired
next.Epoch = current.Epoch + 1
c.layerProfileState.Store(packLayerProfileState(next))
c.layerProfileEvidenceLayer = layer
c.layerProfileEvidenceMsgID = msgID
if supported {
c.setLegacyClientLayer(layer)
} else {
c.setLegacyClientLayer(0)
}
return true, nil
}
// seedOrderedLayerProfile restores exact-session evidence atomically before
// any request on a replacement physical connection is admitted.
func (c *Conn) seedOrderedLayerProfile(profile tg.LayerProfile, msgID int64) error {
if c == nil {
return nil
}
if msgID < 0 {
return fmt.Errorf("invalid restored layer evidence msg_id %d", msgID)
}
if err := validateLayerProfile(profile); err != nil {
return err
}
if msgID > 0 {
_, err := c.freezeRawLayerProfileAt(int(profile), msgID)
return err
}
c.layerProfileMu.Lock()
defer c.layerProfileMu.Unlock()
current := unpackLayerProfileState(c.layerProfileState.Load())
if current.Profile != profile || current.Origin != LayerProfileExplicit {
if current.Epoch == math.MaxUint32 {
return ErrLayerProfileEpochExhausted
}
current = LayerProfileSnapshot{Profile: profile, Origin: LayerProfileExplicit, Epoch: current.Epoch + 1}
c.layerProfileState.Store(packLayerProfileState(current))
c.setLegacyClientLayer(int(profile))
}
c.layerProfileEvidenceMsgID = msgID
c.layerProfileEvidenceLayer = 0
return nil
}
func (c *Conn) seedRawLayerEvidence(layer int, msgID int64) error {
if layer <= 0 || msgID <= 0 {
return fmt.Errorf("invalid restored raw layer evidence layer=%d msg_id=%d", layer, msgID)
}
_, err := c.freezeRawLayerProfileAt(layer, msgID)
return err
}
// refreshInheritedLayerProfile is reserved for identity normalization at
// auth.bindTempAuthKey: once a raw temporary key is resolved to its permanent
// key, the permanent key's default supersedes an older raw-key shadow. Explicit
// evidence on the concrete session is never overwritten.
func (c *Conn) refreshInheritedLayerProfile(profile tg.LayerProfile) (bool, error) {
if c == nil {
return false, nil
}
if err := validateLayerProfile(profile); err != nil {
return false, err
}
c.layerProfileMu.Lock()
defer c.layerProfileMu.Unlock()
current := unpackLayerProfileState(c.layerProfileState.Load())
if current.Origin == LayerProfileExplicit || c.layerProfileEvidenceMsgID > 0 {
return false, nil
}
if current.Origin == LayerProfileInherited && current.Profile == profile {
return false, nil
}
if current.Epoch == math.MaxUint32 {
return false, ErrLayerProfileEpochExhausted
}
next := LayerProfileSnapshot{Profile: profile, Origin: LayerProfileInherited, Epoch: current.Epoch + 1}
c.layerProfileState.Store(packLayerProfileState(next))
c.layerProfileEvidenceMsgID = 0
c.layerProfileEvidenceLayer = 0
c.setLegacyClientLayer(int(profile))
return true, nil
}
func (c *Conn) clearInheritedLayerProfileState() (bool, error) {
if c == nil {
return false, nil
}
c.layerProfileMu.Lock()
defer c.layerProfileMu.Unlock()
current := unpackLayerProfileState(c.layerProfileState.Load())
if current.Origin != LayerProfileInherited {
return false, nil
}
if current.Epoch == math.MaxUint32 {
return false, ErrLayerProfileEpochExhausted
}
c.layerProfileState.Store(packLayerProfileState(LayerProfileSnapshot{Epoch: current.Epoch + 1}))
c.layerProfileEvidenceMsgID = 0
c.layerProfileEvidenceLayer = 0
c.setLegacyClientLayer(0)
return true, nil
}
func (c *Conn) clearInheritedLayerProfile() error {
_, err := c.clearInheritedLayerProfileState()
return err
}
// seedInitialLayerProfile applies recovery sources in descending authority.
// Unsupported metadata never clamps to a nearby generated layer: the Conn stays
// unknown so the client can explicitly renegotiate.
func (s *Server) seedInitialLayerProfile(
ctx context.Context,
c *Conn,
fetchedLayer int,
previous LayerProfileSnapshot,
) error {
if s == nil || c == nil {
return nil
}
durableResolver, hasDurableResolver := s.layerRPC.(LayerRPCDurableSessionProfileResolver)
if hasDurableResolver {
layer, msgID, found, err := durableResolver.ResolveNegotiatedSessionLayerEvidence(ctx, c.authKeyID, c.sessionID)
if err != nil {
if isLayerEvidenceDurabilityUnavailable(err) {
s.log.Warn("Resolve durable exact session Layer during connection seed unavailable; continuing with auth-key default",
zap.String("auth_key_id", c.authKeyHex), zap.Int64("session_id", c.sessionID), zap.Error(err))
// Exact-session proof is the strongest recovery source, but its
// availability failure must not discard a permanent auth_keys.layer
// already loaded with the key on this same first frame. Continue to the
// auth-key default below; never reuse previous connection-local evidence
// in durable mode.
found = false
} else {
return fmt.Errorf("resolve durable exact session Layer during connection seed: %w", err)
}
}
if found {
if layer <= 0 || msgID < 0 {
return fmt.Errorf("invalid durable exact session Layer seed layer=%d msg_id=%d", layer, msgID)
}
if msgID > 0 {
return c.seedRawLayerEvidence(layer, msgID)
}
// Older in-process exact-session registries did not retain a message
// watermark. Keep that compatibility-only seed usable without treating
// it as durable ordered evidence; real durable stores never persist zero.
profile, supported := tg.ResolveLayerProfile(layer)
if !supported {
return nil
}
return c.seedOrderedLayerProfile(profile, 0)
}
} else if resolver, ok := s.layerRPC.(LayerRPCOrderedSessionProfileResolver); ok {
if layer, msgID, found := resolver.NegotiatedSessionLayerEvidence(c.authKeyID, c.sessionID); found {
if layer <= 0 || msgID < 0 {
return nil
}
if msgID > 0 {
return c.seedRawLayerEvidence(layer, msgID)
}
profile, supported := tg.ResolveLayerProfile(layer)
if !supported {
return nil
}
return c.seedOrderedLayerProfile(profile, 0)
}
} else if resolver, ok := s.layerRPC.(LayerRPCSessionProfileResolver); ok {
if layer, found := resolver.NegotiatedSessionLayer(c.authKeyID, c.sessionID); found {
profile, supported := tg.ResolveLayerProfile(layer)
if !supported {
return nil
}
return c.SeedLayerProfile(profile)
}
}
// A freshly fetched permanent-key row is already the canonical auth-key
// record. Prefer its non-zero Layer without making Router query the same PG
// row again. Unsupported metadata remains unknown and must not fall through
// to a weaker mirror.
if c.authKeyExpiresAt == 0 && fetchedLayer != 0 {
profile, supported := tg.ResolveLayerProfile(fetchedLayer)
if !supported {
return nil
}
return c.SeedInheritedLayerProfile(profile)
}
// Temporary keys resolve through their bound permanent key before consulting
// the raw temp-key shadow, which may predate a client upgrade.
if resolver, ok := s.layerRPC.(LayerRPCInheritedAuthKeyProfileResolver); ok {
layer, found, err := resolver.ResolveInheritedAuthKeyLayer(ctx, c.authKeyID)
if err != nil {
if !isLayerEvidenceDurabilityUnavailable(err) {
// A binding conflict, missing/destroyed key, or malformed durable
// value is not an availability hint. Stay unknown and never let a
// stale raw temp-key shadow outrank that structural failure.
if s.log != nil {
s.log.Warn("Resolve inherited auth-key Layer failed; awaiting explicit evidence",
zap.String("auth_key_id", c.authKeyHex), zap.Error(err))
}
return nil
}
// The same first frame already authenticated and loaded fetchedLayer
// from this raw temp key. During a transient permanent-identity lookup
// outage it is safe only as this physical Conn's inherited shadow; it
// is never published back to the shared permanent default.
if s.log != nil {
s.log.Warn("Resolve inherited auth-key Layer unavailable; continuing with raw auth-key shadow",
zap.String("auth_key_id", c.authKeyHex), zap.Error(err))
}
} else if !found {
// Fall through to a raw auth-key shadow when the resolver has no
// canonical permanent-key default (for example an unbound temp key).
} else {
profile, supported := tg.ResolveLayerProfile(layer)
if !supported {
return nil
}
return c.SeedInheritedLayerProfile(profile)
}
}
if fetchedLayer != 0 {
profile, supported := tg.ResolveLayerProfile(fetchedLayer)
if !supported {
return nil
}
return c.SeedInheritedLayerProfile(profile)
}
if hasDurableResolver {
// previous belongs to the old logical Conn which occupied this physical
// transport. In durable mode only an exact session row or auth-key default
// may cross that boundary. In particular, a connection-local selector used
// during a store outage must not leak into a newly selected session.
return nil
}
if previous.Origin != LayerProfileUnknown {
return c.SeedInheritedLayerProfile(previous.Profile)
}
return nil
}
// refreshActivatedInheritedLayerProfile closes the auth.bindTempAuthKey race
// after BeginActivation has made this Conn visible in claimsByAuth. If bind won
// first, the second resolver read sees the permanent-key default; if the claim
// won first, bind's SessionManager refresh sees and updates this Conn. Explicit
// session evidence always wins either ordering.
func (s *Server) refreshActivatedInheritedLayerProfile(ctx context.Context, c *Conn, fetchedLayer int) error {
if s == nil || c == nil || c.LayerProfileState().Origin == LayerProfileExplicit {
return nil
}
if c.authKeyExpiresAt == 0 {
if fetchedLayer == 0 {
return nil
}
profile, ok := tg.ResolveLayerProfile(fetchedLayer)
if !ok {
return c.clearInheritedLayerProfile()
}
_, err := c.refreshInheritedLayerProfile(profile)
return err
}
if resolver, ok := s.layerRPC.(LayerRPCInheritedAuthKeyProfileResolver); ok {
layer, found, err := resolver.ResolveInheritedAuthKeyLayer(ctx, c.authKeyID)
if err != nil {
if !isLayerEvidenceDurabilityUnavailable(err) {
if s.log != nil {
s.log.Warn("Re-resolve inherited auth-key Layer after activation claim failed",
zap.String("auth_key_id", c.authKeyHex), zap.Error(err))
}
// A pre-claim raw temp shadow may have been installed before a
// concurrent bind became visible. Structural identity/key failures
// must revoke that weaker inherited evidence; only an explicit
// selector (guarded above) is allowed to survive this branch.
return c.clearInheritedLayerProfile()
}
if s.log != nil {
s.log.Warn("Re-resolve inherited auth-key Layer unavailable; keeping raw auth-key shadow",
zap.String("auth_key_id", c.authKeyHex), zap.Error(err))
}
} else if found {
profile, supported := tg.ResolveLayerProfile(layer)
if !supported {
return c.clearInheritedLayerProfile()
}
_, err = c.refreshInheritedLayerProfile(profile)
return err
}
}
if fetchedLayer == 0 {
return nil
}
profile, ok := tg.ResolveLayerProfile(fetchedLayer)
if !ok {
return c.clearInheritedLayerProfile()
}
_, err := c.refreshInheritedLayerProfile(profile)
return err
}
// lockSessionLayerBinding linearizes a proactive update's final validation and
// physical write with profile correction. If correction wins first, validation
// observes the new epoch and drops the update. If the writer wins first, the
// correction becomes visible only after those bytes have landed. Returning a
// bool avoids allocating a release closure on the outbound hot path.
func (c *Conn) lockSessionLayerBinding(binding *outboundLayerBinding) bool {
if c == nil || binding == nil || binding.wireInvariant ||
binding.kind == outboundLayerBindingRequest || binding.epoch == 0 {
return false
}
c.layerProfileMu.RLock()
return true
}

7
internal/mtprotoedge/conn_mgmt_fixes_test.go
View file

@ -7,8 +7,8 @@ import (
"go.uber.org/zap/zaptest"
"github.com/gotd/td/proto"
"github.com/gotd/td/tg"
"github.com/iamxvbaba/td/proto"
"github.com/iamxvbaba/td/tg"
)
// TestPushSkipsConnReboundToOtherUser 锁定跨账号投递窗口的修复：pushToUserWithSender 在锁外
@ -28,6 +28,9 @@ func TestPushSkipsConnReboundToOtherUser(t *testing.T) {
c.userID.Store(userA)
c.userIDResolved.Store(true)
c.receivesUpdates.Store(true)
if err := c.FreezeLayerProfile(tg.LayerProfileCanonical); err != nil {
t.Fatal(err)
}
sm.Register(c)
return c
}

4
internal/mtprotoedge/decrypt.go
View file

@ -8,8 +8,8 @@ import (
"github.com/gotd/ige"
"github.com/gotd/td/bin"
"github.com/gotd/td/crypto"
"github.com/iamxvbaba/td/bin"
"github.com/iamxvbaba/td/crypto"
)
// clientFrame 是解密后的单帧客户端消息视图。data/plaintext 引用调用方持有的复用明文

4
internal/mtprotoedge/decrypt_test.go
View file

@ -9,8 +9,8 @@ import (
"github.com/gotd/ige"
"github.com/gotd/td/bin"
"github.com/gotd/td/crypto"
"github.com/iamxvbaba/td/bin"
"github.com/iamxvbaba/td/crypto"
)
func newTestAuthKey(t *testing.T) crypto.AuthKey {

66
internal/mtprotoedge/destroy_auth_key.go
View file

@ -1,9 +1,12 @@
package mtprotoedge
import (
"errors"
"fmt"
"github.com/gotd/td/bin"
"github.com/iamxvbaba/td/bin"
"github.com/iamxvbaba/td/proto"
"github.com/iamxvbaba/td/tg"
)
const (
@ -12,6 +15,39 @@ const (
destroyAuthKeyFailTypeID = 0xea109b13
)
var errDestroyAuthKeyMustBeExclusive = errors.New("wrapped destroy_auth_key must be the only logical message")
// wrappedDestroyAuthKeyTerminal accepts only evidence emitted by the generated
// exact wrapper parser after it has legally reached the innermost non-API
// terminal. It never re-parses wrapper bytes at runtime.
func wrappedDestroyAuthKeyTerminal(err error) (*tg.LayerRPCUnknownTerminalError, bool) {
var terminal *tg.LayerRPCUnknownTerminalError
if !errors.As(err, &terminal) || terminal == nil || terminal.WireID != destroyAuthKeyRequestTypeID {
return nil, false
}
return terminal, true
}
// validWrappedDestroyAuthKeyChain is deliberately narrower than "some generated
// wrapper decoded". invokeAfter*, takeout and update-suppression wrappers carry
// execution semantics which the service-message fast path must not silently
// discard. The official first-connection path is exactly
// invokeWithLayer(initConnection(destroy_auth_key)); an already initialized
// connection sends the bare service message and is classified before Layer RPC
// admission.
func validWrappedDestroyAuthKeyChain(terminal *tg.LayerRPCUnknownTerminalError) bool {
if terminal == nil || terminal.WrapperCount() != 2 {
return false
}
outer, outerOK := terminal.Wrapper(0)
inner, innerOK := terminal.Wrapper(1)
return outerOK && innerOK &&
outer.Profile() == terminal.Profile &&
inner.Profile() == terminal.Profile &&
outer.Semantic() == tg.LayerSemanticMethodInvokeWithLayer &&
inner.Semantic() == tg.LayerSemanticMethodInitConnection
}
type destroyAuthKeyRequest struct{}
func (*destroyAuthKeyRequest) Encode(b *bin.Buffer) error {
@ -26,16 +62,26 @@ func (*destroyAuthKeyRequest) Decode(b *bin.Buffer) error {
return nil
}
type destroyAuthKeyOk struct{}
func (*destroyAuthKeyOk) Encode(b *bin.Buffer) error {
b.PutID(destroyAuthKeyOkTypeID)
return nil
// destroyAuthKeyRPCResult is the only rpc_result envelope admitted as a
// layer-invariant control value. Its closed result set is part of the type, so
// it can never smuggle a profile-dependent API payload past the exact Layer
// binding boundary.
type destroyAuthKeyRPCResult struct {
RequestMessageID int64
ResultTypeID uint32
}
type destroyAuthKeyFail struct{}
func (*destroyAuthKeyFail) Encode(b *bin.Buffer) error {
b.PutID(destroyAuthKeyFailTypeID)
func (r *destroyAuthKeyRPCResult) Encode(b *bin.Buffer) error {
if r == nil {
return fmt.Errorf("encode destroy_auth_key rpc_result: nil result")
}
switch r.ResultTypeID {
case destroyAuthKeyOkTypeID, destroyAuthKeyFailTypeID:
default:
return fmt.Errorf("encode destroy_auth_key rpc_result: invalid inner constructor %#x", r.ResultTypeID)
}
b.PutID(proto.ResultTypeID)
b.PutLong(r.RequestMessageID)
b.PutID(r.ResultTypeID)
return nil
}

8
internal/mtprotoedge/duplicate_admission_test.go
View file

@ -7,9 +7,9 @@ import (
"testing"
"time"
"github.com/gotd/td/bin"
"github.com/gotd/td/proto"
"github.com/gotd/td/tg"
"github.com/iamxvbaba/td/bin"
"github.com/iamxvbaba/td/proto"
"github.com/iamxvbaba/td/tg"
)
type blockingDuplicateRPC struct {
@ -53,7 +53,7 @@ func TestPendingSameConnectionDuplicateDoesNotBlockFreshRequest(t *testing.T) {
defer handler.unblock()
addr, pub, server := startTestServer(t, Options{
DC: dc,
RPC: handler,
legacyRPC: handler,
RPCMaxInflight: 2,
RPCGlobalWorkers: 2,
RPCQueueSize: 8,

18
internal/mtprotoedge/e2e_test.go
View file

@ -11,13 +11,13 @@ import (
"go.uber.org/zap/zaptest"
"github.com/gotd/log/logzap"
"github.com/gotd/td/clock"
"github.com/gotd/td/exchange"
"github.com/gotd/td/session"
"github.com/gotd/td/telegram"
"github.com/gotd/td/telegram/dcs"
"github.com/gotd/td/tg"
"github.com/gotd/td/transport"
"github.com/iamxvbaba/td/clock"
"github.com/iamxvbaba/td/exchange"
"github.com/iamxvbaba/td/session"
"github.com/iamxvbaba/td/telegram"
"github.com/iamxvbaba/td/telegram/dcs"
"github.com/iamxvbaba/td/tg"
"github.com/iamxvbaba/td/transport"
"telesrv/internal/app/auth"
"telesrv/internal/app/updates"
@ -26,7 +26,7 @@ import (
"telesrv/internal/store/memory"
)
// TestTelegramClientEndToEnd 是连接层的最强端到端验证：用 gotd/td 的完整
// TestTelegramClientEndToEnd 是连接层的最强端到端验证：用 iamxvbaba/td 的完整
// telegram.Client（而非底层 cipher）连本地 mtprotoedge，client 自动经
// invokeWithLayer(initConnection(help.getConfig)) 完成初始化，并取得含本地 DC 的 Config。
func TestTelegramClientEndToEnd(t *testing.T) {
@ -51,7 +51,7 @@ func TestTelegramClientEndToEnd(t *testing.T) {
Updates: updates.NewService(memory.NewUpdateStateStore(), memory.NewUpdateEventStore()),
}
router := rpc.New(rpc.Config{DC: dc, IP: tcpAddr.IP.String(), Port: tcpAddr.Port}, deps, zaptest.NewLogger(t), clock.System)
srv := New(Options{Logger: zaptest.NewLogger(t), DC: dc, RSAKey: rsaKey, AuthKeys: authKeyStore, RPC: router})
srv := New(Options{Logger: zaptest.NewLogger(t), DC: dc, RSAKey: rsaKey, AuthKeys: authKeyStore, LayerRPC: router})
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
defer cancel()

492
internal/mtprotoedge/encrypted.go
View file

@ -10,19 +10,20 @@ import (
"fmt"
"io"
"math"
"runtime/debug"
"time"
"go.uber.org/zap"
"github.com/gotd/td/bin"
"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"
"github.com/iamxvbaba/td/bin"
"github.com/iamxvbaba/td/crypto"
"github.com/iamxvbaba/td/mt"
"github.com/iamxvbaba/td/proto"
"github.com/iamxvbaba/td/proto/codec"
"github.com/iamxvbaba/td/tg"
"github.com/iamxvbaba/td/tgerr"
"github.com/iamxvbaba/td/transport"
"telesrv/internal/compat/layerwire"
"telesrv/internal/observability/dbtrace"
"telesrv/internal/postresponse"
"telesrv/internal/store"
@ -140,7 +141,11 @@ func (s *Server) handleEncrypted(ctx context.Context, tc transport.Conn, cs *con
// 临时创建 Conn 会在同一 socket 上启动多个 outbound actor，Android 的启动重试
// 风暴随即变成并发写和重复结果放大。
if current == nil || current.sessionID != frame.sessionID || current.authKeyID != key.ID {
var previousLayer LayerProfileSnapshot
if current != nil {
if current.authKeyID == key.ID {
previousLayer = current.LayerProfileState()
}
cs.reset()
current.beginTerminalShutdown()
s.conns.Unregister(current)
@ -156,15 +161,20 @@ func (s *Server) handleEncrypted(ctx context.Context, tc transport.Conn, cs *con
current = s.newConn(tc, key, frame.sessionID, serverSalt)
}
current.authKeyExpiresAt = authKeyExpiresAt
// 注册即播种协商 layer：新 Conn 的 clientLayer 为 0（=canonical 227），若等到
// 首条 RPC 的 Dispatch 返回后才刷新，重连老客户端在首条 RPC handler 执行期间
// 收到的 pending flush / 并发 push 会漏降级。进程内重连时 rpc 层留有
// (auth_key, session) / auth_key 两级协商记录，这里一次查询即可闭合该空窗。
// Same-session evidence is restored as explicit; auth-key metadata is only
// an inherited default and can be corrected by the next invokeWithLayer.
if s.rpc != nil {
if layer, ok := s.rpc.NegotiatedLayer(current.authKeyID, current.sessionID); ok {
current.SetClientLayer(layer)
current.setLegacyClientLayer(layer)
}
}
fetchedLayer := 0
if fetchedKey != nil {
fetchedLayer = fetchedKey.Layer
}
if err := s.seedInitialLayerProfile(ctx, current, fetchedLayer, previousLayer); err != nil {
return current, fmt.Errorf("seed connection layer profile: %w", err)
}
}
if frame.salt != serverSalt {
@ -194,6 +204,13 @@ func (s *Server) handleEncrypted(ctx context.Context, tc transport.Conn, cs *con
}
defer plan.close()
if err := s.prepareInboundRPCBatch(ctx, current, plan); err != nil {
if errors.Is(err, errDestroyAuthKeyMustBeExclusive) {
s.log.Debug("Rejecting mixed destroy_auth_key container",
zap.Int64("msg_id", frame.messageID),
zap.Int32("seq_no", frame.seqNo),
)
return current, s.sendBadMsg(ctx, current, frame.messageID, frame.seqNo, badMsgContainer)
}
return current, err
}
if err := sendQuickAckIfRequested(ctx, current.transport, key, frame.plaintext, s.writeTimeout); err != nil {
@ -233,6 +250,12 @@ func (s *Server) handleEncrypted(ctx context.Context, tc transport.Conn, cs *con
}
return current, errActivationAuthKeyRejected
}
// Re-resolve inherited Layer only after the activation claim is visible.
// This closes the bind-vs-connect window for temporary keys without ever
// replacing explicit invokeWithLayer evidence admitted above.
if err := s.refreshActivatedInheritedLayerProfile(ctx, current, fresh.Layer); err != nil {
return current, fmt.Errorf("refresh claimed connection layer profile: %w", err)
}
if current.isRetired() || !current.isPhysicalTransportCurrentOpen() {
return current, ErrConnClosed
}
@ -674,29 +697,44 @@ func (s *Server) handleRPC(ctx context.Context, c *Conn, msgID int64, method str
ctx = postresponse.WithCallbacks(ctx)
ctx, dbStats := dbtrace.WithStats(ctx)
// legacyRPC is an unexported package-test hook, but its result still has to
// obey the production exact-codec invariant. Admit a defensive copy using
// the generated current profile before the legacy router consumes b.
admissionBody := &bin.Buffer{Buf: append([]byte(nil), b.Buf...)}
admitted, err := tg.NewServerDispatcher(nil).AdmitDefaultLayerWithLimits(
tg.LayerProfileCanonical,
admissionBody,
inboundLayerDecodeLimits,
)
if err != nil {
return fmt.Errorf("admit legacy test RPC through generated codec: %w", err)
}
start := s.clock.Now()
effectiveMethod := method
var (
result bin.Encoder
err error
result bin.Encoder
dispatchErr error
)
if detailed, ok := s.rpc.(RPCHandlerWithMethod); ok {
if detailed, ok := s.rpc.(legacyRPCHandlerWithMethod); ok {
var innerMethod string
result, innerMethod, err = detailed.DispatchWithMethod(ctx, c.authKeyID, c.sessionID, b)
result, innerMethod, dispatchErr = detailed.DispatchWithMethod(ctx, c.authKeyID, c.sessionID, b)
if innerMethod != "" {
effectiveMethod = innerMethod
}
} else {
result, err = s.rpc.Dispatch(ctx, c.authKeyID, c.sessionID, b)
result, dispatchErr = s.rpc.Dispatch(ctx, c.authKeyID, c.sessionID, b)
}
if dispatchErr == nil && result != nil && !isLayerInvariantRPCResultEncoder(result) {
if _, exact := result.(exactLayerRPCResultEncoder); !exact {
result = &legacyTestRPCResultEncoder{call: admitted.Call(), result: result}
}
}
dur := s.clock.Now().Sub(start)
s.metrics.RPCHandled(effectiveMethod, dur, err)
// 刷新本连接协商 layer（invokeWithLayer/initConnection 已被 Dispatch 处理并登记），
// 供 rpc_result 与后续 push 出站降级使用。仅在确实观测到 layer 时更新——缓存被驱逐
// 时 NegotiatedLayer 返回 ok=false，此时必须保留连接已记住的 layer，绝不覆盖成默认值，
// 否则长连接老客户端的条目被驱逐后会被误降回 227。
s.metrics.RPCHandled(effectiveMethod, dur, dispatchErr)
// 刷新本连接由 invokeWithLayer 证明并冻结的 exact-session layer。ok=false
// 表示仍无协议证据；设备/授权元数据和其它 session 都不具备回填资格。
if layer, ok := s.rpc.NegotiatedLayer(c.authKeyID, c.sessionID); ok {
c.SetClientLayer(layer)
c.setLegacyClientLayer(layer)
}
fields := make([]zap.Field, 0, 12)
@ -725,7 +763,7 @@ func (s *Server) handleRPC(ctx context.Context, c *Conn, msgID int64, method str
// Plain connection cancellation remains retryable on the replacement.
var terminal bin.Encoder
runPostResponse := false
if err == nil && result != nil {
if dispatchErr == nil && result != nil {
terminal = result
runPostResponse = true
} else if errors.Is(ctxErr, context.DeadlineExceeded) {
@ -741,23 +779,23 @@ func (s *Server) handleRPC(ctx context.Context, c *Conn, msgID int64, method str
}
}
cancelFields := append(fields, zap.NamedError("context_error", ctxErr))
if err != nil {
cancelFields = append(cancelFields, zap.NamedError("dispatch_error", err))
if dispatchErr != nil {
cancelFields = append(cancelFields, zap.NamedError("dispatch_error", dispatchErr))
}
s.log.Info("RPC canceled", cancelFields...)
return ctxErr
}
if err != nil {
if dispatchErr != nil {
var rpcErr *tgerr.Error
if errors.As(err, &rpcErr) {
if errors.As(dispatchErr, &rpcErr) {
s.log.Info("RPC error", append(fields, zap.Int("code", rpcErr.Code), zap.String("error", rpcErr.Message))...)
return s.publishRPCResult(c, msgID, effectiveMethod, owner, &mt.RPCError{
ErrorCode: rpcErr.Code,
ErrorMessage: rpcErr.Message,
}, nil)
}
s.log.Info("RPC internal error", append(fields, zap.Error(err))...)
s.log.Info("RPC internal error", append(fields, zap.Error(dispatchErr))...)
return s.publishRPCResult(c, msgID, effectiveMethod, owner, &mt.RPCError{
ErrorCode: 500,
ErrorMessage: "INTERNAL",
@ -768,9 +806,14 @@ func (s *Server) handleRPC(ctx context.Context, c *Conn, msgID int64, method str
return s.publishRPCResult(c, msgID, effectiveMethod, owner, result, postresponse.Take(ctx))
}
var errRPCResultRetentionHandoff = errors.New("mtproto rpc result retention handoff failed")
type rpcResultRetentionHandoff func(*encodedOutboundMessage, error) error
// publishRPCResult ends the inbound worker's ownership at bounded egress
// admission. Physical delivery, fencing, completed-cache publication and the
// post-response hook are thereafter owned by the single outbound actor.
// admission. Physical delivery is thereafter owned either by the single
// outbound actor or, under retained-byte saturation, by a fenced completed-cache
// entry that the replacement connection can replay without rerunning business.
func (s *Server) publishRPCResult(
c *Conn,
reqMsgID int64,
@ -788,45 +831,117 @@ func (s *Server) publishRPCResult(
}
prepareCtx, cancel := context.WithTimeout(context.Background(), prepareTimeout)
defer cancel()
encoded, err := s.encodeRPCResultContext(prepareCtx, c, reqMsgID, result)
prepareEncoded := func(encoded *encodedOutboundMessage) (outboundPriority, bool) {
if owner != nil && owner.Delivery() != nil {
// The owner-level delivery coordinator exists before the handler starts, so
// an initConnection rewrap can retarget even while result encoding is still
// pending. The encoded body itself remains immutable; the actor clones only
// the 12-byte rpc_result prefix when it snapshots the physical target.
encoded.delivery = owner.Delivery()
}
if afterDelivered != nil {
encoded.setDeliveryHook(afterDelivered)
}
priority := rpcResultPriority(method, encoded)
encoded.priority = priority
if metrics, ok := s.metrics.(RPCResultMetrics); ok {
metrics.RPCResultPrepared(method, priority.String(), encoded.uncompressedBytes, len(encoded.body), encoded.compressed)
}
visible := encoded.compressed || priority == outboundPriorityCritical || priority == outboundPriorityBulk
return priority, visible
}
// A successful business result may never leave the encode slot as an
// unaccounted []byte. If the primary 512MiB retained-body budget is full, make
// overload terminal for this physical generation and publish the exact result
// into the independently bounded completed cache before releasing the slot.
retainForReplay := func(encoded *encodedOutboundMessage, admissionErr error) error {
if s == nil || s.rpcResults == nil || c == nil || encoded == nil || reqMsgID == 0 {
return errors.New("rpc result completed cache is unavailable")
}
if int64(len(encoded.body)) > s.rpcResults.completedBytes.max {
// Every transport-legal result fits the production completed cache by the
// compile-time invariant in rpc_result_cache.go. A test/custom cache that
// violates it cannot safely complete this flight, so fail fast while the
// body is still confined to the encode slot.
panic(fmt.Sprintf(
"mtprotoedge: encoded rpc result exceeds completed-cache budget: body=%d max=%d",
len(encoded.body), s.rpcResults.completedBytes.max,
))
}
priority, visible := prepareEncoded(encoded)
if owner != nil && !owner.HandOff() {
return ErrRPCResultFlightInvalid
}
started := time.Now()
encoded.markReplayable()
// Put may expose a completed result only after the old logical connection
// is irreversibly unable to accept another same-generation request.
c.fenceUndeliveredRPCResult()
s.storeRPCResult(c, reqMsgID, encoded)
latency := time.Since(started)
if metrics, ok := s.metrics.(RPCResultMetrics); ok {
metrics.RPCResultDelivered(method, latency, len(encoded.body), admissionErr)
}
resultLogLevel := zap.DebugLevel
if visible {
resultLogLevel = zap.InfoLevel
}
if checked := s.log.Check(resultLogLevel, "RPC result retained for replay after egress saturation"); checked != nil {
checked.Write(
zap.String("method", method), zap.Int64("req_msg_id", reqMsgID),
zap.Int64("delivered_req_msg_id", encoded.writtenRequestID()),
zap.String("auth_key_id", c.authKeyHex), zap.Int64("session_id", c.sessionID),
zap.Int("wire_bytes", len(encoded.body)), zap.Bool("gzip", encoded.compressed),
zap.String("priority", priority.String()), zap.Error(admissionErr))
}
return nil
}
encoded, reserved, retained, err := s.encodeRPCResultReservedWithHandoffContext(
prepareCtx, c, reqMsgID, result, retainForReplay,
)
if retained {
return err
}
if errors.Is(err, errRPCResultRetentionHandoff) {
c.fenceUndeliveredRPCResult()
return err
}
if err != nil {
s.log.Warn("Encode RPC result failed; publishing INTERNAL",
zap.String("method", method), zap.Int64("req_msg_id", reqMsgID), zap.Error(err))
afterDelivered = nil
encoded, err = s.encodeRPCResultContext(prepareCtx, c, reqMsgID, &mt.RPCError{
ErrorCode: 500, ErrorMessage: "INTERNAL",
})
encoded, reserved, retained, err = s.encodeRPCResultReservedWithHandoffContext(
prepareCtx, c, reqMsgID, &mt.RPCError{ErrorCode: 500, ErrorMessage: "INTERNAL"}, retainForReplay,
)
if retained {
return err
}
if err != nil {
c.fenceUndeliveredRPCResult()
return err
}
}
if owner != nil && owner.Delivery() != nil {
// The owner-level delivery coordinator exists before the handler starts, so
// an initConnection rewrap can retarget even while result encoding is still
// pending. The encoded body itself remains immutable; the actor clones only
// the 12-byte rpc_result prefix when it snapshots the physical target.
encoded.delivery = owner.Delivery()
}
if afterDelivered != nil {
encoded.delivery.fn = afterDelivered
if encoded == nil || reserved == nil {
c.fenceUndeliveredRPCResult()
return errors.New("rpc result encode completed without tracked retention")
}
// Until enqueue transfers ownership, every exit must return the retained-byte
// charge. A successful transfer clears the reservation and makes this a no-op.
defer reserved.release()
priority, visible := prepareEncoded(encoded)
if owner != nil && !owner.HandOff() {
return ErrRPCResultFlightInvalid
}
priority := rpcResultPriority(method, encoded)
encoded.priority = priority
resultLogLevel := zap.DebugLevel
if encoded.compressed || priority == outboundPriorityCritical || priority == outboundPriorityBulk {
if visible {
// Keep ordinary small RPCs at debug, but make convergence and bulk/gzip
// delivery visible in default service logs. These are the
// delivery visible in the default service logs. These are the
// responses whose queueing and write latency diagnose startup Updating.
resultLogLevel = zap.InfoLevel
}
if metrics, ok := s.metrics.(RPCResultMetrics); ok {
metrics.RPCResultPrepared(method, priority.String(), encoded.uncompressedBytes, len(encoded.body), encoded.compressed)
}
egressStarted := time.Now()
terminal := func(deliveryErr error) {
latency := time.Since(egressStarted)
@ -848,8 +963,8 @@ func (s *Server) publishRPCResult(
}
return
}
s.storeRPCResult(c, reqMsgID, encoded)
encoded.markDelivered()
s.storeRPCResult(c, reqMsgID, encoded)
if checked := s.log.Check(resultLogLevel, "RPC result delivered"); checked != nil {
checked.Write(
zap.String("method", method), zap.Int64("req_msg_id", reqMsgID),
@ -860,7 +975,7 @@ func (s *Server) publishRPCResult(
}
}
encoded.markQueued()
if err := c.enqueueEncodedDelivery(prepareCtx, proto.MessageServerResponse, encoded, priority, terminal); err != nil {
if err := c.enqueueEncodedDeliveryReserved(prepareCtx, proto.MessageServerResponse, encoded, priority, terminal, reserved); err != nil {
// HandOff already made the egress path the terminal owner. No bytes were
// admitted, so fence this generation before publishing a replayable result.
terminal(err)
@ -905,10 +1020,10 @@ func (s *Server) sendResult(ctx context.Context, c *Conn, reqMsgID int64, result
s.storeRPCResult(c, reqMsgID, encoded)
return err
}
encoded.markDelivered()
// 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)
encoded.markDelivered()
return nil
}
@ -916,24 +1031,114 @@ func (s *Server) sendResult(ctx context.Context, c *Conn, reqMsgID int64, result
// 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 {
return s.sendCachedRPCResultWithHook(ctx, c, encoded, nil)
}
func (s *Server) sendCachedRPCResultWithHook(
ctx context.Context,
c *Conn,
encoded *encodedOutboundMessage,
afterSuccessfulDelivery func() error,
) error {
if encoded == nil {
c.fenceUndeliveredRPCResult()
return errors.New("nil cached rpc_result")
}
if err := c.SendEncoded(ctx, proto.MessageServerResponse, encoded); err != nil {
attempt, reserved, err := c.cloneRPCResultForRequestReserved(encoded, encoded.reqMsgID, false)
if err != nil {
c.failOutboundBudget(err)
c.fenceUndeliveredRPCResult()
encoded.markReplayable()
return err
}
encoded.markDelivered()
return nil
// take clears the producer reservation after actor admission. Every earlier
// return, including a closed connection, must drop the replay pin here.
defer reserved.release()
pendingLogicalRestore := attempt.pendingLogicalDeliveryHook()
var finishRestore func()
if afterSuccessfulDelivery != nil || pendingLogicalRestore {
finishRestore = c.beginRPCReplayRestore()
defer finishRestore()
}
// Cached replay owns its delivery-gated state synchronously. Calling the
// lower send primitive avoids reserving the process-wide asynchronous hook
// executor; the logical hook is claimed only after this physical write wins.
if err := c.sendOutboundWithTerminalReserved(
ctx, proto.MessageServerResponse, nil, attempt, false, nil, reserved,
); err != nil {
c.fenceUndeliveredRPCResult()
attempt.markReplayable()
return err
}
// Physical success is irrevocable even if the caller's send context expires
// at the same instant. Give the ordered restore its own bounded lifetime.
restoreCtx, cancelRestore := boundedRPCReplayRestoreContext(context.Background())
defer cancelRestore()
logicalRestore, claimErr := attempt.claimLogicalDeliveryHook(restoreCtx, false)
attempt.markDelivered()
if claimErr != nil {
// Another replay owns Claimed/InProgress state (or a retarget still owns
// the sticky deferral). Fence before the deferred barrier is released; a
// later physical generation may wait for Done and replay the same bytes.
c.fenceUndeliveredRPCResult()
return fmt.Errorf("wait for cached rpc_result logical restore: %w", claimErr)
}
return s.runBoundedRPCReplayRestore(
restoreCtx, c, "cached rpc_result", logicalRestore, afterSuccessfulDelivery,
)
}
// composeRPCReplayRestore keeps replacement-connection metadata first while
// still guaranteeing that the original handler's delivery-gated cursor/outbox
// work runs after a physical replay even when metadata restoration reports an
// error. runRPCReplayRestore provides panic isolation and terminal fencing for
// the combined ordered transaction.
func composeRPCReplayRestore(logical func(), replacement func() error) func() error {
if logical == nil && replacement == nil {
return nil
}
return func() (err error) {
if logical != nil {
defer logical()
}
if replacement != nil {
return replacement()
}
return nil
}
}
// runRPCReplayRestore is the panic/error boundary executed only by the fixed-
// capacity runner in rpc_replay_restore.go. Replay restore may touch auth/session
// stores and membership state; its caller holds the per-Conn scheduler barrier.
// Any error or panic fences the partially restored physical generation so a
// replacement can retry from the immutable completed result.
func (s *Server) runRPCReplayRestore(c *Conn, source string, restore func() error) (err error) {
if restore == nil {
return nil
}
defer func() {
if recovered := recover(); recovered != nil {
err = fmt.Errorf("restore replay state after %s: panic: %v", source, recovered)
if s != nil && s.log != nil {
s.log.Error("Exact RPC replay state restore panicked",
zap.String("source", source), zap.ByteString("stack", debug.Stack()), zap.Any("panic", recovered))
}
}
if err != nil {
if s != nil && s.log != nil {
s.log.Warn("Exact RPC replay state restore failed", zap.String("source", source), zap.Error(err))
}
c.fenceUndeliveredRPCResult()
}
}()
return restore()
}
// encodeRPCResult 编码 rpc_result。内层对象与 rpc_result 头（type_id + req_msg_id）
// 一次性编码进同一 buffer——旧实现先编码内层、再经 proto.Result.Encode 整体拷贝一遍，
// 每条响应多一份全量 body 拷贝。内层按连接协商 layer 降级（layer==227 直通，零开销），
// 降级改写字节时才重建整条消息。降级失败 fail-safe：记日志并发送 canonical 字节——
// 宁可老客户端对个别长尾对象渲染异常，也不让连接/流崩。
// 每条响应多一份全量 body 拷贝。生成式结果携带完整 Layer profile + result TypeRef
// 绑定，后续发送、缓存和重放均只能用于同一精确 profile；package 测试保留的 legacy
// handler 也必须先绑定 generated admitted call，生产路径不存在旧转码桥。
func (s *Server) encodeRPCResult(c *Conn, reqMsgID int64, result bin.Encoder) (*encodedOutboundMessage, error) {
return s.encodeRPCResultContext(context.Background(), c, reqMsgID, result)
}
@ -942,36 +1147,138 @@ func (s *Server) encodeRPCResultContext(ctx context.Context, c *Conn, reqMsgID i
if ctx == nil {
ctx = context.Background()
}
var inner bin.Buffer
// Terminal result preparation must survive physical-generation retirement:
// an overlapping replacement may already be waiting to replay this owner's
// result. Only the bounded preparation context, not the old socket stop, owns it.
if err := withOutboundEncodeSlot(ctx, nil, func() error {
return result.Encode(&inner)
}); err != nil {
var encoded *encodedOutboundMessage
err := withOutboundEncodeSlot(ctx, nil, func() error {
var err error
encoded, err = s.encodeRPCResultWithoutSlot(ctx, c, reqMsgID, result)
return err
})
return encoded, err
}
// encodeRPCResultReservedContext keeps the process-wide encode slot until the
// completed immutable body is charged to the shared retained-byte budget. This
// closes the otherwise unbounded interval in which every RPC worker could own
// a large encoded result that neither the inbound nor outbound budget tracked.
func (s *Server) encodeRPCResultReservedContext(
ctx context.Context,
c *Conn,
reqMsgID int64,
result bin.Encoder,
) (*encodedOutboundMessage, *outboundBodyReservation, error) {
encoded, reserved, _, err := s.encodeRPCResultReservedWithHandoffContext(ctx, c, reqMsgID, result, nil)
return encoded, reserved, err
}
// encodeRPCResultReservedWithHandoffContext has only two successful ownership
// outcomes for a completed body: a primary outbound reservation, or a caller
// handoff that synchronously installs another bounded owner while the encode slot
// is still held. A failed/no handoff clears encoded before the slot is released.
func (s *Server) encodeRPCResultReservedWithHandoffContext(
ctx context.Context,
c *Conn,
reqMsgID int64,
result bin.Encoder,
handoff rpcResultRetentionHandoff,
) (*encodedOutboundMessage, *outboundBodyReservation, bool, error) {
if ctx == nil {
ctx = context.Background()
}
var (
encoded *encodedOutboundMessage
reserved *outboundBodyReservation
retained bool
)
err := withOutboundEncodeSlot(ctx, nil, func() error {
var err error
encoded, err = s.encodeRPCResultWithoutSlot(ctx, c, reqMsgID, result)
if err != nil {
return err
}
budget := c.outboundMessageBudget(encoded.typeID, false)
bytes := len(encoded.body)
if budget.reserve(bytes) {
reserved = &outboundBodyReservation{budget: budget, bytes: bytes}
return nil
}
if handoff != nil {
admissionErr := fmt.Errorf("reserve encoded rpc result: %w", ErrOutboundTrackedBudget)
if err := handoff(encoded, admissionErr); err != nil {
encoded = nil
return fmt.Errorf("%w: %w", errRPCResultRetentionHandoff, errors.Join(admissionErr, err))
}
retained = true
// The handoff owns the only surviving pointer. Do not return a second
// producer reference after the encode slot releases; the completed cache
// may independently evict the entry under its bounded policy.
encoded = nil
return admissionErr
}
// Non-publish callers have no alternate bounded owner. They may wait for
// the caller's deadline, but on failure the body is discarded in-slot.
if err := budget.waitReserve(ctx, nil, bytes); err != nil {
encoded = nil
return fmt.Errorf("reserve encoded rpc result: %w", err)
}
reserved = &outboundBodyReservation{budget: budget, bytes: bytes}
return nil
})
if err != nil && !retained && reserved == nil {
encoded = nil
}
return encoded, reserved, retained, err
}
func (s *Server) encodeRPCResultWithoutSlot(ctx context.Context, c *Conn, reqMsgID int64, result bin.Encoder) (*encodedOutboundMessage, error) {
var layerBinding *outboundLayerBinding
exactResult, exactLayerResult := result.(exactLayerRPCResultEncoder)
layerInvariantResult := isLayerInvariantRPCResultEncoder(result)
if !exactLayerResult && !layerInvariantResult {
return nil, ErrOutboundLayerBindingRequired
}
if exactLayerResult {
binding := exactResult.exactLayerRPCResultBinding()
layerBinding = &binding
if err := validateOutboundLayerBinding(c, &encodedOutboundMessage{layer: layerBinding}); err != nil {
return nil, fmt.Errorf("bind exact layer rpc result: %w", err)
}
}
// Encode the ordinary exact/no-gzip path directly behind the rpc_result
// prefix. This avoids both the old generated Prepare snapshot and another
// full-body copy merely to prepend the 12-byte envelope.
var envelope bin.Buffer
envelope.PutID(proto.ResultTypeID)
envelope.PutLong(reqMsgID)
if err := result.Encode(&envelope); err != nil {
return nil, fmt.Errorf("encode rpc result: %w", err)
}
innerBody := inner.Raw()
if layer := c.ClientLayer(); layer < layerwire.CanonicalLayer {
if down, err := layerwire.Transcode(innerBody, layer); err != nil {
s.log.Warn("layerwire downgrade failed; sending canonical rpc_result",
zap.Int("layer", layer), zap.Int64("req_msg_id", reqMsgID), zap.Error(err))
} else {
innerBody = down
}
envelopeInner := envelope.Raw()[12:]
innerBody := envelopeInner
// Inputs above gotd's decompression ceiling can never be gzip_packed. Reject
// them before allocating a second transport-envelope-sized buffer.
if len(innerBody) > rpcResultGZIPMaxInputBytes && len(innerBody) > maxOutboundBodyBytes-12 {
return nil, fmt.Errorf("%w: body=%d limit=%d", ErrOutboundMessageTooLarge, len(innerBody)+12, maxOutboundBodyBytes)
}
wireInner, compressed, err := encodeAdaptiveRPCResultInner(ctx, nil, innerBody)
if err != nil {
return nil, fmt.Errorf("compress rpc result: %w", err)
}
var out bin.Buffer
out.PutID(proto.ResultTypeID)
out.PutLong(reqMsgID)
out.Put(wireInner)
if len(wireInner) > maxOutboundBodyBytes-12 {
return nil, fmt.Errorf("%w: body=%d limit=%d", ErrOutboundMessageTooLarge, len(wireInner)+12, maxOutboundBodyBytes)
}
body := envelope.Raw()
if compressed || !sameBacking(wireInner, envelopeInner) {
var out bin.Buffer
out.PutID(proto.ResultTypeID)
out.PutLong(reqMsgID)
out.Put(wireInner)
body = out.Raw()
}
return &encodedOutboundMessage{
typeID: proto.ResultTypeID, body: out.Raw(), reqMsgID: reqMsgID,
compressed: compressed, uncompressedBytes: len(innerBody), delivery: newRPCResultDelivery(),
typeID: proto.ResultTypeID, body: body, reqMsgID: reqMsgID,
compressed: compressed, uncompressedBytes: len(innerBody), delivery: newRPCResultDelivery(0),
layer: layerBinding, layerInvariant: layerInvariantResult,
}, nil
}
@ -1075,11 +1382,20 @@ func (s *Server) sendMsgsStateInfo(ctx context.Context, c *Conn, reqMsgID int64,
}
func (s *Server) sendDestroySession(ctx context.Context, c *Conn, sessionID int64) error {
removed := false
removed, removedDurable := false, false
if sessionID != c.sessionID {
removed = s.conns.DestroySessionForAuthKey(c.authKeyID, sessionID)
if deleter, ok := s.layerRPC.(LayerRPCDurableSessionProfileDeleter); ok {
var err error
removedDurable, err = deleter.DeleteNegotiatedSessionLayerEvidence(ctx, c.authKeyID, sessionID)
if err != nil {
return fmt.Errorf("delete durable exact session Layer evidence: %w", err)
}
}
if s.conns != nil {
removed = s.conns.DestroySessionForAuthKey(c.authKeyID, sessionID)
}
}
if removed {
if removed || removedDurable {
return c.Send(ctx, proto.MessageServerResponse, &mt.DestroySessionOk{SessionID: sessionID})
}
return c.Send(ctx, proto.MessageServerResponse, &mt.DestroySessionNone{SessionID: sessionID})
@ -1113,7 +1429,7 @@ func (s *Server) typeName(id uint32) string {
}
func validateClientEnvelope(now time.Time, msgID int64, seqNo int32, typeID uint32) int {
if msgID == 0 || proto.MessageID(msgID).Type() != proto.MessageFromClient {
if !validClientMessageIDBits(msgID) {
return badMsgIDInvalidBits
}
msgTime := proto.MessageID(msgID).Time()
@ -1137,7 +1453,7 @@ func validateClientEnvelope(now time.Time, msgID int64, seqNo int32, typeID uint
}
func validateClientContainerEnvelope(msgID int64, seqNo int32, typeID uint32) int {
if msgID == 0 || proto.MessageID(msgID).Type() != proto.MessageFromClient {
if !validClientMessageIDBits(msgID) {
return badMsgIDInvalidBits
}
if clientMessageAllowsEitherSeqParity(typeID) {

350
internal/mtprotoedge/encrypted_test.go
View file

@ -2,17 +2,73 @@ package mtprotoedge
import (
"context"
"errors"
"sync"
"testing"
"time"
"github.com/gotd/td/bin"
"github.com/gotd/td/crypto"
"github.com/gotd/td/mt"
"github.com/gotd/td/proto"
"github.com/gotd/td/tg"
"github.com/gotd/td/transport"
"github.com/iamxvbaba/td/bin"
"github.com/iamxvbaba/td/crypto"
"github.com/iamxvbaba/td/mt"
"github.com/iamxvbaba/td/proto"
"github.com/iamxvbaba/td/tg"
"github.com/iamxvbaba/td/transport"
"telesrv/internal/store/memory"
)
type durableDestroyLayerRPC struct {
*admissionOnlyLayerRPC
mu sync.Mutex
deleted bool
err error
authKeyID [8]byte
sessionID int64
}
type deleteFailAuthKeyStore struct {
*memory.AuthKeyStore
err error
}
type trailingDestroyAuthKeyRequest struct{}
func (*trailingDestroyAuthKeyRequest) Encode(b *bin.Buffer) error {
b.PutID(destroyAuthKeyRequestTypeID)
b.PutID(0xdeadbeef)
return nil
}
func (*trailingDestroyAuthKeyRequest) Decode(b *bin.Buffer) error {
if err := b.ConsumeID(destroyAuthKeyRequestTypeID); err != nil {
return err
}
_, err := b.ID()
return err
}
func (s *deleteFailAuthKeyStore) Delete(context.Context, [8]byte) error {
return s.err
}
func (h *durableDestroyLayerRPC) DeleteNegotiatedSessionLayerEvidence(
_ context.Context,
authKeyID [8]byte,
sessionID int64,
) (bool, error) {
h.mu.Lock()
defer h.mu.Unlock()
h.authKeyID = authKeyID
h.sessionID = sessionID
return h.deleted, h.err
}
func (h *durableDestroyLayerRPC) deletion() ([8]byte, int64) {
h.mu.Lock()
defer h.mu.Unlock()
return h.authKeyID, h.sessionID
}
// TestEncryptedPingPong 验证 M2/M4：握手后 client 加密 ping，
// server 回 new_session_created + pong + msgs_ack。
func TestEncryptedPingPong(t *testing.T) {
@ -47,7 +103,7 @@ func TestEncryptedPingPong(t *testing.T) {
func TestDuplicateMsgIDIdempotent(t *testing.T) {
const dc = 2
handler := &admissionCountingRPC{}
addr, pub, _ := startTestServer(t, Options{DC: dc, RPC: handler})
addr, pub, _ := startTestServer(t, Options{DC: dc, legacyRPC: handler})
conn, auth, cipher := dialHandshake(t, addr, dc, pub)
clientMsgID := proto.NewMessageIDGen(time.Now)
@ -231,8 +287,65 @@ func TestDestroySession(t *testing.T) {
}
}
func TestDestroySessionAcknowledgesOfflineDurableEvidenceDeletion(t *testing.T) {
const dc = 2
handler := &durableDestroyLayerRPC{
admissionOnlyLayerRPC: newAdmissionOnlyLayerRPC(),
deleted: true,
}
addr, pub, _ := startTestServer(t, Options{DC: dc, LayerRPC: handler})
conn, auth, cipher := dialHandshake(t, addr, dc, pub)
clientMsgID := proto.NewMessageIDGen(time.Now)
reqMsgID := clientMsgID.New(proto.MessageFromClient)
targetSessionID := auth.SessionID + 4
sendEncrypted(t, conn, cipher, auth, reqMsgID, &mt.DestroySessionRequest{SessionID: targetSessionID})
replies := collectReplies(t, conn, cipher, auth.AuthKey, mt.DestroySessionOkTypeID)
buf := mustHave(t, replies, mt.DestroySessionOkTypeID, "destroy_session_ok")
var res mt.DestroySessionOk
if err := res.Decode(buf); err != nil {
t.Fatal(err)
}
if res.SessionID != targetSessionID {
t.Fatalf("destroy_session_ok.session_id = %d, want %d", res.SessionID, targetSessionID)
}
authKeyID, deletedSessionID := handler.deletion()
if authKeyID == ([8]byte{}) || deletedSessionID != targetSessionID {
t.Fatalf("durable deletion = auth:%x session:%d", authKeyID, deletedSessionID)
}
}
func TestDestroySessionDurabilityFailureDoesNotAcknowledgeOrRetireLiveSession(t *testing.T) {
boom := errors.New("database unavailable")
handler := &durableDestroyLayerRPC{
admissionOnlyLayerRPC: newAdmissionOnlyLayerRPC(),
err: boom,
}
manager := NewSessionManager(nil)
authKeyID := [8]byte{0xd3, 0x57}
target := &Conn{authKeyID: authKeyID, sessionID: 2, metrics: NopMetrics{}}
if err := manager.Register(target); err != nil {
t.Fatal(err)
}
defer manager.Unregister(target)
s := New(Options{DC: 2, LayerRPC: handler, ActiveSessions: manager})
current := &Conn{authKeyID: authKeyID, sessionID: 1, metrics: NopMetrics{}}
err := s.sendDestroySession(context.Background(), current, target.sessionID)
if !errors.Is(err, boom) {
t.Fatalf("destroy durability error = %v, want %v", err, boom)
}
manager.mu.RLock()
stillCurrent := manager.bySession[connSessionKey(target)] == target
manager.mu.RUnlock()
if !stillCurrent {
t.Fatal("durability failure retired the live target session")
}
}
// TestRPCDropAnswer 验证 rpc_drop_answer 以 rpc_result 包装 RpcDropAnswer 返回，
// 与 gotd/td 和 TDesktop 的请求/响应模型对齐。
// 与 iamxvbaba/td 和 TDesktop 的请求/响应模型对齐。
func TestRPCDropAnswer(t *testing.T) {
const dc = 2
addr, pub, _ := startTestServer(t, Options{DC: dc})
@ -411,27 +524,218 @@ func TestPingDelayDisconnectOddSeqAccepted(t *testing.T) {
}
}
// TestDestroyAuthKey 验证 MTProto service message destroy_auth_key 由连接层直接响应，
// 避免 TDesktop 清理旧 key 时落到业务 RPC fallback。
// TestDestroyAuthKey 验证 MTProto service message destroy_auth_key 无论裸发，
// 还是沿官方客户端的 invokeWithLayer/initConnection 路径发送，都由连接层
// 直接处理，并以绑定原请求的 rpc_result 回复。
func TestDestroyAuthKey(t *testing.T) {
tests := []struct {
name string
layer int
wrapped bool
}{
{name: "bare"},
{name: "layer225_wrapped", layer: 225, wrapped: true},
{name: "layer226_wrapped", layer: 226, wrapped: true},
{name: "layer227_wrapped", layer: 227, wrapped: true},
{name: "layer228_wrapped", layer: 228, wrapped: true},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
const dc = 2
addr, pub, srv := startTestServer(t, Options{DC: dc, LayerRPC: newAdmissionOnlyLayerRPC()})
conn, auth, cipher := dialHandshake(t, addr, dc, pub)
var request bin.Encoder = &destroyAuthKeyRequest{}
if test.wrapped {
request = &tg.InvokeWithLayerRequest{
Layer: test.layer,
Query: &tg.InitConnectionRequest{
APIID: 1,
DeviceModel: "destroy-key-test",
SystemVersion: "test",
AppVersion: "test",
SystemLangCode: "en",
LangCode: "en",
Query: &destroyAuthKeyRequest{},
},
}
}
clientMsgID := proto.NewMessageIDGen(time.Now)
reqMsgID := clientMsgID.New(proto.MessageFromClient)
sendEncrypted(t, conn, cipher, auth, reqMsgID, request)
replies := collectReplies(t, conn, cipher, auth.AuthKey, proto.ResultTypeID)
assertDestroyAuthKeyRPCResult(t, mustHave(t, replies, proto.ResultTypeID, "destroy_auth_key rpc_result"), reqMsgID, destroyAuthKeyOkTypeID)
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)
}
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
defer cancel()
var frame bin.Buffer
if err := conn.Recv(ctx, &frame); err == nil {
t.Fatal("destroy_auth_key requester remained readable after required rpc_result(ok)")
}
})
}
}
func assertDestroyAuthKeyRPCResult(t *testing.T, b *bin.Buffer, reqMsgID int64, wantInner uint32) {
t.Helper()
var result proto.Result
if err := result.Decode(b); err != nil {
t.Fatalf("decode destroy_auth_key rpc_result: %v", err)
}
if result.RequestMessageID != reqMsgID {
t.Fatalf("destroy_auth_key rpc_result.req_msg_id = %d, want %d", result.RequestMessageID, reqMsgID)
}
inner := &bin.Buffer{Buf: result.Result}
innerID, err := inner.PeekID()
if err != nil {
t.Fatalf("peek destroy_auth_key rpc_result inner: %v", err)
}
if innerID != wantInner || inner.Len() != bin.Word {
t.Fatalf("destroy_auth_key rpc_result inner = %#x/%d bytes, want %#x/%d", innerID, inner.Len(), wantInner, bin.Word)
}
}
func TestDestroyAuthKeyDeleteFailureReturnsCorrelatedFailAndKeepsConnection(t *testing.T) {
const dc = 2
addr, pub, srv := startTestServer(t, Options{DC: dc})
deleteErr := errors.New("delete auth key failed")
keys := &deleteFailAuthKeyStore{AuthKeyStore: memory.NewAuthKeyStore(), err: deleteErr}
addr, pub, srv := startTestServer(t, Options{DC: dc, AuthKeys: keys, LayerRPC: newAdmissionOnlyLayerRPC()})
conn, auth, cipher := dialHandshake(t, addr, dc, pub)
clientMsgID := proto.NewMessageIDGen(time.Now)
reqMsgID := clientMsgID.New(proto.MessageFromClient)
sendEncrypted(t, conn, cipher, auth, reqMsgID, &destroyAuthKeyRequest{})
ids := proto.NewMessageIDGen(time.Now)
destroyReqMsgID := ids.New(proto.MessageFromClient)
sendEncrypted(t, conn, cipher, auth, destroyReqMsgID, &destroyAuthKeyRequest{})
replies := collectReplies(t, conn, cipher, auth.AuthKey, proto.ResultTypeID)
assertDestroyAuthKeyRPCResult(t, mustHave(t, replies, proto.ResultTypeID, "destroy_auth_key fail rpc_result"), destroyReqMsgID, destroyAuthKeyFailTypeID)
replies := collectReplies(t, conn, cipher, auth.AuthKey, destroyAuthKeyOkTypeID)
mustHave(t, replies, 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)
if _, found, err := srv.authKeys.Get(context.Background(), auth.AuthKey.ID); err != nil || !found {
t.Fatalf("auth key after failed delete: found=%v err=%v", found, err)
}
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
defer cancel()
var frame bin.Buffer
if err := conn.Recv(ctx, &frame); err == nil {
t.Fatal("destroy_auth_key requester remained readable after required ok")
pingReqMsgID := ids.New(proto.MessageFromClient)
sendEncryptedWithSeq(t, conn, cipher, auth, pingReqMsgID, 3, &mt.PingRequest{PingID: 99})
pongReplies := collectReplies(t, conn, cipher, auth.AuthKey, mt.PongTypeID)
var pong mt.Pong
if err := pong.Decode(mustHave(t, pongReplies, mt.PongTypeID, "pong after failed destroy_auth_key")); err != nil {
t.Fatalf("decode pong after failed destroy_auth_key: %v", err)
}
if pong.MsgID != pingReqMsgID || pong.PingID != 99 {
t.Fatalf("pong after failed destroy_auth_key = %+v", pong)
}
}
func TestWrappedDestroyAuthKeyTrailingBytesDoNotDelete(t *testing.T) {
const dc = 2
addr, pub, srv := startTestServer(t, Options{DC: dc, LayerRPC: newAdmissionOnlyLayerRPC()})
conn, auth, cipher := dialHandshake(t, addr, dc, pub)
ids := proto.NewMessageIDGen(time.Now)
reqMsgID := ids.New(proto.MessageFromClient)
request := &tg.InvokeWithLayerRequest{
Layer: 228,
Query: &tg.InitConnectionRequest{
APIID: 1, DeviceModel: "malformed-destroy-key-test", SystemVersion: "test",
AppVersion: "test", SystemLangCode: "en", LangCode: "en",
Query: &trailingDestroyAuthKeyRequest{},
},
}
sendEncrypted(t, conn, cipher, auth, reqMsgID, request)
replies := collectReplies(t, conn, cipher, auth.AuthKey, proto.ResultTypeID)
var result proto.Result
if err := result.Decode(mustHave(t, replies, proto.ResultTypeID, "malformed destroy_auth_key rpc_result")); err != nil {
t.Fatalf("decode malformed destroy_auth_key rpc_result: %v", err)
}
if result.RequestMessageID != reqMsgID {
t.Fatalf("malformed destroy_auth_key req_msg_id = %d, want %d", result.RequestMessageID, reqMsgID)
}
var rpcErr mt.RPCError
if err := rpcErr.Decode(&bin.Buffer{Buf: result.Result}); err != nil {
t.Fatalf("decode malformed destroy_auth_key RPC error: %v", err)
}
if rpcErr.ErrorCode != 400 || rpcErr.ErrorMessage != "INPUT_REQUEST_INVALID" {
t.Fatalf("malformed destroy_auth_key RPC error = %+v", rpcErr)
}
if _, found, err := srv.authKeys.Get(context.Background(), auth.AuthKey.ID); err != nil || !found {
t.Fatalf("auth key after malformed wrapped destroy: found=%v err=%v", found, err)
}
}
func TestWrappedDestroyAuthKeySemanticWrapperDoesNotDelete(t *testing.T) {
const dc = 2
addr, pub, srv := startTestServer(t, Options{DC: dc, LayerRPC: newAdmissionOnlyLayerRPC()})
conn, auth, cipher := dialHandshake(t, addr, dc, pub)
ids := proto.NewMessageIDGen(time.Now)
reqMsgID := ids.New(proto.MessageFromClient)
request := &tg.InvokeWithLayerRequest{
Layer: 228,
Query: &tg.InitConnectionRequest{
APIID: 1, DeviceModel: "semantic-wrapper-destroy-key-test", SystemVersion: "test",
AppVersion: "test", SystemLangCode: "en", LangCode: "en",
Query: &tg.InvokeAfterMsgRequest{
MsgID: 1,
Query: &destroyAuthKeyRequest{},
},
},
}
sendEncrypted(t, conn, cipher, auth, reqMsgID, request)
replies := collectReplies(t, conn, cipher, auth.AuthKey, proto.ResultTypeID)
var result proto.Result
if err := result.Decode(mustHave(t, replies, proto.ResultTypeID, "semantic-wrapper destroy_auth_key rpc_result")); err != nil {
t.Fatalf("decode semantic-wrapper destroy_auth_key rpc_result: %v", err)
}
if result.RequestMessageID != reqMsgID {
t.Fatalf("semantic-wrapper destroy_auth_key req_msg_id = %d, want %d", result.RequestMessageID, reqMsgID)
}
var rpcErr mt.RPCError
if err := rpcErr.Decode(&bin.Buffer{Buf: result.Result}); err != nil {
t.Fatalf("decode semantic-wrapper destroy_auth_key RPC error: %v", err)
}
if rpcErr.ErrorCode != 400 || rpcErr.ErrorMessage != "INPUT_REQUEST_INVALID" {
t.Fatalf("semantic-wrapper destroy_auth_key RPC error = %+v", rpcErr)
}
if _, found, err := srv.authKeys.Get(context.Background(), auth.AuthKey.ID); err != nil || !found {
t.Fatalf("auth key after semantic-wrapper destroy_auth_key: found=%v err=%v", found, err)
}
}
func TestWrappedDestroyAuthKeyMixedContainerIsRejectedAtomically(t *testing.T) {
const dc = 2
addr, pub, srv := startTestServer(t, Options{DC: dc, LayerRPC: newAdmissionOnlyLayerRPC()})
conn, auth, cipher := dialHandshake(t, addr, dc, pub)
ids := proto.NewMessageIDGen(time.Now)
destroyBody := encodeClientMessageBodyForTest(t, &tg.InvokeWithLayerRequest{
Layer: 228,
Query: &tg.InitConnectionRequest{
APIID: 1, DeviceModel: "mixed-destroy-key-test", SystemVersion: "test",
AppVersion: "test", SystemLangCode: "en", LangCode: "en",
Query: &destroyAuthKeyRequest{},
},
})
pingBody := encodeClientMessageBodyForTest(t, &mt.PingRequest{PingID: 7})
destroyMsgID := ids.New(proto.MessageFromClient)
pingMsgID := ids.New(proto.MessageFromClient)
outerMsgID := ids.New(proto.MessageFromClient)
container := &proto.MessageContainer{Messages: []proto.Message{
{ID: destroyMsgID, SeqNo: 1, Bytes: len(destroyBody), Body: destroyBody},
{ID: pingMsgID, SeqNo: 3, Bytes: len(pingBody), Body: pingBody},
}}
sendEncrypted(t, conn, cipher, auth, outerMsgID, container)
replies := collectReplies(t, conn, cipher, auth.AuthKey, mt.BadMsgNotificationTypeID)
var bad mt.BadMsgNotification
if err := bad.Decode(mustHave(t, replies, mt.BadMsgNotificationTypeID, "bad_msg for mixed destroy_auth_key container")); err != nil {
t.Fatalf("decode mixed destroy_auth_key bad_msg: %v", err)
}
if bad.BadMsgID != outerMsgID || bad.ErrorCode != badMsgContainer {
t.Fatalf("mixed destroy_auth_key bad_msg = %+v, want msg_id=%d code=%d", bad, outerMsgID, badMsgContainer)
}
if _, found, err := srv.authKeys.Get(context.Background(), auth.AuthKey.ID); err != nil || !found {
t.Fatalf("auth key after mixed destroy_auth_key container: found=%v err=%v", found, err)
}
}

12
internal/mtprotoedge/exchange.go
View file

@ -8,12 +8,12 @@ import (
"go.uber.org/zap"
"github.com/gotd/td/bin"
"github.com/gotd/td/crypto"
"github.com/gotd/td/exchange"
"github.com/gotd/td/mt"
"github.com/gotd/td/proto/codec"
"github.com/gotd/td/transport"
"github.com/iamxvbaba/td/bin"
"github.com/iamxvbaba/td/crypto"
"github.com/iamxvbaba/td/exchange"
"github.com/iamxvbaba/td/mt"
"github.com/iamxvbaba/td/proto/codec"
"github.com/iamxvbaba/td/transport"
"telesrv/internal/store"
)

18
internal/mtprotoedge/exchange_compat.go
View file

@ -13,14 +13,14 @@ import (
gofaster "github.com/go-faster/errors"
"go.uber.org/zap"
"github.com/gotd/td/bin"
"github.com/gotd/td/clock"
"github.com/gotd/td/crypto"
"github.com/gotd/td/exchange"
"github.com/gotd/td/mt"
"github.com/gotd/td/proto"
"github.com/gotd/td/proto/codec"
"github.com/gotd/td/transport"
"github.com/iamxvbaba/td/bin"
"github.com/iamxvbaba/td/clock"
"github.com/iamxvbaba/td/crypto"
"github.com/iamxvbaba/td/exchange"
"github.com/iamxvbaba/td/mt"
"github.com/iamxvbaba/td/proto"
"github.com/iamxvbaba/td/proto/codec"
"github.com/iamxvbaba/td/transport"
)
// runServerExchange is a gotd server exchange compatibility shim.
@ -425,7 +425,7 @@ func (s serverExchangeCompat) readUnencrypted(ctx context.Context, b *bin.Buffer
if err := msg.Decode(b); err != nil {
return err
}
if proto.MessageID(msg.MessageID).Type() != proto.MessageFromClient {
if !validClientMessageIDBits(msg.MessageID) {
return gofaster.New("bad msg type")
}
b.ResetTo(msg.MessageData)

12
internal/mtprotoedge/exchange_test.go
View file

@ -15,12 +15,12 @@ import (
"go.uber.org/zap/zaptest"
"github.com/gotd/log/logzap"
"github.com/gotd/td/bin"
"github.com/gotd/td/exchange"
"github.com/gotd/td/mt"
tgproto "github.com/gotd/td/proto"
"github.com/gotd/td/proto/codec"
"github.com/gotd/td/transport"
"github.com/iamxvbaba/td/bin"
"github.com/iamxvbaba/td/exchange"
"github.com/iamxvbaba/td/mt"
tgproto "github.com/iamxvbaba/td/proto"
"github.com/iamxvbaba/td/proto/codec"
"github.com/iamxvbaba/td/transport"
"telesrv/internal/store"
"telesrv/internal/store/memory"

4
internal/mtprotoedge/flush_identity_test.go
View file

@ -4,8 +4,8 @@ import (
"context"
"testing"
"github.com/gotd/td/proto"
"github.com/gotd/td/tg"
"github.com/iamxvbaba/td/proto"
"github.com/iamxvbaba/td/tg"
)
// TestRunFlushDiscardsBatchOnIdentitySwitch 验证排空进行中连接易主（登出/换号致

6
internal/mtprotoedge/frame_budget.go
View file

@ -7,9 +7,9 @@ import (
"io"
"sync/atomic"
"github.com/gotd/td/bin"
"github.com/gotd/td/proto/codec"
"github.com/gotd/td/transport"
"github.com/iamxvbaba/td/bin"
"github.com/iamxvbaba/td/proto/codec"
"github.com/iamxvbaba/td/transport"
)
const defaultInboundFrameGlobalMaxBytes int64 = 512 << 20

6
internal/mtprotoedge/frame_budget_test.go
View file

@ -10,9 +10,9 @@ import (
"testing"
"time"
"github.com/gotd/td/bin"
"github.com/gotd/td/proto/codec"
"github.com/gotd/td/transport"
"github.com/iamxvbaba/td/bin"
"github.com/iamxvbaba/td/proto/codec"
"github.com/iamxvbaba/td/transport"
)
type frameBudgetTestConn struct {

8
internal/mtprotoedge/handshake_reqpq_cap_test.go
View file

@ -5,10 +5,10 @@ import (
"errors"
"testing"
"github.com/gotd/td/bin"
"github.com/gotd/td/mt"
"github.com/gotd/td/proto"
"github.com/gotd/td/transport"
"github.com/iamxvbaba/td/bin"
"github.com/iamxvbaba/td/mt"
"github.com/iamxvbaba/td/proto"
"github.com/iamxvbaba/td/transport"
)
// reqPQConn 是只会不断返回同一个 req_pq_multi 帧的假 transport.Conn，用于驱动 bufferedConn

115
internal/mtprotoedge/helpers_test.go
View file

@ -11,13 +11,97 @@ import (
"go.uber.org/zap/zaptest"
"github.com/gotd/log/logzap"
"github.com/gotd/td/bin"
"github.com/gotd/td/crypto"
"github.com/gotd/td/exchange"
"github.com/gotd/td/proto"
"github.com/gotd/td/transport"
"github.com/iamxvbaba/td/bin"
"github.com/iamxvbaba/td/crypto"
"github.com/iamxvbaba/td/exchange"
"github.com/iamxvbaba/td/proto"
"github.com/iamxvbaba/td/tg"
"github.com/iamxvbaba/td/transport"
)
// legacyCanonicalTestConn explicitly declares the canonical-only profile used
// by old connection-state tests. Production exact-path tests must instead call
// FreezeLayerProfile/SeedLayerProfile with protocol evidence.
func legacyCanonicalTestConn(t testing.TB, c *Conn) *Conn {
return legacyLayerWireTestConn(t, c, int(tg.LayerProfileCanonical))
}
// legacyLayerWireTestConn preserves only the old tests' profile setup. It does
// not enable any wire conversion; application values still need an exact
// generated binding at the outbound boundary.
func legacyLayerWireTestConn(t testing.TB, c *Conn, layer int) *Conn {
t.Helper()
if c == nil {
t.Fatal("nil legacy exact-layer test Conn")
}
profile, ok := tg.ResolveLayerProfile(layer)
if !ok {
t.Fatalf("unsupported generated test Layer %d", layer)
}
if err := c.FreezeLayerProfile(profile); err != nil {
t.Fatalf("freeze generated test Layer %d: %v", layer, err)
}
c.setLegacyClientLayer(layer)
return c
}
// exactTestUpdatesEncoded gives transport/state-machine tests an explicit
// generated session binding without invoking the production fan-out cache.
// Tests which assert wire conversion use layerUpdatesFanout directly instead.
func exactTestUpdatesEncoded(t testing.TB, c *Conn, body []byte) *encodedOutboundMessage {
t.Helper()
if c == nil {
t.Fatal("nil exact test Conn")
}
state := c.LayerProfileState()
if state.Origin == LayerProfileUnknown {
t.Fatal("exact test Conn has no generated Layer profile")
}
return &encodedOutboundMessage{
body: append([]byte(nil), body...),
typeID: tg.UpdatesTooLongTypeID,
layer: &outboundLayerBinding{
profile: state.Profile,
typ: tg.LayerClassUpdatesType().Ref(),
epoch: state.Epoch,
},
}
}
func exactTestUpdatesTooLong(t testing.TB, c *Conn) *encodedOutboundMessage {
t.Helper()
var body bin.Buffer
if err := (&tg.UpdatesTooLong{}).Encode(&body); err != nil {
t.Fatalf("encode exact test updatesTooLong: %v", err)
}
return exactTestUpdatesEncoded(t, c, body.Raw())
}
// opaqueExactTestRPCResult is an explicit request-bound capability for tests
// of compression, retention and delivery mechanics. Semantic result conversion
// is covered by generated dispatcher tests; no production path constructs it.
type opaqueExactTestRPCResult struct{ result bin.Encoder }
func (r *opaqueExactTestRPCResult) Encode(b *bin.Buffer) error { return r.result.Encode(b) }
func (r *opaqueExactTestRPCResult) exactLayerRPCResultBinding() outboundLayerBinding {
return outboundLayerBinding{
profile: tg.LayerProfileCanonical,
typ: tg.LayerClassUpdatesType().Ref(),
kind: outboundLayerBindingRequest,
}
}
func exactTestRPCResult(result bin.Encoder) bin.Encoder {
if result == nil || isLayerInvariantRPCResultEncoder(result) {
return result
}
if _, ok := result.(exactLayerRPCResultEncoder); ok {
return result
}
return &opaqueExactTestRPCResult{result: result}
}
// startTestServer 生成 RSA key、监听随机端口并启动 Server，返回监听地址与公钥。
// 通过 t.Cleanup 自动取消并校验优雅退出。opts 的 RSAKey/Logger/DC 会被补默认。
func startTestServer(t *testing.T, opts Options) (addr string, pub exchange.PublicKey, srv *Server) {
@ -92,6 +176,27 @@ func dialTransportOnly(t *testing.T, addr string) transport.Conn {
return conn
}
// freezeActiveTestSessionProfile gives low-level transport fixtures the exact
// profile that a production invokeWithLayer admission would have proven. It is
// intentionally explicit: handshake/new_session_created alone never implies a
// TL Layer, and production push code must keep failing closed in that state.
func freezeActiveTestSessionProfile(t *testing.T, sessions *SessionManager, authKeyID [8]byte, sessionID int64, profile tg.LayerProfile) {
t.Helper()
if sessions == nil {
t.Fatal("freeze test session profile on nil SessionManager")
}
key := sessionKey{authKeyID: authKeyID, sessionID: sessionID}
sessions.mu.RLock()
c := sessions.bySession[key]
sessions.mu.RUnlock()
if c == nil {
t.Fatalf("active test session %x/%d is missing", authKeyID, sessionID)
}
if err := c.FreezeLayerProfile(profile); err != nil {
t.Fatalf("freeze active test session profile %d: %v", profile, err)
}
}
// sendEncrypted 用 client cipher 加密并发送一条带 msgID 的消息。
func sendEncrypted(t *testing.T, conn transport.Conn, cipher crypto.Cipher, auth exchange.ClientExchangeResult, msgID int64, msg bin.Encoder) {
t.Helper()

1109
internal/mtprotoedge/inbound_layer_rpc.go Normal file
View file

File diff suppressed because it is too large Load diff

1899
internal/mtprotoedge/inbound_layer_rpc_test.go Normal file
View file

File diff suppressed because it is too large Load diff

178
internal/mtprotoedge/inbound_preflight.go
View file

@ -10,9 +10,10 @@ import (
"go.uber.org/zap"
"github.com/gotd/td/bin"
"github.com/gotd/td/mt"
"github.com/gotd/td/proto"
"github.com/iamxvbaba/td/bin"
"github.com/iamxvbaba/td/mt"
"github.com/iamxvbaba/td/proto"
"github.com/iamxvbaba/td/tg"
)
type inboundItemKind uint8
@ -34,6 +35,7 @@ const (
inboundItemDestroyAuthKey
inboundItemRPC
inboundItemCapacityError
inboundItemRPCAdmissionError
// inboundItemRewrappedRPC is an initConnection retry whose exact inner TL
// request is already executing (or completed) under the client's old msg_id.
// It never dispatches business code a second time.
@ -48,13 +50,33 @@ const (
)
type inboundItem struct {
kind inboundItemKind
msgID int64
seqNo int32
typeID uint32
content bool
body []byte
payload any
kind inboundItemKind
msgID int64
admissionSeq uint64
seqNo int32
typeID uint32
content bool
body []byte
payload any
admitted tg.LayerRequest
method string
replayAfterSuccessfulDelivery func() error
layerProfileEvidenceFreshness inboundLayerProfileEvidenceFreshness
}
type inboundLayerProfileEvidenceFreshness uint8
const (
// Unspecified is retained for focused force-style unit tests which construct
// inboundItem directly, outside MTProto envelope preflight. Production items
// are always classified from the frame's one clock sample.
inboundLayerProfileEvidenceFreshnessUnspecified inboundLayerProfileEvidenceFreshness = iota
inboundLayerProfileEvidenceFresh
inboundLayerProfileEvidenceRequestBound
)
func (i inboundItem) profileEvidenceFresh() bool {
return i.layerProfileEvidenceFreshness != inboundLayerProfileEvidenceRequestBound
}
type stagedClientMessage struct {
@ -75,13 +97,23 @@ type inboundPlan struct {
rpcTasks []inboundRPC
rpcOwners []*rpcResultOwnerLease
rewrapAliases []*rpcRewrapAlias
rewrapIndices []int
}
func (p *inboundPlan) close() {
if p == nil {
return
}
// Drop exact typed request graphs and uncommitted task closures before their
// materialization reservation becomes reusable. Otherwise an abort could
// advertise the same bytes to another connection while this plan still kept
// the old graph reachable until its caller returned.
for i := range p.items {
p.items[i].admitted = tg.LayerRequest{}
}
for i := range p.rpcTasks {
p.rpcTasks[i] = inboundRPC{}
}
p.rpcTasks = nil
if p.rpcReservation != nil {
p.rpcReservation.abort()
p.rpcReservation = nil
@ -97,7 +129,6 @@ func (p *inboundPlan) close() {
}
}
p.rewrapAliases = nil
p.rewrapIndices = nil
for i := len(p.releases) - 1; i >= 0; i-- {
p.releases[i]()
}
@ -117,15 +148,47 @@ func (p *inboundPlan) commitRewrapAliases(s *Server) error {
}
}
p.rewrapAliases = nil
p.rewrapIndices = nil
return err
}
}
p.rewrapAliases = nil
p.rewrapIndices = nil
return nil
}
// rejectNewRPCOwners turns only ownership acquired by this batch into bounded
// capacity responses. Existing completed replays and pending joins remain
// active: canceling them would either lose a response or publish into another
// request's flight.
func (p *inboundPlan) rejectNewRPCOwners(indices []int) {
if p == nil {
return
}
for _, index := range indices {
if index >= 0 && index < len(p.items) {
p.items[index].kind = inboundItemCapacityError
}
}
kept := p.rewrapAliases[:0]
for _, alias := range p.rewrapAliases {
if alias == nil || alias.newOwner == nil {
kept = append(kept, alias)
continue
}
if alias.itemIndex >= 0 && alias.itemIndex < len(p.items) {
p.items[alias.itemIndex].kind = inboundItemCapacityError
p.items[alias.itemIndex].payload = nil
}
alias.releaseCandidate()
// This owner was acquired by the rejected batch and is not present in
// plan.rpcOwners because it belonged to a rewrap alias. Abort it here
// before dropping the alias, otherwise the exact flight remains pending
// forever with no task or publisher able to complete it.
alias.newOwner.Abort()
alias.newOwner = nil
}
p.rewrapAliases = kept
}
func (p *inboundPlan) commitRPCBatch() error {
if p == nil || p.rpcReservation == nil {
return nil
@ -405,6 +468,14 @@ func (s *Server) walkInbound(
if err != nil {
return err
}
if validateInboundMessageID(budget.now, msgID, false) == 0 {
item.layerProfileEvidenceFreshness = inboundLayerProfileEvidenceFresh
} else {
// Inner container messages deliberately bypass the wall-clock rejection
// above, but old/future ids are request-bound and cannot publish mutable
// Layer/init/readiness/auth-bind evidence.
item.layerProfileEvidenceFreshness = inboundLayerProfileEvidenceRequestBound
}
plan.items = append(plan.items, item)
if content {
plan.ackIDs = append(plan.ackIDs, msgID)
@ -413,7 +484,7 @@ func (s *Server) walkInbound(
}
func validateInboundMessageID(now time.Time, msgID int64, insideContainer bool) int {
if msgID == 0 || proto.MessageID(msgID).Type() != proto.MessageFromClient {
if !validClientMessageIDBits(msgID) {
return badMsgIDInvalidBits
}
// A container's outer envelope supplies the wall-clock admission boundary for
@ -432,6 +503,11 @@ func validateInboundMessageID(now time.Time, msgID int64, insideContainer bool)
return 0
}
func validClientMessageIDBits(msgID int64) bool {
return msgID > 0 && uint32(msgID) != 0 &&
proto.MessageID(msgID).Type() == proto.MessageFromClient
}
func appendInboundDuplicate(plan *inboundPlan, msgID int64, seqNo int32, typeID uint32, record clientMsgRecord) error {
plan.includeLogicalID(msgID)
kind := inboundItemDuplicate
@ -653,6 +729,9 @@ func preflightInboundItem(msgID int64, seqNo int32, typeID uint32, content bool,
// 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 {
if s.layerRPC != nil {
return s.prepareInboundLayerRPCBatch(ctx, c, plan)
}
// 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
@ -676,8 +755,11 @@ func (s *Server) prepareInboundRPCBatch(ctx context.Context, c *Conn, plan *inbo
method := s.typeName(item.typeID)
init, isInitRewrap := decodeRPCRewrapInit(item.body)
if isInitRewrap {
firstInit := !c.rpcRewrapInitialized.Swap(true)
c.SetClientLayer(init.layer)
firstInit := false
if item.profileEvidenceFresh() {
firstInit = !c.rpcRewrapInitialized.Swap(true)
c.setLegacyClientLayer(init.layer)
}
if candidate := s.rpcRewrap.claim(c, init.inner); candidate != nil {
claim, err := s.rpcResults.Acquire(c.authKeyID, c.sessionID, item.msgID)
if errors.Is(err, ErrRPCResultFlightCapacity) {
@ -701,10 +783,9 @@ func (s *Server) prepareInboundRPCBatch(ctx context.Context, c *Conn, plan *inbo
item.kind = inboundItemRewrappedRPC
item.payload = claim.waiter
plan.rewrapAliases = append(plan.rewrapAliases, &rpcRewrapAlias{
conn: c, newReqID: item.msgID, method: candidate.method,
conn: c, itemIndex: i, newReqID: item.msgID, method: candidate.method,
oldWaiter: claim.waiter, observeInit: firstInit, init: init,
})
plan.rewrapIndices = append(plan.rewrapIndices, i)
case rpcResultAcquireOwner:
if ownersInPlan == nil {
ownersInPlan = make(map[int64]*rpcResultOwnerLease)
@ -713,13 +794,12 @@ func (s *Server) prepareInboundRPCBatch(ctx context.Context, c *Conn, plan *inbo
item.kind = inboundItemRewrappedRPC
item.payload = claim.owner
plan.rewrapAliases = append(plan.rewrapAliases, &rpcRewrapAlias{
conn: c, newReqID: item.msgID, method: candidate.method,
conn: c, itemIndex: i, newReqID: item.msgID, method: candidate.method,
oldWaiter: candidate.waiter, newOwner: claim.owner,
sourceConn: candidate.source, sourceOwner: candidate.owner,
observeInit: firstInit, init: init,
candidate: candidate, registry: s.rpcRewrap,
})
plan.rewrapIndices = append(plan.rewrapIndices, i)
default:
s.rpcRewrap.release(candidate)
return ErrRPCResultFlightInvalid
@ -732,7 +812,7 @@ func (s *Server) prepareInboundRPCBatch(ctx context.Context, c *Conn, plan *inbo
zap.String("auth_key_id", c.authKeyHex), zap.Int64("session_id", c.sessionID))
continue
}
} else if c.rpcRewrapInitialized.Load() && !clearedPostInitCandidates {
} else if item.profileEvidenceFresh() && c.rpcRewrapInitialized.Load() && !clearedPostInitCandidates {
// A naked request after this connection has observed initConnection is
// event-level proof that the client finished moving its old running set.
// Retire any unmatched candidates without a timer.
@ -773,7 +853,7 @@ func (s *Server) prepareInboundRPCBatch(ctx context.Context, c *Conn, plan *inbo
item.kind = inboundItemRewrappedRPC
item.payload = claim.waiter
plan.rewrapAliases = append(plan.rewrapAliases, &rpcRewrapAlias{
conn: c, newReqID: item.msgID, method: method, oldWaiter: claim.waiter,
conn: c, itemIndex: i, newReqID: item.msgID, method: method, oldWaiter: claim.waiter,
})
}
case rpcResultAcquireOwner:
@ -807,17 +887,7 @@ func (s *Server) prepareInboundRPCBatch(ctx context.Context, c *Conn, plan *inbo
for _, index := range indices {
plan.items[index].kind = inboundItemCapacityError
}
for _, index := range plan.rewrapIndices {
plan.items[index].kind = inboundItemCapacityError
}
for _, alias := range plan.rewrapAliases {
alias.releaseCandidate()
if alias.newOwner != nil {
alias.newOwner.Abort()
}
}
plan.rewrapAliases = nil
plan.rewrapIndices = nil
plan.rejectNewRPCOwners(indices)
return nil
}
if len(specs) == 0 {
@ -830,17 +900,7 @@ func (s *Server) prepareInboundRPCBatch(ctx context.Context, c *Conn, plan *inbo
for _, index := range indices {
plan.items[index].kind = inboundItemCapacityError
}
for _, index := range plan.rewrapIndices {
plan.items[index].kind = inboundItemCapacityError
}
for _, alias := range plan.rewrapAliases {
alias.releaseCandidate()
if alias.newOwner != nil {
alias.newOwner.Abort()
}
}
plan.rewrapAliases = nil
plan.rewrapIndices = nil
plan.rejectNewRPCOwners(indices)
return nil
}
return err
@ -874,7 +934,7 @@ func (s *Server) executeInboundPlan(ctx context.Context, cs *connState, c *Conn,
}
case inboundItemReplayRPC:
if encoded, _ := item.payload.(*encodedOutboundMessage); encoded != nil {
if err := s.sendCachedRPCResult(ctx, c, encoded); err != nil {
if err := s.sendCachedRPCResultWithHook(ctx, c, encoded, item.replayAfterSuccessfulDelivery); err != nil {
return err
}
} else if err := s.replayRPCResultByRequest(ctx, c, item.msgID); err != nil {
@ -942,13 +1002,22 @@ func (s *Server) executeInboundPlan(ctx context.Context, cs *connState, c *Conn,
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.SendRequiredControl(ctx, proto.MessageServerResponse, &destroyAuthKeyFail{})
return c.SendRequiredControl(ctx, proto.MessageServerResponse, &destroyAuthKeyRPCResult{
RequestMessageID: item.msgID,
ResultTypeID: destroyAuthKeyFailTypeID,
})
}
if registry, ok := s.layerRPC.(LayerRPCSessionProfileRegistry); ok {
registry.ForgetNegotiatedAuthKey(c.authKeyID)
}
// Fence every other active/claiming generation before acknowledging the
// deletion. The exact requester remains writable only long enough to put the
// required destroy_auth_key_ok frame on the wire.
// request-correlated rpc_result(destroy_auth_key_ok) frame on the wire.
s.conns.CloseSessionsForRawAuthKeyExceptConn(c.authKeyID, c)
if err := c.SendRequiredControl(ctx, proto.MessageServerResponse, &destroyAuthKeyOk{}); err != nil {
if err := c.SendRequiredControl(ctx, proto.MessageServerResponse, &destroyAuthKeyRPCResult{
RequestMessageID: item.msgID,
ResultTypeID: destroyAuthKeyOkTypeID,
}); err != nil {
return err
}
c.beginTerminalShutdown()
@ -959,12 +1028,23 @@ func (s *Server) executeInboundPlan(ctx context.Context, cs *connState, c *Conn,
// execution begins; commitRPCBatch publishes them after all protocol barriers.
continue
case inboundItemCapacityError:
if owner, _ := item.payload.(*rpcResultOwnerLease); owner != nil {
owner.CompleteExecution(false)
}
if err := s.sendResult(ctx, c, item.msgID, &mt.RPCError{
ErrorCode: 420,
ErrorMessage: "FLOOD_WAIT_1",
}); err != nil {
return err
}
case inboundItemRPCAdmissionError:
rpcErr, _ := item.payload.(*mt.RPCError)
if rpcErr == nil {
rpcErr = &mt.RPCError{ErrorCode: 400, ErrorMessage: "INPUT_REQUEST_INVALID"}
}
if err := s.sendResult(ctx, c, item.msgID, rpcErr); err != nil {
return err
}
default:
return fmt.Errorf("unknown inbound item kind %d", item.kind)
}

6
internal/mtprotoedge/inbound_preflight_bench_test.go
View file

@ -4,9 +4,9 @@ import (
"testing"
"time"
"github.com/gotd/td/bin"
"github.com/gotd/td/proto"
"github.com/gotd/td/tg"
"github.com/iamxvbaba/td/bin"
"github.com/iamxvbaba/td/proto"
"github.com/iamxvbaba/td/tg"
)
func BenchmarkInboundPlan32RPCContainer(b *testing.B) {

233
internal/mtprotoedge/inbound_rpc.go
View file

@ -12,8 +12,10 @@ import (
// ErrInboundRPCQueueFull 表示 inbound RPC 已触达单连接或进程级预算。
var ErrInboundRPCQueueFull = errors.New("inbound rpc queue full")
// maxInflightRPCBytes 是单连接所有已预留、排队和执行中 RPC body 的总字节上限。
// 进程级预算在 Copy 前先兜底；这里再隔离单个连接，避免一个客户端独占全局内存。
// maxInflightRPCBytes 是单连接所有已预留、排队和执行中 RPC 的内存 charge 上限。
// legacy 路径的 charge 等于 copied body；exact 路径在 typed decode 前按 wire 大小、
// 生成对象/向量/interface/string/bytes 放大保守计费。它不是可接收 wire bytes 上限。
// 进程级预算先兜底；这里再隔离单个连接，避免一个客户端独占全局内存。
const maxInflightRPCBytes = 32 << 20 // 32 MiB
// rpcCloseWaitTimeout 是连接/Server 关闭时等待在途 RPC 或共享 worker 退出的上限。
@ -37,8 +39,49 @@ type inboundRPC struct {
release func()
budget *inboundRPCGlobalReservation
ticket *inboundRPCTicket
gate *inboundRPCGate
}
// inboundRPCGate keeps dependency waits out of the shared worker pool. A
// gated task remains within the ordinary queue/task/byte budgets, but workers
// skip it until every prerequisite publishes its terminal execution outcome.
type inboundRPCGate struct {
remaining atomic.Int32
failed atomic.Bool
ready atomic.Bool
wake func()
}
func newInboundRPCGate(prerequisites int, wake func()) *inboundRPCGate {
if prerequisites < 0 {
prerequisites = 0
}
g := &inboundRPCGate{wake: wake}
// One sentinel keeps the gate closed while subscribers are installed; the
// caller resolves it after registration is complete.
g.remaining.Store(int32(prerequisites + 1))
return g
}
func (g *inboundRPCGate) resolve(success bool) {
if g == nil {
return
}
if !success {
g.failed.Store(true)
}
remaining := g.remaining.Add(-1)
if remaining < 0 {
panic("mtproto inbound RPC gate counter underflow")
}
if remaining == 0 && g.ready.CompareAndSwap(false, true) && g.wake != nil {
g.wake()
}
}
func (g *inboundRPCGate) runnable() bool { return g == nil || g.ready.Load() }
func (g *inboundRPCGate) success() bool { return g == nil || (g.ready.Load() && !g.failed.Load()) }
type inboundRPCTicket struct {
onTimeout func()
}
@ -78,7 +121,8 @@ type inboundRPCGlobalReservation struct {
}
// inboundRPCSpec 是 container preflight 与 RPC scheduler 之间的有界 admission 描述。
// method 仅用于 metrics，size 是在 Copy 之前必须预留的 request body 字节数。
// method 仅用于 metrics；size 是 materialization charge。legacy 路径等于 Copy 前
// request body 字节数，exact 路径是 typed decode 前计算的保守内存上界。
type inboundRPCSpec struct {
method string
size int
@ -102,6 +146,7 @@ type inboundRPCBatchReservation struct {
}
var errInboundRPCBatchTaskCount = errors.New("inbound rpc batch task count mismatch")
var errInboundRPCBatchSelection = errors.New("inbound rpc batch selection is invalid")
func newInboundRPCScheduler(workers, maxTasks int, maxBytes int64) *inboundRPCScheduler {
if workers <= 0 {
@ -499,6 +544,76 @@ func (c *Conn) dropInboundRPCSpecs(specs []inboundRPCSpec, reason string) {
}
}
// retain keeps a subset of a provisional batch on the original connection and
// global reservations. Exact-layer admission uses this after typed decode has
// classified completed replays, pending joins, admission errors, and fresh
// owners. A retained fresh owner therefore never passes through a release then
// reacquire window where another connection could consume its memory/task
// budget. Entries not retained are returned immediately as one batch.
//
// indices are reservation-entry indices, must be strictly increasing, and have
// a one-to-one correspondence with specs. The original conservative byte
// charge is intentionally retained; a typed request remains reachable from its
// scheduler task until completion, so shrinking it to the wire size would make
// the materialized graph unaccounted.
func (r *inboundRPCBatchReservation) retain(indices []int, specs []inboundRPCSpec) error {
if r == nil || len(indices) != len(specs) || len(indices) > len(r.entries) {
return errInboundRPCBatchSelection
}
retained := make([]inboundRPCBatchEntry, len(indices))
removed := make([]*inboundRPCGlobalReservation, 0, len(r.entries)-len(indices))
var (
previous = -1
retainedSize int64
removedSize int64
)
selected := 0
for output, index := range indices {
if index <= previous || index < 0 || index >= len(r.entries) {
return errInboundRPCBatchSelection
}
previous = index
for selected < index {
entry := r.entries[selected]
removed = append(removed, entry.global)
removedSize += int64(entry.size)
selected++
}
entry := r.entries[index]
// The caller may improve the provisional metric label but may not increase
// or replace its byte charge after materialization has started.
if specs[output].size > entry.size {
return errInboundRPCBatchSelection
}
entry.method = specs[output].method
retained[output] = entry
retainedSize += int64(entry.size)
selected = index + 1
}
for selected < len(r.entries) {
entry := r.entries[selected]
removed = append(removed, entry.global)
removedSize += int64(entry.size)
selected++
}
c := r.conn
c.rpcMu.Lock()
c.rpcReserved -= len(removed)
c.inflightRPCBytes.Add(-removedSize)
r.entries = retained
r.totalSize = retainedSize
c.rpcMu.Unlock()
releaseInboundRPCGlobalBatch(removed)
if len(retained) == 0 {
// Finish the one reservation waiter as part of the same ownership
// transition. abort sees an empty batch and is therefore accounting-only.
r.abort()
}
return nil
}
// commit 在一次 rpcMu 临界区内把整批 task append 到队列并立即发布 ready token。
// 协议 barrier 必须在调用 commit 前完成；延迟发布 token 无法阻止已有 worker
// 从同一连接队列取走新任务，因此不提供虚假的 deferred-schedule 模式。
@ -581,7 +696,7 @@ func (r *inboundRPCBatchReservation) commit(tasks []inboundRPC) (result error) {
firstQueueLen = len(c.rpcQueue) + 1
c.rpcQueue = append(c.rpcQueue, prepared...)
queueCap = c.rpcQueueSize
if len(prepared) > 0 && c.rpcRunning < c.rpcMaxInflight && !c.rpcReady {
if c.rpcRunning < c.rpcMaxInflight && !c.rpcReady && c.hasRunnableInboundRPCLocked() {
c.rpcReady = true
reschedule = true
}
@ -636,21 +751,121 @@ func (c *Conn) takeInboundRPC() (task inboundRPC, ok, reschedule bool) {
if c.rpcClosed || len(c.rpcQueue) == 0 || c.rpcRunning >= c.rpcMaxInflight {
return inboundRPC{}, false, false
}
task = c.rpcQueue[0]
c.rpcQueue[0] = inboundRPC{}
c.rpcQueue = c.rpcQueue[1:]
index := -1
for i := range c.rpcQueue {
if c.rpcQueue[i].gate.runnable() {
index = i
break
}
}
if index < 0 {
return inboundRPC{}, false, false
}
task = c.rpcQueue[index]
if index == 0 {
// The overwhelmingly common FIFO path advances the slice head in O(1).
// Clear the departed element first so its request graph and closures are
// not retained by the shared backing array.
c.rpcQueue[0] = inboundRPC{}
c.rpcQueue = c.rpcQueue[1:]
} else {
copy(c.rpcQueue[index:], c.rpcQueue[index+1:])
last := len(c.rpcQueue) - 1
c.rpcQueue[last] = inboundRPC{}
c.rpcQueue = c.rpcQueue[:last]
}
if len(c.rpcQueue) == 0 {
c.rpcQueue = nil
}
c.rpcRunning++
c.rpcWG.Add(1)
if len(c.rpcQueue) > 0 && c.rpcRunning < c.rpcMaxInflight {
if c.rpcRunning < c.rpcMaxInflight && c.hasRunnableInboundRPCLocked() {
c.rpcReady = true
reschedule = true
}
return task, true, reschedule
}
func (c *Conn) hasRunnableInboundRPCLocked() bool {
if c.rpcReplayRestores > 0 {
return false
}
for i := range c.rpcQueue {
if c.rpcQueue[i].gate.runnable() {
return true
}
}
return false
}
// beginRPCReplayRestore installs a scheduler-level barrier without occupying a
// global RPC worker. The returned idempotent completion function wakes queued
// work only after the last overlapping restore has finished.
func (c *Conn) beginRPCReplayRestore() func() {
if c == nil {
return func() {}
}
active := false
unschedule := false
c.rpcMu.Lock()
if !c.rpcClosed && !c.isRetired() {
c.rpcReplayRestores++
active = true
if c.rpcReady {
c.rpcReady = false
unschedule = true
}
}
c.rpcMu.Unlock()
if unschedule && c.rpcScheduler != nil {
c.rpcScheduler.unschedule(c)
}
var once sync.Once
return func() {
once.Do(func() {
if !active {
return
}
reschedule := false
c.rpcMu.Lock()
c.rpcReplayRestores--
if c.rpcReplayRestores < 0 {
c.rpcMu.Unlock()
panic("mtproto inbound RPC replay restore barrier underflow")
}
if c.rpcReplayRestores == 0 && !c.rpcClosed && c.rpcRunning < c.rpcMaxInflight &&
!c.rpcReady && c.hasRunnableInboundRPCLocked() {
c.rpcReady = true
reschedule = true
}
c.rpcMu.Unlock()
if reschedule && c.rpcScheduler != nil {
c.rpcScheduler.schedule(c)
}
})
}
}
// wakeInboundRPC is safe from a flight completion callback. It does no work
// while the gate's task has not been committed yet; commit performs the same
// runnable scan and publishes the initial scheduler token.
func (c *Conn) wakeInboundRPC() {
if c == nil || c.rpcScheduler == nil {
return
}
reschedule := false
c.rpcMu.Lock()
if !c.rpcClosed && c.rpcRunning < c.rpcMaxInflight && !c.rpcReady && c.hasRunnableInboundRPCLocked() {
c.rpcReady = true
reschedule = true
}
c.rpcMu.Unlock()
if reschedule {
c.rpcScheduler.schedule(c)
}
}
func (c *Conn) runInboundRPC(task inboundRPC) {
defer c.finishInboundRPC(task)
@ -681,7 +896,7 @@ func (c *Conn) finishInboundRPC(task inboundRPC) {
c.rpcMu.Lock()
c.rpcRunning--
c.inflightRPCBytes.Add(-int64(task.size))
if !c.rpcClosed && len(c.rpcQueue) > 0 && c.rpcRunning < c.rpcMaxInflight && !c.rpcReady {
if !c.rpcClosed && c.rpcRunning < c.rpcMaxInflight && !c.rpcReady && c.hasRunnableInboundRPCLocked() {
c.rpcReady = true
reschedule = true
}

93
internal/mtprotoedge/inbound_rpc_test.go
View file

@ -108,6 +108,57 @@ func TestInboundRPCSchedulerBoundsConcurrentWork(t *testing.T) {
}
}
func TestInboundRPCSchedulerSkipsUnresolvedDependencyGate(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)
}()
blockedRan := make(chan struct{}, 1)
independentRan := make(chan struct{}, 1)
gate := newInboundRPCGate(1, c.wakeInboundRPC)
gate.resolve(true) // subscriber-installation sentinel; one dependency remains.
if err := c.enqueueInboundRPC(context.Background(), inboundRPC{
method: "invokeAfter",
gate: gate,
run: func(context.Context) error {
blockedRan <- struct{}{}
return nil
},
}); err != nil {
t.Fatal(err)
}
if err := c.enqueueInboundRPC(context.Background(), inboundRPC{
method: "independent",
run: func(context.Context) error {
independentRan <- struct{}{}
return nil
},
}); err != nil {
t.Fatal(err)
}
select {
case <-independentRan:
case <-time.After(time.Second):
t.Fatal("independent task was starved behind unresolved invokeAfter")
}
select {
case <-blockedRan:
t.Fatal("invokeAfter ran before its dependency completed")
default:
}
gate.resolve(true)
select {
case <-blockedRan:
case <-time.After(time.Second):
t.Fatal("resolved invokeAfter was not rescheduled")
}
}
func TestInboundRPCSchedulerFairAcrossConnections(t *testing.T) {
scheduler := newInboundRPCScheduler(1, 16, 1<<20)
c1 := newInboundTestConn(scheduler, 1, 4, time.Second)
@ -491,3 +542,45 @@ func TestInboundRPCCloseDrainsQueueAndReturnsBudgets(t *testing.T) {
t.Fatalf("connection inflight bytes after close = %d, want zero", got)
}
}
func TestInboundRPCReplayRestoreBarrierKeepsFollowingTaskOffWorkers(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)
}()
finishFirst := c.beginRPCReplayRestore()
finishSecond := c.beginRPCReplayRestore()
ran := make(chan struct{})
if err := c.enqueueInboundRPC(context.Background(), inboundRPC{
method: "following.naked.rpc",
size: 4,
run: func(context.Context) error {
close(ran)
return nil
},
}); err != nil {
t.Fatalf("enqueue following task: %v", err)
}
select {
case <-ran:
t.Fatal("following RPC ran before replay restore completed")
case <-time.After(30 * time.Millisecond):
}
finishFirst()
select {
case <-ran:
t.Fatal("one of two replay restores released the scheduler early")
case <-time.After(30 * time.Millisecond):
}
finishSecond()
select {
case <-ran:
case <-time.After(time.Second):
t.Fatal("following RPC did not run after the final replay restore")
}
}

469
internal/mtprotoedge/layer_admission_budget_test.go Normal file
View file

@ -0,0 +1,469 @@
package mtprotoedge
import (
"context"
"errors"
"reflect"
"strings"
"sync/atomic"
"testing"
"time"
"github.com/iamxvbaba/td/bin"
"github.com/iamxvbaba/td/clock"
"github.com/iamxvbaba/td/tg"
"go.uber.org/zap/zaptest"
appfiles "telesrv/internal/app/files"
"telesrv/internal/rpc"
)
func TestLayerRPCAdmissionMaterializationConstants(t *testing.T) {
constructors := make([]reflect.Type, 0, len(tg.TypesConstructorMap()))
for _, constructor := range tg.TypesConstructorMap() {
if typ := reflect.TypeOf(constructor()); typ != nil {
constructors = append(constructors, typ)
}
}
seen := make(map[reflect.Type]struct{})
var (
maxSize uintptr
maxType reflect.Type
visit func(reflect.Type)
)
visit = func(typ reflect.Type) {
if typ == nil {
return
}
if _, ok := seen[typ]; ok {
return
}
seen[typ] = struct{}{}
switch typ.Kind() {
case reflect.Pointer, reflect.Slice, reflect.Array:
visit(typ.Elem())
case reflect.Interface:
for _, candidate := range constructors {
// invokeWithLayer/query uses the deliberately broad bin.Object
// interface. Exact admission restricts that slot to generated RPC
// methods, so only request constructors are reachable there.
broad := typ.PkgPath() != "github.com/iamxvbaba/td/tg"
if candidate.Implements(typ) && (!broad || strings.HasSuffix(candidate.Elem().Name(), "Request")) {
visit(candidate)
}
}
case reflect.Struct:
if typ.PkgPath() == "github.com/iamxvbaba/td/tg" && typ.Size() > maxSize {
maxSize, maxType = typ.Size(), typ
}
for i := 0; i < typ.NumField(); i++ {
visit(typ.Field(i).Type)
}
}
}
for _, typ := range constructors {
if typ.Kind() == reflect.Pointer && strings.HasSuffix(typ.Elem().Name(), "Request") {
visit(typ)
}
}
if maxSize > layerRPCAdmissionStaticObjectBytes {
t.Fatalf("request-reachable generated TL object %v is %d bytes, exceeds admission ceiling %d", maxType, maxSize, layerRPCAdmissionStaticObjectBytes)
}
if layerRPCAdmissionGraphSlack != layerRPCAdmissionStaticObjectBytes*inboundLayerDecodeLimits.MaxDepth {
t.Fatalf("graph slack %d does not cover %d bytes across decode depth %d", layerRPCAdmissionGraphSlack, layerRPCAdmissionStaticObjectBytes, inboundLayerDecodeLimits.MaxDepth)
}
// Preserve enough room for the maximum upload payload plus any supported
// transparent wrapper/client-info envelope on a default connection.
if got := layerRPCAdmissionReservationSize(appfiles.MaxUploadPartBytes + (32 << 10)); got > maxInflightRPCBytes {
t.Fatalf("largest legal upload request charge = %d, exceeds default connection budget %d", got, maxInflightRPCBytes)
}
maxInt := int(^uint(0) >> 1)
if got := layerRPCAdmissionReservationSize(maxInt); got != maxInt {
t.Fatalf("saturating charge = %d, want max int %d", got, maxInt)
}
if got := layerRPCAdmissionReservationSize(-1); got != layerRPCAdmissionGraphSlack {
t.Fatalf("negative wire charge = %d, want fixed slack %d", got, layerRPCAdmissionGraphSlack)
}
}
type countingLayerRPCAdmission struct {
LayerRPCHandler
decodeCalls atomic.Int32
}
type failingReplayLayerRPC struct {
LayerRPCHandler
err error
}
func (h *failingReplayLayerRPC) PrepareAdmittedReplay(
context.Context,
[8]byte,
int64,
int64,
uint64,
tg.LayerRequest,
) (func() error, error) {
return nil, h.err
}
func (h *countingLayerRPCAdmission) AdmitLayer(profile tg.LayerProfile, b *bin.Buffer, limits tg.LayerDecodeLimits) (tg.LayerRequest, error) {
h.decodeCalls.Add(1)
return h.LayerRPCHandler.AdmitLayer(profile, b, limits)
}
func (h *countingLayerRPCAdmission) AdmitUnprofiled(b *bin.Buffer, limits tg.LayerDecodeLimits) (tg.LayerRequest, error) {
h.decodeCalls.Add(1)
return h.LayerRPCHandler.AdmitUnprofiled(b, limits)
}
func TestLayerRPCAdmissionCapacityRejectsBeforeDecoder(t *testing.T) {
for _, test := range []struct {
name string
fillGlobal bool
}{
{name: "connection_queue"},
{name: "global_task", fillGlobal: true},
} {
t.Run(test.name, func(t *testing.T) {
router := rpc.New(rpc.Config{DC: 2}, rpc.Deps{}, zaptest.NewLogger(t), clock.System)
counting := &countingLayerRPCAdmission{LayerRPCHandler: router}
s := New(Options{DC: 2, LayerRPC: counting})
scheduler := newInboundRPCScheduler(1, 1, 1<<30)
s.rpcScheduler = scheduler
target := &Conn{authKeyID: [8]byte{8, 1}, sessionID: 81, metrics: NopMetrics{}}
target.startInboundRPCScheduler(scheduler, 1, 1, time.Second)
holder := target
if test.fillGlobal {
holder = &Conn{authKeyID: [8]byte{8, 2}, sessionID: 82, metrics: NopMetrics{}}
holder.startInboundRPCScheduler(scheduler, 1, 1, time.Second)
}
occupied, err := holder.reserveInboundRPCBatch(context.Background(), []inboundRPCSpec{{method: "occupied", size: 1}})
if err != nil {
t.Fatal(err)
}
defer occupied.abort()
body := exactLayerRPCBody(t, &tg.InvokeWithLayerRequest{Layer: 225, Query: &tg.HelpGetConfigRequest{}})
plan := &inboundPlan{items: []inboundItem{{kind: inboundItemRPC, msgID: 100, body: body}}}
defer plan.close()
if err := s.prepareInboundLayerRPCBatch(context.Background(), target, plan); err != nil {
t.Fatal(err)
}
if got := counting.decodeCalls.Load(); got != 0 {
t.Fatalf("typed decoder entered %d times after capacity rejection", got)
}
if plan.items[0].kind != inboundItemCapacityError || plan.rpcReservation != nil || len(plan.rpcTasks) != 0 {
t.Fatalf("rejected plan = kind:%d reservation:%v tasks:%d", plan.items[0].kind, plan.rpcReservation != nil, len(plan.rpcTasks))
}
})
}
}
func TestLayerRPCAdmissionTransfersOriginalReservationToFreshOwner(t *testing.T) {
router := rpc.New(rpc.Config{DC: 2}, rpc.Deps{}, zaptest.NewLogger(t), clock.System)
s := New(Options{DC: 2, LayerRPC: router})
c := &Conn{authKeyID: [8]byte{8, 3}, sessionID: 83, metrics: NopMetrics{}}
c.startInboundRPCScheduler(s.rpcScheduler, 1, 4, time.Second)
if err := c.FreezeLayerProfile(tg.LayerProfile225); err != nil {
t.Fatal(err)
}
bad := make([]byte, bin.Word)
bad[0], bad[1], bad[2], bad[3] = 0x04, 0x03, 0x02, 0x01
fresh := exactOutboundLayerRPCBody(t, tg.LayerProfile225, &tg.HelpGetConfigRequest{})
plan := &inboundPlan{items: []inboundItem{
{kind: inboundItemRPC, msgID: 100, body: bad},
{kind: inboundItemRPC, msgID: 104, body: fresh},
}}
defer plan.close()
if err := s.prepareInboundLayerRPCBatch(context.Background(), c, plan); err != nil {
t.Fatal(err)
}
if plan.items[0].kind != inboundItemRPCAdmissionError || len(plan.rpcTasks) != 1 || plan.rpcReservation == nil {
t.Fatalf("classified plan = bad:%d tasks:%d reservation:%v", plan.items[0].kind, len(plan.rpcTasks), plan.rpcReservation != nil)
}
wantCharge := int64(layerRPCAdmissionReservationSize(len(fresh)))
if got := c.inflightRPCBytes.Load(); got != wantCharge {
t.Fatalf("connection retained bytes = %d, want %d", got, wantCharge)
}
s.rpcScheduler.budgetMu.Lock()
globalTasks, globalBytes := s.rpcScheduler.tasks, s.rpcScheduler.bytes
s.rpcScheduler.budgetMu.Unlock()
if globalTasks != 1 || globalBytes != wantCharge {
t.Fatalf("global retained budget = %d/%d, want 1/%d", globalTasks, globalBytes, wantCharge)
}
plan.close()
if got := c.inflightRPCBytes.Load(); got != 0 {
t.Fatalf("plan abort leaked %d connection bytes", got)
}
s.rpcScheduler.budgetMu.Lock()
globalTasks, globalBytes = s.rpcScheduler.tasks, s.rpcScheduler.bytes
s.rpcScheduler.budgetMu.Unlock()
if globalTasks != 0 || globalBytes != 0 {
t.Fatalf("plan abort leaked global budget %d/%d", globalTasks, globalBytes)
}
}
func TestLayerRPCAdmissionPendingReplayReleasesProvisionalEntry(t *testing.T) {
router := rpc.New(rpc.Config{DC: 2}, rpc.Deps{}, zaptest.NewLogger(t), clock.System)
s := New(Options{DC: 2, LayerRPC: router})
c := &Conn{authKeyID: [8]byte{8, 4}, sessionID: 84, metrics: NopMetrics{}}
c.startInboundRPCScheduler(s.rpcScheduler, 1, 4, time.Second)
if err := c.FreezeLayerProfile(tg.LayerProfile225); err != nil {
t.Fatal(err)
}
pendingBody := exactOutboundLayerRPCBody(t, tg.LayerProfile225, &tg.HelpGetConfigRequest{})
identityBuffer := &bin.Buffer{Buf: append([]byte(nil), pendingBody...)}
pendingRequest, err := router.AdmitLayer(tg.LayerProfile225, identityBuffer, tg.LayerDecodeLimits{})
if err != nil {
t.Fatal(err)
}
pending, err := s.rpcResults.AcquireIdentified(c.authKeyID, c.sessionID, 100, pendingRequest.Prepared().Identity())
if err != nil || pending.owner == nil {
t.Fatalf("pending owner = %v, %v", pending.owner, err)
}
freshBody := exactOutboundLayerRPCBody(t, tg.LayerProfile225, &tg.HelpGetNearestDCRequest{})
plan := &inboundPlan{items: []inboundItem{
{kind: inboundItemRPC, msgID: 100, body: pendingBody},
{kind: inboundItemRPC, msgID: 104, body: freshBody},
}}
defer plan.close()
if err := s.prepareInboundLayerRPCBatch(context.Background(), c, plan); err != nil {
t.Fatal(err)
}
if plan.items[0].kind != inboundItemRewrappedRPC || len(plan.rewrapAliases) != 1 || len(plan.rpcTasks) != 1 {
t.Fatalf("pending/fresh classification = kind:%d aliases:%d tasks:%d", plan.items[0].kind, len(plan.rewrapAliases), len(plan.rpcTasks))
}
wantCharge := int64(layerRPCAdmissionReservationSize(len(freshBody)))
if got := c.inflightRPCBytes.Load(); got != wantCharge {
t.Fatalf("pending replay retained bytes = %d, want only fresh %d", got, wantCharge)
}
plan.close()
if !pending.owner.Abort() {
t.Fatal("plan cleanup aborted the pre-existing pending replay owner")
}
}
func TestLayerRPCAdmissionCompletedReplayReleasesWholeProvisionalBatch(t *testing.T) {
router := rpc.New(rpc.Config{DC: 2}, rpc.Deps{}, zaptest.NewLogger(t), clock.System)
s := New(Options{DC: 2, LayerRPC: router})
c := &Conn{authKeyID: [8]byte{8, 8}, sessionID: 88, metrics: NopMetrics{}}
c.startInboundRPCScheduler(s.rpcScheduler, 1, 2, time.Second)
if err := c.FreezeLayerProfile(tg.LayerProfile225); err != nil {
t.Fatal(err)
}
body := exactOutboundLayerRPCBody(t, tg.LayerProfile225, &tg.HelpGetConfigRequest{})
identityBuffer := &bin.Buffer{Buf: append([]byte(nil), body...)}
request, err := router.AdmitLayer(tg.LayerProfile225, identityBuffer, tg.LayerDecodeLimits{})
if err != nil {
t.Fatal(err)
}
claim, err := s.rpcResults.AcquireIdentified(c.authKeyID, c.sessionID, 100, request.Prepared().Identity())
if err != nil || claim.owner == nil {
t.Fatalf("completed replay owner = %v, %v", claim.owner, err)
}
if !claim.owner.CompleteExecution(true) {
t.Fatal("complete replay business outcome failed")
}
s.rpcResults.Put(c.authKeyID, c.sessionID, 100, &encodedOutboundMessage{body: []byte{1, 2, 3, 4}})
plan := &inboundPlan{items: []inboundItem{{kind: inboundItemRPC, msgID: 100, body: body}}}
defer plan.close()
if err := s.prepareInboundLayerRPCBatch(context.Background(), c, plan); err != nil {
t.Fatal(err)
}
if plan.items[0].kind != inboundItemReplayRPC || plan.rpcReservation != nil || len(plan.rpcTasks) != 0 {
t.Fatalf("completed replay plan = kind:%d reservation:%v tasks:%d", plan.items[0].kind, plan.rpcReservation != nil, len(plan.rpcTasks))
}
if got := c.inflightRPCBytes.Load(); got != 0 || c.rpcReserved != 0 {
t.Fatalf("completed replay leaked connection budget bytes:%d tasks:%d", got, c.rpcReserved)
}
s.rpcScheduler.budgetMu.Lock()
globalTasks, globalBytes := s.rpcScheduler.tasks, s.rpcScheduler.bytes
s.rpcScheduler.budgetMu.Unlock()
if globalTasks != 0 || globalBytes != 0 {
t.Fatalf("completed replay leaked global budget %d/%d", globalTasks, globalBytes)
}
}
func TestLayerRPCAdmissionReplayPreparationErrorIsNotSilentlyDelivered(t *testing.T) {
router := rpc.New(rpc.Config{DC: 2}, rpc.Deps{}, zaptest.NewLogger(t), clock.System)
prepareErr := errors.New("invalid replay wrapper metadata")
s := New(Options{DC: 2, LayerRPC: &failingReplayLayerRPC{
LayerRPCHandler: router,
err: prepareErr,
}})
c := &Conn{authKeyID: [8]byte{8, 9}, sessionID: 89, metrics: NopMetrics{}}
c.startInboundRPCScheduler(s.rpcScheduler, 1, 2, time.Second)
if err := c.FreezeLayerProfile(tg.LayerProfile225); err != nil {
t.Fatal(err)
}
body := exactOutboundLayerRPCBody(t, tg.LayerProfile225, &tg.HelpGetConfigRequest{})
identityBuffer := &bin.Buffer{Buf: append([]byte(nil), body...)}
request, err := router.AdmitLayer(tg.LayerProfile225, identityBuffer, tg.LayerDecodeLimits{})
if err != nil {
t.Fatal(err)
}
claim, err := s.rpcResults.AcquireIdentified(c.authKeyID, c.sessionID, 100, request.Prepared().Identity())
if err != nil || claim.owner == nil {
t.Fatalf("completed replay owner = %v, %v", claim.owner, err)
}
if !claim.owner.CompleteExecution(true) {
t.Fatal("complete replay business outcome failed")
}
s.rpcResults.Put(c.authKeyID, c.sessionID, 100, &encodedOutboundMessage{body: []byte{1, 2, 3, 4}})
plan := &inboundPlan{items: []inboundItem{{kind: inboundItemRPC, msgID: 100, body: body}}}
if err := s.prepareInboundLayerRPCBatch(context.Background(), c, plan); !errors.Is(err, prepareErr) {
plan.close()
t.Fatalf("replay preparation error = %v, want %v", err, prepareErr)
}
if plan.items[0].kind == inboundItemReplayRPC {
plan.close()
t.Fatal("invalid replay metadata was converted into a deliverable cached result")
}
plan.close()
if got := c.inflightRPCBytes.Load(); got != 0 || c.rpcReserved != 0 {
t.Fatalf("failed replay preparation leaked connection budget bytes:%d tasks:%d", got, c.rpcReserved)
}
s.rpcScheduler.budgetMu.Lock()
globalTasks, globalBytes := s.rpcScheduler.tasks, s.rpcScheduler.bytes
s.rpcScheduler.budgetMu.Unlock()
if globalTasks != 0 || globalBytes != 0 {
t.Fatalf("failed replay preparation leaked global budget %d/%d", globalTasks, globalBytes)
}
}
func TestLayerRPCAdmissionTransferredBatchClosesWithoutLeak(t *testing.T) {
router := rpc.New(rpc.Config{DC: 2}, rpc.Deps{}, zaptest.NewLogger(t), clock.System)
s := New(Options{DC: 2, LayerRPC: router})
c := &Conn{authKeyID: [8]byte{8, 5}, sessionID: 85, metrics: NopMetrics{}}
c.startInboundRPCScheduler(s.rpcScheduler, 1, 2, time.Second)
if err := c.FreezeLayerProfile(tg.LayerProfile225); err != nil {
t.Fatal(err)
}
plan := &inboundPlan{items: []inboundItem{{
kind: inboundItemRPC, msgID: 100,
body: exactOutboundLayerRPCBody(t, tg.LayerProfile225, &tg.HelpGetConfigRequest{}),
}}}
defer plan.close()
if err := s.prepareInboundLayerRPCBatch(context.Background(), c, plan); err != nil {
t.Fatal(err)
}
c.beginCloseInboundRPCScheduler()
if err := plan.commitRPCBatch(); err != ErrConnClosed {
t.Fatalf("commit after connection close = %v, want ErrConnClosed", err)
}
plan.close()
if !c.waitInboundShutdown(time.Second) {
t.Fatal("connection close did not converge after transferred reservation failed commit")
}
if got := c.inflightRPCBytes.Load(); got != 0 {
t.Fatalf("connection close leaked %d admission bytes", got)
}
s.rpcScheduler.budgetMu.Lock()
globalTasks, globalBytes := s.rpcScheduler.tasks, s.rpcScheduler.bytes
s.rpcScheduler.budgetMu.Unlock()
if globalTasks != 0 || globalBytes != 0 {
t.Fatalf("connection close leaked global budget %d/%d", globalTasks, globalBytes)
}
}
func TestLayerRPCAdmissionTransferredBatchCommitsConservativeCharge(t *testing.T) {
router := rpc.New(rpc.Config{DC: 2}, rpc.Deps{}, zaptest.NewLogger(t), clock.System)
s := New(Options{DC: 2, LayerRPC: router})
c := &Conn{authKeyID: [8]byte{8, 6}, sessionID: 86, metrics: NopMetrics{}}
c.startInboundRPCScheduler(s.rpcScheduler, 1, 2, time.Second)
if err := c.FreezeLayerProfile(tg.LayerProfile225); err != nil {
t.Fatal(err)
}
body := exactOutboundLayerRPCBody(t, tg.LayerProfile225, &tg.HelpGetConfigRequest{})
plan := &inboundPlan{items: []inboundItem{{kind: inboundItemRPC, msgID: 100, body: body}}}
defer plan.close()
if err := s.prepareInboundLayerRPCBatch(context.Background(), c, plan); err != nil {
t.Fatal(err)
}
if err := plan.commitRPCBatch(); err != nil {
t.Fatal(err)
}
wantCharge := layerRPCAdmissionReservationSize(len(body))
c.rpcMu.Lock()
queued := len(c.rpcQueue)
gotCharge := 0
if queued == 1 {
gotCharge = c.rpcQueue[0].size
}
c.rpcMu.Unlock()
if queued != 1 || gotCharge != wantCharge {
t.Fatalf("committed queue = len:%d charge:%d, want 1/%d", queued, gotCharge, wantCharge)
}
c.beginCloseInboundRPCScheduler()
if got := c.inflightRPCBytes.Load(); got != 0 {
t.Fatalf("queued exact task close leaked %d bytes", got)
}
s.rpcScheduler.budgetMu.Lock()
globalTasks, globalBytes := s.rpcScheduler.tasks, s.rpcScheduler.bytes
s.rpcScheduler.budgetMu.Unlock()
if globalTasks != 0 || globalBytes != 0 {
t.Fatalf("queued exact task close leaked global budget %d/%d", globalTasks, globalBytes)
}
}
func TestLayerRPCAdmissionLocalDuplicateConsumesNoProvisionalEntry(t *testing.T) {
router := rpc.New(rpc.Config{DC: 2}, rpc.Deps{}, zaptest.NewLogger(t), clock.System)
s := New(Options{DC: 2, LayerRPC: router})
c := &Conn{authKeyID: [8]byte{8, 7}, sessionID: 87, metrics: NopMetrics{}}
c.startInboundRPCScheduler(s.rpcScheduler, 1, 2, time.Second)
if err := c.FreezeLayerProfile(tg.LayerProfile225); err != nil {
t.Fatal(err)
}
body := exactOutboundLayerRPCBody(t, tg.LayerProfile225, &tg.HelpGetConfigRequest{})
plan := &inboundPlan{items: []inboundItem{
{kind: inboundItemDuplicate, msgID: 96, body: body},
{kind: inboundItemRPC, msgID: 100, body: body},
}}
defer plan.close()
if err := s.prepareInboundLayerRPCBatch(context.Background(), c, plan); err != nil {
t.Fatal(err)
}
if plan.items[0].kind != inboundItemDuplicate || len(plan.rpcTasks) != 1 || c.rpcReserved != 1 {
t.Fatalf("duplicate/fresh admission = duplicate:%d tasks:%d reserved:%d", plan.items[0].kind, len(plan.rpcTasks), c.rpcReserved)
}
}
func TestTakeInboundRPCFIFOAdvancesSliceHead(t *testing.T) {
scheduler := newInboundRPCScheduler(1, 8, 1<<20)
c := &Conn{metrics: NopMetrics{}}
c.startInboundRPCScheduler(scheduler, 1, 8, time.Second)
c.rpcQueue = []inboundRPC{{method: "first"}, {method: "second"}, {method: "third"}}
c.rpcReady = true
oldSecond := &c.rpcQueue[1]
task, ok, _ := c.takeInboundRPC()
if !ok || task.method != "first" {
t.Fatalf("take = (%q,%v), want first", task.method, ok)
}
if len(c.rpcQueue) != 2 || &c.rpcQueue[0] != oldSecond {
t.Fatal("FIFO take copied the queue instead of advancing its slice head")
}
c.finishInboundRPC(task)
c.beginCloseInboundRPCScheduler()
}
func BenchmarkLayerRPCAdmissionReservationSize(b *testing.B) {
for _, size := range []int{64, appfiles.MaxUploadPartBytes + 24} {
b.Run(time.Duration(size).String(), func(b *testing.B) {
b.ReportAllocs()
var charge int
for i := 0; i < b.N; i++ {
charge = layerRPCAdmissionReservationSize(size)
}
if charge == 0 {
b.Fatal("zero admission charge")
}
})
}
}

366
internal/mtprotoedge/layer_downgrade_test.go
View file

@ -2,78 +2,82 @@ package mtprotoedge
import (
"bytes"
"context"
"encoding/binary"
"errors"
"fmt"
"sync/atomic"
"testing"
"time"
"github.com/gotd/td/bin"
"github.com/gotd/td/proto"
"github.com/gotd/td/tg"
"github.com/iamxvbaba/td/bin"
"github.com/iamxvbaba/td/proto"
"github.com/iamxvbaba/td/tg"
"go.uber.org/zap/zaptest"
)
// TestConnDowngradedClone verifies the outbound seam downgrades a canonical
// (227) object to the connection's negotiated layer, is a no-op for 227, and
// — critically for push fan-out — never mutates the shared input message (one
// pre-encoded update is reused across many connections of differing layers).
func TestConnDowngradedClone(t *testing.T) {
const (
message227CRC = 0x7600b9d3
message220CRC = 0xb92f76cf
)
msg := &tg.Message{
ID: 2,
FromID: &tg.PeerUser{UserID: 3},
PeerID: &tg.PeerUser{UserID: 3},
Date: 1,
Message: "hi",
}
type countingLayerRPCResult struct {
inner tg.LayerRPCResult
encodeCalls atomic.Int32
prepareCalls atomic.Int32
}
// layer 220: returns a NEW message rewritten to the 220 constructor id,
// leaving the shared input untouched (227).
enc, err := encodeOutboundMessage(msg)
if err != nil {
t.Fatalf("encode: %v", err)
}
c := &Conn{metrics: NopMetrics{}}
c.SetClientLayer(220)
out := c.downgradedClone(enc)
const (
testChannelWireID227 uint32 = 0x1c32b11c
testChannelWireID228 uint32 = 0xd49f34c6
)
if id, _ := (&bin.Buffer{Buf: out.body}).PeekID(); id != message220CRC {
t.Fatalf("downgraded id = %#08x, want %#08x", id, message220CRC)
}
if out.typeID != message220CRC {
t.Fatalf("downgraded typeID = %#08x, want %#08x", out.typeID, message220CRC)
}
// Input must be unmodified — this is what makes shared push fan-out safe.
if id, _ := (&bin.Buffer{Buf: enc.body}).PeekID(); id != message227CRC {
t.Fatalf("input message was mutated: id now %#08x, want 227 %#08x", id, message227CRC)
func testChannelWireID(profile tg.LayerProfile) uint32 {
if profile == tg.LayerProfile228 {
return testChannelWireID228
}
return testChannelWireID227
}
// Two connections sharing one pre-encoded message get independent results.
encShared, _ := encodeOutboundMessage(msg)
c220 := &Conn{metrics: NopMetrics{}}
c220.SetClientLayer(220)
c227 := &Conn{metrics: NopMetrics{}} // ClientLayer() defaults to 227
out220 := c220.downgradedClone(encShared)
out227 := c227.downgradedClone(encShared)
if id, _ := (&bin.Buffer{Buf: out220.body}).PeekID(); id != message220CRC {
t.Fatalf("shared->220 id = %#08x, want %#08x", id, message220CRC)
func testOtherChannelWireID(profile tg.LayerProfile) uint32 {
if profile == tg.LayerProfile228 {
return testChannelWireID227
}
if out227 != encShared {
t.Errorf("227 connection should pass the shared message through unchanged (same pointer)")
}
if !bytes.Equal(encShared.body, out227.body) {
t.Errorf("227 passthrough altered bytes")
return testChannelWireID228
}
func testLayerChannel() *tg.Channel {
return &tg.Channel{
ID: 100,
Title: "layer proof",
Photo: &tg.ChatPhotoEmpty{},
Date: 1,
}
}
func TestEncodeRPCResultDowngradesDifferenceMessagesForNegotiatedLayer225(t *testing.T) {
const (
message227CRC = 0x7600b9d3
message225CRC = 0x95ef6f2b
)
c := &Conn{metrics: NopMetrics{}}
c.SetClientLayer(225)
func (r *countingLayerRPCResult) Encode(b *bin.Buffer) error {
r.encodeCalls.Add(1)
return r.inner.Encode(b)
}
func (r *countingLayerRPCResult) Prepared() tg.LayerPreparedCall { return r.inner.Prepared() }
func (r *countingLayerRPCResult) WireInvariant() bool { return r.inner.WireInvariant() }
func (r *countingLayerRPCResult) Freeze() (tg.LayerFrozenResult, error) {
return r.inner.Freeze()
}
func (r *countingLayerRPCResult) Prepare() (tg.LayerPreparedResult, error) {
r.prepareCalls.Add(1)
return r.inner.Prepare()
}
func TestExactLayerRPCResultEncodesDifferenceWithAdmittedCodec(t *testing.T) {
for _, profile := range []tg.LayerProfile{tg.LayerProfile225, tg.LayerProfile227, tg.LayerProfile228} {
t.Run(fmt.Sprintf("layer_%d", profile), func(t *testing.T) {
testExactLayerRPCResultEncodesDifferenceWithAdmittedCodec(t, profile)
})
}
}
func testExactLayerRPCResultEncodesDifferenceWithAdmittedCodec(t *testing.T, profile tg.LayerProfile) {
t.Helper()
diff := &tg.UpdatesDifference{
NewMessages: []tg.MessageClass{
&tg.Message{
@ -86,116 +90,190 @@ func TestEncodeRPCResultDowngradesDifferenceMessagesForNegotiatedLayer225(t *tes
},
NewEncryptedMessages: []tg.EncryptedMessageClass{},
OtherUpdates: []tg.UpdateClass{},
Chats: []tg.ChatClass{},
Chats: []tg.ChatClass{testLayerChannel()},
Users: []tg.UserClass{},
State: tg.UpdatesState{Pts: 2, Date: 1},
}
dispatcher := tg.NewServerDispatcher(nil)
dispatcher.OnUpdatesGetDifference(func(context.Context, *tg.UpdatesGetDifferenceRequest) (tg.UpdatesDifferenceClass, error) {
return diff, nil
})
outbound, err := tg.PrepareLayerOutboundCall(profile, &tg.UpdatesGetDifferenceRequest{Pts: 1, Date: 1})
if err != nil {
t.Fatal(err)
}
var requestBody bin.Buffer
if err := outbound.Encode(&requestBody); err != nil {
t.Fatal(err)
}
admitted, err := dispatcher.AdmitLayer(profile, &requestBody)
if err != nil {
t.Fatal(err)
}
serverResult, err := dispatcher.DispatchAdmitted(context.Background(), admitted)
if err != nil {
t.Fatal(err)
}
counted := &countingLayerRPCResult{inner: serverResult}
exact := &layerRPCResultEncoder{call: counted.Prepared().Call(), result: counted}
c := &Conn{metrics: NopMetrics{}, msgID: proto.NewMessageIDGen(time.Now)}
if err := c.FreezeLayerProfile(profile); err != nil {
t.Fatal(err)
}
// Simulate an invokeWithLayer correction admitted while this handler was
// still running. The result must retain the request's admitted profile.
corrected := tg.LayerProfile227
if profile == tg.LayerProfile227 {
corrected = tg.LayerProfile225
}
if err := c.FreezeLayerProfile(corrected); err != nil {
t.Fatal(err)
}
s := &Server{log: zaptest.NewLogger(t)}
encoded, err := s.encodeRPCResult(c, 12345, diff)
encoded, err := s.encodeRPCResult(c, 12345, exact)
if err != nil {
t.Fatalf("encode rpc_result: %v", err)
}
var result proto.Result
if err := result.Decode(&bin.Buffer{Buf: encoded.body}); err != nil {
if got := counted.prepareCalls.Load(); got != 0 {
t.Fatalf("generated Prepare calls = %d, want 0; inbound workers must not snapshot result bytes", got)
}
if got := counted.encodeCalls.Load(); got != 1 {
t.Fatalf("generated Encode calls = %d, want exactly 1 under outbound admission", got)
}
if encoded.layer == nil || encoded.layer.profile != profile || encoded.layer.typ != admitted.Call().WireResultType() {
t.Fatalf("result binding = %#v, want profile %d and admitted result TypeRef", encoded.layer, profile)
}
if encoded.layer.kind != outboundLayerBindingRequest {
t.Fatalf("exact RPC result binding kind = %d, want request-bound", encoded.layer.kind)
}
beforeFrame := append([]byte(nil), encoded.body...)
frame, err := c.buildFrame(context.Background(), proto.MessageServerResponse, nil, encoded)
if err != nil {
t.Fatal(err)
}
if !bytes.Equal(frame.body, beforeFrame) {
t.Fatal("exact RPC result was transcoded after generated preparation")
}
var rpcEnvelope proto.Result
if err := rpcEnvelope.Decode(&bin.Buffer{Buf: frame.body}); err != nil {
t.Fatalf("decode rpc_result: %v", err)
}
if result.RequestMessageID != 12345 {
t.Fatalf("req_msg_id = %d, want 12345", result.RequestMessageID)
if rpcEnvelope.RequestMessageID != 12345 {
t.Fatalf("req_msg_id = %d, want 12345", rpcEnvelope.RequestMessageID)
}
if !bytes.Contains(result.Result, littleEndianID(message225CRC)) {
t.Fatalf("rpc_result inner object does not contain layer 225 message id %#08x", message225CRC)
wantChannelID := testChannelWireID(profile)
if !bytes.Contains(rpcEnvelope.Result, littleEndianID(wantChannelID)) {
t.Fatalf("profile %d offline difference lacks channel constructor %#08x", profile, wantChannelID)
}
if bytes.Contains(result.Result, littleEndianID(message227CRC)) {
t.Fatalf("rpc_result inner object still contains canonical message id %#08x", message227CRC)
if otherChannelID := testOtherChannelWireID(profile); bytes.Contains(rpcEnvelope.Result, littleEndianID(otherChannelID)) {
t.Fatalf("profile %d offline difference leaked channel constructor %#08x", profile, otherChannelID)
}
inner := bin.Buffer{Buf: rpcEnvelope.Result}
decoded, err := tg.DecodeLayer(profile, tg.LayerClassUpdatesDifferenceType(), &inner)
if err != nil {
t.Fatalf("decode exact difference: %v", err)
}
if inner.Len() != 0 {
t.Fatalf("exact difference left %d bytes", inner.Len())
}
got, ok := decoded.(*tg.UpdatesDifference)
message, messageOK := func() (*tg.Message, bool) {
if !ok || len(got.NewMessages) != 1 {
return nil, false
}
value, valueOK := got.NewMessages[0].(*tg.Message)
return value, valueOK
}()
if !messageOK || message.ID != 2 {
t.Fatalf("decoded exact difference = %#v", decoded)
}
if len(got.Chats) != 1 {
t.Fatalf("decoded exact difference chats = %#v", got.Chats)
}
channel, channelOK := got.Chats[0].(*tg.Channel)
if !channelOK || channel.ID != 100 {
t.Fatalf("decoded exact difference channel = %#v", got.Chats)
}
}
func TestEncodeRPCResultDowngradesDialogMessagesForNegotiatedLayer225(t *testing.T) {
const (
message227CRC = 0x7600b9d3
message225CRC = 0x95ef6f2b
)
c := &Conn{metrics: NopMetrics{}}
c.SetClientLayer(225)
dialogs := &tg.MessagesDialogs{
Dialogs: []tg.DialogClass{
&tg.Dialog{
Peer: &tg.PeerUser{UserID: 3},
TopMessage: 2,
NotifySettings: tg.PeerNotifySettings{},
},
},
Messages: []tg.MessageClass{
&tg.Message{
ID: 2,
FromID: &tg.PeerUser{UserID: 3},
PeerID: &tg.PeerUser{UserID: 3},
Date: 1,
Message: "hi",
},
},
Chats: []tg.ChatClass{},
Users: []tg.UserClass{
&tg.User{ID: 3, AccessHash: 5, FirstName: "A"},
},
}
s := &Server{log: zaptest.NewLogger(t)}
encoded, err := s.encodeRPCResult(c, 12345, dialogs)
func TestExactLayerRPCResultUsesHistoricalMethodResultType(t *testing.T) {
const profile = tg.LayerProfile225
dispatcher := tg.NewServerDispatcher(nil)
dispatcher.OnChannelsJoinChannel(func(context.Context, tg.InputChannelClass) (tg.MessagesChatInviteJoinResultClass, error) {
return &tg.MessagesChatInviteJoinResultOk{Updates: &tg.UpdatesTooLong{}}, nil
})
outbound, err := tg.PrepareLayerOutboundCall(profile, &tg.ChannelsJoinChannelRequest{Channel: &tg.InputChannelEmpty{}})
if err != nil {
t.Fatalf("encode rpc_result: %v", err)
t.Fatal(err)
}
var result proto.Result
if err := result.Decode(&bin.Buffer{Buf: encoded.body}); err != nil {
t.Fatalf("decode rpc_result: %v", err)
var requestBody bin.Buffer
if err := outbound.Encode(&requestBody); err != nil {
t.Fatal(err)
}
if !bytes.Contains(result.Result, littleEndianID(message225CRC)) {
t.Fatalf("rpc_result inner object does not contain layer 225 message id %#08x", message225CRC)
admitted, err := dispatcher.AdmitLayer(profile, &requestBody)
if err != nil {
t.Fatal(err)
}
if bytes.Contains(result.Result, littleEndianID(message227CRC)) {
t.Fatalf("rpc_result inner object still contains canonical message id %#08x", message227CRC)
if admitted.Call().WireID() == tg.ChannelsJoinChannelRequestTypeID {
t.Fatal("historical request unexpectedly retained canonical method id")
}
serverResult, err := dispatcher.DispatchAdmitted(context.Background(), admitted)
if err != nil {
t.Fatal(err)
}
exact := &layerRPCResultEncoder{call: serverResult.Prepared().Call(), result: serverResult}
c := &Conn{metrics: NopMetrics{}}
if err := c.FreezeLayerProfile(profile); err != nil {
t.Fatal(err)
}
encoded, err := (&Server{log: zaptest.NewLogger(t)}).encodeRPCResult(c, 67890, exact)
if err != nil {
t.Fatal(err)
}
var rpcEnvelope proto.Result
if err := rpcEnvelope.Decode(&bin.Buffer{Buf: encoded.body}); err != nil {
t.Fatal(err)
}
inner := bin.Buffer{Buf: rpcEnvelope.Result}
updates, err := tg.DecodeLayer(profile, tg.LayerClassUpdatesType(), &inner)
if err != nil {
t.Fatalf("decode historical channels.joinChannel result: %v", err)
}
if inner.Len() != 0 {
t.Fatalf("historical result left %d bytes", inner.Len())
}
if _, ok := updates.(*tg.UpdatesTooLong); !ok {
t.Fatalf("historical result = %T, want Updates", updates)
}
}
func TestConnDowngradedCloneDowngradesUpdateNewMessageForLayer225(t *testing.T) {
const (
message227CRC = 0x7600b9d3
message225CRC = 0x95ef6f2b
)
updates := &tg.Updates{
Updates: []tg.UpdateClass{
&tg.UpdateNewMessage{
Message: &tg.Message{
ID: 2,
FromID: &tg.PeerUser{UserID: 3},
PeerID: &tg.PeerUser{UserID: 3},
Date: 1,
Message: "hi",
},
Pts: 2,
PtsCount: 1,
},
},
Users: []tg.UserClass{
&tg.User{ID: 3, AccessHash: 5, FirstName: "A"},
},
Chats: []tg.ChatClass{},
Date: 1,
Seq: 1,
}
enc, err := encodeOutboundMessage(updates)
if err != nil {
t.Fatalf("encode updates: %v", err)
}
func TestProductionUnboundApplicationResultFailsClosedForLayer227(t *testing.T) {
c := &Conn{metrics: NopMetrics{}}
c.SetClientLayer(225)
out := c.downgradedClone(enc)
if !bytes.Contains(out.body, littleEndianID(message225CRC)) {
t.Fatalf("push update does not contain layer 225 message id %#08x", message225CRC)
if err := c.FreezeLayerProfile(tg.LayerProfile227); err != nil {
t.Fatal(err)
}
if bytes.Contains(out.body, littleEndianID(message227CRC)) {
t.Fatalf("push update still contains canonical message id %#08x", message227CRC)
encoded, err := (&Server{log: zaptest.NewLogger(t)}).encodeRPCResult(c, 12345, testLayerChannel())
if !errors.Is(err, ErrOutboundLayerBindingRequired) {
t.Fatalf("unbound Layer 228 result error = %v, want %v", err, ErrOutboundLayerBindingRequired)
}
if encoded != nil {
t.Fatalf("unbound result produced %d wire bytes", len(encoded.body))
}
}
func TestProductionUnboundApplicationPushFailsClosedForLayer227(t *testing.T) {
c := &Conn{metrics: NopMetrics{}}
if err := c.FreezeLayerProfile(tg.LayerProfile227); err != nil {
t.Fatal(err)
}
frame, err := c.buildFrame(context.Background(), proto.MessageFromServer, testLayerChannelUpdatesValue(321), nil)
if !errors.Is(err, ErrOutboundLayerBindingRequired) {
t.Fatalf("unbound Layer 228 push error = %v, want %v", err, ErrOutboundLayerBindingRequired)
}
if frame != nil {
t.Fatalf("unbound push produced frame %#v", frame)
}
}

596
internal/mtprotoedge/layer_profile_test.go Normal file
View file

@ -0,0 +1,596 @@
package mtprotoedge
import (
"context"
"errors"
"fmt"
"sync"
"testing"
"github.com/iamxvbaba/td/tg"
)
type countingInheritedLayerResolver struct {
LayerRPCHandler
calls int
layer int
found bool
err error
}
type orderedSessionLayerResolver struct {
LayerRPCHandler
layer int
msgID int64
found bool
}
func (r *orderedSessionLayerResolver) NegotiatedSessionLayerEvidence([8]byte, int64) (int, int64, bool) {
return r.layer, r.msgID, r.found
}
func (r *countingInheritedLayerResolver) ResolveInheritedAuthKeyLayer(context.Context, [8]byte) (int, bool, error) {
r.calls++
return r.layer, r.found, r.err
}
func TestConnLayerProfileUnknownFreezeAndIdempotence(t *testing.T) {
c := &Conn{}
if profile, ok := c.LayerProfile(); ok || profile != 0 {
t.Fatalf("initial LayerProfile = (%d, %v), want (0, false)", profile, ok)
}
if err := c.FreezeLayerProfile(tg.LayerProfile225); err != nil {
t.Fatalf("freeze layer 225: %v", err)
}
if err := c.FreezeLayerProfile(tg.LayerProfile225); err != nil {
t.Fatalf("repeat freeze layer 225: %v", err)
}
if profile, ok := c.LayerProfile(); !ok || profile != tg.LayerProfile225 {
t.Fatalf("LayerProfile = (%d, %v), want (225, true)", profile, ok)
}
}
func TestConnLayerProfileInheritedCanBeCorrectedExplicitly(t *testing.T) {
c := &Conn{}
if err := c.SeedInheritedLayerProfile(tg.LayerProfile225); err != nil {
t.Fatalf("seed inherited layer 225: %v", err)
}
initial := c.LayerProfileState()
if initial.Profile != tg.LayerProfile225 || initial.Origin != LayerProfileInherited || initial.Epoch != 1 {
t.Fatalf("initial inherited state = %#v", initial)
}
if err := c.SeedInheritedLayerProfile(tg.LayerProfile226); err != nil {
t.Fatalf("repeat inherited seed: %v", err)
}
if got := c.LayerProfileState(); got != initial {
t.Fatalf("second inherited seed replaced selected default: got %#v want %#v", got, initial)
}
if err := c.FreezeLayerProfile(tg.LayerProfile225); err != nil {
t.Fatalf("promote inherited evidence: %v", err)
}
promoted := c.LayerProfileState()
if promoted.Profile != tg.LayerProfile225 || promoted.Origin != LayerProfileExplicit || promoted.Epoch != initial.Epoch+1 {
t.Fatalf("promoted explicit state = %#v", promoted)
}
if err := c.FreezeLayerProfile(tg.LayerProfile227); err != nil {
t.Fatalf("correct explicit layer: %v", err)
}
corrected := c.LayerProfileState()
if corrected.Profile != tg.LayerProfile227 || corrected.Origin != LayerProfileExplicit || corrected.Epoch != promoted.Epoch+1 {
t.Fatalf("corrected explicit state = %#v", corrected)
}
}
func TestConnSeedLayerProfile(t *testing.T) {
c := &Conn{}
if err := c.SeedLayerProfile(tg.LayerProfile226); err != nil {
t.Fatalf("seed layer 226: %v", err)
}
if err := c.SeedLayerProfile(tg.LayerProfile226); err != nil {
t.Fatalf("repeat seed layer 226: %v", err)
}
if err := c.FreezeLayerProfile(tg.LayerProfile226); err != nil {
t.Fatalf("freeze seeded layer 226: %v", err)
}
if profile, ok := c.LayerProfile(); !ok || profile != tg.LayerProfile226 {
t.Fatalf("LayerProfile = (%d, %v), want (226, true)", profile, ok)
}
}
func TestConnLayerProfileRejectsUnsupported(t *testing.T) {
for _, profile := range []tg.LayerProfile{0, 219, 229} {
t.Run(fmt.Sprintf("layer_%d", profile), func(t *testing.T) {
c := &Conn{}
if err := c.FreezeLayerProfile(profile); !errors.Is(err, ErrLayerProfileUnsupported) {
t.Fatalf("FreezeLayerProfile(%d) error = %v, want ErrLayerProfileUnsupported", profile, err)
}
if err := c.SeedLayerProfile(profile); !errors.Is(err, ErrLayerProfileUnsupported) {
t.Fatalf("SeedLayerProfile(%d) error = %v, want ErrLayerProfileUnsupported", profile, err)
}
if got, ok := c.LayerProfile(); ok || got != 0 {
t.Fatalf("LayerProfile after invalid values = (%d, %v), want (0, false)", got, ok)
}
})
}
}
func TestConnLayerProfileConcurrentCorrectionsRemainAtomic(t *testing.T) {
const goroutines = 128
c := &Conn{}
start := make(chan struct{})
errs := make([]error, goroutines)
profiles := make([]tg.LayerProfile, goroutines)
var wg sync.WaitGroup
wg.Add(goroutines)
for i := range goroutines {
profile := tg.LayerProfile225
if i%2 != 0 {
profile = tg.LayerProfile227
}
profiles[i] = profile
go func(index int, requested tg.LayerProfile) {
defer wg.Done()
<-start
errs[index] = c.FreezeLayerProfile(requested)
}(i, profile)
}
close(start)
wg.Wait()
state := c.LayerProfileState()
if state.Origin != LayerProfileExplicit || (state.Profile != tg.LayerProfile225 && state.Profile != tg.LayerProfile227) {
t.Fatalf("concurrent final state = %#v, want supported explicit contender", state)
}
if state.Epoch == 0 || state.Epoch > goroutines {
t.Fatalf("concurrent epoch = %d, want 1..%d", state.Epoch, goroutines)
}
for i, err := range errs {
if err != nil {
t.Fatalf("correction contender %d (%d) returned error: %v", i, profiles[i], err)
}
}
}
func TestConnLayerProfileEvidenceUsesClientMessageOrder(t *testing.T) {
c := &Conn{}
if err := c.seedOrderedLayerProfile(tg.LayerProfile225, 100); err != nil {
t.Fatal(err)
}
if applied, err := c.FreezeLayerProfileAt(tg.LayerProfile227, 104); err != nil || !applied {
t.Fatalf("newer correction applied=%v err=%v", applied, err)
}
corrected := c.LayerProfileState()
if applied, err := c.FreezeLayerProfileAt(tg.LayerProfile225, 100); err != nil || applied {
t.Fatalf("old duplicate applied=%v err=%v", applied, err)
}
if got := c.LayerProfileState(); got != corrected {
t.Fatalf("old duplicate changed profile: got %#v want %#v", got, corrected)
}
if applied, err := c.FreezeLayerProfileAt(tg.LayerProfile225, 104); !errors.Is(err, ErrLayerProfileConflict) || applied {
t.Fatalf("same-msg conflicting evidence applied=%v err=%v", applied, err)
}
if applied, err := c.FreezeLayerProfileAt(tg.LayerProfile227, 108); err != nil || !applied {
t.Fatalf("same-layer newer evidence applied=%v err=%v", applied, err)
}
state, msgID := c.layerProfileEvidenceState()
if state != corrected || msgID != 108 {
t.Fatalf("same-layer cursor advance = state:%#v msgID:%d, want state:%#v msgID:108", state, msgID, corrected)
}
}
func TestSessionManagerSeedsOnlyUnknownRawAuthKeyConnections(t *testing.T) {
m := NewSessionManager(nil)
authKeyID := [8]byte{2, 2, 7}
unknown := &Conn{authKeyID: authKeyID, sessionID: 1}
explicit := &Conn{authKeyID: authKeyID, sessionID: 2}
inherited := &Conn{authKeyID: authKeyID, sessionID: 3}
if err := explicit.FreezeLayerProfile(tg.LayerProfile225); err != nil {
t.Fatal(err)
}
if err := inherited.SeedInheritedLayerProfile(tg.LayerProfile226); err != nil {
t.Fatal(err)
}
for _, c := range []*Conn{unknown, explicit, inherited} {
if err := m.Register(c); err != nil {
t.Fatal(err)
}
}
if seeded := m.SeedInheritedLayerForRawAuthKey(authKeyID, 227); seeded != 1 {
t.Fatalf("seeded connections = %d, want 1", seeded)
}
if got := unknown.LayerProfileState(); got.Profile != tg.LayerProfile227 || got.Origin != LayerProfileInherited {
t.Fatalf("unknown connection seed = %#v", got)
}
if got := explicit.LayerProfileState(); got.Profile != tg.LayerProfile225 || got.Origin != LayerProfileExplicit {
t.Fatalf("explicit connection was overwritten = %#v", got)
}
if got := inherited.LayerProfileState(); got.Profile != tg.LayerProfile226 || got.Origin != LayerProfileInherited {
t.Fatalf("existing inherited connection was overwritten = %#v", got)
}
if seeded := m.SeedInheritedLayerForRawAuthKey(authKeyID, 229); seeded != 0 {
t.Fatalf("unsupported layer seeded %d connections", seeded)
}
}
func TestSessionManagerRefreshesInheritedRawKeyShadowAtBind(t *testing.T) {
m := NewSessionManager(nil)
authKeyID := [8]byte{2, 2, 8}
unknown := &Conn{authKeyID: authKeyID, sessionID: 1}
inherited := &Conn{authKeyID: authKeyID, sessionID: 2}
explicit := &Conn{authKeyID: authKeyID, sessionID: 3}
if err := inherited.SeedInheritedLayerProfile(tg.LayerProfile225); err != nil {
t.Fatal(err)
}
if err := explicit.FreezeLayerProfile(tg.LayerProfile225); err != nil {
t.Fatal(err)
}
for _, c := range []*Conn{unknown, inherited, explicit} {
if err := m.Register(c); err != nil {
t.Fatal(err)
}
}
if refreshed := m.RefreshInheritedLayerForRawAuthKey(authKeyID, 227); refreshed != 2 {
t.Fatalf("refreshed connections = %d, want 2", refreshed)
}
for name, c := range map[string]*Conn{"unknown": unknown, "inherited": inherited} {
if got := c.LayerProfileState(); got.Profile != tg.LayerProfile227 || got.Origin != LayerProfileInherited {
t.Fatalf("%s refresh = %#v", name, got)
}
}
if got := explicit.LayerProfileState(); got.Profile != tg.LayerProfile225 || got.Origin != LayerProfileExplicit {
t.Fatalf("explicit evidence overwritten = %#v", got)
}
}
func TestSessionManagerClearsOnlyInheritedRawKeyShadowAtBind(t *testing.T) {
m := NewSessionManager(nil)
authKeyID := [8]byte{2, 2, 81}
inherited := &Conn{authKeyID: authKeyID, sessionID: 1}
explicit := &Conn{authKeyID: authKeyID, sessionID: 2}
unknown := &Conn{authKeyID: authKeyID, sessionID: 3}
if err := inherited.SeedInheritedLayerProfile(tg.LayerProfile225); err != nil {
t.Fatal(err)
}
if err := explicit.seedOrderedLayerProfile(tg.LayerProfile227, 104); err != nil {
t.Fatal(err)
}
for _, c := range []*Conn{inherited, explicit, unknown} {
if err := m.Register(c); err != nil {
t.Fatal(err)
}
}
if cleared := m.ClearInheritedLayerForRawAuthKey(authKeyID); cleared != 1 {
t.Fatalf("cleared connections = %d, want 1", cleared)
}
if got := inherited.LayerProfileState(); got.Origin != LayerProfileUnknown || got.Profile != 0 {
t.Fatalf("inherited shadow after clear = %#v, want unknown", got)
}
if got := explicit.LayerProfileState(); got.Profile != tg.LayerProfile227 || got.Origin != LayerProfileExplicit {
t.Fatalf("explicit evidence was cleared = %#v", got)
}
if state, msgID := explicit.layerProfileEvidenceState(); state.Origin != LayerProfileExplicit || msgID != 104 {
t.Fatalf("explicit ordered evidence changed = %#v msgID:%d", state, msgID)
}
if got := unknown.LayerProfileState(); got.Origin != LayerProfileUnknown || got.Profile != 0 {
t.Fatalf("unknown state changed = %#v", got)
}
if cleared := m.ClearInheritedLayerForRawAuthKey(authKeyID); cleared != 0 {
t.Fatalf("idempotent clear changed %d connections", cleared)
}
}
func TestSessionManagerSeedsUnknownSessionsAcrossBusinessAuthKey(t *testing.T) {
m := NewSessionManager(nil)
permAuthKeyID := [8]byte{9, 9, 9}
rawOne := [8]byte{9, 9, 1}
rawTwo := [8]byte{9, 9, 2}
first := &Conn{authKeyID: rawOne, sessionID: 1}
second := &Conn{authKeyID: rawTwo, sessionID: 2}
explicit := &Conn{authKeyID: rawTwo, sessionID: 3}
inherited := &Conn{authKeyID: rawOne, sessionID: 4}
if err := explicit.FreezeLayerProfile(tg.LayerProfile225); err != nil {
t.Fatal(err)
}
if err := inherited.SeedInheritedLayerProfile(tg.LayerProfile225); err != nil {
t.Fatal(err)
}
for _, c := range []*Conn{first, second, explicit, inherited} {
if err := m.Register(c); err != nil {
t.Fatal(err)
}
m.BindAuthKeyForSession(c.authKeyID, c.sessionID, permAuthKeyID)
}
if seeded := m.SeedInheritedLayerForBusinessAuthKey(permAuthKeyID, 227); seeded != 2 {
t.Fatalf("business auth-key seeded=%d, want 2", seeded)
}
for name, c := range map[string]*Conn{"first": first, "second": second} {
if got := c.LayerProfileState(); got.Profile != tg.LayerProfile227 || got.Origin != LayerProfileInherited {
t.Fatalf("%s business default = %#v", name, got)
}
}
if got := explicit.LayerProfileState(); got.Profile != tg.LayerProfile225 || got.Origin != LayerProfileExplicit {
t.Fatalf("business seed overwrote explicit = %#v", got)
}
if got := inherited.LayerProfileState(); got.Profile != tg.LayerProfile225 || got.Origin != LayerProfileInherited {
t.Fatalf("business seed overwrote inherited = %#v", got)
}
}
func TestSessionManagerExplicitLayerEvidenceUsesLiveExactSession(t *testing.T) {
m := NewSessionManager(nil)
authKeyID := [8]byte{2, 2, 9}
const sessionID = int64(229)
c := &Conn{authKeyID: authKeyID, sessionID: sessionID}
if err := c.seedOrderedLayerProfile(tg.LayerProfile226, 1234); err != nil {
t.Fatal(err)
}
if err := m.Register(c); err != nil {
t.Fatal(err)
}
if layer, msgID, ok := m.ExplicitLayerEvidenceForAuthKey(authKeyID, sessionID); !ok || layer != 226 || msgID != 1234 {
t.Fatalf("live explicit evidence = (%d,%d,%v)", layer, msgID, ok)
}
inherited := &Conn{authKeyID: authKeyID, sessionID: sessionID + 1}
if err := inherited.SeedInheritedLayerProfile(tg.LayerProfile227); err != nil {
t.Fatal(err)
}
if err := m.Register(inherited); err != nil {
t.Fatal(err)
}
if layer, msgID, ok := m.ExplicitLayerEvidenceForAuthKey(authKeyID, sessionID+1); ok || layer != 0 || msgID != 0 {
t.Fatalf("inherited state exposed as explicit = (%d,%d,%v)", layer, msgID, ok)
}
}
func TestSessionManagerExplicitLayerEvidenceChoosesNewestActiveOrClaim(t *testing.T) {
m := NewSessionManager(nil)
authKeyID := [8]byte{2, 3, 0}
const sessionID = int64(230)
active := &Conn{authKeyID: authKeyID, sessionID: sessionID}
claim := &Conn{authKeyID: authKeyID, sessionID: sessionID}
if err := active.seedOrderedLayerProfile(tg.LayerProfile225, 100); err != nil {
t.Fatal(err)
}
if err := claim.seedOrderedLayerProfile(tg.LayerProfile227, 104); err != nil {
t.Fatal(err)
}
key := sessionKey{authKeyID: authKeyID, sessionID: sessionID}
m.mu.Lock()
m.bySession[key] = active
m.claims[key] = claim
m.mu.Unlock()
if layer, msgID, ok := m.ExplicitLayerEvidenceForAuthKey(authKeyID, sessionID); !ok || layer != 227 || msgID != 104 {
t.Fatalf("newest active/claim evidence = (%d,%d,%v)", layer, msgID, ok)
}
claim.retire()
if layer, msgID, ok := m.ExplicitLayerEvidenceForAuthKey(authKeyID, sessionID); !ok || layer != 225 || msgID != 100 {
t.Fatalf("retired claim shadowed active evidence = (%d,%d,%v)", layer, msgID, ok)
}
}
func TestOrderedSessionLayerBroadcastConvergesAcrossPhysicalGenerations(t *testing.T) {
m := NewSessionManager(nil)
authKeyID := [8]byte{3, 0, 0}
const sessionID = int64(300)
oldPhysical := &Conn{authKeyID: authKeyID, sessionID: sessionID}
current := &Conn{authKeyID: authKeyID, sessionID: sessionID}
m.mu.Lock()
m.bySession[sessionKey{authKeyID: authKeyID, sessionID: sessionID}] = current
m.mu.Unlock()
if applied, err := m.ApplyOrderedLayerProfileForSession(oldPhysical, authKeyID, sessionID, tg.LayerProfile227, 300); err != nil || applied != 2 {
t.Fatalf("newer broadcast applied=%d err=%v", applied, err)
}
if applied, err := m.ApplyOrderedLayerProfileForSession(oldPhysical, authKeyID, sessionID, tg.LayerProfile225, 200); err != nil || applied != 0 {
t.Fatalf("delayed older broadcast applied=%d err=%v", applied, err)
}
for name, c := range map[string]*Conn{"old": oldPhysical, "current": current} {
state, msgID := c.layerProfileEvidenceState()
if state.Profile != tg.LayerProfile227 || state.Origin != LayerProfileExplicit || msgID != 300 {
t.Fatalf("%s physical state = %#v msgID:%d", name, state, msgID)
}
}
}
func TestInitialProfileSeedAvoidsPermanentKeyResolverAndPrefersPermForTemp(t *testing.T) {
t.Run("fetched permanent metadata is authoritative fast path", func(t *testing.T) {
resolver := &countingInheritedLayerResolver{layer: 225, found: true}
s := &Server{layerRPC: resolver}
c := &Conn{authKeyExpiresAt: 0}
if err := s.seedInitialLayerProfile(context.Background(), c, 227, LayerProfileSnapshot{}); err != nil {
t.Fatal(err)
}
if resolver.calls != 0 {
t.Fatalf("permanent resolver calls = %d, want 0", resolver.calls)
}
if got := c.LayerProfileState(); got.Profile != tg.LayerProfile227 || got.Origin != LayerProfileInherited {
t.Fatalf("permanent seed = %#v", got)
}
})
t.Run("temporary key resolves permanent before raw shadow", func(t *testing.T) {
resolver := &countingInheritedLayerResolver{layer: 227, found: true}
s := &Server{layerRPC: resolver}
c := &Conn{authKeyExpiresAt: 1_900_000_000}
if err := s.seedInitialLayerProfile(context.Background(), c, 225, LayerProfileSnapshot{}); err != nil {
t.Fatal(err)
}
if resolver.calls != 1 {
t.Fatalf("temporary resolver calls = %d, want 1", resolver.calls)
}
if got := c.LayerProfileState(); got.Profile != tg.LayerProfile227 || got.Origin != LayerProfileInherited {
t.Fatalf("temporary canonical seed = %#v", got)
}
})
t.Run("unsupported permanent metadata stays unknown", func(t *testing.T) {
resolver := &countingInheritedLayerResolver{layer: 227, found: true}
s := &Server{layerRPC: resolver}
c := &Conn{authKeyExpiresAt: 0}
if err := s.seedInitialLayerProfile(context.Background(), c, 229, LayerProfileSnapshot{}); err != nil {
t.Fatal(err)
}
if resolver.calls != 0 {
t.Fatalf("unsupported permanent metadata fell through to resolver")
}
if got := c.LayerProfileState(); got.Origin != LayerProfileUnknown {
t.Fatalf("unsupported permanent seed = %#v", got)
}
})
t.Run("unsupported bound permanent blocks raw temp shadow", func(t *testing.T) {
resolver := &countingInheritedLayerResolver{layer: 229, found: true}
s := &Server{layerRPC: resolver}
c := &Conn{authKeyExpiresAt: 1_900_000_000}
if err := s.seedInitialLayerProfile(context.Background(), c, 225, LayerProfileSnapshot{}); err != nil {
t.Fatal(err)
}
if got := c.LayerProfileState(); got.Origin != LayerProfileUnknown {
t.Fatalf("authoritative unsupported perm fell back to raw shadow = %#v", got)
}
})
}
func TestInitialProfileSeedRestoresOrderedExactSessionEvidence(t *testing.T) {
resolver := &orderedSessionLayerResolver{layer: 226, msgID: 123456, found: true}
s := &Server{layerRPC: resolver}
c := &Conn{authKeyExpiresAt: 0}
if err := s.seedInitialLayerProfile(context.Background(), c, 227, LayerProfileSnapshot{}); err != nil {
t.Fatal(err)
}
state, msgID := c.layerProfileEvidenceState()
if state.Profile != tg.LayerProfile226 || state.Origin != LayerProfileExplicit || msgID != resolver.msgID {
t.Fatalf("ordered exact seed = state:%#v msgID:%d", state, msgID)
}
}
func TestInheritedLayerResolverFailureLeavesExplicitRecoveryAvailable(t *testing.T) {
resolverErr := errors.New("temporary resolver unavailable")
resolver := &countingInheritedLayerResolver{err: resolverErr}
s := &Server{layerRPC: resolver}
c := &Conn{authKeyExpiresAt: 1_900_000_000}
if err := s.seedInitialLayerProfile(context.Background(), c, 225, LayerProfileSnapshot{}); err != nil {
t.Fatalf("optional inherited resolver error escaped: %v", err)
}
if got := c.LayerProfileState(); got.Origin != LayerProfileUnknown {
t.Fatalf("resolver error selected stale fallback = %#v", got)
}
}
func TestInheritedLayerResolverAvailabilityUsesOnlySupportedRawTempShadow(t *testing.T) {
for _, tt := range []struct {
name string
fetchedLayer int
wantProfile tg.LayerProfile
wantOrigin LayerProfileOrigin
}{
{name: "supported raw shadow", fetchedLayer: 225, wantProfile: tg.LayerProfile225, wantOrigin: LayerProfileInherited},
{name: "future raw shadow stays unknown", fetchedLayer: 229, wantOrigin: LayerProfileUnknown},
} {
t.Run(tt.name, func(t *testing.T) {
resolver := &countingInheritedLayerResolver{err: layerDurabilityUnavailableTestError{}}
s := &Server{layerRPC: resolver}
c := &Conn{authKeyExpiresAt: 1_900_000_000}
if err := s.seedInitialLayerProfile(context.Background(), c, tt.fetchedLayer, LayerProfileSnapshot{}); err != nil {
t.Fatal(err)
}
if resolver.calls != 1 {
t.Fatalf("resolver calls=%d, want 1", resolver.calls)
}
if got := c.LayerProfileState(); got.Profile != tt.wantProfile || got.Origin != tt.wantOrigin {
t.Fatalf("availability seed = %#v, want profile=%d origin=%d", got, tt.wantProfile, tt.wantOrigin)
}
})
}
}
func TestActivationClaimRecheckClosesTempBindLayerRace(t *testing.T) {
authKeyID := [8]byte{4, 0, 0}
t.Run("bind wins before claimed recheck", func(t *testing.T) {
resolver := &countingInheritedLayerResolver{found: false}
s := &Server{layerRPC: resolver}
m := NewSessionManager(nil)
c := &Conn{authKeyID: authKeyID, sessionID: 401, authKeyExpiresAt: 1_900_000_000}
if err := s.seedInitialLayerProfile(context.Background(), c, 225, LayerProfileSnapshot{}); err != nil {
t.Fatal(err)
}
if err := m.BeginActivation(c); err != nil {
t.Fatal(err)
}
defer m.AbortActivation(c)
resolver.layer, resolver.found = 227, true
if err := s.refreshActivatedInheritedLayerProfile(context.Background(), c, 225); err != nil {
t.Fatal(err)
}
if got := c.LayerProfileState(); got.Profile != tg.LayerProfile227 || got.Origin != LayerProfileInherited {
t.Fatalf("post-claim permanent recheck = %#v", got)
}
})
t.Run("claim wins before bind refresh", func(t *testing.T) {
m := NewSessionManager(nil)
c := &Conn{authKeyID: authKeyID, sessionID: 402, authKeyExpiresAt: 1_900_000_000}
if err := c.SeedInheritedLayerProfile(tg.LayerProfile225); err != nil {
t.Fatal(err)
}
if err := m.BeginActivation(c); err != nil {
t.Fatal(err)
}
defer m.AbortActivation(c)
if refreshed := m.RefreshInheritedLayerForRawAuthKey(authKeyID, 227); refreshed != 1 {
t.Fatalf("bind refresh count = %d, want 1", refreshed)
}
if got := c.LayerProfileState(); got.Profile != tg.LayerProfile227 || got.Origin != LayerProfileInherited {
t.Fatalf("claim-visible bind refresh = %#v", got)
}
})
t.Run("unsupported permanent clears preclaim raw shadow", func(t *testing.T) {
resolver := &countingInheritedLayerResolver{layer: 229, found: true}
s := &Server{layerRPC: resolver}
c := &Conn{authKeyID: authKeyID, sessionID: 403, authKeyExpiresAt: 1_900_000_000}
if err := c.SeedInheritedLayerProfile(tg.LayerProfile225); err != nil {
t.Fatal(err)
}
if err := s.refreshActivatedInheritedLayerProfile(context.Background(), c, 225); err != nil {
t.Fatal(err)
}
if got := c.LayerProfileState(); got.Origin != LayerProfileUnknown {
t.Fatalf("unsupported permanent kept stale raw shadow = %#v", got)
}
})
t.Run("availability outage retains current raw shadow", func(t *testing.T) {
resolver := &countingInheritedLayerResolver{err: layerDurabilityUnavailableTestError{}}
s := &Server{layerRPC: resolver}
c := &Conn{authKeyID: authKeyID, sessionID: 404, authKeyExpiresAt: 1_900_000_000}
if err := s.refreshActivatedInheritedLayerProfile(context.Background(), c, 225); err != nil {
t.Fatal(err)
}
if got := c.LayerProfileState(); got.Profile != tg.LayerProfile225 || got.Origin != LayerProfileInherited {
t.Fatalf("availability recheck lost raw shadow = %#v", got)
}
})
t.Run("structural resolver failure clears preclaim raw shadow", func(t *testing.T) {
resolver := &countingInheritedLayerResolver{err: errors.New("invalid binding identity")}
s := &Server{layerRPC: resolver}
c := &Conn{authKeyID: authKeyID, sessionID: 405, authKeyExpiresAt: 1_900_000_000}
if err := c.SeedInheritedLayerProfile(tg.LayerProfile225); err != nil {
t.Fatal(err)
}
if err := s.refreshActivatedInheritedLayerProfile(context.Background(), c, 225); err != nil {
t.Fatal(err)
}
if got := c.LayerProfileState(); got.Origin != LayerProfileUnknown {
t.Fatalf("structural resolver failure retained raw shadow = %#v", got)
}
})
}

282
internal/mtprotoedge/layer_rpc_execution.go Normal file
View file

@ -0,0 +1,282 @@
package mtprotoedge
import (
"context"
"errors"
"time"
"github.com/iamxvbaba/td/bin"
"github.com/iamxvbaba/td/mt"
"github.com/iamxvbaba/td/tg"
"github.com/iamxvbaba/td/tgerr"
"go.uber.org/zap"
"telesrv/internal/observability/dbtrace"
"telesrv/internal/postresponse"
)
// layerRPCResultEncoder keeps the generated result bound to the immutable
// admitted call, but deliberately does not prepare a second byte snapshot in
// the inbound worker. The only Encode call happens later under outbound encode
// and retained-byte admission.
type layerRPCResultEncoder struct {
call tg.LayerCall
result tg.LayerRPCResult
}
func (e *layerRPCResultEncoder) Encode(b *bin.Buffer) error {
if e == nil {
return errors.New("nil layer RPC result")
}
if e.result == nil {
return errors.New("nil generated layer RPC result")
}
return e.result.Encode(b)
}
func (e *layerRPCResultEncoder) exactLayerRPCResultBinding() outboundLayerBinding {
if e == nil {
return outboundLayerBinding{}
}
return outboundLayerBinding{
profile: e.call.Profile(),
typ: e.call.WireResultType(),
wireInvariant: e.call.WireInvariant(),
kind: outboundLayerBindingRequest,
}
}
type exactLayerRPCResultEncoder interface {
bin.Encoder
exactLayerRPCResultBinding() outboundLayerBinding
}
// legacyTestRPCResultEncoder migrates the unexported legacyRPC package-test
// hook onto the generated result codec. The hook may still exercise the old
// scheduling API, but it no longer has a canonical-bytes escape hatch.
type legacyTestRPCResultEncoder struct {
call tg.LayerCall
result bin.Encoder
}
func (e *legacyTestRPCResultEncoder) Encode(b *bin.Buffer) error {
if e == nil || e.result == nil {
return errors.New("nil legacy test RPC result")
}
return e.call.EncodeResult(e.result, b)
}
func (e *legacyTestRPCResultEncoder) exactLayerRPCResultBinding() outboundLayerBinding {
if e == nil {
return outboundLayerBinding{}
}
return outboundLayerBinding{
profile: e.call.Profile(),
typ: e.call.WireResultType(),
wireInvariant: e.call.WireInvariant(),
kind: outboundLayerBindingRequest,
}
}
var errLayerRPCResultIdentityMismatch = errors.New("layer RPC result does not match admitted request identity")
// bindAdmittedLayerRPCResult closes the integration boundary around the
// generated dispatcher. A LayerRPCHandler implementation must return the
// result capability created from this exact admission; accepting a result from
// another request would pair the wrong result TypeRef/profile with this
// flight/cache identity even when both methods happen to share a Go type.
func bindAdmittedLayerRPCResult(request tg.LayerRequest, result tg.LayerRPCResult) (*layerRPCResultEncoder, error) {
if result == nil {
return nil, nil
}
if result.Prepared().Identity() != request.Prepared().Identity() {
return nil, errLayerRPCResultIdentityMismatch
}
return &layerRPCResultEncoder{call: request.Call(), result: result}, nil
}
func (s *Server) newInboundLayerRPCTask(
c *Conn,
msgID int64,
admissionSeq uint64,
method string,
profileEvidenceFresh bool,
request tg.LayerRequest,
dependencies layerRPCDependencySet,
owner *rpcResultOwnerLease,
) inboundRPC {
wireSize := request.Prepared().WireSize()
gate := newLayerRPCExecutionGate(c, dependencies)
timeoutResponse := func() {
writeTimeout := c.writeTimeout
if writeTimeout <= 0 || writeTimeout > 5*time.Second {
writeTimeout = 5 * time.Second
}
responseCtx, cancel := context.WithTimeout(context.Background(), writeTimeout)
defer cancel()
if sendErr := s.sendResult(responseCtx, c, msgID, &mt.RPCError{
ErrorCode: 500, ErrorMessage: layerRPCTimeoutMessage(gate),
}); sendErr != nil && !isClientDisconnect(sendErr) {
s.log.Debug("Send RPC timeout failed",
zap.String("method", method), zap.Int64("msg_id", msgID),
zap.String("auth_key_id", c.authKeyHex), zap.Int64("session_id", c.sessionID),
zap.Error(sendErr))
}
}
return inboundRPC{
method: method,
size: wireSize,
onTimeout: timeoutResponse,
release: func() {
if owner != nil && owner.Abort() {
c.fenceUndeliveredRPCResult()
}
},
gate: gate,
run: func(taskCtx context.Context) error {
if gate != nil && !gate.success() {
return s.publishAdmittedLayerRPCResult(c, msgID, method, owner, false, &mt.RPCError{
ErrorCode: 500, ErrorMessage: "MSG_WAIT_FAILED",
}, nil)
}
if err := s.handleAdmittedLayerRPC(s.withLayerRPCProfileEvidenceFresh(taskCtx, profileEvidenceFresh), c, msgID, admissionSeq, method, request, owner); err != nil {
fields := []zap.Field{
zap.Int64("msg_id", msgID), zap.String("auth_key_id", c.authKeyHex),
zap.Int64("session_id", c.sessionID), zap.Error(err),
}
if isClientDisconnect(err) {
s.log.Debug("RPC async handler canceled", fields...)
} else {
s.log.Info("RPC async handler failed", fields...)
}
return err
}
return nil
},
}
}
func layerRPCTimeoutMessage(gate *inboundRPCGate) string {
if gate != nil && !gate.runnable() {
return "MSG_WAIT_TIMEOUT"
}
return "RPC_TIMEOUT"
}
func newLayerRPCExecutionGate(c *Conn, dependencies layerRPCDependencySet) *inboundRPCGate {
if len(dependencies.waiters) == 0 && !dependencies.failed {
return nil
}
gate := newInboundRPCGate(len(dependencies.waiters), c.wakeInboundRPC)
if dependencies.failed {
gate.failed.Store(true)
}
for _, waiter := range dependencies.waiters {
if err := waiter.SubscribeExecution(gate.resolve); err != nil {
gate.resolve(false)
}
}
// Release the subscriber-installation sentinel.
gate.resolve(true)
return gate
}
func (s *Server) publishAdmittedLayerRPCResult(
c *Conn,
msgID int64,
method string,
owner *rpcResultOwnerLease,
success bool,
result bin.Encoder,
after func(),
) error {
if owner != nil {
owner.CompleteExecution(success)
}
return s.publishRPCResult(c, msgID, method, owner, result, after)
}
func (s *Server) handleAdmittedLayerRPC(
ctx context.Context,
c *Conn,
msgID int64,
admissionSeq uint64,
method string,
request tg.LayerRequest,
owner *rpcResultOwnerLease,
) error {
if s.layerRPC == nil {
return s.publishAdmittedLayerRPCResult(c, msgID, method, owner, false, &mt.RPCError{
ErrorCode: 500, ErrorMessage: "NOT_IMPLEMENTED",
}, nil)
}
ctx = postresponse.WithCallbacks(ctx)
ctx, dbStats := dbtrace.WithStats(ctx)
start := s.clock.Now()
result, effectiveMethod, err := s.layerRPC.DispatchAdmitted(ctx, c.authKeyID, c.sessionID, msgID, admissionSeq, request)
businessSucceeded := err == nil
if effectiveMethod == "" {
effectiveMethod = method
}
var exact *layerRPCResultEncoder
if err == nil && result != nil {
exact, err = bindAdmittedLayerRPCResult(request, result)
}
dur := s.clock.Now().Sub(start)
s.metrics.RPCHandled(effectiveMethod, dur, err)
fields := []zap.Field{
zap.String("method", effectiveMethod), zap.String("auth_key_id", c.authKeyHex),
zap.Int64("session_id", c.sessionID), zap.Int64("msg_id", msgID),
zap.Int("profile", int(request.Call().Profile())), zap.Duration("dur", dur),
}
if effectiveMethod != method {
fields = append(fields, zap.String("outer_method", method))
}
if businessAuthKeyHex, ok := c.BusinessAuthKeyHex(); ok {
fields = append(fields, zap.String("business_auth_key_id", businessAuthKeyHex))
}
if userID := c.UserID(); userID != 0 {
fields = append(fields, zap.Int64("user_id", userID))
}
fields = dbtrace.AppendZapFields(fields, "", dbStats.Snapshot())
if ctxErr := ctx.Err(); ctxErr != nil {
var terminal bin.Encoder
var after func()
if err == nil && exact != nil {
terminal = exact
after = postresponse.Take(context.WithoutCancel(ctx))
} else if errors.Is(ctxErr, context.DeadlineExceeded) {
terminal = &mt.RPCError{ErrorCode: 500, ErrorMessage: "RPC_TIMEOUT"}
}
if terminal != nil {
if sendErr := s.publishAdmittedLayerRPCResult(c, msgID, effectiveMethod, owner, err == nil && exact != nil, terminal, after); sendErr != nil {
s.log.Debug("Publish canceled RPC result failed", append(fields, zap.Error(sendErr))...)
}
}
s.log.Info("RPC canceled", append(fields, zap.NamedError("context_error", ctxErr))...)
return ctxErr
}
if err != nil {
var rpcErr *tgerr.Error
if errors.As(err, &rpcErr) {
s.log.Info("RPC error", append(fields, zap.Int("code", rpcErr.Code), zap.String("error", rpcErr.Message))...)
return s.publishAdmittedLayerRPCResult(c, msgID, effectiveMethod, owner, businessSucceeded, &mt.RPCError{
ErrorCode: rpcErr.Code, ErrorMessage: rpcErr.Message,
}, nil)
}
s.log.Info("RPC internal error", append(fields, zap.Error(err))...)
return s.publishAdmittedLayerRPCResult(c, msgID, effectiveMethod, owner, businessSucceeded, &mt.RPCError{
ErrorCode: 500, ErrorMessage: "INTERNAL",
}, nil)
}
if exact == nil {
return s.publishAdmittedLayerRPCResult(c, msgID, effectiveMethod, owner, businessSucceeded, &mt.RPCError{
ErrorCode: 500, ErrorMessage: "INTERNAL",
}, nil)
}
s.log.Info("RPC handled", fields...)
return s.publishAdmittedLayerRPCResult(c, msgID, effectiveMethod, owner, true, exact, postresponse.Take(ctx))
}
var _ exactLayerRPCResultEncoder = (*layerRPCResultEncoder)(nil)

153
internal/mtprotoedge/layer_rpc_execution_test.go Normal file
View file

@ -0,0 +1,153 @@
package mtprotoedge
import (
"context"
"errors"
"sync/atomic"
"testing"
"time"
"github.com/iamxvbaba/td/bin"
"github.com/iamxvbaba/td/mt"
"github.com/iamxvbaba/td/proto"
"github.com/iamxvbaba/td/tg"
)
// preparedOnlyLayerRPCResult is intentionally incapable of encoding. The
// binding guard must reject it solely from immutable admission identity before
// any result method other than Prepared can be observed.
type preparedOnlyLayerRPCResult struct {
tg.LayerRPCResult
prepared tg.LayerPreparedCall
}
func (r *preparedOnlyLayerRPCResult) Prepared() tg.LayerPreparedCall { return r.prepared }
func TestBindAdmittedLayerRPCResultRequiresExactRequestIdentity(t *testing.T) {
dispatcher := tg.NewServerDispatcher(nil)
admit := func(request bin.Encoder) tg.LayerRequest {
t.Helper()
body := &bin.Buffer{Buf: exactLayerRPCBody(t, request)}
admitted, err := dispatcher.AdmitLayer(tg.LayerProfile227, body)
if err != nil {
t.Fatal(err)
}
return admitted
}
request := admit(&tg.HelpGetConfigRequest{})
other := admit(&tg.HelpGetNearestDCRequest{})
if _, err := bindAdmittedLayerRPCResult(request, &preparedOnlyLayerRPCResult{prepared: other.Prepared()}); !errors.Is(err, errLayerRPCResultIdentityMismatch) {
t.Fatalf("mismatched result error = %v, want %v", err, errLayerRPCResultIdentityMismatch)
}
bound, err := bindAdmittedLayerRPCResult(request, &preparedOnlyLayerRPCResult{prepared: request.Prepared()})
if err != nil {
t.Fatal(err)
}
if bound == nil || bound.call.Identity() != request.Call().Identity() {
t.Fatal("matching result did not retain the admitted call identity")
}
}
type mismatchedProjectionLayerRPC struct {
*admissionOnlyLayerRPC
result tg.LayerRPCResult
calls atomic.Int32
}
func (h *mismatchedProjectionLayerRPC) DispatchAdmitted(
context.Context,
[8]byte,
int64,
int64,
uint64,
tg.LayerRequest,
) (tg.LayerRPCResult, string, error) {
h.calls.Add(1)
return h.result, "help.getConfig", nil
}
func TestProjectionFailureCachesInternalWithoutRepeatingBusiness(t *testing.T) {
dispatcher := tg.NewServerDispatcher(nil)
admit := func(request bin.Encoder) tg.LayerRequest {
t.Helper()
body := &bin.Buffer{Buf: exactLayerRPCBody(t, request)}
admitted, err := dispatcher.AdmitLayer(tg.LayerProfile227, body)
if err != nil {
t.Fatal(err)
}
return admitted
}
request := admit(&tg.HelpGetConfigRequest{})
other := admit(&tg.HelpGetNearestDCRequest{})
handler := &mismatchedProjectionLayerRPC{
admissionOnlyLayerRPC: newAdmissionOnlyLayerRPC(),
result: &preparedOnlyLayerRPCResult{prepared: other.Prepared()},
}
s := New(Options{DC: 2, LayerRPC: handler})
c := newOutboundTestConn(t, &collectingSessionTransport{}, newOutboundTrackedBudget(1<<20))
c.authKeyID = [8]byte{0x41, 0x01}
c.sessionID = 4101
const reqMsgID = int64(410100)
claim, err := s.rpcResults.AcquireLayerIdentified(
c.authKeyID, c.sessionID, reqMsgID,
tg.LayerProfile227, request.Prepared().Identity(),
)
if err != nil || claim.owner == nil {
t.Fatalf("owner acquisition err=%v", err)
}
if err := s.handleAdmittedLayerRPC(
context.Background(), c, reqMsgID, claim.admissionSeq,
"help.getConfig", request, claim.owner,
); err != nil {
t.Fatalf("publish projection failure: %v", err)
}
deadline := time.Now().Add(2 * time.Second)
var completed rpcResultAcquire
for {
completed, err = s.rpcResults.AcquireLayerIdentified(
c.authKeyID, c.sessionID, reqMsgID,
tg.LayerProfile227, request.Prepared().Identity(),
)
if err == nil && completed.state == rpcResultAcquireCompleted {
break
}
if time.Now().After(deadline) {
t.Fatalf("projection failure did not become completed: state=%d err=%v", completed.state, err)
}
time.Sleep(time.Millisecond)
}
if !completed.executionKnown || !completed.executionOK {
t.Fatalf("projection failure lost successful business outcome: known=%v ok=%v", completed.executionKnown, completed.executionOK)
}
if got := handler.calls.Load(); got != 1 {
t.Fatalf("business calls=%d, want 1", got)
}
var envelope proto.Result
if err := envelope.Decode(&bin.Buffer{Buf: completed.encoded.body}); err != nil {
t.Fatal(err)
}
var rpcErr mt.RPCError
if err := rpcErr.Decode(&bin.Buffer{Buf: envelope.Result}); err != nil {
t.Fatal(err)
}
if rpcErr.ErrorCode != 500 || rpcErr.ErrorMessage != "INTERNAL" {
t.Fatalf("projection terminal = %+v", rpcErr)
}
// A same-msg replay is served from the completed exact identity; there is no
// second DispatchAdmitted call even though projection failed after business
// success.
replay, err := s.rpcResults.AcquireLayerIdentified(
c.authKeyID, c.sessionID, reqMsgID,
tg.LayerProfile227, request.Prepared().Identity(),
)
if err != nil || replay.state != rpcResultAcquireCompleted || replay.encoded != completed.encoded {
t.Fatalf("projection replay = state:%d err:%v", replay.state, err)
}
if got := handler.calls.Load(); got != 1 {
t.Fatalf("replay repeated business calls=%d", got)
}
}

33
internal/mtprotoedge/layer_seed_test.go
View file

@ -5,33 +5,12 @@ import (
"testing"
"time"
"go.uber.org/zap/zaptest"
"github.com/gotd/td/bin"
"github.com/gotd/td/mt"
"github.com/gotd/td/proto"
"github.com/gotd/td/tg"
"github.com/iamxvbaba/td/bin"
"github.com/iamxvbaba/td/mt"
"github.com/iamxvbaba/td/proto"
"github.com/iamxvbaba/td/tg"
)
func TestSessionManagerSetClientLayerForAuthKey(t *testing.T) {
sm := NewSessionManager(zaptest.NewLogger(t))
c := &Conn{sessionID: 42, authKeyID: [8]byte{1, 2, 3}}
sm.Register(c)
defer sm.Unregister(c)
sm.SetClientLayerForAuthKey([8]byte{1, 2, 3}, 42, 225)
if got := c.ClientLayer(); got != 225 {
t.Fatalf("ClientLayer = %d, want 225", got)
}
// 未注册 session：no-op，不 panic。
sm.SetClientLayerForAuthKey([8]byte{9}, 1, 220)
// 非法 layer：忽略，保留已有值。
sm.SetClientLayerForAuthKey([8]byte{1, 2, 3}, 42, 0)
if got := c.ClientLayer(); got != 225 {
t.Fatalf("ClientLayer after layer=0 = %d, want 225", got)
}
}
type seededLayerRPC struct{}
func (seededLayerRPC) Dispatch(context.Context, [8]byte, int64, *bin.Buffer) (bin.Encoder, error) {
@ -45,7 +24,7 @@ func (seededLayerRPC) NegotiatedLayer([8]byte, int64) (int, bool) { return 225,
// 而不是等首条 RPC 的 Dispatch 返回后才刷新——否则重连老客户端在首条 RPC handler
// 执行期间收到的 pending flush / 并发 push 会按 canonical 227 漏降级。
func TestRegisterSeedsNegotiatedLayerBeforeFirstRPC(t *testing.T) {
addr, pub, srv := startTestServer(t, Options{DC: 2, RPC: seededLayerRPC{}})
addr, pub, srv := startTestServer(t, Options{DC: 2, legacyRPC: seededLayerRPC{}})
conn, auth, cipher := dialHandshake(t, addr, 2, pub)
clientMsgID := proto.NewMessageIDGen(time.Now)
@ -67,7 +46,7 @@ func TestRegisterSeedsNegotiatedLayerBeforeFirstRPC(t *testing.T) {
if c == nil {
t.Fatal("connection not registered")
}
if got := c.ClientLayer(); got != 225 {
if got := c.legacyClientLayer(); got != 225 {
t.Fatalf("ClientLayer after registration = %d, want seeded 225", got)
}
}

235
internal/mtprotoedge/layer_updates.go Normal file
View file

@ -0,0 +1,235 @@
package mtprotoedge
import (
"context"
"errors"
"fmt"
"sync"
"github.com/iamxvbaba/td/bin"
"github.com/iamxvbaba/td/tg"
)
var (
ErrOutboundLayerProfileUnknown = errors.New("outbound exact layer profile is unknown")
ErrOutboundLayerProfileMismatch = errors.New("outbound exact layer profile mismatch")
// ErrOutboundLayerBindingRequired means application bytes reached the edge
// without the request- or session-owned profile proof required by the native
// generated codec. Never substitute canonical bytes or the retired legacy
// transcoder on a production connection.
ErrOutboundLayerBindingRequired = errors.New("outbound application value requires exact layer binding")
// ErrOutboundLayerProfileStale means a proactive update was prepared for an
// older connection profile epoch. It is deliberately non-terminal: discard
// the online accelerator and let durable difference recovery fill the gap.
ErrOutboundLayerProfileStale = errors.New("outbound exact layer profile epoch is stale")
)
type outboundLayerBindingKind uint8
const (
// The zero value preserves strict connection binding for defensive and
// legacy constructions which do not declare a stronger ownership model.
outboundLayerBindingSession outboundLayerBindingKind = iota
// A request-bound result retains the profile captured by admission. A later
// invokeWithLayer correction must not invalidate that in-flight/cached result.
outboundLayerBindingRequest
)
type outboundLayerBinding struct {
profile tg.LayerProfile
typ *tg.LayerTypeRef
wireInvariant bool
kind outboundLayerBindingKind
// epoch is required for proactive updates. Zero is accepted only for older
// defensive/test bindings which predate mutable profile correction.
epoch uint32
}
type preparedLayerUpdates struct {
done chan struct{}
encoded *encodedOutboundMessage
err error
}
// layerUpdatesFanout is one immutable canonical Updates snapshot plus a
// request-scoped cache of exact prepared bytes. FreezeLayer and
// PrepareFrozenLayer are the same generated TypeRef codec used by RPC results
// and differences; this type adds only fan-out singleflight and ownership.
type layerUpdatesFanout struct {
frozen tg.LayerFrozen[tg.UpdatesClass]
size int
mu sync.Mutex
prepared map[tg.LayerProfile]*preparedLayerUpdates
}
func newLayerUpdatesFanout(value tg.UpdatesClass) (*layerUpdatesFanout, error) {
frozen, err := tg.FreezeLayer(tg.LayerClassUpdatesType(), value)
if err != nil {
return nil, fmt.Errorf("freeze exact layer updates: %w", err)
}
return &layerUpdatesFanout{
frozen: frozen,
size: frozen.CanonicalSize(),
prepared: make(map[tg.LayerProfile]*preparedLayerUpdates),
}, nil
}
func (u *layerUpdatesFanout) canonicalSize() int {
if u == nil {
return 0
}
return u.size
}
func (u *layerUpdatesFanout) prepareForConn(ctx context.Context, c *Conn) (*encodedOutboundMessage, error) {
if u == nil || c == nil {
return nil, errors.New("nil exact layer update or connection")
}
state := c.LayerProfileState()
if state.Origin == LayerProfileUnknown {
return nil, ErrOutboundLayerProfileUnknown
}
base, err := u.prepare(ctx, state.Profile)
if err != nil {
return nil, err
}
if base == nil || base.layer == nil {
return nil, errors.New("prepared exact layer update lost binding")
}
// Exact bytes stay shared per profile. Only the small binding is copied per
// physical target because epoch belongs to that connection.
encoded := *base
binding := *base.layer
binding.kind = outboundLayerBindingSession
binding.epoch = state.Epoch
encoded.layer = &binding
return &encoded, nil
}
func (u *layerUpdatesFanout) prepare(ctx context.Context, profile tg.LayerProfile) (*encodedOutboundMessage, error) {
if ctx == nil {
ctx = context.Background()
}
u.mu.Lock()
entry := u.prepared[profile]
if entry == nil {
entry = &preparedLayerUpdates{done: make(chan struct{})}
u.prepared[profile] = entry
u.mu.Unlock()
entry.encoded, entry.err = prepareFrozenLayerUpdatesContext(ctx, profile, u.frozen)
close(entry.done)
if entry.err != nil {
// Context/encode admission failure is attempt-local. Do not poison this
// semantic update for every later connection or pending-flush retry.
u.mu.Lock()
if u.prepared[profile] == entry {
delete(u.prepared, profile)
}
u.mu.Unlock()
}
} else {
u.mu.Unlock()
// A completed profile is immutable and no longer needs caller budget.
// Prefer it deterministically even when the shared fan-out deadline became
// ready at the same instant; otherwise Go's select may randomly choose the
// canceled branch and skip a healthy later connection of the same profile.
select {
case <-entry.done:
return entry.encoded, entry.err
default:
}
select {
case <-entry.done:
case <-ctx.Done():
return nil, ctx.Err()
}
}
return entry.encoded, entry.err
}
func (u *layerUpdatesFanout) discardPrepared(profile tg.LayerProfile, encoded *encodedOutboundMessage) {
if u == nil || encoded == nil {
return
}
u.mu.Lock()
entry := u.prepared[profile]
if entry != nil {
select {
case <-entry.done:
if entry.encoded == encoded || (entry.encoded != nil &&
entry.encoded.layer != nil && encoded.layer != nil &&
entry.encoded.layer.profile == encoded.layer.profile &&
sameBacking(entry.encoded.body, encoded.body)) {
delete(u.prepared, profile)
}
default:
}
}
u.mu.Unlock()
}
func prepareFrozenLayerUpdatesContext(
ctx context.Context,
profile tg.LayerProfile,
frozen tg.LayerFrozen[tg.UpdatesClass],
) (*encodedOutboundMessage, error) {
var encoded *encodedOutboundMessage
err := withOutboundEncodeSlot(ctx, nil, func() error {
prepared, err := tg.PrepareFrozenLayer(profile, frozen)
if err != nil {
return err
}
var body bin.Buffer
typ := tg.LayerClassUpdatesType()
if err := prepared.Encode(profile, typ, &body); err != nil {
return err
}
id, err := body.PeekID()
if err != nil {
return fmt.Errorf("peek exact updates constructor: %w", err)
}
encoded = &encodedOutboundMessage{
body: body.Copy(), typeID: id,
layer: &outboundLayerBinding{profile: profile, typ: prepared.TypeRef()},
}
return nil
})
if err != nil {
return nil, fmt.Errorf("prepare updates for layer %d: %w", profile, err)
}
return encoded, nil
}
func validateOutboundLayerBinding(c *Conn, encoded *encodedOutboundMessage) error {
if encoded == nil || encoded.layer == nil {
return nil
}
if encoded.layer.typ == nil {
return errors.New("outbound exact layer TypeRef is nil")
}
if encoded.layer.wireInvariant || encoded.layer.kind == outboundLayerBindingRequest {
return nil
}
state := c.LayerProfileState()
if state.Origin == LayerProfileUnknown {
return ErrOutboundLayerProfileUnknown
}
if encoded.layer.epoch != 0 && state.Epoch != encoded.layer.epoch {
return fmt.Errorf("%w: connection=%d encoded=%d", ErrOutboundLayerProfileStale, state.Epoch, encoded.layer.epoch)
}
if state.Profile != encoded.layer.profile {
return fmt.Errorf("%w: connection=%d encoded=%d", ErrOutboundLayerProfileMismatch, state.Profile, encoded.layer.profile)
}
return nil
}
func isOutboundStaleLayerEpoch(err error) bool {
return errors.Is(err, ErrOutboundLayerProfileStale)
}
func isOutboundLayerProfileError(err error) bool {
return errors.Is(err, ErrOutboundLayerProfileUnknown) ||
errors.Is(err, ErrOutboundLayerProfileMismatch)
}

326
internal/mtprotoedge/layer_updates_test.go Normal file
View file

@ -0,0 +1,326 @@
package mtprotoedge
import (
"context"
"errors"
"fmt"
"io"
"sync"
"testing"
"time"
"github.com/iamxvbaba/td/bin"
"github.com/iamxvbaba/td/tg"
)
type epochBlockingTransport struct {
started chan struct{}
release chan struct{}
startedOnce sync.Once
closeOnce sync.Once
}
func newEpochBlockingTransport() *epochBlockingTransport {
return &epochBlockingTransport{started: make(chan struct{}), release: make(chan struct{})}
}
func (t *epochBlockingTransport) Send(context.Context, *bin.Buffer) error {
t.startedOnce.Do(func() { close(t.started) })
<-t.release
return nil
}
func (t *epochBlockingTransport) Recv(context.Context, *bin.Buffer) error { return io.EOF }
func (t *epochBlockingTransport) Close() error {
t.closeOnce.Do(func() { close(t.release) })
return nil
}
func testLayerUpdatesValue(expires int) tg.UpdatesClass {
return &tg.UpdateShort{
Update: &tg.UpdateUserStatus{
UserID: 42,
Status: &tg.UserStatusOnline{
Expires: expires,
},
},
Date: 1_900_000_000,
}
}
func testLayerChannelUpdatesValue(expires int) tg.UpdatesClass {
return &tg.Updates{
Updates: []tg.UpdateClass{
&tg.UpdateUserStatus{
UserID: 42,
Status: &tg.UserStatusOnline{
Expires: expires,
},
},
},
Users: []tg.UserClass{},
Chats: []tg.ChatClass{testLayerChannel()},
Date: 1_900_000_000,
Seq: 1,
}
}
func testConnWithLayerProfile(t *testing.T, profile tg.LayerProfile) *Conn {
t.Helper()
c := &Conn{}
if err := c.FreezeLayerProfile(profile); err != nil {
t.Fatalf("freeze profile %d: %v", profile, err)
}
return c
}
func TestLayerUpdatesFanoutPreparesExactMixedProfiles(t *testing.T) {
fanout, err := newLayerUpdatesFanout(testLayerUpdatesValue(123))
if err != nil {
t.Fatalf("freeze updates: %v", err)
}
for _, profile := range []tg.LayerProfile{tg.LayerProfile225, tg.LayerProfile227, tg.LayerProfile228} {
t.Run(fmt.Sprintf("layer_%d", profile), func(t *testing.T) {
c := testConnWithLayerProfile(t, profile)
encoded, err := fanout.prepareForConn(context.Background(), c)
if err != nil {
t.Fatalf("prepare profile %d: %v", profile, err)
}
if encoded.layer == nil || encoded.layer.profile != profile {
t.Fatalf("binding = %#v, want profile %d", encoded.layer, profile)
}
input := bin.Buffer{Buf: encoded.body}
decoded, err := tg.DecodeLayer(profile, tg.LayerClassUpdatesType(), &input)
if err != nil {
t.Fatalf("decode profile %d: %v", profile, err)
}
if input.Len() != 0 {
t.Fatalf("profile %d left %d trailing bytes", profile, input.Len())
}
short, ok := decoded.(*tg.UpdateShort)
if !ok {
t.Fatalf("decoded %T, want *tg.UpdateShort", decoded)
}
statusUpdate, ok := short.Update.(*tg.UpdateUserStatus)
if !ok {
t.Fatalf("nested update %T, want *tg.UpdateUserStatus", short.Update)
}
status, ok := statusUpdate.Status.(*tg.UserStatusOnline)
if !ok || status.Expires != 123 {
t.Fatalf("decoded status = %#v", statusUpdate.Status)
}
})
}
}
func TestLayerUpdatesFanoutFreezesDefensivelyAndSharesPreparedProfile(t *testing.T) {
value := testLayerUpdatesValue(123)
fanout, err := newLayerUpdatesFanout(value)
if err != nil {
t.Fatalf("freeze updates: %v", err)
}
value.(*tg.UpdateShort).Update.(*tg.UpdateUserStatus).Status.(*tg.UserStatusOnline).Expires = 999
c := testConnWithLayerProfile(t, tg.LayerProfile225)
const workers = 16
prepared := make([]*encodedOutboundMessage, workers)
prepareErrs := make([]error, workers)
var wg sync.WaitGroup
for i := range prepared {
wg.Add(1)
go func(i int) {
defer wg.Done()
prepared[i], prepareErrs[i] = fanout.prepareForConn(context.Background(), c)
}(i)
}
wg.Wait()
for i, prepareErr := range prepareErrs {
if prepareErr != nil {
t.Fatalf("prepare %d: %v", i, prepareErr)
}
}
for i := 1; i < len(prepared); i++ {
if !sameBacking(prepared[i].body, prepared[0].body) {
t.Fatalf("profile preparation %d did not share immutable bytes", i)
}
if prepared[i].layer == prepared[0].layer || prepared[i].layer.epoch != prepared[0].layer.epoch {
t.Fatalf("profile preparation %d did not retain per-target epoch binding", i)
}
}
input := bin.Buffer{Buf: prepared[0].body}
decoded, err := tg.DecodeLayer(tg.LayerProfile225, tg.LayerClassUpdatesType(), &input)
if err != nil {
t.Fatalf("decode frozen value: %v", err)
}
status := decoded.(*tg.UpdateShort).Update.(*tg.UpdateUserStatus).Status.(*tg.UserStatusOnline)
if status.Expires != 123 {
t.Fatalf("frozen value mutated: expires=%d", status.Expires)
}
}
func TestLayerUpdatesEpochBecomesStaleWithoutRetiringProfile(t *testing.T) {
fanout, err := newLayerUpdatesFanout(testLayerUpdatesValue(123))
if err != nil {
t.Fatal(err)
}
c := &Conn{}
if err := c.SeedInheritedLayerProfile(tg.LayerProfile225); err != nil {
t.Fatal(err)
}
encoded, err := fanout.prepareForConn(context.Background(), c)
if err != nil {
t.Fatal(err)
}
oldEpoch := encoded.layer.epoch
if err := c.FreezeLayerProfile(tg.LayerProfile227); err != nil {
t.Fatal(err)
}
if err := validateOutboundLayerBinding(c, encoded); !errors.Is(err, ErrOutboundLayerProfileStale) {
t.Fatalf("old push validation = %v, want ErrOutboundLayerProfileStale", err)
}
state := c.LayerProfileState()
if state.Profile != tg.LayerProfile227 || state.Origin != LayerProfileExplicit || state.Epoch <= oldEpoch {
t.Fatalf("corrected profile state = %#v, old epoch %d", state, oldEpoch)
}
}
func TestRequestBoundLayerResultSurvivesConnectionCorrection(t *testing.T) {
fanout, err := newLayerUpdatesFanout(testLayerUpdatesValue(123))
if err != nil {
t.Fatal(err)
}
c := testConnWithLayerProfile(t, tg.LayerProfile225)
encoded, err := fanout.prepare(context.Background(), tg.LayerProfile225)
if err != nil {
t.Fatal(err)
}
encoded.layer.kind = outboundLayerBindingRequest
if err := c.FreezeLayerProfile(tg.LayerProfile227); err != nil {
t.Fatal(err)
}
if err := validateOutboundLayerBinding(c, encoded); err != nil {
t.Fatalf("request-bound old-profile result rejected after correction: %v", err)
}
}
func TestProfileCorrectionLinearizesAfterStartedPushWrite(t *testing.T) {
transport := newEpochBlockingTransport()
c := newOutboundTestConn(t, transport, nil)
if err := c.FreezeLayerProfile(tg.LayerProfile225); err != nil {
t.Fatal(err)
}
fanout, err := newLayerUpdatesFanout(testLayerUpdatesValue(123))
if err != nil {
t.Fatal(err)
}
encoded, err := fanout.prepareForConn(context.Background(), c)
if err != nil {
t.Fatal(err)
}
if err := c.SendBestEffortEncoded(context.Background(), 0, encoded, 0); err != nil {
t.Fatal(err)
}
select {
case <-transport.started:
case <-time.After(time.Second):
t.Fatal("profile-bound push did not enter physical write")
}
corrected := make(chan error, 1)
go func() { corrected <- c.FreezeLayerProfile(tg.LayerProfile227) }()
select {
case err := <-corrected:
t.Fatalf("profile correction crossed an old-epoch physical write: %v", err)
case <-time.After(50 * time.Millisecond):
}
if err := transport.Close(); err != nil {
t.Fatal(err)
}
select {
case err := <-corrected:
if err != nil {
t.Fatal(err)
}
case <-time.After(time.Second):
t.Fatal("profile correction did not continue after old write completed")
}
if err := validateOutboundLayerBinding(c, encoded); !errors.Is(err, ErrOutboundLayerProfileStale) {
t.Fatalf("completed old push binding = %v, want stale after correction", err)
}
}
func TestStaleLayerPushIsRemovedFromResendTracking(t *testing.T) {
c := &Conn{metrics: NopMetrics{}}
if err := c.FreezeLayerProfile(tg.LayerProfile225); err != nil {
t.Fatal(err)
}
fanout, err := newLayerUpdatesFanout(testLayerUpdatesValue(123))
if err != nil {
t.Fatal(err)
}
encoded, err := fanout.prepareForConn(context.Background(), c)
if err != nil {
t.Fatal(err)
}
if err := c.FreezeLayerProfile(tg.LayerProfile227); err != nil {
t.Fatal(err)
}
frame := &outboundFrame{msgID: 100, body: encoded.body, layer: encoded.layer}
state := &outboundState{
pending: map[int64]*outboundFrame{100: frame},
order: []int64{100},
totalBytes: len(frame.body),
}
if _, err := c.handleOutboundResend(state, context.Background(), []int64{100}); err != nil {
t.Fatal(err)
}
if _, ok := state.pending[100]; ok || frame.body != nil {
t.Fatal("stale profile frame remained resendable")
}
}
func TestOutboundLayerBindingRejectsUnknownAndMismatchedConnections(t *testing.T) {
fanout, err := newLayerUpdatesFanout(testLayerUpdatesValue(123))
if err != nil {
t.Fatalf("freeze updates: %v", err)
}
encoded, err := fanout.prepare(context.Background(), tg.LayerProfile225)
if err != nil {
t.Fatalf("prepare profile 225: %v", err)
}
if _, err := (&Conn{}).buildFrame(context.Background(), 0, nil, encoded); !errors.Is(err, ErrOutboundLayerProfileUnknown) {
t.Fatalf("unknown profile error = %v", err)
}
wrong := testConnWithLayerProfile(t, tg.LayerProfile227)
if _, err := wrong.buildFrame(context.Background(), 0, nil, encoded); !errors.Is(err, ErrOutboundLayerProfileMismatch) {
t.Fatalf("profile mismatch error = %v", err)
}
}
func TestPendingPushReservationAccountsPreparedProfilesOnce(t *testing.T) {
budget := newOutboundTrackedBudget(4096)
if !budget.reserve(100) {
t.Fatal("reserve canonical snapshot")
}
reservation := &pendingPushReservation{budget: budget}
reservation.bytes.Store(100)
reservation.refs.Store(1)
if !reservation.reservePrepared(tg.LayerProfile225, 80) {
t.Fatal("reserve first profile")
}
if !reservation.reservePrepared(tg.LayerProfile225, 80) {
t.Fatal("reuse first profile reservation")
}
if !reservation.reservePrepared(tg.LayerProfile227, 120) {
t.Fatal("reserve second profile")
}
if got := budget.snapshot(); got != 300 {
t.Fatalf("tracked pending bytes = %d, want canonical + unique profiles = 300", got)
}
reservation.release()
if got := budget.snapshot(); got != 0 {
t.Fatalf("tracked pending bytes after final release = %d", got)
}
}

18
internal/mtprotoedge/login_e2e_test.go
View file

@ -12,13 +12,13 @@ import (
"go.uber.org/zap/zaptest"
"github.com/gotd/log/logzap"
"github.com/gotd/td/clock"
"github.com/gotd/td/exchange"
"github.com/gotd/td/session"
"github.com/gotd/td/telegram"
"github.com/gotd/td/telegram/dcs"
"github.com/gotd/td/tg"
"github.com/gotd/td/transport"
"github.com/iamxvbaba/td/clock"
"github.com/iamxvbaba/td/exchange"
"github.com/iamxvbaba/td/session"
"github.com/iamxvbaba/td/telegram"
"github.com/iamxvbaba/td/telegram/dcs"
"github.com/iamxvbaba/td/tg"
"github.com/iamxvbaba/td/transport"
"telesrv/internal/app/account"
"telesrv/internal/app/auth"
@ -93,7 +93,7 @@ func TestLoginRegisterFlow(t *testing.T) {
LangPack: langpack.NewService(langPackStore),
}
router := rpc.New(rpc.Config{DC: dc, IP: tcpAddr.IP.String(), Port: tcpAddr.Port}, deps, zaptest.NewLogger(t), clock.System)
srv := New(Options{Logger: zaptest.NewLogger(t), DC: dc, RSAKey: rsaKey, AuthKeys: authKeyStore, RPC: router})
srv := New(Options{Logger: zaptest.NewLogger(t), DC: dc, RSAKey: rsaKey, AuthKeys: authKeyStore, LayerRPC: router})
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
defer cancel()
@ -314,7 +314,7 @@ func TestPrivateMessageRoundTripFlow(t *testing.T) {
Sessions: activeSessions,
}
router := rpc.New(rpc.Config{DC: dc, IP: tcpAddr.IP.String(), Port: tcpAddr.Port}, deps, zaptest.NewLogger(t), clock.System)
srv := New(Options{Logger: zaptest.NewLogger(t), DC: dc, RSAKey: rsaKey, AuthKeys: authKeyStore, RPC: router, ActiveSessions: activeSessions})
srv := New(Options{Logger: zaptest.NewLogger(t), DC: dc, RSAKey: rsaKey, AuthKeys: authKeyStore, LayerRPC: router, ActiveSessions: activeSessions})
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
defer cancel()

16
internal/mtprotoedge/login_email_e2e_test.go
View file

@ -12,13 +12,13 @@ import (
"go.uber.org/zap/zaptest"
"github.com/gotd/log/logzap"
"github.com/gotd/td/clock"
"github.com/gotd/td/exchange"
"github.com/gotd/td/session"
"github.com/gotd/td/telegram"
"github.com/gotd/td/telegram/dcs"
"github.com/gotd/td/tg"
"github.com/gotd/td/transport"
"github.com/iamxvbaba/td/clock"
"github.com/iamxvbaba/td/exchange"
"github.com/iamxvbaba/td/session"
"github.com/iamxvbaba/td/telegram"
"github.com/iamxvbaba/td/telegram/dcs"
"github.com/iamxvbaba/td/tg"
"github.com/iamxvbaba/td/transport"
"telesrv/internal/app/account"
"telesrv/internal/app/auth"
@ -100,7 +100,7 @@ func TestLoginEmailEndToEnd(t *testing.T) {
Dialogs: dialogs.NewService(dialogStore),
}
router := rpc.New(rpc.Config{DC: dc, IP: tcpAddr.IP.String(), Port: tcpAddr.Port}, deps, zaptest.NewLogger(t), clock.System)
srv := New(Options{Logger: zaptest.NewLogger(t), DC: dc, RSAKey: rsaKey, AuthKeys: authKeyStore, RPC: router})
srv := New(Options{Logger: zaptest.NewLogger(t), DC: dc, RSAKey: rsaKey, AuthKeys: authKeyStore, LayerRPC: router})
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
serveErr := make(chan error, 1)

8
internal/mtprotoedge/metrics_test.go
View file

@ -7,9 +7,9 @@ import (
"go.uber.org/zap/zaptest"
"github.com/gotd/td/clock"
"github.com/gotd/td/proto"
"github.com/gotd/td/tg"
"github.com/iamxvbaba/td/clock"
"github.com/iamxvbaba/td/proto"
"github.com/iamxvbaba/td/tg"
"telesrv/internal/rpc"
)
@ -42,7 +42,7 @@ func TestMetricsHooks(t *testing.T) {
const dc = 2
m := &countingMetrics{}
router := rpc.New(rpc.Config{DC: dc, IP: "127.0.0.1", Port: 2398}, rpc.Deps{}, zaptest.NewLogger(t), clock.System)
addr, pub, _ := startTestServer(t, Options{DC: dc, RPC: router, Metrics: m})
addr, pub, _ := startTestServer(t, Options{DC: dc, legacyRPC: router, Metrics: m})
conn, auth, cipher := dialHandshake(t, addr, dc, pub)
clientMsgID := proto.NewMessageIDGen(time.Now)

6
internal/mtprotoedge/new_session_uid_test.go
View file

@ -4,9 +4,9 @@ import (
"testing"
"time"
"github.com/gotd/td/mt"
"github.com/gotd/td/proto"
"github.com/gotd/td/tg"
"github.com/iamxvbaba/td/mt"
"github.com/iamxvbaba/td/proto"
"github.com/iamxvbaba/td/tg"
)
// TestNewSessionCreatedUniqueIDPerSession 验证两次 session 建立收到的

884
internal/mtprotoedge/outbound.go
View file

File diff suppressed because it is too large Load diff

8
internal/mtprotoedge/outbound_required_control_test.go
View file

@ -10,10 +10,10 @@ import (
"testing"
"time"
"github.com/gotd/td/bin"
"github.com/gotd/td/crypto"
"github.com/gotd/td/mt"
"github.com/gotd/td/proto"
"github.com/iamxvbaba/td/bin"
"github.com/iamxvbaba/td/crypto"
"github.com/iamxvbaba/td/mt"
"github.com/iamxvbaba/td/proto"
)
type gatedRequiredControlTransport struct {

2
internal/mtprotoedge/outbound_scratch.go
View file

@ -4,7 +4,7 @@ import (
"context"
"time"
"github.com/gotd/td/bin"
"github.com/iamxvbaba/td/bin"
)
const (

79
internal/mtprotoedge/outbound_test.go
View file

@ -11,12 +11,12 @@ import (
"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"
"github.com/gotd/td/transport"
"github.com/iamxvbaba/td/bin"
"github.com/iamxvbaba/td/crypto"
"github.com/iamxvbaba/td/mt"
"github.com/iamxvbaba/td/proto"
"github.com/iamxvbaba/td/tg"
"github.com/iamxvbaba/td/transport"
)
type failAfterTransport struct {
@ -28,6 +28,44 @@ type failAfterTransport struct {
last []byte
}
func TestRPCResultReplayAttemptHooksArePhysicalConnectionLocal(t *testing.T) {
const reqMsgID = int64(771)
base := &encodedOutboundMessage{
body: make([]byte, 12),
typeID: proto.ResultTypeID,
reqMsgID: reqMsgID,
delivery: newRPCResultDelivery(reqMsgID),
}
var logical, firstAttempt, secondAttempt atomic.Int32
base.setDeliveryHook(func() { logical.Add(1) })
first, err := cloneRPCResultForRequest(base, reqMsgID, false)
if err != nil {
t.Fatal(err)
}
second, err := cloneRPCResultForRequest(base, reqMsgID, false)
if err != nil {
t.Fatal(err)
}
first.setAttemptDeliveryHook(func() { firstAttempt.Add(1) })
second.setAttemptDeliveryHook(func() { secondAttempt.Add(1) })
first.markDelivered()
deadline := time.Now().Add(time.Second)
for (logical.Load() != 1 || firstAttempt.Load() != 1) && time.Now().Before(deadline) {
time.Sleep(time.Millisecond)
}
if logical.Load() != 1 || firstAttempt.Load() != 1 || secondAttempt.Load() != 0 {
t.Fatalf("first delivery hooks = logical:%d first:%d second:%d", logical.Load(), firstAttempt.Load(), secondAttempt.Load())
}
second.markDelivered()
deadline = time.Now().Add(time.Second)
for secondAttempt.Load() != 1 && time.Now().Before(deadline) {
time.Sleep(time.Millisecond)
}
if logical.Load() != 1 || firstAttempt.Load() != 1 || secondAttempt.Load() != 1 {
t.Fatalf("second delivery hooks = logical:%d first:%d second:%d", logical.Load(), firstAttempt.Load(), secondAttempt.Load())
}
}
type blockingOutboundTransport struct {
started chan struct{}
release chan struct{}
@ -166,7 +204,7 @@ func TestEncodedControlFramesUseIndependentBudgetForQueuedAndPendingLifetime(t *
defer cancel()
// One content frame fills the ordinary body budget and remains pending.
if err := c.Send(ctx, proto.MessageFromServer, &tg.UpdatesTooLong{}); err != nil {
if err := c.SendEncoded(ctx, proto.MessageFromServer, exactTestUpdatesTooLong(t, c)); err != nil {
t.Fatalf("fill body budget: %v", err)
}
first, err := crypto.NewClientCipher(rand.Reader).DecryptFromBuffer(c.key, &bin.Buffer{Buf: tr.lastFrame()})
@ -297,7 +335,7 @@ func TestOutboundScratchAdmissionUsesWriteTimeoutWithoutClosingHealthyConnection
c.writeTimeout = 25 * time.Millisecond
start := time.Now()
err = c.Send(context.Background(), proto.MessageFromServer, &tg.UpdatesTooLong{})
err = c.SendEncoded(context.Background(), proto.MessageFromServer, exactTestUpdatesTooLong(t, c))
elapsed := time.Since(start)
if !errors.Is(err, context.DeadlineExceeded) {
t.Fatalf("scratch admission err = %v, want deadline exceeded", err)
@ -319,7 +357,7 @@ func TestOutboundScratchAdmissionUsesWriteTimeoutWithoutClosingHealthyConnection
pool.release(blocker)
c.writeTimeout = time.Second
if err := c.Send(context.Background(), proto.MessageFromServer, &tg.UpdatesTooLong{}); err != nil {
if err := c.SendEncoded(context.Background(), proto.MessageFromServer, exactTestUpdatesTooLong(t, c)); err != nil {
t.Fatalf("send after scratch capacity returned: %v", err)
}
if got := tr.sends.Load(); got != 1 {
@ -387,6 +425,7 @@ func newOutboundTestConn(t *testing.T, tr transport.Conn, budget *outboundTracke
sessionID: 456,
outboundTrackedBudget: budget,
}
legacyCanonicalTestConn(t, c)
c.startOutbound()
t.Cleanup(c.Close)
return c
@ -506,6 +545,7 @@ func TestOutboundActorSerializesConcurrentSends(t *testing.T) {
clientMsgID := proto.NewMessageIDGen(time.Now)
sendEncrypted(t, conn, cipher, auth, clientMsgID.New(proto.MessageFromClient), &mt.PingRequest{PingID: 1})
collectReplies(t, conn, cipher, auth.AuthKey, mt.MsgsAckTypeID)
freezeActiveTestSessionProfile(t, srv.Conns(), auth.AuthKey.ID, auth.SessionID, tg.LayerProfileCanonical)
srv.Conns().SetReceivesUpdates(auth.SessionID, true)
const sends = 64
@ -556,7 +596,7 @@ func TestOutboundWriteErrorTerminallyClosesWithoutActorDeadlock(t *testing.T) {
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
defer cancel()
if err := c.Send(ctx, proto.MessageFromServer, &tg.UpdatesTooLong{}); err == nil {
if err := c.SendEncoded(ctx, proto.MessageFromServer, exactTestUpdatesTooLong(t, c)); err == nil {
t.Fatal("Send unexpectedly succeeded")
}
select {
@ -567,7 +607,7 @@ func TestOutboundWriteErrorTerminallyClosesWithoutActorDeadlock(t *testing.T) {
if got := tr.closes.Load(); got != 1 {
t.Fatalf("transport closes = %d, want 1", got)
}
if err := c.Send(ctx, proto.MessageFromServer, &tg.UpdatesTooLong{}); !errors.Is(err, ErrConnClosed) {
if err := c.SendEncoded(ctx, proto.MessageFromServer, exactTestUpdatesTooLong(t, c)); !errors.Is(err, ErrConnClosed) {
t.Fatalf("second Send err = %v, want ErrConnClosed", err)
}
if got := tr.sends.Load(); got != 1 {
@ -581,7 +621,7 @@ func TestOutboundResendWriteErrorTerminallyCloses(t *testing.T) {
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
defer cancel()
if err := c.Send(ctx, proto.MessageFromServer, &tg.UpdatesTooLong{}); err != nil {
if err := c.SendEncoded(ctx, proto.MessageFromServer, exactTestUpdatesTooLong(t, c)); err != nil {
t.Fatalf("initial Send: %v", err)
}
data, err := crypto.NewClientCipher(rand.Reader).DecryptFromBuffer(c.key, &bin.Buffer{Buf: tr.lastFrame()})
@ -608,7 +648,7 @@ func TestOutboundTrackedBudgetSharedAcrossConnections(t *testing.T) {
tr2 := &failAfterTransport{}
c1 := newOutboundTestConn(t, tr1, budget)
c2 := newOutboundTestConn(t, tr2, budget)
body := &encodedOutboundMessage{body: make([]byte, 8), typeID: tg.UpdatesTooLongTypeID}
body := exactTestUpdatesEncoded(t, c1, make([]byte, 8))
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
defer cancel()
@ -687,7 +727,7 @@ func TestOutboundGlobalBudgetIncludesQueuedBodies(t *testing.T) {
budget := newOutboundTrackedBudget(24)
tr := newBlockingOutboundTransport()
c := newOutboundTestConn(t, tr, budget)
body := &encodedOutboundMessage{body: make([]byte, 8), typeID: tg.UpdatesTooLongTypeID}
body := exactTestUpdatesEncoded(t, c, make([]byte, 8))
if err := c.SendBestEffortEncoded(context.Background(), proto.MessageFromServer, body, 0); err != nil {
t.Fatalf("enqueue writing body: %v", err)
@ -733,7 +773,7 @@ func TestOutboundOversizedBodyRejectedBeforeEncryption(t *testing.T) {
budget := newOutboundTrackedBudget(64 << 20)
tr := &failAfterTransport{}
c := newOutboundTestConn(t, tr, budget)
body := &encodedOutboundMessage{body: make([]byte, maxOutboundBodyBytes+1), typeID: tg.UpdatesTooLongTypeID}
body := exactTestUpdatesEncoded(t, c, make([]byte, maxOutboundBodyBytes+1))
err := c.SendEncoded(context.Background(), proto.MessageFromServer, body)
if !errors.Is(err, ErrOutboundMessageTooLarge) {
t.Fatalf("oversized outbound err = %v, want ErrOutboundMessageTooLarge", err)
@ -749,7 +789,7 @@ func TestOutboundOversizedBodyRejectedBeforeEncryption(t *testing.T) {
func TestOutboundCloseRaceDrainsEveryProducerReservation(t *testing.T) {
budget := newOutboundTrackedBudget(1 << 20)
c := newOutboundTestConn(t, &failAfterTransport{}, budget)
body := &encodedOutboundMessage{body: make([]byte, 128), typeID: tg.UpdatesTooLongTypeID}
body := exactTestUpdatesEncoded(t, c, make([]byte, 128))
start := make(chan struct{})
var wg sync.WaitGroup
for i := 0; i < 128; i++ {
@ -775,7 +815,7 @@ func TestOutboundTrackedBudgetAckAndCloseReturnExactly(t *testing.T) {
c := newOutboundTestConn(t, tr, budget)
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
defer cancel()
body := &encodedOutboundMessage{body: make([]byte, 12), typeID: tg.UpdatesTooLongTypeID}
body := exactTestUpdatesEncoded(t, c, make([]byte, 12))
if err := c.SendEncoded(ctx, proto.MessageFromServer, body); err != nil {
t.Fatalf("send: %v", err)
}
@ -801,7 +841,7 @@ func TestOutboundTrackedBudgetAckAndCloseReturnExactly(t *testing.T) {
c := newOutboundTestConn(t, &failAfterTransport{}, budget)
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
defer cancel()
body := &encodedOutboundMessage{body: make([]byte, 12), typeID: tg.UpdatesTooLongTypeID}
body := exactTestUpdatesEncoded(t, c, make([]byte, 12))
if err := c.SendEncoded(ctx, proto.MessageFromServer, body); err != nil {
t.Fatalf("send: %v", err)
}
@ -822,7 +862,7 @@ func TestOutboundTrackedBudgetWriteFailureReturnsReservation(t *testing.T) {
c := newOutboundTestConn(t, tr, budget)
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
defer cancel()
body := &encodedOutboundMessage{body: make([]byte, 12), typeID: tg.UpdatesTooLongTypeID}
body := exactTestUpdatesEncoded(t, c, make([]byte, 12))
if err := c.SendEncoded(ctx, proto.MessageFromServer, body); err == nil {
t.Fatal("send unexpectedly succeeded")
}
@ -995,6 +1035,7 @@ func TestOutboundResendAndAckState(t *testing.T) {
clientMsgID := proto.NewMessageIDGen(time.Now)
sendEncrypted(t, conn, cipher, auth, clientMsgID.New(proto.MessageFromClient), &mt.PingRequest{PingID: 1})
collectReplies(t, conn, cipher, auth.AuthKey, mt.MsgsAckTypeID)
freezeActiveTestSessionProfile(t, srv.Conns(), auth.AuthKey.ID, auth.SessionID, tg.LayerProfileCanonical)
srv.Conns().SetReceivesUpdates(auth.SessionID, true)
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)

16
internal/mtprotoedge/passkey_e2e_test.go
View file

@ -14,13 +14,13 @@ import (
"go.uber.org/zap/zaptest"
"github.com/gotd/log/logzap"
"github.com/gotd/td/clock"
"github.com/gotd/td/exchange"
"github.com/gotd/td/session"
"github.com/gotd/td/telegram"
"github.com/gotd/td/telegram/dcs"
"github.com/gotd/td/tg"
"github.com/gotd/td/transport"
"github.com/iamxvbaba/td/clock"
"github.com/iamxvbaba/td/exchange"
"github.com/iamxvbaba/td/session"
"github.com/iamxvbaba/td/telegram"
"github.com/iamxvbaba/td/telegram/dcs"
"github.com/iamxvbaba/td/tg"
"github.com/iamxvbaba/td/transport"
"telesrv/internal/app/account"
"telesrv/internal/app/auth"
@ -115,7 +115,7 @@ func TestPasskeyEndToEnd(t *testing.T) {
Passkey: passkeyService,
}
router := rpc.New(rpc.Config{DC: dc, IP: tcpAddr.IP.String(), Port: tcpAddr.Port}, deps, zaptest.NewLogger(t), clock.System)
srv := New(Options{Logger: zaptest.NewLogger(t), DC: dc, RSAKey: rsaKey, AuthKeys: authKeyStore, RPC: router})
srv := New(Options{Logger: zaptest.NewLogger(t), DC: dc, RSAKey: rsaKey, AuthKeys: authKeyStore, LayerRPC: router})
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
serveErr := make(chan error, 1)

7
internal/mtprotoedge/pending_flush_test.go
View file

@ -5,9 +5,9 @@ import (
"testing"
"time"
"github.com/gotd/td/mt"
"github.com/gotd/td/proto"
"github.com/gotd/td/tg"
"github.com/iamxvbaba/td/mt"
"github.com/iamxvbaba/td/proto"
"github.com/iamxvbaba/td/tg"
)
// TestSetReceivesUpdatesFlushesPendingBeforeActivation 验证置位时先排空暂存推送
@ -23,6 +23,7 @@ func TestSetReceivesUpdatesFlushesPendingBeforeActivation(t *testing.T) {
collectReplies(t, conn, cipher, auth.AuthKey, mt.MsgsAckTypeID)
raw := auth.AuthKey.ID
freezeActiveTestSessionProfile(t, srv.Conns(), raw, auth.SessionID, tg.LayerProfileCanonical)
ctx := context.Background()
// 完全就绪还要求 membership 路由建立（ReceivesUpdatesForAuthKey 的另一半条件）。

39
internal/mtprotoedge/production_compat_test.go
View file

@ -8,8 +8,9 @@ import (
"go.uber.org/zap"
"github.com/gotd/td/bin"
"github.com/gotd/td/proto"
"github.com/iamxvbaba/td/bin"
"github.com/iamxvbaba/td/proto"
"github.com/iamxvbaba/td/tg"
)
var ErrSessionAmbiguous = errors.New("session id is shared by multiple auth keys")
@ -153,7 +154,17 @@ func (m *SessionManager) SetReceivesUpdates(sessionID int64, receives bool) {
}
}
func (m *SessionManager) PushToSession(ctx context.Context, sessionID int64, t proto.MessageType, msg bin.Encoder) error {
func (m *SessionManager) SetLayerProfile(sessionID int64, profile tg.LayerProfile) bool {
m.mu.RLock()
c, _, ok, ambiguous := m.uniqueSessionForTestLocked(sessionID)
m.mu.RUnlock()
if ambiguous || !ok {
return false
}
return c.SeedLayerProfile(profile) == nil
}
func (m *SessionManager) PushToSession(ctx context.Context, sessionID int64, t proto.MessageType, msg tg.UpdatesClass) error {
m.mu.RLock()
c, key, ok, ambiguous := m.uniqueSessionForTestLocked(sessionID)
if ambiguous {
@ -167,20 +178,28 @@ func (m *SessionManager) PushToSession(ctx context.Context, sessionID int64, t p
ready := c.receivesUpdates.Load()
m.mu.RUnlock()
if ready {
return c.Send(ctx, t, msg)
updates, err := newLayerUpdatesFanoutContext(ctx, msg)
if err != nil {
return err
}
encoded, err := updates.prepareForConn(ctx, c)
if err != nil {
return err
}
return c.SendEncoded(ctx, t, encoded)
}
return m.queueOrSendPrepared(ctx, key, t, msg)
}
func (m *SessionManager) PushToUser(ctx context.Context, userID int64, t proto.MessageType, msg bin.Encoder) (int, error) {
func (m *SessionManager) PushToUser(ctx context.Context, userID int64, t proto.MessageType, msg tg.UpdatesClass) (int, error) {
return m.PushToUserExceptAuthKeySession(ctx, userID, [8]byte{}, 0, t, msg)
}
func (m *SessionManager) PushToUserExceptSession(ctx context.Context, userID, excludeSessionID int64, t proto.MessageType, msg bin.Encoder) (int, error) {
func (m *SessionManager) PushToUserExceptSession(ctx context.Context, userID, excludeSessionID int64, t proto.MessageType, msg tg.UpdatesClass) (int, error) {
return m.pushToUser(ctx, userID, nil, excludeSessionID, t, msg)
}
func (m *SessionManager) PushToUserExceptSessionBestEffort(ctx context.Context, userID, excludeSessionID int64, t proto.MessageType, msg bin.Encoder, timeout time.Duration) (int, error) {
func (m *SessionManager) PushToUserExceptSessionBestEffort(ctx context.Context, userID, excludeSessionID int64, t proto.MessageType, msg tg.UpdatesClass, timeout time.Duration) (int, error) {
return m.pushToUserBestEffort(ctx, userID, nil, excludeSessionID, t, msg, timeout)
}
@ -207,8 +226,8 @@ func (m *SessionManager) OnlineChannelIDsAfter(afterChannelID int64, limit int)
return all[start:end]
}
func (m *SessionManager) queueLocked(key sessionKey, t proto.MessageType, msg bin.Encoder) bool {
encoded, reservation, err := m.preparePendingPush(context.Background(), msg)
func (m *SessionManager) queueLocked(key sessionKey, t proto.MessageType, msg tg.UpdatesClass) bool {
updates, reservation, err := m.preparePendingPush(onceLayerUpdatesFanout(context.Background(), msg))
if err != nil {
m.log.Debug("Drop pending push outside byte budget",
zap.String("auth_key_id", sessionKeyLog(key.authKeyID)),
@ -218,7 +237,7 @@ func (m *SessionManager) queueLocked(key sessionKey, t proto.MessageType, msg bi
return false
}
defer reservation.release()
return m.queuePreparedLocked(key, t, encoded, reservation)
return m.queuePreparedLocked(key, t, updates, reservation)
}
func (cs *connState) track(msgID int64, seqNo int32, content bool, state byte) {

23
internal/mtprotoedge/provisional_revocation_test.go
View file

@ -8,11 +8,11 @@ import (
"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"
"github.com/iamxvbaba/td/bin"
"github.com/iamxvbaba/td/mt"
"github.com/iamxvbaba/td/proto"
"github.com/iamxvbaba/td/proto/codec"
"github.com/iamxvbaba/td/tg"
"telesrv/internal/store"
"telesrv/internal/store/memory"
@ -59,7 +59,7 @@ 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})
addr, pub, _ := startTestServer(t, Options{DC: dc, AuthKeys: keys, legacyRPC: handler})
conn, auth, cipher := dialHandshake(t, addr, dc, pub)
ids := proto.NewMessageIDGen(time.Now)
firstID := ids.New(proto.MessageFromClient)
@ -104,7 +104,7 @@ func TestActivationFinalAuthKeyCheckRunsAfterClaim(t *testing.T) {
}
}()
handler := &admissionCountingRPC{}
addr, pub, srv := startTestServer(t, Options{DC: dc, AuthKeys: keys, RPC: handler})
addr, pub, srv := startTestServer(t, Options{DC: dc, AuthKeys: keys, legacyRPC: handler})
conn, auth, cipher := dialHandshake(t, addr, dc, pub)
msgID := proto.NewMessageIDGen(time.Now).New(proto.MessageFromClient)
@ -148,7 +148,7 @@ func TestActivationFinalAuthKeyCheckRunsAfterClaim(t *testing.T) {
func TestBadSaltProvisionalCannotReactivateDeletedAuthKey(t *testing.T) {
const dc = 2
handler := &admissionCountingRPC{}
addr, pub, srv := startTestServer(t, Options{DC: dc, RPC: handler})
addr, pub, srv := startTestServer(t, Options{DC: dc, legacyRPC: handler})
provisional, auth, cipher := dialHandshake(t, addr, dc, pub)
ids := proto.NewMessageIDGen(time.Now)
reqMsgID := ids.New(proto.MessageFromClient)
@ -174,12 +174,13 @@ func TestBadSaltProvisionalCannotReactivateDeletedAuthKey(t *testing.T) {
destroyer := dialTransportOnly(t, addr)
destroySessionID := auth.SessionID ^ 1
destroyBody := encodeClientMessageBodyForTest(t, &destroyAuthKeyRequest{})
destroyReqMsgID := ids.New(proto.MessageFromClient)
sendEncryptedWithSessionSaltAndSeq(
t, destroyer, cipher, auth, destroySessionID, auth.ServerSalt,
ids.New(proto.MessageFromClient), 1, destroyBody,
destroyReqMsgID, 1, destroyBody,
)
destroyReplies := collectReplies(t, destroyer, cipher, auth.AuthKey, destroyAuthKeyOkTypeID)
mustHave(t, destroyReplies, destroyAuthKeyOkTypeID, "destroy_auth_key_ok")
destroyReplies := collectReplies(t, destroyer, cipher, auth.AuthKey, proto.ResultTypeID)
assertDestroyAuthKeyRPCResult(t, mustHave(t, destroyReplies, proto.ResultTypeID, "destroy_auth_key rpc_result"), destroyReqMsgID, destroyAuthKeyOkTypeID)
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)
}

4
internal/mtprotoedge/quick_ack_deadline_test.go
View file

@ -5,8 +5,8 @@ import (
"testing"
"time"
"github.com/gotd/td/bin"
"github.com/gotd/td/crypto"
"github.com/iamxvbaba/td/bin"
"github.com/iamxvbaba/td/crypto"
)
type quickAckDeadlineProbe struct {

95
internal/mtprotoedge/rpc_admission_tracker.go Normal file
View file

@ -0,0 +1,95 @@
package mtprotoedge
import (
"math"
"sync"
"sync/atomic"
)
const rpcAdmissionTrackerShards = 64
// rpcAdmissionTracker retains exactly the admission sequences whose unique
// owner can still publish/complete. Allocation holds the barrier from sequence
// CAS through shard registration, so a stable floor scan can never observe an
// allocated-but-not-yet-active gap.
type rpcAdmissionTracker struct {
allocationBarrier sync.RWMutex
shards [rpcAdmissionTrackerShards]rpcAdmissionTrackerShard
}
type rpcAdmissionTrackerShard struct {
mu sync.Mutex
active map[uint64]struct{}
}
func (t *rpcAdmissionTracker) allocateAndRegister(next *atomic.Uint64) (uint64, error) {
if t == nil || next == nil {
return 0, ErrRPCResultFlightInvalid
}
t.allocationBarrier.RLock()
defer t.allocationBarrier.RUnlock()
var sequence uint64
for {
current := next.Load()
if current == math.MaxUint64 {
return 0, ErrRPCAdmissionSeqExhausted
}
sequence = current + 1
if next.CompareAndSwap(current, sequence) {
break
}
}
shard := &t.shards[sequence&(rpcAdmissionTrackerShards-1)]
shard.mu.Lock()
if shard.active == nil {
shard.active = make(map[uint64]struct{})
}
shard.active[sequence] = struct{}{}
shard.mu.Unlock()
return sequence, nil
}
func (t *rpcAdmissionTracker) retire(sequence uint64) {
if t == nil || sequence == 0 {
return
}
shard := &t.shards[sequence&(rpcAdmissionTrackerShards-1)]
shard.mu.Lock()
if _, ok := shard.active[sequence]; !ok {
shard.mu.Unlock()
panic("mtprotoedge: rpc admission sequence retired more than once")
}
delete(shard.active, sequence)
shard.mu.Unlock()
}
// stableSafeFloor returns the lowest sequence which can still publish, or one
// past the last allocated sequence when no owner remains. The short exclusive
// barrier blocks only admission sequence allocation/registration; handler,
// encoding and delivery stay fully concurrent.
func (t *rpcAdmissionTracker) stableSafeFloor(next *atomic.Uint64) uint64 {
if t == nil || next == nil {
return 0
}
t.allocationBarrier.Lock()
defer t.allocationBarrier.Unlock()
var minimum uint64
for index := range t.shards {
shard := &t.shards[index]
shard.mu.Lock()
for sequence := range shard.active {
if minimum == 0 || sequence < minimum {
minimum = sequence
}
}
shard.mu.Unlock()
}
if minimum != 0 {
return minimum
}
last := next.Load()
if last == math.MaxUint64 {
return math.MaxUint64
}
return last + 1
}

467
internal/mtprotoedge/rpc_delivery_coordinator_test.go Normal file
View file

@ -0,0 +1,467 @@
package mtprotoedge
import (
"context"
"errors"
"sync"
"sync/atomic"
"testing"
"time"
)
func TestRPCDeliveryCoordinatorWaitsFromClaimedThroughDone(t *testing.T) {
coordinator := newRPCResultDelivery(1).coordinator
started := make(chan struct{})
release := make(chan struct{})
coordinator.setHook(func() {
close(started)
<-release
})
claim, err := coordinator.claimReplayDeliveredHook(context.Background(), false)
if err != nil || claim == nil {
t.Fatalf("initial hook claim = %p, err=%v", claim, err)
}
if got := coordinator.hookState(); got != rpcResultDeliveryHookClaimed {
t.Fatalf("hook state after claim = %d, want claimed", got)
}
type claimResult struct {
claim *rpcResultDeliveryHookClaim
err error
}
waited := make(chan claimResult, 1)
go func() {
other, waitErr := coordinator.claimReplayDeliveredHook(context.Background(), false)
waited <- claimResult{claim: other, err: waitErr}
}()
select {
case result := <-waited:
t.Fatalf("second replay passed a claimed hook: claim=%p err=%v", result.claim, result.err)
case <-time.After(20 * time.Millisecond):
}
go claim.run()
select {
case <-started:
case <-time.After(time.Second):
t.Fatal("claimed hook did not enter in-progress")
}
if got := coordinator.hookState(); got != rpcResultDeliveryHookInProgress {
t.Fatalf("hook state while callback blocked = %d, want in-progress", got)
}
select {
case result := <-waited:
t.Fatalf("second replay passed an in-progress hook: claim=%p err=%v", result.claim, result.err)
case <-time.After(20 * time.Millisecond):
}
close(release)
select {
case result := <-waited:
if result.err != nil || result.claim != nil {
t.Fatalf("wait after done = claim:%p err:%v", result.claim, result.err)
}
case <-time.After(time.Second):
t.Fatal("completion did not notify the waiting replay")
}
if got := coordinator.hookState(); got != rpcResultDeliveryHookDone {
t.Fatalf("terminal hook state = %d, want done", got)
}
}
func TestRPCDeliveryCoordinatorAbandonedClaimCannotRunAfterReacquire(t *testing.T) {
coordinator := newRPCResultDelivery(1).coordinator
var calls atomic.Int32
coordinator.setHook(func() { calls.Add(1) })
stale, err := coordinator.claimReplayDeliveredHook(context.Background(), false)
if err != nil || stale == nil || !stale.abandon() {
t.Fatalf("abandon initial claim = %p err=%v", stale, err)
}
current, err := coordinator.claimReplayDeliveredHook(context.Background(), false)
if err != nil || current == nil {
t.Fatalf("reacquire hook = %p err=%v", current, err)
}
stale.run()
if got := calls.Load(); got != 0 {
t.Fatalf("stale claim calls = %d, want 0", got)
}
if got := coordinator.hookState(); got != rpcResultDeliveryHookClaimed {
t.Fatalf("stale claim changed current state to %d", got)
}
current.run()
if got := calls.Load(); got != 1 {
t.Fatalf("current claim calls = %d, want 1", got)
}
if got := coordinator.hookState(); got != rpcResultDeliveryHookDone {
t.Fatalf("terminal hook state = %d, want done", got)
}
}
func TestBoundedRPCReplayRestoreDoesNotRetainWorkerOrBarrier(t *testing.T) {
s := New(Options{})
c := &Conn{metrics: NopMetrics{}}
finishBarrier := c.beginRPCReplayRestore()
coordinator := newRPCResultDelivery(1).coordinator
var hookCalls atomic.Int32
coordinator.setHook(func() { hookCalls.Add(1) })
claim, err := coordinator.claimReplayDeliveredHook(context.Background(), false)
if err != nil || claim == nil {
t.Fatalf("claim hook = %p err=%v", claim, err)
}
replacementStarted := make(chan struct{})
releaseReplacement := make(chan struct{})
t.Cleanup(func() {
select {
case <-releaseReplacement:
default:
close(releaseReplacement)
}
})
ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
defer cancel()
started := time.Now()
restoreResult := make(chan error, 1)
go func() {
restoreResult <- s.runBoundedRPCReplayRestore(ctx, c, "non-cooperative replacement", claim, func() error {
close(replacementStarted)
<-releaseReplacement // Deliberately ignores ctx.
return nil
})
}()
select {
case <-replacementStarted:
case <-time.After(time.Second):
t.Fatal("replacement callback never started")
}
select {
case err = <-restoreResult:
case <-time.After(time.Second):
t.Fatal("bounded restore retained its caller past the watchdog")
}
if !errors.Is(err, context.DeadlineExceeded) {
t.Fatalf("bounded restore error = %v, want deadline", err)
}
if elapsed := time.Since(started); elapsed > 500*time.Millisecond {
t.Fatalf("non-cooperative restore retained caller for %v", elapsed)
}
finishBarrier()
if !c.isRetired() {
t.Fatal("timed-out restore did not fence its physical generation")
}
c.rpcMu.Lock()
pendingBarriers := c.rpcReplayRestores
c.rpcMu.Unlock()
if pendingBarriers != 0 {
t.Fatalf("timed-out restore retained %d scheduler barriers", pendingBarriers)
}
if got := coordinator.hookState(); got != rpcResultDeliveryHookPending {
t.Fatalf("not-yet-started timed-out hook state = %d, want pending", got)
}
// A replacement replay can now acquire the logical hook. The old blocked
// goroutine holds a stale token and must not execute it when it eventually
// returns.
retry, err := coordinator.claimReplayDeliveredHook(context.Background(), false)
if err != nil || retry == nil {
t.Fatalf("retry hook claim = %p err=%v", retry, err)
}
retry.run()
close(releaseReplacement)
deadline := time.Now().Add(time.Second)
for len(rpcReplayRestoreSlots) != 0 && time.Now().Before(deadline) {
time.Sleep(time.Millisecond)
}
if got := hookCalls.Load(); got != 1 {
t.Fatalf("logical hook calls after stale restore return = %d, want 1", got)
}
}
func TestBoundedRPCReplayRestoreKeepsInProgressHookAtMostOnce(t *testing.T) {
s := New(Options{})
c := &Conn{metrics: NopMetrics{}}
finishBarrier := c.beginRPCReplayRestore()
coordinator := newRPCResultDelivery(1).coordinator
hookStarted := make(chan struct{})
releaseHook := make(chan struct{})
t.Cleanup(func() {
select {
case <-releaseHook:
default:
close(releaseHook)
}
})
coordinator.setHook(func() {
close(hookStarted)
<-releaseHook // Deliberately ignores the restore deadline.
})
claim, err := coordinator.claimReplayDeliveredHook(context.Background(), false)
if err != nil || claim == nil {
t.Fatalf("claim hook = %p err=%v", claim, err)
}
ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
defer cancel()
restoreResult := make(chan error, 1)
go func() {
restoreResult <- s.runBoundedRPCReplayRestore(ctx, c, "non-cooperative logical hook", claim, nil)
}()
select {
case <-hookStarted:
case <-time.After(time.Second):
t.Fatal("logical hook never entered in-progress")
}
select {
case err = <-restoreResult:
case <-time.After(time.Second):
t.Fatal("non-cooperative logical hook retained its caller past the watchdog")
}
if !errors.Is(err, context.DeadlineExceeded) {
t.Fatalf("bounded logical-hook error = %v, want deadline", err)
}
if got := coordinator.hookState(); got != rpcResultDeliveryHookInProgress {
t.Fatalf("timed-out logical hook state = %d, want in-progress", got)
}
finishBarrier()
if !c.isRetired() {
t.Fatal("timed-out logical hook did not fence its physical generation")
}
c.rpcMu.Lock()
pendingBarriers := c.rpcReplayRestores
c.rpcMu.Unlock()
if pendingBarriers != 0 {
t.Fatalf("timed-out logical hook retained %d scheduler barriers", pendingBarriers)
}
waitCtx, waitCancel := context.WithTimeout(context.Background(), 20*time.Millisecond)
defer waitCancel()
if retry, retryErr := coordinator.claimReplayDeliveredHook(waitCtx, false); retry != nil ||
!errors.Is(retryErr, context.DeadlineExceeded) {
t.Fatalf("retry during in-progress hook = claim:%p err:%v", retry, retryErr)
}
close(releaseHook)
deadline := time.Now().Add(time.Second)
for coordinator.hookState() != rpcResultDeliveryHookDone && time.Now().Before(deadline) {
time.Sleep(time.Millisecond)
}
if got := coordinator.hookState(); got != rpcResultDeliveryHookDone {
t.Fatalf("released logical hook state = %d, want done", got)
}
}
func TestCachedReplacementReplayBarrierWaitsForLogicalHookDone(t *testing.T) {
s := New(Options{})
const reqMsgID = int64(71001)
encoded := encodedRPCResultForPriorityTest(reqMsgID, 0)
encoded.delivery = newRPCResultDelivery(reqMsgID)
var order atomic.Int32
hookStarted := make(chan struct{})
releaseHook := make(chan struct{})
t.Cleanup(func() {
select {
case <-releaseHook:
default:
close(releaseHook)
}
})
encoded.setDeliveryHook(func() {
close(hookStarted)
<-releaseHook
if !order.CompareAndSwap(1, 2) {
panic("logical hook did not follow first replacement restore")
}
})
firstConn := newOutboundTestConn(t, &collectingSessionTransport{}, newOutboundTrackedBudget(1<<20))
firstResult := make(chan error, 1)
go func() {
firstResult <- s.sendCachedRPCResultWithHook(context.Background(), firstConn, encoded, func() error {
if !order.CompareAndSwap(0, 1) {
return errors.New("first replacement restore ran out of order")
}
return nil
})
}()
select {
case <-hookStarted:
case <-time.After(time.Second):
t.Fatal("first replacement did not enter logical hook")
}
if got := encoded.delivery.coordinator.hookState(); got != rpcResultDeliveryHookInProgress {
t.Fatalf("shared hook state = %d, want in-progress", got)
}
secondTransport := &collectingSessionTransport{}
secondConn := newOutboundTestConn(t, secondTransport, newOutboundTrackedBudget(1<<20))
secondReplacement := make(chan struct{})
secondResult := make(chan error, 1)
go func() {
secondResult <- s.sendCachedRPCResultWithHook(context.Background(), secondConn, encoded, func() error {
if !order.CompareAndSwap(2, 3) {
return errors.New("second replacement restore passed logical hook completion")
}
close(secondReplacement)
return nil
})
}()
deadline := time.Now().Add(time.Second)
for len(secondTransport.snapshot()) == 0 && time.Now().Before(deadline) {
time.Sleep(time.Millisecond)
}
if len(secondTransport.snapshot()) == 0 {
t.Fatal("second replacement did not physically deliver cached result")
}
secondConn.rpcMu.Lock()
secondBarriers := secondConn.rpcReplayRestores
secondConn.rpcMu.Unlock()
if secondBarriers != 1 {
t.Fatalf("second replacement barriers while hook in progress = %d, want 1", secondBarriers)
}
select {
case err := <-secondResult:
t.Fatalf("second replacement returned before hook done: %v", err)
case <-time.After(20 * time.Millisecond):
}
select {
case <-secondReplacement:
t.Fatal("second replacement metadata ran before shared hook done")
default:
}
close(releaseHook)
for name, result := range map[string]<-chan error{
"first": firstResult, "second": secondResult,
} {
select {
case err := <-result:
if err != nil {
t.Fatalf("%s replacement replay: %v", name, err)
}
case <-time.After(time.Second):
t.Fatalf("%s replacement replay did not finish", name)
}
}
if got := order.Load(); got != 3 {
t.Fatalf("replacement/logical terminal order = %d, want 3", got)
}
secondConn.rpcMu.Lock()
secondBarriers = secondConn.rpcReplayRestores
secondConn.rpcMu.Unlock()
if secondBarriers != 0 {
t.Fatalf("second replacement retained %d barriers after hook done", secondBarriers)
}
}
func TestRPCRewrapWatchdogCoversCommittedRestoreWithoutBlockingTimerOrWorker(t *testing.T) {
s := New(Options{})
c := &Conn{metrics: NopMetrics{}}
const reqMsgID = int64(72001)
encoded := encodedRPCResultForPriorityTest(reqMsgID, 0)
encoded.delivery = newRPCResultDelivery(reqMsgID)
var logicalCalls atomic.Int32
encoded.setDeliveryHook(func() { logicalCalls.Add(1) })
replacementStarted := make(chan struct{})
releaseReplacement := make(chan struct{})
t.Cleanup(func() {
select {
case <-releaseReplacement:
default:
close(releaseReplacement)
}
})
alias := &rpcRewrapAlias{
conn: c, newReqID: reqMsgID, method: "help.getConfig",
afterSuccessfulDelivery: func() error {
close(replacementStarted)
<-releaseReplacement // Deliberately ignores every deadline.
return nil
},
}
alias.executionOK.Store(true)
alias.beginReplayRestore()
var (
once sync.Once
jobs chan rpcRewrapDeliveryJob
)
workerDone := make(chan error, 1)
watchdogDone := make(chan struct{})
job := rpcRewrapDeliveryJob{
deadline: time.Now().Add(80 * time.Millisecond),
run: func(control *rpcRewrapDeliveryControl, _ time.Time) {
if !control.commit() {
workerDone <- errors.New("commit lost before restore")
return
}
restoreCtx, cancel := context.WithTimeout(context.Background(), 250*time.Millisecond)
defer cancel()
workerDone <- s.completeDeliveredRPCRewrapResult(
restoreCtx, alias, encoded, "committed restore watchdog regression",
)
},
fail: func(error) {
c.fenceUndeliveredRPCResult()
alias.releaseReplayRestoreBarrier()
close(watchdogDone)
},
}
if !scheduleRPCRewrapJob(job, &once, &jobs, 1, 2) {
t.Fatal("schedule committed restore job")
}
select {
case <-replacementStarted:
case <-time.After(time.Second):
t.Fatal("committed replacement restore never started")
}
select {
case <-watchdogDone:
case <-time.After(500 * time.Millisecond):
t.Fatal("committed-state watchdog did not return promptly")
}
if !c.isRetired() {
t.Fatal("committed-state watchdog did not fence connection")
}
c.rpcMu.Lock()
pendingBarriers := c.rpcReplayRestores
c.rpcMu.Unlock()
if pendingBarriers != 0 {
t.Fatalf("watchdog retained %d scheduler barriers", pendingBarriers)
}
select {
case err := <-workerDone:
if !errors.Is(err, context.DeadlineExceeded) {
t.Fatalf("bounded committed restore error = %v, want deadline", err)
}
case <-time.After(time.Second):
t.Fatal("non-cooperative restore permanently retained rewrap worker")
}
followingRan := make(chan struct{})
if !scheduleRPCRewrapJob(rpcRewrapDeliveryJob{
deadline: time.Now().Add(time.Second),
run: func(*rpcRewrapDeliveryControl, time.Time) { close(followingRan) },
}, &once, &jobs, 1, 2) {
t.Fatal("schedule following rewrap job")
}
select {
case <-followingRan:
case <-time.After(time.Second):
t.Fatal("rewrap worker did not accept following job after restore timeout")
}
if got := encoded.delivery.coordinator.hookState(); got != rpcResultDeliveryHookPending {
t.Fatalf("timed-out pre-hook restore state = %d, want pending", got)
}
if got := logicalCalls.Load(); got != 0 {
t.Fatalf("logical hook ran behind non-cooperative replacement = %d", got)
}
close(releaseReplacement)
deadline := time.Now().Add(time.Second)
for len(rpcReplayRestoreSlots) != 0 && time.Now().Before(deadline) {
time.Sleep(time.Millisecond)
}
}

111
internal/mtprotoedge/rpc_delivery_hook_executor_test.go Normal file
View file

@ -0,0 +1,111 @@
package mtprotoedge
import (
"errors"
"sync/atomic"
"testing"
"time"
)
func hookTestMessage(reqMsgID int64, coordinator *rpcResultDeliveryCoordinator, fn func()) *encodedOutboundMessage {
msg := &encodedOutboundMessage{delivery: newRPCResultDelivery(reqMsgID, coordinator)}
msg.setDeliveryHook(fn)
return msg
}
func TestRPCDeliveryHookExecutorBoundsAdmissionWithoutBlockingDelivery(t *testing.T) {
executor := newRPCDeliveryHookExecutor(1, 1)
started := make(chan struct{})
release := make(chan struct{})
first := hookTestMessage(1, nil, func() {
close(started)
<-release
})
if err := first.prepareDeliveryHook(executor); err != nil {
t.Fatalf("reserve first hook: %v", err)
}
delivered := make(chan struct{})
go func() {
first.markDelivered()
close(delivered)
}()
select {
case <-delivered:
case <-time.After(100 * time.Millisecond):
t.Fatal("physical delivery blocked on hook execution")
}
select {
case <-started:
case <-time.After(time.Second):
t.Fatal("first hook did not start")
}
second := hookTestMessage(2, nil, func() {})
if err := second.prepareDeliveryHook(executor); !errors.Is(err, ErrRPCDeliveryHookCapacity) {
t.Fatalf("second reservation = %v, want capacity error", err)
}
close(release)
deadline := time.Now().Add(time.Second)
for len(executor.slots) != 0 && time.Now().Before(deadline) {
time.Sleep(time.Millisecond)
}
if got := len(executor.slots); got != 0 {
t.Fatalf("executor retained %d capacity slots", got)
}
}
func TestRPCDeliveryHookExecutorIsolatesPanicsAndContinues(t *testing.T) {
executor := newRPCDeliveryHookExecutor(1, 2)
first := hookTestMessage(1, nil, func() { panic("hook boom") })
done := make(chan struct{})
second := hookTestMessage(2, nil, func() { close(done) })
if err := first.prepareDeliveryHook(executor); err != nil {
t.Fatalf("reserve panic hook: %v", err)
}
if err := second.prepareDeliveryHook(executor); err != nil {
t.Fatalf("reserve following hook: %v", err)
}
first.markDelivered()
second.markDelivered()
select {
case <-done:
case <-time.After(time.Second):
t.Fatal("worker stopped after a hook panic")
}
deadline := time.Now().Add(time.Second)
for executor.panics.Load() != 1 && time.Now().Before(deadline) {
time.Sleep(time.Millisecond)
}
if got := executor.panics.Load(); got != 1 {
t.Fatalf("recorded hook panics = %d, want 1", got)
}
}
func TestEquivalentRPCDeliveryAttemptsShareExactlyOnceCoordinator(t *testing.T) {
executor := newRPCDeliveryHookExecutor(1, 2)
var calls atomic.Int32
first := hookTestMessage(11, nil, func() { calls.Add(1) })
second := hookTestMessage(22, first.delivery.coordinator, nil)
if err := first.prepareDeliveryHook(executor); err != nil {
t.Fatalf("reserve first attempt: %v", err)
}
if err := second.prepareDeliveryHook(executor); err != nil {
t.Fatalf("reserve equivalent attempt: %v", err)
}
first.markDelivered()
second.markDelivered()
deadline := time.Now().Add(time.Second)
for calls.Load() != 1 && time.Now().Before(deadline) {
time.Sleep(time.Millisecond)
}
if got := calls.Load(); got != 1 {
t.Fatalf("equivalent delivery hooks = %d, want 1", got)
}
deadline = time.Now().Add(time.Second)
for len(executor.slots) != 0 && time.Now().Before(deadline) {
time.Sleep(time.Millisecond)
}
if got := len(executor.slots); got != 0 {
t.Fatalf("equivalent attempts leaked %d tickets", got)
}
}

460
internal/mtprotoedge/rpc_flight.go
View file

@ -5,6 +5,8 @@ import (
"errors"
"sync"
"sync/atomic"
"github.com/iamxvbaba/td/tg"
)
const rpcResultFlightDefaultMaxPending = 8192
@ -12,10 +14,49 @@ 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")
ErrRPCResultFlightCapacity = errors.New("mtproto rpc result in-flight capacity exhausted")
ErrRPCResultSubscriberCapacity = errors.New("mtproto rpc result subscriber capacity exhausted")
ErrRPCResultFlightInvalid = errors.New("mtproto rpc result in-flight claim is invalid")
ErrRPCResultIdentityMismatch = errors.New("mtproto rpc result request identity mismatch")
ErrRPCAdmissionSeqExhausted = errors.New("mtproto rpc admission sequence exhausted")
)
// rpcResultIdentityMismatchError carries the winner's immutable admission
// profile from inside the cache shard critical section. A replacement Conn can
// re-decode the same naked body under that grammar even if the winner aborts
// immediately after the mismatch is returned.
type rpcResultIdentityMismatchError struct {
profile tg.LayerProfile
hasProfile bool
}
func (e *rpcResultIdentityMismatchError) Error() string { return ErrRPCResultIdentityMismatch.Error() }
func (e *rpcResultIdentityMismatchError) Is(target error) bool {
return target == ErrRPCResultIdentityMismatch
}
func identityMismatch(identity rpcResultRequestIdentity) error {
return &rpcResultIdentityMismatchError{profile: identity.profile, hasProfile: identity.valid && identity.profile != 0}
}
type rpcResultRequestIdentity struct {
exact tg.LayerPreparedCallIdentity
// profile is retained separately because LayerPreparedCallIdentity is opaque.
// It lets a same-msg_id replay be re-admitted with the original request
// grammar after the session default has moved to another Layer.
profile tg.LayerProfile
valid bool
}
func (i rpcResultRequestIdentity) matches(requested rpcResultRequestIdentity) bool {
if !requested.valid {
// Legacy service/test callers carry no API request identity and preserve
// the historical cache lookup behavior. Exact callers must always match.
return true
}
return i.valid && i.exact == requested.exact
}
type rpcResultAcquireState uint8
const (
@ -32,20 +73,35 @@ const (
// - 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
state rpcResultAcquireState
admissionSeq uint64
encoded *encodedOutboundMessage
waiter *rpcResultWaiter
owner *rpcResultOwnerLease
executionKnown bool
executionOK bool
}
// rpcResultFlight is not part of the completed cache LRU/TTL lifecycle. Its
// rpcResultFlight is not part of the completed cache 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
subscribers []func(*encodedOutboundMessage, bool)
done chan struct{}
encoded *encodedOutboundMessage
ok bool
subscribers []func(*encodedOutboundMessage, bool)
executionDone bool
executionOK bool
executionSubscribers []func(bool)
// subscriberSlots counts callbacks retained by this pending flight. Result
// and execution callbacks are charged independently; a replay alias installs
// both atomically so a capacity failure cannot leave half an alias behind.
subscriberSlots int
identity rpcResultRequestIdentity
admissionSeq uint64
// reservation owns one entry and at least one byte at global, raw-auth and
// session scopes. Put transfers it to a result/tombstone; Abort releases it.
reservation *rpcResultBudgetReservation
}
type rpcResultWaiter struct {
@ -86,31 +142,107 @@ func (w *rpcResultWaiter) Wait(ctx context.Context) (encoded *encodedOutboundMes
// occupying an RPC worker. The callback is invoked after the cache shard lock is
// released; it must remain non-blocking.
func (w *rpcResultWaiter) Subscribe(fn func(*encodedOutboundMessage, bool)) error {
if w == nil || w.cache == nil || w.flight == nil || fn == nil {
if fn == nil {
return ErrRPCResultFlightInvalid
}
return w.subscribe(fn, nil)
}
// SubscribeExecution registers a non-blocking callback for the terminal
// business outcome. This is deliberately separate from physical rpc_result
// delivery: invokeAfter* ordering depends on handler completion and must not
// occupy a shared RPC worker while a socket write is pending.
func (w *rpcResultWaiter) SubscribeExecution(fn func(bool)) error {
if fn == nil {
return ErrRPCResultFlightInvalid
}
return w.subscribe(nil, fn)
}
// SubscribeResultAndExecution atomically installs both halves of one pending
// replay alias. Either both callbacks are retained, or neither is; this avoids
// an execution-only closure surviving when the result subscriber hits a hard
// capacity limit.
func (w *rpcResultWaiter) SubscribeResultAndExecution(
resultFn func(*encodedOutboundMessage, bool),
executionFn func(bool),
) error {
if resultFn == nil || executionFn == nil {
return ErrRPCResultFlightInvalid
}
return w.subscribe(resultFn, executionFn)
}
func (w *rpcResultWaiter) subscribe(
resultFn func(*encodedOutboundMessage, bool),
executionFn func(bool),
) error {
if w == nil || w.cache == nil || w.flight == nil || (resultFn == nil && executionFn == nil) {
return ErrRPCResultFlightInvalid
}
s := w.cache.shard(w.key)
var (
encoded *encodedOutboundMessage
ok bool
ready bool
encoded *encodedOutboundMessage
resultOK bool
resultReady = resultFn == nil
executionOK bool
executionReady = executionFn == nil
)
s.mu.Lock()
if flight, exists := s.pending[w.key]; exists && flight == w.flight {
flight.subscribers = append(flight.subscribers, fn)
slots := 0
if resultFn != nil {
slots++
}
if executionFn != nil && !flight.executionDone {
slots++
}
if slots > 0 {
if flight.subscriberSlots > w.cache.subscriberPerFlight-slots ||
!w.cache.subscriberBudget.reserve(w.key, slots) {
s.mu.Unlock()
return ErrRPCResultSubscriberCapacity
}
flight.subscriberSlots += slots
}
if resultFn != nil {
flight.subscribers = append(flight.subscribers, resultFn)
}
if executionFn != nil {
if flight.executionDone {
executionOK, executionReady = flight.executionOK, true
} else {
flight.executionSubscribers = append(flight.executionSubscribers, executionFn)
}
}
s.mu.Unlock()
if executionReady && executionFn != nil {
executionFn(executionOK)
}
return nil
}
select {
case <-w.flight.done:
encoded, ok, ready = w.flight.encoded, w.flight.ok, true
default:
if resultFn != nil {
select {
case <-w.flight.done:
encoded, resultOK, resultReady = w.flight.encoded, w.flight.ok, true
default:
}
}
if executionFn != nil && w.flight.executionDone {
executionOK, executionReady = w.flight.executionOK, true
}
s.mu.Unlock()
if !ready {
// Atomic late subscription also means no callback runs unless every requested
// terminal state is available.
if !resultReady || !executionReady {
return ErrRPCResultFlightInvalid
}
fn(encoded, ok)
if executionFn != nil {
executionFn(executionOK)
}
if resultFn != nil {
resultFn(encoded, resultOK)
}
return nil
}
@ -175,6 +307,33 @@ func (l *rpcResultOwnerLease) Delivery() *rpcResultDelivery {
return l.delivery
}
// CompleteExecution publishes the handler outcome exactly once while this
// lease still owns the flight. success=false includes RPC errors, internal
// failures and dependency failures. Delivery/cache completion remains a
// separate later transition.
func (l *rpcResultOwnerLease) CompleteExecution(success bool) bool {
if l == nil || l.cache == nil || l.flight == nil {
return false
}
s := l.cache.shard(l.key)
s.mu.Lock()
flight, ok := s.pending[l.key]
if !ok || flight != l.flight || flight.executionDone {
s.mu.Unlock()
return false
}
flight.executionDone = true
flight.executionOK = success
subscribers := append([]func(bool){}, flight.executionSubscribers...)
flight.executionSubscribers = nil
l.cache.releaseFlightSubscriberSlotsLocked(l.key, flight, len(subscribers))
s.mu.Unlock()
for _, subscriber := range subscribers {
subscriber(success)
}
return true
}
// HandOff transfers completion responsibility from the inbound RPC task to an
// already-admitted egress operation. The egress terminal callback must resolve
// the flight through Put on both successful delivery and fenced failure.
@ -216,9 +375,23 @@ func (l *rpcResultOwnerLease) Abort() bool {
return false
}
delete(s.pending, l.key)
if flight.reservation != nil {
flight.reservation.release()
flight.reservation = nil
}
l.cache.flightLimit.release()
l.cache.activeAdmissions.retire(flight.admissionSeq)
subscribers := append([]func(*encodedOutboundMessage, bool){}, flight.subscribers...)
flight.subscribers = nil
executionSubscribers := append([]func(bool){}, flight.executionSubscribers...)
flight.executionSubscribers = nil
l.cache.releaseFlightSubscriberSlotsLocked(
l.key, flight, len(subscribers)+len(executionSubscribers),
)
if !flight.executionDone {
flight.executionDone = true
flight.executionOK = false
}
close(flight.done)
s.mu.Unlock()
l.hookMu.Lock()
@ -231,6 +404,9 @@ func (l *rpcResultOwnerLease) Abort() bool {
for _, subscriber := range subscribers {
subscriber(nil, false)
}
for _, subscriber := range executionSubscribers {
subscriber(false)
}
return true
}
@ -240,25 +416,39 @@ type rpcResultFlightLimit struct {
}
func (l *rpcResultFlightLimit) reserve() bool {
return l.reserveN(1)
}
func (l *rpcResultFlightLimit) reserveN(delta int64) bool {
if l == nil || l.max <= 0 {
return false
}
if delta <= 0 {
return false
}
for {
used := l.used.Load()
if used >= l.max {
if used < 0 || used > l.max-delta {
return false
}
if l.used.CompareAndSwap(used, used+1) {
if l.used.CompareAndSwap(used, used+delta) {
return true
}
}
}
func (l *rpcResultFlightLimit) release() {
l.releaseN(1)
}
func (l *rpcResultFlightLimit) releaseN(delta int64) {
if l == nil {
return
}
if remaining := l.used.Add(-1); remaining < 0 {
if delta <= 0 {
panic("mtproto rpc result counter release must be positive")
}
if remaining := l.used.Add(-delta); 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.
@ -275,45 +465,159 @@ func (l *rpcResultFlightLimit) snapshot() int64 {
// 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.
// lifecycle from completed cache TTL but reserve the same global/auth/session
// ownership that Put later transfers to a completed result or tombstone.
func (c *rpcResultCache) Acquire(authKeyID [8]byte, sessionID, reqMsgID int64) (rpcResultAcquire, error) {
return c.acquire(authKeyID, sessionID, reqMsgID, rpcResultRequestIdentity{})
}
func (c *rpcResultCache) AcquireIdentified(
authKeyID [8]byte,
sessionID, reqMsgID int64,
identity tg.LayerPreparedCallIdentity,
) (rpcResultAcquire, error) {
return c.acquire(authKeyID, sessionID, reqMsgID, rpcResultRequestIdentity{exact: identity, valid: true})
}
// AcquireLayerIdentified is the production exact-RPC claim. In addition to the
// immutable full request identity it retains the admission profile required to
// decode a later same-msg_id naked replay under its original grammar.
func (c *rpcResultCache) AcquireLayerIdentified(
authKeyID [8]byte,
sessionID, reqMsgID int64,
profile tg.LayerProfile,
identity tg.LayerPreparedCallIdentity,
) (rpcResultAcquire, error) {
return c.acquire(authKeyID, sessionID, reqMsgID, rpcResultRequestIdentity{
exact: identity, profile: profile, valid: true,
})
}
// ExactAdmissionProfile returns the immutable profile of an existing exact
// owner/result. It does not create or join a flight. Callers still perform
// AcquireLayerIdentified after decode, which atomically rejects a same-msg_id
// body change by comparing the full prepared identity.
func (c *rpcResultCache) ExactAdmissionProfile(authKeyID [8]byte, sessionID, reqMsgID int64) (tg.LayerProfile, bool) {
if c == nil || reqMsgID == 0 {
return 0, false
}
key := rpcResultCacheKey{authKeyID: authKeyID, sessionID: sessionID, reqMsgID: reqMsgID}
s := c.shard(key)
now := s.now()
s.mu.Lock()
defer s.mu.Unlock()
if elem := s.byKey[key]; elem != nil {
entry := elem.Value.(*rpcResultCacheEntry)
if entry.expiresAt.After(now) {
if entry.identity.valid && entry.identity.profile != 0 {
return entry.identity.profile, true
}
return 0, false
}
s.removeElement(elem)
}
if flight := s.pending[key]; flight != nil && flight.identity.valid && flight.identity.profile != 0 {
return flight.identity.profile, true
}
return 0, false
}
func (c *rpcResultCache) acquire(
authKeyID [8]byte,
sessionID, reqMsgID int64,
identity rpcResultRequestIdentity,
) (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()
reclaimedExpired := false
for {
now := s.now()
s.mu.Lock()
s.expireLocked(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
if elem, ok := s.byKey[key]; ok {
entry := elem.Value.(*rpcResultCacheEntry)
if !entry.identity.matches(identity) {
s.mu.Unlock()
return rpcResultAcquire{}, identityMismatch(entry.identity)
}
if entry.capacity || entry.encoded == nil {
s.mu.Unlock()
return rpcResultAcquire{}, ErrRPCResultFlightCapacity
}
result := rpcResultAcquire{
state: rpcResultAcquireCompleted, admissionSeq: entry.admissionSeq, encoded: entry.encoded,
executionKnown: entry.executionKnown, executionOK: entry.executionOK,
}
s.mu.Unlock()
return result, nil
}
s.removeElement(elem)
}
if flight, ok := s.pending[key]; ok {
if flight, ok := s.pending[key]; ok {
if !flight.identity.matches(identity) {
s.mu.Unlock()
return rpcResultAcquire{}, identityMismatch(flight.identity)
}
result := rpcResultAcquire{
state: rpcResultAcquirePending,
admissionSeq: flight.admissionSeq,
waiter: &rpcResultWaiter{cache: c, key: key, flight: flight},
}
s.mu.Unlock()
return result, nil
}
if s.maxEntries > 0 && len(s.byKey)+len(s.pending) >= s.maxEntries {
s.mu.Unlock()
return rpcResultAcquire{}, ErrRPCResultFlightCapacity
}
if !c.flightLimit.reserve() {
s.mu.Unlock()
return rpcResultAcquire{}, ErrRPCResultFlightCapacity
}
reservation := c.fairBudget.reserveOwner(key)
if reservation == nil {
c.flightLimit.release()
s.mu.Unlock()
if reclaimedExpired {
return rpcResultAcquire{}, ErrRPCResultFlightCapacity
}
// Expired rows in another full-key shard may be the only consumers at
// the global, auth or session scope. Reap once, then retry every identity
// and capacity check because another goroutine may have won this key.
c.expireCompletedResults()
reclaimedExpired = true
continue
}
var admissionSeq uint64
if identity.valid {
var err error
admissionSeq, err = c.activeAdmissions.allocateAndRegister(&c.nextAdmissionSeq)
if err != nil {
reservation.release()
c.flightLimit.release()
s.mu.Unlock()
return rpcResultAcquire{}, err
}
}
flight := &rpcResultFlight{
done: make(chan struct{}), identity: identity, admissionSeq: admissionSeq,
reservation: reservation,
}
if s.pending == nil {
s.pending = make(map[rpcResultCacheKey]*rpcResultFlight)
}
s.pending[key] = flight
s.mu.Unlock()
return rpcResultAcquire{
state: rpcResultAcquirePending,
waiter: &rpcResultWaiter{cache: c, key: key, flight: flight},
state: rpcResultAcquireOwner,
admissionSeq: admissionSeq,
owner: &rpcResultOwnerLease{
cache: c, key: key, flight: flight, delivery: newRPCResultDelivery(reqMsgID),
},
}, 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, delivery: newRPCResultDelivery(reqMsgID),
},
}, nil
}
// completeRPCResultFlightLocked publishes encoded to the current owner claim.
@ -322,20 +626,62 @@ func (c *rpcResultCache) completeRPCResultFlightLocked(
s *rpcResultCacheShard,
key rpcResultCacheKey,
encoded *encodedOutboundMessage,
) []func(*encodedOutboundMessage, bool) {
) (
[]func(*encodedOutboundMessage, bool),
[]func(bool),
bool,
) {
if c == nil || s == nil || encoded == nil {
return nil
return nil, nil, false
}
flight, ok := s.pending[key]
if !ok {
return nil
return nil, nil, false
}
delete(s.pending, key)
if flight.reservation != nil {
flight.reservation.releasePending()
// The completed entry now owns the same reservation.
flight.reservation = nil
}
flight.encoded = encoded
flight.ok = true
c.flightLimit.release()
c.activeAdmissions.retire(flight.admissionSeq)
subscribers := append([]func(*encodedOutboundMessage, bool){}, flight.subscribers...)
flight.subscribers = nil
executionSubscribers := append([]func(bool){}, flight.executionSubscribers...)
flight.executionSubscribers = nil
executionOK := flight.executionOK
if !flight.executionDone {
// A result without an explicit handler-completion proof cannot satisfy an
// invokeAfter dependency. Production completes execution before Put; this
// branch is the conservative terminal cleanup for defensive callers.
flight.executionDone = true
flight.executionOK = false
executionOK = false
}
c.releaseFlightSubscriberSlotsLocked(
key, flight, len(subscribers)+len(executionSubscribers),
)
close(flight.done)
return subscribers
return subscribers, executionSubscribers, executionOK
}
// releaseFlightSubscriberSlotsLocked releases callbacks detached from a flight.
// The caller holds that flight's cache shard lock, preserving the only lock
// order used by subscription: shard -> subscriber budget.
func (c *rpcResultCache) releaseFlightSubscriberSlotsLocked(
key rpcResultCacheKey,
flight *rpcResultFlight,
slots int,
) {
if slots == 0 {
return
}
if c == nil || flight == nil || slots < 0 || flight.subscriberSlots < slots {
panic("mtproto rpc result subscriber slot underflow")
}
flight.subscriberSlots -= slots
c.subscriberBudget.release(key, slots)
}

499
internal/mtprotoedge/rpc_flight_test.go
View file

@ -6,12 +6,471 @@ import (
"sync"
"testing"
"time"
"github.com/iamxvbaba/td/bin"
"github.com/iamxvbaba/td/tg"
)
func rpcFlightTestAuthID(seed byte) [8]byte {
return [8]byte{seed, seed + 1, seed + 2, seed + 3}
}
func rpcFlightExactIdentity(t *testing.T, profile tg.LayerProfile, request bin.Encoder) tg.LayerPreparedCallIdentity {
t.Helper()
var body bin.Buffer
if err := request.Encode(&body); err != nil {
t.Fatalf("encode exact request: %v", err)
}
admitted, err := tg.NewServerDispatcher(nil).AdmitLayer(profile, &body)
if err != nil {
t.Fatalf("admit exact request: %v", err)
}
if body.Len() != 0 {
t.Fatalf("exact admission left %d bytes", body.Len())
}
return admitted.Prepared().Identity()
}
func newRPCResultSubscriberTestCache(global, auth, session, perFlight int) *rpcResultCache {
return newRPCResultCacheWithFairCapacity(time.Now, rpcResultCacheCapacity{
maxPending: 64,
maxPendingPerAuth: 64,
globalMaxBytes: rpcResultCacheMaxBytes,
globalMaxEntries: rpcResultCacheMaxEntries,
authMaxBytes: rpcResultCacheAuthMaxBytes,
authMaxEntries: rpcResultCacheAuthMaxEntries,
sessionMaxBytes: rpcResultCacheSessionMaxBytes,
sessionMaxEntries: rpcResultCacheSessionMaxEntries,
subscriberMaxGlobal: global,
subscriberMaxAuth: auth,
subscriberMaxSession: session,
subscriberMaxPerFlight: perFlight,
})
}
func TestRPCResultFlightSubscriberPairCapacityFailureIsAtomic(t *testing.T) {
cache := newRPCResultSubscriberTestCache(8, 8, 8, 1)
authKeyID := rpcFlightTestAuthID(70)
claim, err := cache.Acquire(authKeyID, 70, 700)
if err != nil || claim.owner == nil {
t.Fatalf("Acquire owner: claim=%#v err=%v", claim, err)
}
var resultCalls, executionCalls int
err = claim.owner.Waiter().SubscribeResultAndExecution(
func(*encodedOutboundMessage, bool) { resultCalls++ },
func(bool) { executionCalls++ },
)
if !errors.Is(err, ErrRPCResultSubscriberCapacity) {
t.Fatalf("pair subscription err=%v, want %v", err, ErrRPCResultSubscriberCapacity)
}
s := cache.shard(rpcResultCacheKey{authKeyID: authKeyID, sessionID: 70, reqMsgID: 700})
s.mu.Lock()
if got := claim.owner.flight.subscriberSlots; got != 0 {
t.Fatalf("failed pair retained %d subscriber slots", got)
}
if got := len(claim.owner.flight.subscribers) + len(claim.owner.flight.executionSubscribers); got != 0 {
t.Fatalf("failed pair retained %d callbacks", got)
}
s.mu.Unlock()
claim.owner.CompleteExecution(true)
cache.Put(authKeyID, 70, 700, &encodedOutboundMessage{body: []byte{1}, reqMsgID: 700})
if resultCalls != 0 || executionCalls != 0 {
t.Fatalf("failed pair callbacks ran: result=%d execution=%d", resultCalls, executionCalls)
}
if got := cache.subscriberBudget.global.snapshot(); got != 0 {
t.Fatalf("subscriber global usage=%d after failed pair", got)
}
}
func TestRPCResultFlightSubscriberBudgetsIsolateSessionAndAuth(t *testing.T) {
cache := newRPCResultSubscriberTestCache(3, 2, 1, 4)
authA := rpcFlightTestAuthID(71)
authB := rpcFlightTestAuthID(72)
type ownerKey struct {
auth [8]byte
session int64
msgID int64
owner *rpcResultOwnerLease
}
acquire := func(auth [8]byte, session, msgID int64) ownerKey {
t.Helper()
claim, err := cache.Acquire(auth, session, msgID)
if err != nil || claim.owner == nil {
t.Fatalf("Acquire(%d,%d): claim=%#v err=%v", session, msgID, claim, err)
}
return ownerKey{auth: auth, session: session, msgID: msgID, owner: claim.owner}
}
subscribe := func(owner ownerKey) error {
return owner.owner.Waiter().Subscribe(func(*encodedOutboundMessage, bool) {})
}
a1 := acquire(authA, 1, 101)
a2 := acquire(authA, 2, 102)
a3 := acquire(authA, 3, 103)
b1 := acquire(authB, 4, 104)
b2 := acquire(authB, 5, 105)
if err := subscribe(a1); err != nil {
t.Fatalf("first session subscriber: %v", err)
}
if err := subscribe(a1); !errors.Is(err, ErrRPCResultSubscriberCapacity) {
t.Fatalf("same session overflow err=%v", err)
}
if err := subscribe(a2); err != nil {
t.Fatalf("second session subscriber: %v", err)
}
if err := subscribe(a3); !errors.Is(err, ErrRPCResultSubscriberCapacity) {
t.Fatalf("same auth overflow err=%v", err)
}
if err := subscribe(b1); err != nil {
t.Fatalf("other auth subscriber: %v", err)
}
if err := subscribe(b2); !errors.Is(err, ErrRPCResultSubscriberCapacity) {
t.Fatalf("global overflow err=%v", err)
}
if got := cache.subscriberBudget.authSnapshot(authA); got != 2 {
t.Fatalf("auth A usage=%d, want 2", got)
}
if got := cache.subscriberBudget.authSnapshot(authB); got != 1 {
t.Fatalf("auth B usage=%d, want 1", got)
}
if got := cache.subscriberBudget.global.snapshot(); got != 3 {
t.Fatalf("global usage=%d, want 3", got)
}
for _, owner := range []ownerKey{a1, a2, a3, b1, b2} {
owner.owner.Abort()
}
if got := cache.subscriberBudget.global.snapshot(); got != 0 {
t.Fatalf("global usage=%d after aborts", got)
}
if got := cache.subscriberBudget.authSnapshot(authA); got != 0 {
t.Fatalf("auth A usage=%d after aborts", got)
}
}
func TestRPCResultFlightSubscriberSlotsReleasePerTerminalHalf(t *testing.T) {
cache := newRPCResultSubscriberTestCache(4, 4, 4, 4)
authKeyID := rpcFlightTestAuthID(73)
claim, err := cache.Acquire(authKeyID, 73, 730)
if err != nil || claim.owner == nil {
t.Fatalf("Acquire owner: claim=%#v err=%v", claim, err)
}
result := make(chan bool, 1)
execution := make(chan bool, 1)
if err := claim.owner.Waiter().SubscribeResultAndExecution(
func(_ *encodedOutboundMessage, ok bool) { result <- ok },
func(ok bool) { execution <- ok },
); err != nil {
t.Fatalf("pair subscription: %v", err)
}
if got := cache.subscriberBudget.sessionSnapshot(authKeyID, 73); got != 2 {
t.Fatalf("initial subscriber usage=%d, want 2", got)
}
if !claim.owner.CompleteExecution(true) {
t.Fatal("CompleteExecution lost")
}
if ok := <-execution; !ok {
t.Fatal("execution callback reported failure")
}
if got := cache.subscriberBudget.sessionSnapshot(authKeyID, 73); got != 1 {
t.Fatalf("post-execution subscriber usage=%d, want 1", got)
}
cache.Put(authKeyID, 73, 730, &encodedOutboundMessage{body: []byte{1}, reqMsgID: 730})
if ok := <-result; !ok {
t.Fatal("result callback reported failure")
}
if got := cache.subscriberBudget.sessionSnapshot(authKeyID, 73); got != 0 {
t.Fatalf("post-result subscriber usage=%d, want 0", got)
}
}
func TestRPCResultFlightRepeatedReplayJoinsStayBoundedAndPutCleansExecution(t *testing.T) {
cache := newRPCResultSubscriberTestCache(2, 2, 2, 2)
authKeyID := rpcFlightTestAuthID(74)
claim, err := cache.Acquire(authKeyID, 74, 740)
if err != nil || claim.owner == nil {
t.Fatalf("Acquire owner: claim=%#v err=%v", claim, err)
}
var resultCalls, executionCalls int
if err := claim.owner.Waiter().SubscribeResultAndExecution(
func(*encodedOutboundMessage, bool) { resultCalls++ },
func(success bool) {
if success {
t.Error("Put without execution proof reported dependency success")
}
executionCalls++
},
); err != nil {
t.Fatalf("first pair: %v", err)
}
for i := 0; i < 100; i++ {
err := claim.owner.Waiter().SubscribeResultAndExecution(
func(*encodedOutboundMessage, bool) { resultCalls++ },
func(bool) { executionCalls++ },
)
if !errors.Is(err, ErrRPCResultSubscriberCapacity) {
t.Fatalf("join %d err=%v, want capacity", i, err)
}
}
cache.Put(authKeyID, 74, 740, &encodedOutboundMessage{body: []byte{1}, reqMsgID: 740})
if resultCalls != 1 || executionCalls != 1 {
t.Fatalf("terminal callback counts result=%d execution=%d", resultCalls, executionCalls)
}
if got := cache.subscriberBudget.global.snapshot(); got != 0 {
t.Fatalf("global usage=%d after Put", got)
}
}
func TestRPCResultFlightExactIdentityGuardsPendingAndCompletedReuse(t *testing.T) {
cache := newRPCResultCacheWithFlightLimit(time.Now, 2)
authKeyID := rpcFlightTestAuthID(90)
firstIdentity := rpcFlightExactIdentity(t, tg.LayerProfile225, &tg.HelpGetConfigRequest{})
otherIdentity := rpcFlightExactIdentity(t, tg.LayerProfile225, &tg.HelpGetNearestDCRequest{})
owner, err := cache.AcquireIdentified(authKeyID, 90, 900, firstIdentity)
if err != nil || owner.state != rpcResultAcquireOwner || owner.owner == nil {
t.Fatalf("exact owner Acquire = state:%d err:%v", owner.state, err)
}
if _, err := cache.AcquireIdentified(authKeyID, 90, 900, otherIdentity); !errors.Is(err, ErrRPCResultIdentityMismatch) {
t.Fatalf("pending mismatched Acquire err = %v, want %v", err, ErrRPCResultIdentityMismatch)
}
same, err := cache.AcquireIdentified(authKeyID, 90, 900, firstIdentity)
if err != nil || same.state != rpcResultAcquirePending || same.waiter == nil {
t.Fatalf("pending matched Acquire = state:%d err:%v", same.state, err)
}
want := &encodedOutboundMessage{body: []byte{1, 2, 3, 4}, reqMsgID: 900}
cache.Put(authKeyID, 90, 900, want)
if _, err := cache.AcquireIdentified(authKeyID, 90, 900, otherIdentity); !errors.Is(err, ErrRPCResultIdentityMismatch) {
t.Fatalf("completed mismatched Acquire err = %v, want %v", err, ErrRPCResultIdentityMismatch)
}
completed, err := cache.AcquireIdentified(authKeyID, 90, 900, firstIdentity)
if err != nil || completed.state != rpcResultAcquireCompleted || completed.encoded != want {
t.Fatalf("completed matched Acquire = state:%d encoded:%p err:%v", completed.state, completed.encoded, err)
}
if encoded, ok, waitErr := same.waiter.Wait(context.Background()); waitErr != nil || !ok || encoded != want {
t.Fatalf("matched waiter = encoded:%p ok:%v err:%v", encoded, ok, waitErr)
}
}
func TestRPCResultFlightAdmissionSequenceAllocatedOnceAndReplayed(t *testing.T) {
cache := newRPCResultCacheWithFlightLimit(time.Now, 4)
authKeyID := rpcFlightTestAuthID(89)
identity := rpcFlightExactIdentity(t, tg.LayerProfile225, &tg.HelpGetConfigRequest{})
owner, err := cache.AcquireLayerIdentified(authKeyID, 89, 890, tg.LayerProfile225, identity)
if err != nil || owner.state != rpcResultAcquireOwner || owner.owner == nil || owner.admissionSeq == 0 {
t.Fatalf("owner = state:%d seq:%d err:%v", owner.state, owner.admissionSeq, err)
}
pending, err := cache.AcquireLayerIdentified(authKeyID, 89, 890, tg.LayerProfile225, identity)
if err != nil || pending.state != rpcResultAcquirePending || pending.admissionSeq != owner.admissionSeq {
t.Fatalf("pending = state:%d seq:%d err:%v, want seq:%d", pending.state, pending.admissionSeq, err, owner.admissionSeq)
}
owner.owner.CompleteExecution(true)
encoded := &encodedOutboundMessage{body: []byte{1}, reqMsgID: 890}
cache.Put(authKeyID, 89, 890, encoded)
completed, err := cache.AcquireLayerIdentified(authKeyID, 89, 890, tg.LayerProfile225, identity)
if err != nil || completed.state != rpcResultAcquireCompleted || completed.admissionSeq != owner.admissionSeq {
t.Fatalf("completed = state:%d seq:%d err:%v, want seq:%d", completed.state, completed.admissionSeq, err, owner.admissionSeq)
}
second, err := cache.AcquireLayerIdentified(authKeyID, 89, 894, tg.LayerProfile225, identity)
if err != nil || second.admissionSeq <= owner.admissionSeq {
t.Fatalf("second owner seq=%d err=%v, want > %d", second.admissionSeq, err, owner.admissionSeq)
}
second.owner.Abort()
legacy, err := cache.Acquire(authKeyID, 89, 898)
if err != nil || legacy.admissionSeq != 0 {
t.Fatalf("legacy admission seq=%d err=%v, want 0", legacy.admissionSeq, err)
}
legacy.owner.Abort()
}
func TestRPCAdmissionSafeFloorTracksOwnersUntilPutOrAbort(t *testing.T) {
cache := newRPCResultCacheWithFlightLimit(time.Now, 4)
authKeyID := rpcFlightTestAuthID(86)
identity := rpcFlightExactIdentity(t, tg.LayerProfile225, &tg.HelpGetConfigRequest{})
first, err := cache.AcquireLayerIdentified(authKeyID, 86, 860, tg.LayerProfile225, identity)
if err != nil || first.owner == nil {
t.Fatalf("first owner err=%v", err)
}
second, err := cache.AcquireLayerIdentified(authKeyID, 86, 864, tg.LayerProfile225, identity)
if err != nil || second.owner == nil {
t.Fatalf("second owner err=%v", err)
}
if floor := cache.stableAdmissionSafeFloor(); floor != first.admissionSeq {
t.Fatalf("two-owner safe floor=%d, want %d", floor, first.admissionSeq)
}
if !first.owner.Abort() {
t.Fatal("first owner abort failed")
}
if floor := cache.stableAdmissionSafeFloor(); floor != second.admissionSeq {
t.Fatalf("post-abort safe floor=%d, want %d", floor, second.admissionSeq)
}
second.owner.CompleteExecution(true)
cache.Put(authKeyID, 86, 864, &encodedOutboundMessage{body: []byte{1}, reqMsgID: 864})
if floor := cache.stableAdmissionSafeFloor(); floor != second.admissionSeq+1 {
t.Fatalf("terminal safe floor=%d, want %d", floor, second.admissionSeq+1)
}
}
func TestRPCAdmissionSequenceExhaustionCannotWrap(t *testing.T) {
cache := newRPCResultCacheWithFlightLimit(time.Now, 2)
cache.nextAdmissionSeq.Store(^uint64(0) - 1)
authKeyID := rpcFlightTestAuthID(85)
identity := rpcFlightExactIdentity(t, tg.LayerProfile225, &tg.HelpGetConfigRequest{})
last, err := cache.AcquireLayerIdentified(authKeyID, 85, 850, tg.LayerProfile225, identity)
if err != nil || last.admissionSeq != ^uint64(0) || last.owner == nil {
t.Fatalf("last sequence=%d owner:%v err=%v", last.admissionSeq, last.owner != nil, err)
}
if _, err := cache.AcquireLayerIdentified(authKeyID, 85, 854, tg.LayerProfile225, identity); !errors.Is(err, ErrRPCAdmissionSeqExhausted) {
t.Fatalf("post-max allocation err=%v, want %v", err, ErrRPCAdmissionSeqExhausted)
}
if got := cache.nextAdmissionSeq.Load(); got != ^uint64(0) {
t.Fatalf("exhausted sequence wrapped to %d", got)
}
last.owner.Abort()
}
func TestRPCIdentityMismatchCarriesWinnerProfileAcrossAbort(t *testing.T) {
cache := newRPCResultCacheWithFlightLimit(time.Now, 2)
authKeyID := rpcFlightTestAuthID(88)
request := &tg.MessagesGetHistoryRequest{Peer: &tg.InputPeerSelf{}, Limit: 1}
winnerIdentity := rpcFlightExactIdentity(t, tg.LayerProfile225, request)
loserIdentity := rpcFlightExactIdentity(t, tg.LayerProfile227, request)
winner, err := cache.AcquireLayerIdentified(authKeyID, 88, 880, tg.LayerProfile225, winnerIdentity)
if err != nil || winner.owner == nil {
t.Fatalf("winner owner err=%v", err)
}
_, err = cache.AcquireLayerIdentified(authKeyID, 88, 880, tg.LayerProfile227, loserIdentity)
var mismatch *rpcResultIdentityMismatchError
if !errors.As(err, &mismatch) || !mismatch.hasProfile || mismatch.profile != tg.LayerProfile225 {
t.Fatalf("mismatch = %#v err=%v", mismatch, err)
}
if !winner.owner.Abort() {
t.Fatal("winner abort failed")
}
replacement, err := cache.AcquireLayerIdentified(authKeyID, 88, 880, mismatch.profile, winnerIdentity)
if err != nil || replacement.state != rpcResultAcquireOwner || replacement.owner == nil {
t.Fatalf("replacement under retained winner profile = state:%d err:%v", replacement.state, err)
}
replacement.owner.Abort()
}
func TestRPCAdmissionProfileHintSurvivesCompletedEvictionWindow(t *testing.T) {
now := time.Unix(1_900_000_000, 0)
cache := newRPCResultCacheWithFlightLimit(func() time.Time { return now }, 2)
authKeyID := rpcFlightTestAuthID(87)
identity := rpcFlightExactIdentity(t, tg.LayerProfile225, &tg.MessagesGetHistoryRequest{
Peer: &tg.InputPeerSelf{}, Limit: 1,
})
claim, err := cache.AcquireLayerIdentified(authKeyID, 87, 870, tg.LayerProfile225, identity)
if err != nil || claim.owner == nil {
t.Fatalf("owner err=%v", err)
}
claim.owner.CompleteExecution(true)
cache.Put(authKeyID, 87, 870, &encodedOutboundMessage{body: []byte{1}, reqMsgID: 870})
profile, ok := cache.ExactAdmissionProfile(authKeyID, 87, 870)
if !ok || profile != tg.LayerProfile225 {
t.Fatalf("profile hint = (%d,%v)", profile, ok)
}
// Admission already copied the hint into its local decoder cursor. Expiry
// between that probe and the atomic claim must not make it fall back to the
// connection's newer default; it simply becomes a fresh owner under 225.
now = now.Add(rpcResultCacheTTL + time.Second)
replacement, err := cache.AcquireLayerIdentified(authKeyID, 87, 870, profile, identity)
if err != nil || replacement.state != rpcResultAcquireOwner || replacement.owner == nil {
t.Fatalf("post-eviction owner = state:%d err:%v", replacement.state, err)
}
replacement.owner.Abort()
}
func TestRPCInvariantIdentityDoesNotExposeCanonicalProfileHint(t *testing.T) {
cache := newRPCResultCacheWithFlightLimit(time.Now, 2)
authKeyID := rpcFlightTestAuthID(84)
identity := rpcFlightExactIdentity(t, tg.LayerProfile227, &tg.AuthBindTempAuthKeyRequest{
PermAuthKeyID: 1, Nonce: 2, ExpiresAt: 3, EncryptedMessage: []byte("bind"),
})
claim, err := cache.AcquireLayerIdentified(authKeyID, 84, 840, 0, identity)
if err != nil || claim.owner == nil {
t.Fatalf("invariant owner err=%v", err)
}
if profile, ok := cache.ExactAdmissionProfile(authKeyID, 84, 840); ok || profile != 0 {
t.Fatalf("pending invariant profile hint=(%d,%v), want absent", profile, ok)
}
claim.owner.CompleteExecution(true)
cache.Put(authKeyID, 84, 840, &encodedOutboundMessage{body: []byte{1}, reqMsgID: 840})
if profile, ok := cache.ExactAdmissionProfile(authKeyID, 84, 840); ok || profile != 0 {
t.Fatalf("completed invariant profile hint=(%d,%v), want absent", profile, ok)
}
}
func TestRPCResultExecutionCompletionIsExactlyOnceAndDurable(t *testing.T) {
cache := newRPCResultCacheWithFlightLimit(time.Now, 2)
authKeyID := rpcFlightTestAuthID(91)
claim, err := cache.Acquire(authKeyID, 91, 910)
if err != nil || claim.state != rpcResultAcquireOwner || claim.owner == nil {
t.Fatalf("owner Acquire = state:%d err:%v", claim.state, err)
}
waiter := claim.owner.Waiter()
results := make(chan bool, 2)
if err := waiter.SubscribeExecution(func(success bool) { results <- success }); err != nil {
t.Fatal(err)
}
if !claim.owner.CompleteExecution(true) {
t.Fatal("first execution completion lost")
}
if claim.owner.CompleteExecution(false) {
t.Fatal("contradictory second execution completion won")
}
select {
case success := <-results:
if !success {
t.Fatal("execution callback reported failure")
}
case <-time.After(time.Second):
t.Fatal("execution callback did not run")
}
select {
case success := <-results:
t.Fatalf("execution callback ran twice: %v", success)
default:
}
want := &encodedOutboundMessage{body: []byte{9, 1, 0, 0}, reqMsgID: 910}
cache.Put(authKeyID, 91, 910, want)
dependency, ok := cache.ObserveDependency(authKeyID, 91, 910)
if !ok || !dependency.completed || !dependency.success || dependency.waiter != nil {
t.Fatalf("completed dependency = %#v ok:%v", dependency, ok)
}
late := make(chan bool, 1)
if err := waiter.SubscribeExecution(func(success bool) { late <- success }); err != nil {
t.Fatal(err)
}
if success := <-late; !success {
t.Fatal("late execution subscriber lost durable success")
}
}
func TestRPCResultExecutionAbortPublishesFailure(t *testing.T) {
cache := newRPCResultCacheWithFlightLimit(time.Now, 1)
authKeyID := rpcFlightTestAuthID(92)
claim, err := cache.Acquire(authKeyID, 92, 920)
if err != nil || claim.owner == nil {
t.Fatalf("owner Acquire err = %v", err)
}
result := make(chan bool, 1)
if err := claim.owner.Waiter().SubscribeExecution(func(success bool) { result <- success }); err != nil {
t.Fatal(err)
}
if !claim.owner.Abort() {
t.Fatal("Abort lost")
}
if success := <-result; success {
t.Fatal("aborted execution reported success")
}
if _, ok := cache.ObserveDependency(authKeyID, 92, 920); ok {
t.Fatal("aborted flight remained observable as a completed dependency")
}
}
func TestRPCResultFlightConcurrentAcquireHasUniqueOwner(t *testing.T) {
const callers = 64
cache := newRPCResultCacheWithFlightLimit(time.Now, callers)
@ -180,7 +639,6 @@ func TestRPCResultFlightCompletedCachePressureDoesNotEvictPending(t *testing.T)
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)}})
@ -251,7 +709,7 @@ func TestRPCResultFlightCapacityAndCountReturn(t *testing.T) {
}
}
func TestRPCResultFlightOversizedPutStillResolvesWaiters(t *testing.T) {
func TestRPCResultFlightLargePutPublishesCompletedBeforeResolvingWaiters(t *testing.T) {
cache := newRPCResultCacheWithFlightLimit(time.Now, 1)
authKeyID := rpcFlightTestAuthID(50)
owner, err := cache.Acquire(authKeyID, 60, 600)
@ -263,31 +721,32 @@ func TestRPCResultFlightOversizedPutStillResolvesWaiters(t *testing.T) {
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}
// This is larger than the removed 4 MiB per-shard partition but remains a
// legal outbound result and fits the global/auth/session fair byte budgets.
largeSize := rpcResultCacheMaxBytes/rpcResultCacheShards + 1
want := &encodedOutboundMessage{body: make([]byte, largeSize), 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)
t.Fatalf("large 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, ok := cache.Get(authKeyID, 60, 600); !ok || got != want {
t.Fatalf("large completed Get = encoded:%p ok:%v", got, ok)
}
if got := cache.flightLimit.snapshot(); got != 0 {
t.Fatalf("oversized Put leaked pending count %d", got)
t.Fatalf("large Put leaked pending count %d", got)
}
if owner.owner.Abort() {
t.Fatal("oversized Put left its old owner abortable")
t.Fatal("large 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)
completed, err := cache.Acquire(authKeyID, 60, 600)
if err != nil || completed.state != rpcResultAcquireCompleted || completed.encoded != want {
t.Fatalf("Acquire after large completion = state:%d encoded:%p err:%v", completed.state, completed.encoded, err)
}
if !retry.owner.Abort() {
t.Fatal("retry owner failed to abort")
if completed.owner != nil {
t.Fatal("large completed result incorrectly returned a new owner")
}
if got := cache.completedBytes.snapshot(); got != int64(largeSize) {
t.Fatalf("completed byte budget = %d, want %d", got, largeSize)
}
}
@ -370,11 +829,11 @@ func TestQueuedRPCConnectionCloseAbortsOwnerClaim(t *testing.T) {
}
reservation, err := c.reserveInboundRPC(context.Background(), "test.queuedFlight", 4)
if err != nil {
t.Fatalf("reserve queued RPC: %v", err)
t.Fatalf("reserve queued legacyRPC: %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)
t.Fatalf("commit queued legacyRPC: %v", err)
}
// The scheduler is intentionally not started, so close must drain the queued

87
internal/mtprotoedge/rpc_replay_restore.go Normal file
View file

@ -0,0 +1,87 @@
package mtprotoedge
import (
"context"
"fmt"
"time"
)
const (
rpcReplayRestoreTimeout = 5 * time.Second
rpcReplayRestoreMaxActive = 8
)
// rpcReplayRestoreSlots bounds callbacks that ignore their own store/context
// deadline. The caller waits only until ctx (and at most
// rpcReplayRestoreTimeout); a stuck callback can retain one of these fixed
// slots, but it cannot retain a rewrap worker, a live Conn scheduler barrier,
// or create an unbounded goroutine population.
var rpcReplayRestoreSlots = make(chan struct{}, rpcReplayRestoreMaxActive)
func boundedRPCReplayRestoreContext(parent context.Context) (context.Context, context.CancelFunc) {
if parent == nil {
parent = context.Background()
}
return context.WithTimeout(parent, rpcReplayRestoreTimeout)
}
// runBoundedRPCReplayRestore runs replacement metadata first and the shared
// logical delivery hook second. On timeout it fences the physical generation
// before returning, so releasing that Conn's scheduler barrier cannot expose a
// following RPC to partially restored state. A claim that has not entered the
// hook is abandoned and may be acquired by a later replacement replay; an
// in-progress hook remains at-most-once and every later replay waits for Done.
func (s *Server) runBoundedRPCReplayRestore(
ctx context.Context,
c *Conn,
source string,
claim *rpcResultDeliveryHookClaim,
replacement func() error,
) error {
if claim == nil && replacement == nil {
return nil
}
boundedCtx, cancel := boundedRPCReplayRestoreContext(ctx)
defer cancel()
select {
case rpcReplayRestoreSlots <- struct{}{}:
case <-boundedCtx.Done():
claim.abandon()
if c != nil {
c.fenceUndeliveredRPCResult()
}
return fmt.Errorf("restore replay state after %s: %w", source, boundedCtx.Err())
}
result := make(chan error, 1)
var logical func()
if claim != nil {
logical = claim.run
}
go func() {
defer func() { <-rpcReplayRestoreSlots }()
result <- s.runRPCReplayRestore(c, source, composeRPCReplayRestore(logical, replacement))
}()
select {
case err := <-result:
return err
case <-boundedCtx.Done():
// Prefer a restore that reached its terminal state concurrently with the
// deadline; its result channel publishes coordinator Done before send.
select {
case err := <-result:
return err
default:
}
// If replacement metadata is still blocked, the hook is only Claimed and
// can be safely re-acquired. Once hook execution is InProgress, abandon
// deliberately fails: retrying an unknown partial side effect is forbidden.
claim.abandon()
if c != nil {
c.fenceUndeliveredRPCResult()
}
return fmt.Errorf("restore replay state after %s: %w", source, boundedCtx.Err())
}
}

292
internal/mtprotoedge/rpc_result_budget.go Normal file
View file

@ -0,0 +1,292 @@
package mtprotoedge
import (
"encoding/binary"
"hash/maphash"
"sync"
)
const rpcResultBudgetShards = 64
type rpcResultBudgetLimit struct {
entries int64
bytes int64
}
type rpcResultBudgetUsage struct {
entries int64
bytes int64
pending int64
}
type rpcResultSessionBudgetKey struct {
authKeyID [8]byte
sessionID int64
}
type rpcResultAuthBudgetShard struct {
mu sync.Mutex
usage map[[8]byte]rpcResultBudgetUsage
}
type rpcResultSessionBudgetShard struct {
mu sync.Mutex
usage map[rpcResultSessionBudgetKey]rpcResultBudgetUsage
}
// rpcResultFairBudget accounts one ownership reservation at all three scopes.
// A pending owner and its completed result are the same ownership: admission
// reserves one entry plus one byte, Put resizes that byte reservation and moves
// it to the completed row, while Abort/TTL return the whole reservation.
//
// Auth and session maps are striped independently. Every operation takes the
// auth stripe before the session stripe; global counters remain atomic. This
// keeps unrelated auth keys off one process-wide mutex while preserving hard
// limits at every hierarchy level.
type rpcResultFairBudget struct {
seed maphash.Seed
globalEntries *rpcResultFlightLimit
globalBytes *rpcResultCacheByteBudget
authLimit rpcResultBudgetLimit
sessionLimit rpcResultBudgetLimit
pendingPerAuth int64
authShards [rpcResultBudgetShards]rpcResultAuthBudgetShard
sessionShards [rpcResultBudgetShards]rpcResultSessionBudgetShard
}
type rpcResultBudgetReservation struct {
budget *rpcResultFairBudget
key rpcResultCacheKey
bytes int
pending bool
released bool
}
func newRPCResultFairBudget(
seed maphash.Seed,
globalEntries *rpcResultFlightLimit,
globalBytes *rpcResultCacheByteBudget,
authLimit rpcResultBudgetLimit,
sessionLimit rpcResultBudgetLimit,
pendingPerAuth int,
) *rpcResultFairBudget {
b := &rpcResultFairBudget{
seed: seed,
globalEntries: globalEntries,
globalBytes: globalBytes,
authLimit: authLimit,
sessionLimit: sessionLimit,
pendingPerAuth: int64(pendingPerAuth),
}
for i := range b.authShards {
b.authShards[i].usage = make(map[[8]byte]rpcResultBudgetUsage)
b.sessionShards[i].usage = make(map[rpcResultSessionBudgetKey]rpcResultBudgetUsage)
}
return b
}
func (b *rpcResultFairBudget) reserveOwner(key rpcResultCacheKey) *rpcResultBudgetReservation {
return b.reserve(key, 1, true)
}
func (b *rpcResultFairBudget) reserveCompleted(key rpcResultCacheKey, bytes int) *rpcResultBudgetReservation {
return b.reserve(key, bytes, false)
}
func (b *rpcResultFairBudget) reserve(key rpcResultCacheKey, bytes int, pending bool) *rpcResultBudgetReservation {
if b == nil || b.globalEntries == nil || b.globalBytes == nil || bytes < 1 {
return nil
}
authShard := b.authShard(key.authKeyID)
sessionKey := rpcResultSessionBudgetKey{authKeyID: key.authKeyID, sessionID: key.sessionID}
sessionShard := b.sessionShard(sessionKey)
authShard.mu.Lock()
sessionShard.mu.Lock()
authUsage := authShard.usage[key.authKeyID]
sessionUsage := sessionShard.usage[sessionKey]
bytes64 := int64(bytes)
canReserve := withinRPCResultBudget(authUsage.entries, 1, b.authLimit.entries) &&
withinRPCResultBudget(authUsage.bytes, bytes64, b.authLimit.bytes) &&
withinRPCResultBudget(sessionUsage.entries, 1, b.sessionLimit.entries) &&
withinRPCResultBudget(sessionUsage.bytes, bytes64, b.sessionLimit.bytes)
if pending {
canReserve = canReserve && withinRPCResultBudget(authUsage.pending, 1, b.pendingPerAuth)
}
if !canReserve || !b.globalEntries.reserve() {
sessionShard.mu.Unlock()
authShard.mu.Unlock()
return nil
}
if !b.globalBytes.reserve(bytes) {
b.globalEntries.release()
sessionShard.mu.Unlock()
authShard.mu.Unlock()
return nil
}
authUsage.entries++
authUsage.bytes += bytes64
sessionUsage.entries++
sessionUsage.bytes += bytes64
if pending {
authUsage.pending++
}
authShard.usage[key.authKeyID] = authUsage
sessionShard.usage[sessionKey] = sessionUsage
sessionShard.mu.Unlock()
authShard.mu.Unlock()
return &rpcResultBudgetReservation{budget: b, key: key, bytes: bytes, pending: pending}
}
func withinRPCResultBudget(used, delta, limit int64) bool {
return delta >= 0 && limit > 0 && used >= 0 && used <= limit-delta
}
func (r *rpcResultBudgetReservation) resizeBytes(bytes int) bool {
if r == nil || r.budget == nil || r.released || bytes < 1 {
return false
}
if bytes == r.bytes {
return true
}
b := r.budget
authShard := b.authShard(r.key.authKeyID)
sessionKey := rpcResultSessionBudgetKey{authKeyID: r.key.authKeyID, sessionID: r.key.sessionID}
sessionShard := b.sessionShard(sessionKey)
authShard.mu.Lock()
sessionShard.mu.Lock()
authUsage, authOK := authShard.usage[r.key.authKeyID]
sessionUsage, sessionOK := sessionShard.usage[sessionKey]
if !authOK || !sessionOK || authUsage.entries < 1 || sessionUsage.entries < 1 {
sessionShard.mu.Unlock()
authShard.mu.Unlock()
panic("mtprotoedge: rpc result budget reservation disappeared during resize")
}
delta := int64(bytes) - int64(r.bytes)
if delta > 0 {
if !withinRPCResultBudget(authUsage.bytes, delta, b.authLimit.bytes) ||
!withinRPCResultBudget(sessionUsage.bytes, delta, b.sessionLimit.bytes) ||
!b.globalBytes.reserve(int(delta)) {
sessionShard.mu.Unlock()
authShard.mu.Unlock()
return false
}
} else {
if authUsage.bytes < -delta || sessionUsage.bytes < -delta {
sessionShard.mu.Unlock()
authShard.mu.Unlock()
panic("mtprotoedge: rpc result byte reservation underflow during resize")
}
}
authUsage.bytes += delta
sessionUsage.bytes += delta
authShard.usage[r.key.authKeyID] = authUsage
sessionShard.usage[sessionKey] = sessionUsage
r.bytes = bytes
if delta < 0 {
b.globalBytes.release(int(-delta))
}
sessionShard.mu.Unlock()
authShard.mu.Unlock()
return true
}
func (r *rpcResultBudgetReservation) releasePending() {
if r == nil || r.budget == nil || r.released || !r.pending {
return
}
b := r.budget
authShard := b.authShard(r.key.authKeyID)
authShard.mu.Lock()
authUsage, ok := authShard.usage[r.key.authKeyID]
if !ok || authUsage.pending < 1 {
authShard.mu.Unlock()
panic("mtprotoedge: rpc result per-auth pending budget underflow")
}
authUsage.pending--
authShard.usage[r.key.authKeyID] = authUsage
r.pending = false
authShard.mu.Unlock()
}
func (r *rpcResultBudgetReservation) release() {
if r == nil || r.budget == nil || r.released {
return
}
b := r.budget
authShard := b.authShard(r.key.authKeyID)
sessionKey := rpcResultSessionBudgetKey{authKeyID: r.key.authKeyID, sessionID: r.key.sessionID}
sessionShard := b.sessionShard(sessionKey)
authShard.mu.Lock()
sessionShard.mu.Lock()
authUsage, authOK := authShard.usage[r.key.authKeyID]
sessionUsage, sessionOK := sessionShard.usage[sessionKey]
bytes64 := int64(r.bytes)
if !authOK || !sessionOK || authUsage.entries < 1 || sessionUsage.entries < 1 ||
authUsage.bytes < bytes64 || sessionUsage.bytes < bytes64 ||
(r.pending && authUsage.pending < 1) {
sessionShard.mu.Unlock()
authShard.mu.Unlock()
panic("mtprotoedge: rpc result fair budget underflow")
}
authUsage.entries--
authUsage.bytes -= bytes64
sessionUsage.entries--
sessionUsage.bytes -= bytes64
if r.pending {
authUsage.pending--
}
if authUsage == (rpcResultBudgetUsage{}) {
delete(authShard.usage, r.key.authKeyID)
} else {
authShard.usage[r.key.authKeyID] = authUsage
}
if sessionUsage == (rpcResultBudgetUsage{}) {
delete(sessionShard.usage, sessionKey)
} else {
sessionShard.usage[sessionKey] = sessionUsage
}
r.released = true
r.pending = false
r.bytes = 0
b.globalBytes.release(int(bytes64))
b.globalEntries.release()
sessionShard.mu.Unlock()
authShard.mu.Unlock()
}
func (b *rpcResultFairBudget) authSnapshot(authKeyID [8]byte) rpcResultBudgetUsage {
if b == nil {
return rpcResultBudgetUsage{}
}
shard := b.authShard(authKeyID)
shard.mu.Lock()
usage := shard.usage[authKeyID]
shard.mu.Unlock()
return usage
}
func (b *rpcResultFairBudget) sessionSnapshot(authKeyID [8]byte, sessionID int64) rpcResultBudgetUsage {
if b == nil {
return rpcResultBudgetUsage{}
}
key := rpcResultSessionBudgetKey{authKeyID: authKeyID, sessionID: sessionID}
shard := b.sessionShard(key)
shard.mu.Lock()
usage := shard.usage[key]
shard.mu.Unlock()
return usage
}
func (b *rpcResultFairBudget) authShard(authKeyID [8]byte) *rpcResultAuthBudgetShard {
index := maphash.Bytes(b.seed, authKeyID[:]) & (rpcResultBudgetShards - 1)
return &b.authShards[index]
}
func (b *rpcResultFairBudget) sessionShard(key rpcResultSessionBudgetKey) *rpcResultSessionBudgetShard {
var raw [16]byte
copy(raw[:8], key.authKeyID[:])
binary.LittleEndian.PutUint64(raw[8:], uint64(key.sessionID))
index := maphash.Bytes(b.seed, raw[:]) & (rpcResultBudgetShards - 1)
return &b.sessionShards[index]
}

501
internal/mtprotoedge/rpc_result_cache.go
View file

@ -3,17 +3,39 @@ package mtprotoedge
import (
"container/list"
"encoding/binary"
"hash/maphash"
"sync"
"sync/atomic"
"time"
)
const (
rpcResultCacheTTL = 3 * time.Minute
rpcResultCacheMaxEntries = 4096
rpcResultCacheMaxBytes = 64 << 20
// rpcResultCacheShards 把缓存按 (auth_key_id, session_id) 分片：每条 RPC 都要
// Get（重复检测）+ Put（结果缓存），单把全局锁会让所有连接的 RPC 热路径在
// 一个 mutex 上汇聚（同 P0-5 的 SessionManager 教训）。分片数为 2 的幂。
// Telegram accepts client msg_id values up to five minutes old and up to
// thirty seconds in the future. Retain the result across that complete
// replay horizon, plus one second for boundary/scheduler jitter, so a valid
// duplicate cannot rerun its handler merely because our cache expired first.
rpcResultCacheTTL = 331 * time.Second
// Completed results cover the complete replay horizon under explicit global,
// auth and session hard ceilings. At the default 331-second TTL, the 1<<18
// global entries permit about 792 unique RPC/s process-wide before bounded
// backpressure; lower scopes provide noisy-neighbor isolation.
rpcResultCacheMaxEntries = 1 << 18
rpcResultCacheMaxBytes = 64 << 20
rpcResultCacheAuthMaxEntries = 1 << 15
rpcResultCacheAuthMaxBytes = 32 << 20
rpcResultCacheSessionMaxEntries = 1 << 14
rpcResultCacheSessionMaxBytes = 16 << 20
rpcResultFlightMaxPendingPerAuth = 1 << 11
// Keep every transport-legal rpc_result cacheable. Converting the constant
// difference to uint64 intentionally fails compilation if a future transport
// limit grows beyond the completed-result budget.
_ = uint64(rpcResultCacheMaxBytes - maxOutboundBodyBytes)
_ = uint64(rpcResultCacheAuthMaxBytes - maxOutboundBodyBytes)
_ = uint64(rpcResultCacheSessionMaxBytes - maxOutboundBodyBytes)
// rpcResultCacheShards hashes the complete replay identity with a random
// per-instance maphash seed. Including req_msg_id spreads one hot session's
// independent requests instead of forcing them through one mutex. The shard
// count is a power of two.
rpcResultCacheShards = 16
)
@ -24,10 +46,28 @@ type rpcResultCacheKey struct {
}
type rpcResultCacheEntry struct {
key rpcResultCacheKey
encoded *encodedOutboundMessage
size int
expiresAt time.Time
key rpcResultCacheKey
encoded *encodedOutboundMessage
size int
expiresAt time.Time
identity rpcResultRequestIdentity
admissionSeq uint64
executionKnown bool
executionOK bool
// capacity marks a bounded replay tombstone. The original owner and its
// already-joined waiters received encoded, but the byte budget could not
// retain that body. Keeping the immutable identity until TTL prevents a
// duplicate from rerunning business; Acquire returns a capacity error.
capacity bool
// reservation is the same global+auth+session ownership acquired before the
// handler ran. Put transfers it from the pending flight; TTL returns it.
reservation *rpcResultBudgetReservation
}
type rpcResultDependency struct {
waiter *rpcResultWaiter
completed bool
success bool
}
// rpcResultCache 缓存已有交付证明的 rpc_result（按 auth_key+session+req_msg_id），
@ -36,16 +76,35 @@ type rpcResultCacheEntry struct {
// encodedOutboundMessage 构造后不可变（push fan-out 与 pending resend 均依赖该契约），
// 因此 Get/Put 直接共享指针，不做防御性拷贝。
type rpcResultCache struct {
shards [rpcResultCacheShards]rpcResultCacheShard
flightLimit rpcResultFlightLimit
shards [rpcResultCacheShards]rpcResultCacheShard
hashSeed maphash.Seed
completedBytes rpcResultCacheByteBudget
completedEntries rpcResultFlightLimit
fairBudget *rpcResultFairBudget
flightLimit rpcResultFlightLimit
subscriberBudget *rpcResultSubscriberBudget
subscriberPerFlight int
// nextAdmissionSeq is the process-wide ordering authority for auth-key
// shared Layer defaults. Exact owners allocate once; joins/replays retain the
// owner's value from their flight/completed descriptor.
nextAdmissionSeq atomic.Uint64
activeAdmissions rpcAdmissionTracker
}
func (c *rpcResultCache) stableAdmissionSafeFloor() uint64 {
if c == nil {
return 0
}
return c.activeAdmissions.stableSafeFloor(&c.nextAdmissionSeq)
}
type rpcResultCacheShard struct {
mu sync.Mutex
now func() time.Time
ttl time.Duration
mu sync.Mutex
now func() time.Time
ttl time.Duration
// maxEntries is a focused-test seam for one physical shard. Production leaves
// it zero and uses the explicit global/auth/session fair-budget hierarchy.
maxEntries int
maxBytes int
bytes int
order *list.List
byKey map[rpcResultCacheKey]*list.Element
@ -56,20 +115,129 @@ type rpcResultCacheShard struct {
}
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)
pendingPerAuth := rpcResultFlightMaxPendingPerAuth
if pendingPerAuth > maxPending {
pendingPerAuth = maxPending
}
return newRPCResultCacheWithFairCapacity(now, rpcResultCacheCapacity{
maxPending: maxPending,
maxPendingPerAuth: pendingPerAuth,
globalMaxBytes: rpcResultCacheMaxBytes,
globalMaxEntries: rpcResultCacheMaxEntries,
authMaxBytes: rpcResultCacheAuthMaxBytes,
authMaxEntries: rpcResultCacheAuthMaxEntries,
sessionMaxBytes: rpcResultCacheSessionMaxBytes,
sessionMaxEntries: rpcResultCacheSessionMaxEntries,
})
}
func newRPCResultCacheWithLimits(now func() time.Time, maxPending, maxCompletedBytes int) *rpcResultCache {
return newRPCResultCacheWithCapacity(now, maxPending, int64(maxCompletedBytes), rpcResultCacheMaxEntries)
}
func newRPCResultCacheWithCapacity(
now func() time.Time,
maxPending int,
maxCompletedBytes int64,
maxCompletedEntries int,
) *rpcResultCache {
// Compatibility/test constructor: the caller supplied only global limits, so
// keep every fairness scope equal to that global ceiling. Production always
// calls newRPCResultCacheWithFairCapacity with explicit auth/session limits.
return newRPCResultCacheWithFairCapacity(now, rpcResultCacheCapacity{
maxPending: maxPending,
maxPendingPerAuth: maxPending,
globalMaxBytes: maxCompletedBytes,
globalMaxEntries: maxCompletedEntries,
authMaxBytes: maxCompletedBytes,
authMaxEntries: maxCompletedEntries,
sessionMaxBytes: maxCompletedBytes,
sessionMaxEntries: maxCompletedEntries,
})
}
type rpcResultCacheCapacity struct {
maxPending int
maxPendingPerAuth int
globalMaxBytes int64
globalMaxEntries int
authMaxBytes int64
authMaxEntries int
sessionMaxBytes int64
sessionMaxEntries int
subscriberMaxGlobal int
subscriberMaxAuth int
subscriberMaxSession int
subscriberMaxPerFlight int
}
func newRPCResultCacheWithFairCapacity(now func() time.Time, capacity rpcResultCacheCapacity) *rpcResultCache {
if now == nil {
now = time.Now
}
if capacity.maxPending <= 0 {
capacity.maxPending = rpcResultFlightDefaultMaxPending
}
if capacity.maxPendingPerAuth <= 0 {
capacity.maxPendingPerAuth = capacity.maxPending
}
if capacity.globalMaxBytes <= 0 {
capacity.globalMaxBytes = rpcResultCacheMaxBytes
}
if capacity.globalMaxEntries <= 0 {
capacity.globalMaxEntries = rpcResultCacheMaxEntries
}
if capacity.authMaxBytes <= 0 {
capacity.authMaxBytes = capacity.globalMaxBytes
}
if capacity.authMaxEntries <= 0 {
capacity.authMaxEntries = capacity.globalMaxEntries
}
if capacity.sessionMaxBytes <= 0 {
capacity.sessionMaxBytes = capacity.authMaxBytes
}
if capacity.sessionMaxEntries <= 0 {
capacity.sessionMaxEntries = capacity.authMaxEntries
}
if capacity.subscriberMaxGlobal <= 0 {
capacity.subscriberMaxGlobal = rpcResultSubscriberMaxGlobal
}
if capacity.subscriberMaxAuth <= 0 {
capacity.subscriberMaxAuth = rpcResultSubscriberMaxAuth
}
if capacity.subscriberMaxSession <= 0 {
capacity.subscriberMaxSession = rpcResultSubscriberMaxSession
}
if capacity.subscriberMaxPerFlight <= 0 {
capacity.subscriberMaxPerFlight = rpcResultSubscriberMaxPerFlight
}
c := &rpcResultCache{hashSeed: maphash.MakeSeed()}
c.completedBytes.max = capacity.globalMaxBytes
c.completedEntries.max = int64(capacity.globalMaxEntries)
c.flightLimit.max = int64(capacity.maxPending)
c.fairBudget = newRPCResultFairBudget(
c.hashSeed,
&c.completedEntries,
&c.completedBytes,
rpcResultBudgetLimit{entries: int64(capacity.authMaxEntries), bytes: capacity.authMaxBytes},
rpcResultBudgetLimit{entries: int64(capacity.sessionMaxEntries), bytes: capacity.sessionMaxBytes},
capacity.maxPendingPerAuth,
)
c.subscriberBudget = newRPCResultSubscriberBudget(
c.hashSeed,
capacity.subscriberMaxGlobal,
capacity.subscriberMaxAuth,
capacity.subscriberMaxSession,
)
c.subscriberPerFlight = capacity.subscriberMaxPerFlight
for i := range c.shards {
s := &c.shards[i]
s.now = now
s.ttl = rpcResultCacheTTL
s.maxEntries = rpcResultCacheMaxEntries / rpcResultCacheShards
s.maxBytes = rpcResultCacheMaxBytes / rpcResultCacheShards
s.maxEntries = 0
s.order = list.New()
s.byKey = make(map[rpcResultCacheKey]*list.Element)
s.pending = make(map[rpcResultCacheKey]*rpcResultFlight)
@ -78,10 +246,16 @@ func newRPCResultCacheWithFlightLimit(now func() time.Time, maxPending int) *rpc
}
func (c *rpcResultCache) shard(key rpcResultCacheKey) *rpcResultCacheShard {
// auth_key_id 与 session_id 都是均匀随机的 64-bit 值，异或折叠后取低位即可。
h := binary.LittleEndian.Uint64(key.authKeyID[:]) ^ uint64(key.sessionID)
h ^= h >> 32
return &c.shards[h&(rpcResultCacheShards-1)]
return &c.shards[c.shardIndex(key)]
}
func (c *rpcResultCache) shardIndex(key rpcResultCacheKey) uint64 {
var raw [24]byte
copy(raw[:8], key.authKeyID[:])
binary.LittleEndian.PutUint64(raw[8:16], uint64(key.sessionID))
binary.LittleEndian.PutUint64(raw[16:24], uint64(key.reqMsgID))
return maphash.Bytes(c.hashSeed, raw[:]) & (rpcResultCacheShards - 1)
}
func (c *rpcResultCache) Get(authKeyID [8]byte, sessionID, reqMsgID int64) (*encodedOutboundMessage, bool) {
@ -104,45 +278,210 @@ func (c *rpcResultCache) Get(authKeyID [8]byte, sessionID, reqMsgID int64) (*enc
s.removeElement(elem)
return nil, false
}
if entry.capacity || entry.encoded == nil {
return nil, false
}
return entry.encoded, true
}
// ObserveDependency returns a waiter for an admitted in-flight dependency, a
// nil waiter for an already completed dependency, or ok=false when the
// referenced message never established API-RPC ownership. It never creates a
// flight and therefore cannot turn a forged invokeAfterMsg into authority to
// run another request.
func (c *rpcResultCache) ObserveDependency(authKeyID [8]byte, sessionID, reqMsgID int64) (rpcResultDependency, bool) {
if c == nil || reqMsgID == 0 {
return rpcResultDependency{}, false
}
key := rpcResultCacheKey{authKeyID: authKeyID, sessionID: sessionID, reqMsgID: reqMsgID}
s := c.shard(key)
now := s.now()
s.mu.Lock()
defer s.mu.Unlock()
if elem, exists := s.byKey[key]; exists {
entry := elem.Value.(*rpcResultCacheEntry)
if entry.expiresAt.After(now) {
if !entry.executionKnown {
return rpcResultDependency{}, false
}
return rpcResultDependency{completed: true, success: entry.executionOK}, true
}
s.removeElement(elem)
}
if flight := s.pending[key]; flight != nil {
if flight.executionDone {
return rpcResultDependency{completed: true, success: flight.executionOK}, true
}
return rpcResultDependency{waiter: &rpcResultWaiter{cache: c, key: key, flight: flight}}, true
}
return rpcResultDependency{}, false
}
func (c *rpcResultCache) Put(authKeyID [8]byte, sessionID, reqMsgID int64, encoded *encodedOutboundMessage) {
if c == nil || reqMsgID == 0 || encoded == nil {
return
}
if c.putOnce(authKeyID, sessionID, reqMsgID, encoded) {
return
}
// A direct Put has no pre-reserved owner slot. Expired entries in another
// shard may be its only blocker; reap once without holding a shard and retry.
// Production owner publication already carries both reservations and never
// needs this cold path.
c.expireCompletedResults()
_ = c.putOnce(authKeyID, sessionID, reqMsgID, encoded)
}
// putOnce returns false only when a cross-shard expiry reap may release the
// process-wide entry/body capacity needed by a defensive direct Put.
func (c *rpcResultCache) putOnce(authKeyID [8]byte, sessionID, reqMsgID int64, encoded *encodedOutboundMessage) bool {
key := rpcResultCacheKey{authKeyID: authKeyID, sessionID: sessionID, reqMsgID: reqMsgID}
s := c.shard(key)
size := len(encoded.body)
cacheable := s.maxBytes <= 0 || size <= s.maxBytes
now := s.now()
s.mu.Lock()
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()
accountedSize := len(encoded.body)
if accountedSize < 1 {
// Every owner reserves one byte at admission. Keeping zero-length results
// at the same minimum makes entry and byte capacity linearizable.
accountedSize = 1
}
// 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.
subscribers := c.completeRPCResultFlightLocked(s, key, encoded)
// Publication never evicts another unexpired result. A production owner has
// already reserved its entry slot and one byte. If its actual result cannot
// expand that reservation, publish a one-byte identity tombstone: the owner
// and current waiters still receive the immutable result, while later
// duplicates fail admission instead of rerunning the handler.
s.mu.Lock()
now := s.now()
s.expireLocked(now)
old := s.byKey[key]
flight := s.pending[key]
if old == nil && flight == nil && s.maxEntries > 0 && len(s.byKey) >= s.maxEntries {
// Defensive direct Put callers do not own a reserved admission slot.
// Preserve every existing unexpired result and decline the new cache row.
s.mu.Unlock()
return true
}
identity, admissionSeq, executionKnown, executionOK := rpcResultFlightMetadataLocked(s, key)
var oldEntry *rpcResultCacheEntry
if old != nil {
oldEntry = old.Value.(*rpcResultCacheEntry)
if flight == nil {
// A defensive duplicate terminal publication must never downgrade
// completed dependency/identity metadata after its flight disappeared.
identity = oldEntry.identity
admissionSeq = oldEntry.admissionSeq
executionKnown = oldEntry.executionKnown
executionOK = oldEntry.executionOK
}
}
var reservation *rpcResultBudgetReservation
switch {
case flight != nil:
reservation = flight.reservation
if reservation == nil {
s.mu.Unlock()
panic("mtprotoedge: pending rpc result has no fair-budget reservation")
}
case oldEntry != nil && oldEntry.reservation != nil:
reservation = oldEntry.reservation
default:
reservation = c.fairBudget.reserveCompleted(key, accountedSize)
if reservation == nil {
s.mu.Unlock()
return false
}
}
retainedSize := accountedSize
retained := encoded
capacity := false
if !reservation.resizeBytes(accountedSize) {
if flight == nil {
// A direct replacement cannot discard the prior replay body. Leave it
// untouched and let Put perform one cross-shard expiry reap before its
// final bounded failure.
if oldEntry == nil {
reservation.release()
}
s.mu.Unlock()
return false
}
// Owner admission already reserved one byte at all three scopes. When the
// actual body cannot expand, transfer that reservation to an identity
// tombstone so a duplicate never reruns business.
const tombstoneSize = 1
if !reservation.resizeBytes(tombstoneSize) {
s.mu.Unlock()
panic("mtprotoedge: rpc result owner lost its one-byte tombstone reservation")
}
retainedSize = tombstoneSize
retained = nil
capacity = true
}
if old != nil {
s.unlinkElement(old)
if oldEntry.reservation != nil && oldEntry.reservation != reservation {
oldEntry.reservation.release()
oldEntry.reservation = nil
}
}
entry := &rpcResultCacheEntry{
key: key,
encoded: retained,
size: retainedSize,
expiresAt: now.Add(s.ttl),
identity: identity,
admissionSeq: admissionSeq,
executionKnown: executionKnown,
executionOK: executionOK,
capacity: capacity,
reservation: reservation,
}
elem := s.order.PushBack(entry)
s.byKey[key] = elem
s.bytes += retainedSize
// Resolve the independent in-flight entry only after either the completed
// result or its replay tombstone is published under the same shard lock.
subscribers, executionSubscribers, executionOK := c.completeRPCResultFlightLocked(s, key, encoded)
s.mu.Unlock()
for _, subscriber := range subscribers {
subscriber(encoded, true)
}
for _, subscriber := range executionSubscribers {
subscriber(executionOK)
}
return true
}
func rpcResultFlightMetadataLocked(s *rpcResultCacheShard, key rpcResultCacheKey) (
rpcResultRequestIdentity,
uint64,
bool,
bool,
) {
if flight := s.pending[key]; flight != nil {
return flight.identity, flight.admissionSeq, flight.executionDone, flight.executionOK
}
return rpcResultRequestIdentity{}, 0, false, false
}
// expireCompletedResults performs the cold-path cross-shard reap used only
// after a one-byte admission reservation fails. The caller must hold no shard
// lock. Each shard is reaped independently so ordinary result publication on
// the other shards remains parallel.
func (c *rpcResultCache) expireCompletedResults() {
if c == nil {
return
}
for i := range c.shards {
s := &c.shards[i]
s.mu.Lock()
s.expireLocked(s.now())
s.mu.Unlock()
}
}
func (s *rpcResultCacheShard) expireLocked(now time.Time) {
@ -157,20 +496,17 @@ func (s *rpcResultCacheShard) expireLocked(now time.Time) {
}
}
func (s *rpcResultCacheShard) trimLocked() {
for s.order.Len() > 0 {
tooManyEntries := s.maxEntries > 0 && s.order.Len() > s.maxEntries
tooManyBytes := s.maxBytes > 0 && s.bytes > s.maxBytes
if !tooManyEntries && !tooManyBytes {
return
}
s.removeElement(s.order.Front())
func (s *rpcResultCacheShard) removeElement(elem *list.Element) {
entry := s.unlinkElement(elem)
if entry != nil && entry.reservation != nil {
entry.reservation.release()
entry.reservation = nil
}
}
func (s *rpcResultCacheShard) removeElement(elem *list.Element) {
func (s *rpcResultCacheShard) unlinkElement(elem *list.Element) *rpcResultCacheEntry {
if elem == nil {
return
return nil
}
entry := elem.Value.(*rpcResultCacheEntry)
delete(s.byKey, entry.key)
@ -179,4 +515,45 @@ func (s *rpcResultCacheShard) removeElement(elem *list.Element) {
s.bytes = 0
}
s.order.Remove(elem)
return entry
}
type rpcResultCacheByteBudget struct {
max int64
used atomic.Int64
}
func (b *rpcResultCacheByteBudget) reserve(n int) bool {
if n <= 0 {
return true
}
bytes := int64(n)
if b == nil || bytes > b.max {
return false
}
for {
used := b.used.Load()
if used > b.max-bytes {
return false
}
if b.used.CompareAndSwap(used, used+bytes) {
return true
}
}
}
func (b *rpcResultCacheByteBudget) release(n int) {
if b == nil || n <= 0 {
return
}
if remaining := b.used.Add(-int64(n)); remaining < 0 {
panic("mtprotoedge: rpc result completed-byte budget underflow")
}
}
func (b *rpcResultCacheByteBudget) snapshot() int64 {
if b == nil {
return 0
}
return b.used.Load()
}

790
internal/mtprotoedge/rpc_result_cache_test.go
View file

@ -1,11 +1,476 @@
package mtprotoedge
import (
"errors"
"sync"
"sync/atomic"
"testing"
"time"
)
func TestRPCResultCacheFullSessionDoesNotBlockAnotherAuth(t *testing.T) {
now := time.Unix(1000, 0)
cache := newRPCResultCacheWithFairCapacity(func() time.Time { return now }, rpcResultCacheCapacity{
maxPending: 8, maxPendingPerAuth: 6,
globalMaxEntries: 8, globalMaxBytes: 64,
authMaxEntries: 6, authMaxBytes: 48,
sessionMaxEntries: 2, sessionMaxBytes: 16,
})
authA := [8]byte{0xa1}
authB := [8]byte{0xb1}
const sessionA = int64(77)
for i := 0; i < 2; i++ {
msgID := int64(1000 + i)
claim, err := cache.Acquire(authA, sessionA, msgID)
if err != nil || claim.state != rpcResultAcquireOwner {
t.Fatalf("same-session admission %d = %#v, %v", i, claim, err)
}
cache.Put(authA, sessionA, msgID, &encodedOutboundMessage{body: []byte{1}})
}
if _, err := cache.Acquire(authA, sessionA, 2000); !errors.Is(err, ErrRPCResultFlightCapacity) {
t.Fatalf("admission beyond session entry limit = %v, want capacity", err)
}
otherAuth, err := cache.Acquire(authB, 88, 3000)
if err != nil || otherAuth.state != rpcResultAcquireOwner {
t.Fatalf("other auth blocked by full session: %#v, %v", otherAuth, err)
}
if !otherAuth.owner.Abort() {
t.Fatal("other-auth owner did not abort")
}
if _, ok := cache.Get(authA, sessionA, 1000); !ok {
t.Fatal("session capacity pressure evicted an unexpired result")
}
}
func TestRPCResultCacheFullAuthDoesNotBlockAnotherAuth(t *testing.T) {
cache := newRPCResultCacheWithFairCapacity(time.Now, rpcResultCacheCapacity{
maxPending: 8, maxPendingPerAuth: 4,
globalMaxEntries: 8, globalMaxBytes: 64,
authMaxEntries: 2, authMaxBytes: 32,
sessionMaxEntries: 2, sessionMaxBytes: 16,
})
authA := [8]byte{0xa2}
authB := [8]byte{0xb2}
for i := 0; i < 2; i++ {
claim, err := cache.Acquire(authA, int64(10+i), int64(100+i))
if err != nil || claim.state != rpcResultAcquireOwner {
t.Fatalf("auth A admission %d = %#v, %v", i, claim, err)
}
cache.Put(authA, int64(10+i), int64(100+i), &encodedOutboundMessage{body: []byte{1}})
}
if _, err := cache.Acquire(authA, 12, 102); !errors.Is(err, ErrRPCResultFlightCapacity) {
t.Fatalf("same-auth new session at auth limit = %v, want capacity", err)
}
other, err := cache.Acquire(authB, 20, 200)
if err != nil || other.state != rpcResultAcquireOwner {
t.Fatalf("other auth blocked by full auth A: %#v, %v", other, err)
}
other.owner.Abort()
}
func TestRPCResultCacheAuthAndSessionByteLimitsAreIndependent(t *testing.T) {
cache := newRPCResultCacheWithFairCapacity(time.Now, rpcResultCacheCapacity{
maxPending: 8, maxPendingPerAuth: 6,
globalMaxEntries: 10, globalMaxBytes: 10,
authMaxEntries: 8, authMaxBytes: 4,
sessionMaxEntries: 6, sessionMaxBytes: 2,
})
authA := [8]byte{0xa4}
authB := [8]byte{0xb4}
first, err := cache.Acquire(authA, 1, 101)
if err != nil || first.state != rpcResultAcquireOwner {
t.Fatalf("first owner = %#v, %v", first, err)
}
cache.Put(authA, 1, 101, &encodedOutboundMessage{body: []byte{1, 2}})
if _, err := cache.Acquire(authA, 1, 102); !errors.Is(err, ErrRPCResultFlightCapacity) {
t.Fatalf("same session beyond byte limit = %v, want capacity", err)
}
second, err := cache.Acquire(authA, 2, 201)
if err != nil || second.state != rpcResultAcquireOwner {
t.Fatalf("second session owner = %#v, %v", second, err)
}
cache.Put(authA, 2, 201, &encodedOutboundMessage{body: []byte{3, 4}})
if _, err := cache.Acquire(authA, 3, 301); !errors.Is(err, ErrRPCResultFlightCapacity) {
t.Fatalf("same auth beyond byte limit = %v, want capacity", err)
}
other, err := cache.Acquire(authB, 3, 302)
if err != nil || other.state != rpcResultAcquireOwner {
t.Fatalf("other auth blocked by auth A byte limit: %#v, %v", other, err)
}
other.owner.Abort()
}
func TestRPCResultCachePerAuthPendingLimitIsAdditional(t *testing.T) {
cache := newRPCResultCacheWithFairCapacity(time.Now, rpcResultCacheCapacity{
maxPending: 6, maxPendingPerAuth: 2,
globalMaxEntries: 12, globalMaxBytes: 64,
authMaxEntries: 6, authMaxBytes: 32,
sessionMaxEntries: 4, sessionMaxBytes: 16,
})
authA := [8]byte{0xa3}
authB := [8]byte{0xb3}
owners := make([]*rpcResultOwnerLease, 0, 3)
for i := 0; i < 2; i++ {
claim, err := cache.Acquire(authA, int64(i+1), int64(100+i))
if err != nil || claim.state != rpcResultAcquireOwner {
t.Fatalf("pending auth A %d = %#v, %v", i, claim, err)
}
owners = append(owners, claim.owner)
}
if _, err := cache.Acquire(authA, 3, 103); !errors.Is(err, ErrRPCResultFlightCapacity) {
t.Fatalf("third pending owner for auth A = %v, want capacity", err)
}
other, err := cache.Acquire(authB, 4, 104)
if err != nil || other.state != rpcResultAcquireOwner {
t.Fatalf("auth B blocked by auth A pending limit: %#v, %v", other, err)
}
owners = append(owners, other.owner)
for _, owner := range owners {
if !owner.Abort() {
t.Fatal("pending owner did not abort")
}
}
if usage := cache.fairBudget.authSnapshot(authA); usage != (rpcResultBudgetUsage{}) {
t.Fatalf("auth A budget after abort = %#v", usage)
}
}
func TestRPCResultCacheFairReservationLifecycleReturnsEveryScope(t *testing.T) {
now := time.Unix(1000, 0)
cache := newRPCResultCacheWithFairCapacity(func() time.Time { return now }, rpcResultCacheCapacity{
maxPending: 4, maxPendingPerAuth: 3,
globalMaxEntries: 6, globalMaxBytes: 10,
authMaxEntries: 5, authMaxBytes: 8,
sessionMaxEntries: 3, sessionMaxBytes: 6,
})
auth := [8]byte{0xc1}
aborted, err := cache.Acquire(auth, 1, 101)
if err != nil || aborted.state != rpcResultAcquireOwner {
t.Fatalf("aborted owner = %#v, %v", aborted, err)
}
if usage := cache.fairBudget.authSnapshot(auth); usage.entries != 1 || usage.bytes != 1 || usage.pending != 1 {
t.Fatalf("pending auth reservation = %#v", usage)
}
if !aborted.owner.Abort() {
t.Fatal("owner Abort lost")
}
if usage := cache.fairBudget.authSnapshot(auth); usage != (rpcResultBudgetUsage{}) {
t.Fatalf("Abort leaked auth reservation %#v", usage)
}
body, err := cache.Acquire(auth, 1, 102)
if err != nil || body.state != rpcResultAcquireOwner {
t.Fatalf("body owner = %#v, %v", body, err)
}
cache.Put(auth, 1, 102, &encodedOutboundMessage{body: make([]byte, 4)})
if usage := cache.fairBudget.sessionSnapshot(auth, 1); usage.entries != 1 || usage.bytes != 4 || usage.pending != 0 {
t.Fatalf("body session reservation = %#v", usage)
}
tombstone, err := cache.Acquire(auth, 2, 201)
if err != nil || tombstone.state != rpcResultAcquireOwner {
t.Fatalf("tombstone owner = %#v, %v", tombstone, err)
}
// This cannot fit the 10-byte global or 8-byte auth ceiling. Put must not
// panic or lose ownership; it transfers the one-byte token to a tombstone.
cache.Put(auth, 2, 201, &encodedOutboundMessage{body: make([]byte, 20)})
if usage := cache.fairBudget.sessionSnapshot(auth, 2); usage.entries != 1 || usage.bytes != 1 || usage.pending != 0 {
t.Fatalf("tombstone session reservation = %#v", usage)
}
if got := cache.completedEntries.snapshot(); got != 2 {
t.Fatalf("global entries after body+tombstone = %d, want 2", got)
}
if got := cache.completedBytes.snapshot(); got != 5 {
t.Fatalf("global bytes after body+tombstone = %d, want 5", got)
}
cache.Put(auth, 1, 102, &encodedOutboundMessage{body: make([]byte, 2)})
if got := cache.completedBytes.snapshot(); got != 3 {
t.Fatalf("replacement did not resize global bytes: %d", got)
}
now = now.Add(rpcResultCacheTTL + time.Second)
_, _ = cache.Get(auth, 1, 102)
_, _ = cache.Get(auth, 2, 201)
if got := cache.completedEntries.snapshot(); got != 0 {
t.Fatalf("TTL leaked global entries %d", got)
}
if got := cache.completedBytes.snapshot(); got != 0 {
t.Fatalf("TTL leaked global bytes %d", got)
}
if usage := cache.fairBudget.authSnapshot(auth); usage != (rpcResultBudgetUsage{}) {
t.Fatalf("TTL leaked auth reservation %#v", usage)
}
}
func TestRPCResultCacheFullKeyMaphashSpreadsOneSession(t *testing.T) {
first := newRPCResultCacheWithFlightLimit(time.Now, 64)
second := newRPCResultCacheWithFlightLimit(time.Now, 64)
auth := [8]byte{1, 2, 3, 4, 5, 6, 7, 8}
const sessionID = int64(99)
seen := make(map[uint64]struct{})
differentInstance := false
for msgID := int64(1); msgID <= 256; msgID++ {
key := rpcResultCacheKey{authKeyID: auth, sessionID: sessionID, reqMsgID: msgID}
firstIndex := first.shardIndex(key)
seen[firstIndex] = struct{}{}
if firstIndex != second.shardIndex(key) {
differentInstance = true
}
}
if len(seen) < rpcResultCacheShards/2 {
t.Fatalf("one session used only %d/%d full-key shards", len(seen), rpcResultCacheShards)
}
if !differentInstance {
t.Fatal("two cache instances produced an identical shard stream; seed is not instance-random")
}
}
func TestRPCResultCacheConcurrentFairReservationsNeverOvercommit(t *testing.T) {
cache := newRPCResultCacheWithFairCapacity(time.Now, rpcResultCacheCapacity{
maxPending: 24, maxPendingPerAuth: 4,
globalMaxEntries: 24, globalMaxBytes: 24,
authMaxEntries: 8, authMaxBytes: 8,
sessionMaxEntries: 3, sessionMaxBytes: 3,
})
const callers = 256
start := make(chan struct{})
var (
wg sync.WaitGroup
mu sync.Mutex
owners []*rpcResultOwnerLease
)
for i := 0; i < callers; i++ {
wg.Add(1)
go func(i int) {
defer wg.Done()
<-start
auth := [8]byte{byte(i % 4)}
claim, err := cache.Acquire(auth, int64(i%8), int64(1000+i))
if errors.Is(err, ErrRPCResultFlightCapacity) {
return
}
if err != nil || claim.state != rpcResultAcquireOwner {
t.Errorf("Acquire %d = %#v, %v", i, claim, err)
return
}
mu.Lock()
owners = append(owners, claim.owner)
mu.Unlock()
}(i)
}
close(start)
wg.Wait()
if got := cache.completedEntries.snapshot(); got > 24 || got != int64(len(owners)) {
t.Fatalf("global entry usage=%d owners=%d limit=24", got, len(owners))
}
if got := cache.completedBytes.snapshot(); got > 24 || got != int64(len(owners)) {
t.Fatalf("global byte usage=%d owners=%d limit=24", got, len(owners))
}
for i := 0; i < 4; i++ {
auth := [8]byte{byte(i)}
usage := cache.fairBudget.authSnapshot(auth)
if usage.entries > 8 || usage.bytes > 8 || usage.pending > 4 {
t.Fatalf("auth %d overcommitted: %#v", i, usage)
}
for sessionID := int64(0); sessionID < 8; sessionID++ {
session := cache.fairBudget.sessionSnapshot(auth, sessionID)
if session.entries > 3 || session.bytes > 3 {
t.Fatalf("auth %d session %d overcommitted: %#v", i, sessionID, session)
}
}
}
for _, owner := range owners {
if !owner.Abort() {
t.Fatal("concurrent owner did not abort")
}
}
if cache.completedEntries.snapshot() != 0 || cache.completedBytes.snapshot() != 0 {
t.Fatal("concurrent Abort leaked global fair budget")
}
}
func TestRPCResultCacheConcurrentOwnerPublicationAcrossShards(t *testing.T) {
const publications = 256
now := time.Unix(1000, 0)
cache := newRPCResultCacheWithFairCapacity(func() time.Time { return now }, rpcResultCacheCapacity{
maxPending: publications, maxPendingPerAuth: 4,
globalMaxEntries: publications, globalMaxBytes: publications * 4,
authMaxEntries: 4, authMaxBytes: 16,
sessionMaxEntries: 1, sessionMaxBytes: 4,
})
type publication struct {
auth [8]byte
session int64
msgID int64
owner *rpcResultOwnerLease
}
publicationsByKey := make([]publication, 0, publications)
for i := 0; i < publications; i++ {
auth := [8]byte{byte(i), byte(i >> 8), 0xa5}
sessionID := int64(10_000 + i)
msgID := int64(20_000 + i)
claim, err := cache.Acquire(auth, sessionID, msgID)
if err != nil || claim.state != rpcResultAcquireOwner {
t.Fatalf("Acquire %d = %#v, %v", i, claim, err)
}
publicationsByKey = append(publicationsByKey, publication{
auth: auth, session: sessionID, msgID: msgID, owner: claim.owner,
})
}
start := make(chan struct{})
var wg sync.WaitGroup
for i := range publicationsByKey {
item := publicationsByKey[i]
wg.Add(1)
go func() {
defer wg.Done()
<-start
if !item.owner.CompleteExecution(true) {
t.Errorf("CompleteExecution(%d) lost owner", item.msgID)
return
}
cache.Put(item.auth, item.session, item.msgID, &encodedOutboundMessage{body: []byte{1, 2, 3, 4}})
}()
}
close(start)
wg.Wait()
for _, item := range publicationsByKey {
encoded, ok := cache.Get(item.auth, item.session, item.msgID)
if !ok || encoded == nil || len(encoded.body) != 4 {
t.Fatalf("completed publication %d missing: ok=%v encoded=%#v", item.msgID, ok, encoded)
}
}
if got := cache.completedEntries.snapshot(); got != publications {
t.Fatalf("completed entries=%d, want %d", got, publications)
}
if got := cache.completedBytes.snapshot(); got != publications*4 {
t.Fatalf("completed bytes=%d, want %d", got, publications*4)
}
now = now.Add(rpcResultCacheTTL + time.Second)
cache.expireCompletedResults()
if cache.completedEntries.snapshot() != 0 || cache.completedBytes.snapshot() != 0 {
t.Fatal("parallel publications leaked fair-budget reservations after TTL")
}
}
func BenchmarkRPCResultCacheParallelShardPut(b *testing.B) {
cache := newRPCResultCacheWithFlightLimit(time.Now, rpcResultFlightDefaultMaxPending)
var nextWorker atomic.Uint64
b.ReportAllocs()
b.ResetTimer()
b.RunParallel(func(pb *testing.PB) {
id := nextWorker.Add(1)
auth := [8]byte{
byte(id), byte(id >> 8), byte(id >> 16), byte(id >> 24),
byte(id >> 32), byte(id >> 40), byte(id >> 48), byte(id >> 56),
}
sessionID := int64(id)
msgID := int64(1_000_000 + id)
encoded := &encodedOutboundMessage{body: []byte{1, 2, 3, 4}}
for pb.Next() {
cache.Put(auth, sessionID, msgID, encoded)
}
})
}
func TestRPCResultCacheEntryReservationTransfersAndReturns(t *testing.T) {
now := time.Unix(1000, 0)
cache := newRPCResultCacheWithCapacity(func() time.Time { return now }, 4, 2, 2)
authKeyID := [8]byte{0xa2}
first, err := cache.Acquire(authKeyID, 1, 101)
if err != nil || first.state != rpcResultAcquireOwner || cache.completedEntries.snapshot() != 1 {
t.Fatalf("first pending reservation = %#v entries=%d err=%v", first, cache.completedEntries.snapshot(), err)
}
cache.Put(authKeyID, 1, 101, &encodedOutboundMessage{body: []byte{1}})
if got := cache.completedEntries.snapshot(); got != 1 {
t.Fatalf("pending -> body changed entry count to %d", got)
}
second, err := cache.Acquire(authKeyID, 2, 202)
if err != nil || second.state != rpcResultAcquireOwner || cache.completedEntries.snapshot() != 2 {
t.Fatalf("second pending reservation = %#v entries=%d err=%v", second, cache.completedEntries.snapshot(), err)
}
// The byte budget has only the second owner's one-byte token remaining.
// Publication therefore leaves an identity tombstone, which still owns its
// real process-wide entry slot.
cache.Put(authKeyID, 2, 202, &encodedOutboundMessage{body: []byte{2, 2, 2}})
if got := cache.completedEntries.snapshot(); got != 2 {
t.Fatalf("pending -> tombstone changed entry count to %d", got)
}
if _, err := cache.Acquire(authKeyID, 3, 303); !errors.Is(err, ErrRPCResultFlightCapacity) {
t.Fatalf("third admission at entry limit = %v, want capacity", err)
}
now = now.Add(rpcResultCacheTTL + time.Second)
firstShard := cache.shardIndex(rpcResultCacheKey{authKeyID: authKeyID, sessionID: 1, reqMsgID: 101})
secondShard := cache.shardIndex(rpcResultCacheKey{authKeyID: authKeyID, sessionID: 2, reqMsgID: 202})
thirdMsgID := rpcResultTestMsgIDOutsideShards(t, cache, authKeyID, 3, 303, firstShard, secondShard)
third, err := cache.Acquire(authKeyID, 3, thirdMsgID)
if err != nil || third.state != rpcResultAcquireOwner {
t.Fatalf("admission after global expiry reap = %#v, %v", third, err)
}
if got := cache.completedEntries.snapshot(); got != 1 {
t.Fatalf("expired entries were not returned before new owner: %d", got)
}
if !third.owner.Abort() || cache.completedEntries.snapshot() != 0 {
t.Fatalf("Abort did not return entry reservation: entries=%d", cache.completedEntries.snapshot())
}
}
func TestRPCResultCacheConcurrentGlobalEntryReservationNeverOvercommits(t *testing.T) {
const limit = 8
cache := newRPCResultCacheWithCapacity(time.Now, 128, 1<<20, limit)
authKeyID := [8]byte{0xa3}
var (
wg sync.WaitGroup
mu sync.Mutex
owners []*rpcResultOwnerLease
)
for i := 0; i < 64; i++ {
wg.Add(1)
go func(i int) {
defer wg.Done()
claim, err := cache.Acquire(authKeyID, int64(i+1), int64(1000+i))
if err != nil {
if !errors.Is(err, ErrRPCResultFlightCapacity) {
t.Errorf("Acquire %d: %v", i, err)
}
return
}
if claim.state != rpcResultAcquireOwner {
t.Errorf("Acquire %d state = %d", i, claim.state)
return
}
mu.Lock()
owners = append(owners, claim.owner)
mu.Unlock()
}(i)
}
wg.Wait()
if len(owners) != limit || cache.completedEntries.snapshot() != limit {
t.Fatalf("concurrent owners=%d entries=%d, want %d", len(owners), cache.completedEntries.snapshot(), limit)
}
for _, owner := range owners {
if !owner.Abort() {
t.Fatal("reserved owner failed to abort")
}
}
if got := cache.completedEntries.snapshot(); got != 0 {
t.Fatalf("entry reservations after abort = %d", got)
}
}
func TestRPCResultCacheRoundTripAndTTL(t *testing.T) {
if rpcResultCacheTTL != 331*time.Second {
t.Fatalf("replay TTL = %v, want full 300s past + 30s future window + 1s", rpcResultCacheTTL)
}
now := time.Unix(1000, 0)
cache := newRPCResultCache(func() time.Time { return now })
@ -17,6 +482,12 @@ func TestRPCResultCacheRoundTripAndTTL(t *testing.T) {
t.Fatal("unexpected hit on empty cache")
}
cache.Put(keyID, 5, 7, encoded)
if got := cache.completedEntries.snapshot(); got != 1 {
t.Fatalf("direct Put entry reservation = %d, want 1", got)
}
if usage := cache.fairBudget.sessionSnapshot(keyID, 5); usage.entries != 1 || usage.bytes != 4 || usage.pending != 0 {
t.Fatalf("direct Put session reservation = %#v", usage)
}
got, ok := cache.Get(keyID, 5, 7)
if !ok {
@ -40,29 +511,320 @@ func TestRPCResultCacheRoundTripAndTTL(t *testing.T) {
if _, ok := cache.Get(keyID, 5, 7); ok {
t.Fatal("expected expiry after TTL")
}
if got := cache.completedEntries.snapshot(); got != 0 {
t.Fatalf("direct Put expiry left %d entry reservations", got)
}
if usage := cache.fairBudget.authSnapshot(keyID); usage != (rpcResultBudgetUsage{}) {
t.Fatalf("direct Put expiry leaked auth reservation %#v", usage)
}
}
func TestRPCResultCacheShardTrim(t *testing.T) {
func TestRPCResultCacheDuplicatePutPreservesCompletedExecutionMetadata(t *testing.T) {
cache := newRPCResultCacheWithFlightLimit(time.Now, 1)
keyID := [8]byte{1, 9, 8, 4}
const sessionID, reqMsgID = int64(11), int64(12)
claim, err := cache.Acquire(keyID, sessionID, reqMsgID)
if err != nil || claim.state != rpcResultAcquireOwner {
t.Fatalf("Acquire owner = %#v, %v", claim, err)
}
if !claim.owner.CompleteExecution(true) || !claim.owner.HandOff() {
t.Fatal("complete owner metadata")
}
first := &encodedOutboundMessage{body: []byte{1, 2, 3, 4}, typeID: 42, reqMsgID: reqMsgID}
cache.Put(keyID, sessionID, reqMsgID, first)
second := &encodedOutboundMessage{body: []byte{5, 6, 7, 8}, typeID: 42, reqMsgID: reqMsgID}
cache.Put(keyID, sessionID, reqMsgID, second)
replay, err := cache.Acquire(keyID, sessionID, reqMsgID)
if err != nil || replay.state != rpcResultAcquireCompleted || replay.encoded != second ||
!replay.executionKnown || !replay.executionOK {
t.Fatalf("duplicate Put metadata = %#v, err=%v", replay, err)
}
}
func TestRPCResultCacheShardCapacityNeverEvictsUnexpiredResult(t *testing.T) {
now := time.Unix(1000, 0)
cache := newRPCResultCache(func() time.Time { return now })
var keyID [8]byte
// 同一 (auth_key, session) 固定落在同一 shard；塞超过单 shard 条数上限，最旧的被逐出。
perShard := rpcResultCacheMaxEntries / rpcResultCacheShards
for i := 0; i < perShard+1; i++ {
cache.Put(keyID, 1, int64(100+i), &encodedOutboundMessage{body: []byte{byte(i)}})
firstKey := rpcResultCacheKey{authKeyID: keyID, sessionID: 1, reqMsgID: 100}
shard := cache.shard(firstKey)
shard.mu.Lock()
shard.maxEntries = 1
shard.mu.Unlock()
claim, err := cache.Acquire(keyID, 1, 100)
if err != nil || claim.state != rpcResultAcquireOwner {
t.Fatalf("first admission = %#v, %v", claim, err)
}
if _, ok := cache.Get(keyID, 1, 100); ok {
t.Fatal("oldest entry should have been evicted by per-shard entry limit")
first := &encodedOutboundMessage{body: []byte{1}}
cache.Put(keyID, 1, 100, first)
secondMsgID := rpcResultTestMsgIDForShard(t, cache, keyID, 1, 101, cache.shardIndex(firstKey))
if _, err := cache.Acquire(keyID, 1, secondMsgID); !errors.Is(err, ErrRPCResultFlightCapacity) {
t.Fatalf("full-shard admission = %v, want capacity", err)
}
if _, ok := cache.Get(keyID, 1, int64(100+perShard)); !ok {
t.Fatal("newest entry should survive")
if got, ok := cache.Get(keyID, 1, 100); !ok || got != first {
t.Fatalf("unexpired first result was displaced: got=%p ok=%v", got, ok)
}
// 单条超过单 shard 字节预算的结果不入缓存。
huge := &encodedOutboundMessage{body: make([]byte, rpcResultCacheMaxBytes/rpcResultCacheShards+1)}
cache.Put(keyID, 2, 999, huge)
if _, ok := cache.Get(keyID, 2, 999); ok {
t.Fatal("oversized entry should be rejected")
now = now.Add(rpcResultCacheTTL + time.Second)
claim, err = cache.Acquire(keyID, 1, secondMsgID)
if err != nil || claim.state != rpcResultAcquireOwner {
t.Fatalf("admission after expiry = %#v, %v", claim, err)
}
claim.owner.Abort()
}
func TestRPCResultCacheGlobalByteCapacityNeverEvictsUnexpiredResults(t *testing.T) {
now := time.Unix(1000, 0)
cache := newRPCResultCacheWithLimits(func() time.Time { return now }, 32, 10)
var keyID [8]byte
// Five two-byte results consume the global budget. The sixth admission must
// fail bounded; none of the retained results may be sacrificed for it.
for sessionID := int64(1); sessionID <= 5; sessionID++ {
claim, err := cache.Acquire(keyID, sessionID, 100+sessionID)
if err != nil || claim.state != rpcResultAcquireOwner {
t.Fatalf("admission %d = %#v, %v", sessionID, claim, err)
}
cache.Put(keyID, sessionID, 100+sessionID, &encodedOutboundMessage{body: []byte{1, 2}})
}
if got := cache.completedBytes.snapshot(); got != 10 {
t.Fatalf("completed bytes at capacity = %d, want 10", got)
}
if _, err := cache.Acquire(keyID, 6, 106); !errors.Is(err, ErrRPCResultFlightCapacity) {
t.Fatalf("byte-full admission = %v, want capacity", err)
}
for sessionID := int64(1); sessionID <= 5; sessionID++ {
if _, ok := cache.Get(keyID, sessionID, 100+sessionID); !ok {
t.Fatalf("unexpired result %d was evicted", sessionID)
}
}
}
func TestRPCResultCacheByteBudgetReturnsOnReplaceExpiryAndCapacity(t *testing.T) {
now := time.Unix(1000, 0)
cache := newRPCResultCacheWithLimits(func() time.Time { return now }, 32, 32)
var keyID [8]byte
cache.Put(keyID, 1, 101, &encodedOutboundMessage{body: make([]byte, 4)})
cache.Put(keyID, 1, 101, &encodedOutboundMessage{body: make([]byte, 7)})
if got := cache.completedBytes.snapshot(); got != 7 {
t.Fatalf("completed bytes after growing replacement = %d, want 7", got)
}
if usage := cache.fairBudget.sessionSnapshot(keyID, 1); usage.entries != 1 || usage.bytes != 7 {
t.Fatalf("replacement fair reservation after growth = %#v", usage)
}
cache.Put(keyID, 1, 101, &encodedOutboundMessage{body: make([]byte, 2)})
if got := cache.completedBytes.snapshot(); got != 2 {
t.Fatalf("completed bytes after shrinking replacement = %d, want 2", got)
}
if usage := cache.fairBudget.sessionSnapshot(keyID, 1); usage.entries != 1 || usage.bytes != 2 {
t.Fatalf("replacement fair reservation after shrink = %#v", usage)
}
now = now.Add(rpcResultCacheTTL + time.Second)
if _, ok := cache.Get(keyID, 1, 101); ok {
t.Fatal("replacement should expire")
}
if got := cache.completedBytes.snapshot(); got != 0 {
t.Fatalf("completed bytes after expiry = %d, want 0", got)
}
key := rpcResultCacheKey{authKeyID: keyID, sessionID: 2, reqMsgID: 201}
shard := cache.shard(key)
shard.mu.Lock()
shard.maxEntries = 1
shard.mu.Unlock()
claim, err := cache.Acquire(keyID, 2, 201)
if err != nil || claim.state != rpcResultAcquireOwner {
t.Fatalf("entry-capacity first admission = %#v, %v", claim, err)
}
cache.Put(keyID, 2, 201, &encodedOutboundMessage{body: make([]byte, 3)})
secondMsgID := rpcResultTestMsgIDForShard(t, cache, keyID, 2, 202, cache.shardIndex(key))
if _, err := cache.Acquire(keyID, 2, secondMsgID); !errors.Is(err, ErrRPCResultFlightCapacity) {
t.Fatalf("entry-capacity second admission = %v", err)
}
if got := cache.completedBytes.snapshot(); got != 3 {
t.Fatalf("completed bytes after capacity rejection = %d, want 3", got)
}
if _, ok := cache.Get(keyID, 2, 201); !ok {
t.Fatal("capacity rejection displaced the first result")
}
}
func TestRPCResultCachePublicationOverflowLeavesReplayCapacityTombstone(t *testing.T) {
now := time.Unix(1000, 0)
cache := newRPCResultCacheWithLimits(func() time.Time { return now }, 32, 4)
var keyID [8]byte
claim, err := cache.Acquire(keyID, 1, 101)
if err != nil || claim.state != rpcResultAcquireOwner {
t.Fatalf("owner admission = %#v, %v", claim, err)
}
claim.owner.CompleteExecution(true)
tooLarge := &encodedOutboundMessage{body: make([]byte, 5)}
cache.Put(keyID, 1, 101, tooLarge)
if got := cache.completedBytes.snapshot(); got != 1 {
t.Fatalf("tombstone bytes = %d, want 1", got)
}
if _, ok := cache.Get(keyID, 1, 101); ok {
t.Fatal("capacity tombstone must not masquerade as a replayable body")
}
if _, err := cache.Acquire(keyID, 1, 101); !errors.Is(err, ErrRPCResultFlightCapacity) {
t.Fatalf("duplicate after publication overflow = %v, want capacity", err)
}
now = now.Add(rpcResultCacheTTL + time.Second)
retry, err := cache.Acquire(keyID, 1, 101)
if err != nil || retry.state != rpcResultAcquireOwner {
t.Fatalf("admission after tombstone expiry = %#v, %v", retry, err)
}
retry.owner.Abort()
}
func TestRPCResultCacheByteCapacityReclaimsExpiredAcrossShards(t *testing.T) {
now := time.Unix(1000, 0)
cache := newRPCResultCacheWithLimits(func() time.Time { return now }, 32, 2)
var keyID [8]byte
first, err := cache.Acquire(keyID, 1, 101)
if err != nil || first.state != rpcResultAcquireOwner {
t.Fatalf("first admission = %#v, %v", first, err)
}
cache.Put(keyID, 1, 101, &encodedOutboundMessage{body: []byte{1, 2}})
if got := cache.completedBytes.snapshot(); got != 2 {
t.Fatalf("full budget = %d, want 2", got)
}
// Select a key in another full-key shard. Its failed one-byte reservation
// must trigger the cold-path global expiry reap before returning capacity.
now = now.Add(rpcResultCacheTTL + time.Second)
firstKey := rpcResultCacheKey{authKeyID: keyID, sessionID: 1, reqMsgID: 101}
secondMsgID := rpcResultTestMsgIDOutsideShard(t, cache, keyID, 2, 202, cache.shardIndex(firstKey))
second, err := cache.Acquire(keyID, 2, secondMsgID)
if err != nil || second.state != rpcResultAcquireOwner {
t.Fatalf("cross-shard admission after expiry = %#v, %v", second, err)
}
second.owner.Abort()
if got := cache.completedBytes.snapshot(); got != 0 {
t.Fatalf("bytes after expired reap and abort = %d, want 0", got)
}
}
func TestRPCResultCacheServerOptionsPropagateFairLimits(t *testing.T) {
sessionBytes := int64(maxOutboundBodyBytes)
s := New(Options{
RPCGlobalMaxTasks: 6,
RPCResultCacheMaxEntries: 12,
RPCResultCacheMaxBytes: sessionBytes + 2048,
RPCResultCacheAuthMaxEntries: 8,
RPCResultCacheAuthMaxBytes: sessionBytes + 1024,
RPCResultCacheSessionMaxEntries: 4,
RPCResultCacheSessionMaxBytes: sessionBytes,
RPCResultPendingPerAuth: 3,
})
if s.rpcResults.completedEntries.max != 12 || s.rpcResults.completedBytes.max != sessionBytes+2048 {
t.Fatalf("global option propagation = %d/%d", s.rpcResults.completedEntries.max, s.rpcResults.completedBytes.max)
}
budget := s.rpcResults.fairBudget
if budget.authLimit.entries != 8 || budget.authLimit.bytes != sessionBytes+1024 ||
budget.sessionLimit.entries != 4 || budget.sessionLimit.bytes != sessionBytes || budget.pendingPerAuth != 3 {
t.Fatalf("fair option propagation = auth:%#v session:%#v pending:%d",
budget.authLimit, budget.sessionLimit, budget.pendingPerAuth)
}
}
func TestRPCResultCacheServerOptionsFailFast(t *testing.T) {
base := Options{
RPCGlobalMaxTasks: 6,
RPCResultCacheMaxEntries: 12,
RPCResultCacheMaxBytes: 64 << 20,
RPCResultCacheAuthMaxEntries: 8,
RPCResultCacheAuthMaxBytes: 32 << 20,
RPCResultCacheSessionMaxEntries: 4,
RPCResultCacheSessionMaxBytes: 16 << 20,
RPCResultPendingPerAuth: 3,
}
tests := []struct {
name string
mutate func(*Options)
}{
{name: "entry hierarchy", mutate: func(o *Options) { o.RPCResultCacheAuthMaxEntries = 13 }},
{name: "body does not fit session", mutate: func(o *Options) { o.RPCResultCacheSessionMaxBytes = maxOutboundBodyBytes - 1 }},
{name: "pending hierarchy", mutate: func(o *Options) { o.RPCResultPendingPerAuth = 7 }},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
opts := base
test.mutate(&opts)
defer func() {
if recover() == nil {
t.Fatal("New accepted invalid rpc_result cache options")
}
}()
_ = New(opts)
})
}
}
func rpcResultTestMsgIDForShard(
t *testing.T,
cache *rpcResultCache,
authKeyID [8]byte,
sessionID, start int64,
target uint64,
) int64 {
t.Helper()
for msgID := start; msgID < start+1_000_000; msgID++ {
key := rpcResultCacheKey{authKeyID: authKeyID, sessionID: sessionID, reqMsgID: msgID}
if cache.shardIndex(key) == target {
return msgID
}
}
t.Fatal("failed to find rpc_result key for target shard")
return 0
}
func rpcResultTestMsgIDOutsideShard(
t *testing.T,
cache *rpcResultCache,
authKeyID [8]byte,
sessionID, start int64,
excluded uint64,
) int64 {
t.Helper()
for msgID := start; msgID < start+1_000_000; msgID++ {
key := rpcResultCacheKey{authKeyID: authKeyID, sessionID: sessionID, reqMsgID: msgID}
if cache.shardIndex(key) != excluded {
return msgID
}
}
t.Fatal("failed to find rpc_result key outside excluded shard")
return 0
}
func rpcResultTestMsgIDOutsideShards(
t *testing.T,
cache *rpcResultCache,
authKeyID [8]byte,
sessionID, start int64,
excluded ...uint64,
) int64 {
t.Helper()
for msgID := start; msgID < start+1_000_000; msgID++ {
key := rpcResultCacheKey{authKeyID: authKeyID, sessionID: sessionID, reqMsgID: msgID}
index := cache.shardIndex(key)
allowed := true
for _, blocked := range excluded {
if index == blocked {
allowed = false
break
}
}
if allowed {
return msgID
}
}
t.Fatal("failed to find rpc_result key outside excluded shards")
return 0
}

288
internal/mtprotoedge/rpc_result_delivery_test.go
View file

@ -8,9 +8,16 @@ import (
"testing"
"time"
"github.com/gotd/td/bin"
"github.com/gotd/td/proto"
"github.com/gotd/td/tg"
"github.com/iamxvbaba/td/bin"
"github.com/iamxvbaba/td/clock"
"github.com/iamxvbaba/td/proto"
"github.com/iamxvbaba/td/tg"
"go.uber.org/zap/zaptest"
appupdates "telesrv/internal/app/updates"
"telesrv/internal/domain"
rpchandler "telesrv/internal/rpc"
"telesrv/internal/store/memory"
)
type blockingCloseRPCResultTransport struct {
@ -19,6 +26,60 @@ type blockingCloseRPCResultTransport struct {
once sync.Once
}
type identifiedRouterHandler struct {
router *rpchandler.Router
userID int64
}
type failingRPCResultEncoder struct {
err error
}
func (e failingRPCResultEncoder) Encode(*bin.Buffer) error { return e.err }
type encodeFailingIdentifiedRouterHandler struct {
identifiedRouterHandler
err error
}
func (h encodeFailingIdentifiedRouterHandler) Dispatch(ctx context.Context, authKeyID [8]byte, sessionID int64, b *bin.Buffer) (bin.Encoder, error) {
result, err := h.identifiedRouterHandler.Dispatch(ctx, authKeyID, sessionID, b)
if err != nil {
return nil, err
}
if result == nil {
return nil, nil
}
return failingRPCResultEncoder{err: h.err}, nil
}
func (h encodeFailingIdentifiedRouterHandler) DispatchWithMethod(ctx context.Context, authKeyID [8]byte, sessionID int64, b *bin.Buffer) (bin.Encoder, string, error) {
result, method, err := h.identifiedRouterHandler.DispatchWithMethod(ctx, authKeyID, sessionID, b)
if err != nil {
return nil, method, err
}
if result == nil {
return nil, method, nil
}
return failingRPCResultEncoder{err: h.err}, method, nil
}
func (h identifiedRouterHandler) requestContext(ctx context.Context) context.Context {
return rpchandler.WithClientInfo(rpchandler.WithUserID(ctx, h.userID), rpchandler.ClientInfo{Type: rpchandler.ClientTypeTDesktop})
}
func (h identifiedRouterHandler) Dispatch(ctx context.Context, authKeyID [8]byte, sessionID int64, b *bin.Buffer) (bin.Encoder, error) {
return h.router.Dispatch(h.requestContext(ctx), authKeyID, sessionID, b)
}
func (h identifiedRouterHandler) DispatchWithMethod(ctx context.Context, authKeyID [8]byte, sessionID int64, b *bin.Buffer) (bin.Encoder, string, error) {
return h.router.DispatchWithMethod(h.requestContext(ctx), authKeyID, sessionID, b)
}
func (h identifiedRouterHandler) NegotiatedLayer(authKeyID [8]byte, sessionID int64) (int, bool) {
return h.router.NegotiatedLayer(authKeyID, sessionID)
}
func newBlockingCloseRPCResultTransport() *blockingCloseRPCResultTransport {
return &blockingCloseRPCResultTransport{started: make(chan struct{}), release: make(chan struct{})}
}
@ -45,7 +106,7 @@ func TestRPCResultCachePublishesOnlyAfterPhysicalWrite(t *testing.T) {
}
done := make(chan error, 1)
go func() {
done <- s.sendResult(context.Background(), c, reqMsgID, &tg.Config{ThisDC: 2})
done <- s.sendResult(context.Background(), c, reqMsgID, exactTestRPCResult(&tg.Config{ThisDC: 2}))
}()
select {
case <-tr.started:
@ -78,6 +139,217 @@ func TestRPCResultCachePublishesOnlyAfterPhysicalWrite(t *testing.T) {
}
}
func TestRPCResultPostResponseHookWaitsForPhysicalWrite(t *testing.T) {
tr := newGatedRequiredControlTransport(nil)
s := New(Options{WriteTimeout: time.Second})
key := newTestAuthKey(t)
c := s.newConn(tr, key, 74006, 1)
legacyCanonicalTestConn(t, c)
t.Cleanup(c.ForceClose)
reqMsgID := proto.NewMessageIDGen(time.Now).New(proto.MessageFromClient)
claim, err := s.rpcResults.Acquire(key.ID, c.sessionID, reqMsgID)
if err != nil || claim.state != rpcResultAcquireOwner {
t.Fatalf("initial flight owner = %+v err=%v", claim, err)
}
delivered := make(chan struct{})
if err := s.publishRPCResult(c, reqMsgID, "updates.getState", claim.owner, exactTestRPCResult(&tg.UpdatesState{}), func() {
close(delivered)
}); err != nil {
t.Fatalf("publish rpc_result: %v", err)
}
select {
case <-tr.started:
case <-time.After(time.Second):
t.Fatal("rpc_result did not reach physical writer")
}
select {
case <-delivered:
t.Fatal("post-response hook ran while physical writer was blocked")
case <-time.After(20 * time.Millisecond):
}
tr.unblock()
select {
case <-delivered:
case <-time.After(time.Second):
t.Fatal("post-response hook did not run after physical write")
}
}
func TestRouterUpdateCursorCommitsOnlyAfterHandleRPCPhysicalWrite(t *testing.T) {
const userID int64 = 1000000301
events := memory.NewUpdateEventStore()
states := memory.NewUpdateStateStore()
if err := events.Append(context.Background(), userID, domain.UpdateEvent{
UserID: userID, Type: domain.UpdateEventNewMessage,
Pts: 1, PtsCount: 1, Date: 1700000301,
Message: domain.Message{ID: 1, OwnerUserID: userID},
}); err != nil {
t.Fatalf("seed update event: %v", err)
}
router := rpchandler.New(rpchandler.Config{}, rpchandler.Deps{
Updates: appupdates.NewService(states, events),
}, zaptest.NewLogger(t), clock.System)
handler := identifiedRouterHandler{router: router, userID: userID}
tr := newGatedRequiredControlTransport(nil)
s := New(Options{legacyRPC: handler, WriteTimeout: time.Second})
key := newTestAuthKey(t)
c := s.newConn(tr, key, 74007, 1)
t.Cleanup(c.ForceClose)
reqMsgID := proto.NewMessageIDGen(time.Now).New(proto.MessageFromClient)
claim, err := s.rpcResults.Acquire(key.ID, c.sessionID, reqMsgID)
if err != nil || claim.state != rpcResultAcquireOwner {
t.Fatalf("acquire flight = %+v err=%v", claim, err)
}
var request bin.Buffer
if err := (&tg.UpdatesGetStateRequest{}).Encode(&request); err != nil {
t.Fatalf("encode getState: %v", err)
}
handled := make(chan error, 1)
go func() {
handled <- s.handleRPC(context.Background(), c, reqMsgID, "updates.getState", &request, claim.owner)
}()
select {
case <-tr.started:
case <-time.After(time.Second):
t.Fatal("getState rpc_result did not reach physical writer")
}
select {
case err := <-handled:
if err != nil {
t.Fatalf("handleRPC: %v", err)
}
case <-time.After(time.Second):
t.Fatal("handleRPC remained coupled to physical write")
}
if state, found, err := states.Get(context.Background(), key.ID, userID); err != nil || found {
t.Fatalf("confirmed before physical write = %+v/%v err=%v", state, found, err)
}
if state, found := states.ObservedClientState(key.ID, userID); found {
t.Fatalf("observed before physical write = %+v/%v", state, found)
}
tr.unblock()
deadline := time.Now().Add(time.Second)
for time.Now().Before(deadline) {
confirmed, found, err := states.Get(context.Background(), key.ID, userID)
observed, observedFound := states.ObservedClientState(key.ID, userID)
if err == nil && found && confirmed.Pts == 1 && observedFound && observed.Pts == 1 {
return
}
time.Sleep(time.Millisecond)
}
confirmed, found, err := states.Get(context.Background(), key.ID, userID)
observed, observedFound := states.ObservedClientState(key.ID, userID)
t.Fatalf("post-write cursor = confirmed:%+v/%v err=%v observed:%+v/%v", confirmed, found, err, observed, observedFound)
}
func TestRouterUpdateCursorDoesNotCommitAfterPhysicalWriteFailure(t *testing.T) {
const userID int64 = 1000000302
events := memory.NewUpdateEventStore()
states := memory.NewUpdateStateStore()
if err := events.Append(context.Background(), userID, domain.UpdateEvent{
UserID: userID, Type: domain.UpdateEventNewMessage,
Pts: 1, PtsCount: 1, Date: 1700000302,
Message: domain.Message{ID: 1, OwnerUserID: userID},
}); err != nil {
t.Fatalf("seed update event: %v", err)
}
router := rpchandler.New(rpchandler.Config{}, rpchandler.Deps{
Updates: appupdates.NewService(states, events),
}, zaptest.NewLogger(t), clock.System)
tr := &failAfterTransport{}
tr.failAt.Store(1)
s := New(Options{legacyRPC: identifiedRouterHandler{router: router, userID: userID}, WriteTimeout: time.Second})
key := newTestAuthKey(t)
c := s.newConn(tr, key, 74008, 1)
t.Cleanup(c.ForceClose)
reqMsgID := proto.NewMessageIDGen(time.Now).New(proto.MessageFromClient)
claim, err := s.rpcResults.Acquire(key.ID, c.sessionID, reqMsgID)
if err != nil || claim.state != rpcResultAcquireOwner {
t.Fatalf("acquire flight = %+v err=%v", claim, err)
}
var request bin.Buffer
if err := (&tg.UpdatesGetStateRequest{}).Encode(&request); err != nil {
t.Fatalf("encode getState: %v", err)
}
if err := s.handleRPC(context.Background(), c, reqMsgID, "updates.getState", &request, claim.owner); err != nil {
t.Fatalf("handleRPC admission: %v", err)
}
deadline := time.Now().Add(time.Second)
replayable := false
for time.Now().Before(deadline) {
if cached, ok := s.rpcResults.Get(key.ID, c.sessionID, reqMsgID); ok && cached.deliveryState() == rpcResultDeliveryReplayable {
replayable = true
break
}
time.Sleep(time.Millisecond)
}
if !replayable {
t.Fatal("failed physical result did not become replayable")
}
if state, found, err := states.Get(context.Background(), key.ID, userID); err != nil || found {
t.Fatalf("failed write advanced confirmed = %+v/%v err=%v", state, found, err)
}
if state, found := states.ObservedClientState(key.ID, userID); found {
t.Fatalf("failed write advanced observed = %+v/%v", state, found)
}
}
func TestRouterUpdateCursorDoesNotCommitWhenResultEncodingFails(t *testing.T) {
const userID int64 = 1000000303
events := memory.NewUpdateEventStore()
states := memory.NewUpdateStateStore()
if err := events.Append(context.Background(), userID, domain.UpdateEvent{
UserID: userID, Type: domain.UpdateEventNewMessage,
Pts: 1, PtsCount: 1, Date: 1700000303,
Message: domain.Message{ID: 1, OwnerUserID: userID},
}); err != nil {
t.Fatalf("seed update event: %v", err)
}
router := rpchandler.New(rpchandler.Config{}, rpchandler.Deps{
Updates: appupdates.NewService(states, events),
}, zaptest.NewLogger(t), clock.System)
handler := encodeFailingIdentifiedRouterHandler{
identifiedRouterHandler: identifiedRouterHandler{router: router, userID: userID},
err: errors.New("encode result"),
}
tr := &collectingSessionTransport{}
s := New(Options{legacyRPC: handler, WriteTimeout: time.Second})
key := newTestAuthKey(t)
c := s.newConn(tr, key, 74009, 1)
t.Cleanup(c.ForceClose)
reqMsgID := proto.NewMessageIDGen(time.Now).New(proto.MessageFromClient)
claim, err := s.rpcResults.Acquire(key.ID, c.sessionID, reqMsgID)
if err != nil || claim.state != rpcResultAcquireOwner {
t.Fatalf("acquire flight = %+v err=%v", claim, err)
}
var request bin.Buffer
if err := (&tg.UpdatesGetStateRequest{}).Encode(&request); err != nil {
t.Fatalf("encode getState: %v", err)
}
if err := s.handleRPC(context.Background(), c, reqMsgID, "updates.getState", &request, claim.owner); err != nil {
t.Fatalf("handleRPC encoding fallback: %v", err)
}
deadline := time.Now().Add(time.Second)
for time.Now().Before(deadline) {
if cached, ok := s.rpcResults.Get(key.ID, c.sessionID, reqMsgID); ok && cached.deliveryState() == rpcResultDeliveryDelivered {
break
}
time.Sleep(time.Millisecond)
}
if len(tr.snapshot()) == 0 {
t.Fatal("INTERNAL fallback was not physically delivered")
}
if state, found, err := states.Get(context.Background(), key.ID, userID); err != nil || found {
t.Fatalf("encoding failure advanced confirmed = %+v/%v err=%v", state, found, err)
}
if state, found := states.ObservedClientState(key.ID, userID); found {
t.Fatalf("encoding failure advanced observed = %+v/%v", state, found)
}
}
func TestRPCResultPrewriteFailureFencesConnBeforeCachePublication(t *testing.T) {
tr := &collectingSessionTransport{}
s := New(Options{WriteTimeout: 20 * time.Millisecond})
@ -93,7 +365,7 @@ func TestRPCResultPrewriteFailureFencesConnBeforeCachePublication(t *testing.T)
t.Fatalf("initial flight owner = %+v err=%v", owner, err)
}
err = s.sendResult(context.Background(), c, reqMsgID, &tg.Config{ThisDC: 2})
err = s.sendResult(context.Background(), c, reqMsgID, exactTestRPCResult(&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)
@ -132,7 +404,7 @@ func TestRPCResultFailureAfterIntentionalTerminalDoesNotCloseTransferLease(t *te
// 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})
err = s.sendResult(context.Background(), oldConn, reqMsgID, exactTestRPCResult(&tg.Config{ThisDC: 2}))
if !errors.Is(err, ErrOutboundTrackedBudget) {
t.Fatalf("late result error = %v, want ErrOutboundTrackedBudget", err)
}
@ -171,7 +443,9 @@ func TestRPCResultPublishesBeforePathologicalPhysicalCloseReturns(t *testing.T)
}
done := make(chan error, 1)
go func() { done <- s.sendResult(context.Background(), c, reqMsgID, &tg.Config{ThisDC: 2}) }()
go func() {
done <- s.sendResult(context.Background(), c, reqMsgID, exactTestRPCResult(&tg.Config{ThisDC: 2}))
}()
select {
case <-tr.started:
case <-time.After(time.Second):

159
internal/mtprotoedge/rpc_result_egress.go
View file

@ -2,11 +2,15 @@ package mtprotoedge
import (
"context"
"errors"
"log"
"runtime/debug"
"strings"
"sync"
"sync/atomic"
"github.com/gotd/td/bin"
"github.com/gotd/td/proto"
"github.com/iamxvbaba/td/bin"
"github.com/iamxvbaba/td/proto"
)
const (
@ -21,30 +25,143 @@ const (
var rpcResultGZIPSlots = make(chan struct{}, rpcResultGZIPConcurrency)
var (
rpcDeliveryHooksOnce sync.Once
rpcDeliveryHooks chan func()
var defaultRPCDeliveryHookExecutor = newRPCDeliveryHookExecutor(rpcDeliveryHookConcurrency, rpcDeliveryHookQueueSize)
// ErrRPCDeliveryHookCapacity means an RPC result with a delivery-dependent
// transition cannot reserve reliable executor capacity. The result must not be
// written: its connection is fenced and the immutable result stays replayable.
var ErrRPCDeliveryHookCapacity = errors.New("mtproto rpc delivery hook capacity exhausted")
type rpcDeliveryHookTicketState uint32
const (
rpcDeliveryHookTicketReserved rpcDeliveryHookTicketState = iota + 1
rpcDeliveryHookTicketQueued
rpcDeliveryHookTicketReleased
rpcDeliveryHookTicketDone
)
// scheduleRPCDeliveryHook keeps database/update follow-up work off the sole
// socket writer. Hooks are internal, bounded and timeout-aware at registration
// sites; the fixed worker set prevents one delivered result from stalling every
// subsequent frame on that connection.
func scheduleRPCDeliveryHook(fn func()) {
if fn == nil {
type rpcDeliveryHookTicket struct {
executor *rpcDeliveryHookExecutor
state atomic.Uint32
job rpcDeliveryHookJob
}
type rpcDeliveryHookJob struct {
next *rpcDeliveryHookJob
ticket *rpcDeliveryHookTicket
fn func()
}
// rpcDeliveryHookExecutor has process lifetime. Capacity bounds queued plus
// running hooks; every physical write reserves a ticket before admission. A
// successful writer therefore performs only one short O(1) queue append and
// never waits for capacity or hook work. Failed writes release their ticket,
// while the shared logical coordinator remains eligible for a later replay.
type rpcDeliveryHookExecutor struct {
slots chan struct{}
mu sync.Mutex
cond *sync.Cond
head *rpcDeliveryHookJob
tail *rpcDeliveryHookJob
panics atomic.Uint64
}
func newRPCDeliveryHookExecutor(workers, capacity int) *rpcDeliveryHookExecutor {
if workers <= 0 {
workers = 1
}
if capacity < workers {
capacity = workers
}
e := &rpcDeliveryHookExecutor{slots: make(chan struct{}, capacity)}
e.cond = sync.NewCond(&e.mu)
for range workers {
go e.run()
}
return e
}
func (e *rpcDeliveryHookExecutor) reserve() (*rpcDeliveryHookTicket, bool) {
if e == nil {
return nil, false
}
select {
case e.slots <- struct{}{}:
ticket := &rpcDeliveryHookTicket{executor: e}
ticket.state.Store(uint32(rpcDeliveryHookTicketReserved))
return ticket, true
default:
return nil, false
}
}
func (t *rpcDeliveryHookTicket) release() {
if t == nil || t.executor == nil || !t.state.CompareAndSwap(
uint32(rpcDeliveryHookTicketReserved), uint32(rpcDeliveryHookTicketReleased),
) {
return
}
rpcDeliveryHooksOnce.Do(func() {
rpcDeliveryHooks = make(chan func(), rpcDeliveryHookQueueSize)
for range rpcDeliveryHookConcurrency {
go func() {
for hook := range rpcDeliveryHooks {
hook()
}
}()
<-t.executor.slots
}
func (t *rpcDeliveryHookTicket) submit(fn func()) bool {
if t == nil || t.executor == nil || fn == nil || !t.state.CompareAndSwap(
uint32(rpcDeliveryHookTicketReserved), uint32(rpcDeliveryHookTicketQueued),
) {
return false
}
t.job.ticket = t
t.job.fn = fn
t.executor.enqueue(&t.job)
return true
}
func (e *rpcDeliveryHookExecutor) enqueue(job *rpcDeliveryHookJob) {
e.mu.Lock()
if e.tail == nil {
e.head = job
} else {
e.tail.next = job
}
e.tail = job
e.cond.Signal()
e.mu.Unlock()
}
func (e *rpcDeliveryHookExecutor) run() {
for {
e.mu.Lock()
for e.head == nil {
e.cond.Wait()
}
})
rpcDeliveryHooks <- fn
job := e.head
e.head = job.next
if e.head == nil {
e.tail = nil
}
job.next = nil
e.mu.Unlock()
e.runOne(job)
}
}
func (e *rpcDeliveryHookExecutor) runOne(job *rpcDeliveryHookJob) {
defer func() {
if recovered := recover(); recovered != nil {
e.panics.Add(1)
log.Printf("mtprotoedge: rpc delivery hook panic: %v\n%s", recovered, debug.Stack())
}
if job != nil && job.ticket != nil {
job.ticket.state.Store(uint32(rpcDeliveryHookTicketDone))
<-e.slots
}
}()
if job != nil && job.fn != nil {
job.fn()
}
}
// encodeAdaptiveRPCResultInner returns either the original layer-specific TL

389
internal/mtprotoedge/rpc_result_egress_test.go
View file

@ -4,16 +4,17 @@ import (
"bytes"
"context"
"crypto/rand"
"errors"
"io"
"sync"
"sync/atomic"
"testing"
"time"
"github.com/gotd/td/bin"
"github.com/gotd/td/crypto"
"github.com/gotd/td/proto"
"github.com/gotd/td/tg"
"github.com/iamxvbaba/td/bin"
"github.com/iamxvbaba/td/crypto"
"github.com/iamxvbaba/td/proto"
"github.com/iamxvbaba/td/tg"
)
type opaqueRPCResult struct{ body []byte }
@ -26,9 +27,9 @@ func (o opaqueRPCResult) Encode(b *bin.Buffer) error {
func TestEncodeRPCResultUsesAdaptiveGZIP(t *testing.T) {
s := New(Options{})
c := &Conn{metrics: NopMetrics{}}
c := legacyCanonicalTestConn(t, &Conn{metrics: NopMetrics{}})
large := &tg.DataJSON{Data: string(bytes.Repeat([]byte("sticker-metadata-"), 16<<10))}
encoded, err := s.encodeRPCResult(c, 123, large)
encoded, err := s.encodeRPCResult(c, 123, exactTestRPCResult(large))
if err != nil {
t.Fatalf("encode compressed rpc_result: %v", err)
}
@ -57,12 +58,12 @@ func TestEncodeRPCResultUsesAdaptiveGZIP(t *testing.T) {
func TestEncodeRPCResultKeepsIncompressibleBodyRaw(t *testing.T) {
s := New(Options{})
c := &Conn{metrics: NopMetrics{}}
c := legacyCanonicalTestConn(t, &Conn{metrics: NopMetrics{}})
raw := make([]byte, 96<<10)
if _, err := rand.Read(raw); err != nil {
t.Fatalf("random body: %v", err)
}
encoded, err := s.encodeRPCResult(c, 456, opaqueRPCResult{body: raw})
encoded, err := s.encodeRPCResult(c, 456, exactTestRPCResult(opaqueRPCResult{body: raw}))
if err != nil {
t.Fatalf("encode incompressible rpc_result: %v", err)
}
@ -79,6 +80,349 @@ func TestEncodeRPCResultKeepsIncompressibleBodyRaw(t *testing.T) {
}
}
func TestEncodeRPCResultReservedChargesBodyBeforeReturning(t *testing.T) {
budget := newOutboundTrackedBudget(1 << 20)
c := legacyCanonicalTestConn(t, &Conn{metrics: NopMetrics{}, outboundTrackedBudget: budget})
s := New(Options{})
encoded, reserved, err := s.encodeRPCResultReservedContext(
context.Background(), c, 789, exactTestRPCResult(&tg.DataJSON{Data: "bounded"}),
)
if err != nil {
t.Fatalf("encode reserved rpc_result: %v", err)
}
if reserved == nil {
t.Fatal("encode returned no retained-byte reservation")
}
if got, want := budget.used.Load(), int64(len(encoded.body)); got != want {
t.Fatalf("reserved bytes = %d, want encoded body %d", got, want)
}
reserved.release()
if got := budget.used.Load(); got != 0 {
t.Fatalf("reserved bytes after release = %d, want 0", got)
}
}
func TestEncodeRPCResultReservedDropsBodyOnBudgetTimeout(t *testing.T) {
const maxBytes = 1 << 20
budget := newOutboundTrackedBudget(maxBytes)
if !budget.reserve(maxBytes) {
t.Fatal("saturate outbound body budget")
}
c := legacyCanonicalTestConn(t, &Conn{metrics: NopMetrics{}, outboundTrackedBudget: budget})
s := New(Options{})
ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)
defer cancel()
encoded, reserved, err := s.encodeRPCResultReservedContext(
ctx, c, 790, exactTestRPCResult(&tg.DataJSON{Data: "must-not-escape"}),
)
if !errors.Is(err, context.DeadlineExceeded) {
t.Fatalf("saturated reservation error = %v, want deadline exceeded", err)
}
if encoded != nil || reserved != nil {
t.Fatalf("untracked result escaped encode slot: encoded=%p reserved=%p", encoded, reserved)
}
if got := len(outboundEncodeSlots); got != 0 {
t.Fatalf("encode slots retained after timeout = %d, want 0", got)
}
if got := budget.snapshot(); got != maxBytes {
t.Fatalf("primary budget after timeout = %d, want saturated %d", got, maxBytes)
}
budget.release(maxBytes)
if got := budget.snapshot(); got != 0 {
t.Fatalf("primary budget after release = %d, want 0", got)
}
}
func TestEncodeRPCResultFailedRetentionHandoffDropsBodyInSlot(t *testing.T) {
budget := newOutboundTrackedBudget(1)
if !budget.reserve(1) {
t.Fatal("saturate outbound body budget")
}
c := legacyCanonicalTestConn(t, &Conn{metrics: NopMetrics{}, outboundTrackedBudget: budget})
s := New(Options{})
observedInSlot := false
encoded, reserved, retained, err := s.encodeRPCResultReservedWithHandoffContext(
context.Background(), c, 791, exactTestRPCResult(&tg.DataJSON{Data: "handoff-fails"}),
func(body *encodedOutboundMessage, admissionErr error) error {
observedInSlot = body != nil && len(body.body) > 0 && len(outboundEncodeSlots) > 0 &&
errors.Is(admissionErr, ErrOutboundTrackedBudget)
return errors.New("forced retention failure")
},
)
if !errors.Is(err, errRPCResultRetentionHandoff) {
t.Fatalf("retention error = %v, want handoff sentinel", err)
}
if !observedInSlot {
t.Fatal("retention handoff did not run while encoded body was slot-confined")
}
if retained || encoded != nil || reserved != nil {
t.Fatalf("failed handoff escaped ownership: retained=%v encoded=%p reserved=%p", retained, encoded, reserved)
}
if got := len(outboundEncodeSlots); got != 0 {
t.Fatalf("encode slots retained after failed handoff = %d, want 0", got)
}
if got := budget.snapshot(); got != 1 {
t.Fatalf("primary budget after failed handoff = %d, want 1", got)
}
budget.release(1)
}
const saturatedSlotWaveResultData = "exact-business-success"
type saturatedSlotWaveGate struct {
firstWave int32
encodes atomic.Int32
entered chan struct{}
release chan struct{}
}
type saturatedSlotWaveResult struct{ gate *saturatedSlotWaveGate }
func (r saturatedSlotWaveResult) Encode(b *bin.Buffer) error {
call := r.gate.encodes.Add(1)
if call <= r.gate.firstWave {
r.gate.entered <- struct{}{}
<-r.gate.release
}
return (&tg.DataJSON{Data: saturatedSlotWaveResultData}).Encode(b)
}
type saturatedSlotWaveRPC struct {
calls atomic.Int32
gate *saturatedSlotWaveGate
}
func (h *saturatedSlotWaveRPC) Dispatch(context.Context, [8]byte, int64, *bin.Buffer) (bin.Encoder, error) {
h.calls.Add(1)
return exactTestRPCResult(saturatedSlotWaveResult{gate: h.gate}), nil
}
func (*saturatedSlotWaveRPC) NegotiatedLayer([8]byte, int64) (int, bool) { return 227, true }
func TestPublishRPCResultSaturatedBudgetRetainsExactResultsAcrossSlotWaves(t *testing.T) {
slotCount := cap(outboundEncodeSlots)
requestCount := slotCount*2 + 1
gate := &saturatedSlotWaveGate{
firstWave: int32(slotCount),
entered: make(chan struct{}, slotCount),
release: make(chan struct{}),
}
handler := &saturatedSlotWaveRPC{gate: gate}
s := New(Options{legacyRPC: handler})
now := time.Unix(1_700_000_000, 0)
s.rpcResults = newRPCResultCacheWithFlightLimit(func() time.Time { return now }, requestCount+1)
const primaryMax = 1 << 20
primary := newOutboundTrackedBudget(primaryMax)
if !primary.reserve(primaryMax) {
t.Fatal("saturate shared primary outbound budget")
}
conns := make([]*Conn, requestCount)
tasks := make([]inboundRPC, requestCount)
owners := make([]*rpcResultOwnerLease, requestCount)
reqMsgIDs := make([]int64, requestCount)
requestBody := mustEncodeTL(t, &tg.PhoneGetCallConfigRequest{})
for i := 0; i < requestCount; i++ {
var authKeyID [8]byte
authKeyID[0] = byte(i + 1)
authKeyID[1] = byte((i + 1) >> 8)
reqMsgID := int64(10_000 + i)
c := &Conn{
metrics: NopMetrics{},
writeTimeout: time.Second,
authKeyID: authKeyID,
sessionID: int64(20_000 + i),
outboundTrackedBudget: primary,
}
legacyLayerWireTestConn(t, c, 227)
claim, err := s.rpcResults.Acquire(c.authKeyID, c.sessionID, reqMsgID)
if err != nil || claim.state != rpcResultAcquireOwner {
t.Fatalf("acquire request %d = %+v err=%v", i, claim, err)
}
conns[i] = c
owners[i] = claim.owner
reqMsgIDs[i] = reqMsgID
tasks[i] = s.newInboundRPCTask(c, reqMsgID, "phone.getCallConfig", requestBody, claim.owner)
}
start := make(chan struct{})
errs := make([]error, requestCount)
var wg sync.WaitGroup
wg.Add(requestCount)
for i := range tasks {
go func(i int) {
defer wg.Done()
<-start
errs[i] = tasks[i].run(context.Background())
if tasks[i].release != nil {
tasks[i].release()
}
}(i)
}
close(start)
for i := 0; i < slotCount; i++ {
select {
case <-gate.entered:
case <-time.After(time.Second):
close(gate.release)
t.Fatalf("first encode wave entered %d/%d slots; handlers=%d encodes=%d first_err=%v", i, slotCount, handler.calls.Load(), gate.encodes.Load(), errs[0])
}
}
if got := gate.encodes.Load(); got != int32(slotCount) {
close(gate.release)
t.Fatalf("encodes before releasing first wave = %d, want slot cap %d", got, slotCount)
}
close(gate.release)
done := make(chan struct{})
go func() {
wg.Wait()
close(done)
}()
select {
case <-done:
case <-time.After(3 * time.Second):
t.Fatal("slot waves did not converge after saturated-budget retention")
}
if got := handler.calls.Load(); got != int32(requestCount) {
t.Fatalf("business executions = %d, want %d", got, requestCount)
}
if got := gate.encodes.Load(); got != int32(requestCount) {
t.Fatalf("successful result encodes = %d, want %d", got, requestCount)
}
if got := len(outboundEncodeSlots); got != 0 {
t.Fatalf("encode slots after both waves = %d, want 0", got)
}
if got := primary.snapshot(); got != primaryMax {
t.Fatalf("primary budget changed under saturation = %d, want %d", got, primaryMax)
}
var completedBytes int64
for i, c := range conns {
if !errors.Is(errs[i], ErrOutboundTrackedBudget) {
t.Fatalf("publish request %d error = %v, want terminal budget saturation", i, errs[i])
}
if !c.isRetired() {
t.Fatalf("request %d connection was not explicitly fenced", i)
}
if !owners[i].handedOff.Load() {
t.Fatalf("request %d owner was not handed to completed cache", i)
}
cached, ok := s.rpcResults.Get(c.authKeyID, c.sessionID, reqMsgIDs[i])
if !ok || cached == nil {
t.Fatalf("request %d exact result missing from completed cache", i)
}
completedBytes += int64(len(cached.body))
var envelope proto.Result
if err := envelope.Decode(&bin.Buffer{Buf: cached.body}); err != nil {
t.Fatalf("decode request %d cached rpc_result: %v", i, err)
}
if envelope.RequestMessageID != reqMsgIDs[i] {
t.Fatalf("request %d cached req_msg_id = %d, want %d", i, envelope.RequestMessageID, reqMsgIDs[i])
}
var result tg.DataJSON
if err := result.Decode(&bin.Buffer{Buf: envelope.Result}); err != nil {
t.Fatalf("decode request %d exact business result (possibly INTERNAL): %v", i, err)
}
if result.Data != saturatedSlotWaveResultData {
t.Fatalf("request %d cached result = %q, want %q", i, result.Data, saturatedSlotWaveResultData)
}
retry, err := s.rpcResults.Acquire(c.authKeyID, c.sessionID, reqMsgIDs[i])
if err != nil || retry.state != rpcResultAcquireCompleted || retry.encoded != cached {
t.Fatalf("retry request %d = %+v err=%v, want exact completed result", i, retry, err)
}
}
if got := handler.calls.Load(); got != int32(requestCount) {
t.Fatalf("business executions after retries = %d, want unchanged %d", got, requestCount)
}
if got := s.rpcResults.completedBytes.snapshot(); got != completedBytes {
t.Fatalf("completed-cache charge = %d, want exact retained bytes %d", got, completedBytes)
}
// Expiry is the completed cache's ownership release point. Force it
// deterministically and prove every retained byte is returned exactly once.
now = now.Add(rpcResultCacheTTL + time.Second)
for i, c := range conns {
if _, ok := s.rpcResults.Get(c.authKeyID, c.sessionID, reqMsgIDs[i]); ok {
t.Fatalf("request %d remained cached after forced expiry", i)
}
}
if got := s.rpcResults.completedBytes.snapshot(); got != 0 {
t.Fatalf("completed-cache bytes after expiry = %d, want 0", got)
}
primary.release(primaryMax)
if got := primary.snapshot(); got != 0 {
t.Fatalf("primary budget after release = %d, want 0", got)
}
}
func TestCachedReplayRestoreIsSynchronousAndIndependentOfGlobalHookExecutor(t *testing.T) {
// Occupy the entire executor. The replay-state callback must not reserve a
// ticket there: slow auth/store restoration has its own bounded path.
executor := newRPCDeliveryHookExecutor(1, 1)
started := make(chan struct{})
release := make(chan struct{})
ticket, ok := executor.reserve()
if !ok || !ticket.submit(func() {
close(started)
<-release
}) {
t.Fatal("occupy delivery hook executor")
}
<-started
oldExecutor := defaultRPCDeliveryHookExecutor
defaultRPCDeliveryHookExecutor = executor
defer func() {
defaultRPCDeliveryHookExecutor = oldExecutor
close(release)
}()
s := New(Options{WriteTimeout: time.Second})
transport := &collectingSessionTransport{}
key := newTestAuthKey(t)
c := s.newConn(transport, key, 777, 1)
legacyCanonicalTestConn(t, c)
t.Cleanup(c.ForceClose)
encoded := encodedRPCResultForPriorityTest(9001, 0)
encoded.delivery = newRPCResultDelivery(encoded.reqMsgID)
var restoreOrder atomic.Int32
encoded.setDeliveryHook(func() {
if !restoreOrder.CompareAndSwap(1, 2) {
panic("logical replay hook did not run after replacement metadata restore")
}
})
var restored atomic.Bool
if err := s.sendCachedRPCResultWithHook(context.Background(), c, encoded, func() error {
if got := len(transport.snapshot()); got != 1 {
return errors.New("replay restore ran before physical write")
}
if !restoreOrder.CompareAndSwap(0, 1) {
return errors.New("replacement replay restore ran out of order")
}
restored.Store(true)
return nil
}); err != nil {
t.Fatalf("send cached replay with saturated global executor: %v", err)
}
if !restored.Load() {
t.Fatal("cached replay returned before state restore completed")
}
if got := restoreOrder.Load(); got != 2 {
t.Fatalf("ordered replay restore stage = %d, want replacement then logical hook", got)
}
c.rpcMu.Lock()
pending := c.rpcReplayRestores
c.rpcMu.Unlock()
if pending != 0 {
t.Fatalf("replay restore barriers = %d, want 0", pending)
}
}
func TestBootstrapBarriersAlwaysUseConvergenceLane(t *testing.T) {
large := &encodedOutboundMessage{body: make([]byte, bulkOutboundThreshold)}
for _, method := range []string{
@ -144,13 +488,20 @@ func encodedRPCResultForPriorityTest(reqMsgID int64, payloadBytes int) *encodedO
if payloadBytes > 0 {
b.Put(make([]byte, payloadBytes))
}
return &encodedOutboundMessage{typeID: proto.ResultTypeID, reqMsgID: reqMsgID, body: b.Raw()}
return &encodedOutboundMessage{
typeID: proto.ResultTypeID, reqMsgID: reqMsgID, body: b.Raw(),
layer: &outboundLayerBinding{
profile: tg.LayerProfileCanonical,
typ: tg.LayerClassBoolType().Ref(),
kind: outboundLayerBindingRequest,
},
}
}
func TestConvergenceResultPassesQueuedBulkAfterBlockedWrite(t *testing.T) {
tr := newGatedRecordingTransport()
c := newOutboundTestConn(t, tr, newOutboundTrackedBudget(2<<20))
gate := &encodedOutboundMessage{typeID: tg.UpdatesTooLongTypeID, body: []byte{0x0b, 0xa1, 0x01, 0xe3}}
gate := exactTestUpdatesTooLong(t, c)
if err := c.SendBestEffortEncoded(context.Background(), proto.MessageFromServer, gate, 0); err != nil {
t.Fatalf("enqueue gate: %v", err)
}
@ -270,7 +621,7 @@ func TestRPCResultPipelineExportsPreparationAndDeliveryMetrics(t *testing.T) {
t.Fatalf("acquire flight = %+v err=%v", claim, err)
}
result := &tg.DataJSON{Data: string(bytes.Repeat([]byte("sticker-data"), 12<<10))}
if err := s.publishRPCResult(c, reqMsgID, "updates.getDifference#25939651", claim.owner, result, nil); err != nil {
if err := s.publishRPCResult(c, reqMsgID, "updates.getDifference#25939651", claim.owner, exactTestRPCResult(result), nil); err != nil {
t.Fatalf("publish result: %v", err)
}
select {
@ -293,7 +644,7 @@ func TestRPCResultPipelineExportsPreparationAndDeliveryMetrics(t *testing.T) {
func TestWrappedConvergenceMethodDrivesEgressAndReplayPriority(t *testing.T) {
metrics := &captureRPCResultMetrics{delivered: make(chan error, 1)}
s := New(Options{RPC: wrappedConvergenceRPC{}, Metrics: metrics})
s := New(Options{legacyRPC: wrappedConvergenceRPC{}, Metrics: metrics})
c := newOutboundTestConn(t, &failAfterTransport{}, newOutboundTrackedBudget(1<<20))
const reqMsgID = int64(9051)
claim, err := s.rpcResults.Acquire(c.authKeyID, c.sessionID, reqMsgID)
@ -302,7 +653,7 @@ func TestWrappedConvergenceMethodDrivesEgressAndReplayPriority(t *testing.T) {
}
body := mustEncodeTL(t, &tg.HelpGetConfigRequest{})
if err := s.handleRPC(context.Background(), c, reqMsgID, "invokeWithLayer#da9b0d0d", &bin.Buffer{Buf: body}, claim.owner); err != nil {
t.Fatalf("handle wrapped convergence RPC: %v", err)
t.Fatalf("handle wrapped convergence legacyRPC: %v", err)
}
select {
case err := <-metrics.delivered:
@ -336,7 +687,7 @@ func TestWrappedConvergenceMethodDrivesEgressAndReplayPriority(t *testing.T) {
}
func TestRPCWorkerReleasesAfterEgressAdmissionWhileWriteBlocked(t *testing.T) {
s := New(Options{RPC: immediateLargeRPC{}, WriteTimeout: time.Second})
s := New(Options{legacyRPC: immediateLargeRPC{}, WriteTimeout: time.Second})
tr := newGatedRecordingTransport()
c := newOutboundTestConn(t, tr, newOutboundTrackedBudget(2<<20))
const reqMsgID = int64(9001)
@ -396,7 +747,7 @@ func TestDeliveryHookRunsOnceAfterReplayNotFailedWrite(t *testing.T) {
}
var hooks atomic.Int32
if err := s.publishRPCResult(oldConn, reqMsgID, "updates.getDifference", claim.owner,
&tg.DataJSON{Data: "difference"}, func() { hooks.Add(1) }); err != nil {
exactTestRPCResult(&tg.DataJSON{Data: "difference"}), func() { hooks.Add(1) }); err != nil {
// Admission succeeds; the asynchronous physical failure is observed below.
t.Fatalf("publish result: %v", err)
}
@ -431,8 +782,12 @@ func TestDeliveryHookRunsOnceAfterReplayNotFailedWrite(t *testing.T) {
if got := hooks.Load(); got != 1 {
t.Fatalf("delivery hooks after replay = %d, want 1", got)
}
if got := cached.deliveryState(); got != rpcResultDeliveryDelivered {
t.Fatalf("replayed delivery state = %d, want delivered", got)
if got := cached.deliveryState(); got != rpcResultDeliveryReplayable {
t.Fatalf("cached representation state = %d, want original replayable attempt", got)
}
if cached.delivery == nil || cached.delivery.coordinator == nil ||
cached.delivery.coordinator.hookState() != rpcResultDeliveryHookDone {
t.Fatal("successful replay did not complete shared delivery coordinator")
}
if err := s.sendCachedRPCResult(context.Background(), replayConn, cached); err != nil {
t.Fatalf("second replay: %v", err)

253
internal/mtprotoedge/rpc_result_retention_budget_test.go Normal file
View file

@ -0,0 +1,253 @@
package mtprotoedge
import (
"context"
"crypto/rand"
"encoding/binary"
"errors"
"testing"
"time"
"github.com/iamxvbaba/td/bin"
"github.com/iamxvbaba/td/crypto"
"github.com/iamxvbaba/td/proto"
)
func TestRPCResultCloneReservationIsOneShotUnderReleaseRace(t *testing.T) {
const iterations = 256
for i := 0; i < iterations; i++ {
encoded := encodedRPCResultForPriorityTest(int64(i+1), 0)
budget := newOutboundTrackedBudget(int64(len(encoded.body)))
if !budget.reserve(len(encoded.body)) {
t.Fatal("reserve body")
}
reserved := &outboundBodyReservation{budget: budget, bytes: len(encoded.body)}
start := make(chan struct{})
taken := make(chan outboundOp, 1)
go func() {
<-start
op, _ := reserved.take(encoded)
taken <- op
}()
released := make(chan struct{})
go func() {
<-start
reserved.release()
close(released)
}()
close(start)
op := <-taken
<-released
op.releaseReservation(budget)
reserved.release()
if got := budget.snapshot(); got != 0 {
t.Fatalf("iteration %d retained bytes = %d, want 0", i, got)
}
}
}
func TestRPCResultReservationReleaseWinsAdmissionRollback(t *testing.T) {
encoded := encodedRPCResultForPriorityTest(6999, 0)
budget := newOutboundTrackedBudget(int64(len(encoded.body)))
if !budget.reserve(len(encoded.body)) {
t.Fatal("reserve body")
}
reserved := &outboundBodyReservation{budget: budget, bytes: len(encoded.body)}
op, err := reserved.take(encoded)
if err != nil {
t.Fatalf("take reservation: %v", err)
}
// Model a watchdog that retires the queued owner just before queue admission
// rolls the op back. The rollback must observe the release request and return
// the raw op charge instead of resurrecting an owner nobody will release.
reserved.release()
if !reserved.reclaim(&op) {
t.Fatal("reclaim actor reservation")
}
if got := budget.snapshot(); got != 0 {
t.Fatalf("retained bytes after watchdog/rollback race = %d, want 0", got)
}
reserved.release()
op.releaseReservation(budget)
}
func TestCachedRPCResultReplayUsesPreReservedBodyWithoutDoubleCharge(t *testing.T) {
encoded := encodedRPCResultForPriorityTest(7001, 32<<10)
budget := newOutboundTrackedBudget(int64(len(encoded.body)))
tr := newGatedRecordingTransport()
c := newOutboundTestConn(t, tr, budget)
s := New(Options{WriteTimeout: time.Second})
done := make(chan error, 1)
go func() {
done <- s.sendCachedRPCResult(context.Background(), c, encoded)
}()
select {
case <-tr.started:
case <-time.After(time.Second):
t.Fatal("cached replay did not reach the blocked physical write")
}
if got, want := budget.snapshot(), int64(len(encoded.body)); got != want {
t.Fatalf("blocked replay retained bytes = %d, want exactly one body %d", got, want)
}
tr.once.Do(func() { close(tr.release) })
select {
case err := <-done:
if err != nil {
t.Fatalf("cached replay: %v", err)
}
case <-time.After(time.Second):
t.Fatal("cached replay did not finish")
}
// rpc_result is reliable and remains charged once as resend state until ACK/close.
if got, want := budget.snapshot(), int64(len(encoded.body)); got != want {
t.Fatalf("pending replay retained bytes = %d, want %d", got, want)
}
c.Close()
if got := budget.snapshot(); got != 0 {
t.Fatalf("replay retained bytes after close = %d, want 0", got)
}
}
func TestQueuedRPCRewrapClonesAreBoundedBeforeAllocation(t *testing.T) {
source := encodedRPCResultForPriorityTest(7101, 64<<10)
perBody := len(source.body)
budget := newOutboundTrackedBudget(int64(2 * perBody))
c := &Conn{metrics: NopMetrics{}, outboundTrackedBudget: budget}
first, firstReserved, err := c.cloneRPCResultForRequestReserved(source, 7102, false)
if err != nil || first == nil || firstReserved == nil {
t.Fatalf("first queued clone = %p reservation=%p err=%v", first, firstReserved, err)
}
firstAlias := &rpcRewrapAlias{bodyReservation: firstReserved}
second, secondReserved, err := c.cloneRPCResultForRequestReserved(source, 7103, false)
if err != nil || second == nil || secondReserved == nil {
t.Fatalf("second queued clone = %p reservation=%p err=%v", second, secondReserved, err)
}
secondAlias := &rpcRewrapAlias{bodyReservation: secondReserved}
if got, want := budget.snapshot(), int64(2*perBody); got != want {
t.Fatalf("two queued aliases retained bytes = %d, want %d", got, want)
}
third, thirdReserved, err := c.cloneRPCResultForRequestReserved(source, 7104, false)
if !errors.Is(err, ErrOutboundTrackedBudget) || third != nil || thirdReserved != nil {
t.Fatalf("third queued clone = %p reservation=%p err=%v, want pre-allocation budget rejection", third, thirdReserved, err)
}
if got, want := budget.snapshot(), int64(2*perBody); got != want {
t.Fatalf("budget changed after rejected clone = %d, want %d", got, want)
}
firstAlias.finishReplayRestoreWithoutDelivery()
secondAlias.finishReplayRestoreWithoutDelivery()
if got := budget.snapshot(); got != 0 {
t.Fatalf("queued alias bytes after terminal release = %d, want 0", got)
}
}
func TestOutboundActorRetargetRequiresSecondBodyReservation(t *testing.T) {
const (
oldReqID = int64(7201)
newReqID = int64(7202)
)
encoded := encodedRPCResultForPriorityTest(oldReqID, 32<<10)
encoded.delivery = newRPCResultDelivery(oldReqID)
encoded.markQueued()
if !encoded.tryRetarget(newReqID) {
t.Fatal("retarget prepared rpc_result")
}
budget := newOutboundTrackedBudget(int64(len(encoded.body)))
if !budget.reserve(len(encoded.body)) {
t.Fatal("reserve original queued body")
}
tr := &failAfterTransport{}
c := newOutboundTestConn(t, tr, budget)
state := newOutboundState(budget)
var terminalErr error
var terminalBytes int64
c.handleOutboundSend(state, outboundOp{
kind: outboundSend,
ctx: context.Background(),
msgType: proto.MessageServerResponse,
encoded: encoded,
reservedBytes: len(encoded.body),
reservationBudget: budget,
enqueuedAt: time.Now(),
terminal: func(err error) {
terminalErr = err
terminalBytes = budget.snapshot()
},
})
if !errors.Is(terminalErr, ErrOutboundTrackedBudget) {
t.Fatalf("retarget terminal error = %v, want %v", terminalErr, ErrOutboundTrackedBudget)
}
if terminalBytes != int64(len(encoded.body)) {
t.Fatalf("bytes visible to terminal = %d, want original body %d retained", terminalBytes, len(encoded.body))
}
if got := tr.sends.Load(); got != 0 {
t.Fatalf("retarget under one-body budget wrote %d frames, want 0", got)
}
if got := int64(binary.LittleEndian.Uint64(encoded.body[4:12])); got != oldReqID {
t.Fatalf("source req_msg_id mutated to %d before second-body admission", got)
}
if got := budget.snapshot(); got != 0 {
t.Fatalf("retarget bytes after terminal = %d, want 0", got)
}
}
func TestOutboundActorRetargetTransfersOnlyReplacementToPending(t *testing.T) {
const (
oldReqID = int64(7301)
newReqID = int64(7302)
)
encoded := encodedRPCResultForPriorityTest(oldReqID, 32<<10)
encoded.delivery = newRPCResultDelivery(oldReqID)
encoded.markQueued()
if !encoded.tryRetarget(newReqID) {
t.Fatal("retarget prepared rpc_result")
}
perBody := len(encoded.body)
budget := newOutboundTrackedBudget(int64(2 * perBody))
if !budget.reserve(perBody) {
t.Fatal("reserve original queued body")
}
tr := &failAfterTransport{}
c := newOutboundTestConn(t, tr, budget)
state := newOutboundState(budget)
var terminalErr error
var terminalBytes int64
c.handleOutboundSend(state, outboundOp{
kind: outboundSend,
ctx: context.Background(),
msgType: proto.MessageServerResponse,
encoded: encoded,
reservedBytes: perBody,
reservationBudget: budget,
enqueuedAt: time.Now(),
terminal: func(err error) {
terminalErr = err
terminalBytes = budget.snapshot()
},
})
if terminalErr != nil {
t.Fatalf("retarget terminal error: %v", terminalErr)
}
if terminalBytes != int64(2*perBody) {
t.Fatalf("bytes visible to terminal = %d, want original+replacement %d", terminalBytes, 2*perBody)
}
if got := budget.snapshot(); got != int64(perBody) {
t.Fatalf("bytes after terminal = %d, want one pending replacement %d", got, perBody)
}
data, err := crypto.NewClientCipher(rand.Reader).DecryptFromBuffer(c.key, &bin.Buffer{Buf: tr.lastFrame()})
if err != nil {
t.Fatalf("decrypt retargeted frame: %v", err)
}
var result proto.Result
if err := result.Decode(&bin.Buffer{Buf: append([]byte(nil), data.Data()...)}); err != nil {
t.Fatalf("decode retargeted rpc_result: %v", err)
}
if result.RequestMessageID != newReqID {
t.Fatalf("wire req_msg_id = %d, want %d", result.RequestMessageID, newReqID)
}
state.releaseAll()
if got := budget.snapshot(); got != 0 {
t.Fatalf("retarget bytes after pending release = %d, want 0", got)
}
}

149
internal/mtprotoedge/rpc_result_subscriber_budget.go Normal file
View file

@ -0,0 +1,149 @@
package mtprotoedge
import (
"encoding/binary"
"hash/maphash"
"sync"
)
const (
rpcResultSubscriberMaxGlobal = 1 << 16
rpcResultSubscriberMaxAuth = 1 << 13
rpcResultSubscriberMaxSession = 1 << 11
rpcResultSubscriberMaxPerFlight = 1 << 7
)
type rpcResultSubscriberBudgetShard[K comparable] struct {
mu sync.Mutex
usage map[K]int64
}
// rpcResultSubscriberBudget bounds callbacks retained by pending replay
// flights independently from owner/result bytes. A duplicate does not reserve a
// new result row, so charging only unique owners would otherwise leave an
// unbounded same-msg-id reconnect path.
type rpcResultSubscriberBudget struct {
seed maphash.Seed
global rpcResultFlightLimit
authLimit int64
sessionLimit int64
authShards [rpcResultBudgetShards]rpcResultSubscriberBudgetShard[[8]byte]
sessionShards [rpcResultBudgetShards]rpcResultSubscriberBudgetShard[rpcResultSessionBudgetKey]
}
func newRPCResultSubscriberBudget(
seed maphash.Seed,
globalLimit, authLimit, sessionLimit int,
) *rpcResultSubscriberBudget {
b := &rpcResultSubscriberBudget{
seed: seed,
authLimit: int64(authLimit),
sessionLimit: int64(sessionLimit),
}
b.global.max = int64(globalLimit)
for i := range b.authShards {
b.authShards[i].usage = make(map[[8]byte]int64)
b.sessionShards[i].usage = make(map[rpcResultSessionBudgetKey]int64)
}
return b
}
func (b *rpcResultSubscriberBudget) reserve(key rpcResultCacheKey, slots int) bool {
if b == nil || slots <= 0 {
return false
}
authShard := b.authShard(key.authKeyID)
sessionKey := rpcResultSessionBudgetKey{authKeyID: key.authKeyID, sessionID: key.sessionID}
sessionShard := b.sessionShard(sessionKey)
delta := int64(slots)
authShard.mu.Lock()
sessionShard.mu.Lock()
authUsed := authShard.usage[key.authKeyID]
sessionUsed := sessionShard.usage[sessionKey]
if !withinRPCResultBudget(authUsed, delta, b.authLimit) ||
!withinRPCResultBudget(sessionUsed, delta, b.sessionLimit) ||
!b.global.reserveN(delta) {
sessionShard.mu.Unlock()
authShard.mu.Unlock()
return false
}
authShard.usage[key.authKeyID] = authUsed + delta
sessionShard.usage[sessionKey] = sessionUsed + delta
sessionShard.mu.Unlock()
authShard.mu.Unlock()
return true
}
func (b *rpcResultSubscriberBudget) release(key rpcResultCacheKey, slots int) {
if b == nil || slots <= 0 {
panic("mtproto rpc result subscriber release must be positive")
}
authShard := b.authShard(key.authKeyID)
sessionKey := rpcResultSessionBudgetKey{authKeyID: key.authKeyID, sessionID: key.sessionID}
sessionShard := b.sessionShard(sessionKey)
delta := int64(slots)
authShard.mu.Lock()
sessionShard.mu.Lock()
authUsed, authOK := authShard.usage[key.authKeyID]
sessionUsed, sessionOK := sessionShard.usage[sessionKey]
if !authOK || !sessionOK || authUsed < delta || sessionUsed < delta {
sessionShard.mu.Unlock()
authShard.mu.Unlock()
panic("mtproto rpc result subscriber budget underflow")
}
authUsed -= delta
sessionUsed -= delta
if authUsed == 0 {
delete(authShard.usage, key.authKeyID)
} else {
authShard.usage[key.authKeyID] = authUsed
}
if sessionUsed == 0 {
delete(sessionShard.usage, sessionKey)
} else {
sessionShard.usage[sessionKey] = sessionUsed
}
b.global.releaseN(delta)
sessionShard.mu.Unlock()
authShard.mu.Unlock()
}
func (b *rpcResultSubscriberBudget) authSnapshot(authKeyID [8]byte) int64 {
if b == nil {
return 0
}
shard := b.authShard(authKeyID)
shard.mu.Lock()
used := shard.usage[authKeyID]
shard.mu.Unlock()
return used
}
func (b *rpcResultSubscriberBudget) sessionSnapshot(authKeyID [8]byte, sessionID int64) int64 {
if b == nil {
return 0
}
key := rpcResultSessionBudgetKey{authKeyID: authKeyID, sessionID: sessionID}
shard := b.sessionShard(key)
shard.mu.Lock()
used := shard.usage[key]
shard.mu.Unlock()
return used
}
func (b *rpcResultSubscriberBudget) authShard(
authKeyID [8]byte,
) *rpcResultSubscriberBudgetShard[[8]byte] {
index := maphash.Bytes(b.seed, authKeyID[:]) & (rpcResultBudgetShards - 1)
return &b.authShards[index]
}
func (b *rpcResultSubscriberBudget) sessionShard(
key rpcResultSessionBudgetKey,
) *rpcResultSubscriberBudgetShard[rpcResultSessionBudgetKey] {
var raw [16]byte
copy(raw[:8], key.authKeyID[:])
binary.LittleEndian.PutUint64(raw[8:], uint64(key.sessionID))
index := maphash.Bytes(b.seed, raw[:]) & (rpcResultBudgetShards - 1)
return &b.sessionShards[index]
}

803
internal/mtprotoedge/rpc_rewrap.go
View file

@ -3,15 +3,18 @@ package mtprotoedge
import (
"context"
"crypto/sha256"
"fmt"
"log"
"runtime/debug"
"sync"
"sync/atomic"
"time"
"go.uber.org/zap"
"github.com/gotd/td/bin"
"github.com/gotd/td/proto"
"github.com/gotd/td/tg"
"github.com/iamxvbaba/td/bin"
"github.com/iamxvbaba/td/proto"
"github.com/iamxvbaba/td/tg"
)
// rpcRewrapRegistry links only an explicit official-client transition:
@ -37,6 +40,9 @@ type rpcRewrapSessionKey struct {
type rpcRewrapKey struct {
rpcRewrapSessionKey
fingerprint [sha256.Size]byte
semantic tg.LayerSemanticRequestIdentity
call tg.LayerCallIdentity
exact bool
}
type rpcRewrapRequestKey struct {
@ -102,6 +108,59 @@ func (r *rpcRewrapRegistry) register(c *Conn, body []byte, reqMsgID int64, metho
return true
}
func (r *rpcRewrapRegistry) registerSemantic(
c *Conn,
identity tg.LayerSemanticRequestIdentity,
call tg.LayerCallIdentity,
reqMsgID int64,
method string,
owner *rpcResultOwnerLease,
) bool {
if identity.Method() == 0 || identity.CanonicalSize() <= 0 {
return false
}
return r.registerKey(c, rpcRewrapKey{
rpcRewrapSessionKey: rpcRewrapSessionKey{authKeyID: c.authKeyID, sessionID: c.sessionID},
semantic: identity,
call: call,
exact: true,
}, reqMsgID, method, owner)
}
func (r *rpcRewrapRegistry) registerKey(c *Conn, key rpcRewrapKey, reqMsgID int64, method string, owner *rpcResultOwnerLease) bool {
if r == nil || c == nil || c.rpcRewrapInitialized.Load() || owner == nil {
return false
}
candidate := &rpcRewrapCandidate{
active: true, key: key, source: c, reqMsgID: reqMsgID, method: method,
owner: owner, waiter: owner.Waiter(),
}
if candidate.waiter == nil {
return false
}
session := key.rpcRewrapSessionKey
r.mu.Lock()
if r.total >= r.max {
r.mu.Unlock()
return false
}
r.byKey[key] = append(r.byKey[key], candidate)
set := r.bySession[session]
if set == nil {
set = make(map[*rpcRewrapCandidate]struct{})
r.bySession[session] = set
}
set[candidate] = struct{}{}
r.byRequest[rpcRewrapRequestKey{rpcRewrapSessionKey: session, reqMsgID: reqMsgID}] = candidate
r.total++
r.mu.Unlock()
if !owner.InstallAbortHook(func() { r.remove(candidate) }) {
r.remove(candidate)
return false
}
return true
}
func (r *rpcRewrapRegistry) claim(c *Conn, inner []byte) *rpcRewrapCandidate {
if r == nil || c == nil {
return nil
@ -122,6 +181,37 @@ func (r *rpcRewrapRegistry) claim(c *Conn, inner []byte) *rpcRewrapCandidate {
return nil
}
func (r *rpcRewrapRegistry) claimSemantic(
c *Conn,
identity tg.LayerSemanticRequestIdentity,
call tg.LayerCallIdentity,
) *rpcRewrapCandidate {
if r == nil || c == nil || identity.Method() == 0 || identity.CanonicalSize() <= 0 {
return nil
}
return r.claimKey(rpcRewrapKey{
rpcRewrapSessionKey: rpcRewrapSessionKey{authKeyID: c.authKeyID, sessionID: c.sessionID},
semantic: identity,
call: call,
exact: true,
})
}
func (r *rpcRewrapRegistry) claimKey(key rpcRewrapKey) *rpcRewrapCandidate {
r.mu.Lock()
queue := r.byKey[key]
for _, candidate := range queue {
if !candidate.active || candidate.claimed {
continue
}
candidate.claimed = true
r.mu.Unlock()
return candidate
}
r.mu.Unlock()
return nil
}
func (r *rpcRewrapRegistry) commit(candidate *rpcRewrapCandidate) {
if r == nil || candidate == nil {
return
@ -276,107 +366,636 @@ func decodeRPCRewrapInit(body []byte) (rpcRewrapInit, bool) {
}
type rpcRewrapAlias struct {
conn *Conn
newReqID int64
method string
oldWaiter *rpcResultWaiter
newOwner *rpcResultOwnerLease
sourceConn *Conn
sourceOwner *rpcResultOwnerLease
retargeted atomic.Bool
observeInit bool
init rpcRewrapInit
candidate *rpcRewrapCandidate
registry *rpcRewrapRegistry
conn *Conn
itemIndex int
newReqID int64
method string
oldWaiter *rpcResultWaiter
newOwner *rpcResultOwnerLease
sourceConn *Conn
sourceOwner *rpcResultOwnerLease
retargeted atomic.Bool
observeInit bool
init rpcRewrapInit
candidate *rpcRewrapCandidate
registry *rpcRewrapRegistry
afterSuccessfulDelivery func() error
finishReplayRestore func()
afterOnce sync.Once
deliveredFinalizeOnce sync.Once
deliveredFinalizeErr error
resultStoreClaimed atomic.Bool
executionOK atomic.Bool
// bodyReservation pins the replay/retarget clone from before allocation
// through queue residence. It is concurrency-safe because the watchdog and
// outbound actor race to release or take the same one-shot ownership token.
bodyReservation *outboundBodyReservation
}
func (a *rpcRewrapAlias) beginReplayRestore() {
if a == nil || a.conn == nil || a.finishReplayRestore != nil {
return
}
a.finishReplayRestore = a.conn.beginRPCReplayRestore()
}
func (a *rpcRewrapAlias) runAfterSuccessfulDelivery() (err error) {
if a == nil || a.afterSuccessfulDelivery == nil {
return nil
}
a.afterOnce.Do(func() {
if a.executionOK.Load() {
err = a.afterSuccessfulDelivery()
}
})
return err
}
func (a *rpcRewrapAlias) finishReplayRestoreWithoutDelivery() {
if a == nil {
return
}
a.releaseBodyReservation()
// Win or wait for any concurrent callback before dropping the barrier.
a.afterOnce.Do(func() {})
a.releaseReplayRestoreBarrier()
}
func (a *rpcRewrapAlias) releaseBodyReservation() {
if a != nil && a.bodyReservation != nil {
a.bodyReservation.release()
}
}
func (a *rpcRewrapAlias) releaseReplayRestoreBarrier() {
if a != nil && a.finishReplayRestore != nil {
a.finishReplayRestore()
}
}
func (a *rpcRewrapAlias) releaseDeferredLogicalHook() {
if a == nil || a.sourceOwner == nil || a.sourceOwner.Delivery() == nil ||
a.sourceOwner.Delivery().coordinator == nil {
return
}
a.sourceOwner.Delivery().coordinator.releaseDeferredHook()
}
func (a *rpcRewrapAlias) storeResultOnce(s *Server, encoded *encodedOutboundMessage) {
if a == nil || s == nil || encoded == nil || !a.resultStoreClaimed.CompareAndSwap(false, true) {
return
}
s.storeRPCResult(a.conn, a.newReqID, encoded)
}
func claimRPCRewrapLogicalHook(
ctx context.Context,
encoded *encodedOutboundMessage,
) (*rpcResultDeliveryHookClaim, error) {
if encoded == nil {
return nil, nil
}
// Only this alias is allowed to consume a sticky TryRetarget deferral. If a
// late physical success races another replacement replay, the coordinator
// waits for its Claimed/InProgress hook to publish Done.
return encoded.claimLogicalDeliveryHook(ctx, true)
}
// completeDeliveredRPCRewrapResult is safe after a watchdog has already fenced
// this physical generation. The caller has independent proof that the
// retargeted bytes reached the stream; deliveredFinalizeOnce, the shared hook
// coordinator and cache publication make late/concurrent invocations converge
// while preserving replacement -> logical -> cache -> barrier order.
func (s *Server) completeDeliveredRPCRewrapResult(
ctx context.Context,
a *rpcRewrapAlias,
encoded *encodedOutboundMessage,
source string,
) error {
if s == nil || a == nil || encoded == nil {
return ErrRPCResultFlightInvalid
}
a.deliveredFinalizeOnce.Do(func() {
defer a.releaseBodyReservation()
defer a.releaseReplayRestoreBarrier()
restoreCtx, cancel := boundedRPCReplayRestoreContext(ctx)
defer cancel()
logical, claimErr := claimRPCRewrapLogicalHook(restoreCtx, encoded)
encoded.markDelivered()
if claimErr != nil {
a.conn.fenceUndeliveredRPCResult()
a.deliveredFinalizeErr = fmt.Errorf("wait for rewrapped rpc_result logical restore: %w", claimErr)
a.storeResultOnce(s, encoded)
return
}
a.deliveredFinalizeErr = s.runBoundedRPCReplayRestore(
restoreCtx, a.conn, source, logical, a.runAfterSuccessfulDelivery,
)
a.storeResultOnce(s, encoded)
})
return a.deliveredFinalizeErr
}
var (
rpcRewrapDeliveryOnce sync.Once
rpcRewrapDeliveryJobs chan func()
rpcRewrapDeliveryOnce sync.Once
rpcRewrapDeliveryJobs chan rpcRewrapDeliveryJob
rpcRewrapObservationOnce sync.Once
rpcRewrapObservationJobs chan rpcRewrapDeliveryJob
)
const (
rpcRewrapDeliveryWorkers = 4
rpcRewrapDeliveryQueue = 256
rpcRewrapObserverWorkers = 1
rpcRewrapObserverQueue = 64
// Queue residence, physical delivery and ordered restore share one absolute
// deadline. An admitted alias must never retain a Conn scheduler barrier for
// minutes behind older slow jobs.
rpcRewrapDeliveryQueueTimeout = 5 * time.Second
)
func scheduleRPCRewrapDelivery(fn func()) bool {
if fn == nil {
type rpcRewrapDeliveryJob struct {
run func(*rpcRewrapDeliveryControl, time.Time)
fail func(error)
deadline time.Time
control *rpcRewrapDeliveryControl
}
type rpcRewrapDeliveryJobState uint32
const (
rpcRewrapJobPending rpcRewrapDeliveryJobState = iota
rpcRewrapJobRunning
rpcRewrapJobCommitted
rpcRewrapJobComplete
rpcRewrapJobFailed
)
// rpcRewrapDeliveryControl lets an independent deadline timer retire queued,
// running and physically committed jobs. A late worker cannot enter run after
// the timer wins; a committed non-cooperative restore is fenced by fail, and
// its eventual return cannot report failure or finish the barrier a second time.
type rpcRewrapDeliveryControl struct {
state atomic.Uint32
timerMu sync.Mutex
timer *time.Timer
}
type rpcRewrapPhysicalOutcome struct {
err error
owned bool
}
// waitRPCRewrapPhysicalTerminal deliberately keeps one of the four bounded
// workers attached to an in-progress actor write even after the watchdog fences
// the Conn. A broken transport may therefore strand at most four workers, while
// queued jobs still time out independently. If that transport later reports
// success, the worker cannot lose the logical hook merely because timeout won
// before its goroutine resumed.
func waitRPCRewrapPhysicalTerminal(
c *Conn,
ctx context.Context,
encoded *encodedOutboundMessage,
reserved *outboundBodyReservation,
control *rpcRewrapDeliveryControl,
) rpcRewrapPhysicalOutcome {
terminal := make(chan rpcRewrapPhysicalOutcome, 1)
_ = c.sendOutboundWithTerminalReserved(
ctx, proto.MessageServerResponse, nil, encoded, false,
func(err error) {
terminal <- rpcRewrapPhysicalOutcome{err: err, owned: control.commit()}
},
reserved,
)
return <-terminal
}
func newRPCRewrapDeliveryControl() *rpcRewrapDeliveryControl {
c := &rpcRewrapDeliveryControl{}
c.state.Store(uint32(rpcRewrapJobPending))
return c
}
func (c *rpcRewrapDeliveryControl) transition(from, to rpcRewrapDeliveryJobState) bool {
return c != nil && c.state.CompareAndSwap(uint32(from), uint32(to))
}
func (c *rpcRewrapDeliveryControl) fail() bool {
if c == nil {
return true
}
for {
state := rpcRewrapDeliveryJobState(c.state.Load())
if state == rpcRewrapJobComplete || state == rpcRewrapJobFailed {
return false
}
if c.transition(state, rpcRewrapJobFailed) {
c.stopTimer()
return true
}
}
}
func (c *rpcRewrapDeliveryControl) timeout() bool {
if c == nil {
return true
}
for {
state := rpcRewrapDeliveryJobState(c.state.Load())
if state != rpcRewrapJobPending && state != rpcRewrapJobRunning &&
state != rpcRewrapJobCommitted {
return false
}
if c.transition(state, rpcRewrapJobFailed) {
c.stopTimer()
return true
}
}
}
// commit records successful physical delivery (or an already-proven retarget)
// without disarming the watchdog. The same absolute deadline covers the
// replacement/logical restore and cache/barrier terminal path; complete is the
// only successful transition that stops the timer.
func (c *rpcRewrapDeliveryControl) commit() bool {
if c == nil || !c.transition(rpcRewrapJobRunning, rpcRewrapJobCommitted) {
return false
}
rpcRewrapDeliveryOnce.Do(func() {
rpcRewrapDeliveryJobs = make(chan func(), rpcRewrapDeliveryQueue)
for range rpcRewrapDeliveryWorkers {
return true
}
func (c *rpcRewrapDeliveryControl) running() bool {
return c != nil && rpcRewrapDeliveryJobState(c.state.Load()) == rpcRewrapJobRunning
}
func (c *rpcRewrapDeliveryControl) complete() {
if c == nil {
return
}
for {
state := rpcRewrapDeliveryJobState(c.state.Load())
if state != rpcRewrapJobRunning && state != rpcRewrapJobCommitted {
return
}
if c.transition(state, rpcRewrapJobComplete) {
c.stopTimer()
return
}
}
}
func (c *rpcRewrapDeliveryControl) installTimer(timer *time.Timer) {
if c == nil || timer == nil {
return
}
c.timerMu.Lock()
c.timer = timer
state := rpcRewrapDeliveryJobState(c.state.Load())
terminal := state == rpcRewrapJobComplete || state == rpcRewrapJobFailed
c.timerMu.Unlock()
if terminal {
timer.Stop()
}
}
func (c *rpcRewrapDeliveryControl) stopTimer() {
if c == nil {
return
}
c.timerMu.Lock()
timer := c.timer
c.timer = nil
c.timerMu.Unlock()
if timer != nil {
timer.Stop()
}
}
func (j rpcRewrapDeliveryJob) reportFailure(err error) {
if err == nil {
err = fmt.Errorf("rpc rewrap delivery job failed")
}
defer func() {
if recovered := recover(); recovered != nil {
log.Printf("mtprotoedge: rpc rewrap delivery failure callback panicked: %v\n%s", recovered, debug.Stack())
}
}()
if j.fail != nil {
j.fail(err)
return
}
log.Printf("mtprotoedge: rpc rewrap delivery job failed: %v", err)
}
func runRPCRewrapDeliveryJob(j rpcRewrapDeliveryJob) {
if j.run == nil {
j.reportFailure(fmt.Errorf("nil rpc rewrap delivery job"))
return
}
control := j.control
if control == nil {
control = newRPCRewrapDeliveryControl()
}
if !control.transition(rpcRewrapJobPending, rpcRewrapJobRunning) {
return
}
if !j.deadline.IsZero() && !time.Now().Before(j.deadline) {
if control.fail() {
j.reportFailure(context.DeadlineExceeded)
}
return
}
defer func() {
if recovered := recover(); recovered != nil {
if control.fail() {
j.reportFailure(fmt.Errorf("rpc rewrap delivery panic: %v", recovered))
}
log.Printf("mtprotoedge: rpc rewrap delivery job panicked: %v\n%s", recovered, debug.Stack())
return
}
control.complete()
}()
j.run(control, j.deadline)
}
func scheduleRPCRewrapJob(
job rpcRewrapDeliveryJob,
once *sync.Once,
jobs *chan rpcRewrapDeliveryJob,
workers, queue int,
) bool {
if job.run == nil {
return false
}
if job.deadline.IsZero() {
job.deadline = time.Now().Add(rpcRewrapDeliveryQueueTimeout)
}
job.control = newRPCRewrapDeliveryControl()
once.Do(func() {
*jobs = make(chan rpcRewrapDeliveryJob, queue)
for range workers {
go func() {
for job := range rpcRewrapDeliveryJobs {
job()
for job := range *jobs {
runRPCRewrapDeliveryJob(job)
}
}()
}
})
delay := time.Until(job.deadline)
if delay < 0 {
delay = 0
}
timer := time.AfterFunc(delay, func() {
if job.control.timeout() {
job.reportFailure(context.DeadlineExceeded)
}
})
job.control.installTimer(timer)
select {
case rpcRewrapDeliveryJobs <- fn:
case *jobs <- job:
return true
default:
// If the independent timer already won, it owns the fail callback and the
// caller must not report queue failure a second time.
if job.control.transition(rpcRewrapJobPending, rpcRewrapJobComplete) {
job.control.stopTimer()
return false
}
return true
}
}
func scheduleRPCRewrapDeliveryJob(job rpcRewrapDeliveryJob) bool {
return scheduleRPCRewrapJob(job, &rpcRewrapDeliveryOnce, &rpcRewrapDeliveryJobs,
rpcRewrapDeliveryWorkers, rpcRewrapDeliveryQueue)
}
func scheduleRPCRewrapObservation(fn func()) bool {
if fn == nil {
return false
}
return scheduleRPCRewrapJob(rpcRewrapDeliveryJob{
deadline: time.Now().Add(rpcRewrapDeliveryQueueTimeout),
run: func(*rpcRewrapDeliveryControl, time.Time) { fn() },
}, &rpcRewrapObservationOnce, &rpcRewrapObservationJobs,
rpcRewrapObserverWorkers, rpcRewrapObserverQueue)
}
func (s *Server) rpcRewrapRestoreJob(
a *rpcRewrapAlias,
source string,
run func(*rpcRewrapDeliveryControl, time.Time),
) rpcRewrapDeliveryJob {
return rpcRewrapDeliveryJob{
deadline: time.Now().Add(rpcRewrapDeliveryQueueTimeout),
run: run,
fail: func(err error) {
if a != nil && a.conn != nil {
a.conn.fenceUndeliveredRPCResult()
}
if a != nil {
a.releaseBodyReservation()
a.releaseReplayRestoreBarrier()
}
if s != nil && s.log != nil {
s.log.Warn("RPC rewrap delivery job failed",
zap.String("source", source), zap.Error(err))
}
},
}
}
func (s *Server) failRPCRewrapResultJob(
a *rpcRewrapAlias,
encoded *encodedOutboundMessage,
err error,
) {
if a != nil {
defer a.releaseBodyReservation()
}
if a == nil || a.conn == nil || a.newOwner == nil || encoded == nil {
return
}
a.conn.fenceUndeliveredRPCResult()
publish := a.newOwner.HandOff()
encoded.markReplayable()
encoded.releaseDeferredLogicalDeliveryHook()
// Release the connection-local scheduler before any defensive cache panic;
// the physical generation is already fenced, so no following task can run.
a.releaseReplayRestoreBarrier()
if publish {
a.storeResultOnce(s, encoded)
}
if s != nil && s.log != nil {
s.log.Warn("RPC rewrap result job failed; exact result retained",
zap.String("method", a.method), zap.Int64("req_msg_id", a.newReqID), zap.Error(err))
}
}
func (a *rpcRewrapAlias) activate(s *Server) error {
if a == nil || s == nil || a.conn == nil || a.oldWaiter == nil {
return ErrRPCResultFlightInvalid
}
err := a.oldWaiter.Subscribe(func(encoded *encodedOutboundMessage, ok bool) {
// Install the scheduler barrier synchronously, before this plan publishes
// any following naked RPC tasks. The asynchronous physical replay below is
// then free to use a bounded rewrap worker without an ordering race.
a.beginReplayRestore()
var executionSubscriber func(bool)
if a.newOwner != nil || a.afterSuccessfulDelivery != nil {
executionSubscriber = func(success bool) {
a.executionOK.Store(success)
if a.newOwner != nil {
a.newOwner.CompleteExecution(success)
}
}
}
resultSubscriber := func(encoded *encodedOutboundMessage, ok bool) {
if !ok || encoded == nil {
if a.newOwner != nil {
a.newOwner.Abort()
}
a.conn.fenceUndeliveredRPCResult()
a.releaseDeferredLogicalHook()
a.finishReplayRestoreWithoutDelivery()
return
}
if a.newOwner == nil {
if !scheduleRPCRewrapDelivery(func() {
ctx, cancel := context.WithTimeout(context.Background(), min(5*time.Second, max(time.Second, a.conn.writeTimeout)))
defer cancel()
if err := s.sendCachedRPCResult(ctx, a.conn, encoded); err != nil && !isClientDisconnect(err) {
s.log.Debug("RPC init rewrap pending replay failed", zap.Error(err))
}
}) {
a.conn.fenceUndeliveredRPCResult()
}
return
}
clone, err := cloneRPCResultForRequest(encoded, a.newReqID, false)
if err != nil {
a.newOwner.Abort()
a.conn.fenceUndeliveredRPCResult()
return
}
if a.retargeted.Load() {
if !a.newOwner.HandOff() {
attempt, reserved, cloneErr := a.conn.cloneRPCResultForRequestReserved(encoded, encoded.reqMsgID, false)
if cloneErr != nil {
a.conn.failOutboundBudget(cloneErr)
a.conn.fenceUndeliveredRPCResult()
a.finishReplayRestoreWithoutDelivery()
return
}
clone.markDelivered()
s.storeRPCResult(a.conn, a.newReqID, clone)
a.bodyReservation = reserved
job := s.rpcRewrapRestoreJob(a, "pending init rewrap replay", func(control *rpcRewrapDeliveryControl, deadline time.Time) {
ctx, cancel := context.WithDeadline(context.Background(), deadline)
defer cancel()
outcome := waitRPCRewrapPhysicalTerminal(a.conn, ctx, attempt, a.bodyReservation, control)
if outcome.err != nil {
if !outcome.owned {
return
}
a.conn.fenceUndeliveredRPCResult()
attempt.markReplayable()
a.releaseBodyReservation()
a.releaseReplayRestoreBarrier()
if !isClientDisconnect(outcome.err) {
s.log.Debug("RPC init rewrap pending replay failed", zap.Error(outcome.err))
}
return
}
// A watchdog may win after the transport has already returned physical
// success but before this goroutine resumes. Success is irrevocable: run
// the once-only restore with a fresh bounded lifetime if timeout failure
// already fenced/released this physical generation.
restoreParent := ctx
if !outcome.owned {
restoreParent = context.Background()
}
restoreCtx, cancelRestore := boundedRPCReplayRestoreContext(restoreParent)
defer cancelRestore()
logical, claimErr := attempt.claimLogicalDeliveryHook(restoreCtx, false)
attempt.markDelivered()
if claimErr != nil {
a.conn.fenceUndeliveredRPCResult()
a.releaseBodyReservation()
a.releaseReplayRestoreBarrier()
return
}
restoreErr := s.runBoundedRPCReplayRestore(
restoreCtx, a.conn, "pending init rewrap replay", logical, a.runAfterSuccessfulDelivery,
)
a.releaseBodyReservation()
a.releaseReplayRestoreBarrier()
if restoreErr != nil && !isClientDisconnect(restoreErr) {
s.log.Debug("RPC init rewrap pending replay failed", zap.Error(restoreErr))
}
})
if !scheduleRPCRewrapDeliveryJob(job) {
a.conn.fenceUndeliveredRPCResult()
a.finishReplayRestoreWithoutDelivery()
}
return
}
// A successful physical write under the retargeted req_msg_id is the only
// proof that lets the alias reuse that attempt. A mere TryRetarget success
// is not proof: the original socket may have failed before any bytes landed.
retargetDelivered := a.retargeted.Load() &&
encoded.deliveryState() == rpcResultDeliveryDelivered &&
encoded.writtenRequestID() == a.newReqID
clone, reserved, err := a.conn.cloneRPCResultForRequestReserved(encoded, a.newReqID, retargetDelivered)
if err != nil {
a.conn.failOutboundBudget(err)
a.newOwner.Abort()
a.conn.fenceUndeliveredRPCResult()
a.finishReplayRestoreWithoutDelivery()
return
}
a.bodyReservation = reserved
if retargetDelivered {
if !a.newOwner.HandOff() {
a.conn.fenceUndeliveredRPCResult()
a.finishReplayRestoreWithoutDelivery()
return
}
job := s.rpcRewrapRestoreJob(a, "retargeted init rewrap result", func(control *rpcRewrapDeliveryControl, deadline time.Time) {
if !control.commit() {
return
}
restoreCtx, cancelRestore := context.WithDeadline(context.Background(), deadline)
defer cancelRestore()
restoreErr := s.completeDeliveredRPCRewrapResult(
restoreCtx, a, clone, "retargeted init rewrap result",
)
if restoreErr != nil && !isClientDisconnect(restoreErr) {
s.log.Debug("Retargeted RPC restore failed", zap.Error(restoreErr))
}
})
job.fail = func(err error) {
// Never enter deliveredFinalizeOnce from the timer goroutine: the worker
// may already own a non-cooperative restore. Fence and release its Conn
// barrier first, then retain the immutable delivered result for a later
// replacement replay, which will wait on coordinator Claimed/InProgress.
a.conn.fenceUndeliveredRPCResult()
clone.markReplayable()
clone.releaseDeferredLogicalDeliveryHook()
a.releaseReplayRestoreBarrier()
a.storeResultOnce(s, clone)
a.releaseBodyReservation()
s.log.Warn("Retargeted RPC restore watchdog expired",
zap.String("method", a.method), zap.Int64("req_msg_id", a.newReqID), zap.Error(err))
}
if !scheduleRPCRewrapDeliveryJob(job) {
job.fail(ErrOutboundQueueFull)
}
s.log.Info("RPC init rewrap result retargeted",
zap.String("method", a.method), zap.Int64("new_req_msg_id", a.newReqID),
zap.String("auth_key_id", a.conn.authKeyHex), zap.Int64("session_id", a.conn.sessionID))
return
}
if !scheduleRPCRewrapDelivery(func() {
s.publishRewrappedRPCResult(a.conn, a.newReqID, a.method, a.newOwner, clone)
}) {
job := s.rpcRewrapRestoreJob(a, "pending init rewrap result", func(control *rpcRewrapDeliveryControl, deadline time.Time) {
s.publishRewrappedRPCResult(a.conn, a.newReqID, a.method, a.newOwner, clone, a, control, deadline)
})
// Once this alias consumed the source candidate, expiration or panic of
// the admitted worker job must still publish the immutable result under
// the new msg_id. Otherwise the alias owner would remain pending forever
// (or a reconnect could execute the business request a second time).
job.fail = func(err error) { s.failRPCRewrapResultJob(a, clone, err) }
if !scheduleRPCRewrapDeliveryJob(job) {
// The completed result is durable in memory. Fence before publishing it
// under the new msg_id so a replacement can replay without re-executing.
a.conn.fenceUndeliveredRPCResult()
if a.newOwner.HandOff() {
clone.markReplayable()
s.storeRPCResult(a.conn, a.newReqID, clone)
}
s.failRPCRewrapResultJob(a, clone, ErrOutboundQueueFull)
}
})
}
var err error
if executionSubscriber != nil {
err = a.oldWaiter.SubscribeResultAndExecution(resultSubscriber, executionSubscriber)
} else {
err = a.oldWaiter.Subscribe(resultSubscriber)
}
if err != nil {
a.finishReplayRestoreWithoutDelivery()
s.rpcRewrap.release(a.candidate)
return err
}
@ -407,7 +1026,7 @@ func (s *Server) scheduleRewrappedInitObservation(c *Conn, init rpcRewrapInit) {
if !ok || c == nil {
return
}
if !scheduleRPCRewrapDelivery(func() {
if !scheduleRPCRewrapObservation(func() {
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
if err := observer.ObserveInitConnection(
@ -422,33 +1041,75 @@ func (s *Server) scheduleRewrappedInitObservation(c *Conn, init rpcRewrapInit) {
}
}
func (s *Server) publishRewrappedRPCResult(c *Conn, reqMsgID int64, method string, owner *rpcResultOwnerLease, encoded *encodedOutboundMessage) {
func (s *Server) publishRewrappedRPCResult(
c *Conn,
reqMsgID int64,
method string,
owner *rpcResultOwnerLease,
encoded *encodedOutboundMessage,
alias *rpcRewrapAlias,
control *rpcRewrapDeliveryControl,
deadline time.Time,
) {
if s == nil || c == nil || owner == nil || encoded == nil {
if alias != nil {
alias.finishReplayRestoreWithoutDelivery()
}
return
}
if !owner.HandOff() {
c.fenceUndeliveredRPCResult()
alias.finishReplayRestoreWithoutDelivery()
return
}
if control == nil || !control.running() {
alias.releaseBodyReservation()
return
}
priority := rpcResultPriority(method, encoded)
encoded.priority = priority
terminal := func(deliveryErr error) {
if deliveryErr != nil {
encoded.markReplayable()
c.fenceUndeliveredRPCResult()
} else {
encoded.markDelivered()
}
s.storeRPCResult(c, reqMsgID, encoded)
}
encoded.markQueued()
ctx, cancel := context.WithTimeout(context.Background(), min(5*time.Second, max(time.Second, c.writeTimeout)))
if deadline.IsZero() {
deadline = time.Now().Add(rpcRewrapDeliveryQueueTimeout)
}
ctx, cancel := context.WithDeadline(context.Background(), deadline)
defer cancel()
if err := c.enqueueEncodedDelivery(ctx, proto.MessageServerResponse, encoded, priority, terminal); err != nil {
terminal(err)
// Rewrap delivery is synchronous on this small bounded worker pool. This
// makes the queue deadline cover the physical write and lets the pending
// logical hook join the same per-Conn ordered restore, without touching the
// process-wide asynchronous hook executor.
outcome := waitRPCRewrapPhysicalTerminal(c, ctx, encoded, alias.bodyReservation, control)
if outcome.err != nil {
if !outcome.owned {
return
}
encoded.markReplayable()
encoded.releaseDeferredLogicalDeliveryHook()
c.fenceUndeliveredRPCResult()
alias.storeResultOnce(s, encoded)
alias.releaseBodyReservation()
alias.releaseReplayRestoreBarrier()
return
}
s.log.Info("RPC init rewrap result replay admitted",
// Physical success outranks an already-fired watchdog. The timeout path may
// have fenced and cached a replayable clone, but it cannot revoke bytes; the
// shared once/coordinator below still completes logical state exactly once.
// Run replacement metadata then the original logical hook before publishing
// the alias cache entry. Whole-finalization once also covers a watchdog racing
// a late physical terminal, so completed metadata cannot be overwritten.
restoreParent := ctx
if !outcome.owned {
restoreParent = context.Background()
}
restoreCtx, cancelRestore := boundedRPCReplayRestoreContext(restoreParent)
defer cancelRestore()
restoreErr := s.completeDeliveredRPCRewrapResult(
restoreCtx, alias, encoded, "physically delivered init rewrap result",
)
if restoreErr != nil && !isClientDisconnect(restoreErr) {
s.log.Debug("RPC init rewrap delivered-state restore failed", zap.Error(restoreErr))
}
s.log.Info("RPC init rewrap result replay delivered",
zap.String("method", method), zap.Int64("req_msg_id", reqMsgID),
zap.String("auth_key_id", c.authKeyHex), zap.Int64("session_id", c.sessionID),
zap.Int("wire_bytes", len(encoded.body)))

825
internal/mtprotoedge/rpc_rewrap_test.go
View file

@ -3,14 +3,17 @@ package mtprotoedge
import (
"context"
"encoding/binary"
"errors"
"sync"
"sync/atomic"
"testing"
"time"
"github.com/gotd/td/bin"
"github.com/gotd/td/crypto"
"github.com/gotd/td/mt"
"github.com/gotd/td/tg"
"github.com/iamxvbaba/td/bin"
"github.com/iamxvbaba/td/crypto"
"github.com/iamxvbaba/td/mt"
"github.com/iamxvbaba/td/tg"
"go.uber.org/zap/zaptest"
)
func encodeRewrapTestRequest(t *testing.T) ([]byte, []byte) {
@ -67,6 +70,253 @@ func TestRPCResultDeliveryRetargetHasExactWritingBarrier(t *testing.T) {
}
}
func TestRPCRewrapAliasKeepsHookAcrossFailedSourceAttempt(t *testing.T) {
executor := newRPCDeliveryHookExecutor(1, 2)
var hooks atomic.Int32
source := &encodedOutboundMessage{
typeID: mt.RPCResultTypeID,
reqMsgID: 101,
body: make([]byte, 16),
delivery: newRPCResultDelivery(101),
}
binary.LittleEndian.PutUint32(source.body[:4], mt.RPCResultTypeID)
binary.LittleEndian.PutUint64(source.body[4:12], 101)
source.setDeliveryHook(func() { hooks.Add(1) })
if err := source.prepareDeliveryHook(executor); err != nil {
t.Fatalf("reserve source attempt: %v", err)
}
source.markReplayable()
alias, err := cloneRPCResultForRequest(source, 202, false)
if err != nil {
t.Fatalf("clone alias: %v", err)
}
if alias.delivery == source.delivery {
t.Fatal("alias reused source physical-attempt state")
}
if alias.delivery.coordinator != source.delivery.coordinator {
t.Fatal("alias did not inherit logical delivery coordinator")
}
if err := alias.prepareDeliveryHook(executor); err != nil {
t.Fatalf("reserve alias attempt: %v", err)
}
alias.markDelivered()
deadline := time.Now().Add(time.Second)
for hooks.Load() != 1 && time.Now().Before(deadline) {
time.Sleep(time.Millisecond)
}
if got := hooks.Load(); got != 1 {
t.Fatalf("alias delivery hooks = %d, want 1", got)
}
if got := source.deliveryState(); got != rpcResultDeliveryReplayable {
t.Fatalf("source attempt state = %d, want replayable", got)
}
if got := alias.deliveryState(); got != rpcResultDeliveryDelivered {
t.Fatalf("alias attempt state = %d, want delivered", got)
}
}
func TestRPCRewrapFailedSourceAttemptPhysicallyDeliversAliasOnce(t *testing.T) {
s := New(Options{WriteTimeout: time.Second})
failedTransport := &failAfterTransport{}
failedTransport.failAt.Store(1)
aliasTransport := &collectingSessionTransport{}
key := newTestAuthKey(t)
const sessionID, oldReqID, newReqID = int64(89), int64(5101), int64(5201)
sourceConn := s.newConn(failedTransport, key, sessionID, 1)
aliasConn := s.newConn(aliasTransport, key, sessionID, 1)
legacyCanonicalTestConn(t, sourceConn)
legacyCanonicalTestConn(t, aliasConn)
t.Cleanup(sourceConn.ForceClose)
t.Cleanup(aliasConn.ForceClose)
oldClaim, err := s.rpcResults.Acquire(key.ID, sessionID, oldReqID)
if err != nil || oldClaim.state != rpcResultAcquireOwner {
t.Fatalf("old flight = %+v err=%v", oldClaim, err)
}
requestBody := []byte{1, 2, 3, 4}
if !s.rpcRewrap.register(sourceConn, requestBody, oldReqID, "test.method", oldClaim.owner) {
t.Fatal("register source rewrap candidate")
}
candidate := s.rpcRewrap.claim(aliasConn, requestBody)
if candidate == nil {
t.Fatal("claim source rewrap candidate")
}
newClaim, err := s.rpcResults.Acquire(key.ID, sessionID, newReqID)
if err != nil || newClaim.state != rpcResultAcquireOwner {
t.Fatalf("alias flight = %+v err=%v", newClaim, err)
}
alias := &rpcRewrapAlias{
conn: aliasConn, newReqID: newReqID, method: "test.method",
oldWaiter: oldClaim.owner.Waiter(), newOwner: newClaim.owner,
sourceConn: sourceConn, sourceOwner: oldClaim.owner,
candidate: candidate, registry: s.rpcRewrap,
}
if err := alias.activate(s); err != nil {
t.Fatalf("activate alias: %v", err)
}
var hooks atomic.Int32
if err := s.publishRPCResult(sourceConn, oldReqID, "test.method", oldClaim.owner,
&mt.RPCError{ErrorCode: 400, ErrorMessage: "TEST"}, func() { hooks.Add(1) }); err != nil {
t.Fatalf("publish source result: %v", err)
}
deadline := time.Now().Add(2 * time.Second)
for time.Now().Before(deadline) {
cached, ok := s.rpcResults.Get(key.ID, sessionID, newReqID)
if ok && cached.deliveryState() == rpcResultDeliveryDelivered && hooks.Load() == 1 {
if got := cached.writtenRequestID(); got != newReqID {
t.Fatalf("alias physical request ID = %d, want %d", got, newReqID)
}
if got := len(aliasTransport.snapshot()); got != 1 {
t.Fatalf("alias physical writes = %d, want 1", got)
}
sourceCached, sourceOK := s.rpcResults.Get(key.ID, sessionID, oldReqID)
if !sourceOK || sourceCached.deliveryState() != rpcResultDeliveryReplayable {
t.Fatalf("source physical attempt = cached:%v state:%d, want replayable", sourceOK, sourceCached.deliveryState())
}
return
}
time.Sleep(time.Millisecond)
}
cached, ok := s.rpcResults.Get(key.ID, sessionID, newReqID)
t.Fatalf("alias result = cached:%v state:%v hooks:%d writes:%d", ok, cached.deliveryState(), hooks.Load(), len(aliasTransport.snapshot()))
}
func TestRPCRewrapRepeatedReplacementSubscriberCapacityStaysBounded(t *testing.T) {
s := New(Options{WriteTimeout: time.Second})
s.rpcResults = newRPCResultSubscriberTestCache(2, 2, 2, 2)
transport := &collectingSessionTransport{}
key := newTestAuthKey(t)
const sessionID, oldReqID, newReqID = int64(188), int64(15101), int64(15201)
c := s.newConn(transport, key, sessionID, 1)
legacyCanonicalTestConn(t, c)
t.Cleanup(c.ForceClose)
oldClaim, err := s.rpcResults.Acquire(key.ID, sessionID, oldReqID)
if err != nil || oldClaim.owner == nil {
t.Fatalf("old flight = %#v err=%v", oldClaim, err)
}
var sourceResultCalls, sourceExecutionCalls atomic.Int32
if err := oldClaim.owner.Waiter().SubscribeResultAndExecution(
func(*encodedOutboundMessage, bool) { sourceResultCalls.Add(1) },
func(bool) { sourceExecutionCalls.Add(1) },
); err != nil {
t.Fatalf("fill source subscriber capacity: %v", err)
}
for i := 0; i < 100; i++ {
newClaim, err := s.rpcResults.Acquire(key.ID, sessionID, newReqID)
if err != nil || newClaim.owner == nil {
t.Fatalf("replacement %d Acquire = %#v err=%v", i, newClaim, err)
}
plan := &inboundPlan{rewrapAliases: []*rpcRewrapAlias{{
conn: c, oldWaiter: oldClaim.owner.Waiter(), newOwner: newClaim.owner,
newReqID: newReqID, method: "test.method",
}}}
err = plan.commitRewrapAliases(s)
if !errors.Is(err, ErrRPCResultSubscriberCapacity) {
t.Fatalf("replacement %d activation err=%v, want subscriber capacity", i, err)
}
if got := s.rpcResults.subscriberBudget.global.snapshot(); got != 2 {
t.Fatalf("replacement %d subscriber usage=%d, want 2", i, got)
}
c.rpcMu.Lock()
barriers := c.rpcReplayRestores
c.rpcMu.Unlock()
if barriers != 0 {
t.Fatalf("replacement %d leaked replay barrier=%d", i, barriers)
}
}
if !oldClaim.owner.Abort() {
t.Fatal("abort source owner")
}
if sourceResultCalls.Load() != 1 || sourceExecutionCalls.Load() != 1 {
t.Fatalf("source callbacks result=%d execution=%d", sourceResultCalls.Load(), sourceExecutionCalls.Load())
}
if got := s.rpcResults.subscriberBudget.global.snapshot(); got != 0 {
t.Fatalf("subscriber usage=%d after source abort", got)
}
}
func TestRPCRewrapRetargetFailureRequiresReplacementAliasWrite(t *testing.T) {
s := New(Options{WriteTimeout: time.Second})
failedTransport := &failAfterTransport{}
failedTransport.failAt.Store(1)
key := newTestAuthKey(t)
const sessionID, oldReqID, newReqID = int64(90), int64(6101), int64(6201)
failedConn := s.newConn(failedTransport, key, sessionID, 1)
legacyCanonicalTestConn(t, failedConn)
t.Cleanup(failedConn.ForceClose)
oldClaim, err := s.rpcResults.Acquire(key.ID, sessionID, oldReqID)
if err != nil || oldClaim.state != rpcResultAcquireOwner {
t.Fatalf("old flight = %+v err=%v", oldClaim, err)
}
requestBody := []byte{5, 6, 7, 8}
if !s.rpcRewrap.register(failedConn, requestBody, oldReqID, "test.method", oldClaim.owner) {
t.Fatal("register source rewrap candidate")
}
candidate := s.rpcRewrap.claim(failedConn, requestBody)
if candidate == nil {
t.Fatal("claim source rewrap candidate")
}
newClaim, err := s.rpcResults.Acquire(key.ID, sessionID, newReqID)
if err != nil || newClaim.state != rpcResultAcquireOwner {
t.Fatalf("alias flight = %+v err=%v", newClaim, err)
}
alias := &rpcRewrapAlias{
conn: failedConn, newReqID: newReqID, method: "test.method",
oldWaiter: oldClaim.owner.Waiter(), newOwner: newClaim.owner,
sourceConn: failedConn, sourceOwner: oldClaim.owner,
candidate: candidate, registry: s.rpcRewrap,
}
if err := alias.activate(s); err != nil {
t.Fatalf("activate alias: %v", err)
}
if !alias.retargeted.Load() {
t.Fatal("same-Conn queued source was not retargeted")
}
var hooks atomic.Int32
if err := s.publishRPCResult(failedConn, oldReqID, "test.method", oldClaim.owner,
&mt.RPCError{ErrorCode: 400, ErrorMessage: "TEST"}, func() { hooks.Add(1) }); err != nil {
t.Fatalf("publish source result: %v", err)
}
deadline := time.Now().Add(2 * time.Second)
var aliasCached *encodedOutboundMessage
for time.Now().Before(deadline) {
if cached, ok := s.rpcResults.Get(key.ID, sessionID, newReqID); ok && cached.deliveryState() == rpcResultDeliveryReplayable {
aliasCached = cached
break
}
time.Sleep(time.Millisecond)
}
if aliasCached == nil {
t.Fatal("failed retarget was incorrectly treated as delivered or alias result was not retained")
}
if hooks.Load() != 0 {
t.Fatalf("failed retarget ran delivery hook %d times", hooks.Load())
}
replacementTransport := &collectingSessionTransport{}
replacement := s.newConn(replacementTransport, key, sessionID, 1)
legacyCanonicalTestConn(t, replacement)
t.Cleanup(replacement.ForceClose)
if err := s.sendCachedRPCResult(context.Background(), replacement, aliasCached); err != nil {
t.Fatalf("replacement alias replay: %v", err)
}
deadline = time.Now().Add(time.Second)
for hooks.Load() != 1 && time.Now().Before(deadline) {
time.Sleep(time.Millisecond)
}
if hooks.Load() != 1 || len(replacementTransport.snapshot()) != 1 {
t.Fatalf("replacement delivery = hooks:%d writes:%d, want 1/1", hooks.Load(), len(replacementTransport.snapshot()))
}
}
func TestRPCResultWaiterSubscribeIsEventDriven(t *testing.T) {
cache := newRPCResultCacheWithFlightLimit(time.Now, 8)
claim, err := cache.Acquire([8]byte{1}, 2, 3)
@ -286,7 +536,17 @@ func TestInitRewrapAliasesExecutionAndRetargetsQueuedResult(t *testing.T) {
encoded.markDelivered()
s.rpcResults.Put(c.authKeyID, c.sessionID, oldReqID, encoded)
aliased, ok := s.rpcResults.Get(c.authKeyID, c.sessionID, newReqID)
var (
aliased *encodedOutboundMessage
ok bool
)
deadline := time.Now().Add(time.Second)
for time.Now().Before(deadline) {
if aliased, ok = s.rpcResults.Get(c.authKeyID, c.sessionID, newReqID); ok {
break
}
time.Sleep(time.Millisecond)
}
if !ok {
t.Fatal("new req_msg_id result was not completed")
}
@ -297,3 +557,558 @@ func TestInitRewrapAliasesExecutionAndRetargetsQueuedResult(t *testing.T) {
t.Fatalf("rewrap registry retained %d consumed candidates", s.rpcRewrap.total)
}
}
func TestRPCRewrapDeliveryJobPanicAndDeadlineReleaseBarrier(t *testing.T) {
s := &Server{log: zaptest.NewLogger(t)}
c := &Conn{metrics: NopMetrics{}}
alias := &rpcRewrapAlias{conn: c}
// Even a legacy alias with no replacement metadata callback must hold the
// scheduler barrier until its physical replay reaches a terminal state.
alias.beginReplayRestore()
if alias.finishReplayRestore == nil {
t.Fatal("rewrap alias did not install its replay restore barrier")
}
job := s.rpcRewrapRestoreJob(alias, "panic regression", func(*rpcRewrapDeliveryControl, time.Time) {
panic("boom")
})
runRPCRewrapDeliveryJob(job)
c.rpcMu.Lock()
pending := c.rpcReplayRestores
c.rpcMu.Unlock()
if pending != 0 {
t.Fatalf("replay restore barriers after panic = %d, want 0", pending)
}
if !c.isRetired() {
t.Fatal("panic in a rewrap job did not fence the partially restored connection")
}
var ran atomic.Bool
var deadlineErr error
runRPCRewrapDeliveryJob(rpcRewrapDeliveryJob{
deadline: time.Now().Add(-time.Millisecond),
run: func(*rpcRewrapDeliveryControl, time.Time) { ran.Store(true) },
fail: func(err error) { deadlineErr = err },
})
if ran.Load() {
t.Fatal("expired rewrap job ran after its absolute queue deadline")
}
if !errors.Is(deadlineErr, context.DeadlineExceeded) {
t.Fatalf("expired rewrap job error = %v, want context deadline", deadlineErr)
}
// Recovery is per job: a panic must not poison the worker loop's next item.
runRPCRewrapDeliveryJob(rpcRewrapDeliveryJob{
deadline: time.Now().Add(time.Second),
run: func(*rpcRewrapDeliveryControl, time.Time) { ran.Store(true) },
})
if !ran.Load() {
t.Fatal("rewrap worker did not remain usable after a recovered panic")
}
}
func TestExpiredRPCRewrapResultJobPublishesCompletedAliasExactlyOnce(t *testing.T) {
cache := newRPCResultCacheWithFlightLimit(time.Now, 1)
s := &Server{log: zaptest.NewLogger(t), rpcResults: cache}
c := &Conn{
metrics: NopMetrics{},
authKeyID: [8]byte{1, 2, 3, 4, 5, 6, 7, 8},
sessionID: 7001,
}
const reqMsgID = int64(8001)
claim, err := cache.Acquire(c.authKeyID, c.sessionID, reqMsgID)
if err != nil {
t.Fatalf("acquire aliased result owner: %v", err)
}
if claim.state != rpcResultAcquireOwner || claim.owner == nil {
t.Fatalf("aliased result claim = %#v, want owner", claim)
}
if !claim.owner.CompleteExecution(true) {
t.Fatal("complete aliased result execution")
}
encoded := encodedRPCResultForPriorityTest(reqMsgID, 0)
encoded.delivery = claim.owner.Delivery()
alias := &rpcRewrapAlias{
conn: c,
newReqID: reqMsgID,
method: "help.getConfig",
newOwner: claim.owner,
}
alias.finishReplayRestore = c.beginRPCReplayRestore()
var ran atomic.Bool
job := s.rpcRewrapRestoreJob(alias, "expired aliased result", func(*rpcRewrapDeliveryControl, time.Time) {
ran.Store(true)
})
job.deadline = time.Now().Add(-time.Millisecond)
job.fail = func(err error) {
if !errors.Is(err, context.DeadlineExceeded) {
t.Errorf("expired result job error = %v, want context deadline", err)
}
s.failRPCRewrapResultJob(alias, encoded, err)
}
runRPCRewrapDeliveryJob(job)
if ran.Load() {
t.Fatal("expired aliased result job executed")
}
if !c.isRetired() {
t.Fatal("expired aliased result job did not fence its connection")
}
c.rpcMu.Lock()
pendingRestores := c.rpcReplayRestores
c.rpcMu.Unlock()
if pendingRestores != 0 {
t.Fatalf("replay restore barriers = %d, want 0", pendingRestores)
}
if used := cache.flightLimit.snapshot(); used != 0 {
t.Fatalf("pending result flights = %d, want 0", used)
}
completed, ok := cache.Get(c.authKeyID, c.sessionID, reqMsgID)
if !ok || completed != encoded {
t.Fatalf("completed aliased result = (%p, %v), want (%p, true)", completed, ok, encoded)
}
replay, err := cache.Acquire(c.authKeyID, c.sessionID, reqMsgID)
if err != nil {
t.Fatalf("reacquire completed aliased result: %v", err)
}
if replay.state != rpcResultAcquireCompleted || replay.encoded != encoded ||
!replay.executionKnown || !replay.executionOK {
t.Fatalf("completed aliased result metadata = %#v", replay)
}
// A defensive duplicate failure report must not republish or underflow the
// completed flight. The first handoff/cache completion is the sole winner.
s.failRPCRewrapResultJob(alias, encoded, context.DeadlineExceeded)
if used := cache.flightLimit.snapshot(); used != 0 {
t.Fatalf("pending result flights after duplicate failure = %d, want 0", used)
}
if got, ok := cache.Get(c.authKeyID, c.sessionID, reqMsgID); !ok || got != encoded {
t.Fatalf("completed result changed after duplicate failure = (%p, %v)", got, ok)
}
}
func TestRetargetedRPCRestoreIsOrderedAndIndependentOfGlobalHookExecutor(t *testing.T) {
executor := newRPCDeliveryHookExecutor(1, 1)
started := make(chan struct{})
release := make(chan struct{})
ticket, ok := executor.reserve()
if !ok || !ticket.submit(func() {
close(started)
<-release
}) {
t.Fatal("occupy delivery hook executor")
}
<-started
oldExecutor := defaultRPCDeliveryHookExecutor
defaultRPCDeliveryHookExecutor = executor
defer func() {
defaultRPCDeliveryHookExecutor = oldExecutor
close(release)
}()
s := New(Options{})
c := &Conn{
metrics: NopMetrics{},
authKeyID: [8]byte{9, 8, 7, 6, 5, 4, 3, 2},
sessionID: 9101,
}
const oldReqID, newReqID = int64(9201), int64(9202)
oldClaim, err := s.rpcResults.Acquire(c.authKeyID, c.sessionID, oldReqID)
if err != nil || oldClaim.state != rpcResultAcquireOwner {
t.Fatalf("old claim = %#v, err=%v", oldClaim, err)
}
body := []byte{1, 3, 3, 7}
if !s.rpcRewrap.register(c, body, oldReqID, "help.getConfig", oldClaim.owner) {
t.Fatal("register source rewrap candidate")
}
candidate := s.rpcRewrap.claim(c, body)
if candidate == nil {
t.Fatal("claim source rewrap candidate")
}
newClaim, err := s.rpcResults.Acquire(c.authKeyID, c.sessionID, newReqID)
if err != nil || newClaim.state != rpcResultAcquireOwner {
t.Fatalf("new claim = %#v, err=%v", newClaim, err)
}
var order atomic.Int32
alias := &rpcRewrapAlias{
conn: c, newReqID: newReqID, method: "help.getConfig",
oldWaiter: oldClaim.owner.Waiter(), newOwner: newClaim.owner,
sourceConn: c, sourceOwner: oldClaim.owner,
candidate: candidate, registry: s.rpcRewrap,
afterSuccessfulDelivery: func() error {
if !order.CompareAndSwap(0, 1) {
return errors.New("replacement restore ran out of order")
}
return nil
},
}
if err := alias.activate(s); err != nil {
t.Fatalf("activate retarget alias: %v", err)
}
if !alias.retargeted.Load() {
t.Fatal("source result was not retargeted")
}
encoded := encodedRPCResultForPriorityTest(oldReqID, 0)
encoded.delivery = oldClaim.owner.Delivery()
encoded.setDeliveryHook(func() {
if !order.CompareAndSwap(1, 2) {
panic("logical hook did not run after replacement restore")
}
})
if !oldClaim.owner.CompleteExecution(true) {
t.Fatal("complete source execution")
}
if !oldClaim.owner.HandOff() {
t.Fatal("handoff source result")
}
encoded.markQueued()
if got := encoded.beginWriting(); got != newReqID {
t.Fatalf("retargeted physical req_msg_id = %d, want %d", got, newReqID)
}
encoded.markDelivered()
s.rpcResults.Put(c.authKeyID, c.sessionID, oldReqID, encoded)
deadline := time.Now().Add(time.Second)
for order.Load() != 2 && time.Now().Before(deadline) {
time.Sleep(time.Millisecond)
}
if got := order.Load(); got != 2 {
t.Fatalf("retarget restore order = %d, want replacement then logical", got)
}
c.rpcMu.Lock()
pending := c.rpcReplayRestores
c.rpcMu.Unlock()
if pending != 0 {
t.Fatalf("retarget restore barriers = %d, want 0", pending)
}
if _, ok := s.rpcResults.Get(c.authKeyID, c.sessionID, newReqID); !ok {
t.Fatal("retargeted result was not cached under new req_msg_id")
}
}
func TestRPCRewrapWatchdogReleasesQueuedBarrierWhenAllWorkersAreStuck(t *testing.T) {
var (
once sync.Once
jobs chan rpcRewrapDeliveryJob
)
started := make(chan struct{}, rpcRewrapDeliveryWorkers)
release := make(chan struct{})
done := make(chan struct{}, rpcRewrapDeliveryWorkers)
for range rpcRewrapDeliveryWorkers {
if !scheduleRPCRewrapJob(rpcRewrapDeliveryJob{
deadline: time.Now().Add(time.Second),
run: func(*rpcRewrapDeliveryControl, time.Time) {
started <- struct{}{}
<-release // Deliberately ignores its deadline.
done <- struct{}{}
},
}, &once, &jobs, rpcRewrapDeliveryWorkers, 8) {
t.Fatal("schedule blocking rewrap worker")
}
}
for range rpcRewrapDeliveryWorkers {
select {
case <-started:
case <-time.After(time.Second):
t.Fatal("rewrap workers did not all block")
}
}
s := &Server{log: zaptest.NewLogger(t)}
c := &Conn{metrics: NopMetrics{}}
alias := &rpcRewrapAlias{conn: c}
alias.beginReplayRestore()
var ran atomic.Bool
job := s.rpcRewrapRestoreJob(alias, "queued watchdog regression", func(*rpcRewrapDeliveryControl, time.Time) {
ran.Store(true)
})
job.deadline = time.Now().Add(30 * time.Millisecond)
if !scheduleRPCRewrapJob(job, &once, &jobs, rpcRewrapDeliveryWorkers, 8) {
t.Fatal("schedule watched rewrap job")
}
deadline := time.Now().Add(time.Second)
for time.Now().Before(deadline) {
c.rpcMu.Lock()
pending := c.rpcReplayRestores
c.rpcMu.Unlock()
if c.isRetired() && pending == 0 {
break
}
time.Sleep(time.Millisecond)
}
c.rpcMu.Lock()
pending := c.rpcReplayRestores
c.rpcMu.Unlock()
if !c.isRetired() || pending != 0 {
t.Fatalf("watchdog terminal state = retired:%v barriers:%d", c.isRetired(), pending)
}
close(release)
for range rpcRewrapDeliveryWorkers {
select {
case <-done:
case <-time.After(time.Second):
t.Fatal("blocking rewrap worker did not exit")
}
}
// Give a worker the opportunity to dequeue the expired job. Its atomic
// terminal state must suppress the late run entirely.
time.Sleep(10 * time.Millisecond)
if ran.Load() {
t.Fatal("expired queued rewrap job ran after watchdog failure")
}
}
type deliveredThenBlockedTransport struct {
collectingSessionTransport
delivered chan struct{}
release chan struct{}
once sync.Once
}
func newDeliveredThenBlockedTransport() *deliveredThenBlockedTransport {
return &deliveredThenBlockedTransport{
delivered: make(chan struct{}),
release: make(chan struct{}),
}
}
func (t *deliveredThenBlockedTransport) Send(ctx context.Context, b *bin.Buffer) error {
if err := t.collectingSessionTransport.Send(ctx, b); err != nil {
return err
}
t.once.Do(func() { close(t.delivered) })
<-t.release // Simulate a transport that reports success late and ignores ctx.
return nil
}
func TestRPCRewrapPhysicalSuccessAfterWatchdogStillRunsLogicalRestore(t *testing.T) {
s := New(Options{WriteTimeout: time.Second})
transport := newDeliveredThenBlockedTransport()
defer func() {
select {
case <-transport.release:
default:
close(transport.release)
}
}()
key := newTestAuthKey(t)
c := s.newConn(transport, key, 9401, 1)
legacyCanonicalTestConn(t, c)
t.Cleanup(c.ForceClose)
const reqMsgID = int64(9402)
claim, err := s.rpcResults.Acquire(c.authKeyID, c.sessionID, reqMsgID)
if err != nil || claim.state != rpcResultAcquireOwner {
t.Fatalf("rewrap claim = %#v, err=%v", claim, err)
}
if !claim.owner.CompleteExecution(true) {
t.Fatal("complete rewrap execution")
}
var order atomic.Int32
encoded := encodedRPCResultForPriorityTest(reqMsgID, 0)
encoded.delivery = claim.owner.Delivery()
encoded.setDeliveryHook(func() {
if !order.CompareAndSwap(1, 2) {
panic("late physical success lost ordered logical restore")
}
})
alias := &rpcRewrapAlias{
conn: c, newReqID: reqMsgID, method: "help.getConfig", newOwner: claim.owner,
afterSuccessfulDelivery: func() error {
if !order.CompareAndSwap(0, 1) {
return errors.New("late replacement restore ran out of order")
}
return nil
},
}
alias.executionOK.Store(true)
alias.beginReplayRestore()
var (
once sync.Once
jobs chan rpcRewrapDeliveryJob
)
done := make(chan struct{})
job := s.rpcRewrapRestoreJob(alias, "late physical success regression", func(control *rpcRewrapDeliveryControl, deadline time.Time) {
s.publishRewrappedRPCResult(c, reqMsgID, alias.method, claim.owner, encoded, alias, control, deadline)
close(done)
})
job.deadline = time.Now().Add(30 * time.Millisecond)
job.fail = func(err error) { s.failRPCRewrapResultJob(alias, encoded, err) }
if !scheduleRPCRewrapJob(job, &once, &jobs, 1, 1) {
t.Fatal("schedule late-success rewrap job")
}
select {
case <-transport.delivered:
case <-time.After(time.Second):
t.Fatal("transport did not reach physical delivery gate")
}
deadline := time.Now().Add(time.Second)
for time.Now().Before(deadline) {
c.rpcMu.Lock()
pending := c.rpcReplayRestores
c.rpcMu.Unlock()
if c.isRetired() && pending == 0 {
break
}
time.Sleep(time.Millisecond)
}
if got := order.Load(); got != 0 {
t.Fatalf("restore ran before transport reported terminal success: %d", got)
}
close(transport.release)
select {
case <-done:
case <-time.After(time.Second):
t.Fatal("late successful physical attempt did not finish")
}
if got := order.Load(); got != 2 {
t.Fatalf("late physical success restore order = %d, want 2", got)
}
if got := len(transport.snapshot()); got != 1 {
t.Fatalf("late physical success writes = %d, want 1", got)
}
}
func TestConcurrentRPCRewrapDeliveredFinalizationPublishesOnceWithMetadata(t *testing.T) {
cache := newRPCResultCacheWithFlightLimit(time.Now, 1)
s := &Server{log: zaptest.NewLogger(t), rpcResults: cache}
c := &Conn{
metrics: NopMetrics{},
authKeyID: [8]byte{4, 4, 4, 4, 4, 4, 4, 4},
sessionID: 9501,
}
const reqMsgID = int64(9502)
claim, err := cache.Acquire(c.authKeyID, c.sessionID, reqMsgID)
if err != nil || claim.state != rpcResultAcquireOwner {
t.Fatalf("finalization claim = %#v, err=%v", claim, err)
}
if !claim.owner.CompleteExecution(true) || !claim.owner.HandOff() {
t.Fatal("prepare finalization owner")
}
var subscribers atomic.Int32
if err := claim.owner.Waiter().Subscribe(func(*encodedOutboundMessage, bool) {
subscribers.Add(1)
}); err != nil {
t.Fatalf("subscribe finalization probe: %v", err)
}
var replacement, logical atomic.Int32
encoded := encodedRPCResultForPriorityTest(reqMsgID, 0)
encoded.delivery = claim.owner.Delivery()
encoded.setDeliveryHook(func() { logical.Add(1) })
alias := &rpcRewrapAlias{
conn: c, newReqID: reqMsgID, method: "help.getConfig", newOwner: claim.owner,
afterSuccessfulDelivery: func() error {
replacement.Add(1)
return nil
},
}
alias.executionOK.Store(true)
alias.beginReplayRestore()
start := make(chan struct{})
var wg sync.WaitGroup
for _, source := range []string{"worker", "watchdog"} {
wg.Add(1)
go func(source string) {
defer wg.Done()
<-start
_ = s.completeDeliveredRPCRewrapResult(context.Background(), alias, encoded, source)
}(source)
}
close(start)
wg.Wait()
if got := replacement.Load(); got != 1 {
t.Fatalf("replacement finalizations = %d, want 1", got)
}
if got := logical.Load(); got != 1 {
t.Fatalf("logical finalizations = %d, want 1", got)
}
if got := subscribers.Load(); got != 1 {
t.Fatalf("cache subscriber calls = %d, want 1", got)
}
replay, err := cache.Acquire(c.authKeyID, c.sessionID, reqMsgID)
if err != nil || replay.state != rpcResultAcquireCompleted ||
!replay.executionKnown || !replay.executionOK {
t.Fatalf("completed finalization metadata = %#v, err=%v", replay, err)
}
c.rpcMu.Lock()
pending := c.rpcReplayRestores
c.rpcMu.Unlock()
if pending != 0 {
t.Fatalf("finalization barriers = %d, want 0", pending)
}
}
func TestRPCRewrapSubscriberPanicCannotLoseClaimedLogicalHook(t *testing.T) {
s := New(Options{WriteTimeout: time.Second})
transport := &collectingSessionTransport{}
key := newTestAuthKey(t)
c := s.newConn(transport, key, 9301, 1)
legacyCanonicalTestConn(t, c)
t.Cleanup(c.ForceClose)
const reqMsgID = int64(9302)
claim, err := s.rpcResults.Acquire(c.authKeyID, c.sessionID, reqMsgID)
if err != nil || claim.state != rpcResultAcquireOwner {
t.Fatalf("alias claim = %#v, err=%v", claim, err)
}
if !claim.owner.CompleteExecution(true) {
t.Fatal("complete alias execution")
}
if err := claim.owner.Waiter().Subscribe(func(*encodedOutboundMessage, bool) {
panic("subscriber boom")
}); err != nil {
t.Fatalf("subscribe panic probe: %v", err)
}
var order atomic.Int32
encoded := encodedRPCResultForPriorityTest(reqMsgID, 0)
encoded.delivery = claim.owner.Delivery()
encoded.setDeliveryHook(func() {
if !order.CompareAndSwap(1, 2) {
panic("logical hook did not run after replacement restore")
}
})
alias := &rpcRewrapAlias{
conn: c, newReqID: reqMsgID, method: "help.getConfig", newOwner: claim.owner,
afterSuccessfulDelivery: func() error {
if !order.CompareAndSwap(0, 1) {
return errors.New("replacement restore ran out of order")
}
return nil
},
}
alias.executionOK.Store(true)
alias.beginReplayRestore()
job := s.rpcRewrapRestoreJob(alias, "subscriber panic regression", func(control *rpcRewrapDeliveryControl, deadline time.Time) {
s.publishRewrappedRPCResult(c, reqMsgID, alias.method, claim.owner, encoded, alias, control, deadline)
})
job.fail = func(err error) { s.failRPCRewrapResultJob(alias, encoded, err) }
runRPCRewrapDeliveryJob(job)
if got := order.Load(); got != 2 {
t.Fatalf("restore order after subscriber panic = %d, want 2", got)
}
if !c.isRetired() {
t.Fatal("subscriber panic did not fence the connection")
}
c.rpcMu.Lock()
pending := c.rpcReplayRestores
c.rpcMu.Unlock()
if pending != 0 {
t.Fatalf("restore barriers after subscriber panic = %d, want 0", pending)
}
if cached, ok := s.rpcResults.Get(c.authKeyID, c.sessionID, reqMsgID); !ok || cached != encoded {
t.Fatalf("completed result after subscriber panic = (%p, %v), want (%p, true)", cached, ok, encoded)
}
}

74
internal/mtprotoedge/rpc_test.go
View file

@ -8,14 +8,14 @@ import (
"go.uber.org/zap/zaptest"
"github.com/gotd/td/bin"
"github.com/gotd/td/clock"
"github.com/gotd/td/crypto"
"github.com/gotd/td/mt"
"github.com/gotd/td/proto"
"github.com/gotd/td/tg"
"github.com/gotd/td/tgerr"
"github.com/gotd/td/transport"
"github.com/iamxvbaba/td/bin"
"github.com/iamxvbaba/td/clock"
"github.com/iamxvbaba/td/crypto"
"github.com/iamxvbaba/td/mt"
"github.com/iamxvbaba/td/proto"
"github.com/iamxvbaba/td/tg"
"github.com/iamxvbaba/td/tgerr"
"github.com/iamxvbaba/td/transport"
"telesrv/internal/rpc"
)
@ -29,7 +29,7 @@ func TestRPCGetConfig(t *testing.T) {
advPort = 12345
)
router := rpc.New(rpc.Config{DC: dc, IP: advIP, Port: advPort}, rpc.Deps{}, zaptest.NewLogger(t), clock.System)
addr, pub, _ := startTestServer(t, Options{DC: dc, RPC: router})
addr, pub, _ := startTestServer(t, Options{DC: dc, legacyRPC: router})
conn, auth, cipher := dialHandshake(t, addr, dc, pub)
clientMsgID := proto.NewMessageIDGen(time.Now)
@ -67,6 +67,52 @@ func TestRPCGetConfig(t *testing.T) {
}
}
func TestLayerRPCGetConfigUsesExactAdmittedProfile(t *testing.T) {
const (
dc = 2
advIP = "127.0.0.1"
advPort = 12345
)
router := rpc.New(rpc.Config{DC: dc, IP: advIP, Port: advPort}, rpc.Deps{}, zaptest.NewLogger(t), clock.System)
addr, pub, _ := startTestServer(t, Options{DC: dc, LayerRPC: router})
conn, auth, cipher := dialHandshake(t, addr, dc, pub)
clientMsgID := proto.NewMessageIDGen(time.Now)
reqMsgID := clientMsgID.New(proto.MessageFromClient)
request := &tg.InvokeWithLayerRequest{
Layer: int(tg.LayerProfile225),
Query: &tg.InitConnectionRequest{
APIID: 123,
DeviceModel: "Desktop",
SystemVersion: "Windows",
AppVersion: "test",
SystemLangCode: "en",
LangPack: "tdesktop",
LangCode: "en",
Query: &tg.HelpGetConfigRequest{},
},
}
sendEncrypted(t, conn, cipher, auth, reqMsgID, request)
replies := collectReplies(t, conn, cipher, auth.AuthKey, proto.ResultTypeID)
resultBuf := mustHave(t, replies, proto.ResultTypeID, "rpc_result")
var result proto.Result
if err := result.Decode(resultBuf); err != nil {
t.Fatalf("decode rpc_result: %v", err)
}
if result.RequestMessageID != reqMsgID {
t.Fatalf("rpc_result req_msg_id = %d, want %d", result.RequestMessageID, reqMsgID)
}
exact := &bin.Buffer{Buf: result.Result}
config, err := tg.DecodeLayer(tg.LayerProfile225, tg.LayerConstructorConfigType(), exact)
if err != nil {
t.Fatalf("decode layer 225 config: %v", err)
}
if exact.Len() != 0 || config.ThisDC != dc {
t.Fatalf("layer 225 config = dc:%d remaining:%d", config.ThisDC, exact.Len())
}
}
func TestInboundRPCQueueFullReturnsFloodWait(t *testing.T) {
const dc = 2
handler := &blockingRPC{
@ -75,7 +121,7 @@ func TestInboundRPCQueueFullReturnsFloodWait(t *testing.T) {
}
addr, pub, _ := startTestServer(t, Options{
DC: dc,
RPC: handler,
legacyRPC: handler,
RPCMaxInflight: 1,
RPCQueueSize: 1,
RPCTimeout: 5 * time.Second,
@ -115,7 +161,7 @@ func TestInboundRPCQueuedDeadlineReturnsRPCTimeout(t *testing.T) {
}
addr, pub, _ := startTestServer(t, Options{
DC: dc,
RPC: handler,
legacyRPC: handler,
RPCMaxInflight: 1,
RPCQueueSize: 2,
RPCTimeout: 60 * time.Millisecond,
@ -168,7 +214,7 @@ func TestInboundRPCRunningDeadlineWaitsForHandlerTerminalResult(t *testing.T) {
}
addr, pub, _ := startTestServer(t, Options{
DC: dc,
RPC: handler,
legacyRPC: handler,
RPCMaxInflight: 1,
RPCQueueSize: 1,
RPCTimeout: 60 * time.Millisecond,
@ -215,7 +261,7 @@ func TestInboundRPCRunningDeadlineWaitsForHandlerTerminalResult(t *testing.T) {
func TestDuplicateRPCResultAcrossReconnectUsesSessionCache(t *testing.T) {
const dc = 2
handler := &countingConfigRPC{}
addr, pub, _ := startTestServer(t, Options{DC: dc, RPC: handler})
addr, pub, _ := startTestServer(t, Options{DC: dc, legacyRPC: handler})
conn, auth, cipher := dialHandshake(t, addr, dc, pub)
clientMsgID := proto.NewMessageIDGen(time.Now)
@ -253,7 +299,7 @@ func TestCanceledRPCErrorIsNotCachedAcrossReconnect(t *testing.T) {
firstStarted: make(chan struct{}),
firstDone: make(chan struct{}),
}
addr, pub, _ := startTestServer(t, Options{DC: dc, RPC: handler})
addr, pub, _ := startTestServer(t, Options{DC: dc, legacyRPC: handler})
conn, auth, cipher := dialHandshake(t, addr, dc, pub)
clientMsgID := proto.NewMessageIDGen(time.Now)

288
internal/mtprotoedge/server.go
View file

@ -14,43 +14,191 @@ import (
"go.uber.org/zap"
"github.com/gotd/td/bin"
"github.com/gotd/td/clock"
"github.com/gotd/td/crypto"
"github.com/gotd/td/exchange"
"github.com/gotd/td/mt"
"github.com/gotd/td/proto"
"github.com/gotd/td/proto/codec"
"github.com/gotd/td/tg"
"github.com/gotd/td/tmap"
"github.com/gotd/td/transport"
"github.com/iamxvbaba/td/bin"
"github.com/iamxvbaba/td/clock"
"github.com/iamxvbaba/td/crypto"
"github.com/iamxvbaba/td/exchange"
"github.com/iamxvbaba/td/mt"
"github.com/iamxvbaba/td/proto"
"github.com/iamxvbaba/td/proto/codec"
"github.com/iamxvbaba/td/tg"
"github.com/iamxvbaba/td/tmap"
"github.com/iamxvbaba/td/transport"
"telesrv/internal/store"
"telesrv/internal/store/memory"
)
// RPCHandler 把解密后的 RPC 请求体路由到响应。由 internal/rpc 实现。
// legacyRPCHandler 把解密后的 canonical RPC 请求体路由到响应。
//
// b 是明文 RPC 请求（已剥离 MTProto 外壳）；返回的 bin.Encoder 会被包成 rpc_result。
// 返回 *tgerr.Error 时连接层将其转为 rpc_error 回发；其他 error 视为连接级故障。
type RPCHandler interface {
//
// 它只保留给本包旧连接状态机的回归测试。生产 API RPC 必须走 LayerRPCHandler，
// 使 admission、request/result TypeRef 与 exact profile 不可绕过。
type legacyRPCHandler interface {
Dispatch(ctx context.Context, authKeyID [8]byte, sessionID int64, b *bin.Buffer) (bin.Encoder, error)
// NegotiatedLayer returns the TL layer the session negotiated via
// invokeWithLayer and whether one was ever observed. Used to downgrade
// outbound objects for clients compiled on an older layer. ok=false means
// unknown (cold/evicted) — the caller must keep the connection's last-known
// layer rather than overwrite it.
// NegotiatedLayer returns the TL layer proven via invokeWithLayer for this
// exact (auth_key_id, session_id). It must never infer from API ID, device
// metadata, authorization rows, or another session on the same auth key.
NegotiatedLayer(authKeyID [8]byte, sessionID int64) (int, bool)
}
// RPCHandlerWithMethod returns the canonical innermost RPC method after the
// legacyRPCHandlerWithMethod returns the canonical innermost RPC method after the
// router has peeled invokeWithLayer/initConnection/invokeAfter wrappers. Egress
// scheduling must use this identity: the outer wrapper is not a useful signal
// for prioritizing updates convergence over catalog/media responses.
type RPCHandlerWithMethod interface {
type legacyRPCHandlerWithMethod interface {
DispatchWithMethod(ctx context.Context, authKeyID [8]byte, sessionID int64, b *bin.Buffer) (bin.Encoder, string, error)
}
// LayerRPCHandler is the production API-RPC boundary. Admission is a separate
// allocation-bounded phase so the edge can freeze the connection profile,
// validate wrapper dependencies and establish exact request identity before
// flight/cache/scheduler ownership is acquired.
type LayerRPCHandler interface {
AdmitLayer(profile tg.LayerProfile, b *bin.Buffer, limits tg.LayerDecodeLimits) (tg.LayerRequest, error)
AdmitUnprofiled(b *bin.Buffer, limits tg.LayerDecodeLimits) (tg.LayerRequest, error)
DispatchAdmitted(
ctx context.Context,
authKeyID [8]byte,
sessionID int64,
msgID int64,
admissionSeq uint64,
request tg.LayerRequest,
) (tg.LayerRPCResult, string, error)
}
// LayerRPCDefaultProfileAdmitter decodes with a recoverable inherited/default
// profile. Production handlers should implement it with the same generated
// ServerDispatcher and adapter registry used by AdmitLayer. The split keeps old
// test doubles source-compatible while allowing invokeWithLayer to correct even
// a previously explicit Conn profile.
type LayerRPCDefaultProfileAdmitter interface {
AdmitDefaultLayer(profile tg.LayerProfile, b *bin.Buffer, limits tg.LayerDecodeLimits) (tg.LayerRequest, error)
}
// LayerRPCSessionProfileResolver may restore an exact profile only when it was
// previously proven for this same (auth_key_id, session_id). Auth-key-wide
// device metadata is intentionally ineligible: a client upgrade can reuse its
// auth key while opening a new session at a newer Layer.
type LayerRPCSessionProfileResolver interface {
NegotiatedSessionLayer(authKeyID [8]byte, sessionID int64) (int, bool)
}
// LayerRPCOrderedSessionProfileResolver restores both the selected Layer and
// the newest invokeWithLayer client msg_id which proved it. The cursor prevents
// an old cached request replay on a replacement physical connection from
// rolling the logical session back to an older profile.
type LayerRPCOrderedSessionProfileResolver interface {
NegotiatedSessionLayerEvidence(authKeyID [8]byte, sessionID int64) (layer int, msgID int64, ok bool)
}
// LayerRPCInheritedAuthKeyProfileResolver resolves the best persisted
// auth-key-wide client Layer for a new session. Unlike exact same-session
// evidence this value is only an inherited default: a later invokeWithLayer may
// correct it. Implementations may resolve a temporary raw key through its bound
// permanent key, but must not infer from API ID or device strings.
type LayerRPCInheritedAuthKeyProfileResolver interface {
ResolveInheritedAuthKeyLayer(ctx context.Context, rawAuthKeyID [8]byte) (layer int, found bool, err error)
}
// LayerRPCSessionProfileRegistry atomically freezes and invalidates the exact
// same-session profile. Its bounded retention may survive a TCP reconnect and
// admit a naked replay while the entry remains live. Expiry or capacity eviction
// loses only that proof and therefore fails closed until fresh invokeWithLayer;
// it never falls back to auth-key-wide metadata.
type LayerRPCSessionProfileRegistry interface {
LayerRPCSessionProfileResolver
FreezeNegotiatedSessionLayer(authKeyID [8]byte, sessionID int64, layer int) error
ForgetNegotiatedSessionLayer(authKeyID [8]byte, sessionID int64)
ForgetNegotiatedAuthKey(authKeyID [8]byte)
}
// LayerRPCOrderedSessionProfileRegistry linearizes explicit Layer evidence by
// MTProto client msg_id. applied is false for an older or identical duplicate;
// the same msg_id with another Layer must return ErrLayerProfileConflict (or an
// error wrapping it). Implementations advance the cursor even when Layer is
// unchanged.
type LayerRPCOrderedSessionProfileRegistry interface {
LayerRPCOrderedSessionProfileResolver
FreezeNegotiatedSessionLayerAt(authKeyID [8]byte, sessionID int64, layer int, msgID int64) (applied bool, err error)
}
// LayerRPCDurableSessionProfileResolver restores restart-safe raw Layer
// evidence. layer may be newer than this binary's generated codec universe;
// msgID remains the ordering authority in that case.
type LayerRPCDurableSessionProfileResolver interface {
ResolveNegotiatedSessionLayerEvidence(
ctx context.Context,
rawAuthKeyID [8]byte,
sessionID int64,
) (layer int, msgID int64, found bool, err error)
}
// LayerRPCDurableSessionProfileAdvancer atomically advances exact-session and
// auth-key shared-default evidence. publishShared is true only when this exact
// observation still owns the durable shared default.
type LayerRPCDurableSessionProfileAdvancer interface {
AdvanceNegotiatedSessionLayerEvidence(
ctx context.Context,
rawAuthKeyID [8]byte,
sessionID int64,
layer int,
msgID int64,
) (currentLayer int, currentMsgID int64, publishShared bool, err error)
}
type LayerRPCDurableSessionProfileDeleter interface {
DeleteNegotiatedSessionLayerEvidence(
ctx context.Context,
rawAuthKeyID [8]byte,
sessionID int64,
) (deleted bool, err error)
}
// LayerRPCReplayPreparer reapplies connection-local wrapper state for an
// already-executed exact request without consuming its one-shot business
// dispatch lease. The returned callback is safe to run only after a successful
// cached rpc_result reaches the replacement physical connection.
type LayerRPCReplayPreparer interface {
PrepareAdmittedReplay(
ctx context.Context,
authKeyID [8]byte,
sessionID int64,
msgID int64,
admissionSeq uint64,
request tg.LayerRequest,
) (afterSuccessfulDelivery func() error, err error)
}
// LayerRPCProfileEvidenceContext lets a generated/semantic handler carry the
// edge's MTProto msg_id freshness decision through its existing context-based
// dispatch API. fresh=false means the request remains fully request-bound: it
// may decode, execute and produce an exact result, but it must not publish
// mutable Layer/init/readiness/auth-bind state. Production Router implements
// this optional decorator; older handlers which have no such shared state stay
// source compatible.
type LayerRPCProfileEvidenceContext interface {
WithLayerRPCProfileEvidenceFresh(ctx context.Context, fresh bool) context.Context
}
// LayerRPCAdmissionProfilePublisher advances the auth-key-wide inherited
// default for fresh explicit evidence. admissionSeq is allocated once by the
// edge's exact flight owner and globally orders different MTProto sessions;
// cached joins/replays never call this hook again.
type LayerRPCAdmissionProfilePublisher interface {
PublishAdmittedLayerProfileEvidence(
ctx context.Context,
rawAuthKeyID [8]byte,
sessionID int64,
msgID int64,
admissionSeq uint64,
safeFloor uint64,
layer int,
) error
}
// RPCInitConnectionObserver records wrapper metadata when the edge aliases an
// initConnection reissue to an already-running request and therefore correctly
// skips a second business Dispatch.
@ -115,8 +263,23 @@ type Options struct {
RPCGlobalWorkers int
// RPCGlobalMaxTasks 是全进程已预留、排队和执行中的 RPC 条数上限。默认 8192。
RPCGlobalMaxTasks int
// RPCGlobalMaxBytes 是上述 RPC body 的总字节预算。默认 512 MiB。
// RPCGlobalMaxBytes 是上述 RPC 的进程级 memory charge 预算。legacy charge
// 等于 copied body；exact charge 是 typed decode 前的保守 materialization
// 上界，因此该配置不表示可并发接收 512 MiB wire body。默认 512 MiB。
RPCGlobalMaxBytes int64
// RPCResultCache* limits bound pending ownership and completed rpc_result
// replay state across the full 331-second duplicate horizon. Every owner is
// charged simultaneously at global, raw-auth and session scopes. Defaults:
// global 262144/64 MiB, auth 32768/32 MiB, session 16384/16 MiB.
RPCResultCacheMaxEntries int
RPCResultCacheMaxBytes int64
RPCResultCacheAuthMaxEntries int
RPCResultCacheAuthMaxBytes int64
RPCResultCacheSessionMaxEntries int
RPCResultCacheSessionMaxBytes int64
// RPCResultPendingPerAuth is an additional active-owner bound, independent
// from the retained entry limits and RPCGlobalMaxTasks. Default 2048.
RPCResultPendingPerAuth int
// InboundFrameGlobalMaxBytes 是所有物理连接当前正在处理的 transport wire buffer
// 与最大解密 plaintext buffer 的总预算。长度前缀读取后、payload 分配前预留，默认
// 512 MiB；非正值使用默认值。
@ -143,8 +306,13 @@ type Options struct {
AuthKeys store.AuthKeyStore
// ActiveSessions 管理活跃连接。默认新建；传入时可让 RPC 层共享同一注册表。
ActiveSessions *SessionManager
// RPC 是 typed RPC 路由。nil 时加密 RPC 被丢弃并记录。
RPC RPCHandler
// legacyRPC is an internal test hook for the pre-exact RPC state machine.
// It is deliberately unexported so production callers cannot bypass
// generated Layer admission by configuring the canonical-only route.
legacyRPC legacyRPCHandler
// LayerRPC is the generated exact-profile production path. When configured,
// every API request must complete admission before flight/cache scheduling.
LayerRPC LayerRPCHandler
// Metrics 接收连接层指标。默认 NopMetrics。
Metrics Metrics
// OnServing is called after the connection intake loops have been installed.
@ -200,6 +368,30 @@ func (o *Options) setDefaults() {
if o.RPCGlobalMaxBytes <= 0 {
o.RPCGlobalMaxBytes = 512 << 20
}
if o.RPCResultCacheMaxEntries == 0 {
o.RPCResultCacheMaxEntries = rpcResultCacheMaxEntries
}
if o.RPCResultCacheMaxBytes == 0 {
o.RPCResultCacheMaxBytes = rpcResultCacheMaxBytes
}
if o.RPCResultCacheAuthMaxEntries == 0 {
o.RPCResultCacheAuthMaxEntries = rpcResultCacheAuthMaxEntries
}
if o.RPCResultCacheAuthMaxBytes == 0 {
o.RPCResultCacheAuthMaxBytes = rpcResultCacheAuthMaxBytes
}
if o.RPCResultCacheSessionMaxEntries == 0 {
o.RPCResultCacheSessionMaxEntries = rpcResultCacheSessionMaxEntries
}
if o.RPCResultCacheSessionMaxBytes == 0 {
o.RPCResultCacheSessionMaxBytes = rpcResultCacheSessionMaxBytes
}
if o.RPCResultPendingPerAuth == 0 {
o.RPCResultPendingPerAuth = rpcResultFlightMaxPendingPerAuth
if o.RPCResultPendingPerAuth > o.RPCGlobalMaxTasks {
o.RPCResultPendingPerAuth = o.RPCGlobalMaxTasks
}
}
if o.InboundFrameGlobalMaxBytes <= 0 {
o.InboundFrameGlobalMaxBytes = defaultInboundFrameGlobalMaxBytes
}
@ -232,6 +424,34 @@ func (o *Options) setDefaults() {
}
}
func validateRPCResultCacheOptions(o Options) error {
if o.RPCResultCacheMaxEntries <= 0 || o.RPCResultCacheAuthMaxEntries <= 0 || o.RPCResultCacheSessionMaxEntries <= 0 {
return fmt.Errorf("rpc_result cache entry limits must be positive")
}
if o.RPCResultCacheMaxEntries < o.RPCResultCacheAuthMaxEntries ||
o.RPCResultCacheAuthMaxEntries < o.RPCResultCacheSessionMaxEntries {
return fmt.Errorf("rpc_result cache entry hierarchy must satisfy global >= auth >= session: %d/%d/%d",
o.RPCResultCacheMaxEntries, o.RPCResultCacheAuthMaxEntries, o.RPCResultCacheSessionMaxEntries)
}
if o.RPCResultCacheMaxBytes < int64(maxOutboundBodyBytes) ||
o.RPCResultCacheAuthMaxBytes < int64(maxOutboundBodyBytes) ||
o.RPCResultCacheSessionMaxBytes < int64(maxOutboundBodyBytes) {
return fmt.Errorf("rpc_result cache byte limits must each be at least max outbound body %d: %d/%d/%d",
maxOutboundBodyBytes, o.RPCResultCacheMaxBytes, o.RPCResultCacheAuthMaxBytes, o.RPCResultCacheSessionMaxBytes)
}
if o.RPCResultCacheMaxBytes < o.RPCResultCacheAuthMaxBytes ||
o.RPCResultCacheAuthMaxBytes < o.RPCResultCacheSessionMaxBytes {
return fmt.Errorf("rpc_result cache byte hierarchy must satisfy global >= auth >= session: %d/%d/%d",
o.RPCResultCacheMaxBytes, o.RPCResultCacheAuthMaxBytes, o.RPCResultCacheSessionMaxBytes)
}
if o.RPCResultPendingPerAuth <= 0 || o.RPCResultPendingPerAuth > o.RPCGlobalMaxTasks ||
o.RPCResultPendingPerAuth > o.RPCResultCacheAuthMaxEntries {
return fmt.Errorf("rpc_result per-auth pending limit %d must be positive and <= global pending %d and auth entries %d",
o.RPCResultPendingPerAuth, o.RPCGlobalMaxTasks, o.RPCResultCacheAuthMaxEntries)
}
return nil
}
// Server 是 MTProto 连接层（mtprotoedge）。
//
// 职责见 doc.go。它把原始 TCP 字节流转换为「已解密、已识别 session 的 RPC 请求」：
@ -262,7 +482,8 @@ type Server struct {
key exchange.PrivateKey
authKeys store.AuthKeyStore
conns *SessionManager
rpc RPCHandler
rpc legacyRPCHandler
layerRPC LayerRPCHandler
metrics Metrics
onServing func(net.Addr)
cipher crypto.Cipher
@ -281,6 +502,9 @@ type Server struct {
// New 创建 Server。
func New(opts Options) *Server {
opts.setDefaults()
if err := validateRPCResultCacheOptions(opts); err != nil {
panic(fmt.Sprintf("mtprotoedge: invalid result-cache options: %v", err))
}
conns := opts.ActiveSessions
if conns == nil {
conns = NewSessionManager(opts.Logger.Named("sessions"))
@ -309,16 +533,26 @@ func New(opts Options) *Server {
key: exchange.PrivateKey{RSA: opts.RSAKey},
authKeys: opts.AuthKeys,
conns: conns,
rpc: opts.RPC,
rpc: opts.legacyRPC,
layerRPC: opts.LayerRPC,
metrics: opts.Metrics,
onServing: opts.OnServing,
cipher: crypto.NewServerCipher(opts.Rand),
clock: opts.Clock,
rand: opts.Rand,
types: tmap.New(tg.TypesMap(), mt.TypesMap(), proto.TypesMap()),
rpcResults: newRPCResultCacheWithFlightLimit(opts.Clock.Now, opts.RPCGlobalMaxTasks),
rpcRewrap: newRPCRewrapRegistry(opts.RPCGlobalMaxTasks),
admission: newAdmissionController(opts.MaxConnections, opts.MaxConnectionsPerIP, opts.MaxConcurrentHandshakes),
rpcResults: newRPCResultCacheWithFairCapacity(opts.Clock.Now, rpcResultCacheCapacity{
maxPending: opts.RPCGlobalMaxTasks,
maxPendingPerAuth: opts.RPCResultPendingPerAuth,
globalMaxBytes: opts.RPCResultCacheMaxBytes,
globalMaxEntries: opts.RPCResultCacheMaxEntries,
authMaxBytes: opts.RPCResultCacheAuthMaxBytes,
authMaxEntries: opts.RPCResultCacheAuthMaxEntries,
sessionMaxBytes: opts.RPCResultCacheSessionMaxBytes,
sessionMaxEntries: opts.RPCResultCacheSessionMaxEntries,
}),
rpcRewrap: newRPCRewrapRegistry(opts.RPCGlobalMaxTasks),
admission: newAdmissionController(opts.MaxConnections, opts.MaxConnectionsPerIP, opts.MaxConcurrentHandshakes),
}
}

12
internal/mtprotoedge/server_test.go
View file

@ -16,12 +16,12 @@ import (
"go.uber.org/zap/zaptest"
"github.com/gotd/td/bin"
"github.com/gotd/td/mtproxy"
"github.com/gotd/td/mtproxy/obfuscator"
"github.com/gotd/td/proto/codec"
"github.com/gotd/td/telegram/dcs"
"github.com/gotd/td/transport"
"github.com/iamxvbaba/td/bin"
"github.com/iamxvbaba/td/mtproxy"
"github.com/iamxvbaba/td/mtproxy/obfuscator"
"github.com/iamxvbaba/td/proto/codec"
"github.com/iamxvbaba/td/telegram/dcs"
"github.com/iamxvbaba/td/transport"
)
// TestServerAcceptAndCodec 验证 M0：

6
internal/mtprotoedge/session_activation_test.go
View file

@ -8,8 +8,8 @@ import (
"go.uber.org/zap/zaptest"
"github.com/gotd/td/mt"
"github.com/gotd/td/proto"
"github.com/iamxvbaba/td/mt"
"github.com/iamxvbaba/td/proto"
)
func newSessionActivationTestConn(t *testing.T, authKeyID [8]byte, sessionID int64) *Conn {
@ -222,7 +222,7 @@ func TestDestroyFencesOutboundAndReservedRPCBeforeRemoval(t *testing.T) {
reservation, err := c.reserveInboundRPC(context.Background(), "test.destroyFence", 8)
if err != nil {
t.Fatalf("reserve inbound RPC: %v", err)
t.Fatalf("reserve inbound legacyRPC: %v", err)
}
destroyed := make(chan bool, 1)
go func() {

24
internal/mtprotoedge/session_barrier_integration_test.go
View file

@ -10,11 +10,11 @@ import (
"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"
"github.com/iamxvbaba/td/bin"
"github.com/iamxvbaba/td/crypto"
"github.com/iamxvbaba/td/mt"
"github.com/iamxvbaba/td/proto"
"github.com/iamxvbaba/td/tg"
"telesrv/internal/store"
)
@ -44,7 +44,7 @@ func encryptedRPCFrameWithAuthoritativeSaltForBarrierTest(
MessageDataLen: int32(len(body)),
MessageDataWithPadding: body,
}, &frame); err != nil {
t.Fatalf("encrypt client RPC: %v", err)
t.Fatalf("encrypt client legacyRPC: %v", err)
}
return &frame, store.AuthKeyData{
ID: key.ID,
@ -55,7 +55,7 @@ func encryptedRPCFrameWithAuthoritativeSaltForBarrierTest(
func TestBadServerSaltRetainsOneProvisionalConnUntilCorrected(t *testing.T) {
handler := &admissionCountingRPC{}
s := New(Options{RPC: handler, WriteTimeout: time.Second})
s := New(Options{legacyRPC: handler, WriteTimeout: time.Second})
s.rpcScheduler.start()
t.Cleanup(func() { s.rpcScheduler.stop(time.Second) })
@ -149,7 +149,7 @@ func TestWrongSaltSessionChangeTransfersPhysicalOwnership(t *testing.T) {
// 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 := New(Options{legacyRPC: handler, WriteTimeout: time.Second})
s.rpcScheduler.start()
t.Cleanup(func() { s.rpcScheduler.stop(time.Second) })
@ -294,7 +294,7 @@ func (*cancelThenRetryRPC) NegotiatedLayer([8]byte, int64) (int, bool) { return
func TestHandleEncryptedRequiredSessionBarrierPrecedesStateRegistrationAndRPC(t *testing.T) {
handler := &admissionCountingRPC{}
s := New(Options{RPC: handler, WriteTimeout: time.Second})
s := New(Options{legacyRPC: handler, WriteTimeout: time.Second})
s.rpcScheduler.start()
t.Cleanup(func() { s.rpcScheduler.stop(time.Second) })
@ -374,7 +374,7 @@ func TestHandleEncryptedRequiredSessionBarrierPrecedesStateRegistrationAndRPC(t
func TestHandleEncryptedRequiredSessionBarrierFailureIsAtomic(t *testing.T) {
handler := &admissionCountingRPC{}
s := New(Options{RPC: handler, WriteTimeout: time.Second})
s := New(Options{legacyRPC: handler, WriteTimeout: time.Second})
s.rpcScheduler.start()
t.Cleanup(func() { s.rpcScheduler.stop(time.Second) })
@ -441,7 +441,7 @@ func TestHandleEncryptedRequiredSessionBarrierFailureIsAtomic(t *testing.T) {
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 := New(Options{legacyRPC: handler, WriteTimeout: time.Second, RPCTimeout: 5 * time.Second})
s.rpcScheduler.start()
t.Cleanup(func() { s.rpcScheduler.stop(time.Second) })
@ -549,7 +549,7 @@ func TestCrossConnectionInflightRPCHasOneBusinessOwnerAndReplaysResult(t *testin
func TestCrossConnectionInflightAbortRetriesOnlyAfterOldOwnerStops(t *testing.T) {
handler := &cancelThenRetryRPC{started: make(chan struct{})}
s := New(Options{RPC: handler, WriteTimeout: time.Second, RPCTimeout: 5 * time.Second})
s := New(Options{legacyRPC: handler, WriteTimeout: time.Second, RPCTimeout: 5 * time.Second})
s.rpcScheduler.start()
t.Cleanup(func() { s.rpcScheduler.stop(time.Second) })

16
internal/mtprotoedge/session_boundary_test.go
View file

@ -7,10 +7,10 @@ import (
"testing"
"time"
"github.com/gotd/td/bin"
"github.com/gotd/td/mt"
"github.com/gotd/td/proto"
"github.com/gotd/td/tg"
"github.com/iamxvbaba/td/bin"
"github.com/iamxvbaba/td/mt"
"github.com/iamxvbaba/td/proto"
"github.com/iamxvbaba/td/tg"
)
// TestFirstContainerBoundaryKeepsAndroidRequestMap models DrKLO's
@ -120,7 +120,7 @@ func TestContainerRPCAdmissionFailureIsAtomic(t *testing.T) {
handler := &admissionCountingRPC{}
addr, pub, _ := startTestServer(t, Options{
DC: dc,
RPC: handler,
legacyRPC: handler,
RPCMaxInflight: 1,
RPCQueueSize: 1,
RPCGlobalWorkers: 1,
@ -176,7 +176,7 @@ func TestContainerRPCAdmissionFailureIsAtomic(t *testing.T) {
func TestGZIPWrappedRPCUsesLogicalEnvelopeBoundary(t *testing.T) {
const dc = 2
handler := &admissionCountingRPC{}
addr, pub, _ := startTestServer(t, Options{DC: dc, RPC: handler})
addr, pub, _ := startTestServer(t, Options{DC: dc, legacyRPC: handler})
conn, auth, cipher := dialHandshake(t, addr, dc, pub)
requestBody := mustEncodeTL(t, &tg.HelpGetConfigRequest{})
@ -251,7 +251,7 @@ func TestEmptyContainerUsesOuterBoundary(t *testing.T) {
func TestDuplicateContainerAcksWithoutBusinessReexecution(t *testing.T) {
const dc = 2
handler := &admissionCountingRPC{}
addr, pub, _ := startTestServer(t, Options{DC: dc, RPC: handler})
addr, pub, _ := startTestServer(t, Options{DC: dc, legacyRPC: handler})
conn, auth, cipher := dialHandshake(t, addr, dc, pub)
ids := proto.NewMessageIDGen(time.Now)
@ -322,7 +322,7 @@ func TestAndroidStartupBurstKeepsAllLargeRPCResultsAddressable(t *testing.T) {
requests = 30
)
handler := &largeStartupBurstRPC{body: strings.Repeat("x", 192<<10)}
addr, pub, _ := startTestServer(t, Options{DC: dc, RPC: handler})
addr, pub, _ := startTestServer(t, Options{DC: dc, legacyRPC: handler})
conn, auth, cipher := dialHandshake(t, addr, dc, pub)
ids := proto.NewMessageIDGen(time.Now)

624
internal/mtprotoedge/session_manager.go
View file

@ -12,8 +12,8 @@ import (
"go.uber.org/zap"
"go.uber.org/zap/zapcore"
"github.com/gotd/td/bin"
"github.com/gotd/td/proto"
"github.com/iamxvbaba/td/proto"
"github.com/iamxvbaba/td/tg"
)
// ErrSessionNotFound 表示目标 session 当前无活跃连接。
@ -71,15 +71,18 @@ const maxForceCloseParallelism = 64
type queuedPush struct {
t proto.MessageType
encoded *encodedOutboundMessage
updates *layerUpdatesFanout
reservation *pendingPushReservation
at time.Time
}
type pendingPushReservation struct {
budget *outboundTrackedBudget
bytes int
bytes atomic.Int64
refs atomic.Int32
mu sync.Mutex
profiles map[tg.LayerProfile]struct{}
}
func (r *pendingPushReservation) retain() {
@ -100,10 +103,33 @@ func (r *pendingPushReservation) release() {
panic("mtprotoedge: pending push reservation released more than retained")
}
if refs == 0 {
r.budget.release(r.bytes)
r.budget.release(int(r.bytes.Load()))
}
}
// reservePrepared accounts the profile-specific immutable body retained by the
// semantic pending fanout. Multiple queued sessions sharing this reservation
// and profile share both the bytes and this one budget charge.
func (r *pendingPushReservation) reservePrepared(profile tg.LayerProfile, bytes int) bool {
if r == nil || bytes < 0 {
return false
}
r.mu.Lock()
defer r.mu.Unlock()
if _, ok := r.profiles[profile]; ok {
return true
}
if !r.budget.reserve(bytes) {
return false
}
if r.profiles == nil {
r.profiles = make(map[tg.LayerProfile]struct{})
}
r.profiles[profile] = struct{}{}
r.bytes.Add(int64(bytes))
return true
}
type sessionKey struct {
authKeyID [8]byte
sessionID int64
@ -114,6 +140,19 @@ type SessionLifecycleObserver interface {
SessionOffline(rawAuthKeyID [8]byte, sessionID, userID int64, lastForUser bool)
}
// SessionDestructionObserver is an optional explicit control-plane lifecycle.
// It is separate from SessionOffline because a physical disconnect must retain
// logical-session replay metadata, while destroy_session must invalidate it.
type SessionDestructionObserver interface {
SessionDestroyed(rawAuthKeyID [8]byte, sessionID int64)
}
func notifySessionDestroyed(observer SessionLifecycleObserver, authKeyID [8]byte, sessionID int64) {
if destroyed, ok := observer.(SessionDestructionObserver); ok {
destroyed.SessionDestroyed(authKeyID, sessionID)
}
}
// SessionManager 是活跃连接注册表，支持按 session / auth-key / user 查找并主动 push。
//
// 它只管理进程内运行态，持有可发送的活跃连接；协议可恢复事实由 auth key、客户端重连
@ -171,6 +210,294 @@ func (m *SessionManager) SetLifecycleObserver(observer SessionLifecycleObserver)
m.mu.Unlock()
}
// SeedInheritedLayerForRawAuthKey supplies an auth-key-wide default to every
// currently unknown active/provisional connection for rawAuthKeyID. Existing
// inherited or explicit state is left untouched; only ordered invokeWithLayer
// admission may correct a selected profile. The return value is the number of
// connections which transitioned from unknown to inherited.
func (m *SessionManager) SeedInheritedLayerForRawAuthKey(rawAuthKeyID [8]byte, layer int) int {
return m.applyInheritedLayerForRawAuthKey(rawAuthKeyID, layer, false)
}
// RefreshInheritedLayerForRawAuthKey is the auth.bindTempAuthKey identity-
// normalization path. It replaces unknown and inherited raw-temp-key shadows
// with the resolved permanent-key default, while preserving explicit evidence.
// Ordinary auth-key default publication must use SeedInheritedLayerForRawAuthKey
// so it cannot rewrite live sessions which already selected an inherited value.
func (m *SessionManager) RefreshInheritedLayerForRawAuthKey(rawAuthKeyID [8]byte, layer int) int {
return m.applyInheritedLayerForRawAuthKey(rawAuthKeyID, layer, true)
}
// ClearInheritedLayerForRawAuthKey removes a stale raw-key default after
// identity normalization obtains an authoritative unsupported/unknown
// permanent-key result. Only inherited state is cleared: explicit
// invokeWithLayer evidence belongs to the concrete logical session and remains
// authoritative until newer ordered evidence replaces it.
func (m *SessionManager) ClearInheritedLayerForRawAuthKey(rawAuthKeyID [8]byte) int {
if m == nil || rawAuthKeyID == ([8]byte{}) {
return 0
}
m.mu.RLock()
conns := make([]*Conn, 0, len(m.byAuthKey[rawAuthKeyID])+len(m.claimsByAuth[rawAuthKeyID]))
seen := make(map[*Conn]struct{}, cap(conns))
for _, group := range []map[int64]*Conn{m.byAuthKey[rawAuthKeyID], m.claimsByAuth[rawAuthKeyID]} {
for _, c := range group {
if c == nil {
continue
}
if _, duplicate := seen[c]; duplicate {
continue
}
seen[c] = struct{}{}
conns = append(conns, c)
}
}
m.mu.RUnlock()
cleared := 0
for _, c := range conns {
if c.isRetired() {
continue
}
if changed, err := c.clearInheritedLayerProfileState(); err == nil && changed {
cleared++
}
}
return cleared
}
// SeedInheritedLayerForBusinessAuthKey supplies a canonical permanent-key
// default to every live raw physical key already normalized to that business
// identity. This covers multiple temporary/PFS keys for one authorization;
// explicit and previously-selected inherited session profiles remain stable.
func (m *SessionManager) SeedInheritedLayerForBusinessAuthKey(businessAuthKeyID [8]byte, layer int) int {
if m == nil || businessAuthKeyID == ([8]byte{}) {
return 0
}
profile, ok := tg.ResolveLayerProfile(layer)
if !ok {
return 0
}
m.mu.RLock()
group := m.byBusinessAuthKey[businessAuthKeyID]
conns := make([]*Conn, 0, len(group))
seen := make(map[*Conn]struct{}, len(group))
for _, c := range group {
if c == nil {
continue
}
if _, duplicate := seen[c]; duplicate {
continue
}
seen[c] = struct{}{}
conns = append(conns, c)
}
m.mu.RUnlock()
seeded := 0
for _, c := range conns {
if c.isRetired() {
continue
}
if current, resolved := c.BusinessAuthKeyID(); !resolved || current != businessAuthKeyID {
continue
}
if changed, err := c.setLayerProfile(profile, LayerProfileInherited, false); err == nil && changed {
seeded++
}
}
return seeded
}
func (m *SessionManager) applyInheritedLayerForRawAuthKey(rawAuthKeyID [8]byte, layer int, refresh bool) int {
if m == nil {
return 0
}
profile, ok := tg.ResolveLayerProfile(layer)
if !ok {
return 0
}
m.mu.RLock()
conns := make([]*Conn, 0, len(m.byAuthKey[rawAuthKeyID])+len(m.claimsByAuth[rawAuthKeyID]))
seen := make(map[*Conn]struct{}, cap(conns))
for _, group := range []map[int64]*Conn{m.byAuthKey[rawAuthKeyID], m.claimsByAuth[rawAuthKeyID]} {
for _, c := range group {
if c == nil {
continue
}
if _, duplicate := seen[c]; duplicate {
continue
}
seen[c] = struct{}{}
conns = append(conns, c)
}
}
m.mu.RUnlock()
seeded := 0
for _, c := range conns {
if c.isRetired() {
continue
}
var (
changed bool
err error
)
if refresh {
changed, err = c.refreshInheritedLayerProfile(profile)
} else {
changed, err = c.setLayerProfile(profile, LayerProfileInherited, false)
}
if err == nil && changed {
seeded++
}
}
return seeded
}
// ApplyOrderedLayerProfileForSession converges every physical generation
// currently active or claiming the same logical MTProto session. Per-Conn
// msg_id watermarks make broadcasts commutative: even if profile 300 reaches a
// Conn before a delayed profile 200 broadcast, 200 is inert and final state is
// the exact registry's maximum accepted evidence.
func (m *SessionManager) ApplyOrderedLayerProfileForSession(
primary *Conn,
rawAuthKeyID [8]byte,
sessionID int64,
profile tg.LayerProfile,
msgID int64,
) (int, error) {
if err := validateLayerProfile(profile); err != nil {
return 0, err
}
return m.ApplyOrderedRawLayerForSession(primary, rawAuthKeyID, sessionID, int(profile), msgID)
}
// ApplyOrderedRawLayerForSession also carries future Layers unknown to this
// binary. Their raw watermark converges across physical generations while each
// Conn remains codec-unknown until a newer supported selector is admitted.
func (m *SessionManager) ApplyOrderedRawLayerForSession(
primary *Conn,
rawAuthKeyID [8]byte,
sessionID int64,
layer int,
msgID int64,
) (int, error) {
if msgID <= 0 {
return 0, fmt.Errorf("invalid ordered session layer msg_id %d", msgID)
}
if layer <= 0 {
return 0, fmt.Errorf("invalid ordered session layer %d", layer)
}
conns := make([]*Conn, 0, 3)
seen := make(map[*Conn]struct{}, 3)
if primary != nil {
seen[primary] = struct{}{}
conns = append(conns, primary)
}
if m != nil {
key := sessionKey{authKeyID: rawAuthKeyID, sessionID: sessionID}
m.mu.RLock()
for _, c := range []*Conn{m.bySession[key], m.claims[key]} {
if c == nil {
continue
}
if _, duplicate := seen[c]; duplicate {
continue
}
seen[c] = struct{}{}
conns = append(conns, c)
}
m.mu.RUnlock()
}
applied := 0
for _, c := range conns {
if c == nil || c.isRetired() {
continue
}
changed, err := c.freezeRawLayerProfileAt(layer, msgID)
if err != nil {
return applied, err
}
if changed {
applied++
}
}
return applied, nil
}
// ExplicitLayerEvidenceForAuthKey exposes live exact-session truth to
// auth.bindTempAuthKey. Router's bounded exact registry may expire while a Conn
// remains active; bind must not replace that explicit profile with a permanent
// key's inherited default merely because the cache TTL elapsed.
func (m *SessionManager) ExplicitLayerEvidenceForAuthKey(rawAuthKeyID [8]byte, sessionID int64) (layer int, msgID int64, ok bool) {
if m == nil || rawAuthKeyID == ([8]byte{}) || sessionID == 0 {
return 0, 0, false
}
key := sessionKey{authKeyID: rawAuthKeyID, sessionID: sessionID}
m.mu.RLock()
conns := []*Conn{m.bySession[key], m.claims[key]}
m.mu.RUnlock()
seen := make(map[*Conn]struct{}, len(conns))
for _, c := range conns {
if c == nil || c.isRetired() {
continue
}
if _, duplicate := seen[c]; duplicate {
continue
}
seen[c] = struct{}{}
state, evidenceMsgID := c.layerProfileEvidenceState()
if c.isRetired() || state.Origin != LayerProfileExplicit {
continue
}
profile, supported := tg.ResolveLayerProfile(int(state.Profile))
if !supported || profile != state.Profile {
continue
}
if !ok || evidenceMsgID > msgID {
layer, msgID, ok = int(profile), evidenceMsgID, true
continue
}
if evidenceMsgID == msgID && layer != int(profile) {
// This state contradicts the per-session msg_id ordering invariant;
// do not let bind choose either physical generation arbitrarily.
return 0, 0, false
}
}
return layer, msgID, ok
}
// SetClientLayerForAuthKey implements rpc.ClientLayerBinder without weakening
// ordered evidence. It is a legacy/readiness safety net: only an unknown exact
// Conn receives the value as inherited state. Explicit or already-selected
// inherited profiles are owned by the edge's msg_id-ordered path.
func (m *SessionManager) SetClientLayerForAuthKey(rawAuthKeyID [8]byte, sessionID int64, layer int) {
if m == nil {
return
}
profile, ok := tg.ResolveLayerProfile(layer)
if !ok {
return
}
key := sessionKey{authKeyID: rawAuthKeyID, sessionID: sessionID}
m.mu.RLock()
conns := []*Conn{m.bySession[key], m.claims[key]}
m.mu.RUnlock()
seen := make(map[*Conn]struct{}, len(conns))
for _, c := range conns {
if c == nil || c.isRetired() {
continue
}
if _, duplicate := seen[c]; duplicate {
continue
}
seen[c] = struct{}{}
_, _ = c.setLayerProfile(profile, LayerProfileInherited, false)
}
}
// 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
@ -344,6 +671,7 @@ func (m *SessionManager) Unregister(c *Conn) {
// DestroySessionForAuthKey 精确移除某个 raw auth_key_id 下的 session。
func (m *SessionManager) DestroySessionForAuthKey(authKeyID [8]byte, sessionID int64) bool {
m.mu.Lock()
observer := m.lifecycle
key := sessionKey{authKeyID: authKeyID, sessionID: sessionID}
c, ok := m.bySession[key]
if !ok {
@ -356,15 +684,16 @@ func (m *SessionManager) DestroySessionForAuthKey(authKeyID [8]byte, sessionID i
zap.Int64("session_id", sessionID),
)
}
notifySessionDestroyed(observer, authKeyID, sessionID)
return true
}
m.deletePendingLocked(key)
m.mu.Unlock()
notifySessionDestroyed(observer, authKeyID, sessionID)
return false
}
offlineUser := m.retireConnLocked(c, true)
lastForUser := offlineUser != 0 && len(m.byUser[offlineUser]) == 0
observer := m.lifecycle
m.log.Debug("Session destroyed",
zap.String("auth_key_id", sessionKeyLog(authKeyID)),
zap.Int64("session_id", sessionID),
@ -380,6 +709,7 @@ func (m *SessionManager) DestroySessionForAuthKey(authKeyID [8]byte, sessionID i
if observer != nil && offlineUser != 0 {
observer.SessionOffline(authKeyID, sessionID, offlineUser, lastForUser)
}
notifySessionDestroyed(observer, authKeyID, sessionID)
return true
}
@ -762,6 +1092,18 @@ func (m *SessionManager) setReceivesUpdatesLocked(c *Conn, key sessionKey, recei
delete(m.flushing, key)
return 0, false
}
if _, ok := c.LayerProfile(); !ok {
// A successful wire-invariant bootstrap RPC is not evidence that this
// physical session can decode proactive updates. Keep durable updates
// pending until generated exact admission freezes a real profile; do not
// start a flush which would fail layer binding and retire a healthy socket.
c.receivesUpdates.Store(false)
m.clearChannelInterestsLocked(key)
m.clearChannelMembershipsLocked(c, key)
c.membershipsSynced.Store(false)
delete(m.flushing, key)
return 0, false
}
if c.receivesUpdates.Load() || m.flushing[key] {
// 已就绪，或已有排空协程在跑（完成时会自行取走新增暂存并置位）。
return 0, false
@ -819,12 +1161,38 @@ func (m *SessionManager) runFlush(c *Conn, key sessionKey, owner int64, attempt
// Pending entries are durable account updates. Shared body-budget pressure is not
// evidence that this socket is corrupt, so use the non-terminal enqueue path; after
// bounded retries, getDifference is the authoritative recovery path.
err := c.SendBestEffortEncoded(ctx, item.t, item.encoded, 5*time.Second)
encoded, err := item.updates.prepareForConn(ctx, c)
if err == nil {
if encoded == nil || encoded.layer == nil {
err = errors.New("pending exact updates lost layer binding")
} else if !item.reservation.reservePrepared(encoded.layer.profile, len(encoded.body)) {
item.updates.discardPrepared(encoded.layer.profile, encoded)
err = ErrOutboundTrackedBudget
}
}
if err == nil {
err = c.SendBestEffortEncoded(ctx, item.t, encoded, 5*time.Second)
}
cancel()
if err == nil {
item.release()
continue
}
if isOutboundStaleLayerEpoch(err) {
// This durable online accelerator was prepared before a client layer
// correction. Drop only the stale item; difference remains authoritative.
item.release()
continue
}
if isOutboundLayerProfileError(err) {
// updates-ready without an exact profile, or a mismatched final
// body, violates the physical-connection layer invariant. Never
// guess canonical bytes; retire this writer and let durable
// difference recover after a correctly negotiated reconnect.
c.dropSlowConsumer()
releaseQueuedPushes(batch[i:])
return
}
m.mu.Lock()
if cur, ok := m.bySession[key]; !ok || cur != c || !m.flushing[key] || c.userID.Load() != owner {
// 连接换代/取消/易主：剩余 batch 不属于当前连接当前账号，丢弃。
@ -884,23 +1252,11 @@ func (m *SessionManager) ReceivesUpdatesForAuthKey(authKeyID [8]byte, sessionID
m.mu.RLock()
c, ok := m.bySession[sessionKey{authKeyID: authKeyID, sessionID: sessionID}]
m.mu.RUnlock()
return ok && c.receivesUpdates.Load() && c.membershipsSynced.Load()
}
// SetClientLayerForAuthKey 把协商的 TL layer 即时写到指定连接。由 rpc 层在
// invokeWithLayer 观测到新 layer 时（Dispatch 入口，早于鉴权门与 updates 就绪
// 置位）调用，使同一条请求触发的 pending flush 与并发 push 立即按正确 layer
// 降级，不必等连接层在 Dispatch 返回后的兜底刷新。
func (m *SessionManager) SetClientLayerForAuthKey(authKeyID [8]byte, sessionID int64, layer int) {
if m == nil || layer <= 0 {
return
}
m.mu.RLock()
c, ok := m.bySession[sessionKey{authKeyID: authKeyID, sessionID: sessionID}]
m.mu.RUnlock()
if ok {
c.SetClientLayer(layer)
if !ok {
return false
}
_, hasProfile := c.LayerProfile()
return hasProfile && c.receivesUpdates.Load() && c.membershipsSynced.Load()
}
// SetReceivesUpdatesForAuthKey 标记指定 raw auth_key_id + session_id 是否接收主动 updates。
@ -921,7 +1277,7 @@ func (m *SessionManager) SetReceivesUpdatesForAuthKey(authKeyID [8]byte, session
}
// PushToSessionForAuthKey 向指定 raw auth_key_id + session_id 推送一条消息。
func (m *SessionManager) PushToSessionForAuthKey(ctx context.Context, authKeyID [8]byte, sessionID int64, t proto.MessageType, msg bin.Encoder) error {
func (m *SessionManager) PushToSessionForAuthKey(ctx context.Context, authKeyID [8]byte, sessionID int64, t proto.MessageType, msg tg.UpdatesClass) error {
m.mu.RLock()
key := sessionKey{authKeyID: authKeyID, sessionID: sessionID}
c, ok := m.bySession[key]
@ -932,13 +1288,22 @@ func (m *SessionManager) PushToSessionForAuthKey(ctx context.Context, authKeyID
ready := c.receivesUpdates.Load()
m.mu.RUnlock()
if ready {
return c.Send(ctx, t, msg)
updates, err := newLayerUpdatesFanoutContext(ctx, msg)
if err != nil {
return err
}
encoded, err := updates.prepareForConn(ctx, c)
if err != nil {
return err
}
return c.SendEncoded(ctx, t, encoded)
}
return m.queueOrSendPrepared(ctx, key, t, msg)
}
func (m *SessionManager) queueOrSendPrepared(ctx context.Context, key sessionKey, t proto.MessageType, msg bin.Encoder) error {
encoded, reservation, err := m.preparePendingPush(ctx, msg)
func (m *SessionManager) queueOrSendPrepared(ctx context.Context, key sessionKey, t proto.MessageType, msg tg.UpdatesClass) error {
getUpdates := onceLayerUpdatesFanout(ctx, msg)
updates, reservation, err := m.preparePendingPush(getUpdates)
if err != nil {
return err
}
@ -951,11 +1316,15 @@ func (m *SessionManager) queueOrSendPrepared(ctx context.Context, key sessionKey
return ErrSessionNotFound
}
if !c.receivesUpdates.Load() {
_ = m.queuePreparedLocked(key, t, encoded, reservation)
_ = m.queuePreparedLocked(key, t, updates, reservation)
m.mu.Unlock()
return nil
}
m.mu.Unlock()
encoded, err := updates.prepareForConn(ctx, c)
if err != nil {
return err
}
return c.SendEncoded(ctx, t, encoded)
}
@ -964,7 +1333,7 @@ func (m *SessionManager) queueOrSendPrepared(ctx context.Context, key sessionKey
// 它不等待该 session 进入 updates-ready，也不写 pending 队列。仅用于登录前的握手信号
// （例如 updateLoginToken）：这类消息本身就是让客户端继续完成登录的触发器，若走普通
// durable update 队列会卡在客户端尚未调用 updates.getState 的阶段。
func (m *SessionManager) PushToSessionForAuthKeyImmediate(ctx context.Context, authKeyID [8]byte, sessionID int64, t proto.MessageType, msg bin.Encoder) error {
func (m *SessionManager) PushToSessionForAuthKeyImmediate(ctx context.Context, authKeyID [8]byte, sessionID int64, t proto.MessageType, msg tg.UpdatesClass) error {
m.mu.RLock()
key := sessionKey{authKeyID: authKeyID, sessionID: sessionID}
c, ok := m.bySession[key]
@ -972,11 +1341,19 @@ func (m *SessionManager) PushToSessionForAuthKeyImmediate(ctx context.Context, a
if !ok {
return ErrSessionNotFound
}
return c.SendBestEffort(ctx, t, msg, 2*time.Second)
updates, err := newLayerUpdatesFanoutContext(ctx, msg)
if err != nil {
return err
}
encoded, err := updates.prepareForConn(ctx, c)
if err != nil {
return err
}
return c.SendBestEffortEncoded(ctx, t, encoded, 2*time.Second)
}
// PushToUserExceptAuthKeySession 向某 user 所有活跃连接推送，跳过指定 raw auth_key + session。
func (m *SessionManager) PushToUserExceptAuthKeySession(ctx context.Context, userID int64, excludeAuthKeyID [8]byte, excludeSessionID int64, t proto.MessageType, msg bin.Encoder) (int, error) {
func (m *SessionManager) PushToUserExceptAuthKeySession(ctx context.Context, userID int64, excludeAuthKeyID [8]byte, excludeSessionID int64, t proto.MessageType, msg tg.UpdatesClass) (int, error) {
return m.pushToUser(ctx, userID, &excludeAuthKeyID, excludeSessionID, t, msg)
}
@ -985,19 +1362,25 @@ func (m *SessionManager) PushToUserExceptAuthKeySession(ctx context.Context, use
// businessAuthKeyCandidatesLocked，兼容 temp-key/PFS 连接），不是 byAuthKey（raw 索引会
// 漏 temp-key 设备）。未就绪连接跳过、不进 pending——密聊消息 durable 在 qts 队列，
// 离线设备靠 getDifference 补回（在线推送只是加速器）。c.userID 复查防跨账号泄露。
func (m *SessionManager) PushToUserAuthKey(ctx context.Context, userID int64, businessAuthKeyID [8]byte, t proto.MessageType, msg bin.Encoder) (int, error) {
func (m *SessionManager) PushToUserAuthKey(ctx context.Context, userID int64, businessAuthKeyID [8]byte, t proto.MessageType, msg tg.UpdatesClass) (int, error) {
// Secret-chat qts is the durable source of truth, so online delivery is an accelerator just
// like account pts fan-out. Do not synchronously wait for every PFS/raw connection's socket.
return m.pushToBusinessAuthKeyBestEffort(ctx, userID, businessAuthKeyID, t, msg, 2*time.Second)
}
// PushToUserAuthKeyTransient 是 PushToUserAuthKey 的 transient（typing）best-effort 版本。
func (m *SessionManager) PushToUserAuthKeyTransient(ctx context.Context, userID int64, businessAuthKeyID [8]byte, t proto.MessageType, msg bin.Encoder, timeout time.Duration) (int, error) {
func (m *SessionManager) PushToUserAuthKeyTransient(ctx context.Context, userID int64, businessAuthKeyID [8]byte, t proto.MessageType, msg tg.UpdatesClass, timeout time.Duration) (int, error) {
return m.pushToBusinessAuthKeyBestEffort(ctx, userID, businessAuthKeyID, t, msg, timeout)
}
func (m *SessionManager) pushToBusinessAuthKeyBestEffort(ctx context.Context, userID int64, businessAuthKeyID [8]byte, t proto.MessageType, msg bin.Encoder, timeout time.Duration) (int, error) {
getEncoded := onceEncodedOutbound(ctx, msg)
func (m *SessionManager) pushToBusinessAuthKeyBestEffort(ctx context.Context, userID int64, businessAuthKeyID [8]byte, t proto.MessageType, msg tg.UpdatesClass, timeout time.Duration) (int, error) {
if ctx != nil && ctx.Err() != nil {
return 0, ctx.Err()
}
sendCtx := context.Background()
if ctx != nil {
sendCtx = context.WithoutCancel(ctx)
}
var deadline time.Time
if timeout > 0 {
deadline = time.Now().Add(timeout)
@ -1007,11 +1390,21 @@ func (m *SessionManager) pushToBusinessAuthKeyBestEffort(ctx context.Context, us
deadline = ctxDeadline
}
}
if !deadline.IsZero() {
var cancel context.CancelFunc
sendCtx, cancel = context.WithDeadline(sendCtx, deadline)
defer cancel()
}
getUpdates := onceLayerUpdatesFanout(sendCtx, msg)
return m.pushToBusinessAuthKey(ctx, userID, businessAuthKeyID, func(c *Conn) error {
if c.outbound == nil || c.outboundControl == nil {
return ErrConnClosed
}
encoded, err := getEncoded()
updates, err := getUpdates()
if err != nil {
return err
}
encoded, err := updates.prepareForConn(sendCtx, c)
if err != nil {
return err
}
@ -1022,7 +1415,7 @@ func (m *SessionManager) pushToBusinessAuthKeyBestEffort(ctx context.Context, us
remaining = 0
}
}
return c.SendBestEffortEncoded(ctx, t, encoded, remaining)
return c.SendBestEffortEncoded(sendCtx, t, encoded, remaining)
})
}
@ -1049,6 +1442,15 @@ func (m *SessionManager) pushToBusinessAuthKey(ctx context.Context, userID int64
continue
}
if err := send(c); err != nil {
if isOutboundStaleLayerEpoch(err) {
// A concurrent correction invalidated this prepared push, not the
// connection. Durable qts/difference remains the source of truth.
continue
}
if isOutboundLayerProfileError(err) {
c.dropSlowConsumer()
continue
}
if errors.Is(err, ErrOutboundTrackedBudget) {
// Shared process pressure is not evidence that this particular socket is
// slow. Skip this online accelerator; durable qts/difference is the truth.
@ -1071,13 +1473,17 @@ func (m *SessionManager) pushToBusinessAuthKey(ctx context.Context, userID int64
return sent, firstErr
}
func (m *SessionManager) pushToUser(ctx context.Context, userID int64, excludeAuthKeyID *[8]byte, excludeSessionID int64, t proto.MessageType, msg bin.Encoder) (int, error) {
getEncoded := onceEncodedOutbound(ctx, msg)
return m.pushToUserWithSender(ctx, userID, excludeAuthKeyID, excludeSessionID, t, msg, true, func(c *Conn) error {
func (m *SessionManager) pushToUser(ctx context.Context, userID int64, excludeAuthKeyID *[8]byte, excludeSessionID int64, t proto.MessageType, msg tg.UpdatesClass) (int, error) {
getUpdates := onceLayerUpdatesFanout(ctx, msg)
return m.pushToUserWithSender(ctx, userID, excludeAuthKeyID, excludeSessionID, t, getUpdates, true, func(c *Conn) error {
if c.outbound == nil || c.outboundControl == nil {
return ErrConnClosed
}
encoded, err := getEncoded()
updates, err := getUpdates()
if err != nil {
return err
}
encoded, err := updates.prepareForConn(ctx, c)
if err != nil {
return err
}
@ -1090,13 +1496,17 @@ func (m *SessionManager) pushToUser(ctx context.Context, userID int64, excludeAu
// 跳过该连接、不进 pending——transient 数据 getDifference 无法补，就绪后由 getState 快照 /
// 下一次状态变化重建，囤积过期 transient 既无意义又会被 pending 的老化/溢出/重试耗尽误当
// 「durable 兜底」丢弃。走 best-effort 发送，不阻塞调用方。
func (m *SessionManager) PushToUserTransientExceptAuthKeySession(ctx context.Context, userID int64, excludeAuthKeyID [8]byte, excludeSessionID int64, t proto.MessageType, msg bin.Encoder, timeout time.Duration) (int, error) {
getEncoded := onceEncodedOutbound(ctx, msg)
return m.pushToUserWithSender(ctx, userID, &excludeAuthKeyID, excludeSessionID, t, msg, false, func(c *Conn) error {
func (m *SessionManager) PushToUserTransientExceptAuthKeySession(ctx context.Context, userID int64, excludeAuthKeyID [8]byte, excludeSessionID int64, t proto.MessageType, msg tg.UpdatesClass, timeout time.Duration) (int, error) {
getUpdates := onceLayerUpdatesFanout(ctx, msg)
return m.pushToUserWithSender(ctx, userID, &excludeAuthKeyID, excludeSessionID, t, getUpdates, false, func(c *Conn) error {
if c.outbound == nil || c.outboundControl == nil {
return ErrConnClosed
}
encoded, err := getEncoded()
updates, err := getUpdates()
if err != nil {
return err
}
encoded, err := updates.prepareForConn(ctx, c)
if err != nil {
return err
}
@ -1104,12 +1514,18 @@ func (m *SessionManager) PushToUserTransientExceptAuthKeySession(ctx context.Con
})
}
func (m *SessionManager) PushToUserExceptAuthKeySessionBestEffort(ctx context.Context, userID int64, excludeAuthKeyID [8]byte, excludeSessionID int64, t proto.MessageType, msg bin.Encoder, timeout time.Duration) (int, error) {
func (m *SessionManager) PushToUserExceptAuthKeySessionBestEffort(ctx context.Context, userID int64, excludeAuthKeyID [8]byte, excludeSessionID int64, t proto.MessageType, msg tg.UpdatesClass, timeout time.Duration) (int, error) {
return m.pushToUserBestEffort(ctx, userID, &excludeAuthKeyID, excludeSessionID, t, msg, timeout)
}
func (m *SessionManager) pushToUserBestEffort(ctx context.Context, userID int64, excludeAuthKeyID *[8]byte, excludeSessionID int64, t proto.MessageType, msg bin.Encoder, timeout time.Duration) (int, error) {
getEncoded := onceEncodedOutbound(ctx, msg)
func (m *SessionManager) pushToUserBestEffort(ctx context.Context, userID int64, excludeAuthKeyID *[8]byte, excludeSessionID int64, t proto.MessageType, msg tg.UpdatesClass, timeout time.Duration) (int, error) {
if ctx != nil && ctx.Err() != nil {
return 0, ctx.Err()
}
sendCtx := context.Background()
if ctx != nil {
sendCtx = context.WithoutCancel(ctx)
}
// timeout 是整次 fan-out 的等待预算，不是每个 session 各自一份。健康连接始终先走
// SendBestEffortEncoded 的非阻塞快路径；预算耗尽后 remaining=0，仍会尝试快路径，
// 但不会再为后续慢连接串行等待。
@ -1122,11 +1538,21 @@ func (m *SessionManager) pushToUserBestEffort(ctx context.Context, userID int64,
deadline = ctxDeadline
}
}
return m.pushToUserWithSender(ctx, userID, excludeAuthKeyID, excludeSessionID, t, msg, true, func(c *Conn) error {
if !deadline.IsZero() {
var cancel context.CancelFunc
sendCtx, cancel = context.WithDeadline(sendCtx, deadline)
defer cancel()
}
getUpdates := onceLayerUpdatesFanout(sendCtx, msg)
return m.pushToUserWithSender(ctx, userID, excludeAuthKeyID, excludeSessionID, t, getUpdates, true, func(c *Conn) error {
if c.outbound == nil || c.outboundControl == nil {
return ErrConnClosed
}
encoded, err := getEncoded()
updates, err := getUpdates()
if err != nil {
return err
}
encoded, err := updates.prepareForConn(sendCtx, c)
if err != nil {
return err
}
@ -1137,24 +1563,35 @@ func (m *SessionManager) pushToUserBestEffort(ctx context.Context, userID int64,
remaining = 0
}
}
return c.SendBestEffortEncoded(ctx, t, encoded, remaining)
return c.SendBestEffortEncoded(sendCtx, t, encoded, remaining)
})
}
func onceEncodedOutbound(ctx context.Context, msg bin.Encoder) func() (*encodedOutboundMessage, error) {
func newLayerUpdatesFanoutContext(ctx context.Context, msg tg.UpdatesClass) (*layerUpdatesFanout, error) {
var updates *layerUpdatesFanout
err := withOutboundEncodeSlot(ctx, nil, func() error {
var err error
updates, err = newLayerUpdatesFanout(msg)
return err
})
return updates, err
}
func onceLayerUpdatesFanout(ctx context.Context, msg tg.UpdatesClass) func() (*layerUpdatesFanout, error) {
var (
encoded *encodedOutboundMessage
once sync.Once
updates *layerUpdatesFanout
err error
)
return func() (*encodedOutboundMessage, error) {
if encoded == nil && err == nil {
encoded, err = encodeOutboundMessageContext(ctx, msg)
}
return encoded, err
return func() (*layerUpdatesFanout, error) {
once.Do(func() {
updates, err = newLayerUpdatesFanoutContext(ctx, msg)
})
return updates, err
}
}
func (m *SessionManager) pushToUserWithSender(ctx context.Context, userID int64, excludeAuthKeyID *[8]byte, excludeSessionID int64, t proto.MessageType, msg bin.Encoder, queueWhenNotReady bool, send func(*Conn) error) (int, error) {
func (m *SessionManager) pushToUserWithSender(ctx context.Context, userID int64, excludeAuthKeyID *[8]byte, excludeSessionID int64, t proto.MessageType, getUpdates func() (*layerUpdatesFanout, error), queueWhenNotReady bool, send func(*Conn) error) (int, error) {
// push fan-out 是连接层最热路径之一：debug 日志的字段构造（含 auth_key hex 格式化）
// 在关闭 debug 时也会求值，先查级别一次、按需记日志。
debug := m.log.Core().Enabled(zapcore.DebugLevel)
@ -1191,7 +1628,7 @@ func (m *SessionManager) pushToUserWithSender(ctx context.Context, userID int64,
// TL encoding and the process-wide pending-byte reservation may be expensive or
// briefly block on the global encode gate. Do both before taking SessionManager.mu,
// then share the immutable body across every not-ready session found by the re-scan.
pendingEncoded, pendingReservation, pendingErr := m.preparePendingPush(ctx, msg)
pendingUpdates, pendingReservation, pendingErr := m.preparePendingPush(getUpdates)
// 写锁下完整重扫（读锁释放到此之间状态可能变化，以重扫结果为准）。
conns = conns[:0]
queued, dropped, excluded, skipped = 0, 0, 0, 0
@ -1207,7 +1644,7 @@ func (m *SessionManager) pushToUserWithSender(ctx context.Context, userID int64,
skipped++
continue
}
if pendingErr == nil && m.queuePreparedLocked(key, t, pendingEncoded, pendingReservation) {
if pendingErr == nil && m.queuePreparedLocked(key, t, pendingUpdates, pendingReservation) {
queued++
if debug {
m.log.Debug("Push queued (session not updates-ready)",
@ -1252,6 +1689,16 @@ func (m *SessionManager) pushToUserWithSender(ctx context.Context, userID int64,
continue
}
if err := send(c); err != nil {
if isOutboundStaleLayerEpoch(err) {
// Do not classify profile correction as slow-consumer evidence.
dropped++
continue
}
if isOutboundLayerProfileError(err) {
c.dropSlowConsumer()
dropped++
continue
}
if errors.Is(err, ErrOutboundTrackedBudget) {
// Do not turn pressure owned by other sockets into a reconnect storm on
// healthy recipients. The durable event remains recoverable by difference.
@ -1828,43 +2275,34 @@ func (m *SessionManager) takePendingLocked(key sessionKey, ready bool) []queuedP
return pending
}
// preparePendingPush encodes outside SessionManager.mu and reserves the one physical body before
// releasing the process-wide encode slot. Multiple not-ready sessions may then share this
// immutable body via reservation refs instead of encoding/copying it once per session.
func (m *SessionManager) preparePendingPush(ctx context.Context, msg bin.Encoder) (*encodedOutboundMessage, *pendingPushReservation, error) {
var (
encoded *encodedOutboundMessage
bytes int
)
err := withOutboundEncodeSlot(ctx, nil, func() error {
var err error
encoded, err = encodeOutboundMessageWithoutSlot(msg)
if err != nil {
return err
}
if encoded == nil {
return errors.New("nil encoded pending push")
}
bytes = len(encoded.body)
if bytes > maxOutboundBodyBytes {
return fmt.Errorf("%w: body=%d limit=%d", ErrOutboundMessageTooLarge, bytes, maxOutboundBodyBytes)
}
if !m.pendingBudget.reserve(bytes) {
return ErrOutboundTrackedBudget
}
return nil
})
// preparePendingPush reserves the one frozen canonical semantic snapshot. Exact
// wire bytes are deliberately not retained here: they are prepared only when a
// target physical connection has a frozen profile, then cached once per profile
// by layerUpdatesFanout.
func (m *SessionManager) preparePendingPush(getUpdates func() (*layerUpdatesFanout, error)) (*layerUpdatesFanout, *pendingPushReservation, error) {
updates, err := getUpdates()
if err != nil {
return nil, nil, err
}
reservation := &pendingPushReservation{budget: m.pendingBudget, bytes: bytes}
if updates == nil {
return nil, nil, errors.New("nil pending layer updates")
}
bytes := updates.canonicalSize()
if bytes > maxOutboundBodyBytes {
return nil, nil, fmt.Errorf("%w: body=%d limit=%d", ErrOutboundMessageTooLarge, bytes, maxOutboundBodyBytes)
}
if !m.pendingBudget.reserve(bytes) {
return nil, nil, ErrOutboundTrackedBudget
}
reservation := &pendingPushReservation{budget: m.pendingBudget}
reservation.bytes.Store(int64(bytes))
reservation.refs.Store(1) // producer ownership; queue entries retain below.
return encoded, reservation, nil
return updates, reservation, nil
}
// queuePreparedLocked 暂存一条已编码的主动推送，返回是否实际入队。
// queuePreparedLocked 暂存一条已冻结的主动推送，返回是否实际入队。
// 调用方必须在锁外保持 reservation 的 producer ref，并在全部入队完成后 release。
func (m *SessionManager) queuePreparedLocked(key sessionKey, t proto.MessageType, encoded *encodedOutboundMessage, reservation *pendingPushReservation) bool {
func (m *SessionManager) queuePreparedLocked(key sessionKey, t proto.MessageType, updates *layerUpdatesFanout, reservation *pendingPushReservation) bool {
q := m.pending[key]
// 过期保护：最早一条暂存已超过 pendingPushMaxAge（session 迟迟未 ready）时，丢整批并
// 不再囤这条，记 trace。避免「登录后从不 getState」的连接长期占用 pending 内存。
@ -1877,13 +2315,13 @@ func (m *SessionManager) queuePreparedLocked(key sessionKey, t proto.MessageType
m.deletePendingLocked(key)
return false
}
if encoded == nil || reservation == nil {
if updates == nil || reservation == nil {
return false
}
reservation.retain()
push := queuedPush{
t: t,
encoded: encoded,
updates: updates,
reservation: reservation,
at: time.Now(),
}

4
internal/mtprotoedge/session_manager_bench_test.go
View file

@ -14,8 +14,8 @@ import (
"go.uber.org/zap"
"github.com/gotd/td/proto"
"github.com/gotd/td/tg"
"github.com/iamxvbaba/td/proto"
"github.com/iamxvbaba/td/tg"
)
// 连接层 fan-out / churn 压测：聚焦 SessionManager 的锁争用，不走真实 socket / 加密。

256
internal/mtprotoedge/session_manager_test.go
View file

@ -1,6 +1,7 @@
package mtprotoedge
import (
"bytes"
"context"
"errors"
"sync"
@ -10,25 +11,26 @@ import (
"go.uber.org/zap/zaptest"
"github.com/gotd/td/bin"
"github.com/gotd/td/mt"
"github.com/gotd/td/proto"
"github.com/gotd/td/tg"
"github.com/iamxvbaba/td/bin"
"github.com/iamxvbaba/td/mt"
"github.com/iamxvbaba/td/proto"
"github.com/iamxvbaba/td/tg"
)
type countingOutboundEncoder struct {
count *int
}
func (e *countingOutboundEncoder) Encode(b *bin.Buffer) error {
*e.count++
return (&tg.UpdatesTooLong{}).Encode(b)
}
type closeCountingTransport struct {
closes int
}
type sessionDestructionRecorder struct {
destroyed []sessionKey
}
func (*sessionDestructionRecorder) SessionOffline([8]byte, int64, int64, bool) {}
func (r *sessionDestructionRecorder) SessionDestroyed(authKeyID [8]byte, sessionID int64) {
r.destroyed = append(r.destroyed, sessionKey{authKeyID: authKeyID, sessionID: sessionID})
}
type slowCloseTransport struct {
delay time.Duration
release <-chan struct{}
@ -164,6 +166,8 @@ func TestSessionManagerReplacementClosesOldPhysicalTransport(t *testing.T) {
func TestSessionManagerDestroyClosesPhysicalTransport(t *testing.T) {
sm := NewSessionManager(zaptest.NewLogger(t))
observer := &sessionDestructionRecorder{}
sm.SetLifecycleObserver(observer)
raw := [8]byte{4, 5, 6}
physical := &closeCountingTransport{}
c := &Conn{sessionID: 77, authKeyID: raw, transport: physical}
@ -178,11 +182,38 @@ func TestSessionManagerDestroyClosesPhysicalTransport(t *testing.T) {
if sm.Online() != 0 {
t.Fatalf("online after destroy = %d, want 0", sm.Online())
}
if len(observer.destroyed) != 1 || observer.destroyed[0] != (sessionKey{authKeyID: raw, sessionID: 77}) {
t.Fatalf("destroy callbacks = %+v, want exact destroyed session", observer.destroyed)
}
// An explicit destroy for an already-offline logical session must still
// invalidate retained exact-profile metadata even though the wire response
// remains destroy_session_none.
if sm.DestroySessionForAuthKey(raw, 88) {
t.Fatal("missing session reported destroyed")
}
if len(observer.destroyed) != 2 || observer.destroyed[1] != (sessionKey{authKeyID: raw, sessionID: 88}) {
t.Fatalf("offline destroy callbacks = %+v", observer.destroyed)
}
}
func TestSessionManagerBestEffortFanoutPreencodesOnce(t *testing.T) {
func TestSessionManagerUnregisterIsNotSessionDestruction(t *testing.T) {
sm := NewSessionManager(zaptest.NewLogger(t))
observer := &sessionDestructionRecorder{}
sm.SetLifecycleObserver(observer)
c := &Conn{sessionID: 91, authKeyID: [8]byte{9, 1}}
if err := sm.Register(c); err != nil {
t.Fatal(err)
}
sm.Unregister(c)
if len(observer.destroyed) != 0 {
t.Fatalf("ordinary unregister emitted destruction: %+v", observer.destroyed)
}
}
func TestSessionManagerBestEffortFanoutPreparesOncePerProfile(t *testing.T) {
sm := NewSessionManager(zaptest.NewLogger(t))
const userID = int64(100)
conns := make([]*Conn, 0, 2)
for i := 0; i < 2; i++ {
c := &Conn{
sessionID: int64(i + 1),
@ -195,16 +226,19 @@ func TestSessionManagerBestEffortFanoutPreencodesOnce(t *testing.T) {
c.userID.Store(userID)
c.userIDResolved.Store(true)
c.receivesUpdates.Store(true)
if err := c.FreezeLayerProfile(tg.LayerProfile225); err != nil {
t.Fatalf("freeze profile: %v", err)
}
sm.Register(c)
conns = append(conns, c)
}
encodes := 0
sent, err := sm.PushToUserExceptSessionBestEffort(
context.Background(),
userID,
0,
proto.MessageFromServer,
&countingOutboundEncoder{count: &encodes},
&tg.UpdatesTooLong{},
0,
)
if err != nil {
@ -213,8 +247,91 @@ func TestSessionManagerBestEffortFanoutPreencodesOnce(t *testing.T) {
if sent != 2 {
t.Fatalf("sent = %d, want 2", sent)
}
if encodes != 1 {
t.Fatalf("encoded %d times, want 1", encodes)
first := <-conns[0].outbound
second := <-conns[1].outbound
defer first.releaseReservation(conns[0].outboundTrackedBudget)
defer second.releaseReservation(conns[1].outboundTrackedBudget)
if first.encoded == nil || second.encoded == nil || !sameBacking(first.encoded.body, second.encoded.body) {
t.Fatal("same-profile fanout did not share one exact prepared body")
}
if first.encoded.layer == nil || second.encoded.layer == nil || first.encoded.layer == second.encoded.layer {
t.Fatal("same-profile fanout shared connection-specific epoch binding")
}
}
func TestSessionManagerMixedLayerFanoutUsesProfileBoundBodies(t *testing.T) {
sm := NewSessionManager(zaptest.NewLogger(t))
const userID = int64(103)
authKeyID := [8]byte{0x22, 0x70, 0x22, 0x80}
profiles := []tg.LayerProfile{tg.LayerProfile225, tg.LayerProfile227, tg.LayerProfile228}
conns := make([]*Conn, 0, len(profiles))
for _, profile := range profiles {
c := &Conn{
sessionID: int64(profile),
authKeyID: authKeyID,
outbound: make(chan outboundOp, 1),
outboundControl: make(chan outboundOp, 1),
outboundStop: make(chan struct{}),
metrics: NopMetrics{},
}
c.userID.Store(userID)
c.userIDResolved.Store(true)
c.receivesUpdates.Store(true)
if profile == tg.LayerProfile228 {
if err := c.SeedInheritedLayerProfile(tg.LayerProfile227); err != nil {
t.Fatalf("seed Alice inherited profile: %v", err)
}
}
if err := c.FreezeLayerProfile(profile); err != nil {
t.Fatalf("freeze profile %d: %v", profile, err)
}
sm.Register(c)
conns = append(conns, c)
}
value := testLayerChannelUpdatesValue(321)
sent, err := sm.PushToUserExceptSessionBestEffort(
context.Background(), userID, 0, proto.MessageFromServer, value, 0,
)
if err != nil || sent != len(conns) {
t.Fatalf("mixed fanout = sent:%d err:%v", sent, err)
}
var previous *encodedOutboundMessage
for i, c := range conns {
op := <-c.outbound
defer op.releaseReservation(c.outboundTrackedBudget)
if op.encoded == nil || op.encoded.layer == nil || op.encoded.layer.profile != profiles[i] {
t.Fatalf("connection %d binding = %#v", i, op.encoded)
}
state := c.LayerProfileState()
if op.encoded.layer.kind != outboundLayerBindingSession || op.encoded.layer.epoch != state.Epoch {
t.Fatalf("connection %d target binding = %#v, state=%#v", i, op.encoded.layer, state)
}
if previous == op.encoded {
t.Fatal("different profiles shared one final wire body")
}
previous = op.encoded
wantChannelID := testChannelWireID(profiles[i])
if !bytes.Contains(op.encoded.body, littleEndianID(wantChannelID)) {
t.Fatalf("profile %d push lacks channel constructor %#08x", profiles[i], wantChannelID)
}
if otherChannelID := testOtherChannelWireID(profiles[i]); bytes.Contains(op.encoded.body, littleEndianID(otherChannelID)) {
t.Fatalf("profile %d push leaked channel constructor %#08x", profiles[i], otherChannelID)
}
input := bin.Buffer{Buf: op.encoded.body}
decoded, decodeErr := tg.DecodeLayer(profiles[i], tg.LayerClassUpdatesType(), &input)
if decodeErr != nil || input.Len() != 0 {
t.Fatalf("decode profile %d: remaining=%d err=%v", profiles[i], input.Len(), decodeErr)
}
updates := decoded.(*tg.Updates)
channel, ok := updates.Chats[0].(*tg.Channel)
if !ok || channel.ID != 100 {
t.Fatalf("profile %d decoded channel=%#v", profiles[i], updates.Chats)
}
status := updates.Updates[0].(*tg.UpdateUserStatus).Status.(*tg.UserStatusOnline)
if status.Expires != 321 {
t.Fatalf("profile %d decoded expires=%d", profiles[i], status.Expires)
}
}
}
@ -230,24 +347,23 @@ func TestSessionManagerPendingFanoutSharesOneEncodedBodyAndBudget(t *testing.T)
keys = append(keys, connSessionKey(c))
}
encodes := 0
msg := &countingOutboundEncoder{count: &encodes}
msg := &tg.UpdatesTooLong{}
sent, err := sm.PushToUserExceptSession(context.Background(), userID, 0, proto.MessageFromServer, msg)
if err != nil {
t.Fatalf("push: %v", err)
}
if sent != 2 || encodes != 1 {
t.Fatalf("pending fanout = sent:%d encodes:%d, want 2/1", sent, encodes)
if sent != 2 {
t.Fatalf("pending fanout sent=%d, want 2", sent)
}
sm.mu.Lock()
first := sm.pending[keys[0]][0]
second := sm.pending[keys[1]][0]
if first.encoded != second.encoded || first.reservation != second.reservation {
if first.updates != second.updates || first.reservation != second.reservation {
sm.mu.Unlock()
t.Fatal("pending sessions did not share encoded body/reservation")
t.Fatal("pending sessions did not share frozen updates/reservation")
}
wantBytes := int64(len(first.encoded.body))
wantBytes := int64(first.updates.canonicalSize())
sm.deletePendingLocked(keys[0])
if got := sm.pendingBudget.snapshot(); got != wantBytes {
sm.mu.Unlock()
@ -281,6 +397,9 @@ func TestSessionManagerBestEffortFanoutUsesOneBudgetAndDropsOnlySlowConsumers(t
c.userID.Store(userID)
c.userIDResolved.Store(true)
c.receivesUpdates.Store(true)
if err := c.FreezeLayerProfile(tg.LayerProfile227); err != nil {
t.Fatalf("freeze profile: %v", err)
}
sm.Register(c)
slow = append(slow, c)
}
@ -296,6 +415,9 @@ func TestSessionManagerBestEffortFanoutUsesOneBudgetAndDropsOnlySlowConsumers(t
healthy.userID.Store(userID)
healthy.userIDResolved.Store(true)
healthy.receivesUpdates.Store(true)
if err := healthy.FreezeLayerProfile(tg.LayerProfile227); err != nil {
t.Fatalf("freeze healthy profile: %v", err)
}
sm.Register(healthy)
const budget = 40 * time.Millisecond
@ -625,6 +747,9 @@ func TestPushToUserAuthKeyUsesOneDeadlineAndDropsOnlySlowPFSConnections(t *testi
outboundStop: make(chan struct{}),
}
c.receivesUpdates.Store(true)
if err := c.FreezeLayerProfile(tg.LayerProfile227); err != nil {
t.Fatalf("freeze profile: %v", err)
}
if queueFull {
c.outbound <- outboundOp{}
}
@ -787,6 +912,9 @@ func TestPushToSessionForAuthKeyImmediateBypassesReadinessQueue(t *testing.T) {
outboundControl: make(chan outboundOp, 1),
outboundStop: make(chan struct{}),
}
if err := c.FreezeLayerProfile(tg.LayerProfile227); err != nil {
t.Fatalf("freeze profile: %v", err)
}
sm.Register(c)
msg := &tg.UpdateShort{Update: &tg.UpdateLoginToken{}, Date: 1700000000}
@ -819,6 +947,77 @@ func TestPushToSessionForAuthKeyImmediateBypassesReadinessQueue(t *testing.T) {
}
}
func TestSessionManagerWithholdsUpdatesReadinessUntilExactProfile(t *testing.T) {
sm := NewSessionManager(zaptest.NewLogger(t))
key := sessionKey{authKeyID: [8]byte{0x22, 0x02, 0x27}, sessionID: 220227}
c := &Conn{
authKeyID: key.authKeyID,
sessionID: key.sessionID,
outbound: make(chan outboundOp, 1),
outboundControl: make(chan outboundOp, 1),
outboundStop: make(chan struct{}),
metrics: NopMetrics{},
outboundTrackedBudget: newOutboundTrackedBudget(1 << 20),
}
const userID = int64(1000000227)
c.userID.Store(userID)
c.userIDResolved.Store(true)
c.membershipsSynced.Store(true) // model work performed before a defensive Set(true)
if err := sm.Register(c); err != nil {
t.Fatal(err)
}
update := &tg.UpdateShort{Update: &tg.UpdateUserStatus{UserID: userID, Status: &tg.UserStatusOnline{Expires: 1}}, Date: 1}
if err := sm.PushToSessionForAuthKey(context.Background(), key.authKeyID, key.sessionID, proto.MessageFromServer, update); err != nil {
t.Fatal(err)
}
sm.SetReceivesUpdatesForAuthKey(key.authKeyID, key.sessionID, true)
if c.receivesUpdates.Load() || c.membershipsSynced.Load() || sm.ReceivesUpdatesForAuthKey(key.authKeyID, key.sessionID) {
t.Fatal("unknown-profile connection became updates-ready")
}
if c.isRetired() {
t.Fatal("unknown-profile readiness attempt retired a healthy connection")
}
sm.mu.RLock()
pending, flushing := len(sm.pending[key]), sm.flushing[key]
sm.mu.RUnlock()
if pending != 1 || flushing {
t.Fatalf("unknown-profile readiness changed pending state = pending:%d flushing:%v", pending, flushing)
}
select {
case op := <-c.outbound:
op.releaseReservation(c.outboundTrackedBudget)
t.Fatal("unknown-profile readiness flushed a proactive update")
default:
}
if err := c.FreezeLayerProfile(tg.LayerProfile225); err != nil {
t.Fatal(err)
}
c.membershipsSynced.Store(true)
sm.SetReceivesUpdatesForAuthKey(key.authKeyID, key.sessionID, true)
var op outboundOp
select {
case op = <-c.outbound:
case <-time.After(time.Second):
t.Fatal("profiled readiness did not flush pending update")
}
if op.encoded == nil || op.encoded.layer == nil || op.encoded.layer.profile != tg.LayerProfile225 {
t.Fatalf("flushed update layer binding = %#v", op.encoded)
}
op.releaseReservation(c.outboundTrackedBudget)
deadline := time.Now().Add(time.Second)
for !c.receivesUpdates.Load() && time.Now().Before(deadline) {
time.Sleep(time.Millisecond)
}
if !c.receivesUpdates.Load() || !sm.ReceivesUpdatesForAuthKey(key.authKeyID, key.sessionID) {
t.Fatal("profiled connection did not become updates-ready after ordered flush")
}
if c.isRetired() {
t.Fatal("profiled pending flush retired connection")
}
}
func TestPendingPushBodiesUseGlobalByteBudgetAndReleaseOnDrop(t *testing.T) {
sm := NewSessionManager(zaptest.NewLogger(t))
msg := &tg.UpdateShort{Update: &tg.UpdateLoginToken{}, Date: 1700000000}
@ -863,6 +1062,9 @@ func TestPendingFlushGlobalBodyPressureDoesNotTerminateHealthyConnection(t *test
const userID = int64(606)
c.userID.Store(userID)
c.userIDResolved.Store(true)
if err := c.FreezeLayerProfile(tg.LayerProfileCanonical); err != nil {
t.Fatal(err)
}
sm.Register(c)
msg := &tg.UpdateShort{Update: &tg.UpdateLoginToken{}, Date: 1700000000}
@ -958,6 +1160,10 @@ func TestSessionManagerPush(t *testing.T) {
if got := srv.Conns().Online(); got != 2 {
t.Fatalf("online = %d, want 2", got)
}
if !srv.Conns().SetLayerProfile(auth1.SessionID, tg.LayerProfile227) ||
!srv.Conns().SetLayerProfile(auth2.SessionID, tg.LayerProfile227) {
t.Fatal("seed exact test profiles")
}
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()

4
internal/mtprotoedge/session_membership_gen_test.go
View file

@ -5,6 +5,7 @@ import (
"testing"
"time"
"github.com/iamxvbaba/td/tg"
"go.uber.org/zap/zaptest"
)
@ -66,6 +67,9 @@ func TestSetSessionChannelMembershipsDetectsConcurrentIncrementalUpdates(t *test
sm := NewSessionManager(zaptest.NewLogger(t))
raw := [8]byte{1, 2, 3}
c := &Conn{sessionID: 42, authKeyID: raw}
if err := c.FreezeLayerProfile(tg.LayerProfile227); err != nil {
t.Fatal(err)
}
sm.Register(c)
sm.BindUserForAuthKey(raw, 42, 100)
sm.SetReceivesUpdatesForAuthKey(raw, 42, true)

4
internal/mtprotoedge/session_ready_test.go
View file

@ -3,6 +3,7 @@ package mtprotoedge
import (
"testing"
"github.com/iamxvbaba/td/tg"
"go.uber.org/zap/zaptest"
)
@ -15,6 +16,9 @@ func TestReceivesUpdatesForAuthKeyRequiresMembershipSync(t *testing.T) {
sm := NewSessionManager(zaptest.NewLogger(t))
raw := [8]byte{1, 2, 3}
c := &Conn{sessionID: 42, authKeyID: raw}
if err := c.FreezeLayerProfile(tg.LayerProfile227); err != nil {
t.Fatal(err)
}
sm.Register(c)
sm.BindUserForAuthKey(raw, 42, 100)

7
internal/mtprotoedge/structural_limits_test.go
View file

@ -9,9 +9,9 @@ import (
"testing"
"time"
"github.com/gotd/td/bin"
"github.com/gotd/td/mt"
"github.com/gotd/td/proto"
"github.com/iamxvbaba/td/bin"
"github.com/iamxvbaba/td/mt"
"github.com/iamxvbaba/td/proto"
"go.uber.org/zap/zaptest"
)
@ -296,6 +296,7 @@ func TestInvalidMessageIDRejectsBeforeGZIPExpansion(t *testing.T) {
{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},
{name: "empty fractional bits", msgID: time.Now().Unix() << 32, badCode: badMsgIDInvalidBits},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {

4
internal/mtprotoedge/transient_push_test.go
View file

@ -6,8 +6,8 @@ import (
"go.uber.org/zap/zaptest"
"github.com/gotd/td/proto"
"github.com/gotd/td/tg"
"github.com/iamxvbaba/td/proto"
"github.com/iamxvbaba/td/tg"
)
// TestPushTransientSkipsNotReadySession 锁定不变量：transient 推送（typing/presence）对

8
internal/mtprotoedge/transport_compat.go
View file

@ -14,10 +14,10 @@ import (
"github.com/go-faster/errors"
"go.uber.org/multierr"
"github.com/gotd/td/bin"
tdcrypto "github.com/gotd/td/crypto"
"github.com/gotd/td/proto/codec"
"github.com/gotd/td/transport"
"github.com/iamxvbaba/td/bin"
tdcrypto "github.com/iamxvbaba/td/crypto"
"github.com/iamxvbaba/td/proto/codec"
"github.com/iamxvbaba/td/transport"
)
const maxTransportMessageSize = 1 << 24

2
internal/mtprotoedge/transport_compat_test.go
View file

@ -8,7 +8,7 @@ import (
"testing"
"time"
"github.com/gotd/td/bin"
"github.com/iamxvbaba/td/bin"
)
type messageWriteTestConn struct {

4
internal/mtprotoedge/transport_ownership.go
View file

@ -6,8 +6,8 @@ import (
"sync/atomic"
"time"
"github.com/gotd/td/bin"
"github.com/gotd/td/transport"
"github.com/iamxvbaba/td/bin"
"github.com/iamxvbaba/td/transport"
)
const (

2
internal/mtprotoedge/transport_ownership_test.go
View file

@ -9,7 +9,7 @@ import (
"testing"
"time"
"github.com/gotd/td/bin"
"github.com/iamxvbaba/td/bin"
)
type ownershipTestTransport struct {