merged with fixes

This commit is contained in:
onysd 2026-09-09 02:49:30 +03:00
parent a9e758b712
commit 2f1818d656
176 changed files with 9000 additions and 907 deletions

View file

@ -0,0 +1,190 @@
package mtprotoedge
import "sync"
// bulkRPCScheduler limits runnable file handlers before they enter the shared
// inbound worker pool. Waiters remain behind inboundRPCGate, so bulk backend
// latency cannot occupy every worker needed by bootstrap and control RPCs.
type bulkRPCScheduler struct {
mu sync.Mutex
max int
inUse int
waiters []*bulkRPCLease
closed bool
}
type bulkRPCLease struct {
scheduler *bulkRPCScheduler
granted bool
released bool
notified bool
notify func(bool)
}
// bulkRPCAdmission serializes the two bulk prerequisites: a request may enter
// the process-wide handler scheduler only after its logical session owns an ACK
// window credit. This prevents credit-blocked requests from parking all global
// handler slots.
type bulkRPCAdmission struct {
mu sync.Mutex
scheduler *bulkRPCScheduler
lease *bulkRPCLease
released bool
}
func newBulkRPCScheduler(max int) *bulkRPCScheduler {
if max <= 0 {
max = 1
}
return &bulkRPCScheduler{max: max}
}
func (s *bulkRPCScheduler) reserve() *bulkRPCLease {
if s == nil {
return nil
}
lease := &bulkRPCLease{scheduler: s}
s.mu.Lock()
if s.closed {
lease.released = true
} else if s.inUse < s.max {
s.inUse++
lease.granted = true
} else {
s.waiters = append(s.waiters, lease)
}
s.mu.Unlock()
return lease
}
func (l *bulkRPCLease) subscribe(notify func(bool)) {
if l == nil || l.scheduler == nil || notify == nil {
return
}
s := l.scheduler
s.mu.Lock()
l.notify = notify
granted, released := l.granted, l.released
shouldNotify := (granted || released) && !l.notified
if shouldNotify {
l.notified = true
}
s.mu.Unlock()
if !shouldNotify {
return
}
if granted {
notify(true)
} else {
notify(false)
}
}
func (l *bulkRPCLease) release() {
if l == nil || l.scheduler == nil {
return
}
s := l.scheduler
var notifications []func(bool)
s.mu.Lock()
if l.released {
s.mu.Unlock()
return
}
l.released = true
if l.granted {
l.granted = false
s.inUse--
}
for !s.closed && s.inUse < s.max && len(s.waiters) > 0 {
next := s.waiters[0]
s.waiters[0] = nil
s.waiters = s.waiters[1:]
if next == nil || next.released {
continue
}
next.granted = true
s.inUse++
if next.notify != nil && !next.notified {
next.notified = true
notifications = append(notifications, next.notify)
}
}
s.mu.Unlock()
for _, notify := range notifications {
notify(true)
}
}
func newBulkRPCAdmission(scheduler *bulkRPCScheduler) *bulkRPCAdmission {
if scheduler == nil {
return nil
}
return &bulkRPCAdmission{scheduler: scheduler}
}
func (a *bulkRPCAdmission) subscribeAfter(credit *outboundBulkCredit, notify func(bool)) {
if a == nil || a.scheduler == nil || credit == nil || notify == nil {
if notify != nil {
notify(false)
}
return
}
credit.subscribe(func(success bool) {
if !success {
notify(false)
return
}
lease := a.scheduler.reserve()
a.mu.Lock()
if a.released {
a.mu.Unlock()
lease.release()
notify(false)
return
}
a.lease = lease
a.mu.Unlock()
lease.subscribe(notify)
})
}
func (a *bulkRPCAdmission) release() {
if a == nil {
return
}
a.mu.Lock()
if a.released {
a.mu.Unlock()
return
}
a.released = true
lease := a.lease
a.lease = nil
a.mu.Unlock()
lease.release()
}
func (s *bulkRPCScheduler) close() {
if s == nil {
return
}
var notifications []func(bool)
s.mu.Lock()
s.closed = true
for _, lease := range s.waiters {
if lease == nil || lease.released {
continue
}
lease.released = true
if lease.notify != nil && !lease.notified {
lease.notified = true
notifications = append(notifications, lease.notify)
}
}
s.waiters = nil
s.mu.Unlock()
for _, notify := range notifications {
notify(false)
}
}

View file

@ -0,0 +1,92 @@
package mtprotoedge
import (
"testing"
"time"
)
func TestBulkRPCSchedulerKeepsOverflowBehindGate(t *testing.T) {
scheduler := newBulkRPCScheduler(2)
first := scheduler.reserve()
second := scheduler.reserve()
third := scheduler.reserve()
woken := make(chan bool, 1)
third.subscribe(func(success bool) { woken <- success })
select {
case <-woken:
t.Fatal("overflow bulk handler became runnable before a slot was released")
default:
}
first.release()
select {
case success := <-woken:
if !success {
t.Fatal("overflow bulk handler was canceled instead of admitted")
}
case <-time.After(time.Second):
t.Fatal("overflow bulk handler did not become runnable")
}
second.release()
third.release()
}
func TestBulkRPCSchedulerCloseCancelsWaiters(t *testing.T) {
scheduler := newBulkRPCScheduler(1)
active := scheduler.reserve()
waiting := scheduler.reserve()
woken := make(chan bool, 1)
waiting.subscribe(func(success bool) { woken <- success })
scheduler.close()
select {
case success := <-woken:
if success {
t.Fatal("closed bulk scheduler granted a waiting handler")
}
case <-time.After(time.Second):
t.Fatal("closed bulk scheduler did not cancel waiter")
}
active.release()
}
func TestBulkRPCAdmissionDoesNotReserveGlobalSlotBeforeSessionCredit(t *testing.T) {
scheduler := newBulkRPCScheduler(1)
state := newOutboundStateWithLimits(newOutboundTrackedBudget(1<<20), 128, 1<<20)
state.bulkMax = 1
activeCredit := state.reserveBulkCredit()
waitingCredit := state.reserveBulkCredit()
admission := newBulkRPCAdmission(scheduler)
woken := make(chan bool, 1)
admission.subscribeAfter(waitingCredit.credit, func(success bool) { woken <- success })
scheduler.mu.Lock()
inUseBeforeCredit := scheduler.inUse
scheduler.mu.Unlock()
if inUseBeforeCredit != 0 {
t.Fatalf("credit-blocked request reserved %d global slots, want 0", inUseBeforeCredit)
}
select {
case <-woken:
t.Fatal("credit-blocked request became runnable")
default:
}
activeCredit.releaseIfOwned()
select {
case success := <-woken:
if !success {
t.Fatal("request was canceled after session credit became available")
}
case <-time.After(time.Second):
t.Fatal("request did not enter global scheduler after session credit")
}
scheduler.mu.Lock()
inUseAfterCredit := scheduler.inUse
scheduler.mu.Unlock()
if inUseAfterCredit != 1 {
t.Fatalf("admitted request reserved %d global slots, want 1", inUseAfterCredit)
}
admission.release()
waitingCredit.releaseIfOwned()
state.closeBulkWindow()
}

View file

@ -0,0 +1,29 @@
package mtprotoedge
import "testing"
func TestInboundSeenHistoryUsesStableCircularBacking(t *testing.T) {
state := newConnState()
for id := int64(1); id <= maxTrackedClientMsgIDs; id++ {
state.trackInbound(id, int32(id*2+1), true, false, msgStateReceived)
}
if len(state.order) != maxTrackedClientMsgIDs || len(state.seen) != maxTrackedClientMsgIDs {
t.Fatalf("initial seen history = order:%d map:%d", len(state.order), len(state.seen))
}
backing := &state.order[0]
for id := int64(maxTrackedClientMsgIDs + 1); id <= 4*maxTrackedClientMsgIDs; id++ {
state.trackInbound(id, int32(id*2+1), true, false, msgStateReceived)
}
if &state.order[0] != backing {
t.Fatal("full seen history replaced its circular backing")
}
if len(state.order) != maxTrackedClientMsgIDs || len(state.seen) != maxTrackedClientMsgIDs {
t.Fatalf("steady seen history = order:%d map:%d", len(state.order), len(state.seen))
}
if _, ok := state.seenRecord(1); ok {
t.Fatal("seen history retained its oldest ID")
}
if _, ok := state.seenRecord(4 * maxTrackedClientMsgIDs); !ok {
t.Fatal("seen history lost its newest ID")
}
}

View file

