owpengram-server/internal/mtprotoedge/conn.go

package mtprotoedge

import (
	"bufio"
	"context"
	"encoding/hex"
	"errors"
	"sync"
	"sync/atomic"
	"time"

	"github.com/iamxvbaba/td/bin"
	"github.com/iamxvbaba/td/crypto"
	"github.com/iamxvbaba/td/proto"
	"github.com/iamxvbaba/td/tlprofile"
	"github.com/iamxvbaba/td/transport"
)

// Conn 是一个已识别 session 的客户端连接，持有向其加密发送消息所需的全部上下文。
// 由 SessionManager 管理，供请求响应与主动 push 共用。
//
// Send 并发安全：所有出站消息先进 per-Conn outbound actor，由它串行分配 msg_id/seq_no、
// 加密并写 transport，避免高并发 RPC 响应与 push 交错造成 MTProto 顺序错误。
type outboundWriter interface {
	Send(context.Context, *bin.Buffer) error
}

type connLifecycle uint32

const (
	// The zero value is deliberately provisional so test/embedded Conn values start
	// outside every SessionManager index until they complete activation.
	connLifecycleProvisional connLifecycle = iota
	connLifecycleClaiming
	connLifecycleActive
	// retired is terminal and irreversible. A physical connection that lost an
	// activation claim must never become visible again, even if its read goroutine
	// was already between preflight and publish when a replacement arrived.
	connLifecycleRetired
)

var (
	// ErrLayerProfileUnsupported means the requested profile has no generated
	// exact codec. Profiles are never clamped to the nearest supported layer.
	ErrLayerProfileUnsupported = errors.New("unsupported exact layer profile")
	// ErrLayerProfileConflict means one admitted request carried contradictory
	// profile evidence. A later, well-formed invokeWithLayer is allowed to correct
	// the connection profile and therefore does not use this error.
	ErrLayerProfileConflict = errors.New("connection layer profile conflict")
	// ErrLayerProfileEpochExhausted is a defensive terminal guard. Reaching it
	// would require more than four billion effective layer corrections on one
	// physical connection, so wrapping the epoch and making stale pushes current
	// again is never safe.
	ErrLayerProfileEpochExhausted = errors.New("connection layer profile epoch exhausted")
)

