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

Merge remote-tracking branch 'upstream/main' into merge-gramsrv-9106877

Browse source
This commit is contained in:
onysd 2026-08-03 23:29:20 +03:00
commit ac6a50c5ff
697 changed files with 100880 additions and 8052 deletions
Download patch file Download diff file Expand all files Collapse all files

2
internal/mtprotoedge/auth_key_switch_test.go
View file

@ -15,7 +15,7 @@ func TestEncryptedConnectionSwitchesAuthKeyEvenWhenSessionIDIsReused(t *testing.
_, authB, cipherB := dialHandshake(t, addr, dc, pub)
msgID := proto.NewMessageIDGen(time.Now)
sendEncrypted(t, connA, cipherA, authA, msgID.New(proto.MessageFromClient), &mt.PingRequest{PingID: 1})
sendEncryptedWithSeq(t, connA, cipherA, authA, msgID.New(proto.MessageFromClient), 1, &mt.PingRequest{PingID: 1})
for range 3 { // new_session_created + pong + msgs_ack; leave no A-key frame on the socket.
readServerMessage(t, connA, cipherA, authA.AuthKey)
}

13
internal/mtprotoedge/conn.go
View file

@ -108,6 +108,9 @@ type Conn struct {
outboundControlBudgetOnce sync.Once
outboundScratchPool *outboundScratchPool
outboundScratchOnce sync.Once
// outboundState outlives this physical Conn generation. A replacement
// physical connection for the same auth key/session reuses it.
outboundState *outboundState
// lifecycle is the sole monotonic activation/retirement state machine.
// retired never transitions back to claiming/active; one atomic state avoids
// contradictory activation and shutdown observations.
@ -130,7 +133,7 @@ type Conn struct {
rpcReady bool
rpcClosed bool
// rpcReplayRestores is a per-physical-connection ordering barrier. An exact
// cached/rewrapped init request has already executed its business handler,
// replayed/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
@ -139,6 +142,11 @@ type Conn struct {
// Rewrap aliasing never delays execution. initialized stops collecting
// candidates after the first valid init wrapper on this physical generation.
rpcRewrapInitialized atomic.Bool
// layerRPCAdmissionTraceLogged bounds production INFO diagnostics to the
// first non-unknown exact-admission rejection on this physical generation.
// Repeated malformed requests remain visible at Debug without turning an
// authenticated reconnect/session into an unbounded INFO log source.
layerRPCAdmissionTraceLogged 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)
@ -150,7 +158,8 @@ type Conn struct {
rpcRootCtx context.Context
rpcMaxInflight int
// sentContentMessages 只由 outbound actor 访问，用于生成 MTProto seq_no。
// sentContentMessages is retained only for standalone construction tests.
// Server connections allocate seq_no from logical-session outboundState.
sentContentMessages int32
// outboundRand 只由 outbound actor 访问：对 cipher 随机源的缓冲预读，
// 把每帧 padding 的 getrandom syscall 摊薄成 ~1KiB 一次。

2
internal/mtprotoedge/conn_layer_profile.go
View file

@ -137,7 +137,7 @@ func (c *Conn) layerProfileRawEvidenceState() (LayerProfileSnapshot, int, int64)
// freezeLayerProfileAt is the production explicit-evidence transition. The
// positive client msg_id is the protocol ordering authority across TCP
// reconnects and cached request replays.
// reconnects and retained request replays.
func (c *Conn) freezeLayerProfileAt(profile tlprofile.Profile, msgID int64) (bool, error) {
if c == nil {
return false, fmt.Errorf("nil connection layer profile")

2
internal/mtprotoedge/connection_intake.go
View file

@ -89,7 +89,7 @@ func connLocal(conn net.Conn) string {
func intakeTransport(obfuscated bool) string {
if obfuscated {
return "obfuscated_tcp"
return "tcp_auto"
}
return "tcp"
}

10
internal/mtprotoedge/e2e_test.go
View file

@ -76,9 +76,13 @@ func TestTelegramClientEndToEnd(t *testing.T) {
if cfg.ThisDC != dc {
t.Errorf("config.ThisDC = %d, want %d", cfg.ThisDC, dc)
}
// 不下发 DCOptions：客户端使用自己的 DCList / 写死 static 地址。
if len(cfg.DCOptions) != 0 {
t.Errorf("config.DCOptions = %+v, want empty", cfg.DCOptions)
if len(cfg.DCOptions) != 1 {
t.Errorf("config.DCOptions = %+v, want one reconnect route", cfg.DCOptions)
} else {
option := cfg.DCOptions[0]
if option.ID != dc || option.IPAddress != tcpAddr.IP.String() || option.Port != tcpAddr.Port {
t.Errorf("config.DCOptions[0] = %+v, want dc=%d at %s", option, dc, tcpAddr)
}
}
return nil
}); err != nil {

239
internal/mtprotoedge/encrypted.go
View file

@ -3,6 +3,7 @@ package mtprotoedge
import (
"bytes"
"compress/gzip"
"compress/zlib"
"context"
"crypto/sha256"
"encoding/binary"
@ -349,6 +350,35 @@ func (e *dispatchBadMsgError) Error() string {
return fmt.Sprintf("bad client message %d/%d: code %d", e.msgID, e.seqNo, e.code)
}
var errGZIPExpansionLimit = errors.New("gzip expansion limit exceeded")
type gzipExpansionWorkError struct {
expanded int
cause error
}
func (e *gzipExpansionWorkError) Error() string {
if e == nil || e.cause == nil {
return "gzip expansion failed"
}
return e.cause.Error()
}
func (e *gzipExpansionWorkError) Unwrap() error {
if e == nil {
return nil
}
return e.cause
}
func gzipExpansionWork(err error) int {
var work *gzipExpansionWorkError
if errors.As(err, &work) && work.expanded > 0 {
return work.expanded
}
return 0
}
// decodeGZIPWithGlobalBudget reserves the maximum single-wrapper output before
// decompression starts. Once the actual size is known the excess reservation is
// returned, while the actual output remains charged until the inbound plan is
@ -356,6 +386,16 @@ func (e *dispatchBadMsgError) Error() string {
// This closes the gap where every connection read goroutine could otherwise hold
// an unaccounted 10 MiB expansion before the shared RPC scheduler saw the body.
func (s *Server) decodeGZIPWithGlobalBudget(b *bin.Buffer) ([]byte, func(), error) {
return s.decodeGZIPWithGlobalBudgetLimit(b, maxSingleGZIPExpandedBytes)
}
// decodeGZIPWithGlobalBudgetLimit is the caller-bounded form used by exact
// Layer admission. limit is also capped by the protocol's single-wrapper
// ceiling; the returned bytes remain charged until release is called.
func (s *Server) decodeGZIPWithGlobalBudgetLimit(b *bin.Buffer, limit int) ([]byte, func(), error) {
if limit <= 0 || limit > maxSingleGZIPExpandedBytes {
return nil, func() {}, fmt.Errorf("invalid gzip expansion limit %d", limit)
}
compressed, err := gzipPackedBytesView(b)
if err != nil {
return nil, func() {}, err
@ -368,30 +408,33 @@ func (s *Server) decodeGZIPWithGlobalBudget(b *bin.Buffer) ([]byte, func(), erro
}
}
if s.frameBudget != nil {
reserved, err = s.frameBudget.reserve(maxSingleGZIPExpandedBytes, 0)
reserved, err = s.frameBudget.reserve(int64(limit), 0)
if err != nil {
return nil, func() {}, err
}
}
r, err := gzip.NewReader(bytes.NewReader(compressed))
r, err := newGZIPPackedReader(compressed)
if err != nil {
release()
return nil, func() {}, err
}
data, readErr := io.ReadAll(io.LimitReader(r, maxSingleGZIPExpandedBytes+1))
data, readErr := io.ReadAll(io.LimitReader(r, int64(limit)+1))
closeErr := r.Close()
if readErr != nil {
release()
return nil, func() {}, readErr
return nil, func() {}, &gzipExpansionWorkError{expanded: len(data), cause: readErr}
}
if closeErr != nil {
release()
return nil, func() {}, closeErr
return nil, func() {}, &gzipExpansionWorkError{expanded: len(data), cause: closeErr}
}
if len(data) > maxSingleGZIPExpandedBytes {
if len(data) > limit {
release()
return nil, func() {}, fmt.Errorf("gzip expansion %d exceeds %d", len(data), maxSingleGZIPExpandedBytes)
return nil, func() {}, &gzipExpansionWorkError{
expanded: len(data),
cause: fmt.Errorf("%w: expansion %d exceeds %d", errGZIPExpansionLimit, len(data), limit),
}
}
if reserved > int64(len(data)) {
s.frameBudget.release(reserved - int64(len(data)))
@ -400,6 +443,19 @@ func (s *Server) decodeGZIPWithGlobalBudget(b *bin.Buffer) ([]byte, func(), erro
return data, release, nil
}
// newGZIPPackedReader accepts the two wrapped DEFLATE formats emitted by
// official Telegram clients. TDLib uses a zlib wrapper while DrKLO/gotd use a
// gzip wrapper; raw DEFLATE is deliberately unsupported. Selecting by the gzip
// magic keeps malformed gzip input on the gzip validator instead of silently
// retrying it as another format.
func newGZIPPackedReader(compressed []byte) (io.ReadCloser, error) {
source := bytes.NewReader(compressed)
if len(compressed) >= 2 && compressed[0] == 0x1f && compressed[1] == 0x8b {
return gzip.NewReader(source)
}
return zlib.NewReader(source)
}
// gzipPackedBytesView parses the TL bytes envelope without copying the compressed
// payload. proto.GZIP.Decode calls bin.Buffer.Bytes, which duplicates the compressed
// frame before allocating the decompressed result.
@ -674,10 +730,7 @@ func (s *Server) handleInboundRPCAdmissionError(ctx context.Context, c *Conn, ms
zap.String("auth_key_id", c.authKeyHex),
zap.Int64("session_id", c.sessionID),
)
return s.sendResult(ctx, c, msgID, &mt.RPCError{
ErrorCode: 420,
ErrorMessage: "FLOOD_WAIT_1",
})
return s.sendResult(ctx, c, msgID, rpcWorkerBusyError())
}
return err
}
@ -811,9 +864,9 @@ var errRPCResultRetentionHandoff = errors.New("mtproto rpc result retention hand
type rpcResultRetentionHandoff func(*encodedOutboundMessage, error) error
// publishRPCResult ends the inbound worker's ownership at bounded egress
// 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.
// admission. Physical delivery is thereafter owned by the logical-session
// outbox. Under retained-byte saturation the Conn is fenced and the receipt
// ledger records an unavailable tombstone so business cannot rerun.
func (s *Server) publishRPCResult(
c *Conn,
reqMsgID int64,
@ -851,23 +904,13 @@ func (s *Server) publishRPCResult(
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.
// If the sole logical-outbox body budget cannot admit a completed result,
// fence this physical generation and publish only an execution tombstone.
// There is deliberately no fallback payload cache/spool and no business
// re-execution hidden behind a local capacity error.
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,
))
return errors.New("rpc result receipt ledger is unavailable")
}
priority, visible := prepareEncoded(encoded)
if owner != nil && !owner.HandOff() {
@ -875,10 +918,10 @@ func (s *Server) publishRPCResult(
}
started := time.Now()
encoded.markReplayable()
// Put may expose a completed result only after the old logical connection
// Complete may expose terminal execution only after the old connection
// is irreversibly unable to accept another same-generation request.
c.fenceUndeliveredRPCResult()
s.storeRPCResult(c, reqMsgID, encoded)
s.completeRPCResult(c, reqMsgID, encoded, false)
latency := time.Since(started)
if metrics, ok := s.metrics.(RPCResultMetrics); ok {
metrics.RPCResultDelivered(method, latency, len(encoded.body), admissionErr)
@ -887,7 +930,7 @@ func (s *Server) publishRPCResult(
if visible {
resultLogLevel = zap.InfoLevel
}
if checked := s.log.Check(resultLogLevel, "RPC result retained for replay after egress saturation"); checked != nil {
if checked := s.log.Check(resultLogLevel, "RPC result execution fenced 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()),
@ -952,7 +995,7 @@ func (s *Server) publishRPCResult(
if deliveryErr != nil {
encoded.markReplayable()
c.fenceUndeliveredRPCResult()
s.storeRPCResult(c, reqMsgID, encoded)
s.completeRPCResult(c, reqMsgID, encoded, true)
if checked := s.log.Check(resultLogLevel, "RPC result delivery fenced for replay"); checked != nil {
checked.Write(
zap.String("method", method), zap.Int64("req_msg_id", reqMsgID),
@ -964,7 +1007,7 @@ func (s *Server) publishRPCResult(
return
}
encoded.markDelivered()
s.storeRPCResult(c, reqMsgID, encoded)
s.completeRPCResult(c, reqMsgID, encoded, true)
if checked := s.log.Check(resultLogLevel, "RPC result delivered"); checked != nil {
checked.Write(
zap.String("method", method), zap.Int64("req_msg_id", reqMsgID),
@ -1017,24 +1060,24 @@ func (s *Server) sendResult(ctx context.Context, c *Conn, reqMsgID int64, result
// same-Conn duplicate would be ACKed while no result can ever arrive.
c.fenceUndeliveredRPCResult()
encoded.markReplayable()
s.storeRPCResult(c, reqMsgID, encoded)
s.completeRPCResult(c, reqMsgID, encoded, true)
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)
s.completeRPCResult(c, reqMsgID, encoded, true)
return nil
}
// sendCachedRPCResult preserves the delivery half of the rpc_result invariant
// for completed-flight replays: either the cached result reaches this physical
// sendReplayedRPCResult preserves the delivery half of the rpc_result invariant
// for completed-flight replays: either the logical outbox 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) sendReplayedRPCResult(ctx context.Context, c *Conn, encoded *encodedOutboundMessage) error {
return s.sendReplayedRPCResultWithHook(ctx, c, encoded, nil)
}
func (s *Server) sendCachedRPCResultWithHook(
func (s *Server) sendReplayedRPCResultWithHook(
ctx context.Context,
c *Conn,
encoded *encodedOutboundMessage,
@ -1042,7 +1085,7 @@ func (s *Server) sendCachedRPCResultWithHook(
) error {
if encoded == nil {
c.fenceUndeliveredRPCResult()
return errors.New("nil cached rpc_result")
return errors.New("nil replayed rpc_result")
}
attempt, reserved, err := c.cloneRPCResultForRequestReserved(encoded, encoded.reqMsgID, false)
if err != nil {
@ -1059,7 +1102,7 @@ func (s *Server) sendCachedRPCResultWithHook(
finishRestore = c.beginRPCReplayRestore()
defer finishRestore()
}
// Cached replay owns its delivery-gated state synchronously. Calling the
// Outbox 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(
@ -1080,10 +1123,10 @@ func (s *Server) sendCachedRPCResultWithHook(
// 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 fmt.Errorf("wait for replayed rpc_result logical restore: %w", claimErr)
}
return s.runBoundedRPCReplayRestore(
restoreCtx, c, "cached rpc_result", logicalRestore, afterSuccessfulDelivery,
restoreCtx, c, "replayed rpc_result", logicalRestore, afterSuccessfulDelivery,
)
}
@ -1209,8 +1252,9 @@ func (s *Server) encodeRPCResultReservedWithHandoffContext(
}
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.
// producer reference after the encode slot releases. Production handoff
// either transferred the body to the logical outbox or retained only an
// unavailable receipt tombstone.
encoded = nil
return admissionErr
}
@ -1282,11 +1326,11 @@ func (s *Server) encodeRPCResultWithoutSlot(ctx context.Context, c *Conn, reqMsg
}, nil
}
func (s *Server) cachedRPCResult(c *Conn, reqMsgID int64) (*encodedOutboundMessage, bool) {
func (s *Server) replayableRPCResult(c *Conn, reqMsgID int64) (*encodedOutboundMessage, bool) {
if s == nil || s.rpcResults == nil || c == nil {
return nil, false
}
return s.rpcResults.Get(c.authKeyID, c.sessionID, reqMsgID)
return s.rpcResults.Replay(c.authKeyID, c.sessionID, reqMsgID)
}
func (s *Server) replayRPCResultByRequest(ctx context.Context, c *Conn, reqMsgID int64) error {
@ -1297,23 +1341,26 @@ func (s *Server) replayRPCResultByRequest(ctx context.Context, c *Conn, reqMsgID
c.fenceUndeliveredRPCResult()
return err
} else if resent {
s.log.Debug("Resent connection cached rpc_result for duplicate msg_id", zap.Int64("msg_id", reqMsgID))
s.log.Debug("Resent connection-retained rpc_result for duplicate msg_id", zap.Int64("msg_id", reqMsgID))
return nil
}
if cached, ok := s.cachedRPCResult(c, reqMsgID); ok {
if err := s.sendCachedRPCResult(ctx, c, cached); err != nil {
if replayed, ok := s.replayableRPCResult(c, reqMsgID); ok {
if err := s.sendReplayedRPCResult(ctx, c, replayed); err != nil {
return err
}
s.log.Debug("Resent session cached rpc_result for duplicate msg_id", zap.Int64("msg_id", reqMsgID))
s.log.Debug("Resent logical-outbox rpc_result for duplicate msg_id", zap.Int64("msg_id", reqMsgID))
}
return nil
}
func (s *Server) storeRPCResult(c *Conn, reqMsgID int64, encoded *encodedOutboundMessage) {
func (s *Server) completeRPCResult(c *Conn, reqMsgID int64, encoded *encodedOutboundMessage, replayable bool) {
if s == nil || s.rpcResults == nil || c == nil {
return
}
s.rpcResults.Put(c.authKeyID, c.sessionID, reqMsgID, encoded)
if s.conns != nil {
s.conns.adoptLogicalSession(c)
}
s.rpcResults.Complete(c.authKeyID, c.sessionID, reqMsgID, encoded, replayable)
}
// sendPong 回复 mt.PingRequest / mt.PingDelayDisconnectRequest。
@ -1439,15 +1486,15 @@ func validateClientEnvelope(now time.Time, msgID int64, seqNo int32, typeID uint
if msgTime.After(now.Add(30 * time.Second)) {
return badMsgIDTooHigh
}
if clientMessageAllowsEitherSeqParity(typeID) {
return 0
}
if clientMessageNeedsAck(typeID) {
switch clientMessageContentPolicyFor(typeID) {
case clientMessageContentRequired:
if seqNo%2 == 0 {
return badMsgSeqNotOdd
}
} else if seqNo%2 != 0 {
return badMsgSeqNotEven
case clientMessageContentForbidden:
if seqNo%2 != 0 {
return badMsgSeqNotEven
}
}
return 0
}
@ -1456,48 +1503,68 @@ func validateClientContainerEnvelope(msgID int64, seqNo int32, typeID uint32) in
if !validClientMessageIDBits(msgID) {
return badMsgIDInvalidBits
}
if clientMessageAllowsEitherSeqParity(typeID) {
return 0
}
if clientMessageNeedsAck(typeID) {
switch clientMessageContentPolicyFor(typeID) {
case clientMessageContentRequired:
if seqNo%2 == 0 {
return badMsgSeqNotOdd
}
} else if seqNo%2 != 0 {
return badMsgSeqNotEven
case clientMessageContentForbidden:
if seqNo%2 != 0 {
return badMsgSeqNotEven
}
}
return 0
}
func clientMessageAllowsEitherSeqParity(typeID uint32) bool {
switch typeID {
case mt.PingDelayDisconnectRequestTypeID,
// get_future_salts 的 seqno 奇偶在客户端间不一致：部分客户端按内容消息发奇数，
// gotd 按服务消息发偶数。两者都合法（官方服务器都接受），故不在此卡奇偶，避免
// 误判 bad_msg 触发客户端重连风暴。ack/content 行为仍由 clientMessageNeedsAck 决定。
mt.GetFutureSaltsRequestTypeID:
return true
default:
return false
}
}
type clientMessageContentPolicy uint8
func clientMessageNeedsAck(typeID uint32) bool {
const (
clientMessageContentRequired clientMessageContentPolicy = iota + 1
clientMessageContentForbidden
clientMessageContentOptional
)
// clientMessageContentPolicyFor classifies the client envelope, not merely the
// constructor's usual sending convention. MTProto requires API RPCs to be
// content-related and requires containers/acknowledgements to be irrelevant,
// but clients may mark the other service constructors as either. TDLib uses
// even sequence numbers for its reconnect state/resend/cancel service batch,
// while gotd and DrKLO use odd sequence numbers for some of the same requests.
func clientMessageContentPolicyFor(typeID uint32) clientMessageContentPolicy {
switch typeID {
case proto.MessageContainerTypeID,
mt.MsgsAckTypeID,
mt.MsgCopyTypeID:
return clientMessageContentForbidden
case mt.PingRequestTypeID,
mt.PingDelayDisconnectRequestTypeID,
mt.DestroySessionRequestTypeID,
mt.HTTPWaitRequestTypeID,
mt.BadMsgNotificationTypeID,
mt.BadServerSaltTypeID,
mt.GetFutureSaltsRequestTypeID,
mt.MsgsStateReqTypeID,
mt.MsgResendReqTypeID,
mt.MsgsAllInfoTypeID,
mt.MsgsStateInfoTypeID,
mt.DestroySessionRequestTypeID,
mt.HTTPWaitRequestTypeID,
mt.RPCDropAnswerRequestTypeID,
mt.BadMsgNotificationTypeID,
mt.BadServerSaltTypeID,
mt.MsgDetailedInfoTypeID,
mt.MsgNewDetailedInfoTypeID:
return false
mt.MsgNewDetailedInfoTypeID,
destroyAuthKeyRequestTypeID:
return clientMessageContentOptional
default:
return clientMessageContentRequired
}
}
func clientMessageIsContentRelated(typeID uint32, seqNo int32) bool {
switch clientMessageContentPolicyFor(typeID) {
case clientMessageContentRequired:
return true
case clientMessageContentOptional:
return seqNo%2 != 0
default:
return false
}
}

117
internal/mtprotoedge/encrypted_test.go
View file

@ -69,6 +69,72 @@ func (h *durableDestroyLayerRPC) deletion() ([8]byte, int64) {
return h.authKeyID, h.sessionID
}
func TestClientMessageContentPolicy(t *testing.T) {
tests := []struct {
name string
typeID uint32
want clientMessageContentPolicy
}{
{name: "api_rpc", typeID: tg.HelpGetConfigRequestTypeID, want: clientMessageContentRequired},
{name: "container", typeID: proto.MessageContainerTypeID, want: clientMessageContentForbidden},
{name: "msgs_ack", typeID: mt.MsgsAckTypeID, want: clientMessageContentForbidden},
{name: "msg_copy", typeID: mt.MsgCopyTypeID, want: clientMessageContentForbidden},
{name: "bad_msg_notification", typeID: mt.BadMsgNotificationTypeID, want: clientMessageContentOptional},
{name: "bad_server_salt", typeID: mt.BadServerSaltTypeID, want: clientMessageContentOptional},
{name: "msg_detailed_info", typeID: mt.MsgDetailedInfoTypeID, want: clientMessageContentOptional},
{name: "msg_new_detailed_info", typeID: mt.MsgNewDetailedInfoTypeID, want: clientMessageContentOptional},
{name: "ping", typeID: mt.PingRequestTypeID, want: clientMessageContentOptional},
{name: "ping_delay_disconnect", typeID: mt.PingDelayDisconnectRequestTypeID, want: clientMessageContentOptional},
{name: "get_future_salts", typeID: mt.GetFutureSaltsRequestTypeID, want: clientMessageContentOptional},
{name: "msgs_state_req", typeID: mt.MsgsStateReqTypeID, want: clientMessageContentOptional},
{name: "msg_resend_req", typeID: mt.MsgResendReqTypeID, want: clientMessageContentOptional},
{name: "msgs_all_info", typeID: mt.MsgsAllInfoTypeID, want: clientMessageContentOptional},
{name: "msgs_state_info", typeID: mt.MsgsStateInfoTypeID, want: clientMessageContentOptional},
{name: "destroy_session", typeID: mt.DestroySessionRequestTypeID, want: clientMessageContentOptional},
{name: "http_wait", typeID: mt.HTTPWaitRequestTypeID, want: clientMessageContentOptional},
{name: "rpc_drop_answer", typeID: mt.RPCDropAnswerRequestTypeID, want: clientMessageContentOptional},
{name: "destroy_auth_key", typeID: destroyAuthKeyRequestTypeID, want: clientMessageContentOptional},
}
now := time.Now()
msgID := proto.NewMessageIDGen(func() time.Time { return now }).New(proto.MessageFromClient)
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
if got := clientMessageContentPolicyFor(test.typeID); got != test.want {
t.Fatalf("content policy = %d, want %d", got, test.want)
}
evenCode := validateClientContainerEnvelope(msgID, 8, test.typeID)
oddCode := validateClientContainerEnvelope(msgID, 9, test.typeID)
directEvenCode := validateClientEnvelope(now, msgID, 8, test.typeID)
directOddCode := validateClientEnvelope(now, msgID, 9, test.typeID)
if directEvenCode != evenCode || directOddCode != oddCode {
t.Fatalf(
"top-level/container parity mismatch = top(%d,%d) container(%d,%d)",
directEvenCode, directOddCode, evenCode, oddCode,
)
}
switch test.want {
case clientMessageContentRequired:
if evenCode != badMsgSeqNotOdd || oddCode != 0 {
t.Fatalf("required content parity codes = even:%d odd:%d", evenCode, oddCode)
}
case clientMessageContentForbidden:
if evenCode != 0 || oddCode != badMsgSeqNotEven {
t.Fatalf("forbidden content parity codes = even:%d odd:%d", evenCode, oddCode)
}
case clientMessageContentOptional:
if evenCode != 0 || oddCode != 0 {
t.Fatalf("optional content parity codes = even:%d odd:%d", evenCode, oddCode)
}
if clientMessageIsContentRelated(test.typeID, 8) || !clientMessageIsContentRelated(test.typeID, 9) {
t.Fatal("optional service content bit was not derived from seq_no parity")
}
}
})
}
}
// TestEncryptedPingPong 验证 M2/M4：握手后 client 加密 ping，
// server 回 new_session_created + pong + msgs_ack。
func TestEncryptedPingPong(t *testing.T) {
@ -79,7 +145,7 @@ func TestEncryptedPingPong(t *testing.T) {
clientMsgID := proto.NewMessageIDGen(time.Now)
const pingID int64 = 0x1234beef
pingMsgID := clientMsgID.New(proto.MessageFromClient)
sendEncrypted(t, conn, cipher, auth, pingMsgID, &mt.PingRequest{PingID: pingID})
sendEncryptedWithSeq(t, conn, cipher, auth, pingMsgID, 1, &mt.PingRequest{PingID: pingID})
replies := collectReplies(t, conn, cipher, auth.AuthKey, mt.PongTypeID)
mustHave(t, replies, mt.NewSessionCreatedTypeID, "new_session_created")
@ -230,6 +296,55 @@ func TestMsgsStateReq(t *testing.T) {
}
}
// TestTDLibReconnectRecoveryContainerAcceptsEvenServiceMessages reproduces the
// first container TDLib emits after reopening an authenticated session with
// unknown queries. All service entries and the outer container use the current
// even sequence number. Rejecting msgs_state_req as a content-only constructor
// turns the valid inner message into bad_msg_notification(code=64), after which
// TDLib closes the session and retries the same container forever.
func TestTDLibReconnectRecoveryContainerAcceptsEvenServiceMessages(t *testing.T) {
const dc = 2
addr, pub, _ := startTestServer(t, Options{DC: dc})
conn, auth, cipher := dialHandshake(t, addr, dc, pub)
ids := proto.NewMessageIDGen(time.Now)
pendingMsgID := ids.New(proto.MessageFromClient)
ackMsgID := ids.New(proto.MessageFromClient)
stateMsgID := ids.New(proto.MessageFromClient)
pingMsgID := ids.New(proto.MessageFromClient)
outerMsgID := ids.New(proto.MessageFromClient)
const (
pendingSeqNo int32 = 9
serviceSeqNo int32 = 10
)
pendingBody := mustEncodeTL(t, &mt.PingRequest{PingID: 0xc01d})
ackBody := mustEncodeTL(t, &mt.MsgsAck{MsgIDs: []int64{stateMsgID - 4}})
stateBody := mustEncodeTL(t, &mt.MsgsStateReq{MsgIDs: []int64{pendingMsgID, stateMsgID - 4}})
pingBody := mustEncodeTL(t, &mt.PingDelayDisconnectRequest{PingID: 0x5eed, DisconnectDelay: 60})
container := &proto.MessageContainer{Messages: []proto.Message{
{ID: pendingMsgID, SeqNo: int(pendingSeqNo), Bytes: len(pendingBody), Body: pendingBody},
{ID: ackMsgID, SeqNo: int(serviceSeqNo), Bytes: len(ackBody), Body: ackBody},
{ID: stateMsgID, SeqNo: int(serviceSeqNo), Bytes: len(stateBody), Body: stateBody},
{ID: pingMsgID, SeqNo: int(serviceSeqNo), Bytes: len(pingBody), Body: pingBody},
}}
sendEncryptedWithSeq(t, conn, cipher, auth, outerMsgID, serviceSeqNo, container)
frames := collectReplyFrames(t, conn, cipher, auth.AuthKey, map[uint32]int{
mt.MsgsStateInfoTypeID: 1,
mt.PongTypeID: 2,
})
for _, frame := range frames {
if frame.TypeID == mt.BadMsgNotificationTypeID {
var bad mt.BadMsgNotification
if err := bad.Decode(frame.Plain); err != nil {
t.Fatalf("decode bad_msg_notification: %v", err)
}
t.Fatalf("TDLib reconnect recovery container was rejected: %+v", bad)
}
}
}
// TestMsgResendReq 验证 MTProto msg_resend_req 由连接层按状态查询兜底响应，
// 不会落入业务 RPC fallback。
func TestMsgResendReq(t *testing.T) {

53
internal/mtprotoedge/exchange_compat.go
View file

@ -25,11 +25,12 @@ import (
// runServerExchange is a gotd server exchange compatibility shim.
//
// DrKLO Android marks media temporary auth-key exchange with a negative DC in
// p_q_inner_data_temp_dc (for example DC 2 -> -2). gotd v0.158.0 validates this
// field by exact equality and rejects that legitimate media-temp path. Keep the
// permanent-key check strict, but allow temp-key DC values whose absolute value
// matches this server DC.
// In the default single-backend mode, p_q_inner_data_dc and
// p_q_inner_data_temp_dc carry client routing labels only: every int32 value is
// admitted and the label is not persisted or used for key/session identity.
// This also covers DrKLO Android's negative media-temp labels. StrictDC retains
// exact permanent / absolute-value temporary validation as an explicit,
// default-off diagnostic for a future real multi-DC deployment.
func (s *Server) runServerExchange(ctx context.Context, conn transport.Conn) (exchange.ServerExchangeResult, error) {
ex := serverExchangeCompat{
conn: conn,
@ -355,7 +356,7 @@ func (s serverExchangeCompat) validatePQInnerDataDC(d mt.PQInnerDataClass) error
// configured DC is expected, not an error. dc_id plays no
// role in key derivation, so accepting it doesn't weaken the
// exchange.
s.log.Debug("Accepted permanent auth key DC mismatch (lenient mode)",
s.log.Debug("Accepted permanent auth key DC alias",
zap.Int("server_dc", s.dc),
zap.Int("client_dc", innerDataDC.DC))
return nil
@ -366,8 +367,8 @@ func (s serverExchangeCompat) validatePQInnerDataDC(d mt.PQInnerDataClass) error
if !sameDCByAbs(innerDataDC.DC, s.dc) && s.strictDC {
return wrongDCError(s.dc, innerDataDC.DC)
}
if innerDataDC.DC < 0 {
s.log.Warn("Accepted Android media temp auth key negative DC",
if !sameDCByAbs(innerDataDC.DC, s.dc) {
s.log.Debug("Accepted temporary auth key DC alias",
zap.Int("server_dc", s.dc),
zap.Int("client_dc", innerDataDC.DC),
zap.Int("expires_in", innerDataDC.ExpiresIn))
@ -440,7 +441,7 @@ func (s serverExchangeCompat) readUnencrypted(ctx context.Context, b *bin.Buffer
if err := msg.Decode(b); err != nil {
return err
}
if !validClientMessageIDBits(msg.MessageID) {
if !validUnencryptedHandshakeMessageID(msg.MessageID, msg.MessageData) {
return gofaster.New("bad msg type")
}
b.ResetTo(msg.MessageData)
@ -448,6 +449,40 @@ func (s serverExchangeCompat) readUnencrypted(ctx context.Context, b *bin.Buffer
return data.Decode(b)
}
// validUnencryptedHandshakeMessageID preserves the normal client message-id
// rules while admitting the two sentinel ids emitted by official TDLib:
//
// - PingConnectionReqPQ sends req_pq_multi with message_id=1.
// - HandshakeConnection sends every auth-key exchange request with message_id=0.
//
// The exception is deliberately constructor-scoped and is only called after an
// auth_key_id=0 envelope has been decoded. Encrypted traffic continues through
// validClientMessageIDBits and the full inbound preflight without this carve-out.
func validUnencryptedHandshakeMessageID(messageID int64, messageData []byte) bool {
if validClientMessageIDBits(messageID) {
return true
}
payload := &bin.Buffer{Buf: messageData}
typeID, err := payload.PeekID()
if err != nil {
return false
}
switch messageID {
case 1:
return typeID == mt.ReqPqMultiRequestTypeID
case 0:
switch typeID {
case mt.ReqPqMultiRequestTypeID,
mt.ReqDHParamsRequestTypeID,
mt.SetClientDHParamsRequestTypeID:
return true
}
}
return false
}
type compatReqPQ struct {
Type uint32
Nonce bin.Int128

344
internal/mtprotoedge/exchange_test.go
View file

@ -7,6 +7,7 @@ import (
"crypto/rsa"
"encoding/binary"
"errors"
"fmt"
"math/big"
"net"
"testing"
@ -109,6 +110,221 @@ func TestKeyExchange(t *testing.T) {
}
}
type tdlibZeroHandshakeMessageIDConn struct {
transport.Conn
}
func (c *tdlibZeroHandshakeMessageIDConn) Send(ctx context.Context, frame *bin.Buffer) error {
candidate := &bin.Buffer{Buf: frame.Copy()}
var message tgproto.UnencryptedMessage
if err := message.Decode(candidate); err != nil {
return c.Conn.Send(ctx, frame)
}
payload := &bin.Buffer{Buf: message.MessageData}
typeID, err := payload.PeekID()
if err != nil {
return c.Conn.Send(ctx, frame)
}
switch typeID {
case mt.ReqPqMultiRequestTypeID,
mt.ReqDHParamsRequestTypeID,
mt.SetClientDHParamsRequestTypeID:
message.MessageID = 0
default:
return c.Conn.Send(ctx, frame)
}
// TDLib NoCryptoImpl includes 0-255 random bytes in message_data_length.
// Use its maximum legal alignment-plus-15-block shape so the complete
// permanent and temporary exchanges exercise padded bodies at every stage.
paddingSize := (-len(message.MessageData)) & 15
paddingSize += 16 * 15
message.MessageData = append(
message.MessageData,
bytes.Repeat([]byte{0xa5}, paddingSize)...,
)
var rewritten bin.Buffer
if err := message.Encode(&rewritten); err != nil {
return err
}
return c.Conn.Send(ctx, &rewritten)
}
func TestKeyExchangeAcceptsTDLibZeroMessageIDs(t *testing.T) {
const (
dc = 2
expiresIn = 60
)
tests := []struct {
name string
clientDC int
temporary bool
wantExpiry bool
}{
{name: "permanent key", clientDC: dc},
{
name: "media temporary key",
clientDC: -dc,
temporary: true,
wantExpiry: true,
},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
addr, pub, srv := startTestServer(t, Options{DC: dc})
conn := &tdlibZeroHandshakeMessageIDConn{Conn: dialTransportOnly(t, addr)}
t.Cleanup(func() { _ = conn.Close() })
exchanger := exchange.NewExchanger(conn, test.clientDC).
WithRand(rand.Reader).
WithLogger(logzap.New(zaptest.NewLogger(t).Named("tdlib-zero-msg-id-client")))
if test.temporary {
exchanger = exchanger.WithTempMode(expiresIn)
}
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
defer cancel()
result, err := exchanger.
Client([]exchange.PublicKey{pub}).
Run(ctx)
if err != nil {
t.Fatalf("TDLib-shaped client exchange: %v", err)
}
if result.AuthKey.ID == ([8]byte{}) {
t.Fatal("TDLib-shaped client exchange returned an empty auth key id")
}
if got := result.ExpiresAt > 0; got != test.wantExpiry {
t.Fatalf("client expiry present = %v, want %v", got, test.wantExpiry)
}
var (
saved store.AuthKeyData
found bool
)
deadline := time.Now().Add(5 * time.Second)
for time.Now().Before(deadline) {
saved, found, _ = srv.authKeys.Get(context.Background(), result.AuthKey.ID)
if found {
break
}
time.Sleep(20 * time.Millisecond)
}
if !found {
t.Fatalf("server did not store TDLib-shaped auth key %x", result.AuthKey.ID)
}
if got := saved.ExpiresAt > 0; got != test.wantExpiry {
t.Fatalf("server expiry present = %v, want %v", got, test.wantExpiry)
}
})
}
}
func TestValidUnencryptedHandshakeMessageID(t *testing.T) {
encode := func(value bin.Encoder) []byte {
t.Helper()
var payload bin.Buffer
if err := value.Encode(&payload); err != nil {
t.Fatalf("encode payload: %v", err)
}
return payload.Copy()
}
const ordinaryClientMessageID = int64(1<<32 | 4)
tests := []struct {
name string
messageID int64
payload []byte
want bool
}{
{
name: "ordinary client id retains existing admission",
messageID: ordinaryClientMessageID,
payload: encode(&mt.PingRequest{}),
want: true,
},
{
name: "probe sentinel rejects legacy req pq",
messageID: 1,
payload: encode(&mt.ReqPqRequest{}),
want: false,
},
{
name: "TDLib probe req pq multi",
messageID: 1,
payload: encode(&mt.ReqPqMultiRequest{}),
want: true,
},
{
name: "probe sentinel cannot carry req DH",
messageID: 1,
payload: encode(&mt.ReqDHParamsRequest{}),
want: false,
},
{
name: "zero sentinel rejects legacy req pq",
messageID: 0,
payload: encode(&mt.ReqPqRequest{}),
want: false,
},
{
name: "TDLib handshake req pq multi",
messageID: 0,
payload: encode(&mt.ReqPqMultiRequest{}),
want: true,
},
{
name: "TDLib handshake req DH",
messageID: 0,
payload: encode(&mt.ReqDHParamsRequest{}),
want: true,
},
{
name: "TDLib handshake set client DH",
messageID: 0,
payload: encode(&mt.SetClientDHParamsRequest{}),
want: true,
},
{
name: "zero sentinel cannot carry ack",
messageID: 0,
payload: encode(&mt.MsgsAck{}),
want: false,
},
{
name: "other invalid nonzero id remains rejected",
messageID: 2,
payload: encode(&mt.ReqPqMultiRequest{}),
want: false,
},
{
name: "negative id remains rejected",
messageID: -4,
payload: encode(&mt.ReqPqMultiRequest{}),
want: false,
},
{
name: "sentinel requires a complete constructor id",
messageID: 0,
payload: []byte{1, 2, 3},
want: false,
},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
if got := validUnencryptedHandshakeMessageID(test.messageID, test.payload); got != test.want {
t.Fatalf(
"validUnencryptedHandshakeMessageID(%d) = %v, want %v",
test.messageID,
got,
test.want,
)
}
})
}
}
type authKeySaveContextObservation struct {
hasDeadline bool
deadline time.Time
@ -145,6 +361,52 @@ func (c *ownershipFrameConn) Recv(_ context.Context, b *bin.Buffer) error {
return nil
}
func TestReadUnencryptedAcceptsTDLibProbeMessageID(t *testing.T) {
nonce := bin.Int128{1, 2, 3, 4}
var payload bin.Buffer
if err := (&mt.ReqPqMultiRequest{Nonce: nonce}).Encode(&payload); err != nil {
t.Fatalf("encode req_pq_multi: %v", err)
}
tests := []struct {
name string
padding []byte
}{
{name: "without TDLib no-crypto padding"},
{name: "with minimum TDLib no-crypto padding", padding: bytes.Repeat([]byte{0xa5}, 12)},
{name: "with maximum TDLib no-crypto padding", padding: bytes.Repeat([]byte{0x5a}, 252)},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
messageData := append(payload.Copy(), test.padding...)
var frame bin.Buffer
if err := (tgproto.UnencryptedMessage{
MessageID: 1,
MessageData: messageData,
}).Encode(&frame); err != nil {
t.Fatalf("encode unencrypted probe: %v", err)
}
ex := serverExchangeCompat{
conn: &ownershipFrameConn{frame: frame.Copy()},
timeout: time.Second,
}
var decoded compatReqPQ
var scratch bin.Buffer
if err := ex.readUnencrypted(context.Background(), &scratch, &decoded); err != nil {
t.Fatalf("read TDLib probe: %v", err)
}
if decoded.Type != mt.ReqPqMultiRequestTypeID || decoded.Nonce != nonce {
t.Fatalf(
"decoded probe = {type:%#x nonce:%x}, want req_pq_multi nonce %x",
decoded.Type,
decoded.Nonce,
nonce,
)
}
})
}
}
func TestExchangeEncryptedReplayTransfersFrameOwnership(t *testing.T) {
backing := make([]byte, 64)
copy(backing[:8], []byte{1, 2, 3, 4, 5, 6, 7, 8})
@ -366,9 +628,85 @@ func TestKeyExchangeAcceptsAndroidMediaTempNegativeDC(t *testing.T) {
}
}
func TestKeyExchangeRejectsWrongNegativeTempDCWhenStrict(t *testing.T) {
func TestKeyExchangeAcceptsAnyDCLabelByDefault(t *testing.T) {
ex := serverExchangeCompat{dc: 2, log: zaptest.NewLogger(t)}
labels := []int{
2, // canonical
3, // another production DC
0, // no conventional DC mapping
-2, // Android media-temp convention
10002, // test-environment style label
-10002, // negative test-environment style label
-1 << 31,
1<<31 - 1,
}
for _, label := range labels {
t.Run(fmt.Sprintf("permanent_%d", label), func(t *testing.T) {
if err := ex.validatePQInnerDataDC(&mt.PQInnerDataDC{DC: label}); err != nil {
t.Fatalf("validate permanent DC label %d: %v", label, err)
}
})
t.Run(fmt.Sprintf("temporary_%d", label), func(t *testing.T) {
if err := ex.validatePQInnerDataDC(&mt.PQInnerDataTempDC{DC: label, ExpiresIn: 60}); err != nil {
t.Fatalf("validate temporary DC label %d: %v", label, err)
}
})
}
}
func TestKeyExchangePersistsArbitraryPermanentDCLabelByDefault(t *testing.T) {
const clientDC = 10002
keys := memory.NewAuthKeyStore()
addr, pub, _ := startTestServer(t, Options{DC: 2, AuthKeys: keys})
_, auth, _ := dialHandshake(t, addr, clientDC, pub)
saved, found, err := keys.Get(context.Background(), auth.AuthKey.ID)
if err != nil {
t.Fatalf("get persisted auth key: %v", err)
}
if !found {
t.Fatalf("auth key %x was not persisted", auth.AuthKey.ID)
}
if saved.Value != [256]byte(auth.AuthKey.Value) {
t.Fatal("persisted auth key value mismatch")
}
}
func TestKeyExchangeStrictDCValidation(t *testing.T) {
ex := serverExchangeCompat{dc: 2, strictDC: true, log: zaptest.NewLogger(t)}
err := ex.validatePQInnerDataDC(&mt.PQInnerDataTempDC{DC: -3})
tests := []struct {
name string
data mt.PQInnerDataClass
wantErr bool
}{
{name: "permanent exact", data: &mt.PQInnerDataDC{DC: 2}},
{name: "permanent other", data: &mt.PQInnerDataDC{DC: 3}, wantErr: true},
{name: "permanent zero", data: &mt.PQInnerDataDC{DC: 0}, wantErr: true},
{name: "temporary positive exact", data: &mt.PQInnerDataTempDC{DC: 2}},
{name: "temporary negative exact", data: &mt.PQInnerDataTempDC{DC: -2}},
{name: "temporary other", data: &mt.PQInnerDataTempDC{DC: 3}, wantErr: true},
{name: "temporary negative other", data: &mt.PQInnerDataTempDC{DC: -3}, wantErr: true},
{name: "temporary test label", data: &mt.PQInnerDataTempDC{DC: 10002}, wantErr: true},
{name: "temporary min int32", data: &mt.PQInnerDataTempDC{DC: -1 << 31}, wantErr: true},
}
for _, test := range tests {
t.Run(test.name, func(t *testing.T) {
err := ex.validatePQInnerDataDC(test.data)
if !test.wantErr {
if err != nil {
t.Fatalf("validate: %v", err)
}
return
}
assertWrongDCError(t, err)
})
}
}
func assertWrongDCError(t *testing.T, err error) {
t.Helper()
var exErr *exchange.ServerExchangeError
if !errors.As(err, &exErr) {
t.Fatalf("err = %T %v, want ServerExchangeError", err, err)
@ -738,7 +1076,7 @@ func TestReconnectFakeReqPQThenEncryptedFrame(t *testing.T) {
cancel()
msgGen := tgproto.NewMessageIDGen(time.Now)
sendEncrypted(t, conn, cipher, auth, msgGen.New(tgproto.MessageFromClient), &mt.PingRequest{PingID: 7})
sendEncryptedWithSeq(t, conn, cipher, auth, msgGen.New(tgproto.MessageFromClient), 1, &mt.PingRequest{PingID: 7})
var resPQFrame bin.Buffer
ctx, cancel = context.WithTimeout(context.Background(), 5*time.Second)

8
internal/mtprotoedge/frame_budget_test.go
View file

@ -51,7 +51,13 @@ func (a frameBudgetTestAddr) String() string { return string(a) }
func newFrameBudgetTestTransport(packet []byte, c transport.Codec, budget *inboundFrameBudget) (*compatTransportConn, *frameBudgetTestConn) {
raw := newFrameBudgetTestConn(packet)
return &compatTransportConn{conn: raw, codec: c, budget: budget}, raw
return &compatTransportConn{
conn: raw,
codec: c,
codecKind: classifyInboundFrameCodec(c),
budgetedCodec: unwrapInboundFrameBudgetedCodec(c),
budget: budget,
}, raw
}
func TestInboundFrameBudgetSupportsBuiltInCodecs(t *testing.T) {

2
internal/mtprotoedge/helpers_test.go
View file

@ -250,7 +250,7 @@ func encodeClientMessageForTest(t *testing.T, msg bin.Encoder) ([]byte, int32) {
if container, ok := msg.(*proto.MessageContainer); ok {
return raw, clientContainerSeqNoForTest(container)
}
if clientMessageNeedsAck(typeID) {
if clientMessageContentPolicyFor(typeID) == clientMessageContentRequired {
return raw, 1
}
return raw, 0

334
internal/mtprotoedge/inbound_layer_rpc.go
View file

@ -25,6 +25,8 @@ var inboundLayerDecodeLimits = tlprofile.Limits{
}
var errDefaultLayerAdmission = errors.New("selected layer profile rejected naked RPC")
var errLayerRPCGZIPCapacity = errors.New("exact layer gzip admission capacity exhausted")
var errLayerRPCAdmissionCapability = errors.New("exact layer RPC admission capability unavailable")
const (
maxLayerRPCDependencyIDs = 128
@ -40,6 +42,7 @@ const (
// 32-MiB per-connection budget and therefore remains admissible by default.
layerRPCAdmissionStaticObjectBytes = 512
layerRPCAdmissionWireFactor = 60
layerRPCAdmissionFlatBytesFactor = 2
layerRPCAdmissionGraphSlack = layerRPCAdmissionStaticObjectBytes * 32
)
@ -59,7 +62,7 @@ type layerRPCProfileEvidence struct {
// layerRPCAdmissionCursor is the wire-ordered, side-effect-free view used while
// decoding one MTProto container. evidenceMsgID is the last explicit
// invokeWithLayer proof, not merely the profile used by an arbitrary cached
// invokeWithLayer proof, not merely the profile used by an arbitrary retained
// request. It therefore advances only after generated admission reports
// ProfileEvidence.
type layerRPCAdmissionCursor struct {
@ -162,14 +165,136 @@ func (s *Server) initialLayerRPCAdmissionCursor(ctx context.Context, c *Conn) (l
// Saturation deliberately turns hostile integer-sized inputs into ordinary
// capacity rejection; it must never wrap into a small accepted reservation.
func layerRPCAdmissionReservationSize(wireBytes int) int {
return saturatingLayerRPCAdmissionCharge(
layerRPCAdmissionGraphSlack,
layerRPCAdmissionWireCharge(wireBytes),
)
}
func layerRPCAdmissionWireCharge(wireBytes int) int {
if wireBytes < 0 {
wireBytes = 0
}
maxInt := int(^uint(0) >> 1)
if wireBytes > (maxInt-layerRPCAdmissionGraphSlack)/layerRPCAdmissionWireFactor {
if wireBytes > maxInt/layerRPCAdmissionWireFactor {
return maxInt
}
return wireBytes*layerRPCAdmissionWireFactor + layerRPCAdmissionGraphSlack
return wireBytes * layerRPCAdmissionWireFactor
}
// layerRPCFlatBytesWireCharge keeps the generic factor on fixed TL wire and
// applies a tight copy factor only to a payload which the production Router has
// already proven to be one bounded flat bytes field. The temporary expanded
// buffer remains independently owned by inboundFrameBudget while generated
// admission materializes the request; 2x covers the retained bytes copy plus
// allocator rounding without treating every payload byte as a possible nested
// object/vector node. The per-request graph slack is owned by the base
// reservation and must not be repeated for each nested gzip expansion.
func layerRPCFlatBytesWireCharge(wireBytes, payloadBytes int) int {
if wireBytes < 0 || payloadBytes < 0 || payloadBytes > wireBytes {
return layerRPCAdmissionWireCharge(wireBytes)
}
fixedBytes := wireBytes - payloadBytes
maxInt := int(^uint(0) >> 1)
if fixedBytes > maxInt/layerRPCAdmissionWireFactor {
return maxInt
}
charge := fixedBytes * layerRPCAdmissionWireFactor
if payloadBytes > (maxInt-charge)/layerRPCAdmissionFlatBytesFactor {
return maxInt
}
return charge + payloadBytes*layerRPCAdmissionFlatBytesFactor
}
func saturatingLayerRPCAdmissionCharge(left, right int) int {
if left < 0 {
left = 0
}
if right < 0 {
right = 0
}
maxInt := int(^uint(0) >> 1)
if right > maxInt-left {
return maxInt
}
return left + right
}
func (s *Server) layerRPCExpandedWireCharge(wire []byte) int {
if s != nil {
if sizer, ok := s.layerRPC.(LayerRPCFlatBytesPayloadSizer); ok {
if payloadBytes, flat := sizer.LayerRPCFlatBytesPayloadSize(wire); flat && payloadBytes >= 0 && payloadBytes <= len(wire) {
return layerRPCFlatBytesWireCharge(len(wire), payloadBytes)
}
}
}
return layerRPCAdmissionWireCharge(len(wire))
}
// layerRPCGZIPExpansionBudget bridges transient process-wide expansion memory
// to the durable scheduler charge of one exact typed request. The original
// compressed wire retains the generic graph charge. Every successful expansion
// adds its own charge; only a handler-proven flat bytes terminal can use the
// tighter payload factor. An authoritative-profile re-decode reuses the largest
// already-held charge instead of double-counting sequential attempts.
type layerRPCGZIPExpansionBudget struct {
server *Server
plan *inboundPlan
reservation *inboundRPCBatchReservation
entry int
baseCharge int
attemptExpandedCharge int
chargedSize int
}
func (b *layerRPCGZIPExpansionBudget) beginAttempt() {
if b != nil {
b.attemptExpandedCharge = 0
}
}
func (b *layerRPCGZIPExpansionBudget) expand(wire []byte, admissionLimit int) ([]byte, func(), error) {
noop := func() {}
if b == nil || b.server == nil || b.plan == nil || b.reservation == nil {
return nil, noop, errors.New("invalid exact layer gzip expansion budget")
}
frameRemaining := maxDispatchExpandedBytes - b.plan.gzipExpandedBytes
limit := min(admissionLimit, frameRemaining)
if limit <= 0 {
return nil, noop, errors.Join(
errLayerRPCGZIPCapacity,
fmt.Errorf("cumulative gzip expansion reached %d bytes", maxDispatchExpandedBytes),
)
}
data, release, err := b.server.decodeGZIPWithGlobalBudgetLimit(&bin.Buffer{Buf: wire}, limit)
if err != nil {
if work := gzipExpansionWork(err); work > 0 {
b.plan.gzipExpandedBytes += work
}
if errors.Is(err, ErrInboundFrameBudgetExceeded) {
return nil, release, errors.Join(errLayerRPCGZIPCapacity, err)
}
if limit < admissionLimit && errors.Is(err, errGZIPExpansionLimit) {
return nil, release, errors.Join(errLayerRPCGZIPCapacity, err)
}
return nil, release, err
}
b.plan.gzipExpandedBytes += len(data)
b.attemptExpandedCharge = saturatingLayerRPCAdmissionCharge(
b.attemptExpandedCharge,
b.server.layerRPCExpandedWireCharge(data),
)
targetCharge := saturatingLayerRPCAdmissionCharge(b.baseCharge, b.attemptExpandedCharge)
if targetCharge > b.chargedSize {
if err := b.reservation.growEntry(b.entry, targetCharge); err != nil {
if errors.Is(err, ErrInboundRPCQueueFull) {
return nil, release, errors.Join(errLayerRPCGZIPCapacity, err)
}
return nil, release, err
}
b.chargedSize = targetCharge
}
return data, release, nil
}
// prepareInboundLayerRPCBatch is the production API path. The whole container
@ -217,15 +342,40 @@ func (s *Server) prepareInboundLayerRPCBatch(ctx context.Context, c *Conn, plan
return err
}
evidence := make([]layerRPCProfileEvidence, len(plan.items))
for _, index := range candidateItems {
decodeOptions := make([]tlprofile.AdmissionOptions, len(candidateItems))
expansionBudgets := make([]*layerRPCGZIPExpansionBudget, len(candidateItems))
materializationCapacity := false
for reservationIndex, index := range candidateItems {
item := &plan.items[index]
itemState := admissionCursor.state
existingProfile, existing := s.rpcResults.ExactAdmissionProfile(c.authKeyID, c.sessionID, item.msgID)
if existing {
itemState = LayerProfileSnapshot{Profile: existingProfile, Origin: LayerProfileExplicit}
}
admitted, method, err := s.decodeInboundLayerRPC(itemState, item.body)
expansionBudget := &layerRPCGZIPExpansionBudget{
server: s, plan: plan, reservation: reservation,
entry: reservationIndex,
baseCharge: provisionalSpecs[reservationIndex].size,
chargedSize: provisionalSpecs[reservationIndex].size,
}
expansionBudgets[reservationIndex] = expansionBudget
options := tlprofile.AdmissionOptions{
Limits: inboundLayerDecodeLimits,
ExpandGZIP: expansionBudget.expand,
}
decodeOptions[reservationIndex] = options
admitted, method, err := s.decodeInboundLayerRPCWithOptions(itemState, item.body, options)
if err != nil {
if errors.Is(err, errLayerRPCGZIPCapacity) {
materializationCapacity = true
break
}
if errors.Is(err, ErrConnClosed) || errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
return err
}
if errors.Is(err, errLayerRPCAdmissionCapability) {
return err
}
if terminal, recognized := wrappedDestroyAuthKeyTerminal(err); recognized {
if terminal.WireSize != bin.Word || !validWrappedDestroyAuthKeyChain(terminal) {
s.log.Debug("Wrapped destroy_auth_key terminal rejected",
@ -251,16 +401,11 @@ func (s *Server) prepareInboundLayerRPCBatch(ctx context.Context, c *Conn, plan
if errors.Is(err, ErrLayerProfileConflict) {
return err
}
s.log.Debug("RPC exact admission rejected",
zap.String("method", method),
zap.String("auth_key_id", c.authKeyHex),
zap.Int64("session_id", c.sessionID),
zap.Int64("msg_id", item.msgID),
zap.Error(err),
)
rpcError := layerRPCAdmissionError(err)
s.logLayerRPCAdmissionRejection(c, item, itemState, admissionCursor, method, rpcError, err)
item.kind = inboundItemRPCAdmissionError
item.method = method
item.payload = layerRPCAdmissionError(err)
item.payload = rpcError
c.metrics.InboundRPCDropped(method, "layer_admission")
continue
}
@ -268,7 +413,7 @@ func (s *Server) prepareInboundLayerRPCBatch(ctx context.Context, c *Conn, plan
item.method = method
if profile, hasEvidence := admitted.ProfileEvidence(); hasEvidence {
if existing && profile != existingProfile {
return fmt.Errorf("%w: cached msg_id %d used Layer %d but replay selected Layer %d", ErrLayerProfileConflict, item.msgID, existingProfile, profile)
return fmt.Errorf("%w: retained msg_id %d used Layer %d but replay selected Layer %d", ErrLayerProfileConflict, item.msgID, existingProfile, profile)
}
evidence[index] = layerRPCProfileEvidence{profile: profile, present: true, fresh: item.profileEvidenceFresh()}
if evidence[index].fresh {
@ -278,6 +423,19 @@ func (s *Server) prepareInboundLayerRPCBatch(ctx context.Context, c *Conn, plan
}
}
}
if materializationCapacity {
for _, index := range candidateItems {
item := &plan.items[index]
item.kind = inboundItemCapacityError
item.admitted = tlprofile.Admission{}
c.metrics.InboundRPCDropped(s.typeName(item.typeID), "materialization_capacity")
}
if err := reservation.retain(nil, nil); err != nil {
return err
}
plan.rpcReservation = nil
return nil
}
var indices []int
var reservationIndices []int
@ -304,7 +462,16 @@ func (s *Server) prepareInboundLayerRPCBatch(ctx context.Context, c *Conn, plan
item.admitted.Prepared().SemanticIdentity(),
item.admitted.Call().Identity(),
); candidate != nil {
claim, err := s.acquireAdmittedLayerRPC(c, item, &evidence[index])
claim, err := s.acquireAdmittedLayerRPC(
c, item, &evidence[index], decodeOptions[reservationIndex], expansionBudgets[reservationIndex],
)
if errors.Is(err, errLayerRPCGZIPCapacity) {
s.rpcRewrap.release(candidate)
c.metrics.InboundRPCDropped(candidate.method, "materialization_capacity")
flightCapacity = true
item.kind = inboundItemCapacityError
continue
}
if errors.Is(err, ErrRPCResultFlightCapacity) {
s.rpcRewrap.release(candidate)
c.metrics.InboundRPCDropped(candidate.method, "flight_capacity")
@ -325,6 +492,10 @@ func (s *Server) prepareInboundLayerRPCBatch(ctx context.Context, c *Conn, plan
}
switch claim.state {
case rpcResultAcquireCompleted:
// Transfer the materialization ticket to the plan before any
// profile-restore step can fail; plan.close is the universal abort
// path and sendReplayed releases it after outbound-budget handoff.
item.payload = claim.encoded
after, prepareErr := s.prepareAdmittedLayerRPCReplay(ctx, c, item.msgID, claim.admissionSeq, item.profileEvidenceFresh(), item.admitted)
if prepareErr != nil {
s.rpcRewrap.release(candidate)
@ -332,10 +503,16 @@ func (s *Server) prepareInboundLayerRPCBatch(ctx context.Context, c *Conn, plan
}
s.rpcRewrap.commit(candidate)
item.kind = inboundItemReplayRPC
item.payload = claim.encoded
if claim.executionKnown && claim.executionOK {
item.replayAfterSuccessfulDelivery = after
}
case rpcResultAcquireAcknowledged:
// Explicit ACK is terminal proof for this request msg_id. Keep
// exact admission metadata, retire the old rewrap candidate and
// make the duplicate ACK-only without replay side effects.
s.rpcRewrap.commit(candidate)
item.kind = inboundItemDuplicate
item.payload = nil
case rpcResultAcquirePending:
after, prepareErr := s.prepareAdmittedLayerRPCReplay(ctx, c, item.msgID, claim.admissionSeq, item.profileEvidenceFresh(), item.admitted)
if prepareErr != nil {
@ -387,7 +564,15 @@ func (s *Server) prepareInboundLayerRPCBatch(ctx context.Context, c *Conn, plan
clearedPostInitCandidates = true
}
claim, err := s.acquireAdmittedLayerRPC(c, item, &evidence[index])
claim, err := s.acquireAdmittedLayerRPC(
c, item, &evidence[index], decodeOptions[reservationIndex], expansionBudgets[reservationIndex],
)
if errors.Is(err, errLayerRPCGZIPCapacity) {
c.metrics.InboundRPCDropped(method, "materialization_capacity")
flightCapacity = true
item.kind = inboundItemCapacityError
continue
}
if errors.Is(err, ErrRPCResultFlightCapacity) {
c.metrics.InboundRPCDropped(method, "flight_capacity")
flightCapacity = true
@ -406,15 +591,18 @@ func (s *Server) prepareInboundLayerRPCBatch(ctx context.Context, c *Conn, plan
}
switch claim.state {
case rpcResultAcquireCompleted:
item.payload = claim.encoded
after, prepareErr := s.prepareAdmittedLayerRPCReplay(ctx, c, item.msgID, claim.admissionSeq, item.profileEvidenceFresh(), item.admitted)
if prepareErr != nil {
return prepareErr
}
item.kind = inboundItemReplayRPC
item.payload = claim.encoded
if claim.executionKnown && claim.executionOK {
item.replayAfterSuccessfulDelivery = after
}
case rpcResultAcquireAcknowledged:
item.kind = inboundItemDuplicate
item.payload = nil
case rpcResultAcquirePending:
if ownersInPlan[item.msgID] != nil {
item.kind = inboundItemDuplicate
@ -624,6 +812,8 @@ func (s *Server) acquireAdmittedLayerRPC(
c *Conn,
item *inboundItem,
evidence *layerRPCProfileEvidence,
options tlprofile.AdmissionOptions,
expansionBudget *layerRPCGZIPExpansionBudget,
) (rpcResultAcquire, error) {
if s == nil || c == nil || item == nil {
return rpcResultAcquire{}, ErrRPCResultFlightInvalid
@ -652,9 +842,17 @@ func (s *Server) acquireAdmittedLayerRPC(
// mutation, not an inherited-default race.
return rpcResultAcquire{}, err
}
redecoded, method, decodeErr := s.decodeInboundLayerRPC(
// Drop the losing typed graph before materializing its authoritative-profile
// replacement. The grown reservation therefore needs to cover the larger
// graph, not two simultaneous copies.
item.admitted = tlprofile.Admission{}
if expansionBudget != nil {
expansionBudget.beginAttempt()
}
redecoded, method, decodeErr := s.decodeInboundLayerRPCWithOptions(
LayerProfileSnapshot{Profile: winnerProfile, Origin: LayerProfileExplicit},
item.body,
options,
)
if decodeErr != nil {
return rpcResultAcquire{}, decodeErr
@ -786,6 +984,14 @@ func layerRPCAdmissionHasExplicitSelector(body []byte, admissionErr error) bool
// evidence only after the full request identity has acquired an owner (or a
// genuine new-msg_id rewrap alias).
func (s *Server) decodeInboundLayerRPC(state LayerProfileSnapshot, body []byte) (tlprofile.Admission, string, error) {
return s.decodeInboundLayerRPCWithLimits(state, body, inboundLayerDecodeLimits)
}
func (s *Server) decodeInboundLayerRPCWithLimits(state LayerProfileSnapshot, body []byte, limits tlprofile.Limits) (tlprofile.Admission, string, error) {
return s.decodeInboundLayerRPCWithOptions(state, body, tlprofile.AdmissionOptions{Limits: limits})
}
func (s *Server) decodeInboundLayerRPCWithOptions(state LayerProfileSnapshot, body []byte, options tlprofile.AdmissionOptions) (tlprofile.Admission, string, error) {
if s == nil || s.layerRPC == nil || len(body) < bin.Word {
return tlprofile.Admission{}, "unknown", fmt.Errorf("invalid exact RPC admission input")
}
@ -795,15 +1001,30 @@ func (s *Server) decodeInboundLayerRPC(state LayerProfileSnapshot, body []byte)
err error
)
if state.Origin != LayerProfileUnknown {
if admitter, ok := s.layerRPC.(LayerRPCDefaultProfileAdmitter); ok {
request, err = admitter.AdmitDefaultLayer(state.Profile, b, inboundLayerDecodeLimits)
if admitter, ok := s.layerRPC.(LayerRPCDefaultProfileOptionsAdmitter); ok {
request, err = admitter.AdmitDefaultLayerWithOptions(state.Profile, b, options)
} else if admitter, ok := s.layerRPC.(LayerRPCDefaultProfileAdmitter); ok {
// Preserve the stable pre-options capability first: unlike strict
// exact admission, default admission lets an explicit invokeWithLayer
// correct inherited or restored profile evidence.
request, err = admitter.AdmitDefaultLayer(state.Profile, b, options.Limits)
} else if admitter, ok := s.layerRPC.(LayerRPCOptionsAdmitter); ok {
request, err = admitter.AdmitLayerWithOptions(state.Profile, b, options)
} else {
// Compatibility fallback for old package tests/mocks. Production Router
// implements default admission so explicit invokeWithLayer can correct.
request, err = s.layerRPC.AdmitLayer(state.Profile, b, inboundLayerDecodeLimits)
request, err = s.layerRPC.AdmitLayer(state.Profile, b, options.Limits)
}
} else if admitter, ok := s.layerRPC.(LayerRPCOptionsAdmitter); ok {
request, err = admitter.AdmitUnprofiledWithOptions(b, options)
} else {
request, err = s.layerRPC.AdmitUnprofiled(b, inboundLayerDecodeLimits)
request, err = s.layerRPC.AdmitUnprofiled(b, options.Limits)
}
if err != nil && options.ExpandGZIP != nil && errors.Is(err, tlprofile.ErrGZIPExpanderMissing) {
// A handler compiled against the stable Limits-only boundary remains
// valid for plain requests. Encountering gzip_packed without the optional
// capability is instead a server wiring/programming error: returning the
// generated error as INPUT_REQUEST_INVALID would silently blame a valid
// client envelope and make recovery impossible.
err = fmt.Errorf("%w: handler cannot use caller-owned bounded gzip expansion: %w", errLayerRPCAdmissionCapability, err)
}
method := "unknown"
if err == nil {
@ -847,6 +1068,10 @@ func (s *Server) decodeInboundLayerRPC(state LayerProfileSnapshot, body []byte)
method = s.typeName(codecErr.WireID)
}
}
var unknownTerminal *tlprofile.UnknownTerminalError
if errors.As(err, &unknownTerminal) && unknownTerminal.WireID != 0 {
method = s.typeName(unknownTerminal.WireID)
}
if errors.Is(err, tlprofile.ErrUnknownRPCMethod) && s.log != nil {
if terminal, recognized := wrappedDestroyAuthKeyTerminal(err); recognized {
method = "destroy_auth_key"
@ -867,7 +1092,7 @@ func (s *Server) decodeInboundLayerRPC(state LayerProfileSnapshot, body []byte)
// commitLayerProfileEvidence publishes one generated invokeWithLayer proof.
// The exact-session registry is the cross-physical-connection linearization
// point; the Conn cursor then prevents a concurrent older admission from
// overwriting its local wire epoch. Older cached duplicates remain decodable
// overwriting its local wire epoch. Older retained duplicates remain decodable
// and request-bound, but cannot mutate session/profile state.
func (s *Server) commitLayerProfileEvidence(ctx context.Context, c *Conn, profile tlprofile.Profile, msgID int64) (bool, error) {
if s == nil || c == nil {
@ -994,6 +1219,63 @@ func layerRPCAdmissionError(err error) *mt.RPCError {
return &mt.RPCError{ErrorCode: 400, ErrorMessage: "INPUT_REQUEST_INVALID"}
}
// logLayerRPCAdmissionRejection preserves one production-visible diagnostic for
// an otherwise generic RPC 400 without exposing the request body. The INFO path
// is one-shot per physical Conn; compatibility-traced unknown RPCs already have
// their own warning and therefore do not consume this diagnostic slot.
func (s *Server) logLayerRPCAdmissionRejection(
c *Conn,
item *inboundItem,
state LayerProfileSnapshot,
cursor layerRPCAdmissionCursor,
method string,
rpcError *mt.RPCError,
admissionErr error,
) {
if s == nil || s.log == nil || c == nil || item == nil || rpcError == nil {
return
}
wireID := item.typeID
if wireID == 0 {
if id, err := (&bin.Buffer{Buf: item.body}).PeekID(); err == nil {
wireID = id
}
}
fields := []zap.Field{
zap.String("method", method),
zap.String("auth_key_id", c.authKeyHex),
zap.Int64("session_id", c.sessionID),
zap.Int64("msg_id", item.msgID),
zap.Uint32("top_level_wire_id", wireID),
zap.Int("wire_bytes", len(item.body)),
zap.Int("selected_profile", int(state.Profile)),
zap.String("profile_origin", layerProfileOriginLogName(state.Origin)),
zap.Uint32("profile_epoch", state.Epoch),
zap.Int("raw_layer_evidence", cursor.rawLayer),
zap.Int64("layer_evidence_msg_id", cursor.evidenceMsgID),
zap.Bool("explicit_layer_selector", layerRPCAdmissionHasExplicitSelector(item.body, admissionErr)),
zap.Int("rpc_error_code", rpcError.ErrorCode),
zap.String("rpc_error_message", rpcError.ErrorMessage),
zap.Error(admissionErr),
}
if !errors.Is(admissionErr, tlprofile.ErrUnknownRPCMethod) && c.layerRPCAdmissionTraceLogged.CompareAndSwap(false, true) {
s.log.Info("RPC exact admission rejected", fields...)
return
}
s.log.Debug("RPC exact admission rejected", fields...)
}
func layerProfileOriginLogName(origin LayerProfileOrigin) string {
switch origin {
case LayerProfileInherited:
return "inherited"
case LayerProfileExplicit:
return "explicit"
default:
return "unknown"
}
}
func (s *Server) layerRPCDependencies(c *Conn, msgID int64, request tlprofile.Admission) layerRPCDependencySet {
result := layerRPCDependencySet{}
seen := make(map[int64]struct{})

294
internal/mtprotoedge/inbound_layer_rpc_test.go
View file

@ -1,6 +1,7 @@
package mtprotoedge
import (
"bytes"
"context"
"errors"
"fmt"
@ -46,6 +47,14 @@ type admissionOnlyLayerRPC struct {
published []publishedLayerEvidence
}
// legacyAdmissionOnlyLayerRPC intentionally exposes only the original
// Limits-based admission interfaces. It guards source and runtime compatibility
// for implementations compiled before caller-owned AdmissionOptions existed.
type legacyAdmissionOnlyLayerRPC struct {
dispatcher *tlprofile.Dispatcher
lastRemaining int
}
type orderedAdmissionOnlyLayerRPC struct {
*admissionOnlyLayerRPC
exactMu sync.Mutex
@ -221,6 +230,10 @@ func newAdmissionOnlyLayerRPC() *admissionOnlyLayerRPC {
return &admissionOnlyLayerRPC{dispatcher: tlprofile.NewDispatcher()}
}
func newLegacyAdmissionOnlyLayerRPC() *legacyAdmissionOnlyLayerRPC {
return &legacyAdmissionOnlyLayerRPC{dispatcher: tlprofile.NewDispatcher()}
}
func newOrderedAdmissionOnlyLayerRPC() *orderedAdmissionOnlyLayerRPC {
return &orderedAdmissionOnlyLayerRPC{
admissionOnlyLayerRPC: newAdmissionOnlyLayerRPC(),
@ -265,14 +278,48 @@ func (h *admissionOnlyLayerRPC) AdmitLayer(profile tlprofile.Profile, b *bin.Buf
return h.dispatcher.Admit(profile, b, limits)
}
func (h *admissionOnlyLayerRPC) AdmitLayerWithOptions(profile tlprofile.Profile, b *bin.Buffer, options tlprofile.AdmissionOptions) (tlprofile.Admission, error) {
return h.dispatcher.AdmitWithOptions(profile, b, options)
}
func (h *admissionOnlyLayerRPC) AdmitDefaultLayer(profile tlprofile.Profile, b *bin.Buffer, limits tlprofile.Limits) (tlprofile.Admission, error) {
return h.dispatcher.AdmitDefault(profile, b, limits)
}
func (h *admissionOnlyLayerRPC) AdmitDefaultLayerWithOptions(profile tlprofile.Profile, b *bin.Buffer, options tlprofile.AdmissionOptions) (tlprofile.Admission, error) {
return h.dispatcher.AdmitDefaultWithOptions(profile, b, options)
}
func (h *admissionOnlyLayerRPC) AdmitUnprofiled(b *bin.Buffer, limits tlprofile.Limits) (tlprofile.Admission, error) {
return h.dispatcher.AdmitUnprofiled(b, limits)
}
func (h *admissionOnlyLayerRPC) AdmitUnprofiledWithOptions(b *bin.Buffer, options tlprofile.AdmissionOptions) (tlprofile.Admission, error) {
return h.dispatcher.AdmitUnprofiledWithOptions(b, options)
}
func (h *legacyAdmissionOnlyLayerRPC) AdmitLayer(profile tlprofile.Profile, b *bin.Buffer, limits tlprofile.Limits) (tlprofile.Admission, error) {
request, err := h.dispatcher.Admit(profile, b, limits)
h.lastRemaining = b.Len()
return request, err
}
func (h *legacyAdmissionOnlyLayerRPC) AdmitDefaultLayer(profile tlprofile.Profile, b *bin.Buffer, limits tlprofile.Limits) (tlprofile.Admission, error) {
request, err := h.dispatcher.AdmitDefault(profile, b, limits)
h.lastRemaining = b.Len()
return request, err
}
func (h *legacyAdmissionOnlyLayerRPC) AdmitUnprofiled(b *bin.Buffer, limits tlprofile.Limits) (tlprofile.Admission, error) {
request, err := h.dispatcher.AdmitUnprofiled(b, limits)
h.lastRemaining = b.Len()
return request, err
}
func (*legacyAdmissionOnlyLayerRPC) DispatchAdmitted(context.Context, [8]byte, int64, int64, uint64, tlprofile.Admission) (tlprofile.Result, string, error) {
return nil, "", fmt.Errorf("admission-only handler")
}
func (*admissionOnlyLayerRPC) DispatchAdmitted(context.Context, [8]byte, int64, int64, uint64, tlprofile.Admission) (tlprofile.Result, string, error) {
return nil, "", fmt.Errorf("admission-only handler")
}
@ -301,6 +348,99 @@ func (h *admissionOnlyLayerRPC) publications() []publishedLayerEvidence {
return append([]publishedLayerEvidence(nil), h.published...)
}
func TestLegacyLayerRPCAdmissionInterfacesRemainCompatible(t *testing.T) {
t.Run("inherited default accepts explicit correction", func(t *testing.T) {
handler := newLegacyAdmissionOnlyLayerRPC()
s := New(Options{DC: 2, LayerRPC: handler})
body := exactLayerRPCBody(t, &tg.InvokeWithLayerRequest{
Layer: 225,
Query: &tg.HelpGetConfigRequest{},
})
request, method, err := s.decodeInboundLayerRPCWithOptions(
LayerProfileSnapshot{Profile: tlprofile.Profile228, Origin: LayerProfileInherited},
body,
tlprofile.AdmissionOptions{
Limits: inboundLayerDecodeLimits,
ExpandGZIP: func([]byte, int) ([]byte, func(), error) {
t.Fatal("plain legacy admission unexpectedly requested gzip expansion")
return nil, nil, nil
},
},
)
if err != nil {
t.Fatalf("legacy default admission: %v", err)
}
if method != "help.getConfig" {
t.Fatalf("method = %q, want help.getConfig", method)
}
if request.Call().Profile() != tlprofile.Profile225 {
t.Fatalf("call profile = %d, want 225", request.Call().Profile())
}
if profile, ok := request.ProfileEvidence(); !ok || profile != tlprofile.Profile225 {
t.Fatalf("profile evidence = %d/%v, want 225/true", profile, ok)
}
if handler.lastRemaining != 0 {
t.Fatalf("successful legacy admission left %d bytes", handler.lastRemaining)
}
})
t.Run("unprofiled selector remains supported", func(t *testing.T) {
handler := newLegacyAdmissionOnlyLayerRPC()
s := New(Options{DC: 2, LayerRPC: handler})
body := exactLayerRPCBody(t, &tg.InvokeWithLayerRequest{
Layer: 225,
Query: &tg.HelpGetNearestDCRequest{},
})
request, method, err := s.decodeInboundLayerRPCWithOptions(
LayerProfileSnapshot{},
body,
tlprofile.AdmissionOptions{Limits: inboundLayerDecodeLimits},
)
if err != nil {
t.Fatalf("legacy unprofiled admission: %v", err)
}
if method != "help.getNearestDc" {
t.Fatalf("method = %q, want help.getNearestDc", method)
}
if request.Call().Profile() != tlprofile.Profile225 {
t.Fatalf("call profile = %d, want 225", request.Call().Profile())
}
})
t.Run("nested gzip requires explicit optional capability", func(t *testing.T) {
handler := newLegacyAdmissionOnlyLayerRPC()
s := New(Options{DC: 2, LayerRPC: handler})
body, _ := tdlibNestedGZIPBody(t, tlprofile.Profile228, &tg.HelpGetConfigRequest{})
original := append([]byte(nil), body...)
_, _, err := s.decodeInboundLayerRPCWithOptions(
LayerProfileSnapshot{},
body,
tlprofile.AdmissionOptions{
Limits: inboundLayerDecodeLimits,
ExpandGZIP: func([]byte, int) ([]byte, func(), error) {
t.Fatal("legacy handler unexpectedly received options-only gzip expander")
return nil, nil, nil
},
},
)
if !errors.Is(err, errLayerRPCAdmissionCapability) {
t.Fatalf("nested gzip error = %v, want admission capability error", err)
}
if !errors.Is(err, tlprofile.ErrGZIPExpanderMissing) {
t.Fatalf("nested gzip error = %v, want generated missing-expander cause", err)
}
if handler.lastRemaining != len(body) {
t.Fatalf("failed legacy admission retained %d/%d input bytes", handler.lastRemaining, len(body))
}
if !bytes.Equal(body, original) {
t.Fatal("failed legacy admission mutated caller wire bytes")
}
})
}
func TestNestedExplicitLayerAdmissionErrorsAreNotDefaultFailures(t *testing.T) {
handler := newAdmissionOnlyLayerRPC()
s := New(Options{DC: 2, LayerRPC: handler})
@ -401,7 +541,7 @@ func TestBatchProvisionalCursorKeepsRegistryWatermarkAcrossOldReplay(t *testing.
t.Fatalf("old owner err=%v", err)
}
oldClaim.owner.CompleteExecution(true)
s.rpcResults.Put(authKeyID, sessionID, 100, &encodedOutboundMessage{body: []byte{1}, reqMsgID: 100})
storeLogicalRPCResultForTest(t, s, c, 100, &encodedOutboundMessage{body: []byte{1}, reqMsgID: 100})
nakedBody := exactOutboundLayerRPCBody(t, tlprofile.Profile227, &tg.MessagesGetHistoryRequest{
Peer: &tg.InputPeerSelf{}, Limit: 1,
@ -422,6 +562,131 @@ func TestBatchProvisionalCursorKeepsRegistryWatermarkAcrossOldReplay(t *testing.
}
}
func TestAcknowledgedExactRPCReleasesReplayReceipt(t *testing.T) {
handler := &replayProfileCaptureLayerRPC{admissionOnlyLayerRPC: newAdmissionOnlyLayerRPC()}
s := New(Options{DC: 2, LayerRPC: handler})
authKeyID := [8]byte{0x31, 0xa1}
const sessionID, reqMsgID = int64(3191), int64(100)
body := exactLayerRPCBody(t, &tg.InvokeWithLayerRequest{
Layer: 225, Query: &tg.HelpGetConfigRequest{},
})
admitted, _, err := s.decodeInboundLayerRPC(
LayerProfileSnapshot{Profile: tlprofile.Profile225, Origin: LayerProfileExplicit}, body,
)
if err != nil {
t.Fatal(err)
}
owner, err := s.rpcResults.AcquireLayerIdentified(
authKeyID, sessionID, reqMsgID,
admitted.Call().Profile(), admitted.Prepared().Identity(),
)
if err != nil || owner.state != rpcResultAcquireOwner {
t.Fatalf("owner = %#v, err=%v", owner, err)
}
owner.owner.CompleteExecution(true)
logical := &Conn{authKeyID: authKeyID, sessionID: sessionID}
storeLogicalRPCResultForTest(t, s, logical, reqMsgID, &encodedOutboundMessage{
body: make([]byte, 32), reqMsgID: reqMsgID,
})
acknowledgeLogicalRPCResultForTest(t, s, logical, reqMsgID)
replacement := &Conn{authKeyID: authKeyID, sessionID: sessionID, metrics: NopMetrics{}}
replacement.startInboundRPCScheduler(s.rpcScheduler, 1, 8, time.Second)
defer replacement.Close()
plan := &inboundPlan{items: []inboundItem{{kind: inboundItemRPC, msgID: reqMsgID, body: body}}}
defer plan.close()
if err := s.prepareInboundLayerRPCBatch(context.Background(), replacement, plan); err != nil {
t.Fatal(err)
}
if plan.items[0].kind != inboundItemRPC || plan.items[0].payload == nil {
t.Fatalf("post-ACK request = kind:%d payload:%T", plan.items[0].kind, plan.items[0].payload)
}
if len(plan.rpcTasks) != 1 || len(plan.rpcOwners) != 1 || len(plan.rewrapAliases) != 0 || plan.rpcReservation == nil {
t.Fatalf("post-ACK request scheduling: tasks=%d owners=%d aliases=%d reservation=%v",
len(plan.rpcTasks), len(plan.rpcOwners), len(plan.rewrapAliases), plan.rpcReservation != nil)
}
if profiles, known := handler.capturedProfiles(); len(profiles) != 0 || len(known) != 0 {
t.Fatalf("ACKed duplicate ran replay side effects: profiles=%v known=%v", profiles, known)
}
claim, err := s.rpcResults.AcquireLayerIdentified(
authKeyID, sessionID, reqMsgID,
admitted.Call().Profile(), admitted.Prepared().Identity(),
)
if err != nil || claim.state != rpcResultAcquirePending || claim.admissionSeq == owner.admissionSeq {
t.Fatalf("post-ACK claim after preflight = %#v, err=%v", claim, err)
}
}
func TestAcknowledgedExactInitRewrapReleasesReplayReceipt(t *testing.T) {
handler := &replayProfileCaptureLayerRPC{admissionOnlyLayerRPC: newAdmissionOnlyLayerRPC()}
s := New(Options{DC: 2, LayerRPC: handler})
authKeyID := [8]byte{0x31, 0xa2}
const sessionID, oldReqID, newReqID = int64(3192), int64(100), int64(104)
c := &Conn{authKeyID: authKeyID, sessionID: sessionID, metrics: NopMetrics{}}
if err := c.SeedInheritedLayerProfile(tlprofile.Profile227); err != nil {
t.Fatal(err)
}
c.startInboundRPCScheduler(s.rpcScheduler, 1, 8, time.Second)
defer c.Close()
inner := exactOutboundLayerRPCBody(t, tlprofile.Profile227, &tg.HelpGetConfigRequest{})
oldPlan := &inboundPlan{items: []inboundItem{{kind: inboundItemRPC, msgID: oldReqID, body: inner}}}
defer oldPlan.close()
if err := s.prepareInboundLayerRPCBatch(context.Background(), c, oldPlan); err != nil {
t.Fatal(err)
}
if len(oldPlan.rpcOwners) != 1 || s.rpcRewrap.total != 1 {
t.Fatalf("old exact candidate = owners:%d candidates:%d", len(oldPlan.rpcOwners), s.rpcRewrap.total)
}
wrapped := exactLayerRPCBody(t, &tg.InvokeWithLayerRequest{
Layer: 227,
Query: &tg.InitConnectionRequest{
APIID: 6, DeviceModel: "Pixel", SystemVersion: "SDK 36", AppVersion: "12.8.1",
SystemLangCode: "en", LangPack: "android", LangCode: "en",
Query: &tg.HelpGetConfigRequest{},
},
})
admitted, _, err := s.decodeInboundLayerRPC(
LayerProfileSnapshot{Profile: tlprofile.Profile227, Origin: LayerProfileExplicit}, wrapped,
)
if err != nil {
t.Fatal(err)
}
newOwner, err := s.rpcResults.AcquireLayerIdentified(
authKeyID, sessionID, newReqID,
admitted.Call().Profile(), admitted.Prepared().Identity(),
)
if err != nil || newOwner.state != rpcResultAcquireOwner {
t.Fatalf("new exact receipt owner = %#v, err=%v", newOwner, err)
}
newOwner.owner.CompleteExecution(true)
storeLogicalRPCResultForTest(t, s, c, newReqID, &encodedOutboundMessage{
body: make([]byte, 32), reqMsgID: newReqID,
})
acknowledgeLogicalRPCResultForTest(t, s, c, newReqID)
// In the real delivery path a successful rewrap commits its source candidate.
// Clear the synthetic pending candidate before observing post-ACK admission.
s.rpcRewrap.clearSession(c)
newPlan := &inboundPlan{items: []inboundItem{{kind: inboundItemRPC, msgID: newReqID, body: wrapped}}}
defer newPlan.close()
if err := s.prepareInboundLayerRPCBatch(context.Background(), c, newPlan); err != nil {
t.Fatal(err)
}
if newPlan.items[0].kind != inboundItemRPC || newPlan.items[0].payload == nil ||
len(newPlan.rpcTasks) != 1 || len(newPlan.rpcOwners) != 1 || len(newPlan.rewrapAliases) != 0 {
t.Fatalf("post-ACK init request scheduling: kind=%d tasks=%d owners=%d aliases=%d",
newPlan.items[0].kind, len(newPlan.rpcTasks), len(newPlan.rpcOwners), len(newPlan.rewrapAliases))
}
if s.rpcRewrap.total != 0 {
t.Fatalf("ACKed exact init rewrap left %d candidates", s.rpcRewrap.total)
}
if profiles, known := handler.capturedProfiles(); len(profiles) != 0 || len(known) != 0 {
t.Fatalf("ACKed exact init rewrap ran replay side effects: profiles=%v known=%v", profiles, known)
}
}
func TestBatchProvisionalCursorUsesPendingNewerExplicitEvidence(t *testing.T) {
handler := newAdmissionOnlyLayerRPC()
s := New(Options{DC: 2, LayerRPC: handler})
@ -701,7 +966,7 @@ func TestDurabilityOutageInitializesOnlyCurrentConnection(t *testing.T) {
}
time.Sleep(time.Millisecond)
}
if _, ok := s.rpcResults.Get(c.authKeyID, c.sessionID, msgID); !ok {
if _, ok := s.rpcResults.Replay(c.authKeyID, c.sessionID, msgID); !ok {
t.Fatal("successful outage-local init did not publish its exact RPC result")
}
if layer, found, err := router.ResolveInheritedAuthKeyLayer(ctx, c.authKeyID); err != nil || found || layer != 0 {
@ -781,7 +1046,7 @@ func TestInvariantReplayNeverCachesInternalCanonicalProfile(t *testing.T) {
}
owner.CompleteExecution(true)
s.rpcResults.Put(authKeyID, sessionID, 100, &encodedOutboundMessage{body: []byte{1}, reqMsgID: 100})
storeLogicalRPCResultForTest(t, s, original, 100, &encodedOutboundMessage{body: []byte{1}, reqMsgID: 100})
if profile, ok := s.rpcResults.ExactAdmissionProfile(authKeyID, sessionID, 100); ok || profile != 0 {
t.Fatalf("completed invariant cached profile=(%d,%v)", profile, ok)
}
@ -806,7 +1071,7 @@ func TestInvariantReplayNeverCachesInternalCanonicalProfile(t *testing.T) {
func TestSameMsgIDNakedReplayUsesWinnerAdmissionProfile(t *testing.T) {
handler := newAdmissionOnlyLayerRPC()
s := New(Options{DC: 2, LayerRPC: handler})
s.rpcResults = newRPCResultCacheWithFlightLimit(time.Now, 8)
s.rpcResults = newRPCExecutionLedgerForServerTest(s, time.Now, 8)
authKeyID := [8]byte{0x22, 0x99}
const sessionID = int64(2299)
body := exactOutboundLayerRPCBody(t, tlprofile.Profile225, &tg.MessagesGetHistoryRequest{
@ -833,11 +1098,12 @@ func TestSameMsgIDNakedReplayUsesWinnerAdmissionProfile(t *testing.T) {
}
c220 := &Conn{authKeyID: authKeyID, sessionID: sessionID}
c227 := &Conn{authKeyID: authKeyID, sessionID: sessionID}
winner, err := s.acquireAdmittedLayerRPC(c220, &item220, nil)
options := tlprofile.AdmissionOptions{Limits: inboundLayerDecodeLimits}
winner, err := s.acquireAdmittedLayerRPC(c220, &item220, nil, options, nil)
if err != nil || winner.state != rpcResultAcquireOwner || winner.owner == nil {
t.Fatalf("winner = state:%d err:%v", winner.state, err)
}
loser, err := s.acquireAdmittedLayerRPC(c227, &item227, nil)
loser, err := s.acquireAdmittedLayerRPC(c227, &item227, nil, options, nil)
if err != nil || loser.state != rpcResultAcquirePending || loser.admissionSeq != winner.admissionSeq {
t.Fatalf("loser join = state:%d seq:%d err:%v, winner seq:%d", loser.state, loser.admissionSeq, err, winner.admissionSeq)
}
@ -852,7 +1118,7 @@ func TestSameMsgIDNakedReplayUsesWinnerAdmissionProfile(t *testing.T) {
if err != nil {
t.Fatal(err)
}
if _, err := s.acquireAdmittedLayerRPC(c220, &changed, nil); !errors.Is(err, ErrRPCResultIdentityMismatch) {
if _, err := s.acquireAdmittedLayerRPC(c220, &changed, nil, options, nil); !errors.Is(err, ErrRPCResultIdentityMismatch) {
t.Fatalf("same-msg_id changed body err=%v, want identity mismatch", err)
}
winner.owner.Abort()
@ -935,7 +1201,7 @@ func TestLayerEvidencePublicationBelongsOnlyToFreshFlightOwner(t *testing.T) {
func TestLayerEvidenceNotPublishedWhenBatchFlightCapacityRollsBack(t *testing.T) {
handler := newAdmissionOnlyLayerRPC()
s := New(Options{DC: 2, LayerRPC: handler})
s.rpcResults = newRPCResultCacheWithFlightLimit(time.Now, 1)
s.rpcResults = newRPCExecutionLedgerForServerTest(s, time.Now, 1)
c := &Conn{authKeyID: [8]byte{0x22, 0x97}, sessionID: 2297, metrics: NopMetrics{}}
c.startInboundRPCScheduler(s.rpcScheduler, 1, 8, time.Second)
defer c.Close()
@ -980,7 +1246,7 @@ func TestOldCompletedLayerRequestCannotRollBackCorrectedSession(t *testing.T) {
t.Fatal("old request did not acquire owner")
}
oldOwner.CompleteExecution(true)
s.rpcResults.Put(authKeyID, sessionID, 100, &encodedOutboundMessage{body: []byte{1}, reqMsgID: 100})
storeLogicalRPCResultForTest(t, s, c, 100, &encodedOutboundMessage{body: []byte{1}, reqMsgID: 100})
correctPlan := &inboundPlan{items: []inboundItem{{
kind: inboundItemRPC, msgID: 104,
@ -1019,7 +1285,7 @@ func TestLogicalSessionLayerWatermarkSurvivesResultExpiryAndOldContainer(t *test
now := newExpiryTestClock(time.Unix(1_900_000_000, 0))
router := rpc.New(rpc.Config{DC: 2}, rpc.Deps{}, zaptest.NewLogger(t), now)
s := New(Options{DC: 2, LayerRPC: router, Clock: now})
s.rpcResults = newRPCResultCacheWithFlightLimit(now.Now, 8)
s.rpcResults = newRPCExecutionLedgerForServerTest(s, now.Now, 8)
authKeyID := [8]byte{0x22, 0x95}
const sessionID = int64(2295)
newConn := func() *Conn {
@ -1046,7 +1312,7 @@ func TestLogicalSessionLayerWatermarkSurvivesResultExpiryAndOldContainer(t *test
if oldOwner == nil || !oldOwner.CompleteExecution(true) {
t.Fatal("old Layer 225 request did not establish a completed owner")
}
s.rpcResults.Put(authKeyID, sessionID, oldMsgID, &encodedOutboundMessage{body: []byte{1}, reqMsgID: oldMsgID})
storeLogicalRPCResultForTest(t, s, original, oldMsgID, &encodedOutboundMessage{body: []byte{1}, reqMsgID: oldMsgID})
correctPlan := &inboundPlan{items: []inboundItem{{
kind: inboundItemRPC, msgID: correctedMsgID,
@ -1065,7 +1331,7 @@ func TestLogicalSessionLayerWatermarkSurvivesResultExpiryAndOldContainer(t *test
// inherited default remains Layer 227, while an old inner request may execute
// request-bound but cannot recreate or roll back the exact-session watermark.
now.Advance(10 * time.Minute)
if _, ok := s.rpcResults.Get(authKeyID, sessionID, oldMsgID); ok {
if _, ok := s.rpcResults.Replay(authKeyID, sessionID, oldMsgID); ok {
t.Fatal("old completed result did not expire")
}
@ -1712,7 +1978,7 @@ func TestUnprofiledInvariantBindKeepsProfileUnknownAndReturnsExactBool(t *testin
func TestLayerRPCBatchCapacityKeepsExistingPendingReplay(t *testing.T) {
router := rpc.New(rpc.Config{DC: 2, IP: "127.0.0.1", Port: 2398}, rpc.Deps{}, zaptest.NewLogger(t), clock.System)
s := New(Options{DC: 2, LayerRPC: router})
s.rpcResults = newRPCResultCacheWithFlightLimit(time.Now, 1)
s.rpcResults = newRPCExecutionLedgerForServerTest(s, time.Now, 1)
c := &Conn{
authKeyID: [8]byte{7, 7, 1},
sessionID: 771,
@ -1767,7 +2033,7 @@ func TestLayerRPCBatchCapacityKeepsExistingPendingReplay(t *testing.T) {
}
func TestLayerRPCBatchCapacityAbortsRejectedRewrapOwner(t *testing.T) {
cache := newRPCResultCacheWithFlightLimit(time.Now, 1)
cache := newRPCExecutionLedgerForTest(time.Now, 1)
authKeyID := [8]byte{7, 7, 9}
const (
sessionID = int64(779)

39
internal/mtprotoedge/inbound_preflight.go
View file

@ -41,7 +41,7 @@ const (
// It never dispatches business code a second time.
inboundItemRewrappedRPC
// inboundItemReplayRPC is a request first observed by this physical Conn whose
// terminal result already exists in the cross-connection cache. It is distinct
// terminal result already exists in the cross-connection execution ledger. It is distinct
// from inboundItemDuplicate: a duplicate already present in this Conn's seen
// table must only be ACKed. The original owner/result is already using the same
// reliable TCP stream, so replaying it once per retransmit wave amplifies a
@ -92,6 +92,11 @@ type inboundPlan struct {
ackIDs []int64
logicalMin int64
releases []func()
// gzipExpandedBytes is the non-refundable per-frame decompression work
// already performed by outer and exact-layer nested gzip envelopes. Memory
// reservations are released when their buffers die, but this cumulative
// counter prevents sibling RPCs from recycling the same CPU budget.
gzipExpandedBytes int
rpcReservation *inboundRPCBatchReservation
rpcTasks []inboundRPC
@ -314,6 +319,7 @@ func (s *Server) preflightInbound(cs *connState, msgID int64, seqNo int32, body
plan.close()
return nil, err
}
plan.gzipExpandedBytes = budget.expanded
plan.staged = overlay.staged
if plan.logicalMin == 0 {
plan.close()
@ -448,7 +454,7 @@ func (s *Server) walkInbound(
return &dispatchBadMsgError{msgID: msgID, seqNo: seqNo, code: code}
}
content := clientMessageNeedsAck(typeID)
content := clientMessageIsContentRelated(typeID, seqNo)
if record, seen := overlay.seenRecord(msgID); seen {
if record.seqNo != seqNo || record.content != content {
return &dispatchBadMsgError{msgID: msgID, seqNo: seqNo, code: badMsgContainer}
@ -726,8 +732,8 @@ func preflightInboundItem(msgID int64, seqNo int32, typeID uint32, content bool,
// prepareInboundRPCBatch performs the whole container's count/byte admission
// before copying or scheduling any API RPC. Capacity exhaustion is converted
// into one consistent terminal FLOOD_WAIT result per uncached RPC; no business
// handler from the batch is allowed to start in that case.
// into one consistent local 500 WORKER_BUSY_TOO_LONG_RETRY result per new
// 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)
@ -778,6 +784,13 @@ func (s *Server) prepareInboundRPCBatch(ctx context.Context, c *Conn, plan *inbo
s.rpcRewrap.commit(candidate)
item.kind = inboundItemReplayRPC
item.payload = claim.encoded
case rpcResultAcquireAcknowledged:
// The client already ACKed the correlated rpc_result. Retire the
// stale rewrap candidate and keep this duplicate ACK-only; neither
// the old body nor the business handler may run again.
s.rpcRewrap.commit(candidate)
item.kind = inboundItemDuplicate
item.payload = nil
case rpcResultAcquirePending:
s.rpcRewrap.commit(candidate)
item.kind = inboundItemRewrappedRPC
@ -833,7 +846,7 @@ func (s *Server) prepareInboundRPCBatch(ctx context.Context, c *Conn, plan *inbo
}
switch claim.state {
case rpcResultAcquireCompleted:
s.log.Info("RPC duplicate replay from session cache",
s.log.Info("RPC duplicate replay from logical-session outbox",
zap.String("method", method),
zap.Int64("msg_id", item.msgID),
zap.String("auth_key_id", c.authKeyHex),
@ -841,6 +854,9 @@ func (s *Server) prepareInboundRPCBatch(ctx context.Context, c *Conn, plan *inbo
)
item.kind = inboundItemReplayRPC
item.payload = claim.encoded
case rpcResultAcquireAcknowledged:
item.kind = inboundItemDuplicate
item.payload = nil
case rpcResultAcquirePending:
// A malformed/replayed container may repeat the same msg_id after this
// very plan installed its owner. More generally, any request already in
@ -927,14 +943,14 @@ func (s *Server) executeInboundPlan(ctx context.Context, cs *connState, c *Conn,
case inboundItemServiceDuplicate:
// Classify from the originally committed connState record, never from the
// retransmitted body. This prevents same-id payload replacement. Only an
// already cached answer is eligible for resend (rpc_drop_answer today);
// already retained answer is eligible for resend (rpc_drop_answer today);
// other best-effort service traffic uses a later fresh request.
if err := s.replayRPCResultByRequest(ctx, c, item.msgID); err != nil {
return err
}
case inboundItemReplayRPC:
if encoded, _ := item.payload.(*encodedOutboundMessage); encoded != nil {
if err := s.sendCachedRPCResultWithHook(ctx, c, encoded, item.replayAfterSuccessfulDelivery); err != nil {
if err := s.sendReplayedRPCResultWithHook(ctx, c, encoded, item.replayAfterSuccessfulDelivery); err != nil {
return err
}
} else if err := s.replayRPCResultByRequest(ctx, c, item.msgID); err != nil {
@ -1020,6 +1036,10 @@ func (s *Server) executeInboundPlan(ctx context.Context, cs *connState, c *Conn,
}); err != nil {
return err
}
// The key cannot reconnect to ACK or replay any old answer. Release every
// logical outbox and receipt only after the terminal OK is physically on
// the wire; retaining them for the offline TTL would be pure leakage.
s.conns.ForgetLogicalSessionsForRawAuthKey(c.authKeyID)
c.beginTerminalShutdown()
c.closeTransport()
return nil
@ -1031,10 +1051,7 @@ func (s *Server) executeInboundPlan(ctx context.Context, cs *connState, c *Conn,
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 {
if err := s.sendResult(ctx, c, item.msgID, rpcWorkerBusyError()); err != nil {
return err
}
case inboundItemRPCAdmissionError:

55
internal/mtprotoedge/inbound_rpc.go
View file

@ -4,6 +4,7 @@ import (
"container/list"
"context"
"errors"
"fmt"
"sync"
"sync/atomic"
"time"
@ -544,6 +545,60 @@ func (c *Conn) dropInboundRPCSpecs(specs []inboundRPCSpec, reason string) {
}
}
// growEntry raises one provisional task's materialization charge without a
// release/reacquire window. Exact gzip admission calls this after the expanded
// size is known but before the generated typed decoder can allocate the request
// graph. A failure leaves every existing reservation unchanged so the caller
// can reject and abort the whole container atomically.
func (r *inboundRPCBatchReservation) growEntry(index, targetSize int) error {
if r == nil || index < 0 || index >= len(r.entries) || targetSize < 0 {
return errInboundRPCBatchSelection
}
entry := &r.entries[index]
if targetSize <= entry.size {
return nil
}
if entry.global == nil || entry.global.scheduler == nil || entry.global.released.Load() {
return ErrConnClosed
}
delta := int64(targetSize) - int64(entry.size)
scheduler := entry.global.scheduler
conn := r.conn
// Match initial admission's lock order: global scheduler budget, then the
// connection queue budget. Release paths never hold rpcMu while acquiring
// budgetMu, so this cannot invert task completion or close.
scheduler.budgetMu.Lock()
defer scheduler.budgetMu.Unlock()
select {
case <-scheduler.stopCh:
return ErrConnClosed
default:
}
if delta > scheduler.maxBytes-scheduler.bytes {
return fmt.Errorf("%w: grow exact admission global byte budget by %d", ErrInboundRPCQueueFull, delta)
}
conn.rpcMu.Lock()
defer conn.rpcMu.Unlock()
if err := r.ctx.Err(); err != nil {
return err
}
if conn.rpcClosed || conn.isRetired() {
return ErrConnClosed
}
if delta > int64(maxInflightRPCBytes)-conn.inflightRPCBytes.Load() {
return fmt.Errorf("%w: grow exact admission connection byte budget by %d", ErrInboundRPCQueueFull, delta)
}
scheduler.bytes += delta
entry.global.size += delta
entry.size = targetSize
r.totalSize += delta
conn.inflightRPCBytes.Add(delta)
return nil
}
// 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

74
internal/mtprotoedge/inbound_rpc_batch_test.go
View file

@ -130,6 +130,80 @@ func TestInboundRPCBatchAbortReturnsEveryReservationExactlyOnce(t *testing.T) {
}
}
func TestInboundRPCBatchReservationGrowTransfersBudgetAtomically(t *testing.T) {
scheduler := newInboundRPCScheduler(1, 8, 1<<20)
c := newInboundTestConn(scheduler, 1, 4, time.Second)
defer func() {
c.closeInboundRPCScheduler()
scheduler.stop(time.Second)
}()
reservation, err := c.reserveInboundRPCBatch(context.Background(), []inboundRPCSpec{
{method: "one", size: 3},
{method: "two", size: 5},
})
if err != nil {
t.Fatal(err)
}
if err := reservation.growEntry(0, 11); err != nil {
t.Fatal(err)
}
if tasks, bytes := scheduler.budgetSnapshot(); tasks != 2 || bytes != 16 {
t.Fatalf("grown global budget = %d/%d, want 2/16", tasks, bytes)
}
if got := c.inflightRPCBytes.Load(); got != 16 {
t.Fatalf("grown connection budget = %d, want 16", got)
}
if reservation.entries[0].size != 11 || reservation.totalSize != 16 {
t.Fatalf("grown reservation entry/total = %d/%d", reservation.entries[0].size, reservation.totalSize)
}
reservation.abort()
if tasks, bytes := scheduler.budgetSnapshot(); tasks != 0 || bytes != 0 {
t.Fatalf("grown reservation abort leaked global budget %d/%d", tasks, bytes)
}
if got := c.inflightRPCBytes.Load(); got != 0 {
t.Fatalf("grown reservation abort leaked connection budget %d", got)
}
}
func TestInboundRPCBatchReservationGrowFailureKeepsOriginalBudget(t *testing.T) {
for _, test := range []struct {
name string
globalMax int64
targetSize int
}{
{name: "global", globalMax: 5, targetSize: 6},
{name: "connection", globalMax: int64(maxInflightRPCBytes) * 2, targetSize: maxInflightRPCBytes + 1},
} {
t.Run(test.name, func(t *testing.T) {
scheduler := newInboundRPCScheduler(1, 8, test.globalMax)
c := newInboundTestConn(scheduler, 1, 4, time.Second)
defer func() {
c.closeInboundRPCScheduler()
scheduler.stop(time.Second)
}()
reservation, err := c.reserveInboundRPCBatch(context.Background(), []inboundRPCSpec{{method: "one", size: 3}})
if err != nil {
t.Fatal(err)
}
if err := reservation.growEntry(0, test.targetSize); !errors.Is(err, ErrInboundRPCQueueFull) {
t.Fatalf("grow error = %v, want ErrInboundRPCQueueFull", err)
}
if tasks, bytes := scheduler.budgetSnapshot(); tasks != 1 || bytes != 3 {
t.Fatalf("failed grow changed global budget %d/%d", tasks, bytes)
}
if got := c.inflightRPCBytes.Load(); got != 3 {
t.Fatalf("failed grow changed connection budget %d", got)
}
if reservation.entries[0].size != 3 || reservation.totalSize != 3 {
t.Fatalf("failed grow changed reservation entry/total = %d/%d", reservation.entries[0].size, reservation.totalSize)
}
reservation.abort()
})
}
}
func TestInboundRPCBatchCommitAppendsAllAndSchedulesAtomically(t *testing.T) {
scheduler := newInboundRPCScheduler(1, 8, 1<<20)
c := newInboundTestConn(scheduler, 1, 4, time.Second)

678
internal/mtprotoedge/layer_admission_budget_test.go
View file

@ -2,7 +2,6 @@ package mtprotoedge
import (
"context"
"errors"
"reflect"
"strings"
"sync/atomic"
@ -11,8 +10,12 @@ import (
"github.com/iamxvbaba/td/bin"
"github.com/iamxvbaba/td/clock"
"github.com/iamxvbaba/td/proto"
"github.com/iamxvbaba/td/tg"
"github.com/iamxvbaba/td/tgerr"
"go.uber.org/zap"
"go.uber.org/zap/zaptest"
"go.uber.org/zap/zaptest/observer"
"github.com/iamxvbaba/td/tlprofile"
appfiles "telesrv/internal/app/files"
@ -85,6 +88,18 @@ func TestLayerRPCAdmissionMaterializationConstants(t *testing.T) {
if got := layerRPCAdmissionReservationSize(-1); got != layerRPCAdmissionGraphSlack {
t.Fatalf("negative wire charge = %d, want fixed slack %d", got, layerRPCAdmissionGraphSlack)
}
if got := layerRPCFlatBytesWireCharge(100, 80); got != 20*layerRPCAdmissionWireFactor+80*layerRPCAdmissionFlatBytesFactor {
t.Fatalf("flat bytes wire charge = %d", got)
}
if got := layerRPCFlatBytesWireCharge(10, 11); got != layerRPCAdmissionWireCharge(10) {
t.Fatalf("invalid flat bytes hint charge = %d, want generic %d", got, layerRPCAdmissionWireCharge(10))
}
if got := layerRPCFlatBytesWireCharge(maxInt, maxInt); got != maxInt {
t.Fatalf("saturating flat bytes wire charge = %d, want max int %d", got, maxInt)
}
if got := saturatingLayerRPCAdmissionCharge(maxInt, 1); got != maxInt {
t.Fatalf("saturating addition = %d, want max int %d", got, maxInt)
}
}
type countingLayerRPCAdmission struct {
@ -92,32 +107,32 @@ type countingLayerRPCAdmission struct {
decodeCalls atomic.Int32
}
type failingReplayLayerRPC struct {
LayerRPCHandler
err error
}
func (h *failingReplayLayerRPC) PrepareAdmittedReplay(
context.Context,
[8]byte,
int64,
int64,
uint64,
tlprofile.Admission,
) (func() error, error) {
return nil, h.err
}
func (h *countingLayerRPCAdmission) AdmitLayer(profile tlprofile.Profile, b *bin.Buffer, limits tlprofile.Limits) (tlprofile.Admission, error) {
h.decodeCalls.Add(1)
return h.LayerRPCHandler.AdmitLayer(profile, b, limits)
}
func (h *countingLayerRPCAdmission) AdmitLayerWithOptions(profile tlprofile.Profile, b *bin.Buffer, options tlprofile.AdmissionOptions) (tlprofile.Admission, error) {
h.decodeCalls.Add(1)
if admitter, ok := h.LayerRPCHandler.(LayerRPCOptionsAdmitter); ok {
return admitter.AdmitLayerWithOptions(profile, b, options)
}
return h.LayerRPCHandler.AdmitLayer(profile, b, options.Limits)
}
func (h *countingLayerRPCAdmission) AdmitUnprofiled(b *bin.Buffer, limits tlprofile.Limits) (tlprofile.Admission, error) {
h.decodeCalls.Add(1)
return h.LayerRPCHandler.AdmitUnprofiled(b, limits)
}
func (h *countingLayerRPCAdmission) AdmitUnprofiledWithOptions(b *bin.Buffer, options tlprofile.AdmissionOptions) (tlprofile.Admission, error) {
h.decodeCalls.Add(1)
if admitter, ok := h.LayerRPCHandler.(LayerRPCOptionsAdmitter); ok {
return admitter.AdmitUnprofiledWithOptions(b, options)
}
return h.LayerRPCHandler.AdmitUnprofiled(b, options.Limits)
}
func TestLayerRPCAdmissionCapacityRejectsBeforeDecoder(t *testing.T) {
for _, test := range []struct {
name string
@ -162,6 +177,557 @@ func TestLayerRPCAdmissionCapacityRejectsBeforeDecoder(t *testing.T) {
}
}
func TestLayerRPCAdmissionExpandsTDLibNestedGZIPUnderTransferredBudget(t *testing.T) {
router := rpc.New(rpc.Config{DC: 2}, rpc.Deps{}, zaptest.NewLogger(t), clock.System)
s := New(Options{DC: 2, LayerRPC: router, Logger: zaptest.NewLogger(t)})
c := &Conn{authKeyID: [8]byte{8, 21}, sessionID: 821, metrics: NopMetrics{}}
c.startInboundRPCScheduler(s.rpcScheduler, 1, 4, time.Second)
defer func() {
c.closeInboundRPCScheduler()
s.rpcScheduler.stop(time.Second)
}()
body, expandedBytes := tdlibNestedGZIPBody(t, tlprofile.Profile228, &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 plan.items[0].kind != inboundItemRPC || plan.rpcReservation == nil || len(plan.rpcTasks) != 1 {
t.Fatalf("admitted nested gzip plan = kind:%d reservation:%v tasks:%d", plan.items[0].kind, plan.rpcReservation != nil, len(plan.rpcTasks))
}
wantCharge := int64(layerRPCAdmissionReservationSize(len(body) + expandedBytes))
if got := c.inflightRPCBytes.Load(); got != wantCharge {
t.Fatalf("nested gzip connection charge = %d, want %d", got, wantCharge)
}
if got := plan.rpcReservation.entries[0].size; int64(got) != wantCharge {
t.Fatalf("nested gzip reservation charge = %d, want %d", got, wantCharge)
}
if got := plan.gzipExpandedBytes; got != expandedBytes {
t.Fatalf("nested gzip cumulative expansion = %d, want %d", got, expandedBytes)
}
if got := s.frameBudget.usedBytes(); got != 0 {
t.Fatalf("nested gzip temporary frame budget retained after materialization: %d", got)
}
}
func TestLayerRPCAdmissionAdmitsTDLibUploadParts(t *testing.T) {
tests := []struct {
name string
method string
body bin.Object
withoutGZIP bool
bare bool
}{
{
name: "pixel_9a_small_file_part_negative_file_id",
method: "upload.saveFilePart",
body: &tg.UploadSaveFilePartRequest{
FileID: -3596058967254453060,
FilePart: 0,
Bytes: make([]byte, 1071),
},
},
{
name: "pixel_9a_big_file_part_gzip",
method: "upload.saveBigFilePart",
body: &tg.UploadSaveBigFilePartRequest{
FileID: 92,
FilePart: 0,
FileTotalParts: 364,
Bytes: make([]byte, 64<<10),
},
},
{
name: "pixel_9a_big_file_part_plain",
method: "upload.saveBigFilePart",
withoutGZIP: true,
body: &tg.UploadSaveBigFilePartRequest{
FileID: 93,
FilePart: 7,
FileTotalParts: 364,
Bytes: make([]byte, 64<<10),
},
},
{
name: "pixel_9a_big_file_part_bare_upload_session",
method: "upload.saveBigFilePart",
bare: true,
body: &tg.UploadSaveBigFilePartRequest{
FileID: 94,
FilePart: 15,
FileTotalParts: 364,
Bytes: make([]byte, 64<<10),
},
},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
router := rpc.New(rpc.Config{DC: 2}, rpc.Deps{}, zaptest.NewLogger(t), clock.System)
s := New(Options{DC: 2, LayerRPC: router, Logger: zaptest.NewLogger(t)})
c := &Conn{authKeyID: [8]byte{8, 31}, sessionID: 831, metrics: NopMetrics{}}
c.startInboundRPCScheduler(s.rpcScheduler, 1, 4, time.Second)
defer func() {
c.closeInboundRPCScheduler()
s.rpcScheduler.stop(time.Second)
}()
var (
body []byte
expandedBytes int
)
if tc.bare {
body = exactOutboundLayerRPCBody(t, tlprofile.Profile228, tc.body)
} else if tc.withoutGZIP {
body = tdlibWrappedBody(t, tlprofile.Profile228, tc.body)
} else {
body, expandedBytes = tdlibNestedGZIPBody(t, tlprofile.Profile228, tc.body)
}
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 != inboundItemRPC || plan.rpcReservation == nil || len(plan.rpcTasks) != 1 {
t.Fatalf("admitted nested gzip upload plan = kind:%d reservation:%v tasks:%d payload:%v",
plan.items[0].kind, plan.rpcReservation != nil, len(plan.rpcTasks), plan.items[0].payload)
}
if got := plan.gzipExpandedBytes; got != expandedBytes {
t.Fatalf("nested gzip upload cumulative expansion = %d, want %d", got, expandedBytes)
}
// Admission alone is insufficient: the prepared wrapper chain must remain
// executable after the temporary gzip expansion buffer has been released.
// With no Files dependency configured, reaching the upload handler has the
// stable terminal NOT_IMPLEMENTED; INPUT_REQUEST_INVALID means the wrapper or
// prepared-call boundary corrupted the otherwise valid request.
requestBody := &bin.Buffer{Buf: append([]byte(nil), body...)}
admitted, err := router.AdmitLayerWithOptions(tlprofile.Profile228, requestBody, tlprofile.AdmissionOptions{
Limits: inboundLayerDecodeLimits,
ExpandGZIP: func(wire []byte, limit int) ([]byte, func(), error) {
return s.decodeGZIPWithGlobalBudgetLimit(&bin.Buffer{Buf: wire}, limit)
},
})
if err != nil {
t.Fatal(err)
}
_, method, err := router.DispatchAdmitted(
rpc.WithUserID(context.Background(), 42),
[8]byte{8, 31},
831,
100,
1,
admitted,
)
if method != tc.method || !tgerr.Is(err, "NOT_IMPLEMENTED") {
t.Fatalf("nested gzip upload dispatch = method:%q err:%v, want %s/NOT_IMPLEMENTED", method, err, tc.method)
}
})
}
}
func TestLayerRPCAdmissionRejectionLogsBoundedMetadataWithoutUploadBody(t *testing.T) {
const marker = "DO_NOT_LOG_UPLOAD_BODY_MARKER"
payload := make([]byte, 1071)
copy(payload, marker)
body := tdlibWrappedBody(t, tlprofile.Profile228, &tg.UploadSaveFilePartRequest{
FileID: -3596058967254453060,
FilePart: 0,
Bytes: payload,
})
// Remove the final TL padding byte. Exact admission must reject the malformed
// wire while the edge still records its explicit wrapper and bounded cause.
body = body[:len(body)-1]
core, logs := observer.New(zap.DebugLevel)
router := rpc.New(rpc.Config{DC: 2}, rpc.Deps{}, zap.New(core), clock.System)
s := New(Options{DC: 2, LayerRPC: router, Logger: zap.New(core)})
c := &Conn{
authKeyID: [8]byte{8, 34},
authKeyHex: "0822000000000000",
sessionID: 834,
metrics: NopMetrics{},
}
c.startInboundRPCScheduler(s.rpcScheduler, 1, 2, time.Second)
defer func() {
c.closeInboundRPCScheduler()
s.rpcScheduler.stop(time.Second)
}()
for attempt := 0; attempt < 2; attempt++ {
plan := &inboundPlan{items: []inboundItem{{
kind: inboundItemRPC,
msgID: int64(100 + attempt*4),
typeID: tg.InvokeWithLayerRequestTypeID,
body: body,
}}}
if err := s.prepareInboundLayerRPCBatch(context.Background(), c, plan); err != nil {
plan.close()
t.Fatal(err)
}
if item := plan.items[0]; item.kind != inboundItemRPCAdmissionError {
plan.close()
t.Fatalf("malformed upload attempt %d kind = %d, want admission error", attempt, item.kind)
}
plan.close()
}
entries := logs.FilterMessage("RPC exact admission rejected").All()
if len(entries) != 2 {
t.Fatalf("admission rejection log count = %d, want 2", len(entries))
}
if entries[0].Level != zap.InfoLevel || entries[1].Level != zap.DebugLevel {
t.Fatalf("admission rejection levels = %s/%s, want info/debug", entries[0].Level, entries[1].Level)
}
fields := entries[0].ContextMap()
for key, want := range map[string]any{
"method": "invokeWithLayer#da9b0d0d",
"auth_key_id": "0822000000000000",
"session_id": int64(834),
"msg_id": int64(100),
"top_level_wire_id": uint32(tg.InvokeWithLayerRequestTypeID),
"wire_bytes": int64(len(body)),
"profile_origin": "unknown",
"explicit_layer_selector": true,
"rpc_error_code": int64(400),
"rpc_error_message": "INPUT_REQUEST_INVALID",
} {
if got := fields[key]; got != want {
t.Fatalf("admission rejection field %q = %#v (%T), want %#v (%T); all=%#v", key, got, got, want, want, fields)
}
}
for _, entry := range entries {
if strings.Contains(entry.Message, marker) || strings.Contains(entry.ContextMap()["error"].(string), marker) {
t.Fatalf("admission rejection leaked upload body marker: %#v", entry.ContextMap())
}
if _, ok := entry.ContextMap()["body"]; ok {
t.Fatalf("admission rejection exposed body field: %#v", entry.ContextMap())
}
}
}
func TestLayerRPCAdmissionAdmitsTDLibFirstUploadContainer(t *testing.T) {
router := rpc.New(rpc.Config{DC: 2}, rpc.Deps{}, zaptest.NewLogger(t), clock.System)
s := New(Options{DC: 2, LayerRPC: router, Logger: zaptest.NewLogger(t)})
c := &Conn{authKeyID: [8]byte{8, 32}, sessionID: 832, metrics: NopMetrics{}}
c.startInboundRPCScheduler(s.rpcScheduler, 1, 8, time.Second)
defer func() {
c.closeInboundRPCScheduler()
s.rpcScheduler.stop(time.Second)
}()
// A newly opened TDLib upload Session applies its invokeWithLayer /
// initConnection header to every query in the first MTProto container. The
// Pixel 9a trace contains eight gzip-packed 64 KiB saveBigFilePart requests
// in that first container, all with the same known total and distinct
// parts/message IDs. The fixture is an ELF and its first chunks compress to
// roughly 11-14 KiB each, yielding a 129 KiB encrypted write.
plan, legacyGenericCharge := tdlibFirstUploadPlan(t)
defer plan.close()
if legacyGenericCharge <= maxInflightRPCBytes {
t.Fatalf("test fixture generic charge = %d, must exceed old connection ceiling %d", legacyGenericCharge, maxInflightRPCBytes)
}
if err := s.prepareInboundLayerRPCBatch(context.Background(), c, plan); err != nil {
t.Fatal(err)
}
if plan.rpcReservation == nil || len(plan.rpcTasks) != len(plan.items) {
t.Fatalf("first upload container reservation/tasks = %v/%d, want retained/%d",
plan.rpcReservation != nil, len(plan.rpcTasks), len(plan.items))
}
if got := plan.rpcReservation.totalSize; got <= 0 || got >= legacyGenericCharge || got > maxInflightRPCBytes {
t.Fatalf("first upload container retained charge = %d, want 0 < charge < legacy %d and <= %d",
got, legacyGenericCharge, maxInflightRPCBytes)
}
for index := range plan.items {
item := &plan.items[index]
if item.kind != inboundItemRPC {
t.Fatalf("first upload container item %d kind = %d, payload = %v", index, item.kind, item.payload)
}
if method := plan.rpcTasks[index].method; method != "upload.saveBigFilePart" {
t.Fatalf("first upload container task %d method = %q, want upload.saveBigFilePart", index, method)
}
}
}
func TestLayerRPCAdmissionTDLibFirstUploadContainerRequiresFlatBytesCapability(t *testing.T) {
router := rpc.New(rpc.Config{DC: 2}, rpc.Deps{}, zaptest.NewLogger(t), clock.System)
// The wrapper deliberately exposes exact admission but not the optional
// flat-bytes sizing capability. The edge must therefore retain the generic
// worst-case charge and reject the whole oversized batch atomically.
handler := &countingLayerRPCAdmission{LayerRPCHandler: router}
s := New(Options{DC: 2, LayerRPC: handler, Logger: zaptest.NewLogger(t)})
c := &Conn{authKeyID: [8]byte{8, 33}, sessionID: 833, metrics: NopMetrics{}}
c.startInboundRPCScheduler(s.rpcScheduler, 1, 8, time.Second)
defer func() {
c.closeInboundRPCScheduler()
s.rpcScheduler.stop(time.Second)
}()
plan, _ := tdlibFirstUploadPlan(t)
defer plan.close()
if err := s.prepareInboundLayerRPCBatch(context.Background(), c, plan); err != nil {
t.Fatal(err)
}
for index := range plan.items {
if plan.items[index].kind != inboundItemCapacityError {
t.Fatalf("item %d kind = %d, want capacity error", index, plan.items[index].kind)
}
}
if plan.rpcReservation != nil || len(plan.rpcTasks) != 0 {
t.Fatalf("generic fallback retained reservation/tasks = %v/%d", plan.rpcReservation != nil, len(plan.rpcTasks))
}
if got := c.inflightRPCBytes.Load(); got != 0 {
t.Fatalf("generic fallback leaked connection charge %d", got)
}
if tasks, bytes := s.rpcScheduler.budgetSnapshot(); tasks != 0 || bytes != 0 {
t.Fatalf("generic fallback leaked global budget %d/%d", tasks, bytes)
}
}
func tdlibFirstUploadPlan(t *testing.T) (*inboundPlan, int64) {
t.Helper()
plan := &inboundPlan{items: make([]inboundItem, 8)}
var legacyGenericCharge int64
for index := range plan.items {
payload := make([]byte, 64<<10)
state := uint32(index + 1)
for byteIndex := 0; byteIndex < 13<<10; byteIndex++ {
state ^= state << 13
state ^= state >> 17
state ^= state << 5
payload[byteIndex] = byte(state)
}
body, expandedBytes := tdlibNestedGZIPBody(t, tlprofile.Profile228, &tg.UploadSaveBigFilePartRequest{
FileID: 95,
FilePart: index,
FileTotalParts: 364,
Bytes: payload,
})
plan.items[index] = inboundItem{
kind: inboundItemRPC,
msgID: int64(100 + index*4),
body: body,
}
legacyGenericCharge += int64(layerRPCAdmissionReservationSize(len(body) + expandedBytes))
}
return plan, legacyGenericCharge
}
func TestLayerRPCAdmissionAdmitsTDLibWrappedBindTempAuthKey(t *testing.T) {
router := rpc.New(rpc.Config{DC: 2}, rpc.Deps{}, zaptest.NewLogger(t), clock.System)
s := New(Options{DC: 2, LayerRPC: router, Logger: zaptest.NewLogger(t)})
c := &Conn{authKeyID: [8]byte{8, 41}, sessionID: -7940676790771565328, metrics: NopMetrics{}}
c.startInboundRPCScheduler(s.rpcScheduler, 1, 4, time.Second)
defer func() {
c.closeInboundRPCScheduler()
s.rpcScheduler.stop(time.Second)
}()
request := &tg.AuthBindTempAuthKeyRequest{
PermAuthKeyID: 9179421154451858694,
Nonce: 5318578202586482454,
ExpiresAt: 1785817822,
EncryptedMessage: make([]byte, 104),
}
body := tdlibWrappedBody(t, tlprofile.Profile228, request)
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 != inboundItemRPC || plan.rpcReservation == nil || len(plan.rpcTasks) != 1 {
t.Fatalf("admitted TDLib bind plan = kind:%d reservation:%v tasks:%d payload:%v",
plan.items[0].kind, plan.rpcReservation != nil, len(plan.rpcTasks), plan.items[0].payload)
}
requestBody := &bin.Buffer{Buf: append([]byte(nil), body...)}
admitted, err := router.AdmitUnprofiled(requestBody, inboundLayerDecodeLimits)
if err != nil {
t.Fatal(err)
}
result, method, err := router.DispatchAdmitted(
context.Background(),
c.authKeyID,
c.sessionID,
100,
1,
admitted,
)
if err != nil || method != "auth.bindTempAuthKey" || result == nil || !result.WireInvariant() {
t.Fatalf("TDLib wrapped bind dispatch = method:%q result:%T invariant:%v err:%v",
method, result, result != nil && result.WireInvariant(), err)
}
}
func TestLayerRPCAdmissionNestedGZIPGrowFailureRejectsWholeBatch(t *testing.T) {
router := rpc.New(rpc.Config{DC: 2}, rpc.Deps{}, zaptest.NewLogger(t), clock.System)
s := New(Options{DC: 2, LayerRPC: router, Logger: zaptest.NewLogger(t)})
first, _ := tdlibNestedGZIPBody(t, tlprofile.Profile228, &tg.HelpGetConfigRequest{})
second, _ := tdlibNestedGZIPBody(t, tlprofile.Profile228, &tg.HelpGetNearestDCRequest{})
initialCharge := int64(layerRPCAdmissionReservationSize(len(first)) + layerRPCAdmissionReservationSize(len(second)))
s.rpcScheduler = newInboundRPCScheduler(1, 4, initialCharge)
c := &Conn{authKeyID: [8]byte{8, 22}, sessionID: 822, metrics: NopMetrics{}}
c.startInboundRPCScheduler(s.rpcScheduler, 1, 4, time.Second)
defer func() {
c.closeInboundRPCScheduler()
s.rpcScheduler.stop(time.Second)
}()
plan := &inboundPlan{items: []inboundItem{
{kind: inboundItemRPC, msgID: 100, body: first},
{kind: inboundItemRPC, msgID: 104, body: second},
}}
defer plan.close()
if err := s.prepareInboundLayerRPCBatch(context.Background(), c, plan); err != nil {
t.Fatal(err)
}
for index := range plan.items {
if plan.items[index].kind != inboundItemCapacityError {
t.Fatalf("item %d kind = %d, want capacity error", index, plan.items[index].kind)
}
}
if plan.rpcReservation != nil || len(plan.rpcTasks) != 0 {
t.Fatalf("capacity plan retained reservation/tasks = %v/%d", plan.rpcReservation != nil, len(plan.rpcTasks))
}
if got := c.inflightRPCBytes.Load(); got != 0 {
t.Fatalf("grow failure leaked connection charge %d", got)
}
if tasks, bytes := s.rpcScheduler.budgetSnapshot(); tasks != 0 || bytes != 0 {
t.Fatalf("grow failure leaked global budget %d/%d", tasks, bytes)
}
if got := s.frameBudget.usedBytes(); got != 0 {
t.Fatalf("grow failure leaked temporary frame budget %d", got)
}
}
func TestLayerRPCAdmissionNestedGZIPSiblingsShareFrameExpansionLimit(t *testing.T) {
router := rpc.New(rpc.Config{DC: 2}, rpc.Deps{}, zaptest.NewLogger(t), clock.System)
s := New(Options{DC: 2, LayerRPC: router, Logger: zaptest.NewLogger(t)})
c := &Conn{authKeyID: [8]byte{8, 23}, sessionID: 823, metrics: NopMetrics{}}
c.startInboundRPCScheduler(s.rpcScheduler, 1, 4, time.Second)
defer func() {
c.closeInboundRPCScheduler()
s.rpcScheduler.stop(time.Second)
}()
first, expandedBytes := tdlibNestedGZIPBody(t, tlprofile.Profile228, &tg.HelpGetConfigRequest{})
second, _ := tdlibNestedGZIPBody(t, tlprofile.Profile228, &tg.HelpGetNearestDCRequest{})
plan := &inboundPlan{
gzipExpandedBytes: maxDispatchExpandedBytes - expandedBytes,
items: []inboundItem{
{kind: inboundItemRPC, msgID: 100, body: first},
{kind: inboundItemRPC, msgID: 104, body: second},
},
}
defer plan.close()
if err := s.prepareInboundLayerRPCBatch(context.Background(), c, plan); err != nil {
t.Fatal(err)
}
for index := range plan.items {
if plan.items[index].kind != inboundItemCapacityError {
t.Fatalf("item %d kind = %d, want capacity error", index, plan.items[index].kind)
}
}
if got := plan.gzipExpandedBytes; got != maxDispatchExpandedBytes {
t.Fatalf("shared cumulative expansion = %d, want %d", got, maxDispatchExpandedBytes)
}
if got := c.inflightRPCBytes.Load(); got != 0 {
t.Fatalf("shared-limit rejection leaked connection charge %d", got)
}
if got := s.frameBudget.usedBytes(); got != 0 {
t.Fatalf("shared-limit rejection leaked temporary frame budget %d", got)
}
}
func TestLayerRPCAdmissionNestedGZIPReDecodeReusesMaterializationCharge(t *testing.T) {
handler := newAdmissionOnlyLayerRPC()
s := New(Options{DC: 2, LayerRPC: handler, Logger: zaptest.NewLogger(t)})
s.rpcResults = newRPCExecutionLedgerForServerTest(s, time.Now, 8)
scheduler := newInboundRPCScheduler(1, 4, 1<<30)
s.rpcScheduler = scheduler
authKeyID := [8]byte{8, 24}
const sessionID = int64(824)
c225 := &Conn{authKeyID: authKeyID, sessionID: sessionID, metrics: NopMetrics{}}
c227 := &Conn{authKeyID: authKeyID, sessionID: sessionID, metrics: NopMetrics{}}
c225.startInboundRPCScheduler(scheduler, 1, 2, time.Second)
c227.startInboundRPCScheduler(scheduler, 1, 2, time.Second)
defer func() {
c225.closeInboundRPCScheduler()
c227.closeInboundRPCScheduler()
scheduler.stop(time.Second)
}()
terminal := exactOutboundLayerRPCBody(t, tlprofile.Profile225, &tg.MessagesGetHistoryRequest{
Peer: &tg.InputPeerSelf{}, Limit: 1,
})
body := exactLayerRPCBody(t, &tg.InvokeWithoutUpdatesRequest{Query: &proto.GZIP{Data: terminal}})
initialCharge := layerRPCAdmissionReservationSize(len(body))
reservation225, err := c225.reserveInboundRPCBatch(context.Background(), []inboundRPCSpec{{method: "messages.getHistory", size: initialCharge}})
if err != nil {
t.Fatal(err)
}
defer reservation225.abort()
reservation227, err := c227.reserveInboundRPCBatch(context.Background(), []inboundRPCSpec{{method: "messages.getHistory", size: initialCharge}})
if err != nil {
t.Fatal(err)
}
defer reservation227.abort()
plan225 := &inboundPlan{}
budget225 := &layerRPCGZIPExpansionBudget{
server: s, plan: plan225, reservation: reservation225,
baseCharge: initialCharge, chargedSize: initialCharge,
}
options225 := tlprofile.AdmissionOptions{Limits: inboundLayerDecodeLimits, ExpandGZIP: budget225.expand}
item225 := inboundItem{msgID: 100, body: body}
item225.admitted, item225.method, err = s.decodeInboundLayerRPCWithOptions(
LayerProfileSnapshot{Profile: tlprofile.Profile225, Origin: LayerProfileInherited}, body, options225,
)
if err != nil {
t.Fatal(err)
}
plan227 := &inboundPlan{}
budget227 := &layerRPCGZIPExpansionBudget{
server: s, plan: plan227, reservation: reservation227,
baseCharge: initialCharge, chargedSize: initialCharge,
}
options227 := tlprofile.AdmissionOptions{Limits: inboundLayerDecodeLimits, ExpandGZIP: budget227.expand}
item227 := inboundItem{msgID: 100, body: body}
item227.admitted, item227.method, err = s.decodeInboundLayerRPCWithOptions(
LayerProfileSnapshot{Profile: tlprofile.Profile227, Origin: LayerProfileInherited}, body, options227,
)
if err != nil {
t.Fatal(err)
}
if item225.admitted.Prepared().Identity() == item227.admitted.Prepared().Identity() {
t.Fatal("test request identity is invariant; need authoritative-profile re-decode")
}
winner, err := s.acquireAdmittedLayerRPC(c225, &item225, nil, options225, budget225)
if err != nil || winner.state != rpcResultAcquireOwner || winner.owner == nil {
t.Fatalf("winner = state:%d err:%v", winner.state, err)
}
defer winner.owner.Abort()
loser, err := s.acquireAdmittedLayerRPC(c227, &item227, nil, options227, budget227)
if err != nil || loser.state != rpcResultAcquirePending {
t.Fatalf("loser = state:%d err:%v", loser.state, err)
}
if got := item227.admitted.Call().Profile(); got != tlprofile.Profile225 {
t.Fatalf("loser re-admitted profile = %d, want 225", got)
}
wantCharge := layerRPCAdmissionReservationSize(len(body) + len(terminal))
if got := reservation227.entries[0].size; got != wantCharge {
t.Fatalf("re-decode reservation charge = %d, want single-graph maximum %d", got, wantCharge)
}
if got := plan227.gzipExpandedBytes; got != 2*len(terminal) {
t.Fatalf("re-decode cumulative work = %d, want %d", got, 2*len(terminal))
}
if got := s.frameBudget.usedBytes(); got != 0 {
t.Fatalf("re-decode leaked temporary frame budget %d", got)
}
}
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})
@ -207,6 +773,34 @@ func TestLayerRPCAdmissionTransfersOriginalReservationToFreshOwner(t *testing.T)
}
}
func tdlibNestedGZIPBody(t *testing.T, profile tlprofile.Profile, terminal bin.Object) ([]byte, int) {
t.Helper()
terminalWire := exactOutboundLayerRPCBody(t, profile, terminal)
return tdlibWrappedBody(t, profile, zlibPackedObjectForTest(t, terminalWire)), len(terminalWire)
}
func tdlibWrappedBody(t *testing.T, profile tlprofile.Profile, terminal bin.Object) []byte {
t.Helper()
request := &tg.InvokeWithLayerRequest{
Layer: int(profile),
Query: &tg.InitConnectionRequest{
APIID: 1,
DeviceModel: "android",
SystemVersion: "test",
AppVersion: "1.0",
SystemLangCode: "en",
LangPack: "",
LangCode: "en",
Params: &tg.JSONObject{Value: []tg.JSONObjectValue{{
Key: "tz_offset",
Value: &tg.JSONNumber{Value: 8 * 60 * 60},
}}},
Query: terminal,
},
}
return exactLayerRPCBody(t, request)
}
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})
@ -270,7 +864,7 @@ func TestLayerRPCAdmissionCompletedReplayReleasesWholeProvisionalBatch(t *testin
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}})
storeLogicalRPCResultForTest(t, s, c, 100, &encodedOutboundMessage{body: []byte{1, 2, 3, 4}})
plan := &inboundPlan{items: []inboundItem{{kind: inboundItemRPC, msgID: 100, body: body}}}
defer plan.close()
@ -291,54 +885,6 @@ func TestLayerRPCAdmissionCompletedReplayReleasesWholeProvisionalBatch(t *testin
}
}
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(tlprofile.Profile225); err != nil {
t.Fatal(err)
}
body := exactOutboundLayerRPCBody(t, tlprofile.Profile225, &tg.HelpGetConfigRequest{})
identityBuffer := &bin.Buffer{Buf: append([]byte(nil), body...)}
request, err := router.AdmitLayer(tlprofile.Profile225, identityBuffer, tlprofile.Limits{})
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})

2
internal/mtprotoedge/layer_rpc_execution.go
View file

@ -82,7 +82,7 @@ var errLayerRPCResultIdentityMismatch = errors.New("layer RPC result does not ma
// 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.
// execution-ledger identity even when both methods happen to share a Go type.
func bindAdmittedLayerRPCResult(request tlprofile.Admission, result tlprofile.Result) (*layerRPCResultEncoder, error) {
if result == nil {
return nil, nil

10
internal/mtprotoedge/layer_rpc_execution_test.go
View file

@ -138,6 +138,9 @@ func TestProjectionFailureCachesInternalWithoutRepeatingBusiness(t *testing.T) {
if rpcErr.ErrorCode != 500 || rpcErr.ErrorMessage != "INTERNAL" {
t.Fatalf("projection terminal = %+v", rpcErr)
}
if _, ok := s.rpcResults.Replay(c.authKeyID, c.sessionID, reqMsgID); !ok {
t.Fatal("projection receipt disappeared before duplicate acquire")
}
// A same-msg replay is served from the completed exact identity; there is no
// second DispatchAdmitted call even though projection failed after business
// success.
@ -145,9 +148,14 @@ func TestProjectionFailureCachesInternalWithoutRepeatingBusiness(t *testing.T) {
c.authKeyID, c.sessionID, reqMsgID,
tlprofile.Profile227, request.Prepared().Identity(),
)
if err != nil || replay.state != rpcResultAcquireCompleted || replay.encoded != completed.encoded {
if err != nil || replay.state != rpcResultAcquireCompleted {
t.Fatalf("projection replay = state:%d err:%v", replay.state, err)
}
if replay.encoded == nil || replay.encoded.replayMsgID != completed.encoded.replayMsgID ||
replay.encoded.replaySeqNo != completed.encoded.replaySeqNo ||
!sameBacking(replay.encoded.body, completed.encoded.body) {
t.Fatal("projection replay did not reference the original logical-outbox frame")
}
if got := handler.calls.Load(); got != 1 {
t.Fatalf("replay repeated business calls=%d", got)
}

2
internal/mtprotoedge/layer_seed_test.go
View file

@ -28,7 +28,7 @@ func TestRegisterSeedsNegotiatedLayerBeforeFirstRPC(t *testing.T) {
conn, auth, cipher := dialHandshake(t, addr, 2, pub)
clientMsgID := proto.NewMessageIDGen(time.Now)
sendEncrypted(t, conn, cipher, auth, clientMsgID.New(proto.MessageFromClient), &mt.PingRequest{PingID: 7})
sendEncryptedWithSeq(t, conn, cipher, auth, clientMsgID.New(proto.MessageFromClient), 1, &mt.PingRequest{PingID: 7})
// 等 pong 回来，确保携带注册动作的那一帧已处理完成。
gotPong := false

2
internal/mtprotoedge/layer_updates.go
View file

@ -32,7 +32,7 @@ const (
// 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.
// invokeWithLayer correction must not invalidate that in-flight/retained result.
outboundLayerBindingRequest
)

186
internal/mtprotoedge/logical_session.go Normal file
View file

@ -0,0 +1,186 @@
package mtprotoedge
import "time"
// logicalSession is the process-local MTProto session that survives physical
// transport replacement. The outbound state is the sole owner of every
// content-related server message until msgs_ack, explicit session destruction,
// or the bounded offline-retention window expires.
type logicalSession struct {
key sessionKey
outbound *outboundState
offlineAt time.Time
businessAuthKeyID [8]byte
businessAuthResolved bool
}
const logicalSessionOfflineTTL = 6 * time.Minute
func (m *SessionManager) attachLogicalSession(c *Conn, budget *outboundTrackedBudget) {
if m == nil || c == nil {
return
}
key := connSessionKey(c)
m.mu.Lock()
logical := m.logicalSessions[key]
if logical == nil {
logical = &logicalSession{
key: key,
outbound: newOutboundState(budget),
}
m.logicalSessions[key] = logical
}
logical.outbound.persistent.Store(true)
if businessAuthKeyID, resolved := c.BusinessAuthKeyID(); resolved {
logical.businessAuthKeyID = businessAuthKeyID
logical.businessAuthResolved = true
}
logical.offlineAt = time.Time{}
c.outboundState = logical.outbound
m.mu.Unlock()
}
// adoptLogicalSession is a construction-test/embedded-Conn bridge. Production
// Conns are attached before their actor starts; direct Conn builders may already
// own an actor-local state when a Server terminal callback publishes the receipt.
func (m *SessionManager) adoptLogicalSession(c *Conn) {
if m == nil || c == nil || c.outboundState == nil {
return
}
key := connSessionKey(c)
m.mu.Lock()
logical := m.logicalSessions[key]
if logical == nil {
logical = &logicalSession{key: key, outbound: c.outboundState}
m.logicalSessions[key] = logical
}
logical.outbound.persistent.Store(true)
if businessAuthKeyID, resolved := c.BusinessAuthKeyID(); resolved {
logical.businessAuthKeyID = businessAuthKeyID
logical.businessAuthResolved = true
}
logical.offlineAt = time.Time{}
// The actor's state pointer is immutable after startOutbound. Production
// attaches before start; this adoption bridge only marks that already-owned
// state persistent and must never write the Conn field concurrently.
m.mu.Unlock()
}
func (m *SessionManager) markLogicalSessionOfflineLocked(key sessionKey, now time.Time) {
logical := m.logicalSessions[key]
if logical == nil || m.bySession[key] != nil || m.claims[key] != nil {
return
}
if logical.offlineAt.IsZero() {
logical.offlineAt = now
}
}
func (m *SessionManager) destroyLogicalSessionLocked(key sessionKey) *outboundState {
logical := m.logicalSessions[key]
if logical == nil {
return nil
}
delete(m.logicalSessions, key)
return logical.outbound
}
func (m *SessionManager) bindLogicalSessionAuthKeyLocked(key sessionKey, businessAuthKeyID [8]byte) {
if logical := m.logicalSessions[key]; logical != nil {
logical.businessAuthKeyID = businessAuthKeyID
logical.businessAuthResolved = true
}
}
func (m *SessionManager) releaseLogicalSession(key sessionKey, state *outboundState) {
if state != nil {
state.releaseAll()
}
m.mu.RLock()
hook := m.logicalSessionReleased
m.mu.RUnlock()
if hook != nil {
hook(key)
}
}
// ForgetLogicalSessionsForRawAuthKey is the terminal auth-key destruction path.
// Once destroy_auth_key_ok is on the wire the key can never reconnect to ACK or
// replay an old answer, so retaining any session payload or receipt is useless.
func (m *SessionManager) ForgetLogicalSessionsForRawAuthKey(authKeyID [8]byte) {
if m == nil {
return
}
var release []*logicalSession
m.mu.Lock()
for key, logical := range m.logicalSessions {
if key.authKeyID != authKeyID {
continue
}
delete(m.logicalSessions, key)
if logical != nil {
release = append(release, logical)
}
}
m.mu.Unlock()
for _, logical := range release {
m.releaseLogicalSession(logical.key, logical.outbound)
}
}
func (m *SessionManager) sweepLogicalSessions(now time.Time) {
if m == nil {
return
}
var release []*logicalSession
m.mu.Lock()
for key, logical := range m.logicalSessions {
if logical == nil || logical.offlineAt.IsZero() || now.Sub(logical.offlineAt) < logicalSessionOfflineTTL {
continue
}
if m.bySession[key] != nil || m.claims[key] != nil {
logical.offlineAt = time.Time{}
continue
}
delete(m.logicalSessions, key)
release = append(release, logical)
}
m.mu.Unlock()
for _, logical := range release {
m.releaseLogicalSession(logical.key, logical.outbound)
}
}
func (m *SessionManager) releaseAllLogicalSessions() {
if m == nil {
return
}
var release []*logicalSession
m.mu.Lock()
for key, logical := range m.logicalSessions {
delete(m.logicalSessions, key)
if logical != nil {
release = append(release, logical)
}
}
m.mu.Unlock()
for _, logical := range release {
m.releaseLogicalSession(logical.key, logical.outbound)
}
}
// rpcResult returns the exact unacknowledged wire result owned by the logical
// session. The receipt ledger calls this only after validating request identity.
func (m *SessionManager) rpcResult(authKeyID [8]byte, sessionID, reqMsgID int64) (*encodedOutboundMessage, bool) {
if m == nil || reqMsgID == 0 {
return nil, false
}
key := sessionKey{authKeyID: authKeyID, sessionID: sessionID}
m.mu.RLock()
logical := m.logicalSessions[key]
m.mu.RUnlock()
if logical == nil || logical.outbound == nil {
return nil, false
}
return logical.outbound.rpcResult(reqMsgID)
}

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

@ -0,0 +1,282 @@
package mtprotoedge
import (
"bytes"
"encoding/binary"
"testing"
"time"
"github.com/iamxvbaba/td/proto"
"go.uber.org/zap"
)
// storeLogicalRPCResultForTest publishes the same production invariant as the
// outbound actor: the exact rpc_result frame enters the logical-session outbox
// before its metadata receipt becomes visible. Tests that call the receipt
// ledger directly must not invent the old impossible state where a replay row
// exists without an unacknowledged server frame.
func storeLogicalRPCResultForTest(
t *testing.T,
s *Server,
c *Conn,
reqMsgID int64,
encoded *encodedOutboundMessage,
) {
t.Helper()
if s == nil || c == nil || encoded == nil || len(encoded.body) == 0 {
t.Fatal("invalid logical rpc_result fixture")
}
if c.outboundState == nil {
s.conns.attachLogicalSession(c, s.outboundTrackedBudget)
} else {
s.conns.adoptLogicalSession(c)
}
physicalReqMsgID := encoded.writtenRequestID()
if physicalReqMsgID == 0 {
physicalReqMsgID = reqMsgID
}
frameBody := encoded.body
if physicalReqMsgID != reqMsgID && encoded.typeID == proto.ResultTypeID {
physical, err := cloneRPCResultForRequest(encoded, physicalReqMsgID, true)
if err != nil {
t.Fatalf("retarget logical rpc_result fixture: %v", err)
}
frameBody = physical.body
}
state := c.outboundState
state.mu.Lock()
if state.budget == nil || !state.budget.reserve(len(frameBody)) {
state.mu.Unlock()
t.Fatal("reserve logical rpc_result fixture")
}
candidate := reqMsgID*4 + 1
if candidate <= 0 {
candidate = 1
}
frame := &outboundFrame{
msgID: state.reserveMsgID(candidate), seqNo: state.peekSeqNo(true),
typeID: proto.ResultTypeID, body: frameBody, reqMsgID: physicalReqMsgID,
priority: encoded.priority, delivery: encoded.delivery,
compressed: encoded.compressed, uncompressedBytes: encoded.uncompressedBytes,
layer: encoded.layer, layerInvariant: encoded.layerInvariant,
reservedBytes: len(frameBody), reservationBudget: state.budget,
}
if err := state.admitReserved(frame); err != nil {
frame.releaseReservation(state.budget)
state.mu.Unlock()
t.Fatalf("admit logical rpc_result fixture: %v", err)
}
encoded.replayMsgID = frame.msgID
encoded.replaySeqNo = frame.seqNo
state.mu.Unlock()
s.rpcResults.Complete(c.authKeyID, c.sessionID, reqMsgID, encoded, true)
}
func acknowledgeLogicalRPCResultForTest(t *testing.T, s *Server, c *Conn, reqMsgID int64) {
t.Helper()
if s == nil || c == nil || c.outboundState == nil {
t.Fatal("missing logical rpc_result fixture")
}
state := c.outboundState
state.mu.Lock()
msgID := state.byRequest[reqMsgID]
requestIDs := state.ack([]int64{msgID})
state.mu.Unlock()
if len(requestIDs) != 1 || requestIDs[0] != reqMsgID {
t.Fatalf("logical ACK resolved request ids %v", requestIDs)
}
if !s.rpcResults.Acknowledge(c.authKeyID, c.sessionID, reqMsgID) {
t.Fatal("logical ACK did not release result receipt")
}
}
func TestLogicalSessionOwnsExactResultAcrossPhysicalReconnect(t *testing.T) {
manager := NewSessionManager(zap.NewNop())
budget := newOutboundTrackedBudget(1024)
authKeyID := [8]byte{1, 2, 3, 4, 5, 6, 7, 8}
first := &Conn{authKeyID: authKeyID, sessionID: 42}
manager.attachLogicalSession(first, budget)
body := append([]byte{0xf3, 0x5c, 0x6d, 0xf3}, make([]byte, 28)...)
binary.LittleEndian.PutUint64(body[4:12], uint64(77))
if !budget.reserve(len(body)) {
t.Fatal("reserve result body")
}
frame := &outboundFrame{
msgID: 101, seqNo: 1, typeID: proto.ResultTypeID, body: body,
reqMsgID: 77, reservedBytes: len(body), reservationBudget: budget,
}
first.outboundState.mu.Lock()
if err := first.outboundState.admitReserved(frame); err != nil {
first.outboundState.mu.Unlock()
t.Fatalf("admit result: %v", err)
}
first.outboundState.mu.Unlock()
second := &Conn{authKeyID: authKeyID, sessionID: 42}
manager.attachLogicalSession(second, budget)
if first.outboundState != second.outboundState {
t.Fatal("physical reconnect did not reuse logical outbound state")
}
replay, ok := manager.rpcResult(authKeyID, 42, 77)
if !ok {
t.Fatal("logical result not found after reconnect")
}
if replay.replayMsgID != frame.msgID || replay.replaySeqNo != frame.seqNo || !bytes.Equal(replay.body, body) {
t.Fatalf("replay identity/body = msg:%d seq:%d bytes:%d", replay.replayMsgID, replay.replaySeqNo, len(replay.body))
}
if &replay.body[0] != &frame.body[0] {
t.Fatal("replay created a second payload owner")
}
attempt, err := cloneRPCResultForRequest(replay, replay.reqMsgID, false)
if err != nil {
t.Fatalf("clone same-request replay descriptor: %v", err)
}
if attempt.replayMsgID != frame.msgID || attempt.replaySeqNo != frame.seqNo ||
!sameBacking(attempt.body, frame.body) {
t.Fatal("queued duplicate lost stable logical frame identity")
}
}
func TestLogicalSessionACKReleasesPayloadAndReceipt(t *testing.T) {
manager := NewSessionManager(zap.NewNop())
budget := newOutboundTrackedBudget(1024)
authKeyID := [8]byte{8, 7, 6, 5, 4, 3, 2, 1}
c := &Conn{authKeyID: authKeyID, sessionID: 43}
manager.attachLogicalSession(c, budget)
body := make([]byte, 64)
if !budget.reserve(len(body)) {
t.Fatal("reserve result body")
}
frame := &outboundFrame{
msgID: 201, seqNo: 1, typeID: proto.ResultTypeID, body: body,
reqMsgID: 88, reservedBytes: len(body), reservationBudget: budget,
}
c.outboundState.mu.Lock()
if err := c.outboundState.admitReserved(frame); err != nil {
c.outboundState.mu.Unlock()
t.Fatalf("admit result: %v", err)
}
c.outboundState.mu.Unlock()
ledger := newRPCExecutionLedger(time.Now, rpcExecutionLedgerCapacity{
maxPending: 16, maxPendingPerAuth: 16,
globalMaxEntries: 16, authMaxEntries: 16, sessionMaxEntries: 16,
replayStore: manager,
})
claim, err := ledger.Acquire(authKeyID, 43, 88)
if err != nil || claim.state != rpcResultAcquireOwner {
t.Fatalf("acquire owner = %#v, %v", claim, err)
}
claim.owner.CompleteExecution(true)
claim.owner.HandOff()
ledger.Complete(authKeyID, 43, 88, &encodedOutboundMessage{body: body, typeID: proto.ResultTypeID, reqMsgID: 88}, true)
key := rpcExecutionKey{authKeyID: authKeyID, sessionID: 43, reqMsgID: 88}
shard := ledger.shard(key)
shard.mu.Lock()
entry := shard.byKey[key].Value.(*rpcExecutionReceipt)
if entry.unavailable {
shard.mu.Unlock()
t.Fatal("logical outbox receipt was marked unavailable")
}
shard.mu.Unlock()
if got := ledger.receiptBudgetBytes(); got != rpcExecutionReceiptBudgetBytes {
t.Fatalf("receipt budget bytes = %d, want %d", got, rpcExecutionReceiptBudgetBytes)
}
c.outboundState.mu.Lock()
acked := c.outboundState.ack([]int64{201})
c.outboundState.mu.Unlock()
if len(acked) != 1 || acked[0] != 88 {
t.Fatalf("acked request ids = %v", acked)
}
if !ledger.Acknowledge(authKeyID, 43, 88) {
t.Fatal("ledger did not observe ACK")
}
if got := budget.snapshot(); got != 0 {
t.Fatalf("tracked payload bytes after ACK = %d", got)
}
shard.mu.Lock()
_, retained := shard.byKey[key]
shard.mu.Unlock()
if retained {
t.Fatal("ACK retained a completed receipt")
}
}
func TestLogicalSessionCapacityNeverEvictsUnackedFrame(t *testing.T) {
budget := newOutboundTrackedBudget(1024)
state := newOutboundStateWithLimits(budget, 1, 16)
firstBody := make([]byte, 8)
if !budget.reserve(len(firstBody)) {
t.Fatal("reserve first body")
}
first := &outboundFrame{
msgID: 1, seqNo: 1, typeID: proto.ResultTypeID, body: firstBody,
reqMsgID: 11, reservedBytes: len(firstBody), reservationBudget: budget,
}
if err := state.admitReserved(first); err != nil {
t.Fatalf("admit first: %v", err)
}
second := &outboundFrame{msgID: 2, seqNo: 3, typeID: proto.ResultTypeID, body: make([]byte, 8), reqMsgID: 12}
if err := state.admitReserved(second); err == nil {
t.Fatal("capacity admitted a second unacknowledged frame")
}
if state.pending[first.msgID] != first || state.byRequest[first.reqMsgID] != first.msgID {
t.Fatal("capacity failure evicted or rewired the existing frame")
}
state.releaseAll()
}
func TestLogicalSessionDestroyReleasesPayloadAndCompletedReceipt(t *testing.T) {
s := New(Options{})
authKeyID := [8]byte{9, 1}
c := &Conn{authKeyID: authKeyID, sessionID: 91}
claim, err := s.rpcResults.Acquire(authKeyID, 91, 901)
if err != nil || claim.owner == nil {
t.Fatalf("acquire owner: %#v err=%v", claim, err)
}
claim.owner.CompleteExecution(true)
storeLogicalRPCResultForTest(t, s, c, 901, &encodedOutboundMessage{
body: make([]byte, 32), typeID: proto.ResultTypeID, reqMsgID: 901,
})
if _, ok := s.rpcResults.Replay(authKeyID, 91, 901); !ok {
t.Fatal("logical result fixture was not published")
}
if removed := s.conns.DestroySessionForAuthKey(authKeyID, 91); removed {
t.Fatal("construction-only logical session unexpectedly reported active")
}
if _, ok := s.rpcResults.Replay(authKeyID, 91, 901); ok {
t.Fatal("destroy retained completed receipt")
}
if got := s.outboundTrackedBudget.snapshot(); got != 0 {
t.Fatalf("destroy retained %d payload bytes", got)
}
}
func TestBusinessAuthRevocationReleasesOfflineTempLogicalSession(t *testing.T) {
s := New(Options{})
rawAuthKeyID := [8]byte{9, 2}
businessAuthKeyID := [8]byte{9, 3}
c := &Conn{authKeyID: rawAuthKeyID, sessionID: 92}
c.SetBusinessAuthKeyID(businessAuthKeyID)
claim, err := s.rpcResults.Acquire(rawAuthKeyID, 92, 902)
if err != nil || claim.owner == nil {
t.Fatalf("acquire owner: %#v err=%v", claim, err)
}
claim.owner.CompleteExecution(true)
storeLogicalRPCResultForTest(t, s, c, 902, &encodedOutboundMessage{
body: make([]byte, 48), typeID: proto.ResultTypeID, reqMsgID: 902,
})
if closed := s.conns.CloseSessionsForBusinessAuthKey(businessAuthKeyID); closed != 0 {
t.Fatalf("offline logical revocation closed %d physical conns", closed)
}
if _, ok := s.rpcResults.Replay(rawAuthKeyID, 92, 902); ok {
t.Fatal("business auth revocation retained temp-key receipt")
}
if got := s.outboundTrackedBudget.snapshot(); got != 0 {
t.Fatalf("business auth revocation retained %d payload bytes", got)
}
}

14
internal/mtprotoedge/login_e2e_test.go
View file

@ -188,8 +188,8 @@ func TestLoginRegisterFlow(t *testing.T) {
// 注意：client.Run 回调里 t.Fatalf 只会杀当前 goroutine、测试主协程
// 会等到 ctx 超时——断言失败用 return fmt.Errorf 让 Run 立即返回。
cfg, ok := appConfig.(*tg.HelpAppConfig)
if !ok || cfg.Hash == 0 || cfg.Hash == seedAppConfigHash {
return fmt.Errorf("help.getAppConfig = %T %+v, want authenticated overlay hash distinct from seed=%d", appConfig, appConfig, seedAppConfigHash)
if !ok || cfg.Hash != seedAppConfigHash {
return fmt.Errorf("help.getAppConfig = %T %+v, want seeded base hash=%d", appConfig, appConfig, seedAppConfigHash)
}
object, ok := cfg.Config.(*tg.JSONObject)
if !ok {
@ -199,15 +199,11 @@ func TestLoginRegisterFlow(t *testing.T) {
for _, item := range object.Value {
values[item.Key] = item.Value
}
for _, key := range []string{"freeze_since_date", "freeze_until_date"} {
value, ok := values[key].(*tg.JSONNumber)
if !ok || value.Value != 0 {
return fmt.Errorf("help.getAppConfig %s = %T %+v, want zero clear value", key, values[key], values[key])
for _, key := range []string{"freeze_since_date", "freeze_until_date", "freeze_appeal_url"} {
if _, ok := values[key]; ok {
return fmt.Errorf("help.getAppConfig unexpectedly included inactive freeze field %s", key)
}
}
if value, ok := values["freeze_appeal_url"].(*tg.JSONString); !ok || value.Value != "" {
return fmt.Errorf("help.getAppConfig freeze_appeal_url = %T %+v, want empty clear value", values["freeze_appeal_url"], values["freeze_appeal_url"])
}
if value, ok := values["quote_length_max"].(*tg.JSONNumber); !ok || value.Value != 1024 {
return fmt.Errorf("help.getAppConfig lost seeded base config: quote_length_max=%T %+v", values["quote_length_max"], values["quote_length_max"])
}

15
internal/mtprotoedge/login_email_e2e_test.go
View file

@ -51,7 +51,7 @@ func TestLoginEmailEndToEnd(t *testing.T) {
dc = 2
phone = "+8613800138777"
wantPhone = "8613800138777"
code = "12345"
devCode = "12345"
email = "owner@example.com"
wantMask = "o***r@example.com"
)
@ -80,9 +80,10 @@ func TestLoginEmailEndToEnd(t *testing.T) {
accountService := account.NewService(passwordStore,
account.WithUsers(userStore),
account.WithLoginEmailVerification(codeStore, emailSender, 5*time.Minute, 5, 6))
authService := auth.NewService(userStore, authzStore, codeStore, authKeyStore, memory.NewTempAuthKeyBindingStore(authKeyStore), code,
authService := auth.NewService(userStore, authzStore, codeStore, authKeyStore, memory.NewTempAuthKeyBindingStore(authKeyStore), devCode,
auth.WithLoginMessages(messageStore, dialogStore),
auth.WithLoginCodeDelivery(memory.NewLoginCodeDeliveryStore(messageStore, updateEventStore)),
auth.WithPhoneCodeDelivery(emailSender, 5),
auth.WithPasswords(passwordStore),
auth.WithLoginEmail(auth.LoginEmailOptions{
Enabled: true,
@ -236,11 +237,15 @@ func TestLoginEmailEndToEnd(t *testing.T) {
return err
}
sentCode := sent.(*tg.AuthSentCode)
if _, ok := sentCode.Type.(*tg.AuthSentCodeTypeEmailCode); !ok {
emailType, ok := sentCode.Type.(*tg.AuthSentCodeTypeEmailCode)
if !ok {
return fmt.Errorf("pre-reset sendCode type = %T, want email code", sentCode.Type)
}
if _, ok := emailType.GetResetAvailablePeriod(); !ok {
return fmt.Errorf("pre-reset email code omitted reset_available_period with real SMS sender")
}
// 重置登录邮箱：返回一个新的手机验证码 sentCode（sentCodeTypeApp）。
// 重置登录邮箱：真实 SMS sender 签发随机手机验证码并返回 sentCodeTypeSms。
resetRes, err := raw.AuthResetLoginEmail(ctx, &tg.AuthResetLoginEmailRequest{PhoneNumber: phone, PhoneCodeHash: sentCode.PhoneCodeHash})
if err != nil {
return err
@ -253,7 +258,7 @@ func TestLoginEmailEndToEnd(t *testing.T) {
return fmt.Errorf("resetLoginEmail sentCode type = %T, want *tg.AuthSentCodeTypeSMS (back to phone, real sender configured)", resetSent.Type)
}
// 用手机验证码完成登录（真实投递的随机码，不是固定 dev code）。
// 用 sender 捕获的动态手机验证码完成登录，禁止回退固定 development code。
signInRes, err := raw.AuthSignIn(ctx, &tg.AuthSignInRequest{PhoneNumber: phone, PhoneCodeHash: resetSent.PhoneCodeHash, PhoneCode: phoneSender.code})
if err != nil {
return err

12
internal/mtprotoedge/metrics.go
View file

@ -2,8 +2,8 @@ package mtprotoedge
import "time"
// Metrics 接收连接层运行指标。实现可对接 Prometheus 等监控系统；
// 默认 NopMetrics（零开销）。第一阶段仅预留钩子，正式指标后续接入。
// Metrics 接收连接层运行指标。生产入口接入有界 Prometheus exporter；
// 其它 embedder 可继续使用 NopMetrics（零开销）。
type Metrics interface {
// ConnOpened 在接受一个连接时调用。
ConnOpened()
@ -39,6 +39,14 @@ type RPCResultMetrics interface {
RPCResultDelivered(method string, egressLatency time.Duration, wireBytes int, err error)
}
// LogicalOutboxMetrics observes the sole owner of unacknowledged server frames.
// It is intentionally optional: embedders can keep the small Metrics surface,
// while production capacity tests can distinguish physical delivery from the
// later client ACK that actually releases retained bytes.
type LogicalOutboxMetrics interface {
LogicalOutboxAcknowledged(bytes int, retainedFor time.Duration, rpcResult bool)
}
// ConnectionIntakeMetrics is an optional extension for the pre-session
// connection pipeline. stage is one of raw_accept, mux_sniff, mux_delivery,
// transport_dispatch, transport_promote, or first_frame; outcome is a bounded

290
internal/mtprotoedge/outbound.go
View file

@ -50,8 +50,8 @@ const (
maxTrackedServerMsgIDs = 4096
maxTrackedAckedMsgIDs = 1024
// maxTrackedServerBytes 是 pending（已发送待 ack、用于 resend）总 body 字节上限。
// 与 maxTrackedServerMsgIDs 并列：客户端从不 ack 时，大响应体按字节滚动丢弃，
// 防 pending 被「4096 条 × 大 body」撑爆。
// 与 maxTrackedServerMsgIDs 并列；到达任一上限后拒绝新可靠 frame，绝不滚动
// 丢弃尚未 ACK 的旧 frame。
maxTrackedServerBytes = 64 << 20 // 64 MiB
// Encrypted transport adds auth-key/msg-key, plaintext headers, randomized
// padding and codec framing. Reject before creating the two encryption buffers.
@ -214,6 +214,12 @@ type encodedOutboundMessage struct {
layer *outboundLayerBinding
compressed bool
uncompressedBytes int
// replayMsgID/replaySeqNo identify an existing logical-session frame. They
// are populated only by the receipt ledger's outbox lookup, never by a newly
// encoded result. A replay writes that exact frame instead of allocating a
// second payload owner or a new MTProto message identity.
replayMsgID int64
replaySeqNo int32
}
type rpcResultDeliveryState uint32
@ -263,7 +269,7 @@ type rpcResultDeliveryCoordinator struct {
// deferredToReplay is sticky while a successful initConnection retarget owns
// the logical hook. The ordinary source terminal may mark physical delivery,
// but only the alias restore barrier may claim the hook. On terminal alias
// failure the flag is released so a later completed-cache replay can retry.
// failure the flag is released so a later logical-outbox replay can retry.
deferredToReplay bool
}
@ -722,18 +728,25 @@ func cloneRPCResultForRequest(encoded *encodedOutboundMessage, reqMsgID int64, s
if shareDelivery {
delivery = encoded.delivery
}
var replayMsgID int64
var replaySeqNo int32
if reqMsgID == encoded.reqMsgID {
replayMsgID = encoded.replayMsgID
replaySeqNo = encoded.replaySeqNo
}
return &encodedOutboundMessage{
body: body, typeID: encoded.typeID, reqMsgID: reqMsgID,
priority: encoded.priority, delivery: delivery, compressed: encoded.compressed,
layer: encoded.layer, layerInvariant: encoded.layerInvariant,
uncompressedBytes: encoded.uncompressedBytes,
replayMsgID: replayMsgID, replaySeqNo: replaySeqNo,
}, nil
}
// cloneRPCResultForRequestReserved charges the target connection's retained-body
// budget before a retarget copy can exist. Even the same-req_id zero-copy case
// needs a reservation: the replay/rewrap owner may outlive cache eviction while
// it is queued, so its immutable body must remain independently pinned.
// needs a reservation: the replay/rewrap attempt may outlive ACK/removal of its
// source frame while queued, so its immutable body must remain independently pinned.
func (c *Conn) cloneRPCResultForRequestReserved(
encoded *encodedOutboundMessage,
reqMsgID int64,
@ -776,14 +789,20 @@ type outboundFrame struct {
// but their bytes must remain on the independent control budget for the full lifetime.
reservationBudget *outboundTrackedBudget
reqMsgID int64
priority outboundPriority
delivery *rpcResultDelivery
compressed bool
uncompressedBytes int
// layer is retained only for proactive session-bound frames so a later
// msg_resend_req cannot replay bytes from an obsolete profile epoch.
layer *outboundLayerBinding
sentAt time.Time
sends int
layer *outboundLayerBinding
layerInvariant bool
sentAt time.Time
sends int
}
type outboundState struct {
mu sync.Mutex
pending map[int64]*outboundFrame
order []int64
byRequest map[int64]int64
@ -793,6 +812,16 @@ type outboundState struct {
maxMessages int
maxBytes int
budget *outboundTrackedBudget
// sentContentMessages is logical-session state, not physical-connection
// state. Reserving a new content frame advances it before the first write so
// a failed write can be replayed with the same seq_no after reconnect.
sentContentMessages int32
lastMsgID int64
// persistent becomes true once this state is owned by a logical session.
// Directly constructed Conns may start with actor-local ownership and be
// adopted only from the terminal callback after a failed first write; the
// actor must not release that outbox while the logical session retains it.
persistent atomic.Bool
}
// outboundTrackedBudget 是 body/control/write 三类预算共用的原子 byte-budget primitive。
@ -1001,6 +1030,12 @@ func (c *Conn) startOutbound() {
c.outboundBulk = make(chan outboundOp, bulkSize)
c.outboundStop = make(chan struct{})
c.outboundDone = make(chan struct{})
// Publish the actor state before starting its goroutine. The pointer remains
// immutable for the physical Conn lifetime; logical-session adoption may only
// mark the existing state persistent.
if c.outboundState == nil {
c.outboundState = newOutboundState(c.outboundTrackedBudget)
}
go c.outboundLoop()
}
@ -1669,12 +1704,14 @@ func (c *Conn) endOutboundEnqueue() {
}
func (c *Conn) outboundLoop() {
state := newOutboundState(c.outboundTrackedBudget)
state := c.outboundState
ordinarySinceBulk := 0
defer func() {
// pending frames belong exclusively to this actor. Releasing after drain ensures no
// resend path can race the final budget return and no Conn body survives actor exit.
state.releaseAll()
// A Server Conn leaves pending frames with the logical session. Standalone
// construction/tests retain the old actor-local ownership boundary.
if !state.persistent.Load() {
state.releaseAll()
}
close(c.outboundDone)
}()
for {
@ -1790,29 +1827,46 @@ func (c *Conn) handleOutboundOp(state *outboundState, op outboundOp) {
if op.kind != outboundSend {
defer op.releaseReservation(state.budget)
}
// Physical generations may overlap briefly during activation fencing. The
// logical-session mutex preserves one writer/seq/outbox state machine across
// both actors without transferring payload ownership. Terminal callbacks run
// only after unlock: publication may look the frame up in this same outbox.
state.mu.Lock()
var (
result outboundResult
acked []outboundAcknowledgement
)
switch op.kind {
case outboundSend:
c.handleOutboundSend(state, op)
result.err = c.handleOutboundSend(state, op)
case outboundAck:
for _, reqMsgID := range state.ack(op.ids) {
if c.rpcResultAcked != nil {
c.rpcResultAcked(c, reqMsgID)
}
}
acked = state.ackWithDetails(op.ids)
case outboundQueryState:
op.finish(outboundResult{info: state.stateInfo(op.ids)})
result.info = state.stateInfo(op.ids)
case outboundResend:
info, err := c.handleOutboundResend(state, op.ctx, op.ids)
op.finish(outboundResult{info: info, err: err})
result.info, result.err = c.handleOutboundResend(state, op.ctx, op.ids)
case outboundResendByRequest:
resent, err := c.handleOutboundResendByRequest(state, op.ctx, op.reqMsgID)
op.finish(outboundResult{resent: resent, err: err})
result.resent, result.err = c.handleOutboundResendByRequest(state, op.ctx, op.reqMsgID)
default:
op.finish(outboundResult{err: fmt.Errorf("unknown outbound op %d", op.kind)})
result.err = fmt.Errorf("unknown outbound op %d", op.kind)
}
state.mu.Unlock()
for _, ack := range acked {
if metrics, ok := c.metrics.(LogicalOutboxMetrics); ok {
retainedFor := time.Duration(0)
if !ack.sentAt.IsZero() {
retainedFor = time.Since(ack.sentAt)
}
metrics.LogicalOutboxAcknowledged(ack.bytes, retainedFor, ack.reqMsgID != 0)
}
if ack.reqMsgID != 0 && c.rpcResultAcked != nil {
c.rpcResultAcked(c, ack.reqMsgID)
}
}
op.finish(result)
}
func (c *Conn) handleOutboundSend(state *outboundState, op outboundOp) {
func (c *Conn) handleOutboundSend(state *outboundState, op outboundOp) error {
var binding *outboundLayerBinding
if op.encoded != nil {
binding = op.encoded.layer
@ -1861,7 +1915,7 @@ func (c *Conn) handleOutboundSend(state *outboundState, op outboundOp) {
}
var frame *outboundFrame
if err == nil {
frame, err = c.buildFrame(op.ctx, op.msgType, op.msg, op.encoded)
frame, err = c.buildFrameWithState(op.ctx, op.msgType, op.msg, op.encoded, state)
}
// A profile-bound preparation can allocate a different body. Reserve the
// replacement before dropping the original prepared-body reservation. The
@ -1875,11 +1929,30 @@ func (c *Conn) handleOutboundSend(state *outboundState, op outboundOp) {
reserved = len(frame.body)
}
}
if err == nil && frame != nil && frameNeedsAck(frame.typeID) {
// The queue reservation is transferred to pending after write. A frame larger
// than the per-Conn resend ceiling is rejected before any bytes hit the wire.
needsAck := err == nil && frame != nil && frameNeedsAck(frame.typeID)
replaying := false
if needsAck && frame.replayMsgID() != 0 {
if existing := state.pending[frame.msgID]; existing != nil {
frame = existing
replaying = true
}
}
if needsAck && !replaying {
// Transfer the producer reservation into the logical outbox before the
// first physical write. A write failure therefore remains replayable on a
// replacement connection with the same msg_id/seq_no.
if len(frame.body) > maxTrackedServerBytes {
err = ErrOutboundTrackedBudget
} else {
frame.reservedBytes = reserved
frame.reservationBudget = reservationBudget
if admitErr := state.admitReserved(frame); admitErr != nil {
frame.reservedBytes = 0
frame.reservationBudget = nil
err = admitErr
} else {
reserved = 0
}
}
}
if errors.Is(err, ErrOutboundTrackedBudget) {
@ -1894,18 +1967,6 @@ func (c *Conn) handleOutboundSend(state *outboundState, op outboundOp) {
if err == nil {
err = c.writeFrame(op.ctx, frame)
}
if err == nil && frame != nil && frameNeedsAck(frame.typeID) {
// 写成功后才提交 content seq_no 递增（peekSeqNo 已按当前计数算好本帧 seq_no）。
c.commitContentSeqNo()
frame.reservedBytes = reserved
frame.reservationBudget = reservationBudget
reserved = 0
if dropped := state.addReserved(frame); dropped > 0 {
for i := 0; i < dropped; i++ {
c.metrics.OutboundDropped("tracked_queue_overflow")
}
}
}
queueWait := time.Since(op.enqueuedAt)
bytes := 0
typeID := uint32(0)
@ -1914,7 +1975,7 @@ func (c *Conn) handleOutboundSend(state *outboundState, op outboundOp) {
typeID = frame.typeID
}
c.metrics.OutboundSend(typeID, queueWait, bytes, err)
op.finish(outboundResult{err: err})
return err
}
func (c *Conn) handleOutboundResend(state *outboundState, ctx context.Context, ids []int64) ([]byte, error) {
@ -2141,6 +2202,16 @@ func (c *Conn) ensureOutboundControlTrackedBudget() *outboundTrackedBudget {
}
func (c *Conn) buildFrame(ctx context.Context, t proto.MessageType, msg bin.Encoder, encoded *encodedOutboundMessage) (*outboundFrame, error) {
return c.buildFrameWithState(ctx, t, msg, encoded, c.outboundState)
}
func (c *Conn) buildFrameWithState(
ctx context.Context,
t proto.MessageType,
msg bin.Encoder,
encoded *encodedOutboundMessage,
state *outboundState,
) (*outboundFrame, error) {
if encoded == nil {
var err error
encoded, err = encodeOutboundMessage(msg)
@ -2155,14 +2226,29 @@ func (c *Conn) buildFrame(ctx context.Context, t proto.MessageType, msg bin.Enco
return nil, ErrOutboundLayerBindingRequired
}
content := frameNeedsAck(encoded.typeID)
msgID := c.msgID.New(t)
msgID := encoded.replayMsgID
seqNo := encoded.replaySeqNo
if msgID == 0 {
msgID = c.msgID.New(t)
if state != nil {
msgID = state.reserveMsgID(msgID)
seqNo = state.peekSeqNo(content)
} else {
seqNo = c.peekSeqNo(content)
}
}
return &outboundFrame{
msgID: msgID,
seqNo: c.peekSeqNo(content),
typeID: encoded.typeID,
body: encoded.body,
reqMsgID: encoded.reqMsgID,
layer: encoded.layer,
msgID: msgID,
seqNo: seqNo,
typeID: encoded.typeID,
body: encoded.body,
reqMsgID: encoded.reqMsgID,
layer: encoded.layer,
layerInvariant: encoded.layerInvariant,
priority: encoded.priority,
delivery: encoded.delivery,
compressed: encoded.compressed,
uncompressedBytes: encoded.uncompressedBytes,
}, nil
}
@ -2545,7 +2631,18 @@ func outboundRequestMsgID(msg bin.Encoder) int64 {
// addReserved 接管调用方已经取得的全局 body 预算。pending 的每个元素恰好对应一份
// reservation；后续只有 removePending/releaseAll 能归还。
func (s *outboundState) addReserved(frame *outboundFrame) int {
func (s *outboundState) admitReserved(frame *outboundFrame) error {
if _, exists := s.pending[frame.msgID]; exists {
return fmt.Errorf("mtprotoedge: duplicate outbound msg_id inserted into resend tracking")
}
if len(s.pending) >= s.maxMessages || s.totalBytes > s.maxBytes-len(frame.body) {
return ErrOutboundTrackedBudget
}
s.insertReserved(frame)
return nil
}
func (s *outboundState) insertReserved(frame *outboundFrame) {
if _, exists := s.pending[frame.msgID]; exists {
panic("mtprotoedge: duplicate outbound msg_id inserted into resend tracking")
}
@ -2561,28 +2658,57 @@ func (s *outboundState) addReserved(frame *outboundFrame) int {
}
s.byRequest[frame.reqMsgID] = frame.msgID
}
if frameNeedsAck(frame.typeID) {
s.sentContentMessages++
}
}
// addReserved is retained as a focused-test helper. Production uses
// admitReserved and never evicts an unacknowledged frame.
func (s *outboundState) addReserved(frame *outboundFrame) int {
s.insertReserved(frame)
return s.shrinkPending()
}
type outboundAcknowledgement struct {
reqMsgID int64
bytes int
sentAt time.Time
}
func (s *outboundState) ack(ids []int64) []int64 {
var requestIDs []int64
details := s.ackWithDetails(ids)
requestIDs := make([]int64, 0, len(details))
for _, detail := range details {
if detail.reqMsgID != 0 {
requestIDs = append(requestIDs, detail.reqMsgID)
}
}
return requestIDs
}
func (s *outboundState) ackWithDetails(ids []int64) []outboundAcknowledgement {
var acknowledged []outboundAcknowledgement
for _, id := range ids {
frame, ok := s.pending[id]
if !ok {
continue
}
if frame.reqMsgID != 0 {
requestIDs = append(requestIDs, frame.reqMsgID)
detail := outboundAcknowledgement{
reqMsgID: frame.reqMsgID,
bytes: len(frame.body),
sentAt: frame.sentAt,
}
if !s.removePending(id) {
continue
}
s.markAcked(id)
acknowledged = append(acknowledged, detail)
}
if len(s.order) > s.maxMessages*2 {
s.compactOrder()
}
return requestIDs
return acknowledged
}
func (s *outboundState) stateInfo(ids []int64) []byte {
@ -2619,6 +2745,8 @@ func (s *outboundState) markAcked(id int64) {
}
}
// shrinkPending remains only for focused legacy state tests. Production
// admission never calls it: unacknowledged frames must not be silently evicted.
func (s *outboundState) shrinkPending() int {
dropped := 0
for (len(s.pending) > s.maxMessages || s.totalBytes > s.maxBytes) && len(s.order) > 0 {
@ -2652,6 +2780,8 @@ func (s *outboundState) removePending(id int64) bool {
}
func (s *outboundState) releaseAll() {
s.mu.Lock()
defer s.mu.Unlock()
for _, frame := range s.pending {
frame.body = nil
frame.releaseReservation(s.budget)
@ -2662,6 +2792,58 @@ func (s *outboundState) releaseAll() {
s.totalBytes = 0
}
func (s *outboundState) peekSeqNo(content bool) int32 {
seqNo := s.sentContentMessages * 2
if content {
seqNo++
}
return seqNo
}
func (s *outboundState) reserveMsgID(candidate int64) int64 {
if candidate <= s.lastMsgID {
candidate = s.lastMsgID + 4
}
for s.pending[candidate] != nil {
candidate += 4
}
s.lastMsgID = candidate
return candidate
}
func (s *outboundState) rpcResult(reqMsgID int64) (*encodedOutboundMessage, bool) {
if s == nil || reqMsgID == 0 {
return nil, false
}
s.mu.Lock()
defer s.mu.Unlock()
msgID := s.byRequest[reqMsgID]
frame := s.pending[msgID]
if frame == nil || frame.typeID != proto.ResultTypeID || len(frame.body) == 0 {
return nil, false
}
return &encodedOutboundMessage{
body: frame.body,
typeID: frame.typeID,
reqMsgID: frame.reqMsgID,
priority: frame.priority,
delivery: frame.delivery,
layer: frame.layer,
layerInvariant: frame.layerInvariant,
compressed: frame.compressed,
uncompressedBytes: frame.uncompressedBytes,
replayMsgID: frame.msgID,
replaySeqNo: frame.seqNo,
}, true
}
func (f *outboundFrame) replayMsgID() int64 {
if f == nil {
return 0
}
return f.msgID
}
func (f *outboundFrame) releaseReservation(defaultBudget *outboundTrackedBudget) {
if f == nil || f.reservedBytes <= 0 {
return

36
internal/mtprotoedge/outbound_test.go
View file

@ -29,6 +29,21 @@ type failAfterTransport struct {
last []byte
}
type acknowledgementCaptureMetrics struct {
NopMetrics
count atomic.Int64
bytes atomic.Int64
retainedNS atomic.Int64
rpcResult atomic.Bool
}
func (m *acknowledgementCaptureMetrics) LogicalOutboxAcknowledged(bytes int, retainedFor time.Duration, rpcResult bool) {
m.count.Add(1)
m.bytes.Add(int64(bytes))
m.retainedNS.Store(int64(retainedFor))
m.rpcResult.Store(rpcResult)
}
func TestRPCResultReplayAttemptHooksArePhysicalConnectionLocal(t *testing.T) {
const reqMsgID = int64(771)
base := &encodedOutboundMessage{
@ -544,7 +559,7 @@ func TestOutboundActorSerializesConcurrentSends(t *testing.T) {
conn, auth, cipher := dialHandshake(t, addr, dc, pub)
clientMsgID := proto.NewMessageIDGen(time.Now)
sendEncrypted(t, conn, cipher, auth, clientMsgID.New(proto.MessageFromClient), &mt.PingRequest{PingID: 1})
sendEncryptedWithSeq(t, conn, cipher, auth, clientMsgID.New(proto.MessageFromClient), 1, &mt.PingRequest{PingID: 1})
collectReplies(t, conn, cipher, auth.AuthKey, mt.MsgsAckTypeID)
freezeActiveTestSessionProfile(t, srv.Conns(), auth.AuthKey.ID, auth.SessionID, tlprofile.ProfileCanonical)
srv.Conns().SetReceivesUpdates(auth.SessionID, true)
@ -814,6 +829,8 @@ func TestOutboundTrackedBudgetAckAndCloseReturnExactly(t *testing.T) {
budget := newOutboundTrackedBudget(64)
tr := &failAfterTransport{}
c := newOutboundTestConn(t, tr, budget)
metrics := &acknowledgementCaptureMetrics{}
c.metrics = metrics
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
defer cancel()
body := exactTestUpdatesEncoded(t, c, make([]byte, 12))
@ -827,6 +844,9 @@ func TestOutboundTrackedBudgetAckAndCloseReturnExactly(t *testing.T) {
if err != nil {
t.Fatalf("decrypt frame: %v", err)
}
// Windows wall-clock resolution can otherwise make an immediate ACK look
// like zero retention even though sentAt was populated after the write.
time.Sleep(time.Millisecond)
c.AckServerMessages([]int64{data.MessageID})
deadline := time.Now().Add(time.Second)
for budget.snapshot() != 0 && time.Now().Before(deadline) {
@ -835,6 +855,18 @@ func TestOutboundTrackedBudgetAckAndCloseReturnExactly(t *testing.T) {
if got := budget.snapshot(); got != 0 {
t.Fatalf("tracked bytes after ack = %d, want 0", got)
}
if got := metrics.count.Load(); got != 1 {
t.Fatalf("logical ACK metric count = %d, want 1", got)
}
if got := metrics.bytes.Load(); got != 12 {
t.Fatalf("logical ACK metric bytes = %d, want 12", got)
}
if metrics.retainedNS.Load() <= 0 {
t.Fatal("logical ACK metric did not record positive retention")
}
if metrics.rpcResult.Load() {
t.Fatal("ordinary update ACK was classified as rpc_result")
}
})
t.Run("close", func(t *testing.T) {
@ -1034,7 +1066,7 @@ func TestOutboundResendAndAckState(t *testing.T) {
conn, auth, cipher := dialHandshake(t, addr, dc, pub)
clientMsgID := proto.NewMessageIDGen(time.Now)
sendEncrypted(t, conn, cipher, auth, clientMsgID.New(proto.MessageFromClient), &mt.PingRequest{PingID: 1})
sendEncryptedWithSeq(t, conn, cipher, auth, clientMsgID.New(proto.MessageFromClient), 1, &mt.PingRequest{PingID: 1})
collectReplies(t, conn, cipher, auth.AuthKey, mt.MsgsAckTypeID)
freezeActiveTestSessionProfile(t, srv.Conns(), auth.AuthKey.ID, auth.SessionID, tlprofile.ProfileCanonical)
srv.Conns().SetReceivesUpdates(auth.SessionID, true)

71
internal/mtprotoedge/production_compat_test.go
View file

@ -4,6 +4,7 @@ import (
"context"
"errors"
"sort"
"sync"
"time"
"go.uber.org/zap"
@ -25,10 +26,72 @@ func (c *Conn) ForceClose() {
c.waitOutboundShutdown()
}
// newRPCResultCache keeps older focused cache tests concise without exposing a
// second production constructor.
func newRPCResultCache(now func() time.Time) *rpcResultCache {
return newRPCResultCacheWithFlightLimit(now, rpcResultFlightDefaultMaxPending)
type rpcReplayStoreForTest struct {
mu sync.Mutex
results map[rpcExecutionKey]*encodedOutboundMessage
}
func newRPCReplayStoreForTest() *rpcReplayStoreForTest {
return &rpcReplayStoreForTest{results: make(map[rpcExecutionKey]*encodedOutboundMessage)}
}
func (s *rpcReplayStoreForTest) rpcResult(
authKeyID [8]byte,
sessionID, reqMsgID int64,
) (*encodedOutboundMessage, bool) {
if s == nil {
return nil, false
}
key := rpcExecutionKey{authKeyID: authKeyID, sessionID: sessionID, reqMsgID: reqMsgID}
s.mu.Lock()
result, ok := s.results[key]
s.mu.Unlock()
return result, ok
}
func (s *rpcReplayStoreForTest) put(
authKeyID [8]byte,
sessionID, reqMsgID int64,
encoded *encodedOutboundMessage,
) {
key := rpcExecutionKey{authKeyID: authKeyID, sessionID: sessionID, reqMsgID: reqMsgID}
s.mu.Lock()
s.results[key] = encoded
s.mu.Unlock()
}
func newRPCExecutionLedgerForTest(now func() time.Time, maxPending int) *rpcExecutionLedger {
return newRPCExecutionLedger(now, rpcExecutionLedgerCapacity{
maxPending: maxPending, maxPendingPerAuth: maxPending,
globalMaxEntries: rpcExecutionMaxEntries,
authMaxEntries: rpcExecutionMaxEntries, sessionMaxEntries: rpcExecutionMaxEntries,
replayStore: newRPCReplayStoreForTest(),
})
}
func newRPCExecutionLedgerForServerTest(s *Server, now func() time.Time, maxPending int) *rpcExecutionLedger {
if s == nil || s.conns == nil {
panic("newRPCExecutionLedgerForServerTest requires a server SessionManager")
}
return newRPCExecutionLedger(now, rpcExecutionLedgerCapacity{
maxPending: maxPending, maxPendingPerAuth: maxPending,
globalMaxEntries: rpcExecutionMaxEntries,
authMaxEntries: rpcExecutionMaxEntries, sessionMaxEntries: rpcExecutionMaxEntries,
replayStore: s.conns,
})
}
func (l *rpcExecutionLedger) completeReplayableForTest(
authKeyID [8]byte,
sessionID, reqMsgID int64,
encoded *encodedOutboundMessage,
) {
store, ok := l.replayStore.(*rpcReplayStoreForTest)
if !ok {
panic("completeReplayableForTest requires rpcReplayStoreForTest")
}
store.put(authKeyID, sessionID, reqMsgID, encoded)
l.Complete(authKeyID, sessionID, reqMsgID, encoded, true)
}
// Conns is a white-box test accessor. Production wires the shared manager

20
internal/mtprotoedge/rpc_capacity_error.go Normal file
View file

@ -0,0 +1,20 @@
package mtprotoedge
import "github.com/iamxvbaba/td/mt"
const (
rpcWorkerBusyErrorCode = 500
rpcWorkerBusyErrorMessage = "WORKER_BUSY_TOO_LONG_RETRY"
)
// rpcWorkerBusyError reports transient server admission pressure. A local
// worker/queue/ledger ceiling is not Telegram flood control: returning a 420
// FLOOD_WAIT would falsely blame the account and makes clients surface or cache
// a rate-limit state. Official clients classify 5xx as transient; TDLib and
// DrKlo additionally recognize WORKER_BUSY_TOO_LONG_RETRY and retry with delay.
func rpcWorkerBusyError() *mt.RPCError {
return &mt.RPCError{
ErrorCode: rpcWorkerBusyErrorCode,
ErrorMessage: rpcWorkerBusyErrorMessage,
}
}

4
internal/mtprotoedge/rpc_delivery_coordinator_test.go
View file

@ -279,7 +279,7 @@ func TestCachedReplacementReplayBarrierWaitsForLogicalHookDone(t *testing.T) {
firstConn := newOutboundTestConn(t, &collectingSessionTransport{}, newOutboundTrackedBudget(1<<20))
firstResult := make(chan error, 1)
go func() {
firstResult <- s.sendCachedRPCResultWithHook(context.Background(), firstConn, encoded, func() error {
firstResult <- s.sendReplayedRPCResultWithHook(context.Background(), firstConn, encoded, func() error {
if !order.CompareAndSwap(0, 1) {
return errors.New("first replacement restore ran out of order")
}
@ -300,7 +300,7 @@ func TestCachedReplacementReplayBarrierWaitsForLogicalHookDone(t *testing.T) {
secondReplacement := make(chan struct{})
secondResult := make(chan error, 1)
go func() {
secondResult <- s.sendCachedRPCResultWithHook(context.Background(), secondConn, encoded, func() error {
secondResult <- s.sendReplayedRPCResultWithHook(context.Background(), secondConn, encoded, func() error {
if !order.CompareAndSwap(2, 3) {
return errors.New("second replacement restore passed logical hook completion")
}

206
internal/mtprotoedge/rpc_execution_budget.go Normal file
View file

@ -0,0 +1,206 @@
package mtprotoedge
import (
"encoding/binary"
"hash/maphash"
"sync"
)
const rpcExecutionBudgetShards = 64
type rpcExecutionBudgetUsage struct {
entries int64
pending int64
}
type rpcExecutionSessionBudgetKey struct {
authKeyID [8]byte
sessionID int64
}
type rpcExecutionAuthBudgetShard struct {
mu sync.Mutex
usage map[[8]byte]rpcExecutionBudgetUsage
}
type rpcExecutionSessionBudgetShard struct {
mu sync.Mutex
usage map[rpcExecutionSessionBudgetKey]rpcExecutionBudgetUsage
}
// rpcExecutionFairBudget accounts one bounded ledger slot at global, auth and
// session scopes. Pending owner and completed receipt are the same ownership:
// Complete transfers the reservation; Abort/ACK/TTL return it.
type rpcExecutionFairBudget struct {
seed maphash.Seed
globalEntries *rpcResultFlightLimit
authLimit int64
sessionLimit int64
pendingPerAuth int64
authShards [rpcExecutionBudgetShards]rpcExecutionAuthBudgetShard
sessionShards [rpcExecutionBudgetShards]rpcExecutionSessionBudgetShard
}
type rpcExecutionBudgetReservation struct {
budget *rpcExecutionFairBudget
key rpcExecutionKey
pending bool
released bool
}
func newRPCExecutionFairBudget(
seed maphash.Seed,
globalEntries *rpcResultFlightLimit,
authLimit int64,
sessionLimit int64,
pendingPerAuth int,
) *rpcExecutionFairBudget {
b := &rpcExecutionFairBudget{
seed: seed, globalEntries: globalEntries, authLimit: authLimit,
sessionLimit: sessionLimit, pendingPerAuth: int64(pendingPerAuth),
}
for i := range b.authShards {
b.authShards[i].usage = make(map[[8]byte]rpcExecutionBudgetUsage)
b.sessionShards[i].usage = make(map[rpcExecutionSessionBudgetKey]rpcExecutionBudgetUsage)
}
return b
}
func (b *rpcExecutionFairBudget) reserveOwner(key rpcExecutionKey) *rpcExecutionBudgetReservation {
return b.reserve(key, true)
}
func (b *rpcExecutionFairBudget) reserveCompleted(key rpcExecutionKey) *rpcExecutionBudgetReservation {
return b.reserve(key, false)
}
func (b *rpcExecutionFairBudget) reserve(key rpcExecutionKey, pending bool) *rpcExecutionBudgetReservation {
if b == nil || b.globalEntries == nil {
return nil
}
authShard := b.authShard(key.authKeyID)
sessionKey := rpcExecutionSessionBudgetKey{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]
allowed := withinRPCExecutionBudget(authUsage.entries, 1, b.authLimit) &&
withinRPCExecutionBudget(sessionUsage.entries, 1, b.sessionLimit)
if pending {
allowed = allowed && withinRPCExecutionBudget(authUsage.pending, 1, b.pendingPerAuth)
}
if !allowed || !b.globalEntries.reserve() {
sessionShard.mu.Unlock()
authShard.mu.Unlock()
return nil
}
authUsage.entries++
sessionUsage.entries++
if pending {
authUsage.pending++
}
authShard.usage[key.authKeyID] = authUsage
sessionShard.usage[sessionKey] = sessionUsage
sessionShard.mu.Unlock()
authShard.mu.Unlock()
return &rpcExecutionBudgetReservation{budget: b, key: key, pending: pending}
}
func withinRPCExecutionBudget(used, delta, limit int64) bool {
return delta >= 0 && limit > 0 && used >= 0 && used <= limit-delta
}
func (r *rpcExecutionBudgetReservation) releasePending() {
if r == nil || r.budget == nil || r.released || !r.pending {
return
}
b := r.budget
shard := b.authShard(r.key.authKeyID)
shard.mu.Lock()
usage, ok := shard.usage[r.key.authKeyID]
if !ok || usage.pending < 1 {
shard.mu.Unlock()
panic("mtprotoedge: rpc execution per-auth pending budget underflow")
}
usage.pending--
shard.usage[r.key.authKeyID] = usage
r.pending = false
shard.mu.Unlock()
}
func (r *rpcExecutionBudgetReservation) release() {
if r == nil || r.budget == nil || r.released {
return
}
b := r.budget
authShard := b.authShard(r.key.authKeyID)
sessionKey := rpcExecutionSessionBudgetKey{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 ||
(r.pending && authUsage.pending < 1) {
sessionShard.mu.Unlock()
authShard.mu.Unlock()
panic("mtprotoedge: rpc execution fair budget underflow")
}
authUsage.entries--
sessionUsage.entries--
if r.pending {
authUsage.pending--
}
if authUsage == (rpcExecutionBudgetUsage{}) {
delete(authShard.usage, r.key.authKeyID)
} else {
authShard.usage[r.key.authKeyID] = authUsage
}
if sessionUsage == (rpcExecutionBudgetUsage{}) {
delete(sessionShard.usage, sessionKey)
} else {
sessionShard.usage[sessionKey] = sessionUsage
}
r.released = true
r.pending = false
b.globalEntries.release()
sessionShard.mu.Unlock()
authShard.mu.Unlock()
}
func (b *rpcExecutionFairBudget) authSnapshot(authKeyID [8]byte) rpcExecutionBudgetUsage {
if b == nil {
return rpcExecutionBudgetUsage{}
}
shard := b.authShard(authKeyID)
shard.mu.Lock()
usage := shard.usage[authKeyID]
shard.mu.Unlock()
return usage
}
func (b *rpcExecutionFairBudget) sessionSnapshot(authKeyID [8]byte, sessionID int64) rpcExecutionBudgetUsage {
if b == nil {
return rpcExecutionBudgetUsage{}
}
key := rpcExecutionSessionBudgetKey{authKeyID: authKeyID, sessionID: sessionID}
shard := b.sessionShard(key)
shard.mu.Lock()
usage := shard.usage[key]
shard.mu.Unlock()
return usage
}
func (b *rpcExecutionFairBudget) authShard(authKeyID [8]byte) *rpcExecutionAuthBudgetShard {
index := maphash.Bytes(b.seed, authKeyID[:]) & (rpcExecutionBudgetShards - 1)
return &b.authShards[index]
}
func (b *rpcExecutionFairBudget) sessionShard(key rpcExecutionSessionBudgetKey) *rpcExecutionSessionBudgetShard {
var raw [16]byte
copy(raw[:8], key.authKeyID[:])
binary.LittleEndian.PutUint64(raw[8:], uint64(key.sessionID))
index := maphash.Bytes(b.seed, raw[:]) & (rpcExecutionBudgetShards - 1)
return &b.sessionShards[index]
}

490
internal/mtprotoedge/rpc_execution_ledger.go Normal file
View file

@ -0,0 +1,490 @@
package mtprotoedge
import (
"container/list"
"context"
"encoding/binary"
"hash/maphash"
"sync"
"sync/atomic"
"time"
)
const (
// A valid client msg_id can be up to five minutes old or thirty seconds in
// the future. The extra second covers scheduler and boundary jitter. This is
// a no-ACK execution-receipt horizon, not a payload retention policy.
rpcExecutionReceiptTTL = 331 * time.Second
rpcExecutionMaxEntries = 1 << 18
rpcExecutionAuthMaxEntries = 1 << 15
rpcExecutionSessionMaxEntries = 1 << 14
rpcExecutionPendingPerAuth = 1 << 11
// Receipts contain only fixed-shape identity/outcome metadata. This
// conservative charge covers the receipt, list node, map bucket share and
// reservation bookkeeping. Payload bytes are accounted exclusively by the
// logical-session outbox.
rpcExecutionReceiptBudgetBytes = 384
// The complete replay identity is hashed with an instance-random seed. The
// shard count is a power of two.
rpcExecutionLedgerShards = 16
)
type rpcExecutionKey struct {
authKeyID [8]byte
sessionID int64
reqMsgID int64
}
// rpcReplayStore is the sole source of retained rpc_result payloads. The
// production implementation is SessionManager's logical-session outbox.
type rpcReplayStore interface {
rpcResult(authKeyID [8]byte, sessionID, reqMsgID int64) (*encodedOutboundMessage, bool)
}
type rpcExecutionReceipt struct {
key rpcExecutionKey
expiresAt time.Time
identity rpcResultRequestIdentity
admissionSeq uint64
executionKnown bool
executionOK bool
// acknowledged exists only for the ACK-before-Complete race. A completed
// receipt is removed immediately when the ACK wins.
acknowledged bool
// unavailable is a bounded execution tombstone. It prevents a completed
// business operation from running again when no exact outbox frame exists.
unavailable bool
// The pending owner transfers this same global/auth/session entry
// reservation to the completed receipt.
reservation *rpcExecutionBudgetReservation
}
type rpcResultDependency struct {
waiter *rpcResultWaiter
completed bool
success bool
}
// rpcExecutionLedger owns request execution identity and duplicate
// coordination. It never owns or copies rpc_result payload bytes; replay bytes
// are resolved from replayStore only while the logical-session outbox retains
// the exact unacknowledged frame.
type rpcExecutionLedger struct {
shards [rpcExecutionLedgerShards]rpcExecutionLedgerShard
hashSeed maphash.Seed
reservedEntries rpcResultFlightLimit
receiptCount atomic.Int64
fairBudget *rpcExecutionFairBudget
flightLimit rpcResultFlightLimit
subscriberBudget *rpcResultSubscriberBudget
subscriberPerFlight int
replayStore rpcReplayStore
nextAdmissionSeq atomic.Uint64
activeAdmissions rpcAdmissionTracker
}
func (l *rpcExecutionLedger) stableAdmissionSafeFloor() uint64 {
if l == nil {
return 0
}
return l.activeAdmissions.stableSafeFloor(&l.nextAdmissionSeq)
}
type rpcExecutionLedgerShard struct {
mu sync.Mutex
now func() time.Time
ttl time.Duration
receiptCount *atomic.Int64
order *list.List
byKey map[rpcExecutionKey]*list.Element
// In-flight owners are independent from receipt TTL and cannot disappear
// under completed-receipt pressure.
pending map[rpcExecutionKey]*rpcResultFlight
}
type rpcExecutionLedgerCapacity struct {
maxPending int
maxPendingPerAuth int
globalMaxEntries int
authMaxEntries int
sessionMaxEntries int
subscriberMaxGlobal int
subscriberMaxAuth int
subscriberMaxSession int
subscriberMaxPerFlight int
replayStore rpcReplayStore
}
func newRPCExecutionLedger(now func() time.Time, capacity rpcExecutionLedgerCapacity) *rpcExecutionLedger {
if now == nil {
now = time.Now
}
if capacity.replayStore == nil {
panic("mtprotoedge: rpc execution ledger requires a replay store")
}
if capacity.maxPending <= 0 {
capacity.maxPending = rpcResultFlightDefaultMaxPending
}
if capacity.maxPendingPerAuth <= 0 {
capacity.maxPendingPerAuth = capacity.maxPending
}
if capacity.globalMaxEntries <= 0 {
capacity.globalMaxEntries = rpcExecutionMaxEntries
}
if capacity.authMaxEntries <= 0 {
capacity.authMaxEntries = capacity.globalMaxEntries
}
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
}
l := &rpcExecutionLedger{hashSeed: maphash.MakeSeed(), replayStore: capacity.replayStore}
l.reservedEntries.max = int64(capacity.globalMaxEntries)
l.flightLimit.max = int64(capacity.maxPending)
l.fairBudget = newRPCExecutionFairBudget(
l.hashSeed,
&l.reservedEntries,
int64(capacity.authMaxEntries),
int64(capacity.sessionMaxEntries),
capacity.maxPendingPerAuth,
)
l.subscriberBudget = newRPCResultSubscriberBudget(
l.hashSeed,
capacity.subscriberMaxGlobal,
capacity.subscriberMaxAuth,
capacity.subscriberMaxSession,
)
l.subscriberPerFlight = capacity.subscriberMaxPerFlight
for i := range l.shards {
s := &l.shards[i]
s.now = now
s.ttl = rpcExecutionReceiptTTL
s.receiptCount = &l.receiptCount
s.order = list.New()
s.byKey = make(map[rpcExecutionKey]*list.Element)
s.pending = make(map[rpcExecutionKey]*rpcResultFlight)
}
return l
}
func (l *rpcExecutionLedger) shard(key rpcExecutionKey) *rpcExecutionLedgerShard {
return &l.shards[l.shardIndex(key)]
}
func (l *rpcExecutionLedger) shardIndex(key rpcExecutionKey) 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(l.hashSeed, raw[:]) & (rpcExecutionLedgerShards - 1)
}
// Replay resolves an immutable result descriptor from the logical-session
// outbox. A receipt hit without an outbox frame is never treated as permission
// to execute the business handler again.
func (l *rpcExecutionLedger) Replay(authKeyID [8]byte, sessionID, reqMsgID int64) (*encodedOutboundMessage, bool) {
if l == nil || reqMsgID == 0 {
return nil, false
}
key := rpcExecutionKey{authKeyID: authKeyID, sessionID: sessionID, reqMsgID: reqMsgID}
s := l.shard(key)
now := s.now()
s.mu.Lock()
elem := s.byKey[key]
if elem == nil {
s.mu.Unlock()
return nil, false
}
receipt := elem.Value.(*rpcExecutionReceipt)
if !receipt.expiresAt.After(now) {
s.removeElement(elem)
s.mu.Unlock()
return nil, false
}
if receipt.unavailable || receipt.acknowledged {
s.mu.Unlock()
return nil, false
}
s.mu.Unlock()
return l.replayStore.rpcResult(authKeyID, sessionID, reqMsgID)
}
// Acknowledge removes a completed receipt immediately. reqMsgID must already
// have been resolved from the outbound actor's trusted server-msg-id mapping.
// A pending flight records the ACK so a racing completion cannot resurrect a
// receipt after the outbox body has been released.
func (l *rpcExecutionLedger) Acknowledge(authKeyID [8]byte, sessionID, reqMsgID int64) bool {
if l == nil || reqMsgID == 0 {
return false
}
key := rpcExecutionKey{authKeyID: authKeyID, sessionID: sessionID, reqMsgID: reqMsgID}
s := l.shard(key)
s.mu.Lock()
s.expireLocked(s.now())
if elem := s.byKey[key]; elem != nil {
s.removeElement(elem)
s.mu.Unlock()
return true
}
if flight := s.pending[key]; flight != nil {
flight.acknowledged = true
s.mu.Unlock()
return true
}
s.mu.Unlock()
return false
}
// ObserveDependency returns a waiter for an admitted in-flight dependency, a
// completed execution outcome, or ok=false for an unknown request. It never
// creates execution ownership.
func (l *rpcExecutionLedger) ObserveDependency(authKeyID [8]byte, sessionID, reqMsgID int64) (rpcResultDependency, bool) {
if l == nil || reqMsgID == 0 {
return rpcResultDependency{}, false
}
key := rpcExecutionKey{authKeyID: authKeyID, sessionID: sessionID, reqMsgID: reqMsgID}
s := l.shard(key)
now := s.now()
s.mu.Lock()
defer s.mu.Unlock()
if elem := s.byKey[key]; elem != nil {
receipt := elem.Value.(*rpcExecutionReceipt)
if receipt.expiresAt.After(now) {
if !receipt.executionKnown {
return rpcResultDependency{}, false
}
return rpcResultDependency{completed: true, success: receipt.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{ledger: l, key: key, flight: flight}}, true
}
return rpcResultDependency{}, false
}
// Complete publishes terminal execution metadata and resolves current
// waiters. replayable is only a claim from the egress path; the ledger verifies
// that replayStore actually owns the exact frame before publishing a replayable
// receipt. encoded is passed transiently to joined waiters and is never stored.
func (l *rpcExecutionLedger) Complete(
authKeyID [8]byte,
sessionID, reqMsgID int64,
encoded *encodedOutboundMessage,
replayable bool,
) {
if l == nil || reqMsgID == 0 || encoded == nil {
return
}
if replayable {
_, replayable = l.replayStore.rpcResult(authKeyID, sessionID, reqMsgID)
}
if l.completeOnce(authKeyID, sessionID, reqMsgID, encoded, replayable) {
return
}
// An expired receipt in another shard may be the only global blocker.
l.expireReceipts()
_ = l.completeOnce(authKeyID, sessionID, reqMsgID, encoded, replayable)
}
// completeOnce returns false only when a cross-shard expiry reap may release
// the global reservation required by a defensive completion without an owner.
func (l *rpcExecutionLedger) completeOnce(
authKeyID [8]byte,
sessionID, reqMsgID int64,
encoded *encodedOutboundMessage,
replayable bool,
) bool {
key := rpcExecutionKey{authKeyID: authKeyID, sessionID: sessionID, reqMsgID: reqMsgID}
s := l.shard(key)
s.mu.Lock()
now := s.now()
s.expireLocked(now)
old := s.byKey[key]
flight := s.pending[key]
var oldReceipt *rpcExecutionReceipt
if old != nil {
oldReceipt = old.Value.(*rpcExecutionReceipt)
if flight == nil {
// Terminal publication is immutable. Late callbacks cannot extend TTL,
// replace replay identity or resurrect an acknowledged result.
s.mu.Unlock()
return true
}
}
identity, admissionSeq, executionKnown, executionOK, acknowledged := rpcResultFlightMetadataLocked(s, key)
var reservation *rpcExecutionBudgetReservation
switch {
case flight != nil:
reservation = flight.reservation
if reservation == nil {
s.mu.Unlock()
panic("mtprotoedge: pending rpc execution has no fair-budget reservation")
}
case oldReceipt != nil && oldReceipt.reservation != nil:
reservation = oldReceipt.reservation
default:
reservation = l.fairBudget.reserveCompleted(key)
if reservation == nil {
s.mu.Unlock()
return false
}
}
if old != nil {
s.unlinkElement(old)
if oldReceipt.reservation != nil && oldReceipt.reservation != reservation {
oldReceipt.reservation.release()
oldReceipt.reservation = nil
}
}
receipt := &rpcExecutionReceipt{
key: key,
expiresAt: now.Add(s.ttl),
identity: identity,
admissionSeq: admissionSeq,
executionKnown: executionKnown,
executionOK: executionOK,
acknowledged: acknowledged,
unavailable: !replayable,
reservation: reservation,
}
elem := s.order.PushBack(receipt)
s.byKey[key] = elem
s.incrementReceiptCount()
subscribers, executionSubscribers, terminalExecutionOK := l.completeRPCResultFlightLocked(s, key, encoded)
if acknowledged {
// ACK won before completion. Current subscribers still receive encoded,
// but no post-ACK receipt survives this critical section.
s.removeElement(elem)
}
s.mu.Unlock()
for _, subscriber := range subscribers {
subscriber(encoded, true)
}
for _, subscriber := range executionSubscribers {
subscriber(terminalExecutionOK)
}
return true
}
func rpcResultFlightMetadataLocked(s *rpcExecutionLedgerShard, key rpcExecutionKey) (
rpcResultRequestIdentity,
uint64,
bool,
bool,
bool,
) {
if flight := s.pending[key]; flight != nil {
return flight.identity, flight.admissionSeq, flight.executionDone, flight.executionOK, flight.acknowledged
}
return rpcResultRequestIdentity{}, 0, false, false, false
}
func (l *rpcExecutionLedger) expireReceipts() {
if l == nil {
return
}
for i := range l.shards {
s := &l.shards[i]
s.mu.Lock()
s.expireLocked(s.now())
s.mu.Unlock()
}
}
func (l *rpcExecutionLedger) receiptBudgetBytes() int64 {
if l == nil {
return 0
}
return l.receiptCount.Load() * rpcExecutionReceiptBudgetBytes
}
func (l *rpcExecutionLedger) Close() error { return nil }
func (l *rpcExecutionLedger) CloseContext(context.Context) error { return nil }
// forgetSession removes every terminal receipt for a destroyed logical
// session. SessionManager calls it only after physical producers converge.
func (l *rpcExecutionLedger) forgetSession(authKeyID [8]byte, sessionID int64) {
if l == nil {
return
}
for i := range l.shards {
s := &l.shards[i]
s.mu.Lock()
for elem := s.order.Front(); elem != nil; {
next := elem.Next()
receipt := elem.Value.(*rpcExecutionReceipt)
if receipt.key.authKeyID == authKeyID && receipt.key.sessionID == sessionID {
s.removeElement(elem)
}
elem = next
}
s.mu.Unlock()
}
}
func (s *rpcExecutionLedgerShard) expireLocked(now time.Time) {
for elem := s.order.Front(); elem != nil; {
next := elem.Next()
receipt := elem.Value.(*rpcExecutionReceipt)
if receipt.expiresAt.After(now) {
return
}
s.removeElement(elem)
elem = next
}
}
func (s *rpcExecutionLedgerShard) removeElement(elem *list.Element) {
receipt := s.unlinkElement(elem)
if receipt != nil && receipt.reservation != nil {
receipt.reservation.release()
receipt.reservation = nil
}
}
func (s *rpcExecutionLedgerShard) unlinkElement(elem *list.Element) *rpcExecutionReceipt {
if elem == nil {
return nil
}
receipt := elem.Value.(*rpcExecutionReceipt)
delete(s.byKey, receipt.key)
s.order.Remove(elem)
s.decrementReceiptCount()
return receipt
}
func (s *rpcExecutionLedgerShard) incrementReceiptCount() {
if s.receiptCount == nil {
panic("mtprotoedge: rpc execution receipt counter is unavailable")
}
s.receiptCount.Add(1)
}
func (s *rpcExecutionLedgerShard) decrementReceiptCount() {
if s.receiptCount == nil || s.receiptCount.Add(-1) < 0 {
panic("mtprotoedge: rpc execution receipt counter underflow")
}
}

346
internal/mtprotoedge/rpc_execution_ledger_test.go Normal file
View file

@ -0,0 +1,346 @@
package mtprotoedge
import (
"container/list"
"errors"
"sync"
"testing"
"time"
"unsafe"
)
func newRPCExecutionLedgerWithLimitsForTest(
now func() time.Time,
maxPending, maxPendingPerAuth, global, auth, session int,
) *rpcExecutionLedger {
return newRPCExecutionLedger(now, rpcExecutionLedgerCapacity{
maxPending: maxPending, maxPendingPerAuth: maxPendingPerAuth,
globalMaxEntries: global, authMaxEntries: auth, sessionMaxEntries: session,
replayStore: newRPCReplayStoreForTest(),
})
}
func TestRPCExecutionLedgerSessionCapacityIsolatesAnotherAuth(t *testing.T) {
ledger := newRPCExecutionLedgerWithLimitsForTest(time.Now, 8, 6, 8, 6, 2)
authA := [8]byte{0xa1}
authB := [8]byte{0xb1}
for i := 0; i < 2; i++ {
msgID := int64(1000 + i)
claim, err := ledger.Acquire(authA, 77, msgID)
if err != nil || claim.state != rpcResultAcquireOwner {
t.Fatalf("same-session admission %d = %#v, %v", i, claim, err)
}
ledger.completeReplayableForTest(authA, 77, msgID, &encodedOutboundMessage{body: []byte{1}})
}
if _, err := ledger.Acquire(authA, 77, 2000); !errors.Is(err, ErrRPCResultFlightCapacity) {
t.Fatalf("admission beyond session limit = %v, want capacity", err)
}
other, err := ledger.Acquire(authB, 88, 3000)
if err != nil || other.state != rpcResultAcquireOwner {
t.Fatalf("other auth blocked by full session: %#v, %v", other, err)
}
other.owner.Abort()
}
func TestRPCExecutionLedgerAuthCapacityIsolatesAnotherAuth(t *testing.T) {
ledger := newRPCExecutionLedgerWithLimitsForTest(time.Now, 8, 4, 8, 2, 2)
authA := [8]byte{0xa2}
authB := [8]byte{0xb2}
for i := 0; i < 2; i++ {
claim, err := ledger.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)
}
ledger.completeReplayableForTest(authA, int64(10+i), int64(100+i), &encodedOutboundMessage{body: []byte{1}})
}
if _, err := ledger.Acquire(authA, 12, 102); !errors.Is(err, ErrRPCResultFlightCapacity) {
t.Fatalf("same auth beyond limit = %v, want capacity", err)
}
other, err := ledger.Acquire(authB, 20, 200)
if err != nil || other.state != rpcResultAcquireOwner {
t.Fatalf("other auth blocked by auth A: %#v, %v", other, err)
}
other.owner.Abort()
}
func TestRPCExecutionLedgerPendingLimitIsAdditional(t *testing.T) {
ledger := newRPCExecutionLedgerWithLimitsForTest(time.Now, 6, 2, 12, 6, 4)
authA := [8]byte{0xa3}
authB := [8]byte{0xb3}
owners := make([]*rpcResultOwnerLease, 0, 3)
for i := 0; i < 2; i++ {
claim, err := ledger.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 := ledger.Acquire(authA, 3, 103); !errors.Is(err, ErrRPCResultFlightCapacity) {
t.Fatalf("third pending owner for auth A = %v, want capacity", err)
}
other, err := ledger.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 := ledger.fairBudget.authSnapshot(authA); usage != (rpcExecutionBudgetUsage{}) {
t.Fatalf("auth A budget after abort = %#v", usage)
}
}
func TestRPCExecutionLedgerReceiptLifecycleACKAndTTL(t *testing.T) {
now := time.Unix(1000, 0)
ledger := newRPCExecutionLedgerWithLimitsForTest(func() time.Time { return now }, 4, 4, 6, 5, 3)
auth := [8]byte{0xc1}
claim, err := ledger.Acquire(auth, 1, 101)
if err != nil || claim.state != rpcResultAcquireOwner {
t.Fatalf("owner = %#v, %v", claim, err)
}
if usage := ledger.fairBudget.authSnapshot(auth); usage.entries != 1 || usage.pending != 1 {
t.Fatalf("pending reservation = %#v", usage)
}
claim.owner.CompleteExecution(true)
ledger.completeReplayableForTest(auth, 1, 101, &encodedOutboundMessage{body: make([]byte, 8<<20)})
if ledger.flightLimit.snapshot() != 0 || ledger.receiptCount.Load() != 1 || ledger.reservedEntries.snapshot() != 1 {
t.Fatalf("terminal counts owner=%d receipt=%d reserved=%d", ledger.flightLimit.snapshot(), ledger.receiptCount.Load(), ledger.reservedEntries.snapshot())
}
if got := ledger.receiptBudgetBytes(); got != rpcExecutionReceiptBudgetBytes {
t.Fatalf("8 MiB result charged %d receipt bytes, want fixed %d", got, rpcExecutionReceiptBudgetBytes)
}
if !ledger.Acknowledge(auth, 1, 101) {
t.Fatal("ACK did not remove receipt")
}
if ledger.receiptCount.Load() != 0 || ledger.reservedEntries.snapshot() != 0 || ledger.receiptBudgetBytes() != 0 {
t.Fatal("ACK leaked receipt reservation")
}
second, err := ledger.Acquire(auth, 2, 201)
if err != nil || second.state != rpcResultAcquireOwner {
t.Fatalf("second owner = %#v, %v", second, err)
}
ledger.completeReplayableForTest(auth, 2, 201, &encodedOutboundMessage{body: []byte{1}})
now = now.Add(rpcExecutionReceiptTTL + time.Second)
if _, ok := ledger.Replay(auth, 2, 201); ok {
t.Fatal("expired receipt remained replayable")
}
if ledger.receiptCount.Load() != 0 || ledger.reservedEntries.snapshot() != 0 {
t.Fatal("TTL leaked receipt reservation")
}
}
func TestRPCExecutionLedgerACKBeforeCompleteDoesNotResurrectReceipt(t *testing.T) {
ledger := newRPCExecutionLedgerForTest(time.Now, 2)
auth := [8]byte{0xd1}
claim, err := ledger.Acquire(auth, 1, 101)
if err != nil || claim.state != rpcResultAcquireOwner {
t.Fatalf("owner = %#v, %v", claim, err)
}
joined, err := ledger.Acquire(auth, 1, 101)
if err != nil || joined.state != rpcResultAcquirePending {
t.Fatalf("join = %#v, %v", joined, err)
}
if !ledger.Acknowledge(auth, 1, 101) {
t.Fatal("ACK did not mark pending owner")
}
want := &encodedOutboundMessage{body: []byte{1}, reqMsgID: 101}
ledger.completeReplayableForTest(auth, 1, 101, want)
if got, ok, waitErr := joined.waiter.Wait(t.Context()); waitErr != nil || !ok || got != want {
t.Fatalf("joined waiter = %p/%v/%v", got, ok, waitErr)
}
if ledger.receiptCount.Load() != 0 || ledger.reservedEntries.snapshot() != 0 {
t.Fatal("ACK-before-complete resurrected receipt")
}
newClaim, err := ledger.Acquire(auth, 1, 101)
if err != nil || newClaim.state != rpcResultAcquireOwner {
t.Fatalf("post-ACK request did not get a fresh owner: %#v, %v", newClaim, err)
}
newClaim.owner.Abort()
}
func TestRPCExecutionLedgerUnavailableTombstonePreventsReexecution(t *testing.T) {
now := time.Unix(1000, 0)
ledger := newRPCExecutionLedgerWithLimitsForTest(func() time.Time { return now }, 2, 2, 4, 4, 4)
auth := [8]byte{0xe1}
claim, err := ledger.Acquire(auth, 1, 101)
if err != nil || claim.state != rpcResultAcquireOwner {
t.Fatalf("owner = %#v, %v", claim, err)
}
claim.owner.CompleteExecution(true)
ledger.Complete(auth, 1, 101, &encodedOutboundMessage{body: make([]byte, 8<<20)}, false)
if _, ok := ledger.Replay(auth, 1, 101); ok {
t.Fatal("unavailable tombstone masqueraded as replayable")
}
if _, err := ledger.Acquire(auth, 1, 101); !errors.Is(err, ErrRPCResultFlightCapacity) {
t.Fatalf("duplicate after unavailable completion = %v, want capacity", err)
}
if got := ledger.receiptBudgetBytes(); got != rpcExecutionReceiptBudgetBytes {
t.Fatalf("unavailable receipt budget = %d", got)
}
now = now.Add(rpcExecutionReceiptTTL + time.Second)
retry, err := ledger.Acquire(auth, 1, 101)
if err != nil || retry.state != rpcResultAcquireOwner {
t.Fatalf("admission after tombstone expiry = %#v, %v", retry, err)
}
retry.owner.Abort()
}
func TestRPCExecutionLedgerConcurrentReservationsNeverOvercommit(t *testing.T) {
const limit = 24
ledger := newRPCExecutionLedgerWithLimitsForTest(time.Now, limit, 4, limit, 8, 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 := ledger.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 := ledger.reservedEntries.snapshot(); got > limit || got != int64(len(owners)) {
t.Fatalf("reserved=%d owners=%d limit=%d", got, len(owners), limit)
}
for i := 0; i < 4; i++ {
auth := [8]byte{byte(i)}
usage := ledger.fairBudget.authSnapshot(auth)
if usage.entries > 8 || usage.pending > 4 {
t.Fatalf("auth %d overcommitted: %#v", i, usage)
}
}
for _, owner := range owners {
owner.Abort()
}
if ledger.reservedEntries.snapshot() != 0 {
t.Fatal("concurrent abort leaked reservations")
}
}
func TestRPCExecutionLedgerFullKeyHashSpreadsOneSession(t *testing.T) {
first := newRPCExecutionLedgerForTest(time.Now, 64)
second := newRPCExecutionLedgerForTest(time.Now, 64)
auth := [8]byte{1, 2, 3, 4, 5, 6, 7, 8}
seen := make(map[uint64]struct{})
differentInstance := false
for msgID := int64(1); msgID <= 256; msgID++ {
key := rpcExecutionKey{authKeyID: auth, sessionID: 99, reqMsgID: msgID}
firstIndex := first.shardIndex(key)
seen[firstIndex] = struct{}{}
if firstIndex != second.shardIndex(key) {
differentInstance = true
}
}
if len(seen) < rpcExecutionLedgerShards/2 {
t.Fatalf("one session used only %d/%d shards", len(seen), rpcExecutionLedgerShards)
}
if !differentInstance {
t.Fatal("two ledger instances used an identical shard stream")
}
}
func TestRPCExecutionLedgerForgetSessionReleasesReceipts(t *testing.T) {
ledger := newRPCExecutionLedgerForTest(time.Now, 8)
auth := [8]byte{0xf1}
for _, sessionID := range []int64{1, 1, 2} {
msgID := int64(100 + ledger.receiptCount.Load())
claim, err := ledger.Acquire(auth, sessionID, msgID)
if err != nil || claim.state != rpcResultAcquireOwner {
t.Fatalf("owner session=%d: %#v, %v", sessionID, claim, err)
}
ledger.completeReplayableForTest(auth, sessionID, msgID, &encodedOutboundMessage{body: []byte{1}})
}
ledger.forgetSession(auth, 1)
if got := ledger.receiptCount.Load(); got != 1 {
t.Fatalf("receipts after session forget = %d, want 1", got)
}
if _, ok := ledger.Replay(auth, 1, 100); ok {
t.Fatal("forgotten session remained replayable")
}
}
func TestRPCExecutionLedgerServerOptionsPropagateLimits(t *testing.T) {
s := New(Options{
RPCGlobalMaxTasks: 6,
RPCExecutionMaxEntries: 12,
RPCExecutionAuthMaxEntries: 8,
RPCExecutionSessionMaxEntries: 4,
RPCExecutionPendingPerAuth: 3,
})
if s.rpcResults.reservedEntries.max != 12 {
t.Fatalf("global option propagation = %d", s.rpcResults.reservedEntries.max)
}
budget := s.rpcResults.fairBudget
if budget.authLimit != 8 || budget.sessionLimit != 4 || budget.pendingPerAuth != 3 {
t.Fatalf("fair option propagation = auth:%d session:%d pending:%d", budget.authLimit, budget.sessionLimit, budget.pendingPerAuth)
}
}
func TestRPCExecutionLedgerServerOptionsFailFast(t *testing.T) {
base := Options{
RPCGlobalMaxTasks: 6,
RPCExecutionMaxEntries: 12,
RPCExecutionAuthMaxEntries: 8,
RPCExecutionSessionMaxEntries: 4,
RPCExecutionPendingPerAuth: 3,
}
tests := []struct {
name string
mutate func(*Options)
}{
{name: "entry hierarchy", mutate: func(o *Options) { o.RPCExecutionAuthMaxEntries = 13 }},
{name: "pending hierarchy", mutate: func(o *Options) { o.RPCExecutionPendingPerAuth = 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 execution options")
}
}()
_ = New(opts)
})
}
}
func TestRPCExecutionLedgerRequiresReplayStore(t *testing.T) {
defer func() {
if recover() == nil {
t.Fatal("ledger accepted nil replay store")
}
}()
_ = newRPCExecutionLedger(time.Now, rpcExecutionLedgerCapacity{})
}
func TestRPCExecutionReceiptBudgetCoversOwnedFixedStructures(t *testing.T) {
fixed := unsafe.Sizeof(rpcExecutionReceipt{}) +
unsafe.Sizeof(list.Element{}) +
unsafe.Sizeof(rpcExecutionBudgetReservation{})
if fixed > rpcExecutionReceiptBudgetBytes {
t.Fatalf("fixed receipt structures use %d bytes, budget charge is %d", fixed, rpcExecutionReceiptBudgetBytes)
}
}

194
internal/mtprotoedge/rpc_flight.go
View file

@ -18,11 +18,12 @@ var (
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")
ErrRPCResultReplayUnavailable = errors.New("mtproto rpc result replay payload is unavailable")
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
// profile from inside the ledger 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 {
@ -50,8 +51,8 @@ type rpcResultRequestIdentity struct {
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.
// Service-message callers carry no API request identity. Exact API callers
// must always match the winner's full prepared identity.
return true
}
return i.valid && i.exact == requested.exact
@ -61,6 +62,7 @@ type rpcResultAcquireState uint8
const (
rpcResultAcquireCompleted rpcResultAcquireState = iota + 1
rpcResultAcquireAcknowledged
rpcResultAcquirePending
rpcResultAcquireOwner
)
@ -70,8 +72,10 @@ const (
//
// Exactly one state-specific field is non-nil:
// - completed: encoded contains the immutable completed rpc_result;
// - acknowledged: ACK won while the owner was still pending, so the duplicate
// must be ACK-only until that owner completes;
// - pending: waiter joins the already-running owner;
// - owner: owner must eventually complete through rpcResultCache.Put or Abort.
// - owner: owner must eventually complete through ledger Complete or Abort.
type rpcResultAcquire struct {
state rpcResultAcquireState
admissionSeq uint64
@ -82,8 +86,8 @@ type rpcResultAcquire struct {
executionOK bool
}
// rpcResultFlight is not part of the completed cache TTL lifecycle. Its
// done channel is closed exactly once while holding the owning cache shard lock;
// rpcResultFlight is not part of the completed receipt TTL lifecycle. Its
// done channel is closed exactly once while holding the owning ledger shard lock;
// channel close publishes encoded/ok to all waiters without a waiter goroutine.
type rpcResultFlight struct {
done chan struct{}
@ -93,24 +97,29 @@ type rpcResultFlight struct {
executionDone bool
executionOK bool
executionSubscribers []func(bool)
// acknowledged is set only from the sole outbound actor's server-msg-id to
// request-msg-id mapping. It can win the race with asynchronous completion;
// publication then drops the receipt while still resolving waiters with the
// immutable result that was already written on the wire.
acknowledged 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
// reservation owns one entry at global, raw-auth and session scopes.
// Complete transfers it to a receipt/tombstone; Abort releases it.
reservation *rpcExecutionBudgetReservation
}
type rpcResultWaiter struct {
cache *rpcResultCache
key rpcResultCacheKey
ledger *rpcExecutionLedger
key rpcExecutionKey
flight *rpcResultFlight
}
// Wait blocks until the owner publishes through Put, aborts, or ctx expires.
// Wait blocks until the owner publishes through Complete, aborts, or ctx expires.
// ok=false with err=nil means the owner aborted without a result.
func (w *rpcResultWaiter) Wait(ctx context.Context) (encoded *encodedOutboundMessage, ok bool, err error) {
if w == nil || w.flight == nil || ctx == nil {
@ -139,7 +148,7 @@ func (w *rpcResultWaiter) Wait(ctx context.Context) (encoded *encodedOutboundMes
}
// Subscribe registers an event callback without creating a goroutine or
// occupying an RPC worker. The callback is invoked after the cache shard lock is
// occupying an RPC worker. The callback is invoked after the ledger shard lock is
// released; it must remain non-blocking.
func (w *rpcResultWaiter) Subscribe(fn func(*encodedOutboundMessage, bool)) error {
if fn == nil {
@ -177,10 +186,10 @@ 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) {
if w == nil || w.ledger == nil || w.flight == nil || (resultFn == nil && executionFn == nil) {
return ErrRPCResultFlightInvalid
}
s := w.cache.shard(w.key)
s := w.ledger.shard(w.key)
var (
encoded *encodedOutboundMessage
resultOK bool
@ -198,8 +207,8 @@ func (w *rpcResultWaiter) subscribe(
slots++
}
if slots > 0 {
if flight.subscriberSlots > w.cache.subscriberPerFlight-slots ||
!w.cache.subscriberBudget.reserve(w.key, slots) {
if flight.subscriberSlots > w.ledger.subscriberPerFlight-slots ||
!w.ledger.subscriberBudget.reserve(w.key, slots) {
s.mu.Unlock()
return ErrRPCResultSubscriberCapacity
}
@ -247,8 +256,8 @@ func (w *rpcResultWaiter) subscribe(
}
type rpcResultOwnerLease struct {
cache *rpcResultCache
key rpcResultCacheKey
ledger *rpcExecutionLedger
key rpcExecutionKey
flight *rpcResultFlight
delivery *rpcResultDelivery
hookMu sync.Mutex
@ -269,13 +278,13 @@ func (l *rpcResultOwnerLease) SetAbortHook(fn func()) {
}
// InstallAbortHook installs fn only while this lease still owns the pending
// flight. The shard lock linearizes installation with Abort/Put so a registry
// flight. The shard lock linearizes installation with Abort/Complete so a registry
// cannot publish a candidate after its owner has already disappeared.
func (l *rpcResultOwnerLease) InstallAbortHook(fn func()) bool {
if l == nil || l.cache == nil || l.flight == nil || fn == nil {
if l == nil || l.ledger == nil || l.flight == nil || fn == nil {
return false
}
s := l.cache.shard(l.key)
s := l.ledger.shard(l.key)
s.mu.Lock()
flight, ok := s.pending[l.key]
if !ok || flight != l.flight {
@ -290,10 +299,10 @@ func (l *rpcResultOwnerLease) InstallAbortHook(fn func()) bool {
}
func (l *rpcResultOwnerLease) Waiter() *rpcResultWaiter {
if l == nil || l.cache == nil || l.flight == nil {
if l == nil || l.ledger == nil || l.flight == nil {
return nil
}
return &rpcResultWaiter{cache: l.cache, key: l.key, flight: l.flight}
return &rpcResultWaiter{ledger: l.ledger, key: l.key, flight: l.flight}
}
func (l *rpcResultOwnerLease) TryRetarget(reqMsgID int64) bool {
@ -309,13 +318,13 @@ func (l *rpcResultOwnerLease) Delivery() *rpcResultDelivery {
// 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
// failures and dependency failures. Delivery/ledger completion remains a
// separate later transition.
func (l *rpcResultOwnerLease) CompleteExecution(success bool) bool {
if l == nil || l.cache == nil || l.flight == nil {
if l == nil || l.ledger == nil || l.flight == nil {
return false
}
s := l.cache.shard(l.key)
s := l.ledger.shard(l.key)
s.mu.Lock()
flight, ok := s.pending[l.key]
if !ok || flight != l.flight || flight.executionDone {
@ -326,7 +335,7 @@ func (l *rpcResultOwnerLease) CompleteExecution(success bool) bool {
flight.executionOK = success
subscribers := append([]func(bool){}, flight.executionSubscribers...)
flight.executionSubscribers = nil
l.cache.releaseFlightSubscriberSlotsLocked(l.key, flight, len(subscribers))
l.ledger.releaseFlightSubscriberSlotsLocked(l.key, flight, len(subscribers))
s.mu.Unlock()
for _, subscriber := range subscribers {
subscriber(success)
@ -336,12 +345,12 @@ func (l *rpcResultOwnerLease) CompleteExecution(success bool) bool {
// 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.
// the flight through Complete on both successful delivery and fenced failure.
func (l *rpcResultOwnerLease) HandOff() bool {
if l == nil || l.cache == nil || l.flight == nil {
if l == nil || l.ledger == nil || l.flight == nil {
return false
}
s := l.cache.shard(l.key)
s := l.ledger.shard(l.key)
s.mu.Lock()
defer s.mu.Unlock()
flight, ok := s.pending[l.key]
@ -356,13 +365,13 @@ func (l *rpcResultOwnerLease) HandOff() bool {
// result. Pointer identity prevents an old lease from deleting a later owner
// that reacquired the same key. It returns true only for the winning abort.
func (l *rpcResultOwnerLease) Abort() bool {
if l == nil || l.cache == nil || l.flight == nil {
if l == nil || l.ledger == nil || l.flight == nil {
return false
}
if l.handedOff.Load() {
return false
}
s := l.cache.shard(l.key)
s := l.ledger.shard(l.key)
s.mu.Lock()
if l.handedOff.Load() {
s.mu.Unlock()
@ -379,13 +388,13 @@ func (l *rpcResultOwnerLease) Abort() bool {
flight.reservation.release()
flight.reservation = nil
}
l.cache.flightLimit.release()
l.cache.activeAdmissions.retire(flight.admissionSeq)
l.ledger.flightLimit.release()
l.ledger.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.ledger.releaseFlightSubscriberSlotsLocked(
l.key, flight, len(subscribers)+len(executionSubscribers),
)
if !flight.executionDone {
@ -449,7 +458,7 @@ func (l *rpcResultFlightLimit) releaseN(delta int64) {
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
// Complete/Abort use map removal and lease identity to make double release
// impossible. Fail fast instead of masking a capacity-accounting bug that
// could otherwise admit more owners than the configured hard limit.
panic("mtproto rpc result in-flight counter underflow")
@ -465,30 +474,30 @@ 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 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{})
// lifecycle from completed receipt TTL but reserve the same global/auth/session
// ownership that Complete later transfers to a receipt or tombstone.
func (l *rpcExecutionLedger) Acquire(authKeyID [8]byte, sessionID, reqMsgID int64) (rpcResultAcquire, error) {
return l.acquire(authKeyID, sessionID, reqMsgID, rpcResultRequestIdentity{})
}
func (c *rpcResultCache) AcquireIdentified(
func (l *rpcExecutionLedger) AcquireIdentified(
authKeyID [8]byte,
sessionID, reqMsgID int64,
identity tlprofile.PreparedIdentity,
) (rpcResultAcquire, error) {
return c.acquire(authKeyID, sessionID, reqMsgID, rpcResultRequestIdentity{exact: identity, valid: true})
return l.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(
func (l *rpcExecutionLedger) AcquireLayerIdentified(
authKeyID [8]byte,
sessionID, reqMsgID int64,
profile tlprofile.Profile,
identity tlprofile.PreparedIdentity,
) (rpcResultAcquire, error) {
return c.acquire(authKeyID, sessionID, reqMsgID, rpcResultRequestIdentity{
return l.acquire(authKeyID, sessionID, reqMsgID, rpcResultRequestIdentity{
exact: identity, profile: profile, valid: true,
})
}
@ -497,17 +506,17 @@ func (c *rpcResultCache) AcquireLayerIdentified(
// 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) (tlprofile.Profile, bool) {
if c == nil || reqMsgID == 0 {
func (l *rpcExecutionLedger) ExactAdmissionProfile(authKeyID [8]byte, sessionID, reqMsgID int64) (tlprofile.Profile, bool) {
if l == nil || reqMsgID == 0 {
return 0, false
}
key := rpcResultCacheKey{authKeyID: authKeyID, sessionID: sessionID, reqMsgID: reqMsgID}
s := c.shard(key)
key := rpcExecutionKey{authKeyID: authKeyID, sessionID: sessionID, reqMsgID: reqMsgID}
s := l.shard(key)
now := s.now()
s.mu.Lock()
defer s.mu.Unlock()
if elem := s.byKey[key]; elem != nil {
entry := elem.Value.(*rpcResultCacheEntry)
entry := elem.Value.(*rpcExecutionReceipt)
if entry.expiresAt.After(now) {
if entry.identity.valid && entry.identity.profile != 0 {
return entry.identity.profile, true
@ -522,16 +531,16 @@ func (c *rpcResultCache) ExactAdmissionProfile(authKeyID [8]byte, sessionID, req
return 0, false
}
func (c *rpcResultCache) acquire(
func (l *rpcExecutionLedger) acquire(
authKeyID [8]byte,
sessionID, reqMsgID int64,
identity rpcResultRequestIdentity,
) (rpcResultAcquire, error) {
if c == nil || reqMsgID == 0 {
if l == nil || reqMsgID == 0 {
return rpcResultAcquire{}, ErrRPCResultFlightInvalid
}
key := rpcResultCacheKey{authKeyID: authKeyID, sessionID: sessionID, reqMsgID: reqMsgID}
s := c.shard(key)
key := rpcExecutionKey{authKeyID: authKeyID, sessionID: sessionID, reqMsgID: reqMsgID}
s := l.shard(key)
reclaimedExpired := false
for {
now := s.now()
@ -539,20 +548,33 @@ func (c *rpcResultCache) acquire(
s.expireLocked(now)
if elem, ok := s.byKey[key]; ok {
entry := elem.Value.(*rpcResultCacheEntry)
entry := elem.Value.(*rpcExecutionReceipt)
if !entry.identity.matches(identity) {
s.mu.Unlock()
return rpcResultAcquire{}, identityMismatch(entry.identity)
}
if entry.capacity || entry.encoded == nil {
if entry.acknowledged {
result := rpcResultAcquire{
state: rpcResultAcquireAcknowledged, admissionSeq: entry.admissionSeq,
executionKnown: entry.executionKnown, executionOK: entry.executionOK,
}
s.mu.Unlock()
return result, nil
}
if entry.unavailable {
s.mu.Unlock()
return rpcResultAcquire{}, ErrRPCResultFlightCapacity
}
result := rpcResultAcquire{
state: rpcResultAcquireCompleted, admissionSeq: entry.admissionSeq, encoded: entry.encoded,
state: rpcResultAcquireCompleted, admissionSeq: entry.admissionSeq,
executionKnown: entry.executionKnown, executionOK: entry.executionOK,
}
s.mu.Unlock()
encoded, replayable := l.replayStore.rpcResult(authKeyID, sessionID, reqMsgID)
if !replayable {
return rpcResultAcquire{}, ErrRPCResultReplayUnavailable
}
result.encoded = encoded
return result, nil
}
if flight, ok := s.pending[key]; ok {
@ -560,25 +582,29 @@ func (c *rpcResultCache) acquire(
s.mu.Unlock()
return rpcResultAcquire{}, identityMismatch(flight.identity)
}
if flight.acknowledged {
result := rpcResultAcquire{
state: rpcResultAcquireAcknowledged, admissionSeq: flight.admissionSeq,
executionKnown: flight.executionDone, executionOK: flight.executionOK,
}
s.mu.Unlock()
return result, nil
}
result := rpcResultAcquire{
state: rpcResultAcquirePending,
admissionSeq: flight.admissionSeq,
waiter: &rpcResultWaiter{cache: c, key: key, flight: flight},
waiter: &rpcResultWaiter{ledger: l, key: key, flight: flight},
}
s.mu.Unlock()
return result, nil
}
if s.maxEntries > 0 && len(s.byKey)+len(s.pending) >= s.maxEntries {
if !l.flightLimit.reserve() {
s.mu.Unlock()
return rpcResultAcquire{}, ErrRPCResultFlightCapacity
}
if !c.flightLimit.reserve() {
s.mu.Unlock()
return rpcResultAcquire{}, ErrRPCResultFlightCapacity
}
reservation := c.fairBudget.reserveOwner(key)
reservation := l.fairBudget.reserveOwner(key)
if reservation == nil {
c.flightLimit.release()
l.flightLimit.release()
s.mu.Unlock()
if reclaimedExpired {
return rpcResultAcquire{}, ErrRPCResultFlightCapacity
@ -586,17 +612,17 @@ func (c *rpcResultCache) acquire(
// 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()
l.expireReceipts()
reclaimedExpired = true
continue
}
var admissionSeq uint64
if identity.valid {
var err error
admissionSeq, err = c.activeAdmissions.allocateAndRegister(&c.nextAdmissionSeq)
admissionSeq, err = l.activeAdmissions.allocateAndRegister(&l.nextAdmissionSeq)
if err != nil {
reservation.release()
c.flightLimit.release()
l.flightLimit.release()
s.mu.Unlock()
return rpcResultAcquire{}, err
}
@ -606,7 +632,7 @@ func (c *rpcResultCache) acquire(
reservation: reservation,
}
if s.pending == nil {
s.pending = make(map[rpcResultCacheKey]*rpcResultFlight)
s.pending = make(map[rpcExecutionKey]*rpcResultFlight)
}
s.pending[key] = flight
s.mu.Unlock()
@ -614,24 +640,24 @@ func (c *rpcResultCache) acquire(
state: rpcResultAcquireOwner,
admissionSeq: admissionSeq,
owner: &rpcResultOwnerLease{
cache: c, key: key, flight: flight, delivery: newRPCResultDelivery(reqMsgID),
ledger: l, key: key, flight: flight, delivery: newRPCResultDelivery(reqMsgID),
},
}, nil
}
}
// completeRPCResultFlightLocked publishes encoded to the current owner claim.
// The caller must hold s.mu and must publish the completed cache entry first.
func (c *rpcResultCache) completeRPCResultFlightLocked(
s *rpcResultCacheShard,
key rpcResultCacheKey,
// The caller must hold s.mu and must publish the completed receipt first.
func (l *rpcExecutionLedger) completeRPCResultFlightLocked(
s *rpcExecutionLedgerShard,
key rpcExecutionKey,
encoded *encodedOutboundMessage,
) (
[]func(*encodedOutboundMessage, bool),
[]func(bool),
bool,
) {
if c == nil || s == nil || encoded == nil {
if l == nil || s == nil || encoded == nil {
return nil, nil, false
}
flight, ok := s.pending[key]
@ -646,8 +672,8 @@ func (c *rpcResultCache) completeRPCResultFlightLocked(
}
flight.encoded = encoded
flight.ok = true
c.flightLimit.release()
c.activeAdmissions.retire(flight.admissionSeq)
l.flightLimit.release()
l.activeAdmissions.retire(flight.admissionSeq)
subscribers := append([]func(*encodedOutboundMessage, bool){}, flight.subscribers...)
flight.subscribers = nil
executionSubscribers := append([]func(bool){}, flight.executionSubscribers...)
@ -655,13 +681,13 @@ func (c *rpcResultCache) completeRPCResultFlightLocked(
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
// invokeAfter dependency. Production completes execution before Complete;
// branch is the conservative terminal cleanup for defensive callers.
flight.executionDone = true
flight.executionOK = false
executionOK = false
}
c.releaseFlightSubscriberSlotsLocked(
l.releaseFlightSubscriberSlotsLocked(
key, flight, len(subscribers)+len(executionSubscribers),
)
close(flight.done)
@ -669,19 +695,19 @@ func (c *rpcResultCache) completeRPCResultFlightLocked(
}
// releaseFlightSubscriberSlotsLocked releases callbacks detached from a flight.
// The caller holds that flight's cache shard lock, preserving the only lock
// The caller holds that flight's ledger shard lock, preserving the only lock
// order used by subscription: shard -> subscriber budget.
func (c *rpcResultCache) releaseFlightSubscriberSlotsLocked(
key rpcResultCacheKey,
func (l *rpcExecutionLedger) releaseFlightSubscriberSlotsLocked(
key rpcExecutionKey,
flight *rpcResultFlight,
slots int,
) {
if slots == 0 {
return
}
if c == nil || flight == nil || slots < 0 || flight.subscriberSlots < slots {
if l == nil || flight == nil || slots < 0 || flight.subscriberSlots < slots {
panic("mtproto rpc result subscriber slot underflow")
}
flight.subscriberSlots -= slots
c.subscriberBudget.release(key, slots)
l.subscriberBudget.release(key, slots)
}

121
internal/mtprotoedge/rpc_flight_test.go
View file

@ -32,25 +32,23 @@ func rpcFlightExactIdentity(t *testing.T, profile tlprofile.Profile, request bin
return admitted.Prepared().Identity()
}
func newRPCResultSubscriberTestCache(global, auth, session, perFlight int) *rpcResultCache {
return newRPCResultCacheWithFairCapacity(time.Now, rpcResultCacheCapacity{
func newRPCExecutionSubscriberTestLedger(global, auth, session, perFlight int) *rpcExecutionLedger {
return newRPCExecutionLedger(time.Now, rpcExecutionLedgerCapacity{
maxPending: 64,
maxPendingPerAuth: 64,
globalMaxBytes: rpcResultCacheMaxBytes,
globalMaxEntries: rpcResultCacheMaxEntries,
authMaxBytes: rpcResultCacheAuthMaxBytes,
authMaxEntries: rpcResultCacheAuthMaxEntries,
sessionMaxBytes: rpcResultCacheSessionMaxBytes,
sessionMaxEntries: rpcResultCacheSessionMaxEntries,
globalMaxEntries: rpcExecutionMaxEntries,
authMaxEntries: rpcExecutionAuthMaxEntries,
sessionMaxEntries: rpcExecutionSessionMaxEntries,
subscriberMaxGlobal: global,
subscriberMaxAuth: auth,
subscriberMaxSession: session,
subscriberMaxPerFlight: perFlight,
replayStore: newRPCReplayStoreForTest(),
})
}
func TestRPCResultFlightSubscriberPairCapacityFailureIsAtomic(t *testing.T) {
cache := newRPCResultSubscriberTestCache(8, 8, 8, 1)
cache := newRPCExecutionSubscriberTestLedger(8, 8, 8, 1)
authKeyID := rpcFlightTestAuthID(70)
claim, err := cache.Acquire(authKeyID, 70, 700)
if err != nil || claim.owner == nil {
@ -64,7 +62,7 @@ func TestRPCResultFlightSubscriberPairCapacityFailureIsAtomic(t *testing.T) {
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 := cache.shard(rpcExecutionKey{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)
@ -74,7 +72,7 @@ func TestRPCResultFlightSubscriberPairCapacityFailureIsAtomic(t *testing.T) {
}
s.mu.Unlock()
claim.owner.CompleteExecution(true)
cache.Put(authKeyID, 70, 700, &encodedOutboundMessage{body: []byte{1}, reqMsgID: 700})
cache.completeReplayableForTest(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)
}
@ -84,7 +82,7 @@ func TestRPCResultFlightSubscriberPairCapacityFailureIsAtomic(t *testing.T) {
}
func TestRPCResultFlightSubscriberBudgetsIsolateSessionAndAuth(t *testing.T) {
cache := newRPCResultSubscriberTestCache(3, 2, 1, 4)
cache := newRPCExecutionSubscriberTestLedger(3, 2, 1, 4)
authA := rpcFlightTestAuthID(71)
authB := rpcFlightTestAuthID(72)
type ownerKey struct {
@ -149,7 +147,7 @@ func TestRPCResultFlightSubscriberBudgetsIsolateSessionAndAuth(t *testing.T) {
}
func TestRPCResultFlightSubscriberSlotsReleasePerTerminalHalf(t *testing.T) {
cache := newRPCResultSubscriberTestCache(4, 4, 4, 4)
cache := newRPCExecutionSubscriberTestLedger(4, 4, 4, 4)
authKeyID := rpcFlightTestAuthID(73)
claim, err := cache.Acquire(authKeyID, 73, 730)
if err != nil || claim.owner == nil {
@ -175,7 +173,7 @@ func TestRPCResultFlightSubscriberSlotsReleasePerTerminalHalf(t *testing.T) {
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})
cache.completeReplayableForTest(authKeyID, 73, 730, &encodedOutboundMessage{body: []byte{1}, reqMsgID: 730})
if ok := <-result; !ok {
t.Fatal("result callback reported failure")
}
@ -185,7 +183,7 @@ func TestRPCResultFlightSubscriberSlotsReleasePerTerminalHalf(t *testing.T) {
}
func TestRPCResultFlightRepeatedReplayJoinsStayBoundedAndPutCleansExecution(t *testing.T) {
cache := newRPCResultSubscriberTestCache(2, 2, 2, 2)
cache := newRPCExecutionSubscriberTestLedger(2, 2, 2, 2)
authKeyID := rpcFlightTestAuthID(74)
claim, err := cache.Acquire(authKeyID, 74, 740)
if err != nil || claim.owner == nil {
@ -212,7 +210,7 @@ func TestRPCResultFlightRepeatedReplayJoinsStayBoundedAndPutCleansExecution(t *t
t.Fatalf("join %d err=%v, want capacity", i, err)
}
}
cache.Put(authKeyID, 74, 740, &encodedOutboundMessage{body: []byte{1}, reqMsgID: 740})
cache.completeReplayableForTest(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)
}
@ -222,7 +220,7 @@ func TestRPCResultFlightRepeatedReplayJoinsStayBoundedAndPutCleansExecution(t *t
}
func TestRPCResultFlightExactIdentityGuardsPendingAndCompletedReuse(t *testing.T) {
cache := newRPCResultCacheWithFlightLimit(time.Now, 2)
cache := newRPCExecutionLedgerForTest(time.Now, 2)
authKeyID := rpcFlightTestAuthID(90)
firstIdentity := rpcFlightExactIdentity(t, tlprofile.Profile225, &tg.HelpGetConfigRequest{})
otherIdentity := rpcFlightExactIdentity(t, tlprofile.Profile225, &tg.HelpGetNearestDCRequest{})
@ -240,7 +238,7 @@ func TestRPCResultFlightExactIdentityGuardsPendingAndCompletedReuse(t *testing.T
}
want := &encodedOutboundMessage{body: []byte{1, 2, 3, 4}, reqMsgID: 900}
cache.Put(authKeyID, 90, 900, want)
cache.completeReplayableForTest(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)
}
@ -254,7 +252,7 @@ func TestRPCResultFlightExactIdentityGuardsPendingAndCompletedReuse(t *testing.T
}
func TestRPCResultFlightAdmissionSequenceAllocatedOnceAndReplayed(t *testing.T) {
cache := newRPCResultCacheWithFlightLimit(time.Now, 4)
cache := newRPCExecutionLedgerForTest(time.Now, 4)
authKeyID := rpcFlightTestAuthID(89)
identity := rpcFlightExactIdentity(t, tlprofile.Profile225, &tg.HelpGetConfigRequest{})
owner, err := cache.AcquireLayerIdentified(authKeyID, 89, 890, tlprofile.Profile225, identity)
@ -267,7 +265,7 @@ func TestRPCResultFlightAdmissionSequenceAllocatedOnceAndReplayed(t *testing.T)
}
owner.owner.CompleteExecution(true)
encoded := &encodedOutboundMessage{body: []byte{1}, reqMsgID: 890}
cache.Put(authKeyID, 89, 890, encoded)
cache.completeReplayableForTest(authKeyID, 89, 890, encoded)
completed, err := cache.AcquireLayerIdentified(authKeyID, 89, 890, tlprofile.Profile225, 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)
@ -285,7 +283,7 @@ func TestRPCResultFlightAdmissionSequenceAllocatedOnceAndReplayed(t *testing.T)
}
func TestRPCAdmissionSafeFloorTracksOwnersUntilPutOrAbort(t *testing.T) {
cache := newRPCResultCacheWithFlightLimit(time.Now, 4)
cache := newRPCExecutionLedgerForTest(time.Now, 4)
authKeyID := rpcFlightTestAuthID(86)
identity := rpcFlightExactIdentity(t, tlprofile.Profile225, &tg.HelpGetConfigRequest{})
first, err := cache.AcquireLayerIdentified(authKeyID, 86, 860, tlprofile.Profile225, identity)
@ -306,14 +304,14 @@ func TestRPCAdmissionSafeFloorTracksOwnersUntilPutOrAbort(t *testing.T) {
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})
cache.completeReplayableForTest(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 := newRPCExecutionLedgerForTest(time.Now, 2)
cache.nextAdmissionSeq.Store(^uint64(0) - 1)
authKeyID := rpcFlightTestAuthID(85)
identity := rpcFlightExactIdentity(t, tlprofile.Profile225, &tg.HelpGetConfigRequest{})
@ -331,7 +329,7 @@ func TestRPCAdmissionSequenceExhaustionCannotWrap(t *testing.T) {
}
func TestRPCIdentityMismatchCarriesWinnerProfileAcrossAbort(t *testing.T) {
cache := newRPCResultCacheWithFlightLimit(time.Now, 2)
cache := newRPCExecutionLedgerForTest(time.Now, 2)
authKeyID := rpcFlightTestAuthID(88)
request := &tg.MessagesGetHistoryRequest{Peer: &tg.InputPeerSelf{}, Limit: 1}
winnerIdentity := rpcFlightExactIdentity(t, tlprofile.Profile225, request)
@ -357,7 +355,7 @@ func TestRPCIdentityMismatchCarriesWinnerProfileAcrossAbort(t *testing.T) {
func TestRPCAdmissionProfileHintSurvivesCompletedEvictionWindow(t *testing.T) {
now := time.Unix(1_900_000_000, 0)
cache := newRPCResultCacheWithFlightLimit(func() time.Time { return now }, 2)
cache := newRPCExecutionLedgerForTest(func() time.Time { return now }, 2)
authKeyID := rpcFlightTestAuthID(87)
identity := rpcFlightExactIdentity(t, tlprofile.Profile225, &tg.MessagesGetHistoryRequest{
Peer: &tg.InputPeerSelf{}, Limit: 1,
@ -367,7 +365,7 @@ func TestRPCAdmissionProfileHintSurvivesCompletedEvictionWindow(t *testing.T) {
t.Fatalf("owner err=%v", err)
}
claim.owner.CompleteExecution(true)
cache.Put(authKeyID, 87, 870, &encodedOutboundMessage{body: []byte{1}, reqMsgID: 870})
cache.completeReplayableForTest(authKeyID, 87, 870, &encodedOutboundMessage{body: []byte{1}, reqMsgID: 870})
profile, ok := cache.ExactAdmissionProfile(authKeyID, 87, 870)
if !ok || profile != tlprofile.Profile225 {
t.Fatalf("profile hint = (%d,%v)", profile, ok)
@ -375,7 +373,7 @@ func TestRPCAdmissionProfileHintSurvivesCompletedEvictionWindow(t *testing.T) {
// 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)
now = now.Add(rpcExecutionReceiptTTL + 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)
@ -384,7 +382,7 @@ func TestRPCAdmissionProfileHintSurvivesCompletedEvictionWindow(t *testing.T) {
}
func TestRPCInvariantIdentityDoesNotExposeCanonicalProfileHint(t *testing.T) {
cache := newRPCResultCacheWithFlightLimit(time.Now, 2)
cache := newRPCExecutionLedgerForTest(time.Now, 2)
authKeyID := rpcFlightTestAuthID(84)
identity := rpcFlightExactIdentity(t, tlprofile.Profile227, &tg.AuthBindTempAuthKeyRequest{
PermAuthKeyID: 1, Nonce: 2, ExpiresAt: 3, EncryptedMessage: []byte("bind"),
@ -397,14 +395,14 @@ func TestRPCInvariantIdentityDoesNotExposeCanonicalProfileHint(t *testing.T) {
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})
cache.completeReplayableForTest(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)
cache := newRPCExecutionLedgerForTest(time.Now, 2)
authKeyID := rpcFlightTestAuthID(91)
claim, err := cache.Acquire(authKeyID, 91, 910)
if err != nil || claim.state != rpcResultAcquireOwner || claim.owner == nil {
@ -436,7 +434,7 @@ func TestRPCResultExecutionCompletionIsExactlyOnceAndDurable(t *testing.T) {
}
want := &encodedOutboundMessage{body: []byte{9, 1, 0, 0}, reqMsgID: 910}
cache.Put(authKeyID, 91, 910, want)
cache.completeReplayableForTest(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)
@ -451,7 +449,7 @@ func TestRPCResultExecutionCompletionIsExactlyOnceAndDurable(t *testing.T) {
}
func TestRPCResultExecutionAbortPublishesFailure(t *testing.T) {
cache := newRPCResultCacheWithFlightLimit(time.Now, 1)
cache := newRPCExecutionLedgerForTest(time.Now, 1)
authKeyID := rpcFlightTestAuthID(92)
claim, err := cache.Acquire(authKeyID, 92, 920)
if err != nil || claim.owner == nil {
@ -474,7 +472,7 @@ func TestRPCResultExecutionAbortPublishesFailure(t *testing.T) {
func TestRPCResultFlightConcurrentAcquireHasUniqueOwner(t *testing.T) {
const callers = 64
cache := newRPCResultCacheWithFlightLimit(time.Now, callers)
cache := newRPCExecutionLedgerForTest(time.Now, callers)
authKeyID := rpcFlightTestAuthID(1)
start := make(chan struct{})
results := make(chan rpcResultAcquire, callers)
@ -535,7 +533,7 @@ func TestRPCResultFlightConcurrentAcquireHasUniqueOwner(t *testing.T) {
func TestRPCResultFlightPutPublishesAndWakesAllWaiters(t *testing.T) {
const waiters = 24
cache := newRPCResultCacheWithFlightLimit(time.Now, 32)
cache := newRPCExecutionLedgerForTest(time.Now, 32)
authKeyID := rpcFlightTestAuthID(10)
owner, err := cache.Acquire(authKeyID, 20, 200)
if err != nil || owner.state != rpcResultAcquireOwner || owner.owner == nil {
@ -563,13 +561,13 @@ func TestRPCResultFlightPutPublishesAndWakesAllWaiters(t *testing.T) {
for _, waiter := range waiterClaims {
go func(w *rpcResultWaiter) {
encoded, ok, waitErr := w.Wait(ctx)
cached, _ := cache.Get(authKeyID, 20, 200)
cached, _ := cache.Replay(authKeyID, 20, 200)
results <- waiterResult{encoded: encoded, cached: cached, ok: ok, err: waitErr}
}(waiter)
}
want := &encodedOutboundMessage{body: []byte{1, 2, 3, 4}, typeID: 42, reqMsgID: 200}
cache.Put(authKeyID, 20, 200, want)
cache.completeReplayableForTest(authKeyID, 20, 200, want)
for i := 0; i < waiters; i++ {
got := <-results
if got.err != nil || !got.ok {
@ -592,7 +590,7 @@ func TestRPCResultFlightPutPublishesAndWakesAllWaiters(t *testing.T) {
}
func TestRPCResultFlightAbortWakesAndAllowsReclaim(t *testing.T) {
cache := newRPCResultCacheWithFlightLimit(time.Now, 2)
cache := newRPCExecutionLedgerForTest(time.Now, 2)
authKeyID := rpcFlightTestAuthID(20)
first, err := cache.Acquire(authKeyID, 30, 300)
if err != nil || first.state != rpcResultAcquireOwner {
@ -627,22 +625,17 @@ func TestRPCResultFlightAbortWakesAndAllowsReclaim(t *testing.T) {
}
}
func TestRPCResultFlightCompletedCachePressureDoesNotEvictPending(t *testing.T) {
func TestRPCResultFlightReceiptLifecycleDoesNotEvictPending(t *testing.T) {
now := time.Unix(1_000, 0)
cache := newRPCResultCacheWithFlightLimit(func() time.Time { return now }, 4)
cache := newRPCExecutionLedgerForTest(func() time.Time { return now }, 4)
authKeyID := rpcFlightTestAuthID(30)
pending, err := cache.Acquire(authKeyID, 40, 400)
if err != nil || pending.state != rpcResultAcquireOwner {
t.Fatalf("pending Acquire = state:%d err:%v", pending.state, err)
}
key := rpcResultCacheKey{authKeyID: authKeyID, sessionID: 40, reqMsgID: 400}
shard := cache.shard(key)
shard.mu.Lock()
shard.maxEntries = 2
shard.mu.Unlock()
for i := int64(0); i < 16; i++ {
cache.Put(authKeyID, 40, 500+i, &encodedOutboundMessage{body: []byte{byte(i)}})
cache.completeReplayableForTest(authKeyID, 40, 500+i, &encodedOutboundMessage{body: []byte{byte(i)}})
}
if got := cache.flightLimit.snapshot(); got != 1 {
t.Fatalf("completed trim changed pending count to %d", got)
@ -652,9 +645,9 @@ func TestRPCResultFlightCompletedCachePressureDoesNotEvictPending(t *testing.T)
t.Fatalf("Acquire after completed trim = state:%d err:%v", joined.state, err)
}
// Expire the independent completed cache and prove the pending owner remains.
now = now.Add(rpcResultCacheTTL + time.Second)
_, _ = cache.Get(authKeyID, 40, 515)
// Expire independent completed receipts and prove the pending owner remains.
now = now.Add(rpcExecutionReceiptTTL + time.Second)
_, _ = cache.Replay(authKeyID, 40, 515)
joinedAfterTTL, err := cache.Acquire(authKeyID, 40, 400)
if err != nil || joinedAfterTTL.state != rpcResultAcquirePending {
t.Fatalf("Acquire after completed TTL = state:%d err:%v", joinedAfterTTL.state, err)
@ -665,7 +658,7 @@ func TestRPCResultFlightCompletedCachePressureDoesNotEvictPending(t *testing.T)
}
func TestRPCResultFlightCapacityAndCountReturn(t *testing.T) {
cache := newRPCResultCacheWithFlightLimit(time.Now, 2)
cache := newRPCExecutionLedgerForTest(time.Now, 2)
authKeyID := rpcFlightTestAuthID(40)
first, err := cache.Acquire(authKeyID, 50, 501)
if err != nil || first.state != rpcResultAcquireOwner {
@ -687,7 +680,7 @@ func TestRPCResultFlightCapacityAndCountReturn(t *testing.T) {
}
want := &encodedOutboundMessage{body: []byte{9}, reqMsgID: 501}
cache.Put(authKeyID, 50, 501, want)
cache.completeReplayableForTest(authKeyID, 50, 501, want)
if got := cache.flightLimit.snapshot(); got != 1 {
t.Fatalf("pending count after Put = %d, want 1", got)
}
@ -710,8 +703,8 @@ func TestRPCResultFlightCapacityAndCountReturn(t *testing.T) {
}
}
func TestRPCResultFlightLargePutPublishesCompletedBeforeResolvingWaiters(t *testing.T) {
cache := newRPCResultCacheWithFlightLimit(time.Now, 1)
func TestRPCResultFlightLargeCompletionDoesNotChargeLedgerBodyBytes(t *testing.T) {
cache := newRPCExecutionLedgerForTest(time.Now, 1)
authKeyID := rpcFlightTestAuthID(50)
owner, err := cache.Acquire(authKeyID, 60, 600)
if err != nil || owner.state != rpcResultAcquireOwner {
@ -722,22 +715,22 @@ func TestRPCResultFlightLargePutPublishesCompletedBeforeResolvingWaiters(t *test
t.Fatalf("joined Acquire = state:%d err:%v", joined.state, err)
}
// 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
// Large file results remain owned by the replay store. The ledger retains
// only one fixed-shape receipt regardless of payload size.
largeSize := 8 << 20
want := &encodedOutboundMessage{body: make([]byte, largeSize), reqMsgID: 600}
cache.Put(authKeyID, 60, 600, want)
cache.completeReplayableForTest(authKeyID, 60, 600, want)
if encoded, ok, waitErr := joined.waiter.Wait(context.Background()); waitErr != nil || !ok || encoded != want {
t.Fatalf("large Wait = encoded:%p ok:%v err:%v", encoded, ok, waitErr)
}
if got, ok := cache.Get(authKeyID, 60, 600); !ok || got != want {
t.Fatalf("large completed Get = encoded:%p ok:%v", got, ok)
if got, ok := cache.Replay(authKeyID, 60, 600); !ok || got != want {
t.Fatalf("large completed Replay = encoded:%p ok:%v", got, ok)
}
if got := cache.flightLimit.snapshot(); got != 0 {
t.Fatalf("large Put leaked pending count %d", got)
}
if owner.owner.Abort() {
t.Fatal("large Put left its old owner abortable")
t.Fatal("large completion left its old owner abortable")
}
completed, err := cache.Acquire(authKeyID, 60, 600)
if err != nil || completed.state != rpcResultAcquireCompleted || completed.encoded != want {
@ -746,13 +739,13 @@ func TestRPCResultFlightLargePutPublishesCompletedBeforeResolvingWaiters(t *test
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)
if got := cache.receiptBudgetBytes(); got != rpcExecutionReceiptBudgetBytes {
t.Fatalf("receipt budget bytes = %d, want %d", got, rpcExecutionReceiptBudgetBytes)
}
}
func TestRPCResultFlightWaitContextDoesNotReleaseOwner(t *testing.T) {
cache := newRPCResultCacheWithFlightLimit(time.Now, 1)
cache := newRPCExecutionLedgerForTest(time.Now, 1)
authKeyID := rpcFlightTestAuthID(60)
owner, err := cache.Acquire(authKeyID, 70, 700)
if err != nil {
@ -780,7 +773,7 @@ func TestRPCResultFlightConcurrentCapacityReturnsAllSlots(t *testing.T) {
limit = 32
callers = 512
)
cache := newRPCResultCacheWithFlightLimit(time.Now, limit)
cache := newRPCExecutionLedgerForTest(time.Now, limit)
authKeyID := rpcFlightTestAuthID(70)
start := make(chan struct{})
errs := make(chan error, callers)
@ -804,7 +797,7 @@ func TestRPCResultFlightConcurrentCapacityReturnsAllSlots(t *testing.T) {
return
}
if i%2 == 0 {
cache.Put(authKeyID, 80+int64(i%4), 1_000+int64(i), &encodedOutboundMessage{body: []byte{1}})
cache.completeReplayableForTest(authKeyID, 80+int64(i%4), 1_000+int64(i), &encodedOutboundMessage{body: []byte{1}})
} else if !claim.owner.Abort() {
errs <- errors.New("owner Abort lost")
}

292
internal/mtprotoedge/rpc_result_budget.go
View file

@ -1,292 +0,0 @@
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]
}

559
internal/mtprotoedge/rpc_result_cache.go
View file

@ -1,559 +0,0 @@
package mtprotoedge
import (
"container/list"
"encoding/binary"
"hash/maphash"
"sync"
"sync/atomic"
"time"
)
const (
// 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
)
type rpcResultCacheKey struct {
authKeyID [8]byte
sessionID int64
reqMsgID int64
}
type rpcResultCacheEntry struct {
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），
// 用于跨连接重放重复请求。Put 的调用方必须先证明结果已物理写出，或原 logical Conn
// 已不可逆 fenced；绝不能发布“Conn 仍 current/open 但结果尚未上 wire”的完成态。
// encodedOutboundMessage 构造后不可变（push fan-out 与 pending resend 均依赖该契约），
// 因此 Get/Put 直接共享指针，不做防御性拷贝。
type rpcResultCache struct {
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
// 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
bytes int
order *list.List
byKey map[rpcResultCacheKey]*list.Element
// pending is deliberately independent from the completed-result order/byKey
// cache. In-flight owners and waiters must not disappear when completed
// results expire or are trimmed under entry/byte pressure.
pending map[rpcResultCacheKey]*rpcResultFlight
}
func newRPCResultCacheWithFlightLimit(now func() time.Time, maxPending int) *rpcResultCache {
if maxPending <= 0 {
maxPending = rpcResultFlightDefaultMaxPending
}
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 = 0
s.order = list.New()
s.byKey = make(map[rpcResultCacheKey]*list.Element)
s.pending = make(map[rpcResultCacheKey]*rpcResultFlight)
}
return c
}
func (c *rpcResultCache) shard(key rpcResultCacheKey) *rpcResultCacheShard {
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) {
if c == nil || reqMsgID == 0 {
return nil, false
}
key := rpcResultCacheKey{authKeyID: authKeyID, sessionID: sessionID, reqMsgID: reqMsgID}
s := c.shard(key)
now := s.now()
s.mu.Lock()
defer s.mu.Unlock()
elem, ok := s.byKey[key]
if !ok {
return nil, false
}
entry := elem.Value.(*rpcResultCacheEntry)
if !entry.expiresAt.After(now) {
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)
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
}
// 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) {
for elem := s.order.Front(); elem != nil; {
next := elem.Next()
entry := elem.Value.(*rpcResultCacheEntry)
if entry.expiresAt.After(now) {
return
}
s.removeElement(elem)
elem = next
}
}
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) unlinkElement(elem *list.Element) *rpcResultCacheEntry {
if elem == nil {
return nil
}
entry := elem.Value.(*rpcResultCacheEntry)
delete(s.byKey, entry.key)
s.bytes -= entry.size
if s.bytes < 0 {
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()
}

830
internal/mtprotoedge/rpc_result_cache_test.go
View file

@ -1,830 +0,0 @@
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 })
var keyID [8]byte
keyID[0] = 0xab
encoded := &encodedOutboundMessage{body: []byte{1, 2, 3, 4}, typeID: 42, reqMsgID: 7}
if _, ok := cache.Get(keyID, 5, 7); ok {
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 {
t.Fatal("expected hit")
}
// encodedOutboundMessage 不可变契约下 Get/Put 共享指针，不做防御性拷贝。
if got != encoded {
t.Fatal("expected shared pointer, got clone")
}
// 不同 session / msg_id 不串。
if _, ok := cache.Get(keyID, 6, 7); ok {
t.Fatal("hit with wrong session id")
}
if _, ok := cache.Get(keyID, 5, 8); ok {
t.Fatal("hit with wrong msg id")
}
// TTL 过期。
now = now.Add(rpcResultCacheTTL + time.Second)
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 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
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)
}
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 got, ok := cache.Get(keyID, 1, 100); !ok || got != first {
t.Fatalf("unexpired first result was displaced: got=%p ok=%v", got, ok)
}
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
}

22
internal/mtprotoedge/rpc_result_delivery_test.go
View file

@ -92,7 +92,7 @@ func (t *blockingCloseRPCResultTransport) Close() error {
return nil
}
func TestRPCResultCachePublishesOnlyAfterPhysicalWrite(t *testing.T) {
func TestRPCExecutionLedgerPublishesOnlyAfterPhysicalWrite(t *testing.T) {
tr := newGatedRequiredControlTransport(nil)
s := New(Options{WriteTimeout: time.Second})
key := newTestAuthKey(t)
@ -113,7 +113,7 @@ func TestRPCResultCachePublishesOnlyAfterPhysicalWrite(t *testing.T) {
case <-time.After(time.Second):
t.Fatal("rpc_result did not reach physical writer")
}
if _, ok := s.rpcResults.Get(key.ID, c.sessionID, reqMsgID); ok {
if _, ok := s.rpcResults.Replay(key.ID, c.sessionID, reqMsgID); ok {
t.Fatal("rpc_result became completed before physical write")
}
pending, err := s.rpcResults.Acquire(key.ID, c.sessionID, reqMsgID)
@ -280,7 +280,7 @@ func TestRouterUpdateCursorDoesNotCommitAfterPhysicalWriteFailure(t *testing.T)
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 {
if cached, ok := s.rpcResults.Replay(key.ID, c.sessionID, reqMsgID); ok && cached.deliveryState() == rpcResultDeliveryReplayable {
replayable = true
break
}
@ -334,7 +334,7 @@ func TestRouterUpdateCursorDoesNotCommitWhenResultEncodingFails(t *testing.T) {
}
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 {
if cached, ok := s.rpcResults.Replay(key.ID, c.sessionID, reqMsgID); ok && cached.deliveryState() == rpcResultDeliveryDelivered {
break
}
time.Sleep(time.Millisecond)
@ -401,8 +401,8 @@ func TestRPCResultFailureAfterIntentionalTerminalDoesNotCloseTransferLease(t *te
// Session replacement publishes terminal before transferring the physical
// lease. A late old-generation result sees ErrConnClosed at producer admission;
// it may publish cache-only, but must not upgrade that intentional fence into
// a physical close that makes Transfer fail.
// it publishes a metadata-only capacity tombstone, but must not upgrade that
// intentional fence into a physical close that makes Transfer fail.
oldConn.beginTerminalShutdown()
err = s.sendResult(context.Background(), oldConn, reqMsgID, exactTestRPCResult(&tg.Config{ThisDC: 2}))
if !errors.Is(err, ErrOutboundTrackedBudget) {
@ -423,9 +423,8 @@ func TestRPCResultFailureAfterIntentionalTerminalDoesNotCloseTransferLease(t *te
if tr.closed.Load() || !newConn.isPhysicalTransportCurrentOpen() {
t.Fatalf("stale close after transfer: raw_closed=%v current_open=%v", tr.closed.Load(), newConn.isPhysicalTransportCurrentOpen())
}
completed, acquireErr := s.rpcResults.Acquire(key.ID, oldConn.sessionID, reqMsgID)
if acquireErr != nil || completed.state != rpcResultAcquireCompleted || completed.encoded == nil {
t.Fatalf("late result cache = %+v err=%v", completed, acquireErr)
if _, acquireErr := s.rpcResults.Acquire(key.ID, oldConn.sessionID, reqMsgID); !errors.Is(acquireErr, ErrRPCResultFlightCapacity) {
t.Fatalf("late result tombstone err=%v, want %v", acquireErr, ErrRPCResultFlightCapacity)
}
newConn.ForceClose()
}
@ -459,9 +458,8 @@ func TestRPCResultPublishesBeforePathologicalPhysicalCloseReturns(t *testing.T)
case <-time.After(time.Second):
t.Fatal("pathological physical Close blocked result publication")
}
completed, acquireErr := s.rpcResults.Acquire(key.ID, c.sessionID, reqMsgID)
if acquireErr != nil || completed.state != rpcResultAcquireCompleted || completed.encoded == nil {
t.Fatalf("completed result before raw Close return = %+v err=%v", completed, acquireErr)
if _, acquireErr := s.rpcResults.Acquire(key.ID, c.sessionID, reqMsgID); !errors.Is(acquireErr, ErrRPCResultFlightCapacity) {
t.Fatalf("result tombstone before raw Close return err=%v, want %v", acquireErr, ErrRPCResultFlightCapacity)
}
close(tr.release)
if c.transportLease != nil {

61
internal/mtprotoedge/rpc_result_egress_test.go
View file

@ -203,7 +203,7 @@ func (h *saturatedSlotWaveRPC) Dispatch(context.Context, [8]byte, int64, *bin.Bu
func (*saturatedSlotWaveRPC) NegotiatedLayer([8]byte, int64) (int, bool) { return 227, true }
func TestPublishRPCResultSaturatedBudgetRetainsExactResultsAcrossSlotWaves(t *testing.T) {
func TestPublishRPCResultSaturatedBudgetLeavesExecutionTombstonesAcrossSlotWaves(t *testing.T) {
slotCount := cap(outboundEncodeSlots)
requestCount := slotCount*2 + 1
gate := &saturatedSlotWaveGate{
@ -214,7 +214,7 @@ func TestPublishRPCResultSaturatedBudgetRetainsExactResultsAcrossSlotWaves(t *te
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)
s.rpcResults = newRPCExecutionLedgerForServerTest(s, func() time.Time { return now }, requestCount+1)
const primaryMax = 1 << 20
primary := newOutboundTrackedBudget(primaryMax)
@ -302,7 +302,6 @@ func TestPublishRPCResultSaturatedBudgetRetainsExactResultsAcrossSlotWaves(t *te
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])
@ -311,49 +310,31 @@ func TestPublishRPCResultSaturatedBudgetRetainsExactResultsAcrossSlotWaves(t *te
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)
t.Fatalf("request %d owner was not handed to execution ledger", 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)
if cached, ok := s.rpcResults.Replay(c.authKeyID, c.sessionID, reqMsgIDs[i]); ok || cached != nil {
t.Fatalf("request %d retained a payload despite outbox saturation", 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 _, err := s.rpcResults.Acquire(c.authKeyID, c.sessionID, reqMsgIDs[i]); !errors.Is(err, ErrRPCResultFlightCapacity) {
t.Fatalf("retry request %d err=%v, want execution tombstone capacity", i, 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)
if got := s.rpcResults.receiptBudgetBytes(); got != int64(requestCount*rpcExecutionReceiptBudgetBytes) {
t.Fatalf("receipt budget = %d, want %d fixed metadata bytes", got, requestCount*rpcExecutionReceiptBudgetBytes)
}
// 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)
// Expiry releases only execution receipts; no result body was retained.
now = now.Add(rpcExecutionReceiptTTL + 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 _, ok := s.rpcResults.Replay(c.authKeyID, c.sessionID, reqMsgIDs[i]); ok {
t.Fatalf("request %d remained replayable after forced expiry", i)
}
}
if got := s.rpcResults.completedBytes.snapshot(); got != 0 {
t.Fatalf("completed-cache bytes after expiry = %d, want 0", got)
if got := s.rpcResults.receiptBudgetBytes(); got != 0 {
t.Fatalf("receipt budget after expiry = %d, want 0", got)
}
primary.release(primaryMax)
if got := primary.snapshot(); got != 0 {
@ -398,7 +379,7 @@ func TestCachedReplayRestoreIsSynchronousAndIndependentOfGlobalHookExecutor(t *t
})
var restored atomic.Bool
if err := s.sendCachedRPCResultWithHook(context.Background(), c, encoded, func() error {
if err := s.sendReplayedRPCResultWithHook(context.Background(), c, encoded, func() error {
if got := len(transport.snapshot()); got != 1 {
return errors.New("replay restore ran before physical write")
}
@ -672,7 +653,7 @@ func TestWrappedConvergenceMethodDrivesEgressAndReplayPriority(t *testing.T) {
var cached *encodedOutboundMessage
deadline := time.Now().Add(time.Second)
for time.Now().Before(deadline) {
if got, ok := s.rpcResults.Get(c.authKeyID, c.sessionID, reqMsgID); ok {
if got, ok := s.rpcResults.Replay(c.authKeyID, c.sessionID, reqMsgID); ok {
cached = got
break
}
@ -725,7 +706,7 @@ func TestRPCWorkerReleasesAfterEgressAdmissionWhileWriteBlocked(t *testing.T) {
tr.once.Do(func() { close(tr.release) })
deadline := time.Now().Add(time.Second)
for {
if _, ok := s.rpcResults.Get(c.authKeyID, c.sessionID, reqMsgID); ok {
if _, ok := s.rpcResults.Replay(c.authKeyID, c.sessionID, reqMsgID); ok {
break
}
if time.Now().After(deadline) {
@ -754,7 +735,7 @@ func TestDeliveryHookRunsOnceAfterReplayNotFailedWrite(t *testing.T) {
deadline := time.Now().Add(time.Second)
var cached *encodedOutboundMessage
for time.Now().Before(deadline) {
if got, ok := s.rpcResults.Get(oldConn.authKeyID, oldConn.sessionID, reqMsgID); ok {
if got, ok := s.rpcResults.Replay(oldConn.authKeyID, oldConn.sessionID, reqMsgID); ok {
cached = got
break
}
@ -772,7 +753,7 @@ func TestDeliveryHookRunsOnceAfterReplayNotFailedWrite(t *testing.T) {
replayTransport := &failAfterTransport{}
replayConn := newOutboundTestConn(t, replayTransport, newOutboundTrackedBudget(1<<20))
if err := s.sendCachedRPCResult(context.Background(), replayConn, cached); err != nil {
if err := s.sendReplayedRPCResult(context.Background(), replayConn, cached); err != nil {
t.Fatalf("replay result: %v", err)
}
deadline = time.Now().Add(time.Second)
@ -789,7 +770,7 @@ func TestDeliveryHookRunsOnceAfterReplayNotFailedWrite(t *testing.T) {
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 {
if err := s.sendReplayedRPCResult(context.Background(), replayConn, cached); err != nil {
t.Fatalf("second replay: %v", err)
}
time.Sleep(20 * time.Millisecond)

22
internal/mtprotoedge/rpc_result_retention_budget_test.go
View file

@ -80,7 +80,7 @@ func TestCachedRPCResultReplayUsesPreReservedBodyWithoutDoubleCharge(t *testing.
done := make(chan error, 1)
go func() {
done <- s.sendCachedRPCResult(context.Background(), c, encoded)
done <- s.sendReplayedRPCResult(context.Background(), c, encoded)
}()
select {
case <-tr.started:
@ -162,7 +162,7 @@ func TestOutboundActorRetargetRequiresSecondBodyReservation(t *testing.T) {
state := newOutboundState(budget)
var terminalErr error
var terminalBytes int64
c.handleOutboundSend(state, outboundOp{
op := outboundOp{
kind: outboundSend,
ctx: context.Background(),
msgType: proto.MessageServerResponse,
@ -174,12 +174,14 @@ func TestOutboundActorRetargetRequiresSecondBodyReservation(t *testing.T) {
terminalErr = err
terminalBytes = budget.snapshot()
},
})
}
err := c.handleOutboundSend(state, op)
op.finish(outboundResult{err: err})
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 terminalBytes != 0 {
t.Fatalf("bytes visible to terminal = %d, want producer reservation released", terminalBytes)
}
if got := tr.sends.Load(); got != 0 {
t.Fatalf("retarget under one-body budget wrote %d frames, want 0", got)
@ -213,7 +215,7 @@ func TestOutboundActorRetargetTransfersOnlyReplacementToPending(t *testing.T) {
state := newOutboundState(budget)
var terminalErr error
var terminalBytes int64
c.handleOutboundSend(state, outboundOp{
op := outboundOp{
kind: outboundSend,
ctx: context.Background(),
msgType: proto.MessageServerResponse,
@ -225,12 +227,14 @@ func TestOutboundActorRetargetTransfersOnlyReplacementToPending(t *testing.T) {
terminalErr = err
terminalBytes = budget.snapshot()
},
})
}
err := c.handleOutboundSend(state, op)
op.finish(outboundResult{err: err})
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 terminalBytes != int64(perBody) {
t.Fatalf("bytes visible to terminal = %d, want only pending replacement %d", terminalBytes, perBody)
}
if got := budget.snapshot(); got != int64(perBody) {
t.Fatalf("bytes after terminal = %d, want one pending replacement %d", got, perBody)

28
internal/mtprotoedge/rpc_result_subscriber_budget.go
View file

@ -27,8 +27,8 @@ type rpcResultSubscriberBudget struct {
global rpcResultFlightLimit
authLimit int64
sessionLimit int64
authShards [rpcResultBudgetShards]rpcResultSubscriberBudgetShard[[8]byte]
sessionShards [rpcResultBudgetShards]rpcResultSubscriberBudgetShard[rpcResultSessionBudgetKey]
authShards [rpcExecutionBudgetShards]rpcResultSubscriberBudgetShard[[8]byte]
sessionShards [rpcExecutionBudgetShards]rpcResultSubscriberBudgetShard[rpcExecutionSessionBudgetKey]
}
func newRPCResultSubscriberBudget(
@ -43,25 +43,25 @@ func newRPCResultSubscriberBudget(
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)
b.sessionShards[i].usage = make(map[rpcExecutionSessionBudgetKey]int64)
}
return b
}
func (b *rpcResultSubscriberBudget) reserve(key rpcResultCacheKey, slots int) bool {
func (b *rpcResultSubscriberBudget) reserve(key rpcExecutionKey, slots int) bool {
if b == nil || slots <= 0 {
return false
}
authShard := b.authShard(key.authKeyID)
sessionKey := rpcResultSessionBudgetKey{authKeyID: key.authKeyID, sessionID: key.sessionID}
sessionKey := rpcExecutionSessionBudgetKey{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) ||
if !withinRPCExecutionBudget(authUsed, delta, b.authLimit) ||
!withinRPCExecutionBudget(sessionUsed, delta, b.sessionLimit) ||
!b.global.reserveN(delta) {
sessionShard.mu.Unlock()
authShard.mu.Unlock()
@ -74,12 +74,12 @@ func (b *rpcResultSubscriberBudget) reserve(key rpcResultCacheKey, slots int) bo
return true
}
func (b *rpcResultSubscriberBudget) release(key rpcResultCacheKey, slots int) {
func (b *rpcResultSubscriberBudget) release(key rpcExecutionKey, 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}
sessionKey := rpcExecutionSessionBudgetKey{authKeyID: key.authKeyID, sessionID: key.sessionID}
sessionShard := b.sessionShard(sessionKey)
delta := int64(slots)
authShard.mu.Lock()
@ -123,7 +123,7 @@ func (b *rpcResultSubscriberBudget) sessionSnapshot(authKeyID [8]byte, sessionID
if b == nil {
return 0
}
key := rpcResultSessionBudgetKey{authKeyID: authKeyID, sessionID: sessionID}
key := rpcExecutionSessionBudgetKey{authKeyID: authKeyID, sessionID: sessionID}
shard := b.sessionShard(key)
shard.mu.Lock()
used := shard.usage[key]
@ -134,16 +134,16 @@ func (b *rpcResultSubscriberBudget) sessionSnapshot(authKeyID [8]byte, sessionID
func (b *rpcResultSubscriberBudget) authShard(
authKeyID [8]byte,
) *rpcResultSubscriberBudgetShard[[8]byte] {
index := maphash.Bytes(b.seed, authKeyID[:]) & (rpcResultBudgetShards - 1)
index := maphash.Bytes(b.seed, authKeyID[:]) & (rpcExecutionBudgetShards - 1)
return &b.authShards[index]
}
func (b *rpcResultSubscriberBudget) sessionShard(
key rpcResultSessionBudgetKey,
) *rpcResultSubscriberBudgetShard[rpcResultSessionBudgetKey] {
key rpcExecutionSessionBudgetKey,
) *rpcResultSubscriberBudgetShard[rpcExecutionSessionBudgetKey] {
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)
index := maphash.Bytes(b.seed, raw[:]) & (rpcExecutionBudgetShards - 1)
return &b.sessionShards[index]
}

76
internal/mtprotoedge/rpc_rewrap.go
View file

@ -446,7 +446,7 @@ func (a *rpcRewrapAlias) storeResultOnce(s *Server, encoded *encodedOutboundMess
if a == nil || s == nil || encoded == nil || !a.resultStoreClaimed.CompareAndSwap(false, true) {
return
}
s.storeRPCResult(a.conn, a.newReqID, encoded)
s.completeRPCResult(a.conn, a.newReqID, encoded, true)
}
func claimRPCRewrapLogicalHook(
@ -465,8 +465,8 @@ func claimRPCRewrapLogicalHook(
// 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.
// coordinator and ledger publication make late/concurrent invocations converge
// while preserving replacement -> logical -> ledger -> barrier order.
func (s *Server) completeDeliveredRPCRewrapResult(
ctx context.Context,
a *rpcRewrapAlias,
@ -510,8 +510,10 @@ const (
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.
// control deadline. It prevents later stages from starting and fences logical
// ownership, but cannot cancel non-cooperative filesystem, transport or restore
// work. Failure publication may wait for replay preparation to leave its
// ownership transition before releasing the Conn scheduler barrier.
rpcRewrapDeliveryQueueTimeout = 5 * time.Second
)
@ -532,10 +534,10 @@ const (
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.
// rpcRewrapDeliveryControl lets the deadline timer independently transition
// queued, running and physically committed jobs to Failed. A late worker cannot
// enter run after the timer wins. The timer cannot cancel a physical transport
// write; a late return still cannot report failure or finish the barrier twice.
type rpcRewrapDeliveryControl struct {
state atomic.Uint32
timerMu sync.Mutex
@ -550,9 +552,9 @@ type rpcRewrapPhysicalOutcome struct {
// 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.
// queued jobs still transition to Failed at their control deadlines. 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,
@ -616,7 +618,7 @@ func (c *rpcRewrapDeliveryControl) timeout() bool {
// 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
// replacement/logical restore and ledger/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) {
@ -700,12 +702,6 @@ func runRPCRewrapDeliveryJob(j rpcRewrapDeliveryJob) {
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() {
@ -716,6 +712,12 @@ func runRPCRewrapDeliveryJob(j rpcRewrapDeliveryJob) {
}
control.complete()
}()
if !j.deadline.IsZero() && !time.Now().Before(j.deadline) {
if control.fail() {
j.reportFailure(context.DeadlineExceeded)
}
return
}
j.run(control, j.deadline)
}
@ -747,9 +749,10 @@ func scheduleRPCRewrapJob(
delay = 0
}
timer := time.AfterFunc(delay, func() {
if job.control.timeout() {
job.reportFailure(context.DeadlineExceeded)
if !job.control.timeout() {
return
}
job.reportFailure(context.DeadlineExceeded)
})
job.control.installTimer(timer)
select {
@ -766,6 +769,22 @@ func scheduleRPCRewrapJob(
}
}
func (s *Server) attachRPCRewrapReplayPreparation(
job *rpcRewrapDeliveryJob,
c *Conn,
reqMsgID int64,
method string,
encoded *encodedOutboundMessage,
) {
if job == nil || s == nil || c == nil || encoded == nil {
return
}
// Freeze scheduling metadata before the watchdog can publish the compact
// receipt. Exact wire bytes are owned only by the logical-session outbox, so
// there is no pre-send spool preparation or I/O gate.
encoded.priority = rpcResultPriority(method, encoded)
}
func scheduleRPCRewrapDeliveryJob(job rpcRewrapDeliveryJob) bool {
return scheduleRPCRewrapJob(job, &rpcRewrapDeliveryOnce, &rpcRewrapDeliveryJobs,
rpcRewrapDeliveryWorkers, rpcRewrapDeliveryQueue)
@ -821,7 +840,7 @@ func (s *Server) failRPCRewrapResultJob(
publish := a.newOwner.HandOff()
encoded.markReplayable()
encoded.releaseDeferredLogicalDeliveryHook()
// Release the connection-local scheduler before any defensive cache panic;
// Release the connection-local scheduler before any defensive ledger panic;
// the physical generation is already fenced, so no following task can run.
a.releaseReplayRestoreBarrier()
if publish {
@ -953,6 +972,7 @@ func (a *rpcRewrapAlias) activate(s *Server) error {
s.log.Debug("Retargeted RPC restore failed", zap.Error(restoreErr))
}
})
s.attachRPCRewrapReplayPreparation(&job, a.conn, a.newReqID, a.method, clone)
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
@ -978,6 +998,7 @@ func (a *rpcRewrapAlias) activate(s *Server) error {
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)
})
s.attachRPCRewrapReplayPreparation(&job, a.conn, a.newReqID, a.method, clone)
// 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
@ -1067,14 +1088,15 @@ func (s *Server) publishRewrappedRPCResult(
alias.releaseBodyReservation()
return
}
priority := rpcResultPriority(method, encoded)
encoded.priority = priority
encoded.markQueued()
if deadline.IsZero() {
deadline = time.Now().Add(rpcRewrapDeliveryQueueTimeout)
}
ctx, cancel := context.WithDeadline(context.Background(), deadline)
defer cancel()
// attachRPCRewrapReplayPreparation froze replay-copied scheduling metadata
// before the watchdog started. The delivery path must not mutate plain fields
// that a concurrent timeout publication can copy into the replay ledger.
encoded.markQueued()
// 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
@ -1093,10 +1115,10 @@ func (s *Server) publishRewrappedRPCResult(
return
}
// Physical success outranks an already-fired watchdog. The timeout path may
// have fenced and cached a replayable clone, but it cannot revoke bytes; the
// have fenced and published a replayable receipt, 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
// the alias receipt. Whole-finalization once also covers a watchdog racing
// a late physical terminal, so completed metadata cannot be overwritten.
restoreParent := ctx
if !outcome.owned {

68
internal/mtprotoedge/rpc_rewrap_test.go
View file

@ -1,6 +1,7 @@
package mtprotoedge
import (
"bytes"
"context"
"encoding/binary"
"errors"
@ -164,7 +165,7 @@ func TestRPCRewrapFailedSourceAttemptPhysicallyDeliversAliasOnce(t *testing.T) {
deadline := time.Now().Add(2 * time.Second)
for time.Now().Before(deadline) {
cached, ok := s.rpcResults.Get(key.ID, sessionID, newReqID)
cached, ok := s.rpcResults.Replay(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)
@ -172,7 +173,7 @@ func TestRPCRewrapFailedSourceAttemptPhysicallyDeliversAliasOnce(t *testing.T) {
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)
sourceCached, sourceOK := s.rpcResults.Replay(key.ID, sessionID, oldReqID)
if !sourceOK || sourceCached.deliveryState() != rpcResultDeliveryReplayable {
t.Fatalf("source physical attempt = cached:%v state:%d, want replayable", sourceOK, sourceCached.deliveryState())
}
@ -180,13 +181,13 @@ func TestRPCRewrapFailedSourceAttemptPhysicallyDeliversAliasOnce(t *testing.T) {
}
time.Sleep(time.Millisecond)
}
cached, ok := s.rpcResults.Get(key.ID, sessionID, newReqID)
cached, ok := s.rpcResults.Replay(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)
s.rpcResults = newRPCExecutionSubscriberTestLedger(2, 2, 2, 2)
transport := &collectingSessionTransport{}
key := newTestAuthKey(t)
const sessionID, oldReqID, newReqID = int64(188), int64(15101), int64(15201)
@ -288,7 +289,7 @@ func TestRPCRewrapRetargetFailureRequiresReplacementAliasWrite(t *testing.T) {
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 {
if cached, ok := s.rpcResults.Replay(key.ID, sessionID, newReqID); ok && cached.deliveryState() == rpcResultDeliveryReplayable {
aliasCached = cached
break
}
@ -305,7 +306,7 @@ func TestRPCRewrapRetargetFailureRequiresReplacementAliasWrite(t *testing.T) {
replacement := s.newConn(replacementTransport, key, sessionID, 1)
legacyCanonicalTestConn(t, replacement)
t.Cleanup(replacement.ForceClose)
if err := s.sendCachedRPCResult(context.Background(), replacement, aliasCached); err != nil {
if err := s.sendReplayedRPCResult(context.Background(), replacement, aliasCached); err != nil {
t.Fatalf("replacement alias replay: %v", err)
}
deadline = time.Now().Add(time.Second)
@ -318,7 +319,7 @@ func TestRPCRewrapRetargetFailureRequiresReplacementAliasWrite(t *testing.T) {
}
func TestRPCResultWaiterSubscribeIsEventDriven(t *testing.T) {
cache := newRPCResultCacheWithFlightLimit(time.Now, 8)
cache := newRPCExecutionLedgerForTest(time.Now, 8)
claim, err := cache.Acquire([8]byte{1}, 2, 3)
if err != nil || claim.state != rpcResultAcquireOwner {
t.Fatalf("Acquire owner = %+v, %v", claim, err)
@ -336,14 +337,14 @@ func TestRPCResultWaiterSubscribeIsEventDriven(t *testing.T) {
t.Fatal("Subscribe waited for or fabricated a result")
}
encoded := &encodedOutboundMessage{typeID: mt.RPCResultTypeID, body: make([]byte, 16), reqMsgID: 3}
cache.Put([8]byte{1}, 2, 3, encoded)
cache.completeReplayableForTest([8]byte{1}, 2, 3, encoded)
if !called.Load() {
t.Fatal("completion event did not invoke subscriber")
}
}
func TestRPCResultOwnerAbortHookInstallationIsFlightBound(t *testing.T) {
cache := newRPCResultCacheWithFlightLimit(time.Now, 8)
cache := newRPCExecutionLedgerForTest(time.Now, 8)
claim, err := cache.Acquire([8]byte{2}, 3, 4)
if err != nil || claim.state != rpcResultAcquireOwner {
t.Fatalf("Acquire owner = %+v, %v", claim, err)
@ -361,7 +362,7 @@ func TestRPCResultOwnerAbortHookInstallationIsFlightBound(t *testing.T) {
}
func TestRPCRewrapRegistryIsPlatformAgnosticAndAckBound(t *testing.T) {
cache := newRPCResultCacheWithFlightLimit(time.Now, 8)
cache := newRPCExecutionLedgerForTest(time.Now, 8)
claim, err := cache.Acquire([8]byte{4}, 5, 6)
if err != nil || claim.state != rpcResultAcquireOwner {
t.Fatalf("Acquire owner = %+v, %v", claim, err)
@ -445,12 +446,12 @@ func TestInitRewrapAfterWritingReplaysWithoutBusinessExecution(t *testing.T) {
if !oldOwner.HandOff() {
t.Fatal("old owner handoff failed")
}
s.rpcResults.Put(c.authKeyID, c.sessionID, oldReqID, encoded)
storeLogicalRPCResultForTest(t, s, c, oldReqID, encoded)
deadline := time.Now().Add(2 * time.Second)
var replayed *encodedOutboundMessage
for time.Now().Before(deadline) {
if got, ok := s.rpcResults.Get(c.authKeyID, c.sessionID, newReqID); ok {
if got, ok := s.rpcResults.Replay(c.authKeyID, c.sessionID, newReqID); ok {
replayed = got
break
}
@ -534,7 +535,7 @@ func TestInitRewrapAliasesExecutionAndRetargetsQueuedResult(t *testing.T) {
t.Fatal("old owner handoff failed")
}
encoded.markDelivered()
s.rpcResults.Put(c.authKeyID, c.sessionID, oldReqID, encoded)
storeLogicalRPCResultForTest(t, s, c, oldReqID, encoded)
var (
aliased *encodedOutboundMessage
@ -542,7 +543,7 @@ func TestInitRewrapAliasesExecutionAndRetargetsQueuedResult(t *testing.T) {
)
deadline := time.Now().Add(time.Second)
for time.Now().Before(deadline) {
if aliased, ok = s.rpcResults.Get(c.authKeyID, c.sessionID, newReqID); ok {
if aliased, ok = s.rpcResults.Replay(c.authKeyID, c.sessionID, newReqID); ok {
break
}
time.Sleep(time.Millisecond)
@ -608,8 +609,8 @@ func TestRPCRewrapDeliveryJobPanicAndDeadlineReleaseBarrier(t *testing.T) {
}
}
func TestExpiredRPCRewrapResultJobPublishesCompletedAliasExactlyOnce(t *testing.T) {
cache := newRPCResultCacheWithFlightLimit(time.Now, 1)
func TestExpiredRPCRewrapResultJobPublishesUnavailableAliasExactlyOnce(t *testing.T) {
cache := newRPCExecutionLedgerForTest(time.Now, 1)
s := &Server{log: zaptest.NewLogger(t), rpcResults: cache}
c := &Conn{
metrics: NopMetrics{},
@ -667,27 +668,22 @@ func TestExpiredRPCRewrapResultJobPublishesCompletedAliasExactlyOnce(t *testing.
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)
completed, ok := cache.Replay(c.authKeyID, c.sessionID, reqMsgID)
if ok || completed != nil {
t.Fatalf("failed alias retained payload = (%p, %v)", completed, ok)
}
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)
if _, err := cache.Acquire(c.authKeyID, c.sessionID, reqMsgID); !errors.Is(err, ErrRPCResultFlightCapacity) {
t.Fatalf("reacquire unavailable aliased result = %v, want capacity", err)
}
// A defensive duplicate failure report must not republish or underflow the
// completed flight. The first handoff/cache completion is the sole winner.
// completed flight. The first handoff/ledger 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)
if got, ok := cache.Replay(c.authKeyID, c.sessionID, reqMsgID); ok || got != nil {
t.Fatalf("unavailable result changed after duplicate failure = (%p, %v)", got, ok)
}
}
@ -772,7 +768,7 @@ func TestRetargetedRPCRestoreIsOrderedAndIndependentOfGlobalHookExecutor(t *test
t.Fatalf("retargeted physical req_msg_id = %d, want %d", got, newReqID)
}
encoded.markDelivered()
s.rpcResults.Put(c.authKeyID, c.sessionID, oldReqID, encoded)
storeLogicalRPCResultForTest(t, s, c, oldReqID, encoded)
deadline := time.Now().Add(time.Second)
for order.Load() != 2 && time.Now().Before(deadline) {
@ -787,7 +783,7 @@ func TestRetargetedRPCRestoreIsOrderedAndIndependentOfGlobalHookExecutor(t *test
if pending != 0 {
t.Fatalf("retarget restore barriers = %d, want 0", pending)
}
if _, ok := s.rpcResults.Get(c.authKeyID, c.sessionID, newReqID); !ok {
if _, ok := s.rpcResults.Replay(c.authKeyID, c.sessionID, newReqID); !ok {
t.Fatal("retargeted result was not cached under new req_msg_id")
}
}
@ -979,7 +975,7 @@ func TestRPCRewrapPhysicalSuccessAfterWatchdogStillRunsLogicalRestore(t *testing
}
func TestConcurrentRPCRewrapDeliveredFinalizationPublishesOnceWithMetadata(t *testing.T) {
cache := newRPCResultCacheWithFlightLimit(time.Now, 1)
cache := newRPCExecutionLedgerForTest(time.Now, 1)
s := &Server{log: zaptest.NewLogger(t), rpcResults: cache}
c := &Conn{
metrics: NopMetrics{},
@ -1005,6 +1001,7 @@ func TestConcurrentRPCRewrapDeliveredFinalizationPublishesOnceWithMetadata(t *te
encoded := encodedRPCResultForPriorityTest(reqMsgID, 0)
encoded.delivery = claim.owner.Delivery()
encoded.setDeliveryHook(func() { logical.Add(1) })
cache.replayStore.(*rpcReplayStoreForTest).put(c.authKeyID, c.sessionID, reqMsgID, encoded)
alias := &rpcRewrapAlias{
conn: c, newReqID: reqMsgID, method: "help.getConfig", newOwner: claim.owner,
afterSuccessfulDelivery: func() error {
@ -1035,7 +1032,7 @@ func TestConcurrentRPCRewrapDeliveredFinalizationPublishesOnceWithMetadata(t *te
t.Fatalf("logical finalizations = %d, want 1", got)
}
if got := subscribers.Load(); got != 1 {
t.Fatalf("cache subscriber calls = %d, want 1", got)
t.Fatalf("ledger subscriber calls = %d, want 1", got)
}
replay, err := cache.Acquire(c.authKeyID, c.sessionID, reqMsgID)
if err != nil || replay.state != rpcResultAcquireCompleted ||
@ -1108,7 +1105,8 @@ func TestRPCRewrapSubscriberPanicCannotLoseClaimedLogicalHook(t *testing.T) {
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)
if cached, ok := s.rpcResults.Replay(c.authKeyID, c.sessionID, reqMsgID); !ok ||
cached.delivery != encoded.delivery || !bytes.Equal(cached.body, encoded.body) || cached.replayMsgID == 0 {
t.Fatalf("logical outbox result after subscriber panic = (%p, %v), source=%p", cached, ok, encoded)
}
}

15
internal/mtprotoedge/rpc_test.go
View file

@ -62,9 +62,12 @@ func TestRPCGetConfig(t *testing.T) {
if cfg.ThisDC != dc {
t.Fatalf("config.ThisDC = %d, want %d", cfg.ThisDC, dc)
}
// 不下发 DCOptions：客户端使用写死的 static DC 地址（空列表令其保留本地地址）。
if len(cfg.DCOptions) != 0 {
t.Fatalf("config.DCOptions = %+v, want empty (client uses pinned static address)", cfg.DCOptions)
if len(cfg.DCOptions) != 1 {
t.Fatalf("config.DCOptions = %+v, want one reconnect route", cfg.DCOptions)
}
option := cfg.DCOptions[0]
if option.ID != dc || option.IPAddress != advIP || option.Port != advPort {
t.Fatalf("config.DCOptions[0] = %+v, want dc=%d at %s:%d", option, dc, advIP, advPort)
}
}
@ -118,7 +121,7 @@ func TestLayerRPCGetConfigUsesExactAdmittedProfile(t *testing.T) {
}
}
func TestInboundRPCQueueFullReturnsFloodWait(t *testing.T) {
func TestInboundRPCQueueFullReturnsWorkerBusy(t *testing.T) {
const dc = 2
handler := &blockingRPC{
started: make(chan struct{}, 1),
@ -152,8 +155,8 @@ func TestInboundRPCQueueFullReturnsFloodWait(t *testing.T) {
if err := rpcErr.Decode(&bin.Buffer{Buf: result.Result}); err != nil {
t.Fatalf("decode rpc_error: %v", err)
}
if rpcErr.ErrorCode != 420 || rpcErr.ErrorMessage != "FLOOD_WAIT_1" {
t.Fatalf("rpc_error = %d %q, want 420 FLOOD_WAIT_1", rpcErr.ErrorCode, rpcErr.ErrorMessage)
if rpcErr.ErrorCode != rpcWorkerBusyErrorCode || rpcErr.ErrorMessage != rpcWorkerBusyErrorMessage {
t.Fatalf("rpc_error = %d %q, want %d %s", rpcErr.ErrorCode, rpcErr.ErrorMessage, rpcWorkerBusyErrorCode, rpcWorkerBusyErrorMessage)
}
close(handler.release)
}

188
internal/mtprotoedge/runtime_metrics.go Normal file
View file

@ -0,0 +1,188 @@
package mtprotoedge
// RuntimeSnapshot is a point-in-time, identity-free view of the MTProto edge.
// It deliberately exposes only bounded aggregate values so callers can publish
// it through a metrics endpoint without leaking auth keys, sessions or remote
// addresses. Values from independently locked components can differ by one
// concurrent transition; every individual budget/count remains internally
// consistent.
type RuntimeSnapshot struct {
RawConnections int64
RawConnectionLimit int64
Handshakes int64
HandshakeLimit int64
ActiveSessions int64
ProvisionalSessions int64
LogicalSessions int64
OfflineLogicalSessions int64
LogicalOutboxFrames int64
LogicalOutboxBytes int64
PendingPushBytes int64
InboundRPCTasks int64
InboundRPCBytes int64
InboundRPCReadyConnections int64
InboundRPCMaxTasks int64
InboundRPCMaxBytes int64
InboundFrameBytes int64
InboundFrameMaxBytes int64
OutboundTrackedBytes int64
OutboundTrackedMaxBytes int64
OutboundControlBytes int64
OutboundControlMaxBytes int64
OutboundWriteBytes int64
OutboundWriteMaxBytes int64
RPCExecutionOwners int64
RPCExecutionReservedEntries int64
RPCExecutionReceipts int64
RPCExecutionReceiptBudgetBytes int64
RPCExecutionSubscribers int64
}
type sessionManagerRuntimeSnapshot struct {
active int64
provisional int64
logical int64
offlineLogical int64
frames int64
bytes int64
pendingBytes int64
}
func (m *SessionManager) runtimeSnapshot() sessionManagerRuntimeSnapshot {
if m == nil {
return sessionManagerRuntimeSnapshot{}
}
// Never hold SessionManager.mu while taking an outbound-state mutex. The
// physical actor can publish/retire a Conn next to an outbox transition, and
// metrics must not add a new cross-component lock order.
m.mu.RLock()
states := make([]*outboundState, 0, len(m.logicalSessions))
result := sessionManagerRuntimeSnapshot{
active: int64(len(m.bySession)),
provisional: int64(len(m.claims)),
logical: int64(len(m.logicalSessions)),
}
if m.pendingBudget != nil {
result.pendingBytes = m.pendingBudget.snapshot()
}
for _, logical := range m.logicalSessions {
if logical == nil {
continue
}
if !logical.offlineAt.IsZero() {
result.offlineLogical++
}
if logical.outbound != nil {
states = append(states, logical.outbound)
}
}
m.mu.RUnlock()
for _, state := range states {
state.mu.Lock()
result.frames += int64(len(state.pending))
result.bytes += int64(state.totalBytes)
state.mu.Unlock()
}
return result
}
type admissionRuntimeSnapshot struct {
connections int64
connectionLimit int64
handshakes int64
handshakeLimit int64
}
func (a *admissionController) runtimeSnapshot() admissionRuntimeSnapshot {
if a == nil {
return admissionRuntimeSnapshot{}
}
a.mu.Lock()
result := admissionRuntimeSnapshot{
connections: int64(a.connections),
connectionLimit: int64(a.maxConnections),
}
a.mu.Unlock()
if a.handshakes != nil {
result.handshakes = int64(len(a.handshakes))
result.handshakeLimit = int64(cap(a.handshakes))
}
return result
}
type inboundRPCRuntimeSnapshot struct {
tasks int64
bytes int64
ready int64
}
func (s *inboundRPCScheduler) runtimeSnapshot() inboundRPCRuntimeSnapshot {
if s == nil {
return inboundRPCRuntimeSnapshot{}
}
s.budgetMu.Lock()
result := inboundRPCRuntimeSnapshot{tasks: int64(s.tasks), bytes: s.bytes}
s.budgetMu.Unlock()
s.readyMu.Lock()
result.ready = int64(s.ready.Len())
s.readyMu.Unlock()
return result
}
// RuntimeSnapshot returns aggregate MTProto ownership and capacity state.
func (s *Server) RuntimeSnapshot() RuntimeSnapshot {
if s == nil {
return RuntimeSnapshot{}
}
sessions := s.conns.runtimeSnapshot()
admission := s.admission.runtimeSnapshot()
inbound := s.rpcScheduler.runtimeSnapshot()
result := RuntimeSnapshot{
RawConnections: admission.connections,
RawConnectionLimit: admission.connectionLimit,
Handshakes: admission.handshakes,
HandshakeLimit: admission.handshakeLimit,
ActiveSessions: sessions.active,
ProvisionalSessions: sessions.provisional,
LogicalSessions: sessions.logical,
OfflineLogicalSessions: sessions.offlineLogical,
LogicalOutboxFrames: sessions.frames,
LogicalOutboxBytes: sessions.bytes,
PendingPushBytes: sessions.pendingBytes,
InboundRPCTasks: inbound.tasks,
InboundRPCBytes: inbound.bytes,
InboundRPCReadyConnections: inbound.ready,
}
if s.rpcScheduler != nil {
result.InboundRPCMaxTasks = int64(s.rpcScheduler.maxTasks)
result.InboundRPCMaxBytes = s.rpcScheduler.maxBytes
}
if s.frameBudget != nil {
result.InboundFrameBytes = s.frameBudget.usedBytes()
result.InboundFrameMaxBytes = s.frameBudget.max
}
if s.outboundTrackedBudget != nil {
result.OutboundTrackedBytes = s.outboundTrackedBudget.snapshot()
result.OutboundTrackedMaxBytes = s.outboundTrackedBudget.maxBytes
}
if s.outboundControlBudget != nil {
result.OutboundControlBytes = s.outboundControlBudget.snapshot()
result.OutboundControlMaxBytes = s.outboundControlBudget.maxBytes
}
if s.outboundScratchPool != nil && s.outboundScratchPool.budget != nil {
result.OutboundWriteBytes = s.outboundScratchPool.snapshot()
result.OutboundWriteMaxBytes = s.outboundScratchPool.budget.maxBytes
}
if s.rpcResults != nil {
result.RPCExecutionOwners = s.rpcResults.flightLimit.snapshot()
result.RPCExecutionReservedEntries = s.rpcResults.reservedEntries.snapshot()
result.RPCExecutionReceipts = s.rpcResults.receiptCount.Load()
result.RPCExecutionReceiptBudgetBytes = s.rpcResults.receiptBudgetBytes()
if s.rpcResults.subscriberBudget != nil {
result.RPCExecutionSubscribers = s.rpcResults.subscriberBudget.global.snapshot()
}
}
return result
}

57
internal/mtprotoedge/runtime_metrics_test.go Normal file
View file

@ -0,0 +1,57 @@
package mtprotoedge
import (
"testing"
"time"
)
func TestRuntimeSnapshotIsNilSafeAndReportsConfiguredLimits(t *testing.T) {
if got := (*Server)(nil).RuntimeSnapshot(); got != (RuntimeSnapshot{}) {
t.Fatalf("nil server snapshot = %#v, want zero", got)
}
if got := (&Server{}).RuntimeSnapshot(); got != (RuntimeSnapshot{}) {
t.Fatalf("partial server snapshot = %#v, want zero", got)
}
server := New(Options{})
snapshot := server.RuntimeSnapshot()
if snapshot.RawConnectionLimit <= 0 || snapshot.HandshakeLimit <= 0 {
t.Fatalf("admission limits not reported: %#v", snapshot)
}
if snapshot.InboundRPCMaxTasks <= 0 || snapshot.InboundRPCMaxBytes <= 0 {
t.Fatalf("inbound RPC limits not reported: %#v", snapshot)
}
if snapshot.InboundFrameMaxBytes <= 0 || snapshot.OutboundTrackedMaxBytes <= 0 || snapshot.OutboundWriteMaxBytes <= 0 {
t.Fatalf("byte limits not reported: %#v", snapshot)
}
if snapshot.RawConnections != 0 || snapshot.ActiveSessions != 0 || snapshot.LogicalOutboxBytes != 0 {
t.Fatalf("fresh server reported live ownership: %#v", snapshot)
}
}
func TestRuntimeSnapshotSeparatesExecutionOwnersReceiptsAndBudget(t *testing.T) {
server := New(Options{})
server.rpcResults = newRPCExecutionLedgerForTest(time.Now, 4)
auth := [8]byte{1, 2, 3, 4}
claim, err := server.rpcResults.Acquire(auth, 5, 6)
if err != nil || claim.state != rpcResultAcquireOwner {
t.Fatalf("owner = %#v, %v", claim, err)
}
pending := server.RuntimeSnapshot()
if pending.RPCExecutionOwners != 1 || pending.RPCExecutionReservedEntries != 1 ||
pending.RPCExecutionReceipts != 0 || pending.RPCExecutionReceiptBudgetBytes != 0 {
t.Fatalf("pending execution snapshot = %#v", pending)
}
server.rpcResults.completeReplayableForTest(auth, 5, 6, &encodedOutboundMessage{body: make([]byte, 4<<20)})
completed := server.RuntimeSnapshot()
if completed.RPCExecutionOwners != 0 || completed.RPCExecutionReservedEntries != 1 ||
completed.RPCExecutionReceipts != 1 || completed.RPCExecutionReceiptBudgetBytes != rpcExecutionReceiptBudgetBytes {
t.Fatalf("completed execution snapshot = %#v", completed)
}
server.rpcResults.Acknowledge(auth, 5, 6)
released := server.RuntimeSnapshot()
if released.RPCExecutionReservedEntries != 0 || released.RPCExecutionReceipts != 0 ||
released.RPCExecutionReceiptBudgetBytes != 0 {
t.Fatalf("released execution snapshot = %#v", released)
}
}

19
internal/mtprotoedge/same_port_mux.go
View file

@ -1,7 +1,6 @@
package mtprotoedge
import (
"bytes"
"context"
"errors"
"io"
@ -195,10 +194,9 @@ func (m *samePortMux) dispatch(ctx context.Context, conn net.Conn) {
return
}
wrapped := &prefixedNetConn{
Conn: conn,
reader: io.MultiReader(bytes.NewReader(header[:]), conn),
}
// Fixed eight-byte replay storage avoids bytes.Reader + MultiReader allocations
// on every accepted same-port connection.
wrapped := newReplayNetConn(conn, header[:])
target := m.tcp
transport := "tcp"
@ -331,17 +329,6 @@ func isSamePortMuxClosed(ch <-chan struct{}) bool {
}
}
// prefixedNetConn 把被窥探掉的前缀字节回放在数据流最前面，使下游(去混淆/codec 探测/
// http.Server)看到完整原始字节流。
type prefixedNetConn struct {
reader io.Reader
net.Conn
}
func (p *prefixedNetConn) Read(b []byte) (int, error) {
return p.reader.Read(b)
}
// samePortMuxListener 是一个内存 listener：dispatch 把分流后的连接投递进来，下游
// (serveMixed 的 accept 循环 / http.Server) 从这里 Accept。
type samePortMuxListener struct {

200
internal/mtprotoedge/server.go
View file

@ -56,7 +56,7 @@ type legacyRPCHandlerWithMethod interface {
// 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.
// execution-ledger/scheduler ownership is acquired.
type LayerRPCHandler interface {
AdmitLayer(profile tlprofile.Profile, b *bin.Buffer, limits tlprofile.Limits) (tlprofile.Admission, error)
AdmitUnprofiled(b *bin.Buffer, limits tlprofile.Limits) (tlprofile.Admission, error)
@ -70,6 +70,14 @@ type LayerRPCHandler interface {
) (tlprofile.Result, string, error)
}
// LayerRPCOptionsAdmitter extends the stable handler boundary with
// caller-owned admission capabilities. Implementations that do not expose it
// remain usable for requests that need only Limits.
type LayerRPCOptionsAdmitter interface {
AdmitLayerWithOptions(profile tlprofile.Profile, b *bin.Buffer, options tlprofile.AdmissionOptions) (tlprofile.Admission, error)
AdmitUnprofiledWithOptions(b *bin.Buffer, options tlprofile.AdmissionOptions) (tlprofile.Admission, error)
}
// LayerRPCDefaultProfileAdmitter decodes with a recoverable inherited/default
// profile. Production handlers should implement it with the same sparse
// tlprofile dispatcher and semantic adapter registry used by AdmitLayer. The split keeps old
@ -79,6 +87,23 @@ type LayerRPCDefaultProfileAdmitter interface {
AdmitDefaultLayer(profile tlprofile.Profile, b *bin.Buffer, limits tlprofile.Limits) (tlprofile.Admission, error)
}
// LayerRPCDefaultProfileOptionsAdmitter is the capability-aware form used when
// exact admission needs caller-owned resources such as bounded gzip expansion.
type LayerRPCDefaultProfileOptionsAdmitter interface {
AdmitDefaultLayerWithOptions(profile tlprofile.Profile, b *bin.Buffer, options tlprofile.AdmissionOptions) (tlprofile.Admission, error)
}
// LayerRPCFlatBytesPayloadSizer is an optional, allocation-free admission
// capability for exact terminal requests whose generated object graph contains
// one already-bounded flat bytes payload. The handler may return ok only after
// proving the complete terminal wire shape and every semantic field cap it
// relies on. The edge still owns the multiplier, graph slack and all process /
// connection budgets; an absent or invalid hint falls back to the conservative
// generic graph charge.
type LayerRPCFlatBytesPayloadSizer interface {
LayerRPCFlatBytesPayloadSize(wire []byte) (payloadBytes int, ok bool)
}
// 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
@ -89,7 +114,7 @@ type LayerRPCSessionProfileResolver interface {
// 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
// an old retained 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)
@ -161,7 +186,7 @@ type LayerRPCDurableSessionProfileDeleter interface {
// 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.
// replayed rpc_result reaches the replacement physical connection.
type LayerRPCReplayPreparer interface {
PrepareAdmittedReplay(
ctx context.Context,
@ -187,7 +212,7 @@ type LayerRPCProfileEvidenceContext interface {
// 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.
// joined/replayed requests never call this hook again.
type LayerRPCAdmissionProfilePublisher interface {
PublishAdmittedLayerProfileEvidence(
ctx context.Context,
@ -221,12 +246,14 @@ type Options struct {
// codec 必须是 gotd 内置四种 codec（可包 NoHeader），或实现 InboundFrameBudgetedCodec；
// 无法在 payload 分配前预检长度的 codec 会 fail-closed。
Codec func() transport.Codec
// ObfuscatedTCP 先按 MTProto TCP obfuscation 解包，再自动探测 codec。
// Telegram Desktop 的 tcpo_only endpoint 会走这个 64 字节前缀流程。
// ObfuscatedTCP 允许裸 TCP 使用 MTProto transport obfuscation。开启时按每条
// 物理连接的首 1/4/8 字节自动区分明文 transport 与 64-byte obfuscated2，
// 随后把 wire mode + codec 冻结到同一条双向连接；Telegram Desktop 的
// tcpo_only 与未开启混淆的第三方客户端可共用同一端口。
ObfuscatedTCP bool
// WebSocket 在同一个 listener 上接受 MTProto over WebSocket(/apiws*)。
// 开启后仅在连接建立时读取前 4 字节做 HTTP/TCP 分流；MTProto TCP
// 后续仍走原 ObfuscatedTCP + codec 热路径。
// 后续仍走原 TCP wire-mode + codec 探测路径。
WebSocket bool
// WebSocketAllowedOrigins 是允许浏览器发起 WebSocket upgrade 的页面 origin。
// 空列表表示只接受无 Origin 的非浏览器客户端；"*" 表示允许所有来源（仅调试）。
@ -268,19 +295,16 @@ type Options struct {
// 等于 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
// RPCExecution*Entries bound in-flight owners and compact completed
// receipts. Payload bytes are not charged here: the logical-session
// outbox owns them under OutboundTrackedGlobalMaxBytes until ACK. ACK removes
// the receipt immediately; 331 seconds is only the no-ACK safety horizon.
RPCExecutionMaxEntries int
RPCExecutionAuthMaxEntries int
RPCExecutionSessionMaxEntries int
// RPCExecutionPendingPerAuth is an additional active-owner bound, independent
// from the retained entry limits and RPCGlobalMaxTasks. Default 2048.
RPCResultPendingPerAuth int
RPCExecutionPendingPerAuth int
// InboundFrameGlobalMaxBytes 是所有物理连接当前正在处理的 transport wire buffer
// 与最大解密 plaintext buffer 的总预算。长度前缀读取后、payload 分配前预留，默认
// 512 MiB；非正值使用默认值。
@ -328,7 +352,7 @@ type Options struct {
// 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.
// every API request must complete admission before execution-ledger scheduling.
LayerRPC LayerRPCHandler
// Metrics 接收连接层指标。默认 NopMetrics。
Metrics Metrics
@ -385,28 +409,19 @@ func (o *Options) setDefaults() {
if o.RPCGlobalMaxBytes <= 0 {
o.RPCGlobalMaxBytes = 512 << 20
}
if o.RPCResultCacheMaxEntries == 0 {
o.RPCResultCacheMaxEntries = rpcResultCacheMaxEntries
if o.RPCExecutionMaxEntries == 0 {
o.RPCExecutionMaxEntries = rpcExecutionMaxEntries
}
if o.RPCResultCacheMaxBytes == 0 {
o.RPCResultCacheMaxBytes = rpcResultCacheMaxBytes
if o.RPCExecutionAuthMaxEntries == 0 {
o.RPCExecutionAuthMaxEntries = rpcExecutionAuthMaxEntries
}
if o.RPCResultCacheAuthMaxEntries == 0 {
o.RPCResultCacheAuthMaxEntries = rpcResultCacheAuthMaxEntries
if o.RPCExecutionSessionMaxEntries == 0 {
o.RPCExecutionSessionMaxEntries = rpcExecutionSessionMaxEntries
}
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.RPCExecutionPendingPerAuth == 0 {
o.RPCExecutionPendingPerAuth = rpcExecutionPendingPerAuth
if o.RPCExecutionPendingPerAuth > o.RPCGlobalMaxTasks {
o.RPCExecutionPendingPerAuth = o.RPCGlobalMaxTasks
}
}
if o.InboundFrameGlobalMaxBytes <= 0 {
@ -441,30 +456,19 @@ 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")
func validateRPCExecutionOptions(o Options) error {
if o.RPCExecutionMaxEntries <= 0 || o.RPCExecutionAuthMaxEntries <= 0 || o.RPCExecutionSessionMaxEntries <= 0 {
return fmt.Errorf("rpc execution ledger 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.RPCExecutionMaxEntries < o.RPCExecutionAuthMaxEntries ||
o.RPCExecutionAuthMaxEntries < o.RPCExecutionSessionMaxEntries {
return fmt.Errorf("rpc execution ledger entry hierarchy must satisfy global >= auth >= session: %d/%d/%d",
o.RPCExecutionMaxEntries, o.RPCExecutionAuthMaxEntries, o.RPCExecutionSessionMaxEntries)
}
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)
if o.RPCExecutionPendingPerAuth <= 0 || o.RPCExecutionPendingPerAuth > o.RPCGlobalMaxTasks ||
o.RPCExecutionPendingPerAuth > o.RPCExecutionAuthMaxEntries {
return fmt.Errorf("rpc execution per-auth pending limit %d must be positive and <= global pending %d and auth entries %d",
o.RPCExecutionPendingPerAuth, o.RPCGlobalMaxTasks, o.RPCExecutionAuthMaxEntries)
}
return nil
}
@ -510,7 +514,7 @@ type Server struct {
types *tmap.Map
admission *admissionController
rpcResults *rpcResultCache
rpcResults *rpcExecutionLedger
rpcRewrap *rpcRewrapRegistry
// onFrame 是测试钩子：收到一帧时回调其字节数；生产为 nil。
@ -520,14 +524,14 @@ 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))
if err := validateRPCExecutionOptions(opts); err != nil {
panic(fmt.Sprintf("mtprotoedge: invalid rpc execution options: %v", err))
}
conns := opts.ActiveSessions
if conns == nil {
conns = NewSessionManager(opts.Logger.Named("sessions"))
}
return &Server{
server := &Server{
log: opts.Logger,
codec: opts.Codec,
obfuscated: opts.ObfuscatedTCP,
@ -560,19 +564,21 @@ func New(opts Options) *Server {
clock: opts.Clock,
rand: opts.Rand,
types: tmap.New(tg.TypesMap(), mt.TypesMap(), proto.TypesMap()),
rpcResults: newRPCResultCacheWithFairCapacity(opts.Clock.Now, rpcResultCacheCapacity{
rpcResults: newRPCExecutionLedger(opts.Clock.Now, rpcExecutionLedgerCapacity{
maxPending: opts.RPCGlobalMaxTasks,
maxPendingPerAuth: opts.RPCResultPendingPerAuth,
globalMaxBytes: opts.RPCResultCacheMaxBytes,
globalMaxEntries: opts.RPCResultCacheMaxEntries,
authMaxBytes: opts.RPCResultCacheAuthMaxBytes,
authMaxEntries: opts.RPCResultCacheAuthMaxEntries,
sessionMaxBytes: opts.RPCResultCacheSessionMaxBytes,
sessionMaxEntries: opts.RPCResultCacheSessionMaxEntries,
maxPendingPerAuth: opts.RPCExecutionPendingPerAuth,
globalMaxEntries: opts.RPCExecutionMaxEntries,
authMaxEntries: opts.RPCExecutionAuthMaxEntries,
sessionMaxEntries: opts.RPCExecutionSessionMaxEntries,
replayStore: conns,
}),
rpcRewrap: newRPCRewrapRegistry(opts.RPCGlobalMaxTasks),
admission: newAdmissionController(opts.MaxConnections, opts.MaxConnectionsPerIP, opts.MaxConcurrentHandshakes),
}
conns.setLogicalSessionReleaseHook(func(key sessionKey) {
server.rpcResults.forgetSession(key.authKeyID, key.sessionID)
})
return server
}
// ListenAndServe binds the public MTProto socket and immediately enters Serve.
@ -634,8 +640,16 @@ func (s *Server) buildConn(tc transport.Conn, lease *physicalTransportLease, key
outboundTrackedBudget: s.outboundTrackedBudget,
outboundControlTrackedBudget: s.outboundControlBudget,
outboundScratchPool: s.outboundScratchPool,
rpcResultAcked: s.rpcRewrap.acknowledge,
rpcResultAcked: func(conn *Conn, reqMsgID int64) {
// The sole outbound actor invokes this only after resolving a client
// msgs_ack server msg_id through its tracked resend frame. The actor has
// already removed the sole outbox frame; now delete its receipt and
// retire any init-rewrap bookkeeping.
s.rpcResults.Acknowledge(conn.authKeyID, conn.sessionID, reqMsgID)
s.rpcRewrap.acknowledge(conn, reqMsgID)
},
}
s.conns.attachLogicalSession(c, s.outboundTrackedBudget)
c.startOutbound()
c.startInboundRPCScheduler(s.rpcScheduler, s.rpcInflight, s.rpcQueueSize, s.rpcTimeout)
return c
@ -648,6 +662,7 @@ func (s *Server) Serve(ctx context.Context, ln net.Listener) error {
// serveTCP/serveMixed 返回前会等待连接 goroutine 收敛，各 Conn 已先排空/取消任务；
// 最后再停止全局池，避免关闭过程中留下无人消费但仍占预算的队列。
s.rpcScheduler.start()
defer s.conns.releaseAllLogicalSessions()
defer s.rpcScheduler.stop(rpcCloseWaitTimeout)
// 只在最外层 listener 包一次，确保 same-port mux 的 sniff/HTTP upgrade 也计入
// raw admission，而不是等连接已经分流后才计数。
@ -662,7 +677,11 @@ func (s *Server) serveTCP(ctx context.Context, ln net.Listener) error {
ctx, cancel := context.WithCancel(ctx)
defer cancel()
s.log.Info("Serving", zap.String("addr", ln.Addr().String()), zap.Int("dc", s.dc), zap.Bool("obfuscated_tcp", s.obfuscated))
s.log.Info("Serving",
zap.String("addr", ln.Addr().String()),
zap.Int("dc", s.dc),
zap.String("tcp_transport_mode", intakeTransport(s.obfuscated)),
)
defer s.log.Info("Stopped")
errCh := make(chan error, 1)
go func() {
@ -698,7 +717,7 @@ func (s *Server) serveMixed(ctx context.Context, ln net.Listener) error {
s.log.Info("Serving",
zap.String("addr", ln.Addr().String()),
zap.Int("dc", s.dc),
zap.Bool("obfuscated_tcp", s.obfuscated),
zap.String("tcp_transport_mode", intakeTransport(s.obfuscated)),
zap.Bool("websocket", true),
zap.Strings("websocket_origins", s.websocketOrigins),
)
@ -723,7 +742,7 @@ func (s *Server) serveMixed(ctx context.Context, ln net.Listener) error {
defer wg.Done()
errCh <- mux.Serve(ctx)
}()
// 裸 MTProto TCP：每条连接在自己的 goroutine 里完成去混淆 + codec 探测。
// 裸 MTProto TCP：每条连接在自己的 goroutine 里完成 wire mode + codec 探测。
go func() {
defer wg.Done()
errCh <- s.acceptLoop(ctx, mux.TCP(), s.obfuscated)
@ -764,11 +783,11 @@ func (s *Server) serveMixed(ctx context.Context, ln net.Listener) error {
return firstErr
}
// acceptLoop 接受裸连接，并为每条连接单独起 goroutine 完成「去混淆 + codec 探测 +
// acceptLoop 接受裸连接，并为每条连接单独起 goroutine 完成「wire mode + codec 探测 +
// serveConn」。探测在 accept 循环之外、带握手超时进行——慢/半开/坏 init 的客户端只占用
// 自己的 goroutine，绝不阻塞其他连接的接入；单条连接的握手失败也只关闭该连接，不会拖垮
// 整个监听循环。obfuscated 为 true 时先走 obfuscated2 去混淆（裸 MTProto TCP）；WebSocket
// 连接传 false（gotd 升级处理器已完成去混淆）。
// 整个监听循环。obfuscated 为 true 时自动区分 plain 与 obfuscated2（裸 MTProto TCP）；
// WebSocket 连接传 false（gotd 升级处理器已完成去混淆）。
func (s *Server) acceptLoop(ctx context.Context, ln net.Listener, obfuscated bool) error {
return s.acceptLoopTransport(ctx, ln, obfuscated, intakeTransport(obfuscated))
}
@ -829,7 +848,7 @@ func (s *Server) acceptLoopTransport(ctx context.Context, ln net.Listener, obfus
func (s *Server) serveDetectedConn(ctx context.Context, raw net.Conn, obfuscated bool, transportName string) {
started := time.Now()
remote, local := connRemote(raw), connLocal(raw)
// 握手读超时只覆盖去混淆 + codec 探测这一小段；用真实墙钟时间（SetReadDeadline 语义），
// 握手读超时只覆盖 wire-mode + codec 探测这一小段；用真实墙钟时间（SetReadDeadline 语义），
// 不走可能被测试注入的逻辑 clock。
if err := raw.SetReadDeadline(time.Now().Add(s.handshakeTimeout)); err != nil {
_ = raw.Close()
@ -848,7 +867,10 @@ func (s *Server) serveDetectedConn(ctx context.Context, raw net.Conn, obfuscated
}
}()
conn, err := s.promoteConn(raw, obfuscated)
conn, detectedTransport, err := s.promoteConn(raw, obfuscated)
if detectedTransport != "" {
transportName = detectedTransport
}
close(promoted)
if err != nil {
outcome := "error"
@ -885,15 +907,15 @@ func (s *Server) serveDetectedConn(ctx context.Context, raw net.Conn, obfuscated
}
}
// promoteConn 复用与 listener 组合完全一致的「obfuscated2 去混淆 + codec 探测」管线，但针对
// 单条连接，使其可在 accept 循环之外执行。obfuscated 对 WebSocket 连接必须为 false（gotd
// 升级处理器已剥离 obfuscated2 并补回 codec tag）。
func (s *Server) promoteConn(raw net.Conn, obfuscated bool) (transport.Conn, error) {
var ln net.Listener = newSingleConnListener(raw)
// promoteConn 针对单条连接执行一次 transport 提升。obfuscated=true 表示裸 TCP
// 允许混淆并自动区分 plain/obfuscated2；WebSocket 必须传 false，因为 gotd upgrade
// handler 已剥离 obfuscated2 并补回 codec tag。
func (s *Server) promoteConn(raw net.Conn, obfuscated bool) (transport.Conn, string, error) {
if obfuscated {
ln = transport.ObfuscatedListener(ln)
return s.promoteMixedTCP(raw)
}
return newCompatTransportListener(s.codec, ln, s.frameBudget).Accept()
conn, err := newCompatTransportConn(s.codec, raw, s.frameBudget)
return conn, "", err
}
// serveConn 处理单个传输连接：读帧并按 auth_key_id 分流。

25
internal/mtprotoedge/server_test.go
View file

@ -232,13 +232,13 @@ func TestServerAcceptObfuscatedAbridgedQuickAckFrame(t *testing.T) {
}
}
func TestServerSamePortWebSocketAndObfuscatedTCP(t *testing.T) {
func TestServerSamePortWebSocketAndMixedTCP(t *testing.T) {
ln, err := net.Listen("tcp", "127.0.0.1:0")
if err != nil {
t.Fatalf("listen: %v", err)
}
frames := make(chan int, 2)
frames := make(chan int, 3)
srv := New(Options{Logger: zaptest.NewLogger(t), ObfuscatedTCP: true, WebSocket: true})
srv.onFrame = func(n int) {
select {
@ -302,6 +302,27 @@ func TestServerSamePortWebSocketAndObfuscatedTCP(t *testing.T) {
expectFrameLen(t, frames, tcpPayload.Len())
_ = tcpConn.Close()
plainRaw, err := net.Dial("tcp", ln.Addr().String())
if err != nil {
t.Fatalf("plain tcp dial: %v", err)
}
plainConn, err := transport.Intermediate.Handshake(plainRaw)
if err != nil {
_ = plainRaw.Close()
t.Fatalf("plain tcp transport handshake: %v", err)
}
var plainPayload bin.Buffer
plainPayload.PutInt32(0x33445566)
plainPayload.PutInt32(0x77889900)
sendCtx, sc = context.WithTimeout(context.Background(), 5*time.Second)
if err := plainConn.Send(sendCtx, &plainPayload); err != nil {
sc()
t.Fatalf("plain tcp send: %v", err)
}
sc()
expectFrameLen(t, frames, plainPayload.Len())
_ = plainConn.Close()
cancel()
select {
case err := <-serveErr:

2
internal/mtprotoedge/session_boundary_test.go
View file

@ -156,7 +156,7 @@ func TestContainerRPCAdmissionFailureIsAtomic(t *testing.T) {
if err := rpcErr.Decode(&bin.Buffer{Buf: result.Result}); err != nil {
t.Fatalf("decode capacity rpc_error: %v", err)
}
if rpcErr.ErrorCode != 420 || rpcErr.ErrorMessage != "FLOOD_WAIT_1" {
if rpcErr.ErrorCode != rpcWorkerBusyErrorCode || rpcErr.ErrorMessage != rpcWorkerBusyErrorMessage {
t.Fatalf("capacity rpc_error = %+v", rpcErr)
}
delete(requestIDs, result.RequestMessageID)

310
internal/mtprotoedge/session_manager.go
View file

@ -56,6 +56,11 @@ const (
// 登记的 channel 数上限。membership 源于真实成员关系（大账号可能很多），interest 受客户端
// 直接控制；两者都设一个宽松上界防内存放大，超出即截断并记日志。
maxChannelIndexPerSession = 8192
// Official clients short-poll at most ten opened channels per session. Keep the
// server-side passive subscription index at the same hard bound.
maxChannelSubscriptionsPerSession = 10
defaultChannelSubscriptionTTL = 75 * time.Second
maxChannelSubscriptionTTL = 2 * time.Minute
)
// forceCloseBatchTimeout is one deadline for a whole revoke/replace/eviction batch. Conn.Close
@ -136,6 +141,11 @@ type sessionKey struct {
sessionID int64
}
type channelSubscription struct {
userID int64
expiresAt int64
}
// SessionLifecycleObserver receives active connection lifecycle events.
type SessionLifecycleObserver interface {
SessionOffline(rawAuthKeyID [8]byte, sessionID, userID int64, lastForUser bool)
@ -161,21 +171,28 @@ func notifySessionDestroyed(observer SessionLifecycleObserver, authKeyID [8]byte
type SessionManager struct {
mu sync.RWMutex
bySession map[sessionKey]*Conn
// logicalSessions owns MTProto resend state independently of physical Conn
// generations. It is bounded by the Server-wide tracked-body budget and a
// short offline retention window; ACK and destroy release bodies immediately.
logicalSessions map[sessionKey]*logicalSession
// claims owns the provisional -> active gap. A claimant is intentionally
// absent from every push/online index until its required session control frame
// is on the wire and PublishActivation validates the same owner.
claims map[sessionKey]*Conn
claimsByAuth map[[8]byte]map[int64]*Conn // raw authKeyID -> sessionID -> provisional claim
byAuthKey map[[8]byte]map[int64]*Conn // raw authKeyID → sessionID → Conn
byBusinessAuthKey map[[8]byte]map[sessionKey]*Conn
byUser map[int64]map[sessionKey]*Conn
byChannel map[int64]map[sessionKey]int64 // channelID → session → userID，用于频道 active-viewer 临时推送
bySessionChannels map[sessionKey]map[int64]struct{}
byMemberChannel map[int64]map[sessionKey]int64 // channelID → session → userID，用于已上线成员持久 update 推送
bySessionMembers map[sessionKey]map[int64]struct{}
pending map[sessionKey][]queuedPush // updates-ready 前暂存的主动推送
flushing map[sessionKey]bool // 置位时暂存正在排空的 session；排空完成前推送继续进 pending 保序
pendingBudget *outboundTrackedBudget // 未就绪 session 暂存 encoded body 的进程级上限
claims map[sessionKey]*Conn
claimsByAuth map[[8]byte]map[int64]*Conn // raw authKeyID -> sessionID -> provisional claim
byAuthKey map[[8]byte]map[int64]*Conn // raw authKeyID → sessionID → Conn
byBusinessAuthKey map[[8]byte]map[sessionKey]*Conn
byUser map[int64]map[sessionKey]*Conn
byChannel map[int64]map[sessionKey]int64 // channelID → session → userID，用于频道 active-viewer 临时推送
bySessionChannels map[sessionKey]map[int64]struct{}
bySubscribedChannel map[int64]map[sessionKey]channelSubscription
bySessionSubscriptions map[sessionKey]map[int64]int64
byMemberChannel map[int64]map[sessionKey]int64 // channelID → session → userID，用于已上线成员持久 update 推送
bySessionMembers map[sessionKey]map[int64]struct{}
pending map[sessionKey][]queuedPush // updates-ready 前暂存的主动推送
flushing map[sessionKey]bool // 置位时暂存正在排空的 session；排空完成前推送继续进 pending 保序
pendingBudget *outboundTrackedBudget // 未就绪 session 暂存 encoded body 的进程级上限
logicalSessionReleased func(sessionKey)
// pushSessions 记录经 account.registerDevice(token_type=7) 登记的「MTProto 内部
// 推送通道」session：raw auth_key_id → 该 auth_key 下已登记的 session_id 集合。
// 这类连接只发 ping，永远不会调 updates.getState（receivesUpdates 恒 false），
@ -196,21 +213,24 @@ func NewSessionManager(log *zap.Logger) *SessionManager {
log = zap.NewNop()
}
return &SessionManager{
bySession: make(map[sessionKey]*Conn),
claims: make(map[sessionKey]*Conn),
claimsByAuth: make(map[[8]byte]map[int64]*Conn),
byAuthKey: make(map[[8]byte]map[int64]*Conn),
byBusinessAuthKey: make(map[[8]byte]map[sessionKey]*Conn),
byUser: make(map[int64]map[sessionKey]*Conn),
byChannel: make(map[int64]map[sessionKey]int64),
bySessionChannels: make(map[sessionKey]map[int64]struct{}),
byMemberChannel: make(map[int64]map[sessionKey]int64),
bySessionMembers: make(map[sessionKey]map[int64]struct{}),
pending: make(map[sessionKey][]queuedPush),
flushing: make(map[sessionKey]bool),
pendingBudget: newOutboundTrackedBudget(defaultPendingPushMaxBytes),
pushSessions: make(map[[8]byte]map[int64]struct{}),
log: log,
bySession: make(map[sessionKey]*Conn),
logicalSessions: make(map[sessionKey]*logicalSession),
claims: make(map[sessionKey]*Conn),
claimsByAuth: make(map[[8]byte]map[int64]*Conn),
byAuthKey: make(map[[8]byte]map[int64]*Conn),
byBusinessAuthKey: make(map[[8]byte]map[sessionKey]*Conn),
byUser: make(map[int64]map[sessionKey]*Conn),
byChannel: make(map[int64]map[sessionKey]int64),
bySessionChannels: make(map[sessionKey]map[int64]struct{}),
bySubscribedChannel: make(map[int64]map[sessionKey]channelSubscription),
bySessionSubscriptions: make(map[sessionKey]map[int64]int64),
byMemberChannel: make(map[int64]map[sessionKey]int64),
bySessionMembers: make(map[sessionKey]map[int64]struct{}),
pending: make(map[sessionKey][]queuedPush),
flushing: make(map[sessionKey]bool),
pendingBudget: newOutboundTrackedBudget(defaultPendingPushMaxBytes),
pushSessions: make(map[[8]byte]map[int64]struct{}),
log: log,
}
}
@ -260,6 +280,12 @@ func (m *SessionManager) SetLifecycleObserver(observer SessionLifecycleObserver)
m.mu.Unlock()
}
func (m *SessionManager) setLogicalSessionReleaseHook(hook func(sessionKey)) {
m.mu.Lock()
m.logicalSessionReleased = hook
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
@ -480,7 +506,7 @@ func (m *SessionManager) ApplyOrderedRawLayerForSession(
// 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.
// key's inherited default merely because the execution-receipt 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
@ -669,6 +695,7 @@ func (m *SessionManager) AbortActivation(c *Conn) {
m.deletePendingLocked(key)
delete(m.flushing, key)
}
m.markLogicalSessionOfflineLocked(key, time.Now())
owned = true
}
m.mu.Unlock()
@ -712,6 +739,7 @@ func (m *SessionManager) Unregister(c *Conn) {
zap.Int("online", len(m.bySession)),
)
}
m.markLogicalSessionOfflineLocked(key, time.Now())
m.mu.Unlock()
if observer != nil {
observer.SessionOffline(c.authKeyID, c.sessionID, offlineUser, lastForUser)
@ -727,6 +755,7 @@ func (m *SessionManager) DestroySessionForAuthKey(authKeyID [8]byte, sessionID i
if !ok {
if claim := m.claims[key]; claim != nil {
m.retireClaimLocked(key, claim, true)
outbound := m.destroyLogicalSessionLocked(key)
m.mu.Unlock()
if !forceCloseConnBatch([]*Conn{claim}, forceCloseBatchTimeout) {
m.log.Warn("Claimed session close exceeded shared deadline",
@ -734,15 +763,23 @@ func (m *SessionManager) DestroySessionForAuthKey(authKeyID [8]byte, sessionID i
zap.Int64("session_id", sessionID),
)
}
if outbound != nil {
m.releaseLogicalSession(key, outbound)
}
notifySessionDestroyed(observer, authKeyID, sessionID)
return true
}
m.deletePendingLocked(key)
outbound := m.destroyLogicalSessionLocked(key)
m.mu.Unlock()
if outbound != nil {
m.releaseLogicalSession(key, outbound)
}
notifySessionDestroyed(observer, authKeyID, sessionID)
return false
}
offlineUser := m.retireConnLocked(c, true)
outbound := m.destroyLogicalSessionLocked(key)
lastForUser := offlineUser != 0 && len(m.byUser[offlineUser]) == 0
m.log.Debug("Session destroyed",
zap.String("auth_key_id", sessionKeyLog(authKeyID)),
@ -756,6 +793,9 @@ func (m *SessionManager) DestroySessionForAuthKey(authKeyID [8]byte, sessionID i
zap.Int64("session_id", sessionID),
)
}
if outbound != nil {
m.releaseLogicalSession(key, outbound)
}
if observer != nil && offlineUser != 0 {
observer.SessionOffline(authKeyID, sessionID, offlineUser, lastForUser)
}
@ -779,8 +819,7 @@ func (m *SessionManager) bindUserLocked(c *Conn, key sessionKey, userID int64) {
if old := c.userID.Swap(userID); old != 0 {
removeUserIndex(m.byUser, old, key)
if old != userID {
m.clearChannelInterestsLocked(key)
m.clearChannelMembershipsLocked(c, key)
m.clearSessionChannelIndexesLocked(c, key)
c.membershipsSynced.Store(false)
// 身份变化即丢弃暂存推送：它们属于前一个账号，flush 给新账号是跨账号泄露。
// 同时取消进行中的排空（runFlush 还另有 owner 校验做批内兜底）。
@ -792,8 +831,7 @@ func (m *SessionManager) bindUserLocked(c *Conn, key sessionKey, userID int64) {
if userID != 0 {
addUserIndex(m.byUser, userID, key, c)
} else {
m.clearChannelInterestsLocked(key)
m.clearChannelMembershipsLocked(c, key)
m.clearSessionChannelIndexesLocked(c, key)
c.membershipsSynced.Store(false)
m.deletePendingLocked(key)
delete(m.flushing, key)
@ -852,13 +890,13 @@ func (m *SessionManager) bindAuthKeyLocked(c *Conn, key sessionKey, authKeyID [8
removeBusinessAuthKeyIndex(m.byBusinessAuthKey, oldAuthKeyID, key)
}
c.SetBusinessAuthKeyID(authKeyID)
m.bindLogicalSessionAuthKeyLocked(key, authKeyID)
addBusinessAuthKeyIndex(m.byBusinessAuthKey, authKeyID, key, c)
if changed {
if oldUserID != 0 {
removeUserIndex(m.byUser, oldUserID, key)
}
m.clearChannelInterestsLocked(key)
m.clearChannelMembershipsLocked(c, key)
m.clearSessionChannelIndexesLocked(c, key)
c.membershipsSynced.Store(false)
m.deletePendingLocked(key)
delete(m.flushing, key)
@ -903,6 +941,7 @@ func (m *SessionManager) CloseSessionsForBusinessAuthKey(authKeyID [8]byte) int
m.mu.Lock()
var conns []*Conn
var events []offlineEvent
var logicalRelease []*logicalSession
for key, c := range m.businessAuthKeyCandidatesLocked(authKeyID) {
if !connUsesBusinessAuthKey(c, authKeyID) {
continue
@ -918,6 +957,14 @@ func (m *SessionManager) CloseSessionsForBusinessAuthKey(authKeyID [8]byte) int
m.retireClaimLocked(key, c, true)
conns = append(conns, c)
}
for key, logical := range m.logicalSessions {
if logical == nil || (key.authKeyID != authKeyID &&
(!logical.businessAuthResolved || logical.businessAuthKeyID != authKeyID)) {
continue
}
delete(m.logicalSessions, key)
logicalRelease = append(logicalRelease, logical)
}
observer := m.lifecycle
if len(conns) > 0 {
m.log.Debug("Force close sessions for revoked auth key",
@ -932,6 +979,9 @@ func (m *SessionManager) CloseSessionsForBusinessAuthKey(authKeyID [8]byte) int
zap.Int("sessions", len(conns)),
)
}
for _, logical := range logicalRelease {
m.releaseLogicalSession(logical.key, logical.outbound)
}
if observer != nil {
for _, e := range events {
observer.SessionOffline(e.key.authKeyID, e.key.sessionID, e.userID, e.last)
@ -1114,8 +1164,7 @@ func (m *SessionManager) UnbindAuthKey(authKeyID [8]byte) int {
if old := c.userID.Swap(0); old != 0 {
removeUserIndex(m.byUser, old, key)
}
m.clearChannelInterestsLocked(key)
m.clearChannelMembershipsLocked(c, key)
m.clearSessionChannelIndexesLocked(c, key)
c.membershipsSynced.Store(false)
// 授权解除后暂存推送属于已登出的账号，不能等下一个登录者置位时 flush 出去。
m.deletePendingLocked(key)
@ -1134,8 +1183,7 @@ func (m *SessionManager) UnbindAuthKey(authKeyID [8]byte) int {
func (m *SessionManager) setReceivesUpdatesLocked(c *Conn, key sessionKey, receives bool) (int64, bool) {
if !receives {
c.receivesUpdates.Store(false)
m.clearChannelInterestsLocked(key)
m.clearChannelMembershipsLocked(c, key)
m.clearSessionChannelIndexesLocked(c, key)
c.membershipsSynced.Store(false)
// 取消进行中的排空激活：runFlush 在置位前会复查该标志，标志已删则放弃置位，
// 避免把刚置 false 的开关翻回 true。
@ -1148,8 +1196,7 @@ func (m *SessionManager) setReceivesUpdatesLocked(c *Conn, key sessionKey, recei
// 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)
m.clearSessionChannelIndexesLocked(c, key)
c.membershipsSynced.Store(false)
delete(m.flushing, key)
return 0, false
@ -1952,6 +1999,102 @@ func (m *SessionManager) OnlineChannelUserIDs(channelID int64, limit int) []int6
return m.onlineChannelUsers(m.byChannel, channelID, limit)
}
// RefreshChannelSubscription refreshes one public-channel passive-update
// subscription without replacing the other short-polled channels of the same
// session. The index is runtime-only and bounded to the official client limit.
func (m *SessionManager) RefreshChannelSubscription(rawAuthKeyID [8]byte, sessionID, userID, channelID int64, ttl time.Duration) {
if userID == 0 || channelID == 0 {
return
}
if ttl <= 0 {
ttl = defaultChannelSubscriptionTTL
} else if ttl > maxChannelSubscriptionTTL {
ttl = maxChannelSubscriptionTTL
}
key := sessionKey{authKeyID: rawAuthKeyID, sessionID: sessionID}
now := time.Now().UnixNano()
expiresAt := now + int64(ttl)
m.mu.Lock()
defer m.mu.Unlock()
c, ok := m.bySession[key]
if !ok || c.userID.Load() != userID {
return
}
m.pruneSessionSubscriptionsLocked(key, now)
channels := m.bySessionSubscriptions[key]
if channels == nil {
channels = make(map[int64]int64, 1)
m.bySessionSubscriptions[key] = channels
}
if _, exists := channels[channelID]; !exists && len(channels) >= maxChannelSubscriptionsPerSession {
m.log.Warn("Channel passive subscription ignored at per-session cap",
zap.String("auth_key_id", sessionKeyLog(rawAuthKeyID)),
zap.Int64("session_id", sessionID),
zap.Int64("channel_id", channelID),
zap.Int("cap", maxChannelSubscriptionsPerSession))
return
}
channels[channelID] = expiresAt
sessions := m.bySubscribedChannel[channelID]
if sessions == nil {
sessions = make(map[sessionKey]channelSubscription)
m.bySubscribedChannel[channelID] = sessions
}
sessions[key] = channelSubscription{userID: userID, expiresAt: expiresAt}
}
// OnlineChannelSubscriberUserIDs returns users for which at least one live
// session still holds an unexpired short-poll subscription. The user is
// deduplicated because passive updates are subsequently pushed account-wide.
func (m *SessionManager) OnlineChannelSubscriberUserIDs(channelID int64, limit int) []int64 {
return m.onlineChannelSubscriberUserIDsExcluding(channelID, nil, limit)
}
func (m *SessionManager) OnlineChannelSubscriberUserIDsExcluding(channelID int64, exclude map[int64]struct{}, limit int) []int64 {
return m.onlineChannelSubscriberUserIDsExcluding(channelID, exclude, limit)
}
func (m *SessionManager) onlineChannelSubscriberUserIDsExcluding(channelID int64, exclude map[int64]struct{}, limit int) []int64 {
if channelID == 0 {
return nil
}
now := time.Now().UnixNano()
m.mu.Lock()
defer m.mu.Unlock()
sessions := m.bySubscribedChannel[channelID]
if len(sessions) == 0 {
return nil
}
out := make([]int64, 0, positiveLimitOrLen(limit, len(sessions)))
seen := make(map[int64]struct{}, positiveLimitOrLen(limit, len(sessions)))
for key, subscription := range sessions {
if subscription.expiresAt <= now {
m.removeChannelSubscriptionLocked(key, channelID)
continue
}
if subscription.userID == 0 {
continue
}
if _, ok := exclude[subscription.userID]; ok {
continue
}
c, ok := m.bySession[key]
if !ok || c.userID.Load() != subscription.userID {
m.removeChannelSubscriptionLocked(key, channelID)
continue
}
if _, ok := seen[subscription.userID]; ok {
continue
}
seen[subscription.userID] = struct{}{}
out = append(out, subscription.userID)
if limit > 0 && len(out) >= limit {
break
}
}
return out
}
// ChannelMembershipGeneration 返回该 session 的 membership 索引修订号。
// 全量同步方必须在读取持久成员列表【之前】采样，并经 SetSessionChannelMemberships
// 带回比对；session 不在线时返回 0（后续 Set 也会因查不到连接而放弃）。
@ -2083,14 +2226,17 @@ func (m *SessionManager) OnlineChannelMemberUserIDsExcluding(channelID int64, ex
return out
}
// OnlineChannelIDsSnapshot returns every channel with at least one live joined-member session in
// strictly ascending order. The global SessionManager lock is held only while copying map keys;
// OnlineChannelIDsSnapshot returns every channel with at least one live joined-member session or
// unexpired passive subscriber in strictly ascending order. The global SessionManager lock is held
// only while copying map keys and pruning expired subscription entries;
// sorting and all recovery database work happen after unlock. The fixed saturation-recovery actor
// is the sole caller, so its exceptional-path temporary memory is one int64 slice (peak about 8*C
// bytes) rather than repeated O(C) scans under the connection/membership lock.
func (m *SessionManager) OnlineChannelIDsSnapshot() []int64 {
m.mu.RLock()
out := make([]int64, 0, len(m.byMemberChannel))
m.mu.Lock()
now := time.Now().UnixNano()
seen := make(map[int64]struct{}, len(m.byMemberChannel)+len(m.bySubscribedChannel))
out := make([]int64, 0, len(m.byMemberChannel)+len(m.bySubscribedChannel))
for channelID, sessions := range m.byMemberChannel {
if channelID <= 0 || len(sessions) == 0 {
continue
@ -2105,9 +2251,35 @@ func (m *SessionManager) OnlineChannelIDsSnapshot() []int64 {
if !live {
continue
}
seen[channelID] = struct{}{}
out = append(out, channelID)
}
m.mu.RUnlock()
for channelID, sessions := range m.bySubscribedChannel {
if channelID <= 0 || len(sessions) == 0 {
continue
}
live := false
for key, subscription := range sessions {
if subscription.expiresAt <= now {
m.removeChannelSubscriptionLocked(key, channelID)
continue
}
c, ok := m.bySession[key]
if !ok || c.userID.Load() != subscription.userID {
m.removeChannelSubscriptionLocked(key, channelID)
continue
}
live = true
}
if !live {
continue
}
if _, exists := seen[channelID]; exists {
continue
}
out = append(out, channelID)
}
m.mu.Unlock()
sort.Slice(out, func(i, j int) bool { return out[i] < out[j] })
return out
}
@ -2157,8 +2329,7 @@ func (m *SessionManager) removeLocked(c *Conn, dropPending bool) int64 {
if uid != 0 {
removeUserIndex(m.byUser, uid, key)
}
m.clearChannelInterestsLocked(key)
m.clearChannelMembershipsLocked(c, key)
m.clearSessionChannelIndexesLocked(c, key)
if dropPending {
m.deletePendingLocked(key)
}
@ -2255,6 +2426,50 @@ func (m *SessionManager) clearChannelInterestsLocked(key sessionKey) {
m.clearChannelIndexLocked(m.byChannel, m.bySessionChannels, key)
}
func (m *SessionManager) clearChannelSubscriptionsLocked(key sessionKey) {
channels := m.bySessionSubscriptions[key]
if len(channels) == 0 {
delete(m.bySessionSubscriptions, key)
return
}
for channelID := range channels {
sessions := m.bySubscribedChannel[channelID]
delete(sessions, key)
if len(sessions) == 0 {
delete(m.bySubscribedChannel, channelID)
}
}
delete(m.bySessionSubscriptions, key)
}
func (m *SessionManager) pruneSessionSubscriptionsLocked(key sessionKey, now int64) {
channels := m.bySessionSubscriptions[key]
for channelID, expiresAt := range channels {
if expiresAt <= now {
m.removeChannelSubscriptionLocked(key, channelID)
}
}
}
func (m *SessionManager) removeChannelSubscriptionLocked(key sessionKey, channelID int64) {
channels := m.bySessionSubscriptions[key]
delete(channels, channelID)
if len(channels) == 0 {
delete(m.bySessionSubscriptions, key)
}
sessions := m.bySubscribedChannel[channelID]
delete(sessions, key)
if len(sessions) == 0 {
delete(m.bySubscribedChannel, channelID)
}
}
func (m *SessionManager) clearSessionChannelIndexesLocked(c *Conn, key sessionKey) {
m.clearChannelInterestsLocked(key)
m.clearChannelSubscriptionsLocked(key)
m.clearChannelMembershipsLocked(c, key)
}
// clearChannelMembershipsLocked 整体清除某连接的 membership 索引并递增其修订号，
// 使在飞的全量同步（SetSessionChannelMemberships）能检测到清除并放弃过期替换。
func (m *SessionManager) clearChannelMembershipsLocked(c *Conn, key sessionKey) {
@ -2443,6 +2658,7 @@ func (m *SessionManager) RunPendingSweeper(ctx context.Context, interval time.Du
case <-ticker.C:
}
m.sweepStalePending()
m.sweepLogicalSessions(time.Now())
}
}

56
internal/mtprotoedge/session_manager_test.go
View file

@ -858,6 +858,51 @@ func TestSessionManagerChannelInterestIndex(t *testing.T) {
}
}
func TestSessionManagerChannelSubscriptionsAreBoundedDeduplicatedAndExpire(t *testing.T) {
sm := NewSessionManager(zaptest.NewLogger(t))
rawOne := [8]byte{1, 2, 3}
rawTwo := [8]byte{4, 5, 6}
first := &Conn{sessionID: 41, authKeyID: rawOne}
second := &Conn{sessionID: 42, authKeyID: rawTwo}
sm.Register(first)
sm.Register(second)
sm.BindUserForAuthKey(rawOne, 41, 100)
sm.BindUserForAuthKey(rawTwo, 42, 100)
sm.RefreshChannelSubscription(rawOne, 41, 100, 10, time.Second)
sm.RefreshChannelSubscription(rawTwo, 42, 100, 10, time.Second)
if got := sm.OnlineChannelSubscriberUserIDs(10, 10); len(got) != 1 || got[0] != 100 {
t.Fatalf("deduplicated subscribers = %v, want [100]", got)
}
if got := sm.OnlineChannelIDsSnapshot(); len(got) != 1 || got[0] != 10 {
t.Fatalf("subscribed channel snapshot = %v, want [10]", got)
}
for channelID := int64(20); channelID < 20+maxChannelSubscriptionsPerSession; channelID++ {
sm.RefreshChannelSubscription(rawOne, 41, 100, channelID, time.Second)
}
// Channel 10 already consumes one of the ten slots, so channel 29 must be
// refused rather than growing the session-controlled index to eleven.
if got := sm.OnlineChannelSubscriberUserIDs(29, 10); len(got) != 0 {
t.Fatalf("subscriber beyond per-session cap = %v, want empty", got)
}
if got := sm.OnlineChannelSubscriberUserIDsExcluding(10, map[int64]struct{}{100: {}}, 10); len(got) != 0 {
t.Fatalf("excluded subscribers = %v, want empty", got)
}
sm.RefreshChannelSubscription(rawTwo, 42, 100, 99, 10*time.Millisecond)
time.Sleep(30 * time.Millisecond)
if got := sm.OnlineChannelSubscriberUserIDs(99, 10); len(got) != 0 {
t.Fatalf("expired subscribers = %v, want empty", got)
}
sm.Unregister(first)
sm.Unregister(second)
if got := sm.OnlineChannelSubscriberUserIDs(10, 10); len(got) != 0 {
t.Fatalf("subscribers after unregister = %v, want empty", got)
}
}
func TestSessionManagerClearsChannelIndexesOnAuthAndReadinessChanges(t *testing.T) {
sm := NewSessionManager(zaptest.NewLogger(t))
raw := [8]byte{1, 2, 3}
@ -869,6 +914,7 @@ func TestSessionManagerClearsChannelIndexesOnAuthAndReadinessChanges(t *testing.
track := func() {
sm.TrackChannelInterest(raw, 42, 100, []int64{10})
sm.RefreshChannelSubscription(raw, 42, 100, 10, time.Second)
sm.SetSessionChannelMemberships(raw, 42, 100, []int64{10}, sm.ChannelMembershipGeneration(raw, 42))
if got := sm.OnlineChannelUserIDs(10, 10); len(got) != 1 || got[0] != 100 {
t.Fatalf("channel viewers before cleanup = %v, want [100]", got)
@ -876,6 +922,9 @@ func TestSessionManagerClearsChannelIndexesOnAuthAndReadinessChanges(t *testing.
if got := sm.OnlineChannelMemberUserIDs(10, 10); len(got) != 1 || got[0] != 100 {
t.Fatalf("channel members before cleanup = %v, want [100]", got)
}
if got := sm.OnlineChannelSubscriberUserIDs(10, 10); len(got) != 1 || got[0] != 100 {
t.Fatalf("channel subscribers before cleanup = %v, want [100]", got)
}
}
assertCleared := func(label string) {
if got := sm.OnlineChannelUserIDs(10, 10); len(got) != 0 {
@ -884,6 +933,9 @@ func TestSessionManagerClearsChannelIndexesOnAuthAndReadinessChanges(t *testing.
if got := sm.OnlineChannelMemberUserIDs(10, 10); len(got) != 0 {
t.Fatalf("%s members = %v, want empty", label, got)
}
if got := sm.OnlineChannelSubscriberUserIDs(10, 10); len(got) != 0 {
t.Fatalf("%s subscribers = %v, want empty", label, got)
}
}
track()
@ -1153,9 +1205,9 @@ func TestSessionManagerPush(t *testing.T) {
// 各发一个 ping 建立 session，触发注册（并清掉 new_session_created/pong/ack）。
msgGen := proto.NewMessageIDGen(time.Now)
sendEncrypted(t, conn1, cipher1, auth1, msgGen.New(proto.MessageFromClient), &mt.PingRequest{PingID: 1})
sendEncryptedWithSeq(t, conn1, cipher1, auth1, msgGen.New(proto.MessageFromClient), 1, &mt.PingRequest{PingID: 1})
collectReplies(t, conn1, cipher1, auth1.AuthKey, mt.PongTypeID)
sendEncrypted(t, conn2, cipher2, auth2, msgGen.New(proto.MessageFromClient), &mt.PingRequest{PingID: 2})
sendEncryptedWithSeq(t, conn2, cipher2, auth2, msgGen.New(proto.MessageFromClient), 1, &mt.PingRequest{PingID: 2})
collectReplies(t, conn2, cipher2, auth2.AuthKey, mt.PongTypeID)
if got := srv.Conns().Online(); got != 2 {

141
internal/mtprotoedge/structural_limits_test.go
View file

@ -1,6 +1,9 @@
package mtprotoedge
import (
"bytes"
"compress/flate"
"compress/zlib"
"context"
"encoding/binary"
"errors"
@ -664,4 +667,142 @@ func TestGZIPExpansionUsesProcessBudgetBeforeDecode(t *testing.T) {
if got := s.frameBudget.usedBytes(); got != 0 {
t.Fatalf("released expansion budget = %d, want zero", got)
}
s.frameBudget = newInboundFrameBudget(2 * maxSingleGZIPExpandedBytes)
if _, release, err := s.decodeGZIPWithGlobalBudgetLimit(&wrapped, len(payload)-1); err == nil {
release()
t.Fatal("caller-bounded gzip decode accepted an oversized expansion")
} else if got := gzipExpansionWork(err); got != len(payload) {
t.Fatalf("caller-bounded rejected expansion work = %d, want %d", got, len(payload))
}
if got := s.frameBudget.usedBytes(); got != 0 {
t.Fatalf("caller-bounded rejection leaked %d bytes", got)
}
s.frameBudget = newInboundFrameBudget(int64(len(payload) - 1))
if _, release, err := s.decodeGZIPWithGlobalBudgetLimit(&wrapped, len(payload)); !errors.Is(err, ErrInboundFrameBudgetExceeded) {
release()
t.Fatalf("caller-bounded process budget error = %v, want ErrInboundFrameBudgetExceeded", err)
}
if got := s.frameBudget.usedBytes(); got != 0 {
t.Fatalf("caller-bounded reservation failure leaked %d bytes", got)
}
}
func TestGZIPPackedAcceptsOfficialClientWrappers(t *testing.T) {
payload := []byte("official Telegram gzip_packed payload")
tests := []struct {
name string
packed bin.Encoder
}{
{name: "tdlib_zlib", packed: zlibPackedObjectForTest(t, payload)},
{name: "drklo_gotd_gzip", packed: &proto.GZIP{Data: payload}},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
var wrapped bin.Buffer
if err := tt.packed.Encode(&wrapped); err != nil {
t.Fatalf("encode gzip_packed: %v", err)
}
s := New(Options{Logger: zaptest.NewLogger(t)})
s.frameBudget = newInboundFrameBudget(maxSingleGZIPExpandedBytes)
decoded, release, err := s.decodeGZIPWithGlobalBudget(&wrapped)
if err != nil {
t.Fatalf("decode gzip_packed: %v", err)
}
if !bytes.Equal(decoded, payload) {
t.Fatalf("decoded payload = %q, want %q", decoded, payload)
}
if got := s.frameBudget.usedBytes(); got != int64(len(payload)) {
t.Fatalf("held expansion budget = %d, want %d", got, len(payload))
}
release()
if got := s.frameBudget.usedBytes(); got != 0 {
t.Fatalf("released expansion budget = %d, want zero", got)
}
})
}
}
func TestGZIPPackedRejectsCorruptZLIBChecksumWithoutBudgetLeak(t *testing.T) {
packed := zlibPackedObjectForTest(t, []byte("checksum-protected payload"))
packed.PackedData[len(packed.PackedData)-1] ^= 0xff
var wrapped bin.Buffer
if err := packed.Encode(&wrapped); err != nil {
t.Fatalf("encode gzip_packed: %v", err)
}
s := New(Options{Logger: zaptest.NewLogger(t)})
s.frameBudget = newInboundFrameBudget(maxSingleGZIPExpandedBytes)
if _, release, err := s.decodeGZIPWithGlobalBudget(&wrapped); err == nil {
release()
t.Fatal("corrupt zlib checksum unexpectedly accepted")
}
if got := s.frameBudget.usedBytes(); got != 0 {
t.Fatalf("corrupt zlib checksum leaked %d budget bytes", got)
}
}
func TestGZIPPackedZLIBHonorsExpansionLimit(t *testing.T) {
payload := make([]byte, 1<<20)
packed := zlibPackedObjectForTest(t, payload)
var wrapped bin.Buffer
if err := packed.Encode(&wrapped); err != nil {
t.Fatalf("encode gzip_packed: %v", err)
}
s := New(Options{Logger: zaptest.NewLogger(t)})
s.frameBudget = newInboundFrameBudget(2 * maxSingleGZIPExpandedBytes)
if _, release, err := s.decodeGZIPWithGlobalBudgetLimit(&wrapped, len(payload)-1); err == nil {
release()
t.Fatal("caller-bounded zlib decode accepted an oversized expansion")
} else if got := gzipExpansionWork(err); got != len(payload) {
t.Fatalf("caller-bounded zlib expansion work = %d, want %d", got, len(payload))
}
if got := s.frameBudget.usedBytes(); got != 0 {
t.Fatalf("caller-bounded zlib rejection leaked %d bytes", got)
}
}
func TestGZIPPackedRejectsRawDEFLATEWithoutBudgetLeak(t *testing.T) {
var compressed bytes.Buffer
w, err := flate.NewWriter(&compressed, flate.DefaultCompression)
if err != nil {
t.Fatalf("create raw deflate writer: %v", err)
}
if _, err := w.Write([]byte("raw deflate is not a supported Telegram wrapper")); err != nil {
t.Fatalf("write raw deflate payload: %v", err)
}
if err := w.Close(); err != nil {
t.Fatalf("close raw deflate payload: %v", err)
}
packed := &mt.GzipPacked{PackedData: compressed.Bytes()}
var wrapped bin.Buffer
if err := packed.Encode(&wrapped); err != nil {
t.Fatalf("encode gzip_packed: %v", err)
}
s := New(Options{Logger: zaptest.NewLogger(t)})
s.frameBudget = newInboundFrameBudget(maxSingleGZIPExpandedBytes)
if _, release, err := s.decodeGZIPWithGlobalBudget(&wrapped); err == nil {
release()
t.Fatal("raw deflate unexpectedly accepted")
}
if got := s.frameBudget.usedBytes(); got != 0 {
t.Fatalf("raw deflate rejection leaked %d budget bytes", got)
}
}
func zlibPackedObjectForTest(t testing.TB, payload []byte) *mt.GzipPacked {
t.Helper()
var compressed bytes.Buffer
w := zlib.NewWriter(&compressed)
if _, err := w.Write(payload); err != nil {
t.Fatalf("write zlib payload: %v", err)
}
if err := w.Close(); err != nil {
t.Fatalf("close zlib payload: %v", err)
}
return &mt.GzipPacked{PackedData: compressed.Bytes()}
}

133
internal/mtprotoedge/transport_compat.go
View file

@ -89,8 +89,7 @@ func newCompatTransportListener(codec func() transport.Codec, listener net.Liste
}
// singleConnListener 是一个只产出一条「已接受」连接、随后阻塞到关闭的 net.Listener。
// 它让单条裸连接可以走 listener 形态的去混淆/codec 管线（ObfuscatedListener +
// compatTransportListener），从而把这部分阻塞读取从 accept 循环挪到每连接 goroutine。
// 生产接入已直接提升单连接；这个适配器仅保留给仍需 listener 形态的测试/调用方。
type singleConnListener struct {
addr net.Addr
ch chan net.Conn
@ -131,38 +130,13 @@ func (l *compatTransportListener) Accept() (_ transport.Conn, rErr error) {
}
}()
var (
connCodec transport.Codec
reader io.Reader = conn
)
if l.codec != nil {
connCodec = l.codec()
if classifyInboundFrameCodec(connCodec) == inboundFrameCodecUnknown {
// Unknown codecs are rejected before their header or first frame is read. Without an
// explicit preflight contract, calling Codec.Read could allocate from an attacker-
// controlled length before the process-wide budget can be reserved.
return nil, errInboundFrameCodecUnsupported
}
if err := connCodec.ReadHeader(conn); err != nil {
return nil, errors.Wrap(err, "read codec header")
}
} else {
var err error
connCodec, reader, err = detectCompatCodec(conn)
if err != nil {
return nil, errors.Wrap(err, "detect codec")
}
promoted, err := newCompatTransportConn(l.codec, conn, l.budget)
if err != nil {
// Avoid returning a typed nil *compatTransportConn as a non-nil
// transport.Conn interface on admission failure.
return nil, err
}
return &compatTransportConn{
conn: wrappedCompatConn{
reader: reader,
Conn: conn,
},
codec: connCodec,
budget: l.budget,
transportPacketMessages: isTransportPacketMessageConn(conn),
}, nil
return promoted, nil
}
func isTransportPacketMessageConn(conn net.Conn) bool {
@ -187,10 +161,87 @@ func (w wrappedCompatConn) Read(p []byte) (int, error) {
return w.reader.Read(p)
}
// newCompatTransportConn promotes an already accepted stream without allocating the
// single-connection listener/channel used by the historical listener composition.
// Detection happens once per physical connection; the selected codec remains bound to
// the returned transport for both reads and writes.
func newCompatTransportConn(
codecFactory func() transport.Codec,
conn net.Conn,
budget *inboundFrameBudget,
) (*compatTransportConn, error) {
if budget == nil {
panic("mtprotoedge: nil inbound frame budget")
}
var (
connCodec transport.Codec
reader io.Reader = conn
)
if codecFactory != nil {
connCodec = codecFactory()
if classifyInboundFrameCodec(connCodec) == inboundFrameCodecUnknown {
// Unknown codecs are rejected before their header or first frame is read. Without an
// explicit preflight contract, calling Codec.Read could allocate from an attacker-
// controlled length before the process-wide budget can be reserved.
return nil, errInboundFrameCodecUnsupported
}
if err := connCodec.ReadHeader(conn); err != nil {
return nil, errors.Wrap(err, "read codec header")
}
} else {
var err error
connCodec, reader, err = detectCompatCodec(conn)
if err != nil {
return nil, errors.Wrap(err, "detect codec")
}
}
return newCompatTransportConnWithCodec(
wrappedCompatConn{reader: reader, Conn: conn},
connCodec,
budget,
isTransportPacketMessageConn(conn),
)
}
// newCompatTransportConnWithCodec binds a codec whose client-side transport header
// was already consumed by the one-time wire detector.
func newCompatTransportConnWithCodec(
conn net.Conn,
connCodec transport.Codec,
budget *inboundFrameBudget,
transportPacketMessages bool,
) (*compatTransportConn, error) {
if budget == nil {
panic("mtprotoedge: nil inbound frame budget")
}
codecKind := classifyInboundFrameCodec(connCodec)
if codecKind == inboundFrameCodecUnknown {
return nil, errInboundFrameCodecUnsupported
}
var budgetedCodec InboundFrameBudgetedCodec
if codecKind == inboundFrameCodecCustom {
budgetedCodec = unwrapInboundFrameBudgetedCodec(connCodec)
if budgetedCodec == nil {
return nil, errInboundFrameCodecUnsupported
}
}
return &compatTransportConn{
conn: conn,
codec: connCodec,
codecKind: codecKind,
budgetedCodec: budgetedCodec,
budget: budget,
transportPacketMessages: transportPacketMessages,
}, nil
}
type compatTransportConn struct {
conn net.Conn
codec transport.Codec
budget *inboundFrameBudget
conn net.Conn
codec transport.Codec
codecKind inboundFrameCodecKind
budgetedCodec InboundFrameBudgetedCodec
budget *inboundFrameBudget
transportPacketMessages bool
directMessageScratch []byte
@ -427,7 +478,10 @@ func (c *compatTransportConn) Close() error {
}
func (c *compatTransportConn) readInboundFrame(b *bin.Buffer) error {
kind := classifyInboundFrameCodec(c.codec)
// codecKind and budgetedCodec are frozen when the physical connection is
// promoted. The per-frame hot path never re-detects wire mode or type-switches
// the already selected codec.
kind := c.codecKind
if kind == inboundFrameCodecUnknown {
return errInboundFrameCodecUnsupported
}
@ -447,11 +501,10 @@ func (c *compatTransportConn) readInboundFrame(b *bin.Buffer) error {
var err error
if kind == inboundFrameCodecCustom {
custom := unwrapInboundFrameBudgetedCodec(c.codec)
if custom == nil {
if c.budgetedCodec == nil {
return errInboundFrameCodecUnsupported
}
err = custom.ReadWithInboundFrameBudget(c.conn, b, reserve)
err = c.budgetedCodec.ReadWithInboundFrameBudget(c.conn, b, reserve)
} else {
preflight := &inboundFramePreflightReader{r: c.conn, kind: kind, reserve: reserve}
err = c.codec.Read(preflight, b)

234
internal/mtprotoedge/transport_detection.go Normal file
View file

@ -0,0 +1,234 @@
package mtprotoedge
import (
"encoding/binary"
"errors"
"fmt"
"io"
"net"
"github.com/iamxvbaba/td/mtproxy/obfuscated2"
"github.com/iamxvbaba/td/proto/codec"
"github.com/iamxvbaba/td/transport"
)
var errInvalidTCPTransportPrefix = errors.New("invalid MTProto TCP transport prefix")
var errInvalidObfuscatedProtocol = errors.New("invalid obfuscated MTProto protocol tag")
type tcpWireMode uint8
const (
tcpWireModeUnknown tcpWireMode = iota
tcpWireModePlain
tcpWireModeObfuscated
)
type detectedTCPCodec uint8
const (
detectedTCPCodecUnknown detectedTCPCodec = iota
detectedTCPCodecAbridged
detectedTCPCodecIntermediate
detectedTCPCodecPaddedIntermediate
detectedTCPCodecFull
)
type tcpTransportProbe struct {
raw net.Conn
prefix [8]byte
prefixLen int
mode tcpWireMode
codec detectedTCPCodec
}
// detectTCPTransport reads at most the first eight bytes once per physical TCP
// connection. Telegram's obfuscated2 nonce generation explicitly excludes all
// plaintext codec tags and requires bytes 4..8 to be non-zero, while the first
// Full frame has transport sequence number zero. Those disjoint invariants let
// one public port admit both wire modes without probabilistic guessing.
func detectTCPTransport(raw net.Conn) (tcpTransportProbe, error) {
probe := tcpTransportProbe{raw: raw}
if _, err := io.ReadFull(raw, probe.prefix[:1]); err != nil {
return probe, fmt.Errorf("read first transport byte: %w", err)
}
probe.prefixLen = 1
if probe.prefix[0] == codec.AbridgedClientStart[0] {
probe.mode = tcpWireModePlain
probe.codec = detectedTCPCodecAbridged
return probe, nil
}
if _, err := io.ReadFull(raw, probe.prefix[1:4]); err != nil {
return probe, fmt.Errorf("read transport prefix: %w", err)
}
probe.prefixLen = 4
var firstFour [4]byte
copy(firstFour[:], probe.prefix[:4])
switch firstFour {
case codec.IntermediateClientStart:
probe.mode = tcpWireModePlain
probe.codec = detectedTCPCodecIntermediate
return probe, nil
case codec.PaddedIntermediateClientStart:
probe.mode = tcpWireModePlain
probe.codec = detectedTCPCodecPaddedIntermediate
return probe, nil
}
first := binary.LittleEndian.Uint32(firstFour[:])
if isHTTPHeaderPrefix(firstFour) || first == 0x02010316 {
return probe, fmt.Errorf("%w: reserved prefix %x", errInvalidTCPTransportPrefix, firstFour)
}
if _, err := io.ReadFull(raw, probe.prefix[4:8]); err != nil {
return probe, fmt.Errorf("read transport discriminator: %w", err)
}
probe.prefixLen = 8
if binary.LittleEndian.Uint32(probe.prefix[4:8]) == 0 {
length := binary.LittleEndian.Uint32(probe.prefix[:4])
if length < 3*4 || length > maxTransportMessageSize || length%4 != 0 {
return probe, fmt.Errorf("%w: invalid full header length %d", errInvalidTCPTransportPrefix, length)
}
probe.mode = tcpWireModePlain
probe.codec = detectedTCPCodecFull
return probe, nil
}
probe.mode = tcpWireModeObfuscated
return probe, nil
}
func (p tcpTransportProbe) replayConn() net.Conn {
return newReplayNetConn(p.raw, p.prefix[:p.prefixLen])
}
func (p tcpTransportProbe) plainFrameConn() net.Conn {
if p.codec == detectedTCPCodecFull {
// Full has no standalone codec tag: the first eight bytes are already the
// length and sequence number of its first frame and must be replayed.
return p.replayConn()
}
// Abridged/intermediate/padded-intermediate tags are client-only connection
// headers. They were consumed by detection and are not part of the first frame.
return p.raw
}
type replayNetConn struct {
net.Conn
prefix [8]byte
n uint8
offset uint8
}
func newReplayNetConn(conn net.Conn, prefix []byte) *replayNetConn {
if len(prefix) > 8 {
panic("mtprotoedge: replay prefix exceeds fixed transport discriminator")
}
replayed := &replayNetConn{Conn: conn, n: uint8(len(prefix))}
copy(replayed.prefix[:], prefix)
return replayed
}
func (c *replayNetConn) Read(p []byte) (int, error) {
if c.offset < c.n {
n := copy(p, c.prefix[c.offset:c.n])
c.offset += uint8(n)
return n, nil
}
return c.Conn.Read(p)
}
type serverObfuscatedConn struct {
net.Conn
rw io.ReadWriter
}
func (c *serverObfuscatedConn) Read(p []byte) (int, error) {
return c.rw.Read(p)
}
func (c *serverObfuscatedConn) Write(p []byte) (int, error) {
return c.rw.Write(p)
}
func detectedCompatCodec(kind detectedTCPCodec) (transport.Codec, error) {
switch kind {
case detectedTCPCodecAbridged:
return &quickAckAbridgedCodec{}, nil
case detectedTCPCodecIntermediate:
return &quickAckIntermediateCodec{}, nil
case detectedTCPCodecPaddedIntermediate:
return &quickAckPaddedIntermediateCodec{}, nil
case detectedTCPCodecFull:
return transport.Full.Codec(), nil
default:
return nil, fmt.Errorf("%w: unknown detected codec %d", errInvalidTCPTransportPrefix, kind)
}
}
func detectObfuscatedProtocol(tag [4]byte) (detectedTCPCodec, int, error) {
switch tag {
case (codec.Abridged{}).ObfuscatedTag():
return detectedTCPCodecAbridged, 1, nil
case codec.IntermediateClientStart:
return detectedTCPCodecIntermediate, 4, nil
case codec.PaddedIntermediateClientStart:
return detectedTCPCodecPaddedIntermediate, 4, nil
default:
return detectedTCPCodecUnknown, 0, fmt.Errorf("%w: %x", errInvalidObfuscatedProtocol, tag)
}
}
func (s *Server) promoteMixedTCP(raw net.Conn) (transport.Conn, string, error) {
probe, err := detectTCPTransport(raw)
if err != nil {
return nil, "tcp_auto", err
}
if probe.mode == tcpWireModePlain {
if s.codec != nil {
conn, err := newCompatTransportConn(s.codec, probe.replayConn(), s.frameBudget)
return conn, "tcp", err
}
connCodec, err := detectedCompatCodec(probe.codec)
if err != nil {
return nil, "tcp", err
}
conn, err := newCompatTransportConnWithCodec(
probe.plainFrameConn(),
connCodec,
s.frameBudget,
false,
)
return conn, "tcp", err
}
replayed := probe.replayConn()
rw, metadata, err := obfuscated2.Accept(replayed, nil)
if err != nil {
return nil, "obfuscated_tcp", fmt.Errorf("accept obfuscated2: %w", err)
}
obfuscated := &serverObfuscatedConn{Conn: replayed, rw: rw}
detected, tagLen, err := detectObfuscatedProtocol(metadata.Protocol)
if err != nil {
return nil, "obfuscated_tcp", err
}
if s.codec != nil {
conn, err := newCompatTransportConn(
s.codec,
newReplayNetConn(obfuscated, metadata.Protocol[:tagLen]),
s.frameBudget,
)
return conn, "obfuscated_tcp", err
}
connCodec, err := detectedCompatCodec(detected)
if err != nil {
return nil, "obfuscated_tcp", err
}
conn, err := newCompatTransportConnWithCodec(
obfuscated,
connCodec,
s.frameBudget,
false,
)
return conn, "obfuscated_tcp", err
}

458
internal/mtprotoedge/transport_detection_test.go Normal file
View file

@ -0,0 +1,458 @@
package mtprotoedge
import (
"context"
"crypto/rand"
"errors"
"io"
"net"
"testing"
"time"
"go.uber.org/zap/zaptest"
"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/transport"
)
func TestDetectTCPTransportFragmentedPrefixes(t *testing.T) {
tests := []struct {
name string
prefix []byte
mode tcpWireMode
codec detectedTCPCodec
}{
{name: "abridged", prefix: []byte{0xef}, mode: tcpWireModePlain, codec: detectedTCPCodecAbridged},
{name: "intermediate", prefix: []byte{0xee, 0xee, 0xee, 0xee}, mode: tcpWireModePlain, codec: detectedTCPCodecIntermediate},
{name: "padded", prefix: []byte{0xdd, 0xdd, 0xdd, 0xdd}, mode: tcpWireModePlain, codec: detectedTCPCodecPaddedIntermediate},
{name: "full", prefix: []byte{12, 0, 0, 0, 0, 0, 0, 0}, mode: tcpWireModePlain, codec: detectedTCPCodecFull},
{name: "obfuscated", prefix: []byte{1, 2, 3, 4, 5, 6, 7, 8}, mode: tcpWireModeObfuscated, codec: detectedTCPCodecUnknown},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
server, client := net.Pipe()
defer func() { _ = server.Close() }()
defer func() { _ = client.Close() }()
writeErr := make(chan error, 1)
go func() {
for _, b := range tt.prefix {
if _, err := client.Write([]byte{b}); err != nil {
writeErr <- err
return
}
}
writeErr <- nil
}()
probe, err := detectTCPTransport(server)
if err != nil {
t.Fatalf("detect: %v", err)
}
if probe.mode != tt.mode || probe.codec != tt.codec || probe.prefixLen != len(tt.prefix) {
t.Fatalf("probe = mode:%d codec:%d prefix:%d, want %d/%d/%d",
probe.mode, probe.codec, probe.prefixLen, tt.mode, tt.codec, len(tt.prefix))
}
if err := <-writeErr; err != nil {
t.Fatalf("write prefix: %v", err)
}
})
}
}
func TestDetectTCPTransportRejectsReservedAndInvalidFullPrefixes(t *testing.T) {
tests := []struct {
name string
prefix []byte
}{
{name: "http", prefix: []byte("GET ")},
{name: "reserved", prefix: []byte{0x16, 0x03, 0x01, 0x02}},
{name: "full_too_short", prefix: []byte{8, 0, 0, 0, 0, 0, 0, 0}},
{name: "full_unaligned", prefix: []byte{13, 0, 0, 0, 0, 0, 0, 0}},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
server, client := net.Pipe()
defer func() { _ = server.Close() }()
defer func() { _ = client.Close() }()
go func() {
_, _ = client.Write(tt.prefix)
}()
_, err := detectTCPTransport(server)
if !errors.Is(err, errInvalidTCPTransportPrefix) {
t.Fatalf("detect error = %v, want invalid prefix", err)
}
})
}
}
func TestServerMixedTCPAcceptsEveryPlainCodec(t *testing.T) {
tests := []struct {
name string
protocol transport.Protocol
}{
{name: "abridged", protocol: transport.Abridged},
{name: "intermediate", protocol: transport.Intermediate},
{name: "padded_intermediate", protocol: transport.PaddedIntermediate},
{name: "full", protocol: transport.Full},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
addr, frames := startTransportFrameServer(t, Options{
Logger: zaptest.NewLogger(t),
ObfuscatedTCP: true,
})
raw, err := net.Dial("tcp", addr)
if err != nil {
t.Fatalf("dial: %v", err)
}
conn, err := tt.protocol.Handshake(raw)
if err != nil {
_ = raw.Close()
t.Fatalf("transport handshake: %v", err)
}
t.Cleanup(func() { _ = conn.Close() })
var payload bin.Buffer
payload.PutInt32(0x12345678)
payload.PutInt32(0x0badf00d)
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
if err := conn.Send(ctx, &payload); err != nil {
t.Fatalf("send: %v", err)
}
select {
case n := <-frames:
if n < payload.Len() || n > payload.Len()+15 {
t.Fatalf("frame len = %d, want payload %d plus at most padded-intermediate padding", n, payload.Len())
}
case <-ctx.Done():
t.Fatal("server did not receive plain frame in mixed mode")
}
})
}
}
func TestServerMixedTCPAcceptsObfuscatedCodecs(t *testing.T) {
tests := []struct {
name string
tag [4]byte
newCodec func() transport.Codec
}{
{
name: "abridged",
tag: (codec.Abridged{}).ObfuscatedTag(),
newCodec: func() transport.Codec {
return codec.NoHeader{Codec: codec.Abridged{}}
},
},
{
name: "intermediate",
tag: codec.IntermediateClientStart,
newCodec: func() transport.Codec {
return codec.NoHeader{Codec: codec.Intermediate{}}
},
},
{
name: "padded_intermediate",
tag: codec.PaddedIntermediateClientStart,
newCodec: func() transport.Codec {
return codec.NoHeader{Codec: codec.PaddedIntermediate{}}
},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
addr, frames := startTransportFrameServer(t, Options{
Logger: zaptest.NewLogger(t),
ObfuscatedTCP: true,
})
raw, err := net.Dial("tcp", addr)
if err != nil {
t.Fatalf("dial: %v", err)
}
obfuscated := obfuscator.Obfuscated2(rand.Reader, raw)
if err := obfuscated.Handshake(tt.tag, 2, mtproxy.Secret{}); err != nil {
_ = raw.Close()
t.Fatalf("obfuscated handshake: %v", err)
}
conn, err := transport.NewProtocol(tt.newCodec).Handshake(obfuscated)
if err != nil {
_ = raw.Close()
t.Fatalf("transport handshake: %v", err)
}
t.Cleanup(func() { _ = conn.Close() })
var payload bin.Buffer
payload.PutInt32(0x12345678)
payload.PutInt32(0x0badf00d)
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
if err := conn.Send(ctx, &payload); err != nil {
t.Fatalf("send: %v", err)
}
select {
case n := <-frames:
if n < payload.Len() || n > payload.Len()+15 {
t.Fatalf("frame len = %d, want payload %d plus at most padding", n, payload.Len())
}
case <-ctx.Done():
t.Fatal("server did not receive obfuscated frame")
}
})
}
}
func TestServerMixedTCPPlainIntermediateKeyExchange(t *testing.T) {
addr, pub, _ := startTestServer(t, Options{DC: 2, ObfuscatedTCP: true})
conn, _, _ := dialHandshake(t, addr, 2, pub)
_ = conn.Close()
}
func TestServerMixedTCPResponseUsesDetectedWireMode(t *testing.T) {
plainDial := func(protocol transport.Protocol) func(*testing.T, string) transport.Conn {
return func(t *testing.T, addr string) transport.Conn {
t.Helper()
raw, err := net.Dial("tcp", addr)
if err != nil {
t.Fatalf("dial: %v", err)
}
conn, err := protocol.Handshake(raw)
if err != nil {
_ = raw.Close()
t.Fatalf("transport handshake: %v", err)
}
return conn
}
}
obfuscatedDial := func(
tag [4]byte,
newCodec func() transport.Codec,
) func(*testing.T, string) transport.Conn {
return func(t *testing.T, addr string) transport.Conn {
t.Helper()
raw, err := net.Dial("tcp", addr)
if err != nil {
t.Fatalf("dial: %v", err)
}
obfuscated := obfuscator.Obfuscated2(rand.Reader, raw)
if err := obfuscated.Handshake(tag, 2, mtproxy.Secret{}); err != nil {
_ = raw.Close()
t.Fatalf("obfuscated handshake: %v", err)
}
conn, err := transport.NewProtocol(newCodec).Handshake(obfuscated)
if err != nil {
_ = raw.Close()
t.Fatalf("transport handshake: %v", err)
}
return conn
}
}
tests := []struct {
name string
dial func(t *testing.T, addr string) transport.Conn
}{
{
name: "plain_abridged",
dial: plainDial(transport.Abridged),
},
{
name: "plain_intermediate",
dial: plainDial(transport.Intermediate),
},
{
name: "plain_padded_intermediate",
dial: plainDial(transport.PaddedIntermediate),
},
{
name: "plain_full",
dial: plainDial(transport.Full),
},
{
name: "obfuscated_abridged",
dial: obfuscatedDial(
(codec.Abridged{}).ObfuscatedTag(),
func() transport.Codec {
return codec.NoHeader{Codec: codec.Abridged{}}
},
),
},
{
name: "obfuscated_intermediate",
dial: obfuscatedDial(
codec.IntermediateClientStart,
func() transport.Codec {
return codec.NoHeader{Codec: codec.Intermediate{}}
},
),
},
{
name: "obfuscated_padded_intermediate",
dial: obfuscatedDial(
codec.PaddedIntermediateClientStart,
func() transport.Codec {
return codec.NoHeader{Codec: codec.PaddedIntermediate{}}
},
),
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
addr, _, _ := startTestServer(t, Options{DC: 2, ObfuscatedTCP: true})
conn := tt.dial(t, addr)
t.Cleanup(func() { _ = conn.Close() })
var payload bin.Buffer
payload.PutLong(0x1020304050607080) // deliberately unknown auth_key_id
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
defer cancel()
if err := conn.Send(ctx, &payload); err != nil {
t.Fatalf("send: %v", err)
}
var response bin.Buffer
err := conn.Recv(ctx, &response)
var protocolErr *codec.ProtocolErr
if !errors.As(err, &protocolErr) || protocolErr.Code != codec.CodeAuthKeyNotFound {
t.Fatalf("response error = %T %v, want transport -404 in the detected wire mode", err, err)
}
})
}
}
func TestServerMixedTCPRejectsInvalidPrefixesWithoutObfuscationWait(t *testing.T) {
addr, _, _ := startTestServer(t, Options{
ObfuscatedTCP: true,
HandshakeIdleTimeout: 5 * time.Second,
})
raw, err := net.Dial("tcp", addr)
if err != nil {
t.Fatalf("dial: %v", err)
}
defer func() { _ = raw.Close() }()
if _, err := raw.Write([]byte{8, 0, 0, 0, 0, 0, 0, 0}); err != nil {
t.Fatalf("write invalid full prefix: %v", err)
}
if err := raw.SetReadDeadline(time.Now().Add(time.Second)); err != nil {
t.Fatalf("set read deadline: %v", err)
}
_, err = raw.Read(make([]byte, 1))
if err == nil {
t.Fatal("invalid transport prefix unexpectedly kept connection open")
}
var netErr net.Error
if errors.As(err, &netErr) && netErr.Timeout() {
t.Fatalf("invalid prefix waited for obfuscation timeout instead of failing fast: %v", err)
}
}
func TestServerMixedTCPRejectsUnknownObfuscatedProtocolTag(t *testing.T) {
addr, _, _ := startTestServer(t, Options{ObfuscatedTCP: true})
raw, err := net.Dial("tcp", addr)
if err != nil {
t.Fatalf("dial: %v", err)
}
defer func() { _ = raw.Close() }()
obfuscated := obfuscator.Obfuscated2(rand.Reader, raw)
if err := obfuscated.Handshake([4]byte{1, 2, 3, 4}, 2, mtproxy.Secret{}); err != nil {
t.Fatalf("obfuscated handshake: %v", err)
}
if err := raw.SetReadDeadline(time.Now().Add(time.Second)); err != nil {
t.Fatalf("set read deadline: %v", err)
}
_, err = obfuscated.Read(make([]byte, 1))
if err == nil {
t.Fatal("unknown obfuscated protocol tag unexpectedly kept connection open")
}
var netErr net.Error
if errors.As(err, &netErr) && netErr.Timeout() {
t.Fatalf("unknown obfuscated protocol tag did not fail closed: %v", err)
}
if !errors.Is(err, io.EOF) {
// Windows may report a reset instead of EOF; any non-timeout terminal
// network error is the same fail-closed outcome.
t.Logf("terminal read error after invalid obfuscated tag: %v", err)
}
}
func startTransportFrameServer(t *testing.T, opts Options) (string, <-chan int) {
t.Helper()
ln, err := net.Listen("tcp", "127.0.0.1:0")
if err != nil {
t.Fatalf("listen: %v", err)
}
frames := make(chan int, 1)
srv := New(opts)
srv.onFrame = func(n int) {
select {
case frames <- n:
default:
}
}
ctx, cancel := context.WithCancel(context.Background())
serveErr := make(chan error, 1)
go func() { serveErr <- srv.Serve(ctx, ln) }()
t.Cleanup(func() {
cancel()
select {
case err := <-serveErr:
if err != nil {
t.Errorf("serve: %v", err)
}
case <-time.After(5 * time.Second):
t.Error("server did not stop")
}
})
return ln.Addr().String(), frames
}
func BenchmarkDetectTCPTransport(b *testing.B) {
tests := []struct {
name string
prefix []byte
}{
{name: "plain_abridged", prefix: []byte{0xef}},
{name: "plain_intermediate", prefix: []byte{0xee, 0xee, 0xee, 0xee}},
{name: "plain_full", prefix: []byte{12, 0, 0, 0, 0, 0, 0, 0}},
{name: "obfuscated", prefix: []byte{1, 2, 3, 4, 5, 6, 7, 8}},
}
for _, tt := range tests {
b.Run(tt.name, func(b *testing.B) {
conn := &transportProbeBenchmarkConn{payload: tt.prefix}
b.ReportAllocs()
for i := 0; i < b.N; i++ {
conn.offset = 0
if _, err := detectTCPTransport(conn); err != nil {
b.Fatal(err)
}
}
})
}
}
type transportProbeBenchmarkConn struct {
payload []byte
offset int
}
func (c *transportProbeBenchmarkConn) Read(p []byte) (int, error) {
if c.offset >= len(c.payload) {
return 0, io.EOF
}
n := copy(p, c.payload[c.offset:])
c.offset += n
return n, nil
}
func (*transportProbeBenchmarkConn) Write(p []byte) (int, error) { return len(p), nil }
func (*transportProbeBenchmarkConn) Close() error { return nil }
func (*transportProbeBenchmarkConn) LocalAddr() net.Addr { return nil }
func (*transportProbeBenchmarkConn) RemoteAddr() net.Addr { return nil }
func (*transportProbeBenchmarkConn) SetDeadline(time.Time) error { return nil }
func (*transportProbeBenchmarkConn) SetReadDeadline(time.Time) error {
return nil
}
func (*transportProbeBenchmarkConn) SetWriteDeadline(time.Time) error {
return nil
}