@ -0,0 +1,169 @@
package mtprotoedge
import (
"context"
"crypto/rand"
"crypto/rsa"
"fmt"
"net"
"testing"
"time"
"go.uber.org/zap/zaptest"
"github.com/gotd/log/logzap"
"github.com/iamxvbaba/td/clock"
"github.com/iamxvbaba/td/exchange"
"github.com/iamxvbaba/td/session"
"github.com/iamxvbaba/td/telegram"
"github.com/iamxvbaba/td/telegram/dcs"
"github.com/iamxvbaba/td/tg"
"github.com/iamxvbaba/td/transport"
"telesrv/internal/app/account"
"telesrv/internal/app/auth"
"telesrv/internal/app/contacts"
"telesrv/internal/app/dialogs"
"telesrv/internal/app/help"
"telesrv/internal/app/langpack"
"telesrv/internal/app/updates"
"telesrv/internal/app/users"
"telesrv/internal/domain"
"telesrv/internal/rpc"
"telesrv/internal/store/memory"
)
// TestClientIPPersistsToAuthorization verifies the real connection-to-persistence
// path: the accepted connection's remote IP is carried as neutral transport
// metadata by mtprotoedge, flows through RPC routing, and is persisted into the
// device authorization (authorizations.ip). It runs a full login over a real
// MTProto connection and asserts the stored authorization carries the client's
// loopback IP.
func TestClientIPPersistsToAuthorization(t *testing.T) {
const (
dc = 2
phone = "+8613800138100"
code = "12345"
clientIP = "127.0.0.1"
)
rsaKey, err := rsa.GenerateKey(rand.Reader, 2048)
if err != nil {
t.Fatalf("gen rsa: %v", err)
}
ln, err := net.Listen("tcp", "127.0.0.1:0")
if err != nil {
t.Fatalf("listen: %v", err)
}
tcpAddr := ln.Addr().(*net.TCPAddr)
if tcpAddr.IP.String() != clientIP {
t.Fatalf("test listener bound to %s, want loopback %s", tcpAddr.IP, clientIP)
}
userStore := memory.NewUserStore()
authzStore := memory.NewAuthorizationStore()
authKeyStore := memory.NewAuthKeyStore()
helpStore := memory.NewHelpStore()
if err := helpStore.UpsertAppConfig(context.Background(), domain.AppConfig{
Client: "tdesktop", Hash: 1_000_000,
JSON: []byte(`{"chat_read_mark_expire_period":604800,"chat_read_mark_size_threshold":50,"pm_read_date_expire_period":604800,"quote_length_max":1024,"telegram_antispam_group_size_min":200,"telegram_antispam_user_id":"5434988373"}`),
}); err != nil {
t.Fatalf("seed app config: %v", err)
}
if err := helpStore.UpsertCountries(context.Background(), []domain.Country{
{ISO2: "US", DefaultName: "United States", CountryCodes: []domain.CountryCode{{CountryCode: "1", Prefixes: []string{"1"}}}},
}); err != nil {
t.Fatalf("seed countries: %v", err)
}
langPackStore := memory.NewLangPackStore()
if err := langPackStore.UpsertPack(context.Background(), domain.LangPack{
LangPack: "tdesktop", LangCode: "en", Version: 1,
Strings: []domain.LangPackString{{Key: "lng_language_name", Value: "English"}},
}); err != nil {
t.Fatalf("seed langpack: %v", err)
}
deps := rpc.Deps{
Auth: auth.NewService(userStore, authzStore, memory.NewCodeStore(), authKeyStore, memory.NewTempAuthKeyBindingStore(authKeyStore), code),
Account: account.NewService(memory.NewPasswordStore()),
Help: help.NewService(helpStore, helpStore),
Users: users.NewService(userStore),
Updates: updates.NewService(memory.NewUpdateStateStore(), memory.NewUpdateEventStore()),
Contacts: contacts.NewService(memory.NewContactStore()),
Dialogs: dialogs.NewService(memory.NewDialogStore()),
LangPack: langpack.NewService(langPackStore),
}
router := rpc.New(rpc.Config{DC: dc, IP: tcpAddr.IP.String(), Port: tcpAddr.Port}, deps, zaptest.NewLogger(t), clock.System)
srv := New(Options{Logger: zaptest.NewLogger(t), DC: dc, RSAKey: rsaKey, AuthKeys: authKeyStore, LayerRPC: router})
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
defer cancel()
serveErr := make(chan error, 1)
go func() { serveErr <- srv.Serve(ctx, ln) }()
opts := telegram.Options{
PublicKeys: []exchange.PublicKey{{RSA: &rsaKey.PublicKey}},
Resolver: dcs.Plain(dcs.PlainOptions{Protocol: transport.Intermediate}),
DCList: dcs.List{Options: []tg.DCOption{{ID: dc, IPAddress: tcpAddr.IP.String(), Port: tcpAddr.Port, Static: true}}},
Logger: logzap.New(zaptest.NewLogger(t).Named("client")),
SessionStorage: &session.StorageMemory{},
UpdateHandler: telegram.UpdateHandlerFunc(func(context.Context, tg.UpdatesClass) error { return nil }),
}
client := telegram.NewClient(1, "hash", opts)
var newUserID int64
if err := client.Run(ctx, func(ctx context.Context) error {
raw := tg.NewClient(client)
sent, err := raw.AuthSendCode(ctx, &tg.AuthSendCodeRequest{PhoneNumber: phone, APIID: 1, APIHash: "hash", Settings: tg.CodeSettings{}})
if err != nil {
return err
}
sentCode, ok := sent.(*tg.AuthSentCode)
if !ok {
return fmt.Errorf("sendCode result = %T, want *tg.AuthSentCode", sent)
}
if _, err := raw.AuthSignIn(ctx, &tg.AuthSignInRequest{PhoneNumber: phone, PhoneCodeHash: sentCode.PhoneCodeHash, PhoneCode: code}); err != nil {
return err
}
signUpRes, err := raw.AuthSignUp(ctx, &tg.AuthSignUpRequest{PhoneNumber: phone, PhoneCodeHash: sentCode.PhoneCodeHash, FirstName: "IP", LastName: "Test"})
if err != nil {
return err
}
authz, ok := signUpRes.(*tg.AuthAuthorization)
if !ok {
return fmt.Errorf("signUp result = %T, want *tg.AuthAuthorization", signUpRes)
}
newUser, ok := authz.User.(*tg.User)
if !ok {
return fmt.Errorf("signUp user = %T, want *tg.User", authz.User)
}
newUserID = newUser.ID
return nil
}); err != nil {
t.Fatalf("client login flow: %v", err)
}
auths, err := authzStore.ListByUser(ctx, newUserID)
if err != nil || len(auths) == 0 {
t.Fatalf("authorizations for user %d = %d (err=%v), want >=1", newUserID, len(auths), err)
}
var gotIP string
for _, a := range auths {
if a.IP != "" {
gotIP = a.IP
break
}
}
if gotIP != clientIP {
t.Fatalf("persisted authorization IP = %q, want %q", gotIP, clientIP)
}
select {
case err := <-serveErr:
t.Fatalf("server stopped unexpectedly: %v", err)
default:
}
}

View file