type Conn struct {
	transport transport.Conn
	// transportLease owns exactly one generation of the physical transport.
	// It is nil only for directly constructed test/embedded Conns that retain
	// the legacy raw-transport close fallback.
	transportLease *physicalTransportLease
	writer         outboundWriter
	cipher         crypto.Cipher
	msgID          *proto.MessageIDGen
	writeTimeout   time.Duration
	metrics        Metrics
	// now shares the Server protocol clock with inbound expiry admission. Tests may
	// advance it without sleeping; construction-only Conns fall back to time.Now.
	now func() time.Time

	authKeyID [8]byte
	// authKeyHex 是 authKeyID 的 hex 缓存：每条 RPC 的结构化日志都会带它，
	// 建连时算一次，避免热路径反复 hex 编码分配。
	authKeyHex string
	// authKeyExpiresAt=0 表示 permanent key；正值是 temporary/media-temporary
	// key 在握手时确定的绝对协议失效时间；-1 是仅供迁移的 legacy-unknown
	// sentinel（edge 会在创建 Conn 前以 -404 拒绝）。Conn 创建后不可变。
	authKeyExpiresAt int
	sessionID        int64
	salt             int64
	key              crypto.AuthKey

	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
	outboundStop     chan struct{}
	outboundDone     chan struct{}
	outboundClose    sync.Once
	// outboundEnqueueMu orders producer registration against terminal close. Close
	// flips closing under this lock before waiting, so no WaitGroup Add can race Wait.
	outboundEnqueueMu sync.Mutex
	outboundEnqueueWG sync.WaitGroup
	outboundClosing   bool
	// Queue backing is intentionally small and bounded per Conn; control has a separate queue
	// and strict actor priority. Server-created connections share outboundTrackedBudget.
	outboundQueueSize        int
	outboundControlQueueSize int
	outboundTrackedBudget    *outboundTrackedBudget
	outboundBudgetOnce       sync.Once
	// 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
	// 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.
	lifecycle      atomic.Uint32
	transportClose sync.Once

	rpcScheduler *inboundRPCScheduler
	rpcCancel    context.CancelFunc
	rpcClose     sync.Once
	rpcMu        sync.Mutex
	rpcWG        sync.WaitGroup
	// rpcReservationWG 跟踪 Copy 前预算到 commit/abort 的短窗口，使 Close 返回时
	// 全局/单连接预算都已归还或转交给明确的 queued/running task。
	rpcReservationWG sync.WaitGroup
	rpcTimeout       time.Duration
	rpcQueue         []inboundRPC
	rpcQueueSize     int
	rpcReserved      int
	rpcRunning       int
	rpcReady         bool
	rpcClosed        bool
	// rpcReplayRestores is a per-physical-connection ordering barrier. An exact
	// 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
	// restore finishes or the connection is fenced.
	rpcReplayRestores int
	// Rewrap aliasing never delays execution. initialized stops collecting
	// candidates after the first valid init wrapper on this physical generation.
	rpcRewrapInitialized atomic.Bool
	// rpcResultAcked is invoked by the sole outbound actor after it resolves an
	// acknowledged server frame back to the rpc_result request msg_id.
	rpcResultAcked func(*Conn, int64)
	// inflightRPCBytes 跟踪已预留/入队/执行中 inbound RPC 的 memory charge；legacy
	// 等于 copied body，exact 是 typed materialization 的保守放大值。它配合
	// maxInflightRPCBytes 给 RPC 队列设内存预算（不止限条数）。
	inflightRPCBytes atomic.Int64
	// 单连接只保留并发配额；实际 worker 来自 Server 共享池，避免每连接预留 goroutine。
	rpcRootCtx     context.Context
	rpcMaxInflight int

	// 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 一次。
	outboundRand *bufio.Reader

	identityMu              sync.RWMutex
	businessAuthKeyID       [8]byte
	businessAuthKeyHex      string
	businessAuthKeyResolved bool
	userID                  atomic.Int64
	userIDResolved          atomic.Bool
	receivesUpdates         atomic.Bool
	// membershipsSynced 表示该连接的 channel membership 推送路由（byMemberChannel）
	// 已成功建立。它与 receivesUpdates 共同构成「session 完全就绪」：membership
	// 同步失败时保持 false，让置位短路放行、下一条 RPC 重试同步，避免
	// 「已置位但 channel 路由缺失」的 session 静默漏收超级群推送。
	membershipsSynced atomic.Bool
	// membershipGen 是本连接 channel membership 索引的修订号：任何增量修订
	// （join/leave/kick 的 Add/Remove、身份切换/下线的整体清除）都递增。全量同步方
	// 在读取持久成员列表前采样、落地时带回比对，检测「读取窗口内发生增量修订」的
	// 丢失更新竞态（SetSessionChannelMemberships 改走合并路径并保持未就绪重试）。
	membershipGen atomic.Int64
	// createdAt 是连接建立时刻，供同 auth_key session 数触顶时驱逐真正最旧的连接。
	createdAt time.Time
	// layerProfileState atomically packs profile, provenance and epoch. A profile
	// inherited from auth-key metadata is only a default; a later well-formed
	// invokeWithLayer may correct it. The epoch fences proactive updates prepared
	// before that correction without invalidating request-bound RPC results.
	layerProfileMu sync.RWMutex
	// layerProfileEvidenceMsgID is protected by layerProfileMu. Zero means the
	// selected profile came from inherited/legacy recovery and therefore has no
	// ordered client-message cursor yet. Positive values are the newest accepted
	// invokeWithLayer message for this exact MTProto session.
	layerProfileEvidenceMsgID int64
	// layerProfileEvidenceLayer retains the raw negotiated Layer even when this
	// binary has no generated codec for it. In that case the packed profile stays
	// Unknown, but msg_id ordering can still admit a newer supported correction
	// and reject an older/same-id rollback.
	layerProfileEvidenceLayer int
	layerProfileState         atomic.Uint64
	// clientLayer is the package-internal mirror used only by legacy state-machine
	// regression tests. Production application RPC/result/update encoding uses
	// the structured profile state. Zero is unknown and must fail closed if a legacy
	// application value reaches an outbound boundary.
	clientLayer atomic.Int32
}

func (c *Conn) lifecycleState() connLifecycle {
	if c == nil {
		return connLifecycleRetired
	}
	return connLifecycle(c.lifecycle.Load())
}

func (c *Conn) isRetired() bool {
	return c == nil || c.lifecycleState() == connLifecycleRetired
}

// retire irreversibly fences the logical connection. The caller that wins the
// transition may additionally own one-shot physical cleanup.
func (c *Conn) retire() bool {
	if c == nil {
		return false
	}
	for {
		state := c.lifecycle.Load()
		if connLifecycle(state) == connLifecycleRetired {
			return false
		}
		if c.lifecycle.CompareAndSwap(state, uint32(connLifecycleRetired)) {
			return true
		}
	}
}

func (c *Conn) beginActivationClaim() bool {
	if c == nil || !c.isPhysicalTransportCurrentOpen() {
		return false
	}
	if !c.lifecycle.CompareAndSwap(uint32(connLifecycleProvisional), uint32(connLifecycleClaiming)) {
		return false
	}
	// Physical close can win after the pre-check but before the lifecycle CAS.
	// Do not let a doomed claimant enter SessionManager and retire a healthy old
	// owner for the same logical session.
	if c.lifecycleState() != connLifecycleClaiming || !c.isPhysicalTransportCurrentOpen() {
		c.retire()
		return false
	}
	return true
}

func (c *Conn) publishActivation() bool {
	if c == nil || !c.isPhysicalTransportCurrentOpen() {
		return false
	}
	if !c.lifecycle.CompareAndSwap(uint32(connLifecycleClaiming), uint32(connLifecycleActive)) {
		return false
	}
	// A concurrent transport failure can retire the Conn between the first
	// physical check and the CAS. Never let that intermediate active value escape.
	if c.lifecycleState() != connLifecycleActive || !c.isPhysicalTransportCurrentOpen() {
		c.retire()
		return false
	}
	return true
}

func (c *Conn) isActive() bool {
	return c != nil && c.lifecycleState() == connLifecycleActive && c.isPhysicalTransportCurrentOpen()
}

// transferTransportOwnership hands this Conn's physical socket to the next
// logical generation. The caller must have fenced and drained the old writer.
func (c *Conn) transferTransportOwnership() (*physicalTransportLease, bool) {
	if c == nil || c.transportLease == nil {
		return nil, false
	}
	return c.transportLease.Transfer()
}

func (c *Conn) isPhysicalTransportCurrentOpen() bool {
	return c != nil && (c.transportLease == nil || c.transportLease.IsCurrentOpen())
}

// LayerProfile returns the exact TL profile currently selected for this
// connection. ok is false until admission or an inherited auth-key default
// supplies a supported generated profile.
func (c *Conn) LayerProfile() (profile tlprofile.Profile, ok bool) {
	state := c.LayerProfileState()
	return state.Profile, state.Origin != LayerProfileUnknown
}

// FreezeLayerProfile records explicit protocol evidence observed during ordered
// admission. Repeating the same value is idempotent. A later well-formed
// invokeWithLayer may replace either an inherited default or older explicit
// evidence; already-admitted requests retain their own immutable profile.
func (c *Conn) FreezeLayerProfile(profile tlprofile.Profile) error {
	_, err := c.setLayerProfile(profile, LayerProfileExplicit, true)
	return err
}

// FreezeLayerProfileAt applies explicit protocol evidence in client msg_id
// order. A duplicate older than the last accepted evidence is inert; the same
// msg_id carrying another Layer is a protocol conflict. Advancing the evidence
// cursor at an unchanged Layer does not rotate the outbound epoch because the
// wire profile itself did not change.
func (c *Conn) FreezeLayerProfileAt(profile tlprofile.Profile, msgID int64) (bool, error) {
	return c.freezeLayerProfileAt(profile, msgID)
}

