416 lines
12 KiB
Go
416 lines
12 KiB
Go
package rpc
|
||
|
||
import (
|
||
"context"
|
||
"encoding/hex"
|
||
"fmt"
|
||
"sync"
|
||
"time"
|
||
|
||
"go.uber.org/zap"
|
||
"golang.org/x/sync/singleflight"
|
||
|
||
"github.com/gotd/td/bin"
|
||
"github.com/gotd/td/clock"
|
||
"github.com/gotd/td/tg"
|
||
)
|
||
|
||
// maxWrapperDepth 限制 invokeWithLayer/initConnection 等 wrapper 的嵌套深度，防御恶意构造。
|
||
const maxWrapperDepth = 4
|
||
|
||
var (
|
||
tlTypeNamesOnce sync.Once
|
||
tlTypeNames map[uint32]string
|
||
)
|
||
|
||
// Config 是 Router 所需的服务端信息。
|
||
type Config struct {
|
||
DC int
|
||
IP string // 对外公布的 DC IP（写入 DCOptions）
|
||
Port int // 对外公布的 DC 端口
|
||
OutboundPushTimeout time.Duration
|
||
}
|
||
|
||
// Router 把解密后的 RPC 请求按 TypeID 路由到 typed handler（tg.ServerDispatcher）。
|
||
//
|
||
// handler 输入输出均为 gotd/td/tg 类型，各业务域的 handler
|
||
// 与注册见 help.go / auth.go / users.go / updates.go。Router 本身只负责协议外壳：
|
||
// 剥离 invokeWithLayer / initConnection / invokeWithoutUpdates，并兜底未注册 RPC。
|
||
type Router struct {
|
||
cfg Config
|
||
log *zap.Logger
|
||
clock clock.Clock
|
||
deps Deps
|
||
dispatcher *tg.ServerDispatcher
|
||
clientInfoMu sync.RWMutex
|
||
clientInfo map[clientInfoSessionKey]ClientInfo
|
||
authUserMu sync.RWMutex
|
||
authUsers map[[8]byte]authUserCacheEntry
|
||
authUserSF singleflight.Group
|
||
presence *presenceTracker
|
||
}
|
||
|
||
type clientInfoSessionKey struct {
|
||
rawAuthKeyID [8]byte
|
||
sessionID int64
|
||
}
|
||
|
||
type authUserCacheEntry struct {
|
||
userID int64
|
||
found bool
|
||
}
|
||
|
||
// New 创建 Router，由各业务域自行注册其 RPC handler（registerHelp/Auth/Users/Updates）。
|
||
func New(cfg Config, deps Deps, log *zap.Logger, clk clock.Clock) *Router {
|
||
r := &Router{cfg: cfg, log: log, clock: clk, deps: deps, presence: newPresenceTracker()}
|
||
d := tg.NewServerDispatcher(r.fallback)
|
||
|
||
r.registerHelp(d)
|
||
r.registerAuth(d)
|
||
r.registerUsers(d)
|
||
r.registerUpdates(d)
|
||
r.registerAccount(d)
|
||
r.registerMessages(d)
|
||
r.registerChannels(d)
|
||
r.registerUpload(d)
|
||
r.registerPhotos(d)
|
||
r.registerFolders(d)
|
||
r.registerContacts(d)
|
||
r.registerLangpack(d)
|
||
r.registerStories(d)
|
||
r.registerPayments(d)
|
||
r.registerStats(d)
|
||
r.registerPremium(d)
|
||
r.registerAiCompose(d)
|
||
|
||
r.dispatcher = d
|
||
return r
|
||
}
|
||
|
||
// Dispatch 路由一条 RPC 请求：先剥离 invokeWithLayer / initConnection /
|
||
// invokeWithoutUpdates 等 wrapper（注入 layer / 客户端信息到 ctx），
|
||
// 再按 TypeID 路由到 typed handler。满足 mtprotoedge.RPCHandler。
|
||
func (r *Router) Dispatch(ctx context.Context, authKeyID [8]byte, sessionID int64, b *bin.Buffer) (bin.Encoder, error) {
|
||
ctx = WithRawAuthKeyID(ctx, authKeyID)
|
||
effectiveAuthKeyID, err := r.effectiveAuthKeyID(ctx, authKeyID, sessionID)
|
||
if err != nil {
|
||
return nil, internalErr()
|
||
}
|
||
ctx = WithAuthKeyID(ctx, effectiveAuthKeyID)
|
||
ctx = WithSessionID(ctx, sessionID)
|
||
userID, hasUserID, err := r.effectiveUserID(ctx, authKeyID, effectiveAuthKeyID, sessionID)
|
||
if err != nil {
|
||
return nil, internalErr()
|
||
}
|
||
if hasUserID {
|
||
ctx = WithUserID(ctx, userID)
|
||
}
|
||
if info, ok := r.clientInfoForSession(ctx); ok {
|
||
ctx = WithClientInfo(ctx, info)
|
||
}
|
||
return r.dispatch(ctx, b, 0)
|
||
}
|
||
|
||
func (r *Router) effectiveAuthKeyID(ctx context.Context, rawAuthKeyID [8]byte, sessionID int64) ([8]byte, error) {
|
||
if r.deps.Sessions != nil {
|
||
if scoped, ok := r.deps.Sessions.(ScopedSessionBinder); ok {
|
||
if id, ok := scoped.AuthKeyIDForSession(rawAuthKeyID, sessionID); ok {
|
||
return id, nil
|
||
}
|
||
} else if id, ok := r.deps.Sessions.AuthKeyID(sessionID); ok {
|
||
return id, nil
|
||
}
|
||
}
|
||
effective := rawAuthKeyID
|
||
if r.deps.Auth != nil {
|
||
resolved, ok, err := r.deps.Auth.ResolveAuthKey(ctx, rawAuthKeyID)
|
||
if err != nil {
|
||
return [8]byte{}, err
|
||
}
|
||
if ok {
|
||
effective = resolved
|
||
}
|
||
}
|
||
if r.deps.Sessions != nil {
|
||
if scoped, ok := r.deps.Sessions.(ScopedSessionBinder); ok {
|
||
scoped.BindAuthKeyForSession(rawAuthKeyID, sessionID, effective)
|
||
} else {
|
||
r.deps.Sessions.BindAuthKey(sessionID, effective)
|
||
}
|
||
}
|
||
return effective, nil
|
||
}
|
||
|
||
func (r *Router) effectiveUserID(ctx context.Context, rawAuthKeyID, authKeyID [8]byte, sessionID int64) (int64, bool, error) {
|
||
if userID, ok := UserIDFrom(ctx); ok {
|
||
if scoped, ok := r.scopedSessions(); ok {
|
||
scoped.BindUserForAuthKey(rawAuthKeyID, sessionID, userID)
|
||
} else if r.deps.Sessions != nil {
|
||
r.deps.Sessions.BindUser(sessionID, userID)
|
||
}
|
||
return userID, true, nil
|
||
}
|
||
if r.deps.Sessions != nil {
|
||
if scoped, ok := r.deps.Sessions.(ScopedSessionBinder); ok {
|
||
if userID, resolved := scoped.UserIDResolvedForAuthKey(rawAuthKeyID, sessionID); resolved {
|
||
return userID, userID != 0, nil
|
||
}
|
||
} else if userID, resolved := r.deps.Sessions.UserIDResolved(sessionID); resolved {
|
||
return userID, userID != 0, nil
|
||
}
|
||
}
|
||
if r.deps.Auth == nil {
|
||
return 0, false, nil
|
||
}
|
||
userID, found, err := r.lookupAuthUser(ctx, authKeyID)
|
||
if err != nil {
|
||
return 0, false, err
|
||
}
|
||
if r.deps.Sessions != nil {
|
||
if scoped, ok := r.deps.Sessions.(ScopedSessionBinder); ok {
|
||
if cachedUserID, resolved := scoped.UserIDResolvedForAuthKey(rawAuthKeyID, sessionID); resolved {
|
||
return cachedUserID, cachedUserID != 0, nil
|
||
}
|
||
} else if cachedUserID, resolved := r.deps.Sessions.UserIDResolved(sessionID); resolved {
|
||
return cachedUserID, cachedUserID != 0, nil
|
||
}
|
||
if found {
|
||
if scoped, ok := r.deps.Sessions.(ScopedSessionBinder); ok {
|
||
scoped.BindUserForAuthKey(rawAuthKeyID, sessionID, userID)
|
||
} else {
|
||
r.deps.Sessions.BindUser(sessionID, userID)
|
||
}
|
||
r.announceSessionOnline(ctx, userID)
|
||
} else {
|
||
if scoped, ok := r.deps.Sessions.(ScopedSessionBinder); ok {
|
||
scoped.BindUserForAuthKey(rawAuthKeyID, sessionID, 0)
|
||
} else {
|
||
r.deps.Sessions.BindUser(sessionID, 0)
|
||
}
|
||
}
|
||
}
|
||
return userID, found, nil
|
||
}
|
||
|
||
func (r *Router) lookupAuthUser(ctx context.Context, authKeyID [8]byte) (int64, bool, error) {
|
||
if userID, found, ok := r.cachedAuthUser(authKeyID); ok {
|
||
return userID, found, nil
|
||
}
|
||
key := string(authKeyID[:])
|
||
v, err, _ := r.authUserSF.Do(key, func() (any, error) {
|
||
if userID, found, ok := r.cachedAuthUser(authKeyID); ok {
|
||
return authUserCacheEntry{userID: userID, found: found}, nil
|
||
}
|
||
userID, found, err := r.deps.Auth.UserID(ctx, authKeyID)
|
||
if err != nil {
|
||
return authUserCacheEntry{}, err
|
||
}
|
||
r.setAuthUserCache(authKeyID, userID, found)
|
||
return authUserCacheEntry{userID: userID, found: found}, nil
|
||
})
|
||
if err != nil {
|
||
return 0, false, err
|
||
}
|
||
entry := v.(authUserCacheEntry)
|
||
return entry.userID, entry.found, nil
|
||
}
|
||
|
||
func (r *Router) cachedAuthUser(authKeyID [8]byte) (int64, bool, bool) {
|
||
r.authUserMu.RLock()
|
||
defer r.authUserMu.RUnlock()
|
||
entry, ok := r.authUsers[authKeyID]
|
||
if !ok {
|
||
return 0, false, false
|
||
}
|
||
return entry.userID, entry.found, true
|
||
}
|
||
|
||
func (r *Router) setAuthUserCache(authKeyID [8]byte, userID int64, found bool) {
|
||
r.authUserMu.Lock()
|
||
defer r.authUserMu.Unlock()
|
||
if r.authUsers == nil {
|
||
r.authUsers = make(map[[8]byte]authUserCacheEntry)
|
||
}
|
||
r.authUsers[authKeyID] = authUserCacheEntry{userID: userID, found: found}
|
||
}
|
||
|
||
func (r *Router) invalidateAuthUserCache(authKeyID [8]byte) {
|
||
r.authUserMu.Lock()
|
||
delete(r.authUsers, authKeyID)
|
||
r.authUserMu.Unlock()
|
||
r.authUserSF.Forget(string(authKeyID[:]))
|
||
}
|
||
|
||
func (r *Router) scopedSessions() (ScopedSessionBinder, bool) {
|
||
if r.deps.Sessions == nil {
|
||
return nil, false
|
||
}
|
||
scoped, ok := r.deps.Sessions.(ScopedSessionBinder)
|
||
return scoped, ok
|
||
}
|
||
|
||
func (r *Router) dispatch(ctx context.Context, b *bin.Buffer, depth int) (bin.Encoder, error) {
|
||
if depth > maxWrapperDepth {
|
||
return nil, wrapperTooDeepErr()
|
||
}
|
||
|
||
id, err := b.PeekID()
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
switch id {
|
||
case tg.InvokeWithLayerRequestTypeID:
|
||
if err := b.ConsumeID(id); err != nil {
|
||
return nil, err
|
||
}
|
||
layer, err := b.Int()
|
||
if err != nil {
|
||
return nil, fmt.Errorf("decode invokeWithLayer layer: %w", err)
|
||
}
|
||
// query 紧跟 layer，buffer 剩余即内层请求。
|
||
return r.dispatch(WithLayer(ctx, layer), b, depth+1)
|
||
|
||
case tg.InvokeWithoutUpdatesRequestTypeID:
|
||
if err := b.ConsumeID(id); err != nil {
|
||
return nil, err
|
||
}
|
||
return r.dispatch(ctx, b, depth+1)
|
||
|
||
case tg.InitConnectionRequestTypeID:
|
||
req := &tg.InitConnectionRequest{Query: &rawObject{}}
|
||
if err := req.Decode(b); err != nil {
|
||
return nil, fmt.Errorf("decode initConnection: %w", err)
|
||
}
|
||
info := ClientInfo{
|
||
APIID: req.APIID,
|
||
DeviceModel: req.DeviceModel,
|
||
SystemVersion: req.SystemVersion,
|
||
AppVersion: req.AppVersion,
|
||
SystemLangCode: req.SystemLangCode,
|
||
LangPack: req.LangPack,
|
||
LangCode: req.LangCode,
|
||
}
|
||
ctx = WithClientInfo(ctx, info)
|
||
r.rememberClientInfo(ctx, info)
|
||
r.log.Debug("initConnection",
|
||
zap.Int("api_id", req.APIID),
|
||
zap.String("device", req.DeviceModel),
|
||
zap.String("app", req.AppVersion),
|
||
zap.Int("layer", LayerFrom(ctx)),
|
||
)
|
||
inner, ok := req.Query.(*rawObject)
|
||
if !ok {
|
||
return nil, fmt.Errorf("initConnection query: unexpected type %T", req.Query)
|
||
}
|
||
return r.dispatch(ctx, &bin.Buffer{Buf: inner.data}, depth+1)
|
||
|
||
default:
|
||
start := time.Now()
|
||
enc, err := r.dispatcher.Handle(ctx, b)
|
||
dur := time.Since(start)
|
||
fields := append([]zap.Field{
|
||
zap.String("method", tlTypeName(id)),
|
||
zap.String("type_id", fmt.Sprintf("%#x", id)),
|
||
zap.Duration("dur", dur),
|
||
}, r.contextLogFields(ctx)...)
|
||
if err != nil || dur > 100*time.Millisecond {
|
||
if err != nil {
|
||
fields = append(fields, zap.Error(err))
|
||
}
|
||
r.log.Info("RPC inner handled", fields...)
|
||
} else {
|
||
r.log.Debug("RPC inner handled", fields...)
|
||
}
|
||
return enc, err
|
||
}
|
||
}
|
||
|
||
func tlTypeName(id uint32) string {
|
||
tlTypeNamesOnce.Do(func() {
|
||
names := tg.NamesMap()
|
||
tlTypeNames = make(map[uint32]string, len(names))
|
||
for name, typeID := range names {
|
||
tlTypeNames[typeID] = name
|
||
}
|
||
})
|
||
if name, ok := tlTypeNames[id]; ok {
|
||
return name
|
||
}
|
||
return fmt.Sprintf("%#x", id)
|
||
}
|
||
|
||
func (r *Router) rememberClientInfo(ctx context.Context, info ClientInfo) {
|
||
rawAuthKeyID, ok := RawAuthKeyIDFrom(ctx)
|
||
if !ok {
|
||
return
|
||
}
|
||
sessionID, ok := SessionIDFrom(ctx)
|
||
if !ok {
|
||
return
|
||
}
|
||
r.clientInfoMu.Lock()
|
||
defer r.clientInfoMu.Unlock()
|
||
if r.clientInfo == nil {
|
||
r.clientInfo = make(map[clientInfoSessionKey]ClientInfo)
|
||
}
|
||
r.clientInfo[clientInfoSessionKey{rawAuthKeyID: rawAuthKeyID, sessionID: sessionID}] = info
|
||
}
|
||
|
||
func (r *Router) clientInfoForSession(ctx context.Context) (ClientInfo, bool) {
|
||
rawAuthKeyID, ok := RawAuthKeyIDFrom(ctx)
|
||
if !ok {
|
||
return ClientInfo{}, false
|
||
}
|
||
sessionID, ok := SessionIDFrom(ctx)
|
||
if !ok {
|
||
return ClientInfo{}, false
|
||
}
|
||
r.clientInfoMu.RLock()
|
||
defer r.clientInfoMu.RUnlock()
|
||
info, ok := r.clientInfo[clientInfoSessionKey{rawAuthKeyID: rawAuthKeyID, sessionID: sessionID}]
|
||
return info, ok
|
||
}
|
||
|
||
// fallback 处理未注册的 RPC：记录到 compatibility trace（落兼容矩阵），
|
||
// 返回 NOT_IMPLEMENTED rpc_error 让客户端继续运行而非断连。
|
||
func (r *Router) fallback(ctx context.Context, b *bin.Buffer) (bin.Encoder, error) {
|
||
id, _ := b.PeekID()
|
||
fields := append([]zap.Field{zap.String("type_id", fmt.Sprintf("%#x", id))}, r.contextLogFields(ctx)...)
|
||
r.log.Warn("Unhandled RPC (compatibility trace)", fields...)
|
||
return nil, notImplementedErr()
|
||
}
|
||
|
||
func (r *Router) contextLogFields(ctx context.Context) []zap.Field {
|
||
fields := []zap.Field{zap.Int("layer", LayerFrom(ctx))}
|
||
if sessionID, ok := SessionIDFrom(ctx); ok {
|
||
fields = append(fields, zap.Int64("session_id", sessionID))
|
||
}
|
||
if rawAuthKeyID, ok := RawAuthKeyIDFrom(ctx); ok {
|
||
fields = append(fields, zap.String("raw_auth_key_id", hex.EncodeToString(rawAuthKeyID[:])))
|
||
}
|
||
if authKeyID, ok := AuthKeyIDFrom(ctx); ok {
|
||
fields = append(fields, zap.String("auth_key_id", hex.EncodeToString(authKeyID[:])))
|
||
}
|
||
if userID, ok := UserIDFrom(ctx); ok {
|
||
fields = append(fields, zap.Int64("user_id", userID))
|
||
}
|
||
return fields
|
||
}
|
||
|
||
// rawObject 在解码 wrapper 时按原样捕获内层 query 的 TL 字节，供递归分发。
|
||
// 它实现 bin.Object（Encode/Decode），但只搬运字节、不解释内容。
|
||
type rawObject struct {
|
||
data []byte
|
||
}
|
||
|
||
func (o *rawObject) Decode(b *bin.Buffer) error {
|
||
o.data = append(o.data[:0], b.Buf...)
|
||
b.Skip(len(b.Buf))
|
||
return nil
|
||
}
|
||
|
||
func (o *rawObject) Encode(b *bin.Buffer) error {
|
||
b.Put(o.data)
|
||
return nil
|
||
}
|