@ -5,6 +5,7 @@ import (
"context"
"encoding/hex"
"errors"
"net"
"sync"
"sync/atomic"
"time"
@ -81,12 +82,12 @@ type Conn struct {
salt int64
key crypto.AuthKey
outbound chan outboundOp
outboundControl chan outboundOp
outbound chan *outboundOp
outboundControl chan *outboundOp
// Critical RPC results (session/difference convergence) and large bulk
// responses have independent bounded lanes. The actor remains the sole writer.
outboundCritical chan outboundOp
outboundBulk chan outboundOp
outboundCritical chan *outboundOp
outboundBulk chan *outboundOp
outboundStop chan struct{}
outboundDone chan struct{}
outboundClose sync.Once
@ -104,10 +105,14 @@ type Conn struct {
// Encoded MTProto service frames and control vectors use independent headroom: pong,
// new_session_created, bad_msg and msgs_ack must remain admissible when the body budget is
// full. Content-related control frames keep this budget while pending for resend.
outboundControlTrackedBudget *outboundTrackedBudget
outboundControlBudgetOnce sync.Once
outboundScratchPool *outboundScratchPool
outboundScratchOnce sync.Once
outboundControlTrackedBudget *outboundTrackedBudget
outboundControlBudgetOnce sync.Once
outboundCriticalTrackedBudget *outboundTrackedBudget
outboundCriticalBudgetOnce sync.Once
outboundScratchPool *outboundScratchPool
outboundScratchOnce sync.Once
outboundOpPool *outboundOpPool
outboundReplayBodyPool *outboundReplayBodyPool
// outboundState outlives this physical Conn generation. A replacement
// physical connection for the same auth key/session reuses it.
outboundState *outboundState
@ -160,6 +165,9 @@ type Conn struct {
// 单连接只保留并发配额;实际 worker 来自 Server 共享池,避免每连接预留 goroutine。
rpcRootCtx context.Context
rpcMaxInflight int
// remoteAddr 是 MTProto 连接的对端地址(来自 net.Conn.RemoteAddr仅保留 host
// 部分;绑定设备授权时写入 authorizations.ip便于在 admin 面板看到登录 IP。
remoteAddr string
// sentContentMessages is retained only for standalone construction tests.
// Server connections allocate seq_no from logical-session outboundState.
@ -419,3 +427,21 @@ func (c *Conn) ReceivesUpdates() bool { return c.receivesUpdates.Load() }
// SetReceivesUpdates 设置该连接是否接收主动推送的 updates。
// 登录后的主连接在 updates.getState/getDifference 建立同步基线后置为 true。
func (c *Conn) SetReceivesUpdates(v bool) { c.receivesUpdates.Store(v) }
// setRemoteAddrStr 记录 MTProto 连接的对端地址remote 形如 host:port
// 只保留 host 部分(去掉端口),用于绑定设备授权时写入 authorizations.ip。
func (c *Conn) setRemoteAddrStr(remote string) {
if remote == "" {
return
}
host, _, err := net.SplitHostPort(remote)
if err != nil {
// 已经是纯 host例如 unix socket 或 IPv6 无端口形式)。
c.remoteAddr = remote
return
}
c.remoteAddr = host
}
// clientIP 返回连接的对端 IPhost 部分),未设置时为空字符串。
func (c *Conn) clientIP() string { return c.remoteAddr }

View file

@ -22,8 +22,8 @@ func TestPushSkipsConnReboundToOtherUser(t *testing.T) {
c := &Conn{
sessionID: sid,
authKeyID: [8]byte{authKey},
outbound: make(chan outboundOp, 4),
outboundControl: make(chan outboundOp, 4),
outbound: make(chan *outboundOp, 4),
outboundControl: make(chan *outboundOp, 4),
outboundStop: make(chan struct{}),
}
c.userID.Store(userA)

View file

@ -27,6 +27,7 @@ import (
"github.com/iamxvbaba/td/tlprofile"
"telesrv/internal/observability/dbtrace"
"telesrv/internal/postresponse"
"telesrv/internal/rpcresult"
"telesrv/internal/store"
)
@ -40,6 +41,7 @@ type connState struct {
createdFloor int64
seen map[int64]clientMsgRecord // 已处理的 client msg_id用于幂等和 msgs_state_req
order []int64
orderHead int
minSeen int64
maxSeen int64
// maxContentMsgID/maxContentSeqNo 是已接受 content 消息的 msg_id / seq_no 高水位,
@ -116,7 +118,7 @@ var errActivationAuthKeyRejected = errors.New("activation auth key no longer exi
// 直接复用 current.key/current.salt 解密。任何 provisional 在 claim 建立后、发 required
// control 前都会最终回查 AuthKeyStore使外部撤销与 activation 线性化。
// plain 是 serveConn 持有的复用明文缓冲frame 的 slice 仅在下一帧解密前有效。
func (s *Server) handleEncrypted(ctx context.Context, tc transport.Conn, cs *connState, current *Conn, fetchedKey *store.AuthKeyData, b, plain *bin.Buffer) (*Conn, error) {
func (s *Server) handleEncrypted(ctx context.Context, tc transport.Conn, cs *connState, current *Conn, remote string, fetchedKey *store.AuthKeyData, b, plain *bin.Buffer) (*Conn, error) {
var key crypto.AuthKey
var serverSalt int64
var authKeyExpiresAt int
@ -161,6 +163,8 @@ func (s *Server) handleEncrypted(ctx context.Context, tc transport.Conn, cs *con
} else {
current = s.newConn(tc, key, frame.sessionID, serverSalt)
}
// 记录对端 IP供绑定设备授权时写入 authorizations.ipadmin 面板可见)。
current.setRemoteAddrStr(remote)
current.authKeyExpiresAt = authKeyExpiresAt
// Same-session evidence is restored as explicit; auth-key metadata is only
// an inherited default and can be corrected by the next invokeWithLayer.
@ -904,7 +908,7 @@ func (s *Server) publishRPCResult(
priority := rpcResultPriority(method, encoded)
encoded.priority = priority
if metrics, ok := s.metrics.(RPCResultMetrics); ok {
metrics.RPCResultPrepared(method, priority.String(), encoded.uncompressedBytes, len(encoded.body), encoded.compressed)
metrics.RPCResultPrepared(method, priority.String(), encoded.uncompressedBytes, encoded.wireSize(), encoded.compressed)
}
visible := encoded.compressed || priority == outboundPriorityCritical || priority == outboundPriorityBulk
return priority, visible
@ -916,21 +920,26 @@ func (s *Server) publishRPCResult(
// 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 {
if encoded != nil {
encoded.releaseBulkCredit()
}
return errors.New("rpc result receipt ledger is unavailable")
}
priority, visible := prepareEncoded(encoded)
if owner != nil && !owner.HandOff() {
encoded.releaseBulkCredit()
return ErrRPCResultFlightInvalid
}
started := time.Now()
encoded.markReplayable()
encoded.releaseBulkCredit()
// Complete may expose terminal execution only after the old connection
// is irreversibly unable to accept another same-generation request.
c.fenceUndeliveredRPCResult()
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)
metrics.RPCResultDelivered(method, latency, encoded.wireSize(), admissionErr)
}
resultLogLevel := zap.DebugLevel
if visible {
@ -941,14 +950,15 @@ func (s *Server) publishRPCResult(
zap.String("method", method), zap.Int64("req_msg_id", reqMsgID),
zap.Int64("delivered_req_msg_id", encoded.writtenRequestID()),
zap.String("auth_key_id", c.authKeyHex), zap.Int64("session_id", c.sessionID),
zap.Int("wire_bytes", len(encoded.body)), zap.Bool("gzip", encoded.compressed),
zap.Int("wire_bytes", encoded.wireSize()), zap.Bool("gzip", encoded.compressed),
zap.String("priority", priority.String()), zap.Error(admissionErr))
}
return nil
}
encoded, reserved, retained, err := s.encodeRPCResultReservedWithHandoffContext(
prepareCtx, c, reqMsgID, result, retainForReplay,
methodPriority := rpcMethodPriority(method)
encoded, reserved, retained, err := s.encodeRPCResultReservedWithPriorityAndHandoffContext(
prepareCtx, c, reqMsgID, result, methodPriority, retainForReplay,
)
if retained {
return err
@ -958,11 +968,14 @@ func (s *Server) publishRPCResult(
return err
}
if err != nil {
if provider, ok := result.(interface{ exactRPCBulkCredit() *outboundBulkCredit }); ok {
provider.exactRPCBulkCredit().release()
}
s.log.Warn("Encode RPC result failed; publishing INTERNAL",
zap.String("method", method), zap.Int64("req_msg_id", reqMsgID), zap.Error(err))
afterDelivered = nil
encoded, reserved, retained, err = s.encodeRPCResultReservedWithHandoffContext(
prepareCtx, c, reqMsgID, &mt.RPCError{ErrorCode: 500, ErrorMessage: "INTERNAL"}, retainForReplay,
encoded, reserved, retained, err = s.encodeRPCResultReservedWithPriorityAndHandoffContext(
prepareCtx, c, reqMsgID, &mt.RPCError{ErrorCode: 500, ErrorMessage: "INTERNAL"}, methodPriority, retainForReplay,
)
if retained {
return err
@ -973,6 +986,9 @@ func (s *Server) publishRPCResult(
}
}
if encoded == nil || reserved == nil {
if provider, ok := result.(interface{ exactRPCBulkCredit() *outboundBulkCredit }); ok {
provider.exactRPCBulkCredit().release()
}
c.fenceUndeliveredRPCResult()
return errors.New("rpc result encode completed without tracked retention")
}
@ -981,6 +997,7 @@ func (s *Server) publishRPCResult(
defer reserved.release()
priority, _ := prepareEncoded(encoded)
if owner != nil && !owner.HandOff() {
encoded.releaseBulkCredit()
return ErrRPCResultFlightInvalid
}
@ -989,7 +1006,7 @@ func (s *Server) publishRPCResult(
latency := time.Since(egressStarted)
deliveredReqMsgID := encoded.writtenRequestID()
if metrics, ok := s.metrics.(RPCResultMetrics); ok {
metrics.RPCResultDelivered(method, latency, len(encoded.body), deliveryErr)
metrics.RPCResultDelivered(method, latency, encoded.wireSize(), deliveryErr)
}
if deliveryErr != nil {
encoded.markReplayable()
@ -1000,7 +1017,7 @@ func (s *Server) publishRPCResult(
zap.String("method", method), zap.Int64("req_msg_id", reqMsgID),
zap.Int64("delivered_req_msg_id", deliveredReqMsgID),
zap.String("auth_key_id", c.authKeyHex), zap.Int64("session_id", c.sessionID),
zap.Int("wire_bytes", len(encoded.body)), zap.Bool("gzip", encoded.compressed),
zap.Int("wire_bytes", encoded.wireSize()), zap.Bool("gzip", encoded.compressed),
zap.Error(deliveryErr))
}
return
@ -1012,7 +1029,7 @@ func (s *Server) publishRPCResult(
zap.String("method", method), zap.Int64("req_msg_id", reqMsgID),
zap.Int64("delivered_req_msg_id", deliveredReqMsgID),
zap.String("auth_key_id", c.authKeyHex), zap.Int64("session_id", c.sessionID),
zap.Int("wire_bytes", len(encoded.body)), zap.Bool("gzip", encoded.compressed),
zap.Int("wire_bytes", encoded.wireSize()), zap.Bool("gzip", encoded.compressed),
zap.Duration("egress_latency", latency))
}
}
@ -1026,7 +1043,7 @@ func (s *Server) publishRPCResult(
if checked := s.log.Check(zap.DebugLevel, "RPC result admitted"); checked != nil {
checked.Write(
zap.String("method", method), zap.Int64("req_msg_id", reqMsgID),
zap.Int("wire_bytes", len(encoded.body)), zap.Int("inner_bytes", encoded.uncompressedBytes),
zap.Int("wire_bytes", encoded.wireSize()), zap.Int("inner_bytes", encoded.uncompressedBytes),
zap.Bool("gzip", encoded.compressed), zap.String("priority", priority.String()))
}
return nil
@ -1086,7 +1103,7 @@ func (s *Server) sendReplayedRPCResultWithHook(
c.fenceUndeliveredRPCResult()
return errors.New("nil replayed rpc_result")
}
attempt, reserved, err := c.cloneRPCResultForRequestReserved(encoded, encoded.reqMsgID, false)
attempt, reserved, err := c.cloneRPCResultForRequestReservedContext(ctx, encoded, encoded.reqMsgID, false)
if err != nil {
c.failOutboundBudget(err)
c.fenceUndeliveredRPCResult()
@ -1222,6 +1239,19 @@ func (s *Server) encodeRPCResultReservedWithHandoffContext(
reqMsgID int64,
result bin.Encoder,
handoff rpcResultRetentionHandoff,
) (*encodedOutboundMessage, *outboundBodyReservation, bool, error) {
return s.encodeRPCResultReservedWithPriorityAndHandoffContext(
ctx, c, reqMsgID, result, outboundPriorityNormal, handoff,
)
}
func (s *Server) encodeRPCResultReservedWithPriorityAndHandoffContext(
ctx context.Context,
c *Conn,
reqMsgID int64,
result bin.Encoder,
priority outboundPriority,
handoff rpcResultRetentionHandoff,
) (*encodedOutboundMessage, *outboundBodyReservation, bool, error) {
if ctx == nil {
ctx = context.Background()
@ -1237,7 +1267,10 @@ func (s *Server) encodeRPCResultReservedWithHandoffContext(
if err != nil {
return err
}
budget := c.outboundMessageBudget(encoded.typeID, false)
if priority != outboundPriorityNormal {
encoded.priority = priority
}
budget := c.outboundMessageBudgetForPriority(encoded.typeID, encoded.priority, false)
bytes := len(encoded.body)
if budget.reserve(bytes) {
reserved = &outboundBodyReservation{budget: budget, bytes: bytes}
@ -1286,6 +1319,14 @@ func (s *Server) encodeRPCResultWithoutSlot(ctx context.Context, c *Conn, reqMsg
return nil, fmt.Errorf("bind exact layer rpc result: %w", err)
}
}
var replaySource rpcresult.ReplaySource
if provider, ok := result.(interface{ exactRPCReplaySource() rpcresult.ReplaySource }); ok {
replaySource = provider.exactRPCReplaySource()
}
var bulkCredit *outboundBulkCredit
if provider, ok := result.(interface{ exactRPCBulkCredit() *outboundBulkCredit }); ok {
bulkCredit = provider.exactRPCBulkCredit()
}
// Encode the ordinary exact/no-gzip path directly behind the rpc_result
// prefix. This avoids both the old generated Prepare snapshot and another
// full-body copy merely to prepend the 12-byte envelope.
@ -1303,9 +1344,14 @@ func (s *Server) encodeRPCResultWithoutSlot(ctx context.Context, c *Conn, reqMsg
return nil, fmt.Errorf("%w: body=%d limit=%d", ErrOutboundMessageTooLarge, len(innerBody)+12, maxOutboundBodyBytes)
}
wireInner, compressed, err := encodeAdaptiveRPCResultInner(ctx, nil, innerBody)
if err != nil {
return nil, fmt.Errorf("compress rpc result: %w", err)
wireInner := innerBody
compressed := false
if replaySource == nil {
var err error
wireInner, compressed, err = encodeAdaptiveRPCResultInner(ctx, nil, innerBody)
if err != nil {
return nil, fmt.Errorf("compress rpc result: %w", err)
}
}
if len(wireInner) > maxOutboundBodyBytes-12 {
return nil, fmt.Errorf("%w: body=%d limit=%d", ErrOutboundMessageTooLarge, len(wireInner)+12, maxOutboundBodyBytes)
@ -1322,6 +1368,7 @@ func (s *Server) encodeRPCResultWithoutSlot(ctx context.Context, c *Conn, reqMsg
typeID: proto.ResultTypeID, body: body, reqMsgID: reqMsgID,
compressed: compressed, uncompressedBytes: len(innerBody), delivery: newRPCResultDelivery(0),
layer: layerBinding, layerInvariant: layerInvariantResult,
replaySource: replaySource, innerDigest: sha256.Sum256(innerBody), logicalBytes: len(body), bulkCredit: bulkCredit,
}, nil
}
@ -1620,18 +1667,23 @@ func (cs *connState) trackInbound(msgID int64, seqNo int32, content, service boo
cs.maxContentSeqNo = seqNo
}
}
cs.order = append(cs.order, msgID)
var evicted int64
if len(cs.order) < maxTrackedClientMsgIDs {
cs.order = append(cs.order, msgID)
} else {
evicted = cs.order[cs.orderHead]
cs.order[cs.orderHead] = msgID
cs.orderHead = (cs.orderHead + 1) % len(cs.order)
}
if msgID < cs.minSeen {
cs.minSeen = msgID
}
if msgID > cs.maxSeen {
cs.maxSeen = msgID
}
if len(cs.order) > maxTrackedClientMsgIDs {
oldest := cs.order[0]
cs.order = cs.order[1:]
delete(cs.seen, oldest)
if oldest == cs.minSeen || oldest == cs.maxSeen {
if evicted != 0 {
delete(cs.seen, evicted)
if evicted == cs.minSeen || evicted == cs.maxSeen {
cs.recomputeRange()
}
}

View file

@ -19,8 +19,8 @@ func TestRunFlushDiscardsBatchOnIdentitySwitch(t *testing.T) {
c := &Conn{
sessionID: sessionID,
authKeyID: raw,
outbound: make(chan outboundOp, 4),
outboundControl: make(chan outboundOp, 4),
outbound: make(chan *outboundOp, 4),
outboundControl: make(chan *outboundOp, 4),
outboundStop: make(chan struct{}),
}
sm.Register(c)

View file

@ -2098,7 +2098,7 @@ func TestLayerRPCDependencyGateUsesBusinessOutcome(t *testing.T) {
if dependencies.failed || len(dependencies.waiters) != 1 {
t.Fatalf("dependencies before completion = %+v", dependencies)
}
gate := newLayerRPCExecutionGate(c, dependencies)
gate := newLayerRPCExecutionGate(c, dependencies, nil, nil)
if gate == nil || gate.runnable() {
t.Fatal("dependency gate was runnable before business completion")
}
@ -2123,7 +2123,7 @@ func TestLayerRPCDependencyGateUsesBusinessOutcome(t *testing.T) {
t.Fatal(err)
}
dependencies := s.layerRPCDependencies(c, 304, missing)
gate := newLayerRPCExecutionGate(c, dependencies)
gate := newLayerRPCExecutionGate(c, dependencies, nil, nil)
if !dependencies.failed || gate == nil || !gate.runnable() || gate.success() {
t.Fatalf("missing dependency gate = deps:%+v runnable:%v success:%v", dependencies, gate.runnable(), gate.success())
}

View file

@ -8,6 +8,8 @@ import (
"sync"
"sync/atomic"
"time"
"telesrv/internal/transport"
)
// ErrInboundRPCQueueFull 表示 inbound RPC 已触达单连接或进程级预算。
@ -940,6 +942,11 @@ func (c *Conn) runInboundRPC(task inboundRPC) {
c.metrics.InboundRPCStarted(task.method, now.Sub(task.enqueuedAt))
ctx := task.ctx
// 把连接对端 IP 作为中立传输元数据注入,绑定设备授权时写 authorizations.ip。
// Edge 只产生传输事实;由 RPC 层的 transport.ClientIPFrom 消费,避免反向依赖。
if ip := c.clientIP(); ip != "" {
ctx = transport.WithClientIP(ctx, ip)
}
if task.run != nil {
_ = task.run(ctx)
}

View file

@ -13,6 +13,7 @@ import (
"github.com/iamxvbaba/td/tlprofile"
"telesrv/internal/observability/dbtrace"
"telesrv/internal/postresponse"
"telesrv/internal/rpcresult"
)
// layerRPCResultEncoder keeps the generated result bound to the immutable
@ -22,6 +23,8 @@ import (
type layerRPCResultEncoder struct {
call tlprofile.Call
result tlprofile.Result
source rpcresult.ReplaySource
bulk *outboundBulkCredit
}
func (e *layerRPCResultEncoder) Encode(b *bin.Buffer) error {
@ -45,6 +48,20 @@ func (e *layerRPCResultEncoder) exactLayerRPCResultBinding() outboundLayerBindin
}
}
func (e *layerRPCResultEncoder) exactRPCReplaySource() rpcresult.ReplaySource {
if e == nil {
return nil
}
return e.source
}
func (e *layerRPCResultEncoder) exactRPCBulkCredit() *outboundBulkCredit {
if e == nil {
return nil
}
return e.bulk
}
type exactLayerRPCResultEncoder interface {
bin.Encoder
exactLayerRPCResultBinding() outboundLayerBinding
@ -90,7 +107,11 @@ func bindAdmittedLayerRPCResult(request tlprofile.Admission, result tlprofile.Re
if result.Prepared().Identity() != request.Prepared().Identity() {
return nil, errLayerRPCResultIdentityMismatch
}
return &layerRPCResultEncoder{call: request.Call(), result: result}, nil
var source rpcresult.ReplaySource
if carrier, ok := result.(rpcresult.Carrier); ok {
source = carrier.ExactReplaySource()
}
return &layerRPCResultEncoder{call: request.Call(), result: result, source: source}, nil
}
func (s *Server) newInboundLayerRPCTask(
@ -104,7 +125,15 @@ func (s *Server) newInboundLayerRPCTask(
owner *rpcResultOwnerLease,
) inboundRPC {
wireSize := request.Prepared().WireSize()
gate := newLayerRPCExecutionGate(c, dependencies)
var bulkLease *outboundBulkCreditLease
var bulkExecution *bulkRPCAdmission
if method == "upload.getFile" && c != nil && c.outboundState != nil {
bulkLease = c.outboundState.reserveBulkCredit()
if s.bulkRPCScheduler != nil {
bulkExecution = newBulkRPCAdmission(s.bulkRPCScheduler)
}
}
gate := newLayerRPCExecutionGate(c, dependencies, bulkLease, bulkExecution)
timeoutResponse := func() {
writeTimeout := c.writeTimeout
if writeTimeout <= 0 || writeTimeout > 5*time.Second {
@ -126,6 +155,8 @@ func (s *Server) newInboundLayerRPCTask(
size: wireSize,
onTimeout: timeoutResponse,
release: func() {
bulkExecution.release()
bulkLease.releaseIfOwned()
if owner != nil && owner.Abort() {
c.fenceUndeliveredRPCResult()
}
@ -137,7 +168,7 @@ func (s *Server) newInboundLayerRPCTask(
ErrorCode: 500, ErrorMessage: "MSG_WAIT_FAILED",
}, nil)
}
if err := s.handleAdmittedLayerRPC(s.withLayerRPCProfileEvidenceFresh(taskCtx, profileEvidenceFresh), c, msgID, admissionSeq, method, request, owner); err != nil {
if err := s.handleAdmittedLayerRPC(s.withLayerRPCProfileEvidenceFresh(taskCtx, profileEvidenceFresh), c, msgID, admissionSeq, method, request, owner, bulkLease); err != nil {
fields := []zap.Field{
zap.Int64("msg_id", msgID), zap.String("auth_key_id", c.authKeyHex),
zap.Int64("session_id", c.sessionID), zap.Error(err),
@ -161,11 +192,15 @@ func layerRPCTimeoutMessage(gate *inboundRPCGate) string {
return "RPC_TIMEOUT"
}
func newLayerRPCExecutionGate(c *Conn, dependencies layerRPCDependencySet) *inboundRPCGate {
if len(dependencies.waiters) == 0 && !dependencies.failed {
func newLayerRPCExecutionGate(c *Conn, dependencies layerRPCDependencySet, bulk *outboundBulkCreditLease, execution *bulkRPCAdmission) *inboundRPCGate {
if len(dependencies.waiters) == 0 && !dependencies.failed && bulk == nil && execution == nil {
return nil
}
gate := newInboundRPCGate(len(dependencies.waiters), c.wakeInboundRPC)
prerequisites := len(dependencies.waiters)
if bulk != nil {
prerequisites++
}
gate := newInboundRPCGate(prerequisites, c.wakeInboundRPC)
if dependencies.failed {
gate.failed.Store(true)
}
@ -174,6 +209,11 @@ func newLayerRPCExecutionGate(c *Conn, dependencies layerRPCDependencySet) *inbo
gate.resolve(false)
}
}
if bulk != nil && bulk.credit != nil && execution != nil {
execution.subscribeAfter(bulk.credit, gate.resolve)
} else if bulk != nil && bulk.credit != nil {
bulk.credit.subscribe(gate.resolve)
}
// Release the subscriber-installation sentinel.
gate.resolve(true)
return gate
@ -202,6 +242,7 @@ func (s *Server) handleAdmittedLayerRPC(
method string,
request tlprofile.Admission,
owner *rpcResultOwnerLease,
bulkLease *outboundBulkCreditLease,
) error {
if s.layerRPC == nil {
return s.publishAdmittedLayerRPCResult(c, msgID, method, owner, false, &mt.RPCError{
@ -219,6 +260,9 @@ func (s *Server) handleAdmittedLayerRPC(
var exact *layerRPCResultEncoder
if err == nil && result != nil {
exact, err = bindAdmittedLayerRPCResult(request, result)
if err == nil && exact != nil && bulkLease != nil {
exact.bulk = bulkLease.transfer()
}
}
dur := s.clock.Now().Sub(start)
s.metrics.RPCHandled(effectiveMethod, dur, err)

View file

@ -101,7 +101,7 @@ func TestProjectionFailureCachesInternalWithoutRepeatingBusiness(t *testing.T) {
}
if err := s.handleAdmittedLayerRPC(
context.Background(), c, reqMsgID, claim.admissionSeq,
"help.getConfig", request, claim.owner,
"help.getConfig", request, claim.owner, nil,
); err != nil {
t.Fatalf("publish projection failure: %v", err)
}

File diff suppressed because it is too large Load diff

View file

@ -0,0 +1,102 @@
package mtprotoedge
// outboundReplayBodyPool owns the short-lived plaintext buffers used to
// materialize descriptor-backed rpc_result frames for an exact resend. It is
// deliberately bounded and Server-owned: a burst may warm the size classes,
// but it cannot make the process retain an unbounded sync.Pool tail.
type outboundReplayBodyPool struct {
classes []outboundReplayBodyClass
}
type outboundReplayBodyClass struct {
size int
idle chan []byte
}
type outboundReplayBodyClassSpec struct {
size int
maxIdle int
}
var defaultOutboundReplayBodyClasses = []outboundReplayBodyClassSpec{
{size: 4<<10 + 64, maxIdle: 32},
{size: 16<<10 + 64, maxIdle: 32},
{size: 64<<10 + 64, maxIdle: 32},
{size: 256<<10 + 64, maxIdle: 32},
{size: 512<<10 + 64, maxIdle: 32},
{size: 1<<20 + 64, maxIdle: 32},
{size: 2<<20 + 64, maxIdle: 8},
}
func newOutboundReplayBodyPool(specs []outboundReplayBodyClassSpec) *outboundReplayBodyPool {
classes := make([]outboundReplayBodyClass, 0, len(specs))
for _, spec := range specs {
if spec.size <= 0 || spec.maxIdle <= 0 {
continue
}
classes = append(classes, outboundReplayBodyClass{
size: spec.size,
idle: make(chan []byte, spec.maxIdle),
})
}
return &outboundReplayBodyPool{classes: classes}
}
// acquire returns an empty buffer and its owning class. Bodies larger than the
// largest reusable class use ordinary GC ownership and return class -1.
func (p *outboundReplayBodyPool) acquire(size int) ([]byte, int) {
if p == nil || size <= 0 {
return nil, -1
}
for class := range p.classes {
bucket := &p.classes[class]
if size > bucket.size {
continue
}
select {
case buf := <-bucket.idle:
return buf[:0], class
default:
return make([]byte, 0, bucket.size), class
}
}
return nil, -1
}
func (p *outboundReplayBodyPool) release(class int, buf []byte) {
if p == nil || class < 0 || class >= len(p.classes) {
return
}
bucket := &p.classes[class]
if cap(buf) != bucket.size {
return
}
buf = buf[:0]
select {
case bucket.idle <- buf:
default:
}
}
type outboundReplayBodyLease struct {
pool *outboundReplayBodyPool
class int
buf []byte
}
func (l *outboundReplayBodyLease) release() {
if l == nil || l.pool == nil {
return
}
l.pool.release(l.class, l.buf)
*l = outboundReplayBodyLease{}
}
var fallbackOutboundReplayBodyPool = newOutboundReplayBodyPool(defaultOutboundReplayBodyClasses)
func (c *Conn) replayBodyPool() *outboundReplayBodyPool {
if c != nil && c.outboundReplayBodyPool != nil {
return c.outboundReplayBodyPool
}
return fallbackOutboundReplayBodyPool
}

View file

@ -215,13 +215,13 @@ func TestSendRequiredControlQueueDeadlineTerminatesAndReturnsBudget(t *testing.T
writeTimeout: time.Second,
outboundTrackedBudget: newOutboundTrackedBudget(1 << 20),
outboundControlTrackedBudget: controlBudget,
outbound: make(chan outboundOp, 1),
outboundControl: make(chan outboundOp, 1),
outbound: make(chan *outboundOp, 1),
outboundControl: make(chan *outboundOp, 1),
outboundStop: make(chan struct{}),
}
// No actor is running and the bounded control queue is full, so the parent
// deadline must cover queue admission and make the failure terminal.
c.outboundControl <- outboundOp{kind: outboundAck}
c.outboundControl <- &outboundOp{kind: outboundAck}
ctx, cancel := context.WithTimeout(context.Background(), 25*time.Millisecond)
defer cancel()

View file

@ -4,12 +4,15 @@ import (
"bytes"
"context"
"crypto/rand"
"crypto/sha256"
"encoding/binary"
"errors"
"io"
"sync"
"sync/atomic"
"testing"
"time"
"unsafe"
"github.com/iamxvbaba/td/bin"
"github.com/iamxvbaba/td/crypto"
@ -20,6 +23,17 @@ import (
"github.com/iamxvbaba/td/transport"
)
type staticRPCReplaySource struct {
inner []byte
}
func (s *staticRPCReplaySource) EncodeInner(_ context.Context, out *bin.Buffer) error {
out.Put(s.inner)
return nil
}
func (*staticRPCReplaySource) RetainedBytes() int { return 128 }
type failAfterTransport struct {
failAt atomic.Int32
sends atomic.Int32
@ -448,6 +462,9 @@ func newOutboundTestConn(t *testing.T, tr transport.Conn, budget *outboundTracke
}
func TestOutboundQueueBackingUsesSmallConfigurableBounds(t *testing.T) {
if slot, wide := unsafe.Sizeof((*outboundOp)(nil)), unsafe.Sizeof(outboundOp{}); slot >= wide {
t.Fatalf("indirect queue slot = %d bytes, wide outbound op = %d bytes", slot, wide)
}
t.Run("defaults", func(t *testing.T) {
c := &Conn{metrics: NopMetrics{}}
c.startOutbound()
@ -477,6 +494,64 @@ func TestOutboundQueueBackingUsesSmallConfigurableBounds(t *testing.T) {
})
}
func TestOutboundOpPoolClearsReferencesAndBoundsIdle(t *testing.T) {
pool := newOutboundOpPool(1)
op := pool.acquire()
op.ctx = context.Background()
op.msg = &mt.PingRequest{PingID: 1}
op.encoded = &encodedOutboundMessage{body: []byte("payload")}
op.ids = []int64{1, 2, 3}
op.done = make(chan outboundResult, 1)
op.terminal = func(error) {}
pool.release(op)
reused := pool.acquire()
if reused != op {
t.Fatal("idle outbound op was not reused")
}
if reused.ctx != nil || reused.msg != nil || reused.encoded != nil || reused.ids != nil || reused.done != nil || reused.terminal != nil {
t.Fatalf("reused outbound op retained references: %+v", reused)
}
pool.release(reused)
pool.release(&outboundOp{})
if got := len(pool.idle); got != 1 {
t.Fatalf("idle outbound op count = %d, want bounded 1", got)
}
}
func BenchmarkOutboundOpPool(b *testing.B) {
pool := newOutboundOpPool(1)
b.ReportAllocs()
for b.Loop() {
op := pool.acquire()
op.kind = outboundSend
pool.release(op)
}
}
func TestOutboundAckHistoryUsesStableCircularBacking(t *testing.T) {
state := newOutboundState(newOutboundTrackedBudget(1 << 20))
for id := int64(1); id <= maxTrackedAckedMsgIDs; id++ {
state.markAcked(id)
}
if len(state.ackOrder) != maxTrackedAckedMsgIDs || len(state.acked) != maxTrackedAckedMsgIDs {
t.Fatalf("initial ack history = order:%d map:%d", len(state.ackOrder), len(state.acked))
}
backing := &state.ackOrder[0]
for id := int64(maxTrackedAckedMsgIDs + 1); id <= 4*maxTrackedAckedMsgIDs; id++ {
state.markAcked(id)
}
if &state.ackOrder[0] != backing {
t.Fatal("full ack history replaced its circular backing")
}
if len(state.ackOrder) != maxTrackedAckedMsgIDs || len(state.acked) != maxTrackedAckedMsgIDs {
t.Fatalf("steady ack history = order:%d map:%d", len(state.ackOrder), len(state.acked))
}
if state.isKnown(1) || !state.isKnown(4*maxTrackedAckedMsgIDs) {
t.Fatal("ack history did not evict oldest and retain newest IDs")
}
}
func TestOutboundOptionsDefaults(t *testing.T) {
opts := Options{}
opts.setDefaults()
@ -486,13 +561,17 @@ func TestOutboundOptionsDefaults(t *testing.T) {
if opts.OutboundTrackedGlobalMaxBytes != 512<<20 {
t.Fatalf("outbound tracked default = %d, want %d", opts.OutboundTrackedGlobalMaxBytes, 512<<20)
}
if opts.OutboundCriticalGlobalMaxBytes != 64<<20 {
t.Fatalf("outbound critical default = %d, want %d", opts.OutboundCriticalGlobalMaxBytes, 64<<20)
}
}
func TestServerNewConnectionsShareOutboundBudgetAndQueueLimits(t *testing.T) {
srv := New(Options{
OutboundQueueSize: 7,
OutboundControlQueueSize: 3,
OutboundTrackedGlobalMaxBytes: 20,
OutboundQueueSize: 7,
OutboundControlQueueSize: 3,
OutboundTrackedGlobalMaxBytes: 20,
OutboundCriticalGlobalMaxBytes: 30,
})
var rawKey crypto.Key
key := rawKey.WithID()
@ -513,6 +592,12 @@ func TestServerNewConnectionsShareOutboundBudgetAndQueueLimits(t *testing.T) {
if got := srv.outboundTrackedBudget.maxBytes; got != 20 {
t.Fatalf("server outbound tracked max = %d, want 20", got)
}
if c1.outboundCriticalTrackedBudget != srv.outboundCriticalBudget || c2.outboundCriticalTrackedBudget != srv.outboundCriticalBudget {
t.Fatal("server connections did not receive the shared critical tracking budget")
}
if got := srv.outboundCriticalBudget.maxBytes; got != 30 {
t.Fatalf("server outbound critical max = %d, want 30", got)
}
}
func TestEncryptOutboundFrameDecryptsWithGotdCipher(t *testing.T) {
@ -909,6 +994,162 @@ func TestOutboundTrackedBudgetWriteFailureReturnsReservation(t *testing.T) {
}
}
func TestOutboundStateCompactsImmutableRPCResultAndReplaysExactBody(t *testing.T) {
budget := newOutboundTrackedBudget(1 << 20)
state := newOutboundStateWithLimits(budget, 64, 1<<20)
inner := bytes.Repeat([]byte{0x5a}, 4096)
var body bin.Buffer
body.PutID(proto.ResultTypeID)
body.PutLong(7001)
body.Put(inner)
wire := body.Raw()
if !budget.reserve(len(wire)) {
t.Fatal("reserve first-write body")
}
frame := &outboundFrame{
msgID: 9001,
seqNo: 1,
typeID: proto.ResultTypeID,
body: wire,
reservedBytes: len(wire),
reservationBudget: budget,
reqMsgID: 7001,
replaySource: &staticRPCReplaySource{inner: append([]byte(nil), inner...)},
innerDigest: sha256.Sum256(inner),
uncompressedBytes: len(inner),
logicalBytes: len(wire),
}
if err := state.admitReserved(frame); err != nil {
t.Fatalf("admit frame: %v", err)
}
if !state.compactImmutableFrame(frame) {
t.Fatal("immutable frame was not compacted")
}
if frame.body != nil {
t.Fatal("compacted frame retained full body")
}
if got := budget.snapshot(); got != outboundReplayDescriptorCharge {
t.Fatalf("retained bytes = %d, want descriptor charge %d", got, outboundReplayDescriptorCharge)
}
replay, ok := state.rpcResult(7001)
if !ok || replay.replaySource == nil || len(replay.body) != 0 {
t.Fatalf("replay descriptor = %+v ok=%v", replay, ok)
}
materialized, err := replay.materializeRPCResultBody(context.Background(), 7001)
if err != nil {
t.Fatalf("materialize replay: %v", err)
}
if !bytes.Equal(materialized, wire) {
t.Fatal("materialized replay differs from first-write body")
}
state.ack([]int64{9001})
if got := budget.snapshot(); got != 0 {
t.Fatalf("retained bytes after ACK = %d, want 0", got)
}
}
func TestImmutableRPCResultMaterializesDirectlyIntoScratch(t *testing.T) {
inner := bytes.Repeat([]byte{0x6b}, 1<<20)
logicalBytes := 12 + len(inner)
replay := &encodedOutboundMessage{
typeID: proto.ResultTypeID,
reqMsgID: 7101,
replaySource: &staticRPCReplaySource{inner: inner},
innerDigest: sha256.Sum256(inner),
uncompressedBytes: len(inner),
logicalBytes: logicalBytes,
}
pool := newOutboundReplayBodyPool([]outboundReplayBodyClassSpec{{size: logicalBytes, maxIdle: 1}})
scratch, class := pool.acquire(logicalBytes)
body, usedScratch, err := replay.materializeRPCResultBodyInto(context.Background(), replay.reqMsgID, scratch)
if err != nil {
t.Fatalf("materialize replay: %v", err)
}
if !usedScratch {
t.Fatal("descriptor replay did not use the supplied scratch buffer")
}
if len(body) != logicalBytes || &body[0] != &scratch[:cap(scratch)][0] {
t.Fatal("materialized body does not alias the supplied scratch buffer")
}
if got := int64(binary.LittleEndian.Uint64(body[4:12])); got != replay.reqMsgID {
t.Fatalf("materialized req_msg_id = %d, want %d", got, replay.reqMsgID)
}
pool.release(class, body)
if got := len(pool.classes[class].idle); got != 1 {
t.Fatalf("idle pooled bodies = %d, want 1", got)
}
}
func TestOutboundReplayBodyPoolBoundsIdleBuffers(t *testing.T) {
pool := newOutboundReplayBodyPool([]outboundReplayBodyClassSpec{{size: 4096, maxIdle: 1}})
first, firstClass := pool.acquire(4000)
second, secondClass := pool.acquire(4000)
if firstClass != 0 || secondClass != 0 || cap(first) != 4096 || cap(second) != 4096 {
t.Fatalf("acquired classes/capacities = (%d,%d) (%d,%d)", firstClass, cap(first), secondClass, cap(second))
}
pool.release(firstClass, first)
pool.release(secondClass, second)
if got := len(pool.classes[0].idle); got != 1 {
t.Fatalf("idle pooled bodies = %d, want bounded at 1", got)
}
oversized, class := pool.acquire(4097)
if oversized != nil || class != -1 {
t.Fatalf("oversized acquisition = len:%d class:%d, want GC-owned nil/-1", len(oversized), class)
}
}
func BenchmarkImmutableRPCResultMaterializePooled(b *testing.B) {
inner := bytes.Repeat([]byte{0x6b}, 1<<20)
logicalBytes := 12 + len(inner)
replay := &encodedOutboundMessage{
typeID: proto.ResultTypeID,
reqMsgID: 7101,
replaySource: &staticRPCReplaySource{inner: inner},
innerDigest: sha256.Sum256(inner),
uncompressedBytes: len(inner),
logicalBytes: logicalBytes,
}
pool := newOutboundReplayBodyPool([]outboundReplayBodyClassSpec{{size: logicalBytes, maxIdle: 1}})
b.ReportAllocs()
b.SetBytes(int64(logicalBytes))
b.ResetTimer()
for range b.N {
scratch, class := pool.acquire(logicalBytes)
body, usedScratch, err := replay.materializeRPCResultBodyInto(context.Background(), replay.reqMsgID, scratch)
if err != nil || !usedScratch {
b.Fatalf("materialize replay: used=%v err=%v", usedScratch, err)
}
pool.release(class, body)
}
}
func TestOutboundBulkACKWindowWakesNextWaiter(t *testing.T) {
state := newOutboundStateWithLimits(newOutboundTrackedBudget(1<<20), 128, 1<<20)
leasing := make([]*outboundBulkCreditLease, 0, defaultBulkACKWindow+1)
for range defaultBulkACKWindow + 1 {
leasing = append(leasing, state.reserveBulkCredit())
}
woken := make(chan bool, 1)
leasing[len(leasing)-1].credit.subscribe(func(success bool) { woken <- success })
select {
case <-woken:
t.Fatal("window overflow waiter woke before ACK credit release")
default:
}
leasing[0].releaseIfOwned()
select {
case success := <-woken:
if !success {
t.Fatal("window waiter was canceled instead of granted")
}
case <-time.After(time.Second):
t.Fatal("window waiter did not wake after credit release")
}
for _, lease := range leasing[1:] {
lease.releaseIfOwned()
}
}
func TestOutboundStateEvictionReturnsTrackedBudget(t *testing.T) {
budget := newOutboundTrackedBudget(64)
state := newOutboundStateWithLimits(budget, 2, 8)
@ -991,11 +1232,11 @@ func TestOutboundStateReleasesMixedBodyAndControlBudgets(t *testing.T) {
func TestSendBestEffortQueueFullBehavior(t *testing.T) {
c := &Conn{metrics: NopMetrics{}, outboundTrackedBudget: newOutboundTrackedBudget(1 << 20)}
c.outbound = make(chan outboundOp, 1)
c.outboundControl = make(chan outboundOp, 1)
c.outbound = make(chan *outboundOp, 1)
c.outboundControl = make(chan *outboundOp, 1)
c.outboundStop = make(chan struct{})
// 占满普通队列,模拟出站拥塞。
c.outbound <- outboundOp{}
c.outbound <- &outboundOp{}
if err := c.SendBestEffort(context.Background(), proto.MessageFromServer, &mt.MsgsAck{}, 0); err != ErrOutboundQueueFull {
t.Fatalf("timeout=0 on full queue: err = %v, want ErrOutboundQueueFull", err)
@ -1027,10 +1268,10 @@ func TestSendBestEffortQueueFullBehavior(t *testing.T) {
func TestSendAsyncControlQueueBoundary(t *testing.T) {
c := &Conn{metrics: NopMetrics{}, outboundTrackedBudget: newOutboundTrackedBudget(1 << 20)}
c.outbound = make(chan outboundOp, 1)
c.outboundControl = make(chan outboundOp, 1)
c.outbound = make(chan *outboundOp, 1)
c.outboundControl = make(chan *outboundOp, 1)
c.outboundStop = make(chan struct{})
c.outboundControl <- outboundOp{kind: outboundAck}
c.outboundControl <- &outboundOp{kind: outboundAck}
if err := c.SendAsync(context.Background(), proto.MessageFromServer, &mt.MsgsAck{}); err != nil {
t.Fatalf("SendAsync on full control queue: %v", err)

View file

@ -325,16 +325,24 @@ func encodeAdaptiveRPCResultInner(ctx context.Context, stop <-chan struct{}, inn
// Android. These bootstrap barriers must pass background prefetch regardless of
// platform or their own encoded size.
func rpcResultPriority(method string, encoded *encodedOutboundMessage) outboundPriority {
if priority := rpcMethodPriority(method); priority != outboundPriorityNormal {
return priority
}
return classifyOutboundPriority(encoded, false)
}
func rpcMethodPriority(method string) outboundPriority {
base := method
if i := strings.IndexByte(base, '#'); i >= 0 {
base = base[:i]
}
switch base {
case "updates.getDifference", "updates.getChannelDifference", "updates.getState",
"messages.getDialogs", "messages.getPinnedDialogs":
"messages.getDialogs", "messages.getPinnedDialogs",
"auth.bindTempAuthKey", "help.getConfig", "users.getUsers":
return outboundPriorityCritical
}
return classifyOutboundPriority(encoded, false)
return outboundPriorityNormal
}
func (p outboundPriority) String() string {

View file

@ -104,6 +104,30 @@ func TestEncodeRPCResultReservedChargesBodyBeforeReturning(t *testing.T) {
}
}
func TestCriticalRPCResultUsesIndependentRetainedBudget(t *testing.T) {
ordinary := newOutboundTrackedBudget(1)
critical := newOutboundTrackedBudget(1 << 20)
c := legacyCanonicalTestConn(t, &Conn{
metrics: NopMetrics{},
outboundTrackedBudget: ordinary,
outboundCriticalTrackedBudget: critical,
})
s := New(Options{})
encoded, reserved, retained, err := s.encodeRPCResultReservedWithPriorityAndHandoffContext(
context.Background(), c, 791, exactTestRPCResult(&tg.DataJSON{Data: "bootstrap"}), outboundPriorityCritical, nil,
)
if err != nil || retained || encoded == nil || reserved == nil {
t.Fatalf("critical encode encoded=%p reserved=%p retained=%v err=%v", encoded, reserved, retained, err)
}
if got := ordinary.snapshot(); got != 0 {
t.Fatalf("ordinary budget used by critical result = %d", got)
}
if got, want := critical.snapshot(), int64(len(encoded.body)); got != want {
t.Fatalf("critical budget = %d, want %d", got, want)
}
reserved.release()
}
func TestEncodeRPCResultReservedDropsBodyOnBudgetTimeout(t *testing.T) {
const maxBytes = 1 << 20
budget := newOutboundTrackedBudget(maxBytes)

View file

@ -23,10 +23,10 @@ func TestRPCResultCloneReservationIsOneShotUnderReleaseRace(t *testing.T) {
}
reserved := &outboundBodyReservation{budget: budget, bytes: len(encoded.body)}
start := make(chan struct{})
taken := make(chan outboundOp, 1)
taken := make(chan *outboundOp, 1)
go func() {
<-start
op, _ := reserved.take(encoded)
op, _ := reserved.take(encoded, fallbackOutboundOpPool)
taken <- op
}()
released := make(chan struct{})
@ -53,7 +53,7 @@ func TestRPCResultReservationReleaseWinsAdmissionRollback(t *testing.T) {
t.Fatal("reserve body")
}
reserved := &outboundBodyReservation{budget: budget, bytes: len(encoded.body)}
op, err := reserved.take(encoded)
op, err := reserved.take(encoded, fallbackOutboundOpPool)
if err != nil {
t.Fatalf("take reservation: %v", err)
}
@ -61,7 +61,7 @@ func TestRPCResultReservationReleaseWinsAdmissionRollback(t *testing.T) {
// rolls the op back. The rollback must observe the release request and return
// the raw op charge instead of resurrecting an owner nobody will release.
reserved.release()
if !reserved.reclaim(&op) {
if !reserved.reclaim(op) {
t.Fatal("reclaim actor reservation")
}
if got := budget.snapshot(); got != 0 {
@ -175,7 +175,7 @@ func TestOutboundActorRetargetRequiresSecondBodyReservation(t *testing.T) {
terminalBytes = budget.snapshot()
},
}
err := c.handleOutboundSend(state, op)
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)
@ -228,7 +228,7 @@ func TestOutboundActorRetargetTransfersOnlyReplacementToPending(t *testing.T) {
terminalBytes = budget.snapshot()
},
}
err := c.handleOutboundSend(state, op)
err := c.handleOutboundSend(state, &op)
op.finish(outboundResult{err: err})
if terminalErr != nil {
t.Fatalf("retarget terminal error: %v", terminalErr)

View file

@ -1135,5 +1135,5 @@ func (s *Server) publishRewrappedRPCResult(
s.log.Info("RPC init rewrap result replay delivered",
zap.String("method", method), zap.Int64("req_msg_id", reqMsgID),
zap.String("auth_key_id", c.authKeyHex), zap.Int64("session_id", c.sessionID),
zap.Int("wire_bytes", len(encoded.body)))
zap.Int("wire_bytes", encoded.wireSize()))
}

View file

@ -38,6 +38,8 @@ type RuntimeSnapshot struct {
OutboundTrackedMaxBytes int64
OutboundControlBytes int64
OutboundControlMaxBytes int64
OutboundCriticalBytes int64
OutboundCriticalMaxBytes int64
OutboundWriteBytes int64
OutboundWriteMaxBytes int64
RPCExecutionOwners int64
@ -190,6 +192,10 @@ func (s *Server) RuntimeSnapshot() RuntimeSnapshot {
result.OutboundControlBytes = s.outboundControlBudget.snapshot()
result.OutboundControlMaxBytes = s.outboundControlBudget.maxBytes
}
if s.outboundCriticalBudget != nil {
result.OutboundCriticalBytes = s.outboundCriticalBudget.snapshot()
result.OutboundCriticalMaxBytes = s.outboundCriticalBudget.maxBytes
}
if s.outboundScratchPool != nil && s.outboundScratchPool.budget != nil {
result.OutboundWriteBytes = s.outboundScratchPool.snapshot()
result.OutboundWriteMaxBytes = s.outboundScratchPool.budget.maxBytes

View file

@ -330,27 +330,21 @@ type Options struct {
// 阻断连接维持消息。可靠响应无法 tracking 时终止该连接durable best-effort update
// 则只丢在线加速并由 difference 恢复。
OutboundTrackedGlobalMaxBytes int64
// OutboundCriticalGlobalMaxBytes is an independent retained-body reserve for
// bootstrap/convergence RPC results. Bulk traffic cannot consume it.
OutboundCriticalGlobalMaxBytes int64
// OutboundWriteGlobalMaxBytes bounds concurrent encrypted wire/codec/obfuscation scratch.
// Scratch is shared and pooled across connections; default 512 MiB.
OutboundWriteGlobalMaxBytes int64
// DC 是本 server 的 DC ID。默认 2。
DC int
// StrictDC turns on exact DC-ID validation for the permanent-key exchange
// (default off = lenient). telesrv is always a single physical backend —
// there is no real multi-DC federation behind it — but the OwpenGram
// client forks intentionally run in "single-server backend" mode, where
// dc_id 1..5 all alias to this one server (see owpengram_servers.cpp /
// ApplyServerToDcOptions in the desktop client) so that any old data
// referencing a specific dc_id still resolves correctly. When tdesktop
// adds a new local account it picks its own starting dc_id (its usual
// multi-DC load-spreading behavior, unrelated to which physical server
// it's actually talking to) — that choice is not guaranteed to equal our
// configured DC. Strict validation would reject those accounts with
// "-444 wrong dc_id" even though they are connecting to the right (and
// only) server; dc_id is a client-side routing label here, not part of
// key derivation, so accepting the mismatch does not weaken the exchange.
// The switch exists for a hypothetical future real multi-DC deployment.
// StrictDC enables DC-label validation during key exchange. It is false by
// default: this single physical backend accepts every wire int32 label for
// permanent and temporary keys, and the label never changes auth-key
// persistence, session identity, or business state. When enabled,
// permanent labels must equal DC and temporary labels may equal +/-DC.
// This diagnostic switch does not itself provide multi-DC isolation.
StrictDC bool
// RSAKey 是 server RSA 私钥用于密钥交换。nil 时无法完成握手。
RSAKey *rsa.PrivateKey
@ -458,6 +452,9 @@ func (o *Options) setDefaults() {
if o.OutboundTrackedGlobalMaxBytes <= 0 {
o.OutboundTrackedGlobalMaxBytes = defaultOutboundTrackedMaxBytes
}
if o.OutboundCriticalGlobalMaxBytes <= 0 {
o.OutboundCriticalGlobalMaxBytes = defaultOutboundCriticalMaxBytes
}
if o.OutboundWriteGlobalMaxBytes <= 0 {
o.OutboundWriteGlobalMaxBytes = defaultOutboundWriteMaxBytes
}
@ -518,13 +515,17 @@ type Server struct {
rpcQueueSize int
rpcTimeout time.Duration
rpcScheduler *inboundRPCScheduler
bulkRPCScheduler *bulkRPCScheduler
rpcDeliveryHooks *rpcDeliveryHookExecutor
frameBudget *inboundFrameBudget
outboundQueueSize int
outboundControlQueueSize int
outboundTrackedBudget *outboundTrackedBudget
outboundControlBudget *outboundTrackedBudget
outboundCriticalBudget *outboundTrackedBudget
outboundScratchPool *outboundScratchPool
outboundOpPool *outboundOpPool
outboundReplayBodyPool *outboundReplayBodyPool
dc int
strictDC bool
@ -581,17 +582,20 @@ func New(opts Options) *Server {
rpcQueueSize: opts.RPCQueueSize,
rpcTimeout: opts.RPCTimeout,
rpcScheduler: newInboundRPCScheduler(opts.RPCGlobalWorkers, opts.RPCGlobalMaxTasks, opts.RPCGlobalMaxBytes),
bulkRPCScheduler: newBulkRPCScheduler(max(1, opts.RPCGlobalWorkers/2)),
rpcDeliveryHooks: newRPCDeliveryHookExecutor(opts.RPCDeliveryHookWorkers, opts.RPCDeliveryHookMaxPending),
frameBudget: newInboundFrameBudget(opts.InboundFrameGlobalMaxBytes),
outboundQueueSize: opts.OutboundQueueSize,
outboundControlQueueSize: opts.OutboundControlQueueSize,
outboundTrackedBudget: newOutboundTrackedBudget(opts.OutboundTrackedGlobalMaxBytes),
outboundControlBudget: newOutboundTrackedBudget(defaultOutboundControlMaxBytes),
outboundCriticalBudget: newOutboundTrackedBudget(opts.OutboundCriticalGlobalMaxBytes),
outboundScratchPool: newOutboundScratchPool(opts.OutboundWriteGlobalMaxBytes),
outboundOpPool: newOutboundOpPool(defaultOutboundOpPoolSize),
outboundReplayBodyPool: newOutboundReplayBodyPool(defaultOutboundReplayBodyClasses),
dc: opts.DC,
strictDC: opts.StrictDC,
key: exchange.PrivateKey{RSA: opts.RSAKey},
pubKeyPEM: rsaPublicKeyPEM(opts.RSAKey),
authKeys: opts.AuthKeys,
conns: conns,
rpc: opts.legacyRPC,
@ -612,6 +616,7 @@ func New(opts Options) *Server {
}),
rpcRewrap: newRPCRewrapRegistry(opts.RPCGlobalMaxTasks),
admission: newAdmissionController(opts.MaxConnections, opts.MaxConnectionsPerIP, opts.MaxConcurrentHandshakes),
pubKeyPEM: rsaPublicKeyPEM(opts.RSAKey),
}
if opts.IdentityDir != "" {
server.identityStore = identity.NewStore(opts.IdentityDir)
@ -662,26 +667,29 @@ func (s *Server) newConnWithLease(lease *physicalTransportLease, key crypto.Auth
func (s *Server) buildConn(tc transport.Conn, lease *physicalTransportLease, key crypto.AuthKey, sessionID, salt int64) *Conn {
c := &Conn{
transport: tc,
transportLease: lease,
writer: tc,
cipher: s.cipher,
msgID: proto.NewMessageIDGen(s.clock.Now),
writeTimeout: s.writeTimeout,
metrics: s.metrics,
now: s.clock.Now,
authKeyID: key.ID,
authKeyHex: hex.EncodeToString(key.ID[:]),
sessionID: sessionID,
salt: salt,
key: key,
createdAt: s.clock.Now(),
outboundQueueSize: s.outboundQueueSize,
outboundControlQueueSize: s.outboundControlQueueSize,
outboundTrackedBudget: s.outboundTrackedBudget,
outboundControlTrackedBudget: s.outboundControlBudget,
outboundScratchPool: s.outboundScratchPool,
rpcDeliveryHooks: s.rpcDeliveryHooks,
transport: tc,
transportLease: lease,
writer: tc,
cipher: s.cipher,
msgID: proto.NewMessageIDGen(s.clock.Now),
writeTimeout: s.writeTimeout,
metrics: s.metrics,
now: s.clock.Now,
authKeyID: key.ID,
authKeyHex: hex.EncodeToString(key.ID[:]),
sessionID: sessionID,
salt: salt,
key: key,
createdAt: s.clock.Now(),
outboundQueueSize: s.outboundQueueSize,
outboundControlQueueSize: s.outboundControlQueueSize,
outboundTrackedBudget: s.outboundTrackedBudget,
outboundControlTrackedBudget: s.outboundControlBudget,
outboundCriticalTrackedBudget: s.outboundCriticalBudget,
outboundScratchPool: s.outboundScratchPool,
outboundOpPool: s.outboundOpPool,
outboundReplayBodyPool: s.outboundReplayBodyPool,
rpcDeliveryHooks: s.rpcDeliveryHooks,
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
@ -708,6 +716,7 @@ func (s *Server) Serve(ctx context.Context, ln net.Listener) error {
defer s.rpcDeliveryHooks.stop(rpcCloseWaitTimeout)
defer s.conns.releaseAllLogicalSessions()
defer s.rpcScheduler.stop(rpcCloseWaitTimeout)
defer s.bulkRPCScheduler.close()
// 只在最外层 listener 包一次,确保 same-port mux 的 sniff/HTTP upgrade 也计入
// raw admission而不是等连接已经分流后才计数。
ln = s.observeRawAccepts(s.admission.wrapListener(ln))
@ -1147,7 +1156,7 @@ func (s *Server) serveConn(ctx context.Context, raw transport.Conn, remote, loca
fetchedKey = &d
}
current, err = s.handleEncrypted(ctx, conn, cs, current, fetchedKey, &b, &plain)
current, err = s.handleEncrypted(ctx, conn, cs, current, remote, fetchedKey, &b, &plain)
if errors.Is(err, errActivationAuthKeyRejected) {
// handleEncrypted writes -404 while its activation claim still owns the
// physical writer, then its deferred abort removes/closes the claim.

View file

@ -77,7 +77,7 @@ func TestBadServerSaltRetainsOneProvisionalConnUntilCorrected(t *testing.T) {
cs := newConnState()
var plain bin.Buffer
firstConn, err := s.handleEncrypted(context.Background(), tr, cs, nil, &stored, firstWrong, &plain)
firstConn, err := s.handleEncrypted(context.Background(), tr, cs, nil, "", &stored, firstWrong, &plain)
if err != nil {
t.Fatalf("first bad salt: %v", err)
}
@ -92,7 +92,7 @@ func TestBadServerSaltRetainsOneProvisionalConnUntilCorrected(t *testing.T) {
secondWrong, _ := encryptedRPCFrameWithAuthoritativeSaltForBarrierTest(
t, key, wrongSalt, serverSalt, sessionID, secondID, 3,
)
secondConn, err := s.handleEncrypted(context.Background(), tr, cs, firstConn, nil, secondWrong, &plain)
secondConn, err := s.handleEncrypted(context.Background(), tr, cs, firstConn, "", nil, secondWrong, &plain)
if err != nil {
t.Fatalf("second bad salt: %v", err)
}
@ -123,7 +123,7 @@ func TestBadServerSaltRetainsOneProvisionalConnUntilCorrected(t *testing.T) {
corrected, _ := encryptedRPCFrameWithAuthoritativeSaltForBarrierTest(
t, key, serverSalt, serverSalt, sessionID, firstID, 1,
)
activeConn, err := s.handleEncrypted(context.Background(), tr, cs, secondConn, nil, corrected, &plain)
activeConn, err := s.handleEncrypted(context.Background(), tr, cs, secondConn, "", nil, corrected, &plain)
if err != nil {
t.Fatalf("corrected retry: %v", err)
}
@ -171,7 +171,7 @@ func TestWrongSaltSessionChangeTransfersPhysicalOwnership(t *testing.T) {
}
cs := newConnState()
var plain bin.Buffer
oldConn, err := s.handleEncrypted(context.Background(), tr, cs, nil, &stored, firstFrame, &plain)
oldConn, err := s.handleEncrypted(context.Background(), tr, cs, nil, "", &stored, firstFrame, &plain)
if err != nil {
t.Fatalf("activate first session: %v", err)
}
@ -185,7 +185,7 @@ func TestWrongSaltSessionChangeTransfersPhysicalOwnership(t *testing.T) {
wrongFrame, _ := encryptedRPCFrameWithAuthoritativeSaltForBarrierTest(
t, key, wrongSalt, serverSalt, secondSID, secondID, 1,
)
newConn, err := s.handleEncrypted(context.Background(), tr, cs, oldConn, nil, wrongFrame, &plain)
newConn, err := s.handleEncrypted(context.Background(), tr, cs, oldConn, "", nil, wrongFrame, &plain)
if err != nil {
t.Fatalf("new session bad salt: %v", err)
}
@ -205,7 +205,7 @@ func TestWrongSaltSessionChangeTransfersPhysicalOwnership(t *testing.T) {
corrected, _ := encryptedRPCFrameWithAuthoritativeSaltForBarrierTest(
t, key, serverSalt, serverSalt, secondSID, secondID, 1,
)
activated, err := s.handleEncrypted(context.Background(), tr, cs, newConn, nil, corrected, &plain)
activated, err := s.handleEncrypted(context.Background(), tr, cs, newConn, "", nil, corrected, &plain)
if err != nil {
t.Fatalf("activate transferred session: %v", err)
}
@ -318,7 +318,7 @@ func TestHandleEncryptedRequiredSessionBarrierPrecedesStateRegistrationAndRPC(t
}
done := make(chan result, 1)
go func() {
conn, err := s.handleEncrypted(context.Background(), tr, cs, nil, &stored, frame, &plain)
conn, err := s.handleEncrypted(context.Background(), tr, cs, nil, "", &stored, frame, &plain)
done <- result{conn: conn, err: err}
}()
@ -394,7 +394,7 @@ func TestHandleEncryptedRequiredSessionBarrierFailureIsAtomic(t *testing.T) {
done := make(chan error, 1)
go func() {
_, err := s.handleEncrypted(context.Background(), tr, cs, nil, &stored, frame, &plain)
_, err := s.handleEncrypted(context.Background(), tr, cs, nil, "", &stored, frame, &plain)
done <- err
}()
select {
@ -458,7 +458,7 @@ func TestCrossConnectionInflightRPCHasOneBusinessOwnerAndReplaysResult(t *testin
firstTransport := &collectingSessionTransport{}
firstState := newConnState()
var firstPlain bin.Buffer
firstConn, err := s.handleEncrypted(context.Background(), firstTransport, firstState, nil, &stored, firstFrame, &firstPlain)
firstConn, err := s.handleEncrypted(context.Background(), firstTransport, firstState, nil, "", &stored, firstFrame, &firstPlain)
if err != nil {
t.Fatalf("first handleEncrypted: %v", err)
}
@ -478,7 +478,7 @@ func TestCrossConnectionInflightRPCHasOneBusinessOwnerAndReplaysResult(t *testin
}
secondDone := make(chan handleResult, 1)
go func() {
conn, handleErr := s.handleEncrypted(context.Background(), secondTransport, secondState, nil, &stored, secondFrame, &secondPlain)
conn, handleErr := s.handleEncrypted(context.Background(), secondTransport, secondState, nil, "", &stored, secondFrame, &secondPlain)
secondDone <- handleResult{conn: conn, err: handleErr}
}()
@ -566,7 +566,7 @@ func TestCrossConnectionInflightAbortRetriesOnlyAfterOldOwnerStops(t *testing.T)
firstTransport := &collectingSessionTransport{}
firstState := newConnState()
var firstPlain bin.Buffer
firstConn, err := s.handleEncrypted(context.Background(), firstTransport, firstState, nil, &stored, firstFrame, &firstPlain)
firstConn, err := s.handleEncrypted(context.Background(), firstTransport, firstState, nil, "", &stored, firstFrame, &firstPlain)
if err != nil {
t.Fatalf("first handleEncrypted: %v", err)
}
@ -580,7 +580,7 @@ func TestCrossConnectionInflightAbortRetriesOnlyAfterOldOwnerStops(t *testing.T)
secondTransport := &collectingSessionTransport{}
secondState := newConnState()
var secondPlain bin.Buffer
secondConn, err := s.handleEncrypted(context.Background(), secondTransport, secondState, nil, &stored, secondFrame, &secondPlain)
secondConn, err := s.handleEncrypted(context.Background(), secondTransport, secondState, nil, "", &stored, secondFrame, &secondPlain)
if err != nil && !errors.Is(err, ErrConnClosed) {
t.Fatalf("second handleEncrypted: %v", err)
}
@ -605,7 +605,7 @@ func TestCrossConnectionInflightAbortRetriesOnlyAfterOldOwnerStops(t *testing.T)
thirdTransport := &collectingSessionTransport{}
thirdState := newConnState()
var thirdPlain bin.Buffer
thirdConn, err := s.handleEncrypted(context.Background(), thirdTransport, thirdState, nil, &stored, thirdFrame, &thirdPlain)
thirdConn, err := s.handleEncrypted(context.Background(), thirdTransport, thirdState, nil, "", &stored, thirdFrame, &thirdPlain)
if err != nil {
t.Fatalf("third handleEncrypted: %v", err)
}

View file

@ -219,8 +219,8 @@ func TestSessionManagerBestEffortFanoutPreparesOncePerProfile(t *testing.T) {
c := &Conn{
sessionID: int64(i + 1),
authKeyID: [8]byte{byte(i + 1)},
outbound: make(chan outboundOp, 1),
outboundControl: make(chan outboundOp, 1),
outbound: make(chan *outboundOp, 1),
outboundControl: make(chan *outboundOp, 1),
outboundStop: make(chan struct{}),
metrics: NopMetrics{},
}
@ -270,8 +270,8 @@ func TestSessionManagerMixedLayerFanoutUsesProfileBoundBodies(t *testing.T) {
c := &Conn{
sessionID: int64(profile),
authKeyID: authKeyID,
outbound: make(chan outboundOp, 1),
outboundControl: make(chan outboundOp, 1),
outbound: make(chan *outboundOp, 1),
outboundControl: make(chan *outboundOp, 1),
outboundStop: make(chan struct{}),
metrics: NopMetrics{},
}
@ -390,11 +390,11 @@ func TestSessionManagerBestEffortFanoutUsesOneBudgetAndDropsOnlySlowConsumers(t
authKeyID: [8]byte{byte(i + 1)},
transport: tr,
metrics: NopMetrics{},
outbound: make(chan outboundOp, 1),
outboundControl: make(chan outboundOp, 1),
outbound: make(chan *outboundOp, 1),
outboundControl: make(chan *outboundOp, 1),
outboundStop: make(chan struct{}),
}
c.outbound <- outboundOp{}
c.outbound <- &outboundOp{}
c.userID.Store(userID)
c.userIDResolved.Store(true)
c.receivesUpdates.Store(true)
@ -409,8 +409,8 @@ func TestSessionManagerBestEffortFanoutUsesOneBudgetAndDropsOnlySlowConsumers(t
sessionID: 99,
authKeyID: [8]byte{99},
metrics: NopMetrics{},
outbound: make(chan outboundOp, 1),
outboundControl: make(chan outboundOp, 1),
outbound: make(chan *outboundOp, 1),
outboundControl: make(chan *outboundOp, 1),
outboundStop: make(chan struct{}),
}
healthy.userID.Store(userID)
@ -656,8 +656,8 @@ func TestForceCloseBatchTimeoutStillClosesProducerAndRPCGates(t *testing.T) {
c := &Conn{
transport: tr,
metrics: NopMetrics{},
outbound: make(chan outboundOp, 1),
outboundControl: make(chan outboundOp, 1),
outbound: make(chan *outboundOp, 1),
outboundControl: make(chan *outboundOp, 1),
outboundStop: make(chan struct{}),
}
c.startInboundRPCScheduler(scheduler, 1, 1, time.Second)
@ -743,8 +743,8 @@ func TestPushToUserAuthKeyUsesOneDeadlineAndDropsOnlySlowPFSConnections(t *testi
sessionID: sessionID,
metrics: NopMetrics{},
transport: transport,
outbound: make(chan outboundOp, 1),
outboundControl: make(chan outboundOp, 1),
outbound: make(chan *outboundOp, 1),
outboundControl: make(chan *outboundOp, 1),
outboundStop: make(chan struct{}),
}
c.receivesUpdates.Store(true)
@ -752,7 +752,7 @@ func TestPushToUserAuthKeyUsesOneDeadlineAndDropsOnlySlowPFSConnections(t *testi
t.Fatalf("freeze profile: %v", err)
}
if queueFull {
c.outbound <- outboundOp{}
c.outbound <- &outboundOp{}
}
sm.Register(c)
sm.BindAuthKeyForSession(raw, sessionID, business)
@ -961,8 +961,8 @@ func TestPushToSessionForAuthKeyImmediateBypassesReadinessQueue(t *testing.T) {
c := &Conn{
sessionID: 42,
authKeyID: raw,
outbound: make(chan outboundOp, 1),
outboundControl: make(chan outboundOp, 1),
outbound: make(chan *outboundOp, 1),
outboundControl: make(chan *outboundOp, 1),
outboundStop: make(chan struct{}),
}
if err := c.FreezeLayerProfile(tlprofile.Profile227); err != nil {
@ -1006,8 +1006,8 @@ func TestSessionManagerWithholdsUpdatesReadinessUntilExactProfile(t *testing.T)
c := &Conn{
authKeyID: key.authKeyID,
sessionID: key.sessionID,
outbound: make(chan outboundOp, 1),
outboundControl: make(chan outboundOp, 1),
outbound: make(chan *outboundOp, 1),
outboundControl: make(chan *outboundOp, 1),
outboundStop: make(chan struct{}),
metrics: NopMetrics{},
outboundTrackedBudget: newOutboundTrackedBudget(1 << 20),
@ -1049,7 +1049,7 @@ func TestSessionManagerWithholdsUpdatesReadinessUntilExactProfile(t *testing.T)
}
c.membershipsSynced.Store(true)
sm.SetReceivesUpdatesForAuthKey(key.authKeyID, key.sessionID, true)
var op outboundOp
var op *outboundOp
select {
case op = <-c.outbound:
case <-time.After(time.Second):
@ -1210,8 +1210,8 @@ func TestPendingFlushGlobalBodyPressureDoesNotTerminateHealthyConnection(t *test
c := &Conn{
authKeyID: key.authKeyID,
sessionID: key.sessionID,
outbound: make(chan outboundOp, 1),
outboundControl: make(chan outboundOp, 1),
outbound: make(chan *outboundOp, 1),
outboundControl: make(chan *outboundOp, 1),
outboundStop: make(chan struct{}),
metrics: NopMetrics{},
outboundTrackedBudget: newOutboundTrackedBudget(1),

View file

@ -22,8 +22,8 @@ func TestTerminalFailurePathsCloseGatesBeforeBlockingTransportClose(t *testing.T
c := &Conn{
transport: tr,
metrics: NopMetrics{},
outbound: make(chan outboundOp, 1),
outboundControl: make(chan outboundOp, 1),
outbound: make(chan *outboundOp, 1),
outboundControl: make(chan *outboundOp, 1),
outboundStop: make(chan struct{}),
}
c.startInboundRPCScheduler(scheduler, 1, 1, time.Second)

View file

@ -224,8 +224,8 @@ func TestContainerInvalidSequenceTailIsAtomic(t *testing.T) {
cs := newConnState()
c := &Conn{
metrics: NopMetrics{},
outbound: make(chan outboundOp, 4),
outboundControl: make(chan outboundOp, 4),
outbound: make(chan *outboundOp, 4),
outboundControl: make(chan *outboundOp, 4),
outboundStop: make(chan struct{}),
}
var acks []int64

View file

@ -21,8 +21,8 @@ func TestPushTransientSkipsNotReadySession(t *testing.T) {
c := &Conn{
sessionID: 7,
authKeyID: [8]byte{7},
outbound: make(chan outboundOp, 4),
outboundControl: make(chan outboundOp, 4),
outbound: make(chan *outboundOp, 4),
outboundControl: make(chan *outboundOp, 4),
outboundStop: make(chan struct{}),
}
c.userID.Store(userID)
@ -63,7 +63,7 @@ func TestPushTransientCompatibleSkipsUnavailableAndUnknownProfiles(t *testing.T)
makeConn := func(sessionID int64, profile tlprofile.Profile, known bool) *Conn {
c := &Conn{
sessionID: sessionID, authKeyID: [8]byte{byte(sessionID)},
outbound: make(chan outboundOp, 2), outboundControl: make(chan outboundOp, 2),
outbound: make(chan *outboundOp, 2), outboundControl: make(chan *outboundOp, 2),
outboundStop: make(chan struct{}),
}
c.userID.Store(userID)