// SeedLayerProfile restores explicit evidence previously proven for this exact
// logical session. It is kept as the compatible same-session restore API;
// auth-key-wide metadata must use SeedInheritedLayerProfile instead.
func (c *Conn) SeedLayerProfile(profile tlprofile.Profile) error {
	_, err := c.setLayerProfile(profile, LayerProfileExplicit, true)
	return err
}

// SeedInheritedLayerProfile installs an auth-key-wide default only while the
// connection is still unknown. It never overwrites explicit evidence or an
// already selected inherited default; client protocol evidence owns correction.
func (c *Conn) SeedInheritedLayerProfile(profile tlprofile.Profile) error {
	_, err := c.setLayerProfile(profile, LayerProfileInherited, false)
	return err
}

// legacyClientLayer returns the test-only canonical-transcoder profile. Zero
// deliberately remains unknown; it must never be converted into an implicit
// canonical application profile.
func (c *Conn) legacyClientLayer() int { return int(c.clientLayer.Load()) }

// setLegacyClientLayer records the package-internal legacy mirror.
func (c *Conn) setLegacyClientLayer(layer int) { c.clientLayer.Store(int32(layer)) }

// AuthKeyID 返回连接的 auth_key_id。
func (c *Conn) AuthKeyID() [8]byte { return c.authKeyID }

// AuthKeyExpiresAt 返回 raw 协议 key 的失效时间；0 表示 permanent key。
func (c *Conn) AuthKeyExpiresAt() int { return c.authKeyExpiresAt }

func (c *Conn) authKeyProtocolUnavailableNow() bool {
	if c == nil {
		return true
	}
	now := time.Now()
	if c.now != nil {
		now = c.now()
	}
	return authKeyProtocolUnavailable(c.authKeyExpiresAt, now)
}

// BusinessAuthKeyID 返回业务视角的 auth_key_id。
//
// temp auth_key 绑定后解析为 perm auth_key；第二个返回值表示本连接是否已完成解析，
// 即便解析结果等于原始 auth_key_id 也会返回 true，以避免每个 RPC 重复查绑定表。
func (c *Conn) BusinessAuthKeyID() ([8]byte, bool) {
	c.identityMu.RLock()
	defer c.identityMu.RUnlock()
	return c.businessAuthKeyID, c.businessAuthKeyResolved
}

// BusinessAuthKeyHex 返回业务视角 auth_key_id 的 hex 缓存（每 RPC 日志用，免重复编码）。
func (c *Conn) BusinessAuthKeyHex() (string, bool) {
	c.identityMu.RLock()
	defer c.identityMu.RUnlock()
	return c.businessAuthKeyHex, c.businessAuthKeyResolved
}

// SetBusinessAuthKeyID 缓存业务视角 auth_key_id。
func (c *Conn) SetBusinessAuthKeyID(id [8]byte) {
	c.identityMu.Lock()
	changed := !c.businessAuthKeyResolved || c.businessAuthKeyID != id
	if changed || c.businessAuthKeyHex == "" {
		c.businessAuthKeyHex = hex.EncodeToString(id[:])
	}
	c.businessAuthKeyID = id
	c.businessAuthKeyResolved = true
	c.identityMu.Unlock()
	if changed {
		c.userID.Store(0)
		c.userIDResolved.Store(false)
	}
}

// SessionID 返回连接的 session_id。
func (c *Conn) SessionID() int64 { return c.sessionID }

// UserID 返回绑定的用户 id；未登录为 0。
func (c *Conn) UserID() int64 { return c.userID.Load() }

// UserIDResolved 返回 user_id 授权状态是否已为当前连接解析过。
//
// resolved=true 且 userID=0 表示该 auth_key 当前未登录；这样登录前的多次 RPC
// 不会反复查询授权表，后续登录成功会由 BindUser 覆盖为真实用户。
func (c *Conn) UserIDResolved() (userID int64, resolved bool) {
	return c.userID.Load(), c.userIDResolved.Load()
}

// ReceivesUpdates 报告该连接是否接收主动推送的 updates。
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) }