Initial open source release
This commit is contained in:
A
commit
74992e893f
377 changed files with 118084 additions and 0 deletions
416
internal/rpc/router.go
Normal file
416
internal/rpc/router.go
Normal file
|
|
@ -0,0 +1,416 @@
|
|||
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
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue