perf: sync protocol and core hardening updates
This commit is contained in:
parent
152fed3b87
commit
4390ebf5a9
283 changed files with 29231 additions and 2295 deletions
40
.env.example
40
.env.example
|
|
@ -7,6 +7,31 @@ TELESRV_DC=2
|
|||
TELESRV_DEV_AUTH_CODE=12345
|
||||
TELESRV_AUTH_CODE_TTL=5m
|
||||
TELESRV_AUTH_CODE_MAX_ATTEMPTS=5
|
||||
# Unauthenticated login-code issuance uses the same limits for existing and
|
||||
# unknown phones. Phone numbers are SHA-256 digested before becoming Redis keys.
|
||||
TELESRV_AUTH_CODE_PHONE_RATE_LIMIT=5
|
||||
TELESRV_AUTH_CODE_AUTH_KEY_RATE_LIMIT=20
|
||||
TELESRV_AUTH_CODE_RATE_WINDOW=10m
|
||||
|
||||
# MTProto admission and shared inbound RPC budgets. Negative connection/handshake limits disable
|
||||
# that gate; non-positive RPC values fall back to the built-in safe defaults.
|
||||
TELESRV_MTPROTO_MAX_CONNECTIONS=200000
|
||||
TELESRV_MTPROTO_MAX_CONNECTIONS_PER_IP=4096
|
||||
TELESRV_MTPROTO_MAX_CONCURRENT_HANDSHAKES=256
|
||||
TELESRV_MTPROTO_RPC_MAX_INFLIGHT=32
|
||||
TELESRV_MTPROTO_RPC_QUEUE_SIZE=64
|
||||
TELESRV_MTPROTO_RPC_TIMEOUT=30s
|
||||
TELESRV_MTPROTO_RPC_GLOBAL_WORKERS=256
|
||||
TELESRV_MTPROTO_RPC_GLOBAL_MAX_TASKS=8192
|
||||
TELESRV_MTPROTO_RPC_GLOBAL_MAX_BYTES=536870912
|
||||
# Process-wide in-flight transport wire + decrypted plaintext reservation.
|
||||
TELESRV_MTPROTO_INBOUND_FRAME_GLOBAL_MAX_BYTES=536870912
|
||||
# Per-connection outbound mailboxes (normal/control) and process-wide resend pending bodies.
|
||||
TELESRV_MTPROTO_OUTBOUND_QUEUE_SIZE=128
|
||||
TELESRV_MTPROTO_OUTBOUND_CONTROL_QUEUE_SIZE=32
|
||||
TELESRV_MTPROTO_OUTBOUND_TRACKED_GLOBAL_MAX_BYTES=536870912
|
||||
# Concurrent encrypted wire/codec/obfuscation scratch (shared bounded pool, not per connection).
|
||||
TELESRV_MTPROTO_OUTBOUND_WRITE_GLOBAL_MAX_BYTES=536870912
|
||||
|
||||
# Optional login-email verification. When enabled, accounts with a confirmed
|
||||
# login email receive login codes by email; REQUIRE_SETUP also forces new/legacy
|
||||
|
|
@ -56,6 +81,21 @@ TELESRV_REDIS_ADDR=127.0.0.1:6399
|
|||
TELESRV_REDIS_PASSWORD=
|
||||
TELESRV_REDIS_DB=0
|
||||
|
||||
# Bounded retention/GC. User/channel update rows are only pruned behind protocol-safe floors.
|
||||
TELESRV_UPDATE_EVENT_RETENTION=168h
|
||||
TELESRV_BOT_API_UPDATE_RETENTION=24h
|
||||
TELESRV_ORPHAN_AUTH_KEY_RETENTION=24h
|
||||
# Terminal failed outbox heads are kept briefly for diagnosis, then only the online
|
||||
# delivery task is removed. The durable update remains available to getDifference.
|
||||
TELESRV_OUTBOX_POISON_RETENTION=1m
|
||||
TELESRV_OUTBOX_POISON_CLEANUP_INTERVAL=15s
|
||||
TELESRV_RETENTION_INTERVAL=1h
|
||||
TELESRV_RETENTION_BATCH=10000
|
||||
|
||||
# PFS temp->perm binding cache; write-side revoke/rebind invalidates entries precisely.
|
||||
TELESRV_TEMP_KEY_CACHE_MAX_ENTRIES=262144
|
||||
TELESRV_TEMP_KEY_CACHE_TTL=30m
|
||||
|
||||
# Optional. Enables Mapbox-backed map previews and TDesktop map picker config.
|
||||
TELESRV_MAPBOX_TOKEN=
|
||||
TELESRV_MAPTILE_CACHE_DIR=data/maptiles
|
||||
|
|
|
|||
|
|
@ -394,7 +394,7 @@ func run(logger *zap.Logger) error {
|
|||
return fmt.Errorf("seed appearance: %w", err)
|
||||
} else if !stats.Skipped {
|
||||
logger.Info("外观种子导入完成",
|
||||
zap.String("source", "orange-live"),
|
||||
zap.String("source", "default-seed"),
|
||||
zap.Int("wallpapers", stats.Wallpapers),
|
||||
zap.Int("documents", stats.Documents),
|
||||
zap.Int("blobs", stats.Blobs),
|
||||
|
|
@ -434,7 +434,13 @@ func run(logger *zap.Logger) error {
|
|||
cfg.UpdateEventRetention,
|
||||
cfg.RetentionInterval,
|
||||
cfg.RetentionBatch,
|
||||
).WithBotAPIUpdateRetention(botAPIUpdateStore, cfg.BotAPIUpdateRetention).Run(ctx)
|
||||
).WithDispatchOutboxPoisonPolicy(cfg.OutboxPoisonRetention, cfg.OutboxPoisonCleanupInterval).
|
||||
WithBotAPIUpdateRetention(botAPIUpdateStore, cfg.BotAPIUpdateRetention).
|
||||
WithLoginCodeDeliveryRetention(messageStore).
|
||||
WithUserUpdateRetention(updateEventStore).
|
||||
WithChannelUpdateRetention(channelStore).
|
||||
WithOrphanAuthKeyRetention(authKeyStore, activeSessions, cfg.OrphanAuthKeyRetention).
|
||||
Run(ctx)
|
||||
go filesapp.NewUploadPartGCWorker(filesService, logger.Named("files").Named("upload_gc"),
|
||||
cfg.UploadPartTTL,
|
||||
cfg.UploadPartGCInterval,
|
||||
|
|
@ -634,6 +640,7 @@ func run(logger *zap.Logger) error {
|
|||
)
|
||||
authService := auth.NewService(userStore, authzStore, codeStore, authKeyStore, tempAuthKeyStore, cfg.DevAuthCode,
|
||||
auth.WithLoginMessages(messageStore, dialogStore),
|
||||
auth.WithLoginCodeDelivery(messageStore),
|
||||
auth.WithPasswords(passwordStore),
|
||||
auth.WithBotLogin(botStore),
|
||||
auth.WithPremiumGrant(cfg.PremiumGrantMonths),
|
||||
|
|
@ -654,6 +661,9 @@ func run(logger *zap.Logger) error {
|
|||
OutboundPushTimeout: cfg.OutboundPushTimeout,
|
||||
SendRateLimit: cfg.SendRateLimit,
|
||||
SendRateWindow: cfg.SendRateWindow,
|
||||
AuthCodePhoneRateLimit: cfg.AuthCodePhoneRateLimit,
|
||||
AuthCodeAuthKeyRateLimit: cfg.AuthCodeAuthKeyRateLimit,
|
||||
AuthCodeRateWindow: cfg.AuthCodeRateWindow,
|
||||
CatchupRateLimit: cfg.CatchupRateLimit,
|
||||
CatchupRateWindow: cfg.CatchupRateWindow,
|
||||
ChannelNudgeMaxTargets: cfg.ChannelNudgeMaxTargets,
|
||||
|
|
@ -662,9 +672,9 @@ func run(logger *zap.Logger) error {
|
|||
GroupCallMaxParticipants: cfg.GroupCallMaxParticipants,
|
||||
RtmpIngestURL: cfg.LiveStreamRtmpURL,
|
||||
PublicBaseURL: cfg.PublicBaseURL,
|
||||
// PFS temp→perm 解析缓存 5s:削减每帧 ResolveAuthKey 的 PG 查询。显式撤销会清缓存并
|
||||
// 断开连接;re-bind 即时失效(onAuthBindTempAuthKey)。
|
||||
TempKeyResolveCacheTTL: 5 * time.Second,
|
||||
// PFS temp→perm 解析缓存:显式撤销会清缓存并断开连接,re-bind 即时失效;
|
||||
// 配置 TTL 只承担跨进程/异常失效兜底,避免大连接数周期性打满 PG。
|
||||
TempKeyResolveCacheTTL: cfg.TempKeyResolveCacheTTL,
|
||||
TempKeyResolveCacheMaxEntries: cfg.TempKeyResolveCacheMaxEntries,
|
||||
}, rpc.Deps{
|
||||
Auth: authService,
|
||||
|
|
@ -772,16 +782,30 @@ func run(logger *zap.Logger) error {
|
|||
}
|
||||
|
||||
srv := mtprotoedge.New(mtprotoedge.Options{
|
||||
Logger: logger.Named("mtprotoedge"),
|
||||
DC: cfg.DC,
|
||||
RSAKey: rsaKey,
|
||||
RPC: router,
|
||||
AuthKeys: authKeyStore,
|
||||
Sessions: sessionStore,
|
||||
ActiveSessions: activeSessions,
|
||||
ObfuscatedTCP: true,
|
||||
WebSocket: cfg.WebSocketEnable,
|
||||
WebSocketAllowedOrigins: cfg.WebSocketAllowedOrigins,
|
||||
Logger: logger.Named("mtprotoedge"),
|
||||
DC: cfg.DC,
|
||||
RSAKey: rsaKey,
|
||||
RPC: router,
|
||||
AuthKeys: authKeyStore,
|
||||
Sessions: sessionStore,
|
||||
ActiveSessions: activeSessions,
|
||||
ObfuscatedTCP: true,
|
||||
WebSocket: cfg.WebSocketEnable,
|
||||
WebSocketAllowedOrigins: cfg.WebSocketAllowedOrigins,
|
||||
MaxConnections: cfg.MTProtoMaxConnections,
|
||||
MaxConnectionsPerIP: cfg.MTProtoMaxConnectionsPerIP,
|
||||
MaxConcurrentHandshakes: cfg.MTProtoMaxConcurrentHandshakes,
|
||||
RPCMaxInflight: cfg.MTProtoRPCMaxInflight,
|
||||
RPCQueueSize: cfg.MTProtoRPCQueueSize,
|
||||
RPCTimeout: cfg.MTProtoRPCTimeout,
|
||||
RPCGlobalWorkers: cfg.MTProtoRPCGlobalWorkers,
|
||||
RPCGlobalMaxTasks: cfg.MTProtoRPCGlobalMaxTasks,
|
||||
RPCGlobalMaxBytes: cfg.MTProtoRPCGlobalMaxBytes,
|
||||
InboundFrameGlobalMaxBytes: cfg.MTProtoInboundFrameGlobalMaxBytes,
|
||||
OutboundQueueSize: cfg.MTProtoOutboundQueueSize,
|
||||
OutboundControlQueueSize: cfg.MTProtoOutboundControlQueueSize,
|
||||
OutboundTrackedGlobalMaxBytes: cfg.MTProtoOutboundTrackedGlobalMaxBytes,
|
||||
OutboundWriteGlobalMaxBytes: cfg.MTProtoOutboundWriteGlobalMaxBytes,
|
||||
})
|
||||
logger.Info("telesrv 服务就绪",
|
||||
zap.String("listen", cfg.ListenAddr),
|
||||
|
|
|
|||
|
|
@ -0,0 +1,3 @@
|
|||
ALTER TABLE private_messages
|
||||
DROP COLUMN IF EXISTS recipient_delivered,
|
||||
DROP COLUMN IF EXISTS request_fingerprint;
|
||||
25
deploy/migrations/0062_private_message_idempotency.up.sql
Normal file
25
deploy/migrations/0062_private_message_idempotency.up.sql
Normal file
|
|
@ -0,0 +1,25 @@
|
|||
ALTER TABLE private_messages
|
||||
ADD COLUMN request_fingerprint bytea NOT NULL DEFAULT '\x',
|
||||
ADD COLUMN recipient_delivered boolean NOT NULL DEFAULT false;
|
||||
|
||||
-- 旧行没有保存原始请求指纹,保留空 fingerprint,使后续重放显式返回
|
||||
-- RANDOM_ID_DUPLICATE,禁止从可能已编辑的消息投影猜测/修复原请求。
|
||||
-- recipient_delivered 可由同一私聊事实表精确回填,用于区分「被 block 后本就
|
||||
-- 不投递」与「声明已投递但 recipient box 丢失」两种状态。
|
||||
UPDATE private_messages AS p
|
||||
SET recipient_delivered = true
|
||||
WHERE p.sender_user_id <> p.recipient_user_id
|
||||
AND EXISTS (
|
||||
SELECT 1
|
||||
FROM message_boxes AS b
|
||||
WHERE b.private_message_id = p.id
|
||||
AND b.owner_user_id = p.recipient_user_id
|
||||
);
|
||||
|
||||
-- Keep these defaults after the expand step. A pre-0062 process does not name
|
||||
-- either column in INSERT, so dropping them while that process can still serve
|
||||
-- traffic turns an otherwise compatible rolling deployment into a NOT NULL
|
||||
-- failure. The empty fingerprint is an explicit "unknown legacy receipt"
|
||||
-- sentinel: the new replay path rejects it as RANDOM_ID_DUPLICATE before it can
|
||||
-- interpret recipient_delivered. It must never be reconstructed from mutable
|
||||
-- message_boxes.
|
||||
19
deploy/migrations/0063_channel_update_retention.down.sql
Normal file
19
deploy/migrations/0063_channel_update_retention.down.sql
Normal file
|
|
@ -0,0 +1,19 @@
|
|||
-- Once a retained floor advances, channel_update_events below it are physically gone. Dropping
|
||||
-- the checkpoint would make an old client pts look like an ordinary empty difference and silently
|
||||
-- lose the required channelDifferenceTooLong recovery boundary. Refuse that irreversible down.
|
||||
DO $$
|
||||
BEGIN
|
||||
IF EXISTS (
|
||||
SELECT 1
|
||||
FROM public.channel_update_checkpoints
|
||||
WHERE retained_through_pts > 0
|
||||
) THEN
|
||||
RAISE EXCEPTION
|
||||
'cannot roll back channel update retention: retained floor has advanced'
|
||||
USING ERRCODE = '55000';
|
||||
END IF;
|
||||
END
|
||||
$$;
|
||||
|
||||
DROP INDEX IF EXISTS public.channel_update_events_retention_seek_idx;
|
||||
DROP TABLE IF EXISTS public.channel_update_checkpoints;
|
||||
35
deploy/migrations/0063_channel_update_retention.up.sql
Normal file
35
deploy/migrations/0063_channel_update_retention.up.sql
Normal file
|
|
@ -0,0 +1,35 @@
|
|||
-- Channel-scoped update retention checkpoint.
|
||||
--
|
||||
-- channel_update_events may now be pruned in bounded per-channel transactions. The checkpoint is
|
||||
-- the durable protocol boundary: request pts below retained_through_pts must receive
|
||||
-- updates.channelDifferenceTooLong; latest_event_date/latest_pts remain after row deletion so an
|
||||
-- account-level updates.getDifference can still emit UpdateChannelTooLong for an offline member.
|
||||
CREATE TABLE IF NOT EXISTS public.channel_update_checkpoints (
|
||||
channel_id bigint PRIMARY KEY REFERENCES public.channels(id) ON DELETE CASCADE,
|
||||
retained_through_pts integer NOT NULL DEFAULT 0 CHECK (retained_through_pts >= 0),
|
||||
latest_event_date integer NOT NULL DEFAULT 0 CHECK (latest_event_date >= 0),
|
||||
latest_pts integer NOT NULL DEFAULT 0 CHECK (latest_pts >= 0),
|
||||
updated_at timestamp with time zone NOT NULL DEFAULT now(),
|
||||
CONSTRAINT channel_update_checkpoints_floor_check CHECK (retained_through_pts <= latest_pts)
|
||||
);
|
||||
|
||||
-- Backfill every existing channel, including channels without an event row, so read paths can use
|
||||
-- the checkpoint directly without a full event-log fallback.
|
||||
INSERT INTO public.channel_update_checkpoints (
|
||||
channel_id, retained_through_pts, latest_event_date, latest_pts
|
||||
)
|
||||
SELECT c.id,
|
||||
0,
|
||||
COALESCE(MAX(e.date), 0)::integer,
|
||||
c.pts
|
||||
FROM public.channels c
|
||||
LEFT JOIN public.channel_update_events e ON e.channel_id = c.id
|
||||
GROUP BY c.id, c.pts
|
||||
ON CONFLICT (channel_id) DO UPDATE SET
|
||||
latest_event_date = GREATEST(channel_update_checkpoints.latest_event_date, EXCLUDED.latest_event_date),
|
||||
latest_pts = GREATEST(channel_update_checkpoints.latest_pts, EXCLUDED.latest_pts),
|
||||
updated_at = now();
|
||||
|
||||
-- Global retention candidate seek. Exact per-channel deletion continues to use PK(channel_id,pts).
|
||||
CREATE INDEX IF NOT EXISTS channel_update_events_retention_seek_idx
|
||||
ON public.channel_update_events (date, channel_id, pts);
|
||||
20
deploy/migrations/0064_user_update_retention.down.sql
Normal file
20
deploy/migrations/0064_user_update_retention.down.sql
Normal file
|
|
@ -0,0 +1,20 @@
|
|||
-- TDesktop has no account-level differenceTooLong fallback. If any confirmed prefix has already
|
||||
-- been deleted, removing this floor/observed state would turn a durable history hole into a false
|
||||
-- empty difference. A rollback is safe only before retention has advanced anywhere.
|
||||
DO $$
|
||||
BEGIN
|
||||
IF EXISTS (
|
||||
SELECT 1
|
||||
FROM public.user_update_retention
|
||||
WHERE retained_through_pts > 0
|
||||
) THEN
|
||||
RAISE EXCEPTION
|
||||
'cannot roll back user update retention: retained floor has advanced'
|
||||
USING ERRCODE = '55000';
|
||||
END IF;
|
||||
END
|
||||
$$;
|
||||
|
||||
DROP INDEX IF EXISTS public.user_update_events_retention_global_idx;
|
||||
DROP TABLE IF EXISTS public.user_update_retention;
|
||||
ALTER TABLE public.update_states DROP COLUMN IF EXISTS observed_pts;
|
||||
22
deploy/migrations/0064_user_update_retention.up.sql
Normal file
22
deploy/migrations/0064_user_update_retention.up.sql
Normal file
|
|
@ -0,0 +1,22 @@
|
|||
-- 服务端已发送/构造的 update_states.pts 不是客户端确认:difference 响应可能在网络中
|
||||
-- 丢失。observed_pts 只由客户端后续请求实际带回的 pts(或 getState 显式建立的快照
|
||||
-- baseline)推进,retention 只能使用这个水位。
|
||||
ALTER TABLE public.update_states
|
||||
ADD COLUMN observed_pts integer NOT NULL DEFAULT 0 CHECK (observed_pts >= 0);
|
||||
|
||||
-- 账号级 durable update 安全前缀回收水位。
|
||||
-- 这里只记录“所有当前授权设备都已明确确认”的连续前缀;TDesktop 不支持
|
||||
-- updates.differenceTooLong,故该表绝不能用于任意 TTL 硬裁剪。
|
||||
CREATE TABLE public.user_update_retention (
|
||||
user_id bigint PRIMARY KEY REFERENCES public.users(id) ON DELETE CASCADE,
|
||||
retained_through_pts integer NOT NULL DEFAULT 0 CHECK (retained_through_pts >= 0),
|
||||
retained_through_date integer NOT NULL DEFAULT 0 CHECK (retained_through_date >= 0),
|
||||
updated_at timestamp with time zone NOT NULL DEFAULT now()
|
||||
);
|
||||
|
||||
CREATE INDEX user_update_retention_updated_idx
|
||||
ON public.user_update_retention (updated_at, user_id);
|
||||
|
||||
-- retention 先按时间挑全局最老的安全候选,再按 (user_id, pts) 删除连续前缀。
|
||||
CREATE INDEX user_update_events_retention_global_idx
|
||||
ON public.user_update_events (date, user_id, pts);
|
||||
|
|
@ -0,0 +1 @@
|
|||
DROP INDEX IF EXISTS public.auth_keys_orphan_retention_idx;
|
||||
4
deploy/migrations/0065_auth_key_orphan_retention.up.sql
Normal file
4
deploy/migrations/0065_auth_key_orphan_retention.up.sql
Normal file
|
|
@ -0,0 +1,4 @@
|
|||
-- Bounded orphan auth-key GC seeks by creation time. Authorization and temp-key references are
|
||||
-- rechecked in the DELETE statement; active raw keys are supplied by the connection registry.
|
||||
CREATE INDEX IF NOT EXISTS auth_keys_orphan_retention_idx
|
||||
ON public.auth_keys (created_at, auth_key_id);
|
||||
|
|
@ -0,0 +1,2 @@
|
|||
DROP INDEX IF EXISTS dispatch_outbox_logical_shard_head_idx;
|
||||
DROP INDEX IF EXISTS dispatch_outbox_user_pts_uidx;
|
||||
39
deploy/migrations/0066_dispatch_outbox_user_lanes.up.sql
Normal file
39
deploy/migrations/0066_dispatch_outbox_user_lanes.up.sql
Normal file
|
|
@ -0,0 +1,39 @@
|
|||
-- 一条 durable user update 只能有一个在线投递任务。历史 schema 的
|
||||
-- ON CONFLICT DO NOTHING 没有对应唯一键,先显式保留最早任务并删除重复项。
|
||||
WITH duplicates AS (
|
||||
SELECT id
|
||||
FROM (
|
||||
SELECT
|
||||
id,
|
||||
row_number() OVER (
|
||||
PARTITION BY target_user_id, pts
|
||||
-- 若历史重复中仍有可投递任务,优先保留它;不能让一个更早的
|
||||
-- failed 副本覆盖健康 pending 副本并人为阻塞该用户 lane。
|
||||
ORDER BY
|
||||
CASE status
|
||||
WHEN 'pending' THEN 0
|
||||
WHEN 'dispatching' THEN 1
|
||||
ELSE 2
|
||||
END,
|
||||
id ASC
|
||||
) AS rn
|
||||
FROM dispatch_outbox
|
||||
) ranked
|
||||
WHERE ranked.rn > 1
|
||||
)
|
||||
DELETE FROM dispatch_outbox d
|
||||
USING duplicates x
|
||||
WHERE d.id = x.id;
|
||||
|
||||
CREATE UNIQUE INDEX dispatch_outbox_user_pts_uidx
|
||||
ON dispatch_outbox (target_user_id, pts);
|
||||
|
||||
-- ClaimDispatchOutboxShards 固定用 256 logical shards;表达式必须与查询
|
||||
-- 完全一致,避免每个 worker 为筛自己的 lane 全表扫描。
|
||||
CREATE INDEX dispatch_outbox_logical_shard_head_idx
|
||||
ON dispatch_outbox (
|
||||
mod(target_user_id, 256::bigint),
|
||||
target_user_id,
|
||||
pts,
|
||||
id
|
||||
);
|
||||
7
deploy/migrations/0067_auth_key_last_used.down.sql
Normal file
7
deploy/migrations/0067_auth_key_last_used.down.sql
Normal file
|
|
@ -0,0 +1,7 @@
|
|||
DROP INDEX IF EXISTS public.temp_auth_key_bindings_perm_idx;
|
||||
DROP INDEX IF EXISTS public.auth_keys_orphan_last_used_idx;
|
||||
ALTER TABLE public.auth_keys DROP COLUMN IF EXISTS last_used_at;
|
||||
|
||||
-- Restore the 0065 schema used by the pre-last_used orphan collector.
|
||||
CREATE INDEX IF NOT EXISTS auth_keys_orphan_retention_idx
|
||||
ON public.auth_keys (created_at, auth_key_id);
|
||||
17
deploy/migrations/0067_auth_key_last_used.up.sql
Normal file
17
deploy/migrations/0067_auth_key_last_used.up.sql
Normal file
|
|
@ -0,0 +1,17 @@
|
|||
-- A physical connection touches last_used_at atomically while loading its key. Orphan GC uses
|
||||
-- this watermark in addition to the in-memory active-key snapshot, closing the Get->Register
|
||||
-- race without turning the per-frame encrypted fast path into a database write.
|
||||
ALTER TABLE public.auth_keys
|
||||
ADD COLUMN IF NOT EXISTS last_used_at timestamptz NOT NULL DEFAULT now();
|
||||
|
||||
CREATE INDEX IF NOT EXISTS auth_keys_orphan_last_used_idx
|
||||
ON public.auth_keys (last_used_at, auth_key_id);
|
||||
|
||||
-- DeleteOrphaned now seeks exclusively by last_used_at. Keeping the transitional 0065
|
||||
-- created_at index would duplicate auth-key insert/delete maintenance without serving a query.
|
||||
DROP INDEX IF EXISTS public.auth_keys_orphan_retention_idx;
|
||||
|
||||
-- Delete/revoke and orphan-retention probes both resolve perm->temp bindings. The original
|
||||
-- primary key only covers temp_auth_key_id, so the reverse predicate otherwise scans the table.
|
||||
CREATE INDEX IF NOT EXISTS temp_auth_key_bindings_perm_idx
|
||||
ON public.temp_auth_key_bindings (perm_auth_key_id);
|
||||
|
|
@ -0,0 +1,5 @@
|
|||
ALTER TABLE public.private_messages
|
||||
DROP COLUMN IF EXISTS recipient_pts,
|
||||
DROP COLUMN IF EXISTS recipient_box_id,
|
||||
DROP COLUMN IF EXISTS sender_pts,
|
||||
DROP COLUMN IF EXISTS sender_box_id;
|
||||
15
deploy/migrations/0068_private_message_send_receipt.up.sql
Normal file
15
deploy/migrations/0068_private_message_send_receipt.up.sql
Normal file
|
|
@ -0,0 +1,15 @@
|
|||
-- Exact random_id replay must not rebuild the original send result from mutable message_boxes:
|
||||
-- edit advances box pts/body and delete hides the row. Keep the immutable allocation receipt on
|
||||
-- the shared private message instead; 0073 adds the immutable first snapshot/delete receipt used
|
||||
-- to build the client acknowledgement without allocating new pts or update facts.
|
||||
ALTER TABLE public.private_messages
|
||||
ADD COLUMN sender_box_id integer NOT NULL DEFAULT 0 CHECK (sender_box_id >= 0),
|
||||
ADD COLUMN sender_pts integer NOT NULL DEFAULT 0 CHECK (sender_pts >= 0),
|
||||
ADD COLUMN recipient_box_id integer NOT NULL DEFAULT 0 CHECK (recipient_box_id >= 0),
|
||||
ADD COLUMN recipient_pts integer NOT NULL DEFAULT 0 CHECK (recipient_pts >= 0);
|
||||
|
||||
-- Keep zero defaults for writers released before this migration. Zero is not
|
||||
-- a valid immutable receipt: when a row has a valid request fingerprint but a
|
||||
-- legacy writer omitted these columns, the new replay path fails fast instead
|
||||
-- of deriving the first response from mutable message_boxes. New writers
|
||||
-- always persist positive sender receipt values before commit.
|
||||
12
deploy/migrations/0069_dispatch_outbox_user_heads.down.sql
Normal file
12
deploy/migrations/0069_dispatch_outbox_user_heads.down.sql
Normal file
|
|
@ -0,0 +1,12 @@
|
|||
CREATE INDEX IF NOT EXISTS dispatch_outbox_logical_shard_head_idx
|
||||
ON dispatch_outbox (
|
||||
mod(target_user_id, 256::bigint),
|
||||
target_user_id,
|
||||
pts,
|
||||
id
|
||||
);
|
||||
|
||||
DROP TRIGGER IF EXISTS dispatch_outbox_delete_user_head ON dispatch_outbox;
|
||||
DROP TRIGGER IF EXISTS dispatch_outbox_insert_user_head ON dispatch_outbox;
|
||||
DROP FUNCTION IF EXISTS dispatch_outbox_maintain_user_head();
|
||||
DROP TABLE IF EXISTS dispatch_outbox_user_heads;
|
||||
80
deploy/migrations/0069_dispatch_outbox_user_heads.up.sql
Normal file
80
deploy/migrations/0069_dispatch_outbox_user_heads.up.sql
Normal file
|
|
@ -0,0 +1,80 @@
|
|||
-- Claim 只需要查看每个用户当前未完成 head。把 head 持久化后,领取复杂度由
|
||||
-- “扫描全部 outbox 积压并 DISTINCT ON”降为“扫描有积压的用户 lane”。
|
||||
-- 迁移期间阻止并发写,保证 backfill 与随后安装的触发器之间没有缺口。
|
||||
LOCK TABLE dispatch_outbox IN SHARE ROW EXCLUSIVE MODE;
|
||||
|
||||
CREATE TABLE dispatch_outbox_user_heads (
|
||||
target_user_id bigint PRIMARY KEY,
|
||||
head_id bigint NOT NULL,
|
||||
head_pts integer NOT NULL CHECK (head_pts >= 0),
|
||||
logical_shard smallint GENERATED ALWAYS AS (
|
||||
mod(target_user_id, 256::bigint)::smallint
|
||||
) STORED,
|
||||
CHECK (logical_shard >= 0 AND logical_shard < 256)
|
||||
);
|
||||
|
||||
CREATE INDEX dispatch_outbox_user_heads_shard_idx
|
||||
ON dispatch_outbox_user_heads (logical_shard, target_user_id);
|
||||
|
||||
INSERT INTO dispatch_outbox_user_heads (target_user_id, head_id, head_pts)
|
||||
SELECT DISTINCT ON (target_user_id)
|
||||
target_user_id,
|
||||
id,
|
||||
pts
|
||||
FROM dispatch_outbox
|
||||
ORDER BY target_user_id ASC, pts ASC, id ASC;
|
||||
|
||||
CREATE FUNCTION dispatch_outbox_maintain_user_head()
|
||||
RETURNS trigger
|
||||
LANGUAGE plpgsql
|
||||
AS $$
|
||||
DECLARE
|
||||
removed_head bigint;
|
||||
BEGIN
|
||||
IF TG_OP = 'INSERT' THEN
|
||||
INSERT INTO dispatch_outbox_user_heads (target_user_id, head_id, head_pts)
|
||||
VALUES (NEW.target_user_id, NEW.id, NEW.pts)
|
||||
ON CONFLICT (target_user_id) DO UPDATE
|
||||
SET head_id = EXCLUDED.head_id,
|
||||
head_pts = EXCLUDED.head_pts
|
||||
WHERE (EXCLUDED.head_pts, EXCLUDED.head_id) <
|
||||
(dispatch_outbox_user_heads.head_pts, dispatch_outbox_user_heads.head_id);
|
||||
RETURN NULL;
|
||||
END IF;
|
||||
|
||||
-- 删除非 head 不需要重算。删除 head 时用 (target_user_id, pts, id)
|
||||
-- 索引找下一条;failed head 同样会一直阻塞,直到被显式删除。
|
||||
DELETE FROM dispatch_outbox_user_heads
|
||||
WHERE target_user_id = OLD.target_user_id
|
||||
AND head_id = OLD.id
|
||||
RETURNING head_id INTO removed_head;
|
||||
|
||||
IF removed_head IS NOT NULL THEN
|
||||
INSERT INTO dispatch_outbox_user_heads (target_user_id, head_id, head_pts)
|
||||
SELECT target_user_id, id, pts
|
||||
FROM dispatch_outbox
|
||||
WHERE target_user_id = OLD.target_user_id
|
||||
ORDER BY pts ASC, id ASC
|
||||
LIMIT 1
|
||||
ON CONFLICT (target_user_id) DO UPDATE
|
||||
SET head_id = EXCLUDED.head_id,
|
||||
head_pts = EXCLUDED.head_pts
|
||||
WHERE (EXCLUDED.head_pts, EXCLUDED.head_id) <
|
||||
(dispatch_outbox_user_heads.head_pts, dispatch_outbox_user_heads.head_id);
|
||||
END IF;
|
||||
RETURN NULL;
|
||||
END;
|
||||
$$;
|
||||
|
||||
CREATE TRIGGER dispatch_outbox_insert_user_head
|
||||
AFTER INSERT ON dispatch_outbox
|
||||
FOR EACH ROW
|
||||
EXECUTE FUNCTION dispatch_outbox_maintain_user_head();
|
||||
|
||||
CREATE TRIGGER dispatch_outbox_delete_user_head
|
||||
AFTER DELETE ON dispatch_outbox
|
||||
FOR EACH ROW
|
||||
EXECUTE FUNCTION dispatch_outbox_maintain_user_head();
|
||||
|
||||
-- Claim 已不再读取表达式 shard 索引;移除它避免每次 enqueue 的重复写放大。
|
||||
DROP INDEX IF EXISTS dispatch_outbox_logical_shard_head_idx;
|
||||
|
|
@ -0,0 +1,46 @@
|
|||
DROP TRIGGER IF EXISTS dispatch_outbox_update_user_head ON dispatch_outbox;
|
||||
|
||||
CREATE OR REPLACE FUNCTION dispatch_outbox_maintain_user_head()
|
||||
RETURNS trigger
|
||||
LANGUAGE plpgsql
|
||||
AS $$
|
||||
DECLARE
|
||||
removed_head bigint;
|
||||
BEGIN
|
||||
IF TG_OP = 'INSERT' THEN
|
||||
INSERT INTO dispatch_outbox_user_heads (target_user_id, head_id, head_pts)
|
||||
VALUES (NEW.target_user_id, NEW.id, NEW.pts)
|
||||
ON CONFLICT (target_user_id) DO UPDATE
|
||||
SET head_id = EXCLUDED.head_id,
|
||||
head_pts = EXCLUDED.head_pts
|
||||
WHERE (EXCLUDED.head_pts, EXCLUDED.head_id) <
|
||||
(dispatch_outbox_user_heads.head_pts, dispatch_outbox_user_heads.head_id);
|
||||
RETURN NULL;
|
||||
END IF;
|
||||
|
||||
DELETE FROM dispatch_outbox_user_heads
|
||||
WHERE target_user_id = OLD.target_user_id
|
||||
AND head_id = OLD.id
|
||||
RETURNING head_id INTO removed_head;
|
||||
|
||||
IF removed_head IS NOT NULL THEN
|
||||
INSERT INTO dispatch_outbox_user_heads (target_user_id, head_id, head_pts)
|
||||
SELECT target_user_id, id, pts
|
||||
FROM dispatch_outbox
|
||||
WHERE target_user_id = OLD.target_user_id
|
||||
ORDER BY pts ASC, id ASC
|
||||
LIMIT 1
|
||||
ON CONFLICT (target_user_id) DO UPDATE
|
||||
SET head_id = EXCLUDED.head_id,
|
||||
head_pts = EXCLUDED.head_pts
|
||||
WHERE (EXCLUDED.head_pts, EXCLUDED.head_id) <
|
||||
(dispatch_outbox_user_heads.head_pts, dispatch_outbox_user_heads.head_id);
|
||||
END IF;
|
||||
RETURN NULL;
|
||||
END;
|
||||
$$;
|
||||
|
||||
ALTER TABLE dispatch_outbox_user_heads
|
||||
DROP COLUMN updated_at,
|
||||
DROP COLUMN next_attempt_at,
|
||||
DROP COLUMN status;
|
||||
82
deploy/migrations/0070_dispatch_outbox_head_readiness.up.sql
Normal file
82
deploy/migrations/0070_dispatch_outbox_head_readiness.up.sql
Normal file
|
|
@ -0,0 +1,82 @@
|
|||
-- Keep claim readiness on the one-row-per-user head itself. Otherwise PostgreSQL may
|
||||
-- legally reorder the head/outbox join and start from every eligible backlog row,
|
||||
-- reintroducing the full-backlog scan that 0069 is meant to remove.
|
||||
LOCK TABLE dispatch_outbox IN SHARE ROW EXCLUSIVE MODE;
|
||||
|
||||
ALTER TABLE dispatch_outbox_user_heads
|
||||
ADD COLUMN status varchar(16) NOT NULL DEFAULT 'pending'
|
||||
CHECK (status IN ('pending', 'dispatching', 'failed')),
|
||||
ADD COLUMN next_attempt_at timestamptz NOT NULL DEFAULT now(),
|
||||
ADD COLUMN updated_at timestamptz NOT NULL DEFAULT now();
|
||||
|
||||
UPDATE dispatch_outbox_user_heads h
|
||||
SET status = d.status,
|
||||
next_attempt_at = d.next_attempt_at,
|
||||
updated_at = d.updated_at
|
||||
FROM dispatch_outbox d
|
||||
WHERE d.target_user_id = h.target_user_id
|
||||
AND d.id = h.head_id;
|
||||
|
||||
CREATE OR REPLACE FUNCTION dispatch_outbox_maintain_user_head()
|
||||
RETURNS trigger
|
||||
LANGUAGE plpgsql
|
||||
AS $$
|
||||
DECLARE
|
||||
removed_head bigint;
|
||||
BEGIN
|
||||
IF TG_OP = 'INSERT' THEN
|
||||
INSERT INTO dispatch_outbox_user_heads (
|
||||
target_user_id, head_id, head_pts, status, next_attempt_at, updated_at
|
||||
) VALUES (
|
||||
NEW.target_user_id, NEW.id, NEW.pts, NEW.status, NEW.next_attempt_at, NEW.updated_at
|
||||
)
|
||||
ON CONFLICT (target_user_id) DO UPDATE
|
||||
SET head_id = EXCLUDED.head_id,
|
||||
head_pts = EXCLUDED.head_pts,
|
||||
status = EXCLUDED.status,
|
||||
next_attempt_at = EXCLUDED.next_attempt_at,
|
||||
updated_at = EXCLUDED.updated_at
|
||||
WHERE (EXCLUDED.head_pts, EXCLUDED.head_id) <
|
||||
(dispatch_outbox_user_heads.head_pts, dispatch_outbox_user_heads.head_id);
|
||||
RETURN NULL;
|
||||
ELSIF TG_OP = 'UPDATE' THEN
|
||||
UPDATE dispatch_outbox_user_heads
|
||||
SET status = NEW.status,
|
||||
next_attempt_at = NEW.next_attempt_at,
|
||||
updated_at = NEW.updated_at
|
||||
WHERE target_user_id = NEW.target_user_id
|
||||
AND head_id = NEW.id;
|
||||
RETURN NULL;
|
||||
END IF;
|
||||
|
||||
DELETE FROM dispatch_outbox_user_heads
|
||||
WHERE target_user_id = OLD.target_user_id
|
||||
AND head_id = OLD.id
|
||||
RETURNING head_id INTO removed_head;
|
||||
|
||||
IF removed_head IS NOT NULL THEN
|
||||
INSERT INTO dispatch_outbox_user_heads (
|
||||
target_user_id, head_id, head_pts, status, next_attempt_at, updated_at
|
||||
)
|
||||
SELECT target_user_id, id, pts, status, next_attempt_at, updated_at
|
||||
FROM dispatch_outbox
|
||||
WHERE target_user_id = OLD.target_user_id
|
||||
ORDER BY pts ASC, id ASC
|
||||
LIMIT 1
|
||||
ON CONFLICT (target_user_id) DO UPDATE
|
||||
SET head_id = EXCLUDED.head_id,
|
||||
head_pts = EXCLUDED.head_pts,
|
||||
status = EXCLUDED.status,
|
||||
next_attempt_at = EXCLUDED.next_attempt_at,
|
||||
updated_at = EXCLUDED.updated_at
|
||||
WHERE (EXCLUDED.head_pts, EXCLUDED.head_id) <
|
||||
(dispatch_outbox_user_heads.head_pts, dispatch_outbox_user_heads.head_id);
|
||||
END IF;
|
||||
RETURN NULL;
|
||||
END;
|
||||
$$;
|
||||
|
||||
CREATE TRIGGER dispatch_outbox_update_user_head
|
||||
AFTER UPDATE OF status, next_attempt_at, updated_at ON dispatch_outbox
|
||||
FOR EACH ROW
|
||||
EXECUTE FUNCTION dispatch_outbox_maintain_user_head();
|
||||
|
|
@ -0,0 +1 @@
|
|||
DROP INDEX IF EXISTS public.temp_auth_key_bindings_expiry_idx;
|
||||
5
deploy/migrations/0071_temp_auth_key_expiry_seek.up.sql
Normal file
5
deploy/migrations/0071_temp_auth_key_expiry_seek.up.sql
Normal file
|
|
@ -0,0 +1,5 @@
|
|||
-- DeleteExpiredTempAuthKeys orders/filters the binding table by expiry and then deletes the parent
|
||||
-- auth_keys rows in a bounded batch. The PK starts with temp_auth_key_id, so it cannot serve this
|
||||
-- maintenance seek.
|
||||
CREATE INDEX IF NOT EXISTS temp_auth_key_bindings_expiry_idx
|
||||
ON public.temp_auth_key_bindings (expires_at, temp_auth_key_id);
|
||||
|
|
@ -0,0 +1,16 @@
|
|||
CREATE INDEX IF NOT EXISTS dispatch_outbox_pending_ready_idx
|
||||
ON dispatch_outbox (next_attempt_at, target_user_id, pts, id)
|
||||
WHERE status = 'pending';
|
||||
|
||||
CREATE INDEX IF NOT EXISTS dispatch_outbox_dispatching_stale_ready_idx
|
||||
ON dispatch_outbox (updated_at, target_user_id, pts, id)
|
||||
WHERE status = 'dispatching';
|
||||
|
||||
CREATE INDEX IF NOT EXISTS dispatch_outbox_user_heads_shard_idx
|
||||
ON dispatch_outbox_user_heads (logical_shard, target_user_id);
|
||||
|
||||
ALTER TABLE dispatch_outbox_user_heads
|
||||
DROP CONSTRAINT IF EXISTS dispatch_outbox_user_heads_outbox_fkey;
|
||||
|
||||
DROP INDEX IF EXISTS dispatch_outbox_user_heads_dispatching_shard_idx;
|
||||
DROP INDEX IF EXISTS dispatch_outbox_user_heads_pending_shard_idx;
|
||||
|
|
@ -0,0 +1,35 @@
|
|||
-- Claim only eligible durable heads. Partial indexes keep work proportional to ready
|
||||
-- user lanes instead of every account that has a deferred/failed backlog.
|
||||
CREATE INDEX dispatch_outbox_user_heads_pending_shard_idx
|
||||
ON dispatch_outbox_user_heads (
|
||||
logical_shard,
|
||||
next_attempt_at,
|
||||
target_user_id,
|
||||
head_pts,
|
||||
head_id
|
||||
)
|
||||
WHERE status = 'pending';
|
||||
|
||||
CREATE INDEX dispatch_outbox_user_heads_dispatching_shard_idx
|
||||
ON dispatch_outbox_user_heads (
|
||||
logical_shard,
|
||||
updated_at,
|
||||
target_user_id,
|
||||
head_pts,
|
||||
head_id
|
||||
)
|
||||
WHERE status = 'dispatching';
|
||||
|
||||
-- A head must always reference the exact outbox row it represents. The base schema already has
|
||||
-- dispatch_outbox_target_user_id_id_key on these columns; reuse it instead of maintaining a
|
||||
-- duplicate unique index on every enqueue/delete. Deferred validation lets the AFTER DELETE
|
||||
-- trigger promote/delete the head in the same tx.
|
||||
ALTER TABLE dispatch_outbox_user_heads
|
||||
ADD CONSTRAINT dispatch_outbox_user_heads_outbox_fkey
|
||||
FOREIGN KEY (target_user_id, head_id)
|
||||
REFERENCES dispatch_outbox (target_user_id, id)
|
||||
DEFERRABLE INITIALLY DEFERRED;
|
||||
|
||||
DROP INDEX IF EXISTS dispatch_outbox_user_heads_shard_idx;
|
||||
DROP INDEX IF EXISTS dispatch_outbox_pending_ready_idx;
|
||||
DROP INDEX IF EXISTS dispatch_outbox_dispatching_stale_ready_idx;
|
||||
17
deploy/migrations/0073_send_replay_snapshots.down.sql
Normal file
17
deploy/migrations/0073_send_replay_snapshots.down.sql
Normal file
|
|
@ -0,0 +1,17 @@
|
|||
ALTER TABLE public.channel_messages
|
||||
DROP CONSTRAINT IF EXISTS channel_messages_delete_ids_array,
|
||||
DROP CONSTRAINT IF EXISTS channel_messages_send_snapshot_object,
|
||||
DROP COLUMN IF EXISTS delete_message_ids,
|
||||
DROP COLUMN IF EXISTS delete_date,
|
||||
DROP COLUMN IF EXISTS delete_pts_count,
|
||||
DROP COLUMN IF EXISTS delete_pts,
|
||||
DROP COLUMN IF EXISTS send_snapshot;
|
||||
|
||||
ALTER TABLE public.private_messages
|
||||
DROP CONSTRAINT IF EXISTS private_messages_sender_delete_ids_array,
|
||||
DROP CONSTRAINT IF EXISTS private_messages_sender_snapshot_object,
|
||||
DROP COLUMN IF EXISTS sender_delete_message_ids,
|
||||
DROP COLUMN IF EXISTS sender_delete_date,
|
||||
DROP COLUMN IF EXISTS sender_delete_pts_count,
|
||||
DROP COLUMN IF EXISTS sender_delete_pts,
|
||||
DROP COLUMN IF EXISTS sender_snapshot;
|
||||
23
deploy/migrations/0073_send_replay_snapshots.up.sql
Normal file
23
deploy/migrations/0073_send_replay_snapshots.up.sql
Normal file
|
|
@ -0,0 +1,23 @@
|
|||
-- Lost-response random_id replay must be able to acknowledge the original send
|
||||
-- after mutable message rows were edited or deleted. Keep the first sender echo
|
||||
-- and the exact durable delete event alongside the idempotency key.
|
||||
ALTER TABLE public.private_messages
|
||||
ADD COLUMN sender_snapshot jsonb NOT NULL DEFAULT '{}'::jsonb,
|
||||
ADD COLUMN sender_delete_pts integer NOT NULL DEFAULT 0 CHECK (sender_delete_pts >= 0),
|
||||
ADD COLUMN sender_delete_pts_count integer NOT NULL DEFAULT 0 CHECK (sender_delete_pts_count >= 0),
|
||||
ADD COLUMN sender_delete_date integer NOT NULL DEFAULT 0 CHECK (sender_delete_date >= 0),
|
||||
ADD COLUMN sender_delete_message_ids jsonb NOT NULL DEFAULT '[]'::jsonb,
|
||||
ADD CONSTRAINT private_messages_sender_snapshot_object CHECK (jsonb_typeof(sender_snapshot) IS NOT DISTINCT FROM 'object'),
|
||||
ADD CONSTRAINT private_messages_sender_delete_ids_array CHECK (jsonb_typeof(sender_delete_message_ids) IS NOT DISTINCT FROM 'array');
|
||||
|
||||
ALTER TABLE public.channel_messages
|
||||
ADD COLUMN send_snapshot jsonb NOT NULL DEFAULT '{}'::jsonb,
|
||||
ADD COLUMN delete_pts integer NOT NULL DEFAULT 0 CHECK (delete_pts >= 0),
|
||||
ADD COLUMN delete_pts_count integer NOT NULL DEFAULT 0 CHECK (delete_pts_count >= 0),
|
||||
ADD COLUMN delete_date integer NOT NULL DEFAULT 0 CHECK (delete_date >= 0),
|
||||
ADD COLUMN delete_message_ids jsonb NOT NULL DEFAULT '[]'::jsonb,
|
||||
ADD CONSTRAINT channel_messages_send_snapshot_object CHECK (jsonb_typeof(send_snapshot) IS NOT DISTINCT FROM 'object'),
|
||||
ADD CONSTRAINT channel_messages_delete_ids_array CHECK (jsonb_typeof(delete_message_ids) IS NOT DISTINCT FROM 'array');
|
||||
|
||||
-- Existing rows predate immutable replay snapshots. Leaving them as {} makes a
|
||||
-- replay fail fast instead of guessing the first response from mutable state.
|
||||
|
|
@ -0,0 +1 @@
|
|||
DROP INDEX IF EXISTS public.dispatch_outbox_user_heads_failed_cleanup_idx;
|
||||
|
|
@ -0,0 +1,10 @@
|
|||
-- Terminal failed rows are short-lived diagnostic quarantine entries. Cleanup starts from the
|
||||
-- one-row-per-user durable head so it can lock in the same head→outbox order as claim/completion
|
||||
-- without scanning healthy lanes.
|
||||
CREATE INDEX dispatch_outbox_user_heads_failed_cleanup_idx
|
||||
ON public.dispatch_outbox_user_heads (
|
||||
updated_at,
|
||||
target_user_id,
|
||||
head_id
|
||||
)
|
||||
WHERE status = 'failed';
|
||||
1
deploy/migrations/0075_uploaded_media_receipts.down.sql
Normal file
1
deploy/migrations/0075_uploaded_media_receipts.down.sql
Normal file
|
|
@ -0,0 +1 @@
|
|||
DROP TABLE IF EXISTS public.uploaded_media_receipts;
|
||||
12
deploy/migrations/0075_uploaded_media_receipts.up.sql
Normal file
12
deploy/migrations/0075_uploaded_media_receipts.up.sql
Normal file
|
|
@ -0,0 +1,12 @@
|
|||
-- upload.saveFilePart data is transient, while messages.sendMedia may be replayed after its first
|
||||
-- response is lost. Preserve the immutable Photo/Document materialization selected for each
|
||||
-- (owner,file_id) so InputMediaUploaded* remains idempotent after part cleanup.
|
||||
CREATE TABLE public.uploaded_media_receipts (
|
||||
owner_user_id bigint NOT NULL REFERENCES public.users(id) ON DELETE CASCADE,
|
||||
file_id bigint NOT NULL CHECK (file_id <> 0),
|
||||
intent_hash bytea NOT NULL CHECK (octet_length(intent_hash) = 32),
|
||||
media_kind text NOT NULL CHECK (media_kind IN ('photo', 'document')),
|
||||
media_id bigint NOT NULL CHECK (media_id <> 0),
|
||||
created_at timestamptz NOT NULL DEFAULT now(),
|
||||
PRIMARY KEY (owner_user_id, file_id)
|
||||
);
|
||||
1
deploy/migrations/0076_album_group_reservations.down.sql
Normal file
1
deploy/migrations/0076_album_group_reservations.down.sql
Normal file
|
|
@ -0,0 +1 @@
|
|||
DROP TABLE IF EXISTS album_group_reservations;
|
||||
17
deploy/migrations/0076_album_group_reservations.up.sql
Normal file
17
deploy/migrations/0076_album_group_reservations.up.sql
Normal file
|
|
@ -0,0 +1,17 @@
|
|||
-- sendMultiMedia 必须在解析上传媒体、逐条发送之前持久预留 grouped_id。
|
||||
-- 这张表把每个发送 random_id 固定到会话作用域内的相册组,使中途失败后
|
||||
-- 客户端只重试失败子集时仍能恢复首次整包使用的 grouped_id;intent_hash
|
||||
-- 同时阻止内容已改变的相同 random_id 借旧预留错误合并。
|
||||
CREATE TABLE album_group_reservations (
|
||||
sender_user_id BIGINT NOT NULL REFERENCES users(id) ON DELETE CASCADE,
|
||||
peer_type TEXT NOT NULL CHECK (peer_type IN ('user', 'channel')),
|
||||
peer_id BIGINT NOT NULL CHECK (peer_id > 0),
|
||||
random_id BIGINT NOT NULL CHECK (random_id <> 0),
|
||||
intent_hash BYTEA NOT NULL CHECK (octet_length(intent_hash) = 32),
|
||||
grouped_id BIGINT NOT NULL CHECK (grouped_id <> 0),
|
||||
created_at TIMESTAMPTZ NOT NULL DEFAULT now(),
|
||||
PRIMARY KEY (sender_user_id, peer_type, peer_id, random_id)
|
||||
);
|
||||
|
||||
COMMENT ON TABLE album_group_reservations IS
|
||||
'Durable pre-send binding from album item random_id to grouped_id; never reconstructed from a retry subset.';
|
||||
|
|
@ -0,0 +1,8 @@
|
|||
-- Contract rollback: only safe after every pre-0062/0068 writer has drained.
|
||||
ALTER TABLE public.private_messages
|
||||
ALTER COLUMN request_fingerprint DROP DEFAULT,
|
||||
ALTER COLUMN recipient_delivered DROP DEFAULT,
|
||||
ALTER COLUMN sender_box_id DROP DEFAULT,
|
||||
ALTER COLUMN sender_pts DROP DEFAULT,
|
||||
ALTER COLUMN recipient_box_id DROP DEFAULT,
|
||||
ALTER COLUMN recipient_pts DROP DEFAULT;
|
||||
17
deploy/migrations/0077_correct_private_send_defaults.up.sql
Normal file
17
deploy/migrations/0077_correct_private_send_defaults.up.sql
Normal file
|
|
@ -0,0 +1,17 @@
|
|||
-- 0062 and 0068 originally removed their expand-phase defaults immediately.
|
||||
-- Databases that already applied those revisions therefore reject INSERTs from
|
||||
-- an older telesrv process during a rolling deployment. Restore the permanent
|
||||
-- legacy-writer sentinels; fresh databases receive the same defaults directly
|
||||
-- from the corrected original migrations.
|
||||
ALTER TABLE public.private_messages
|
||||
ALTER COLUMN request_fingerprint SET DEFAULT '\x'::bytea,
|
||||
ALTER COLUMN recipient_delivered SET DEFAULT false,
|
||||
ALTER COLUMN sender_box_id SET DEFAULT 0,
|
||||
ALTER COLUMN sender_pts SET DEFAULT 0,
|
||||
ALTER COLUMN recipient_box_id SET DEFAULT 0,
|
||||
ALTER COLUMN recipient_pts SET DEFAULT 0;
|
||||
|
||||
-- Empty fingerprint and zero receipt fields mean "unknown legacy send". They
|
||||
-- are deliberately not backfilled: the immutable first response cannot be
|
||||
-- reconstructed from message_boxes after edit/delete. Replay must reject or
|
||||
-- fail fast through the store invariants instead of normalizing these values.
|
||||
3
deploy/migrations/0078_channel_send_fingerprint.down.sql
Normal file
3
deploy/migrations/0078_channel_send_fingerprint.down.sql
Normal file
|
|
@ -0,0 +1,3 @@
|
|||
ALTER TABLE public.channel_messages
|
||||
DROP CONSTRAINT IF EXISTS channel_messages_request_fingerprint_size,
|
||||
DROP COLUMN IF EXISTS request_fingerprint;
|
||||
16
deploy/migrations/0078_channel_send_fingerprint.up.sql
Normal file
16
deploy/migrations/0078_channel_send_fingerprint.up.sql
Normal file
|
|
@ -0,0 +1,16 @@
|
|||
-- Persist the immutable client intent beside a channel random_id receipt.
|
||||
--
|
||||
-- The empty default is intentional for rolling deploys: binaries that predate
|
||||
-- this migration can continue to insert channel/service messages. Empty
|
||||
-- fingerprints are legacy/unknown receipts and the replay path rejects them;
|
||||
-- it never guesses intent from an editable message projection.
|
||||
ALTER TABLE public.channel_messages
|
||||
ADD COLUMN request_fingerprint bytea NOT NULL DEFAULT '\x';
|
||||
|
||||
ALTER TABLE public.channel_messages
|
||||
ADD CONSTRAINT channel_messages_request_fingerprint_size
|
||||
CHECK (octet_length(request_fingerprint) IN (0, 32)) NOT VALID;
|
||||
|
||||
-- NOT VALID avoids a blocking historical-table validation scan during the
|
||||
-- rolling migration while PostgreSQL still enforces the check for every new
|
||||
-- or updated row. A later maintenance window may validate it online.
|
||||
|
|
@ -0,0 +1 @@
|
|||
DROP TABLE IF EXISTS login_code_message_deliveries;
|
||||
22
deploy/migrations/0079_login_code_message_deliveries.up.sql
Normal file
22
deploy/migrations/0079_login_code_message_deliveries.up.sql
Normal file
|
|
@ -0,0 +1,22 @@
|
|||
-- A phone_code_hash identifies exactly one account-visible 777000 login-code
|
||||
-- notification. Store only its SHA-256 digest plus compact immutable allocation
|
||||
-- facts; the secret code body remains solely in private_messages/message_boxes.
|
||||
CREATE TABLE login_code_message_deliveries (
|
||||
delivery_key bytea PRIMARY KEY,
|
||||
code_fingerprint bytea NOT NULL,
|
||||
user_id bigint NOT NULL REFERENCES users(id) ON DELETE CASCADE,
|
||||
private_message_id bigint NOT NULL CHECK (private_message_id > 0),
|
||||
message_box_id integer NOT NULL CHECK (message_box_id > 0),
|
||||
pts integer NOT NULL CHECK (pts > 0),
|
||||
message_date integer NOT NULL CHECK (message_date > 0),
|
||||
created_at timestamptz NOT NULL DEFAULT now(),
|
||||
CONSTRAINT login_code_message_deliveries_key_size CHECK (octet_length(delivery_key) = 32),
|
||||
CONSTRAINT login_code_message_deliveries_fingerprint_size CHECK (octet_length(code_fingerprint) = 32),
|
||||
CONSTRAINT login_code_message_deliveries_user_box_unique UNIQUE (user_id, message_box_id),
|
||||
CONSTRAINT login_code_message_deliveries_user_pts_unique UNIQUE (user_id, pts)
|
||||
);
|
||||
|
||||
COMMENT ON COLUMN login_code_message_deliveries.delivery_key IS
|
||||
'SHA-256(phone_code_hash); raw phone_code_hash is never persisted';
|
||||
COMMENT ON COLUMN login_code_message_deliveries.code_fingerprint IS
|
||||
'HMAC-SHA-256(code), keyed by the non-persisted raw phone_code_hash';
|
||||
|
|
@ -0,0 +1,4 @@
|
|||
DROP INDEX IF EXISTS login_code_message_deliveries_expiry_idx;
|
||||
|
||||
ALTER TABLE login_code_message_deliveries
|
||||
DROP COLUMN IF EXISTS expires_at;
|
||||
15
deploy/migrations/0080_login_code_delivery_expiry.up.sql
Normal file
15
deploy/migrations/0080_login_code_delivery_expiry.up.sql
Normal file
|
|
@ -0,0 +1,15 @@
|
|||
-- Compact login-code idempotency receipts are only needed while the opaque
|
||||
-- code can still be used/replayed. Keep them seek-prunable instead of growing
|
||||
-- one row per login attempt forever.
|
||||
ALTER TABLE login_code_message_deliveries
|
||||
ADD COLUMN expires_at timestamptz;
|
||||
|
||||
UPDATE login_code_message_deliveries
|
||||
SET expires_at = created_at + interval '24 hours'
|
||||
WHERE expires_at IS NULL;
|
||||
|
||||
ALTER TABLE login_code_message_deliveries
|
||||
ALTER COLUMN expires_at SET NOT NULL;
|
||||
|
||||
CREATE INDEX login_code_message_deliveries_expiry_idx
|
||||
ON login_code_message_deliveries (expires_at, delivery_key);
|
||||
|
|
@ -3,6 +3,8 @@ package account
|
|||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
|
|
@ -32,6 +34,88 @@ type captureMailSender struct {
|
|||
code string
|
||||
}
|
||||
|
||||
type blockingCodeCAS struct {
|
||||
store.CodeStore
|
||||
mu sync.Mutex
|
||||
blockRevision string
|
||||
blockUpdate bool
|
||||
blockDelete bool
|
||||
entered chan struct{}
|
||||
release chan struct{}
|
||||
once sync.Once
|
||||
}
|
||||
|
||||
type switchableEmailOwnerStore struct {
|
||||
store.UserStore
|
||||
mu sync.RWMutex
|
||||
phone string
|
||||
override bool
|
||||
owner domain.User
|
||||
found bool
|
||||
}
|
||||
|
||||
func (s *switchableEmailOwnerStore) ByPhone(ctx context.Context, phone string) (domain.User, bool, error) {
|
||||
s.mu.RLock()
|
||||
if s.override && domain.NormalizePhone(phone) == s.phone {
|
||||
owner, found := s.owner, s.found
|
||||
s.mu.RUnlock()
|
||||
return owner, found, nil
|
||||
}
|
||||
s.mu.RUnlock()
|
||||
return s.UserStore.ByPhone(ctx, phone)
|
||||
}
|
||||
|
||||
func (s *switchableEmailOwnerStore) switchOwner(phone string, owner domain.User) {
|
||||
s.mu.Lock()
|
||||
s.phone = domain.NormalizePhone(phone)
|
||||
s.owner = owner
|
||||
s.found = true
|
||||
s.override = true
|
||||
s.mu.Unlock()
|
||||
}
|
||||
|
||||
type afterSavePasswordStore struct {
|
||||
store.PasswordStore
|
||||
once sync.Once
|
||||
afterSave func(userID int64, settings domain.PasswordSettings)
|
||||
}
|
||||
|
||||
func (s *afterSavePasswordStore) Save(ctx context.Context, userID int64, settings domain.PasswordSettings) error {
|
||||
if err := s.PasswordStore.Save(ctx, userID, settings); err != nil {
|
||||
return err
|
||||
}
|
||||
if s.afterSave != nil {
|
||||
s.once.Do(func() { s.afterSave(userID, settings) })
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *blockingCodeCAS) shouldBlock(revision string, update bool) bool {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
return revision == s.blockRevision && ((update && s.blockUpdate) || (!update && s.blockDelete))
|
||||
}
|
||||
|
||||
func (s *blockingCodeCAS) waitIfBlocked(revision string, update bool) {
|
||||
if !s.shouldBlock(revision, update) {
|
||||
return
|
||||
}
|
||||
s.once.Do(func() {
|
||||
close(s.entered)
|
||||
<-s.release
|
||||
})
|
||||
}
|
||||
|
||||
func (s *blockingCodeCAS) CompareAndUpdate(ctx context.Context, key, revision string, next store.PhoneCode) (bool, error) {
|
||||
s.waitIfBlocked(revision, true)
|
||||
return s.CodeStore.CompareAndUpdate(ctx, key, revision, next)
|
||||
}
|
||||
|
||||
func (s *blockingCodeCAS) CompareAndDelete(ctx context.Context, key, revision string) (bool, error) {
|
||||
s.waitIfBlocked(revision, false)
|
||||
return s.CodeStore.CompareAndDelete(ctx, key, revision)
|
||||
}
|
||||
|
||||
func (s *captureMailSender) SendLoginCode(_ context.Context, to, code string, _ time.Duration) error {
|
||||
s.to = to
|
||||
s.code = code
|
||||
|
|
@ -66,22 +150,22 @@ func TestSetLoginEmailPersistsAndMasks(t *testing.T) {
|
|||
}
|
||||
}
|
||||
|
||||
// TestLoginEmailByPhoneAndClear 验证按手机号读取/清除登录邮箱(sendCode 检测 + reset 用)。
|
||||
// TestLoginEmailByPhoneAndClear 验证 sendCode 可按手机号读取,但 reset 只按已锁定 userID 清除。
|
||||
func TestLoginEmailByPhoneAndClear(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
svc, users := newLoginEmailService(t)
|
||||
createUser(t, users, "15550010002")
|
||||
u := createUser(t, users, "15550010002")
|
||||
|
||||
if err := svc.SetLoginEmailByPhone(ctx, "+1 555 001 0002", "bob@mail.com"); err != nil {
|
||||
t.Fatalf("SetLoginEmailByPhone: %v", err)
|
||||
if err := svc.SetLoginEmail(ctx, u.ID, "bob@mail.com"); err != nil {
|
||||
t.Fatalf("SetLoginEmail: %v", err)
|
||||
}
|
||||
email, found, err := svc.LoginEmailByPhone(ctx, "15550010002")
|
||||
if err != nil || !found || email != "bob@mail.com" {
|
||||
t.Fatalf("LoginEmailByPhone = %q found=%v err=%v", email, found, err)
|
||||
}
|
||||
|
||||
if err := svc.ClearLoginEmailByPhone(ctx, "15550010002"); err != nil {
|
||||
t.Fatalf("ClearLoginEmailByPhone: %v", err)
|
||||
if err := svc.ClearLoginEmail(ctx, u.ID); err != nil {
|
||||
t.Fatalf("ClearLoginEmail: %v", err)
|
||||
}
|
||||
if _, found, _ := svc.LoginEmailByPhone(ctx, "15550010002"); found {
|
||||
t.Fatal("login email still present after clear")
|
||||
|
|
@ -182,7 +266,7 @@ func TestLoginEmailSetupRejectsAlreadyOwnedEmailForNewPhone(t *testing.T) {
|
|||
if err := svc.SetLoginEmail(ctx, owner.ID, "owner@example.test"); err != nil {
|
||||
t.Fatalf("SetLoginEmail owner: %v", err)
|
||||
}
|
||||
if err := codes.Set(ctx, "new-phone-hash", store.PhoneCode{Phone: "15550010108", Channel: "email_setup_required", MaxAttempts: 2}, time.Minute); err != nil {
|
||||
if err := codes.Set(ctx, "new-phone-hash", store.PhoneCode{Version: store.PhoneCodeVersionCurrent, Phone: "15550010108", Channel: "email_setup_required", MaxAttempts: 2}, time.Minute); err != nil {
|
||||
t.Fatalf("seed phone code: %v", err)
|
||||
}
|
||||
|
||||
|
|
@ -234,7 +318,7 @@ func TestLoginEmailSetupStoresPendingEmailOnPhoneCodeHash(t *testing.T) {
|
|||
svc := NewService(memory.NewPasswordStore(),
|
||||
WithUsers(memory.NewUserStore()),
|
||||
WithLoginEmailVerification(codes, sender, time.Minute, 2, 6))
|
||||
if err := codes.Set(ctx, "phone-hash", store.PhoneCode{Phone: "15550010006", Channel: "email_setup_required", MaxAttempts: 2}, time.Minute); err != nil {
|
||||
if err := codes.Set(ctx, "phone-hash", store.PhoneCode{Version: store.PhoneCodeVersionCurrent, Phone: "15550010006", Channel: "email_setup_required", MaxAttempts: 2}, time.Minute); err != nil {
|
||||
t.Fatalf("seed phone code: %v", err)
|
||||
}
|
||||
|
||||
|
|
@ -252,7 +336,7 @@ func TestLoginEmailSetupStoresPendingEmailOnPhoneCodeHash(t *testing.T) {
|
|||
if err != nil || !found {
|
||||
t.Fatalf("phone code found=%v err=%v", found, err)
|
||||
}
|
||||
if rec.Channel != "email_login" || rec.Code != sender.code || rec.Email != "new@example.test" || !rec.VerifiedEmail || rec.PendingEmail != "new@example.test" {
|
||||
if rec.Channel != "email_login" || rec.Code != sender.code || rec.Email != "new@example.test" || !rec.VerifiedEmail || !rec.SignUpVerified || rec.PendingEmail != "new@example.test" {
|
||||
t.Fatalf("phone code after verify = %+v", rec)
|
||||
}
|
||||
}
|
||||
|
|
@ -291,3 +375,184 @@ func TestVerifyLoginEmailDeletesCodeAfterMaxAttempts(t *testing.T) {
|
|||
t.Fatal("login email was set after exhausted verification code")
|
||||
}
|
||||
}
|
||||
|
||||
func TestStaleLoginEmailVerificationCannotMutateResentCode(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
blockUpdate bool
|
||||
blockDelete bool
|
||||
verificationCode func(old string) string
|
||||
}{
|
||||
{
|
||||
name: "wrong-code-update",
|
||||
blockUpdate: true,
|
||||
verificationCode: func(old string) string {
|
||||
if old != "000000" {
|
||||
return "000000"
|
||||
}
|
||||
return "111111"
|
||||
},
|
||||
},
|
||||
{
|
||||
name: "correct-code-delete",
|
||||
blockDelete: true,
|
||||
verificationCode: func(old string) string { return old },
|
||||
},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
users := memory.NewUserStore()
|
||||
baseCodes := memory.NewCodeStore()
|
||||
codes := &blockingCodeCAS{CodeStore: baseCodes}
|
||||
passwords := memory.NewPasswordStore()
|
||||
sender := &captureMailSender{}
|
||||
svc := NewService(passwords,
|
||||
WithUsers(users),
|
||||
WithLoginEmailVerification(codes, sender, time.Minute, 3, 6))
|
||||
u := createUser(t, users, "155500102"+fmt.Sprint(10+len(tc.name)))
|
||||
if _, _, err := svc.SendLoginEmailCode(ctx, u.ID, "", "", "cas@example.test", false); err != nil {
|
||||
t.Fatalf("first SendLoginEmailCode: %v", err)
|
||||
}
|
||||
oldCode := sender.code
|
||||
key := loginEmailVerifyChangePrefix + fmt.Sprint(u.ID)
|
||||
oldSnapshot, found, err := baseCodes.GetSnapshot(ctx, key)
|
||||
if err != nil || !found {
|
||||
t.Fatalf("old snapshot found=%v err=%v", found, err)
|
||||
}
|
||||
codes.blockRevision = oldSnapshot.Revision
|
||||
codes.blockUpdate = tc.blockUpdate
|
||||
codes.blockDelete = tc.blockDelete
|
||||
codes.entered = make(chan struct{})
|
||||
codes.release = make(chan struct{})
|
||||
|
||||
verifyErr := make(chan error, 1)
|
||||
go func() {
|
||||
_, err := svc.VerifyLoginEmail(ctx, u.ID, "", "", tc.verificationCode(oldCode), false)
|
||||
verifyErr <- err
|
||||
}()
|
||||
<-codes.entered
|
||||
for attempts := 0; attempts < 5; attempts++ {
|
||||
if _, _, err := svc.SendLoginEmailCode(ctx, u.ID, "", "", "cas@example.test", false); err != nil {
|
||||
t.Fatalf("resent SendLoginEmailCode: %v", err)
|
||||
}
|
||||
if sender.code != oldCode {
|
||||
break
|
||||
}
|
||||
}
|
||||
newCode := sender.code
|
||||
if newCode == oldCode {
|
||||
t.Fatal("random resend repeatedly produced the old code")
|
||||
}
|
||||
newSnapshot, found, err := baseCodes.GetSnapshot(ctx, key)
|
||||
if err != nil || !found || newSnapshot.Revision == oldSnapshot.Revision {
|
||||
t.Fatalf("new snapshot=%+v found=%v err=%v", newSnapshot, found, err)
|
||||
}
|
||||
close(codes.release)
|
||||
if err := <-verifyErr; !errors.Is(err, domain.ErrEmailCodeInvalid) {
|
||||
t.Fatalf("stale verification err=%v, want ErrEmailCodeInvalid", err)
|
||||
}
|
||||
current, found, err := baseCodes.GetSnapshot(ctx, key)
|
||||
if err != nil || !found || current.Revision != newSnapshot.Revision || current.Record.Code != newCode || current.Record.Attempts != 0 {
|
||||
t.Fatalf("current code after stale verifier=%+v found=%v err=%v", current, found, err)
|
||||
}
|
||||
if _, err := svc.VerifyLoginEmail(ctx, u.ID, "", "", newCode, false); err != nil {
|
||||
t.Fatalf("VerifyLoginEmail new code: %v", err)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestConcurrentWrongLoginEmailCodesNeverAuthorize(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
users := memory.NewUserStore()
|
||||
codes := memory.NewCodeStore()
|
||||
passwords := memory.NewPasswordStore()
|
||||
sender := &captureMailSender{}
|
||||
svc := NewService(passwords,
|
||||
WithUsers(users),
|
||||
WithLoginEmailVerification(codes, sender, time.Minute, 2, 6))
|
||||
u := createUser(t, users, "15550010231")
|
||||
if _, _, err := svc.SendLoginEmailCode(ctx, u.ID, "", "", "wrong@example.test", false); err != nil {
|
||||
t.Fatalf("SendLoginEmailCode: %v", err)
|
||||
}
|
||||
correct := sender.code
|
||||
wrong := "000000"
|
||||
if wrong == correct {
|
||||
wrong = "111111"
|
||||
}
|
||||
const workers = 32
|
||||
start := make(chan struct{})
|
||||
errs := make(chan error, workers)
|
||||
for i := 0; i < workers; i++ {
|
||||
go func() {
|
||||
<-start
|
||||
_, err := svc.VerifyLoginEmail(ctx, u.ID, "", "", wrong, false)
|
||||
errs <- err
|
||||
}()
|
||||
}
|
||||
close(start)
|
||||
for i := 0; i < workers; i++ {
|
||||
if err := <-errs; !errors.Is(err, domain.ErrEmailCodeInvalid) {
|
||||
t.Fatalf("wrong concurrent verification err=%v", err)
|
||||
}
|
||||
}
|
||||
if _, found, err := svc.LoginEmail(ctx, u.ID); err != nil || found {
|
||||
t.Fatalf("LoginEmail after wrong codes found=%v err=%v, want absent", found, err)
|
||||
}
|
||||
key := loginEmailVerifyChangePrefix + fmt.Sprint(u.ID)
|
||||
for attempts := 0; attempts < 3; attempts++ {
|
||||
if _, found, err := codes.GetSnapshot(ctx, key); err != nil {
|
||||
t.Fatalf("GetSnapshot after concurrent attempts: %v", err)
|
||||
} else if !found {
|
||||
break
|
||||
}
|
||||
if _, err := svc.VerifyLoginEmail(ctx, u.ID, "", "", wrong, false); !errors.Is(err, domain.ErrEmailCodeInvalid) {
|
||||
t.Fatalf("final wrong verification err=%v", err)
|
||||
}
|
||||
}
|
||||
if _, err := svc.VerifyLoginEmail(ctx, u.ID, "", "", correct, false); !errors.Is(err, domain.ErrEmailCodeInvalid) {
|
||||
t.Fatalf("correct code after exhausted attempts err=%v, want invalid", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEmailSetupOwnerTransferDuringSaveNeverWritesFactorToNewOwner(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
baseUsers := memory.NewUserStore()
|
||||
ownerA := createUser(t, baseUsers, "15550010241")
|
||||
ownerB := createUser(t, baseUsers, "15550010242")
|
||||
users := &switchableEmailOwnerStore{UserStore: baseUsers}
|
||||
basePasswords := memory.NewPasswordStore()
|
||||
passwords := &afterSavePasswordStore{PasswordStore: basePasswords}
|
||||
codes := memory.NewCodeStore()
|
||||
sender := &captureMailSender{}
|
||||
svc := NewService(passwords,
|
||||
WithUsers(users),
|
||||
WithLoginEmailVerification(codes, sender, time.Minute, 3, 6))
|
||||
hash := "owner-save-race"
|
||||
if err := codes.Set(ctx, hash, store.PhoneCode{
|
||||
Version: store.PhoneCodeVersionCurrent,
|
||||
IssuedUserID: ownerA.ID,
|
||||
Phone: ownerA.Phone,
|
||||
Channel: codeChannelEmailSetupRequired,
|
||||
MaxAttempts: 3,
|
||||
}, time.Minute); err != nil {
|
||||
t.Fatalf("seed phone code: %v", err)
|
||||
}
|
||||
if _, _, err := svc.SendLoginEmailCode(ctx, 0, ownerA.Phone, hash, "owner-a@example.test", true); err != nil {
|
||||
t.Fatalf("SendLoginEmailCode: %v", err)
|
||||
}
|
||||
passwords.afterSave = func(userID int64, settings domain.PasswordSettings) {
|
||||
if userID == ownerA.ID && settings.LoginEmail == "owner-a@example.test" {
|
||||
users.switchOwner(ownerA.Phone, ownerB)
|
||||
}
|
||||
}
|
||||
if _, err := svc.VerifyLoginEmail(ctx, 0, ownerA.Phone, hash, sender.code, true); !errors.Is(err, domain.ErrEmailCodeInvalid) {
|
||||
t.Fatalf("VerifyLoginEmail across save-time owner transfer err=%v, want invalid", err)
|
||||
}
|
||||
if settings, found, err := basePasswords.GetByUser(ctx, ownerB.ID); err != nil || (found && settings.LoginEmail != "") {
|
||||
t.Fatalf("new owner settings=%+v found=%v err=%v, SMTP factor leaked to B", settings, found, err)
|
||||
}
|
||||
if _, found, err := codes.GetSnapshot(ctx, hash); err != nil || found {
|
||||
t.Fatalf("owner-drift phone hash found=%v err=%v, want invalidated", found, err)
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -3,7 +3,6 @@ package account
|
|||
import (
|
||||
"context"
|
||||
"crypto/rand"
|
||||
"crypto/subtle"
|
||||
"encoding/hex"
|
||||
"fmt"
|
||||
"strings"
|
||||
|
|
@ -51,9 +50,10 @@ func (s *Service) SendChangePhoneCode(ctx context.Context, userID int64, authKey
|
|||
return "", domain.AuthCodeDelivery{}, err
|
||||
}
|
||||
rec := store.PhoneCode{
|
||||
Version: store.PhoneCodeVersionCurrent,
|
||||
Phone: phone,
|
||||
Code: s.phoneChangeCode,
|
||||
Channel: "phone",
|
||||
Channel: store.PhoneCodeChannelPhone,
|
||||
Purpose: store.PhoneCodePurposeChangePhone,
|
||||
UserID: userID,
|
||||
AuthKeyID: authKeyID,
|
||||
|
|
@ -68,7 +68,7 @@ func (s *Service) SendChangePhoneCode(ctx context.Context, userID int64, authKey
|
|||
|
||||
// ChangePhone 验证作用域和验证码后执行原子改号。返回事件用于当前 session 的
|
||||
// pts 簿记;其它 session 由 transactional outbox 投递 updateUserPhone。
|
||||
func (s *Service) ChangePhone(ctx context.Context, userID int64, authKeyID [8]byte, sessionID int64, phone, phoneCodeHash, code string, date int) (domain.PhoneChangeResult, error) {
|
||||
func (s *Service) ChangePhone(ctx context.Context, userID int64, authKeyID, originRawAuthKeyID [8]byte, sessionID int64, phone, phoneCodeHash, code string, date int) (domain.PhoneChangeResult, error) {
|
||||
if strings.TrimSpace(phoneCodeHash) == "" || strings.TrimSpace(code) == "" {
|
||||
return domain.PhoneChangeResult{}, domain.ErrPhoneCodeEmpty
|
||||
}
|
||||
|
|
@ -82,51 +82,45 @@ func (s *Service) ChangePhone(ctx context.Context, userID int64, authKeyID [8]by
|
|||
if s.codes == nil || s.phoneChanges == nil {
|
||||
return domain.PhoneChangeResult{}, fmt.Errorf("phone change service is not configured")
|
||||
}
|
||||
rec, found, err := s.codes.Get(ctx, phoneCodeHash)
|
||||
scope := store.PhoneCodeScope{
|
||||
Purpose: store.PhoneCodePurposeChangePhone,
|
||||
UserID: userID,
|
||||
AuthKeyID: authKeyID,
|
||||
Phone: phone,
|
||||
}
|
||||
verified, err := s.codes.VerifyScoped(ctx, phoneCodeHash, scope, strings.TrimSpace(code), s.phoneChangeMaxAttempts)
|
||||
if err != nil {
|
||||
return domain.PhoneChangeResult{}, err
|
||||
}
|
||||
if !found {
|
||||
switch verified.Status {
|
||||
case store.LoginCodeVerifyMissing:
|
||||
return domain.PhoneChangeResult{}, domain.ErrPhoneCodeExpired
|
||||
}
|
||||
if rec.Purpose != store.PhoneCodePurposeChangePhone || rec.Phone != phone || rec.UserID != userID || rec.AuthKeyID != authKeyID {
|
||||
case store.LoginCodeVerifyInvalid:
|
||||
return domain.PhoneChangeResult{}, domain.ErrPhoneCodeInvalid
|
||||
case store.LoginCodeVerifyAccepted:
|
||||
default:
|
||||
return domain.PhoneChangeResult{}, domain.ErrPhoneCodeInvalid
|
||||
}
|
||||
code = strings.TrimSpace(code)
|
||||
if subtle.ConstantTimeCompare([]byte(rec.Code), []byte(code)) != 1 {
|
||||
return domain.PhoneChangeResult{}, s.rejectPhoneChangeCode(ctx, phoneCodeHash, rec)
|
||||
}
|
||||
if existing, occupied, err := s.users.ByPhone(ctx, phone); err != nil {
|
||||
return domain.PhoneChangeResult{}, err
|
||||
} else if occupied && existing.ID != userID {
|
||||
return domain.PhoneChangeResult{}, domain.ErrPhoneNumberOccupied
|
||||
}
|
||||
// 正确 code 必须在进入持久化事务前原子消费。并发重放中只有一个请求能
|
||||
// 获得记录,其余请求不得再次推进 pts 或追加 user_phone event。
|
||||
consumed, found, err := s.codes.ConsumeScoped(ctx, phoneCodeHash, store.PhoneCodeScope{
|
||||
Purpose: store.PhoneCodePurposeChangePhone,
|
||||
UserID: userID,
|
||||
AuthKeyID: authKeyID,
|
||||
Phone: phone,
|
||||
})
|
||||
if err != nil {
|
||||
return domain.PhoneChangeResult{}, err
|
||||
}
|
||||
if !found {
|
||||
return domain.PhoneChangeResult{}, domain.ErrPhoneCodeExpired
|
||||
}
|
||||
if consumed.Purpose != rec.Purpose || consumed.UserID != rec.UserID || consumed.AuthKeyID != rec.AuthKeyID || consumed.Phone != rec.Phone ||
|
||||
subtle.ConstantTimeCompare([]byte(consumed.Code), []byte(code)) != 1 {
|
||||
consumed := verified.Record
|
||||
if consumed.Version != store.PhoneCodeVersionCurrent || consumed.Scope() != scope || consumed.Channel != store.PhoneCodeChannelPhone {
|
||||
return domain.PhoneChangeResult{}, domain.ErrPhoneCodeInvalid
|
||||
}
|
||||
if date == 0 {
|
||||
date = int(time.Now().Unix())
|
||||
}
|
||||
result, err := s.phoneChanges.ChangePhone(ctx, domain.PhoneChangeRequest{
|
||||
UserID: userID,
|
||||
Phone: phone,
|
||||
Date: date,
|
||||
ExcludeAuthKeyID: authKeyID,
|
||||
UserID: userID,
|
||||
Phone: phone,
|
||||
Date: date,
|
||||
// Authorization/code scope is the stable business (perm) key, while dispatch exclusion
|
||||
// must use the physical raw key. They differ on PFS/temp connections; conflating them
|
||||
// echoes updateUserPhone back to the initiating device and suppresses the wrong session.
|
||||
ExcludeAuthKeyID: originRawAuthKeyID,
|
||||
ExcludeSessionID: sessionID,
|
||||
})
|
||||
if err != nil {
|
||||
|
|
@ -162,20 +156,6 @@ func (s *Service) phoneChangeCaller(ctx context.Context, userID int64, authKeyID
|
|||
return u, nil
|
||||
}
|
||||
|
||||
func (s *Service) rejectPhoneChangeCode(ctx context.Context, hash string, rec store.PhoneCode) error {
|
||||
rec.Attempts++
|
||||
max := rec.MaxAttempts
|
||||
if max <= 0 {
|
||||
max = s.phoneChangeMaxAttempts
|
||||
}
|
||||
if max > 0 && rec.Attempts >= max {
|
||||
_ = s.codes.Del(ctx, hash)
|
||||
return domain.ErrPhoneCodeInvalid
|
||||
}
|
||||
_ = s.codes.Update(ctx, hash, rec)
|
||||
return domain.ErrPhoneCodeInvalid
|
||||
}
|
||||
|
||||
func phoneChangeHash() (string, error) {
|
||||
var raw [8]byte
|
||||
if _, err := rand.Read(raw[:]); err != nil {
|
||||
|
|
|
|||
|
|
@ -21,6 +21,26 @@ type phoneChangeFixture struct {
|
|||
events *memory.UpdateEventStore
|
||||
user domain.User
|
||||
authKeyID [8]byte
|
||||
changes *recordingPhoneChangeStore
|
||||
}
|
||||
|
||||
type recordingPhoneChangeStore struct {
|
||||
mu sync.Mutex
|
||||
inner store.PhoneChangeStore
|
||||
last domain.PhoneChangeRequest
|
||||
}
|
||||
|
||||
func (s *recordingPhoneChangeStore) ChangePhone(ctx context.Context, req domain.PhoneChangeRequest) (domain.PhoneChangeResult, error) {
|
||||
s.mu.Lock()
|
||||
s.last = req
|
||||
s.mu.Unlock()
|
||||
return s.inner.ChangePhone(ctx, req)
|
||||
}
|
||||
|
||||
func (s *recordingPhoneChangeStore) lastRequest() domain.PhoneChangeRequest {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
return s.last
|
||||
}
|
||||
|
||||
func newPhoneChangeFixture(t *testing.T) phoneChangeFixture {
|
||||
|
|
@ -38,12 +58,13 @@ func newPhoneChangeFixture(t *testing.T) phoneChangeFixture {
|
|||
if err := auths.Bind(ctx, domain.Authorization{AuthKeyID: authKeyID, UserID: u.ID, CreatedAt: time.Now().Add(-48 * time.Hour)}); err != nil {
|
||||
t.Fatalf("bind auth: %v", err)
|
||||
}
|
||||
changes := &recordingPhoneChangeStore{inner: memory.NewPhoneChangeStore(users, events)}
|
||||
service := NewService(
|
||||
memory.NewPasswordStore(),
|
||||
WithUsers(users),
|
||||
WithPhoneChange(memory.NewPhoneChangeStore(users, events), auths, codes, nil, "12345", time.Minute, 3),
|
||||
WithPhoneChange(changes, auths, codes, nil, "12345", time.Minute, 3),
|
||||
)
|
||||
return phoneChangeFixture{ctx: ctx, service: service, users: users, auths: auths, codes: codes, events: events, user: u, authKeyID: authKeyID}
|
||||
return phoneChangeFixture{ctx: ctx, service: service, users: users, auths: auths, codes: codes, events: events, user: u, authKeyID: authKeyID, changes: changes}
|
||||
}
|
||||
|
||||
func TestPhoneChangeScopesCodeAndPersistsDurableEvent(t *testing.T) {
|
||||
|
|
@ -59,17 +80,21 @@ func TestPhoneChangeScopesCodeAndPersistsDurableEvent(t *testing.T) {
|
|||
if err != nil || !found {
|
||||
t.Fatalf("load code found=%v err=%v", found, err)
|
||||
}
|
||||
if rec.Purpose != store.PhoneCodePurposeChangePhone || rec.Phone != "15550012002" || rec.UserID != f.user.ID || rec.AuthKeyID != f.authKeyID || rec.SessionID != 77 {
|
||||
if rec.Version != store.PhoneCodeVersionCurrent || rec.Purpose != store.PhoneCodePurposeChangePhone || rec.Phone != "15550012002" || rec.UserID != f.user.ID || rec.AuthKeyID != f.authKeyID || rec.SessionID != 77 {
|
||||
t.Fatalf("scoped code = %+v", rec)
|
||||
}
|
||||
|
||||
result, err := f.service.ChangePhone(f.ctx, f.user.ID, f.authKeyID, 88, "+1 555 001 2002", hash, "12345", 1700000000)
|
||||
rawAuthKeyID := [8]byte{8, 8, 8, 8}
|
||||
result, err := f.service.ChangePhone(f.ctx, f.user.ID, f.authKeyID, rawAuthKeyID, 88, "+1 555 001 2002", hash, "12345", 1700000000)
|
||||
if err != nil {
|
||||
t.Fatalf("change phone after session reconnect: %v", err)
|
||||
}
|
||||
if !result.Changed || result.User.Phone != "15550012002" || result.Event.Type != domain.UpdateEventUserPhone || result.Event.Phone != "15550012002" || result.Event.Pts != 1 {
|
||||
t.Fatalf("change result = %+v", result)
|
||||
}
|
||||
if got := f.changes.lastRequest().ExcludeAuthKeyID; got != rawAuthKeyID {
|
||||
t.Fatalf("outbox exclusion auth key = %x, want physical raw %x", got, rawAuthKeyID)
|
||||
}
|
||||
if _, found, _ := f.users.ByPhone(f.ctx, "15550012001"); found {
|
||||
t.Fatal("old phone still resolves")
|
||||
}
|
||||
|
|
@ -103,7 +128,7 @@ func TestPhoneChangeRejectsOccupiedAndCrossAuthCode(t *testing.T) {
|
|||
if err := f.auths.Bind(f.ctx, domain.Authorization{AuthKeyID: otherKey, UserID: occupied.ID}); err != nil {
|
||||
t.Fatalf("bind other auth: %v", err)
|
||||
}
|
||||
if _, err := f.service.ChangePhone(f.ctx, occupied.ID, otherKey, 99, "15550012004", hash, "12345", 0); !errors.Is(err, domain.ErrPhoneCodeInvalid) {
|
||||
if _, err := f.service.ChangePhone(f.ctx, occupied.ID, otherKey, otherKey, 99, "15550012004", hash, "12345", 0); !errors.Is(err, domain.ErrPhoneCodeExpired) {
|
||||
t.Fatalf("cross-auth change err = %v", err)
|
||||
}
|
||||
if got, found, _ := f.users.ByID(f.ctx, occupied.ID); !found || got.Phone != "15550012003" {
|
||||
|
|
@ -118,11 +143,11 @@ func TestPhoneChangeWrongCodeExhaustsAttempts(t *testing.T) {
|
|||
t.Fatalf("send code: %v", err)
|
||||
}
|
||||
for i := 0; i < 3; i++ {
|
||||
if _, err := f.service.ChangePhone(f.ctx, f.user.ID, f.authKeyID, 77, "15550012005", hash, "00000", 0); !errors.Is(err, domain.ErrPhoneCodeInvalid) {
|
||||
if _, err := f.service.ChangePhone(f.ctx, f.user.ID, f.authKeyID, f.authKeyID, 77, "15550012005", hash, "00000", 0); !errors.Is(err, domain.ErrPhoneCodeInvalid) {
|
||||
t.Fatalf("wrong attempt %d err = %v", i+1, err)
|
||||
}
|
||||
}
|
||||
if _, err := f.service.ChangePhone(f.ctx, f.user.ID, f.authKeyID, 77, "15550012005", hash, "12345", 0); !errors.Is(err, domain.ErrPhoneCodeExpired) {
|
||||
if _, err := f.service.ChangePhone(f.ctx, f.user.ID, f.authKeyID, f.authKeyID, 77, "15550012005", hash, "12345", 0); !errors.Is(err, domain.ErrPhoneCodeExpired) {
|
||||
t.Fatalf("exhausted code err = %v", err)
|
||||
}
|
||||
if got, _, _ := f.users.ByID(f.ctx, f.user.ID); got.Phone != "15550012001" {
|
||||
|
|
@ -143,10 +168,10 @@ func TestPhoneChangeNewSendInvalidatesPreviousHash(t *testing.T) {
|
|||
if oldHash == newHash {
|
||||
t.Fatalf("hash was not rotated: %q", oldHash)
|
||||
}
|
||||
if _, err := f.service.ChangePhone(f.ctx, f.user.ID, f.authKeyID, 99, "15550012006", oldHash, "12345", 1700000001); !errors.Is(err, domain.ErrPhoneCodeExpired) {
|
||||
if _, err := f.service.ChangePhone(f.ctx, f.user.ID, f.authKeyID, f.authKeyID, 99, "15550012006", oldHash, "12345", 1700000001); !errors.Is(err, domain.ErrPhoneCodeExpired) {
|
||||
t.Fatalf("old hash replay err = %v", err)
|
||||
}
|
||||
if _, err := f.service.ChangePhone(f.ctx, f.user.ID, f.authKeyID, 99, "15550012006", newHash, "12345", 1700000002); err != nil {
|
||||
if _, err := f.service.ChangePhone(f.ctx, f.user.ID, f.authKeyID, f.authKeyID, 99, "15550012006", newHash, "12345", 1700000002); err != nil {
|
||||
t.Fatalf("new hash change: %v", err)
|
||||
}
|
||||
events, err := f.events.ListAfter(f.ctx, f.user.ID, 0, 10)
|
||||
|
|
@ -168,7 +193,7 @@ func TestPhoneChangeConcurrentReplayAppendsOneEvent(t *testing.T) {
|
|||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
_, err := f.service.ChangePhone(f.ctx, f.user.ID, f.authKeyID, 88, "15550012007", hash, "12345", 1700000003)
|
||||
_, err := f.service.ChangePhone(f.ctx, f.user.ID, f.authKeyID, f.authKeyID, 88, "15550012007", hash, "12345", 1700000003)
|
||||
errs <- err
|
||||
}()
|
||||
}
|
||||
|
|
|
|||
|
|
@ -17,13 +17,14 @@ import (
|
|||
var defaultSecureRandom = []byte("telesrv-tdesktop-dev-secure-rand")
|
||||
|
||||
const (
|
||||
passwordResetWait = 7 * 24 * time.Hour
|
||||
passwordResetRetry = 24 * time.Hour
|
||||
loginEmailVerifyChangePrefix = "login-email-change:"
|
||||
loginEmailVerifySetupPrefix = "login-email-setup:"
|
||||
codeChannelEmailSetup = "email_setup"
|
||||
codeChannelEmailChange = "email_change"
|
||||
codeChannelEmailLogin = "email_login"
|
||||
passwordResetWait = 7 * 24 * time.Hour
|
||||
passwordResetRetry = 24 * time.Hour
|
||||
loginEmailVerifyChangePrefix = "login-email-change:"
|
||||
loginEmailVerifySetupPrefix = "login-email-setup:"
|
||||
codeChannelEmailSetup = "email_setup"
|
||||
codeChannelEmailChange = "email_change"
|
||||
codeChannelEmailLogin = "email_login"
|
||||
codeChannelEmailSetupRequired = "email_setup_required"
|
||||
)
|
||||
|
||||
// Service 提供账号安全配置查询。
|
||||
|
|
@ -568,12 +569,16 @@ func (s *Service) SendLoginEmailCode(ctx context.Context, userID int64, phone, p
|
|||
}
|
||||
key := loginEmailVerifyChangePrefix + fmt.Sprint(userID)
|
||||
rec := store.PhoneCode{
|
||||
Version: store.PhoneCodeVersionCurrent,
|
||||
Code: "",
|
||||
Channel: codeChannelEmailChange,
|
||||
PendingEmail: email,
|
||||
MaxAttempts: s.loginEmailCodeMaxAttempts,
|
||||
}
|
||||
if setup {
|
||||
if s.users == nil {
|
||||
return "", 0, domain.ErrEmailNotAllowed
|
||||
}
|
||||
phone = domain.NormalizePhone(phone)
|
||||
phoneRec, found, err := s.codes.Get(ctx, phoneCodeHash)
|
||||
if err != nil {
|
||||
|
|
@ -582,7 +587,8 @@ func (s *Service) SendLoginEmailCode(ctx context.Context, userID int64, phone, p
|
|||
if !found {
|
||||
return "", 0, domain.ErrEmailCodeInvalid
|
||||
}
|
||||
if phoneRec.Phone != phone {
|
||||
if phoneRec.Version != store.PhoneCodeVersionCurrent || phoneRec.Purpose != "" || phoneRec.Phone != phone ||
|
||||
phoneRec.Channel != codeChannelEmailSetupRequired || phoneRec.SignUpVerified {
|
||||
return "", 0, domain.ErrEmailInvalid
|
||||
}
|
||||
targetUserID := int64(0)
|
||||
|
|
@ -591,6 +597,9 @@ func (s *Service) SendLoginEmailCode(ctx context.Context, userID int64, phone, p
|
|||
} else if found {
|
||||
targetUserID = existingUserID
|
||||
}
|
||||
if phoneRec.IssuedUserID != targetUserID {
|
||||
return "", 0, domain.ErrEmailInvalid
|
||||
}
|
||||
if err := s.ensureLoginEmailAvailable(ctx, targetUserID, email); err != nil {
|
||||
return "", 0, err
|
||||
}
|
||||
|
|
@ -611,7 +620,9 @@ func (s *Service) SendLoginEmailCode(ctx context.Context, userID int64, phone, p
|
|||
return "", 0, err
|
||||
}
|
||||
if err := s.loginEmailSender.SendLoginCode(ctx, email, code, s.loginEmailCodeTTL); err != nil {
|
||||
_ = s.codes.Del(ctx, key)
|
||||
// Set does not expose its generated revision. A blind Del here could
|
||||
// remove a newer concurrent resend; leave the unreachable random code
|
||||
// to expire or be replaced by the retry instead.
|
||||
return "", 0, err
|
||||
}
|
||||
return emailPattern(email), len(code), nil
|
||||
|
|
@ -625,28 +636,43 @@ func (s *Service) VerifyLoginEmail(ctx context.Context, userID int64, phone, pho
|
|||
if setup {
|
||||
key = loginEmailVerifySetupPrefix + phoneCodeHash
|
||||
}
|
||||
rec, found, err := s.codes.Get(ctx, key)
|
||||
snapshot, found, err := s.codes.GetSnapshot(ctx, key)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if !found {
|
||||
return "", domain.ErrEmailCodeInvalid
|
||||
}
|
||||
rec := snapshot.Record
|
||||
if strings.TrimSpace(code) == "" || subtle.ConstantTimeCompare([]byte(rec.Code), []byte(strings.TrimSpace(code))) != 1 {
|
||||
return "", s.rejectEmailCode(ctx, key, rec)
|
||||
return "", s.rejectEmailCode(ctx, key, snapshot)
|
||||
}
|
||||
email := normalizeLoginEmail(rec.PendingEmail)
|
||||
if !validLoginEmail(email) {
|
||||
_ = s.codes.Del(ctx, key)
|
||||
applied, deleteErr := s.codes.CompareAndDelete(ctx, key, snapshot.Revision)
|
||||
if deleteErr != nil {
|
||||
return "", deleteErr
|
||||
}
|
||||
if !applied {
|
||||
return "", domain.ErrEmailCodeInvalid
|
||||
}
|
||||
return "", domain.ErrEmailInvalid
|
||||
}
|
||||
if setup {
|
||||
if s.users == nil || rec.Channel != codeChannelEmailSetup {
|
||||
return "", domain.ErrEmailCodeInvalid
|
||||
}
|
||||
phone = domain.NormalizePhone(phone)
|
||||
phoneRec, found, err := s.codes.Get(ctx, phoneCodeHash)
|
||||
if rec.Phone != phone {
|
||||
return "", domain.ErrEmailCodeInvalid
|
||||
}
|
||||
phoneSnapshot, found, err := s.codes.GetSnapshot(ctx, phoneCodeHash)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if !found || phoneRec.Phone != phone {
|
||||
phoneRec := phoneSnapshot.Record
|
||||
if !found || phoneRec.Version != store.PhoneCodeVersionCurrent || phoneRec.Purpose != "" ||
|
||||
phoneRec.Phone != phone || phoneRec.Channel != codeChannelEmailSetupRequired || phoneRec.SignUpVerified {
|
||||
return "", domain.ErrEmailCodeInvalid
|
||||
}
|
||||
targetUserID := int64(0)
|
||||
|
|
@ -655,11 +681,23 @@ func (s *Service) VerifyLoginEmail(ctx context.Context, userID int64, phone, pho
|
|||
} else if found {
|
||||
targetUserID = existingUserID
|
||||
}
|
||||
if phoneRec.IssuedUserID != targetUserID {
|
||||
s.invalidateLoginCode(ctx, phoneCodeHash, phone)
|
||||
return "", domain.ErrEmailCodeInvalid
|
||||
}
|
||||
if err := s.ensureLoginEmailAvailable(ctx, targetUserID, email); err != nil {
|
||||
_ = s.codes.Del(ctx, key)
|
||||
return "", err
|
||||
}
|
||||
_ = s.codes.Del(ctx, key)
|
||||
// Claim this exact email-code revision before mutating the phone login
|
||||
// state. A concurrent resend rotates the revision, so an old verifier
|
||||
// can neither consume the new code nor authorize the phone hash.
|
||||
claimed, err := s.codes.CompareAndDelete(ctx, key, snapshot.Revision)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if !claimed {
|
||||
return "", domain.ErrEmailCodeInvalid
|
||||
}
|
||||
phoneRec.Channel = codeChannelEmailLogin
|
||||
phoneRec.Code = strings.TrimSpace(code)
|
||||
phoneRec.Email = email
|
||||
|
|
@ -667,43 +705,94 @@ func (s *Service) VerifyLoginEmail(ctx context.Context, userID int64, phone, pho
|
|||
phoneRec.VerifiedEmail = true
|
||||
phoneRec.Attempts = 0
|
||||
phoneRec.MaxAttempts = s.loginEmailCodeMaxAttempts
|
||||
if err := s.codes.Update(ctx, phoneCodeHash, phoneRec); err != nil {
|
||||
updated, err := s.codes.CompareAndUpdate(ctx, phoneCodeHash, phoneSnapshot.Revision, phoneRec)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if _, found, err := s.userIDByPhone(ctx, phone); err != nil {
|
||||
if !updated {
|
||||
return "", domain.ErrEmailCodeInvalid
|
||||
}
|
||||
if targetUserID == 0 {
|
||||
verified, err := s.codes.VerifyLogin(ctx, phoneCodeHash, phone, phoneRec.Code, true, s.loginEmailCodeMaxAttempts)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if verified.Status != store.LoginCodeVerifyAccepted || verified.Record.IssuedUserID != 0 || !verified.Record.SignUpVerified {
|
||||
return "", domain.ErrEmailCodeInvalid
|
||||
}
|
||||
phoneRec = verified.Record
|
||||
}
|
||||
afterUserID := int64(0)
|
||||
if existingUserID, found, err := s.userIDByPhone(ctx, phone); err != nil {
|
||||
return "", err
|
||||
} else if found {
|
||||
if err := s.SetLoginEmailByPhone(ctx, phone, email); err != nil {
|
||||
afterUserID = existingUserID
|
||||
}
|
||||
if afterUserID != targetUserID || phoneRec.IssuedUserID != afterUserID {
|
||||
s.invalidateLoginCode(ctx, phoneCodeHash, phone)
|
||||
return "", domain.ErrEmailCodeInvalid
|
||||
}
|
||||
if targetUserID != 0 {
|
||||
// Keep the identity selected before SMTP verification. Re-resolving
|
||||
// phone at this write boundary would let an A→B transfer attach A's
|
||||
// verified factor to B.
|
||||
if err := s.SetLoginEmail(ctx, targetUserID, email); err != nil {
|
||||
return "", err
|
||||
}
|
||||
finalUserID := int64(0)
|
||||
if existingUserID, found, err := s.userIDByPhone(ctx, phone); err != nil {
|
||||
return "", err
|
||||
} else if found {
|
||||
finalUserID = existingUserID
|
||||
}
|
||||
if finalUserID != targetUserID {
|
||||
s.invalidateLoginCode(ctx, phoneCodeHash, phone)
|
||||
return "", domain.ErrEmailCodeInvalid
|
||||
}
|
||||
}
|
||||
return email, nil
|
||||
}
|
||||
if err := s.ensureLoginEmailAvailable(ctx, userID, email); err != nil {
|
||||
_ = s.codes.Del(ctx, key)
|
||||
return "", err
|
||||
}
|
||||
_ = s.codes.Del(ctx, key)
|
||||
claimed, err := s.codes.CompareAndDelete(ctx, key, snapshot.Revision)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if !claimed {
|
||||
return "", domain.ErrEmailCodeInvalid
|
||||
}
|
||||
if err := s.SetLoginEmail(ctx, userID, email); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return email, nil
|
||||
}
|
||||
|
||||
func (s *Service) rejectEmailCode(ctx context.Context, key string, rec store.PhoneCode) error {
|
||||
func (s *Service) rejectEmailCode(ctx context.Context, key string, snapshot store.PhoneCodeSnapshot) error {
|
||||
rec := snapshot.Record
|
||||
rec.Attempts++
|
||||
max := rec.MaxAttempts
|
||||
if max <= 0 {
|
||||
max = s.loginEmailCodeMaxAttempts
|
||||
}
|
||||
if max > 0 && rec.Attempts >= max {
|
||||
_ = s.codes.Del(ctx, key)
|
||||
if _, err := s.codes.CompareAndDelete(ctx, key, snapshot.Revision); err != nil {
|
||||
return err
|
||||
}
|
||||
return domain.ErrEmailCodeInvalid
|
||||
}
|
||||
_ = s.codes.Update(ctx, key, rec)
|
||||
if _, err := s.codes.CompareAndUpdate(ctx, key, snapshot.Revision, rec); err != nil {
|
||||
return err
|
||||
}
|
||||
return domain.ErrEmailCodeInvalid
|
||||
}
|
||||
|
||||
func (s *Service) invalidateLoginCode(ctx context.Context, hash, phone string) {
|
||||
cleanupCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), 2*time.Second)
|
||||
defer cancel()
|
||||
_, _ = s.codes.InvalidateLoginCode(cleanupCtx, hash, phone)
|
||||
}
|
||||
|
||||
// SetLoginEmail 为已登录用户写入登录邮箱(authed 的 emailVerifyPurposeLoginChange)。
|
||||
// 账号无 2FA 也可设置:account_passwords 行可在 has_password=false 下仅承载登录邮箱。
|
||||
func (s *Service) SetLoginEmail(ctx context.Context, userID int64, email string) error {
|
||||
|
|
@ -726,19 +815,6 @@ func (s *Service) SetLoginEmail(ctx context.Context, userID int64, email string)
|
|||
return s.passwords.Save(ctx, userID, settings)
|
||||
}
|
||||
|
||||
// SetLoginEmailByPhone 为某手机号对应的账号写入登录邮箱(登录流程中的
|
||||
// emailVerifyPurposeLoginSetup,此时尚未鉴权,只能凭 phone 定位用户)。
|
||||
func (s *Service) SetLoginEmailByPhone(ctx context.Context, phone, email string) error {
|
||||
userID, found, err := s.userIDByPhone(ctx, phone)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if !found {
|
||||
return domain.ErrEmailInvalid
|
||||
}
|
||||
return s.SetLoginEmail(ctx, userID, email)
|
||||
}
|
||||
|
||||
// LoginEmail 返回已登录用户的登录邮箱原始地址(用于 verifyEmail 回显 emailVerified.email)。
|
||||
func (s *Service) LoginEmail(ctx context.Context, userID int64) (string, bool, error) {
|
||||
if s == nil || s.passwords == nil || userID == 0 {
|
||||
|
|
@ -754,8 +830,7 @@ func (s *Service) LoginEmail(ctx context.Context, userID int64) (string, bool, e
|
|||
return normalizeLoginEmail(settings.LoginEmail), true, nil
|
||||
}
|
||||
|
||||
// LoginEmailByPhone 按手机号返回登录邮箱原始地址(供 auth.sendCode 检测是否改投邮箱、
|
||||
// login-setup 回显、reset 回显使用)。
|
||||
// LoginEmailByPhone 按手机号返回登录邮箱原始地址,供 auth.sendCode 检测是否改投邮箱。
|
||||
func (s *Service) LoginEmailByPhone(ctx context.Context, phone string) (string, bool, error) {
|
||||
userID, found, err := s.userIDByPhone(ctx, phone)
|
||||
if err != nil || !found {
|
||||
|
|
@ -764,14 +839,12 @@ func (s *Service) LoginEmailByPhone(ctx context.Context, phone string) (string,
|
|||
return s.LoginEmail(ctx, userID)
|
||||
}
|
||||
|
||||
// ClearLoginEmailByPhone 清除某手机号账号的登录邮箱(auth.resetLoginEmail)。
|
||||
func (s *Service) ClearLoginEmailByPhone(ctx context.Context, phone string) error {
|
||||
userID, found, err := s.userIDByPhone(ctx, phone)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if !found {
|
||||
return nil
|
||||
// ClearLoginEmail clears the factor on the exact account selected by the
|
||||
// preceding reset-code consume. Authentication factors must never be mutated
|
||||
// through a second phone→user lookup.
|
||||
func (s *Service) ClearLoginEmail(ctx context.Context, userID int64) error {
|
||||
if s == nil || s.passwords == nil || userID == 0 {
|
||||
return domain.ErrEmailInvalid
|
||||
}
|
||||
settings, found, err := s.passwords.GetByUser(ctx, userID)
|
||||
if err != nil || !found {
|
||||
|
|
|
|||
|
|
@ -15,6 +15,7 @@ func TestResendCodePreservesChangePhoneScopeAndSMSDelivery(t *testing.T) {
|
|||
codes := memory.NewCodeStore()
|
||||
authKeyID := [8]byte{8, 7, 6}
|
||||
rec := store.PhoneCode{
|
||||
Version: store.PhoneCodeVersionCurrent,
|
||||
Phone: "15550014001",
|
||||
Code: "old",
|
||||
Channel: codeChannelPhone,
|
||||
|
|
@ -69,3 +70,34 @@ func TestResendCodePreservesChangePhoneScopeAndSMSDelivery(t *testing.T) {
|
|||
t.Fatal("scoped cancel left hash valid")
|
||||
}
|
||||
}
|
||||
|
||||
func TestResendAndCancelRejectLegacyChangePhoneCode(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
codes := memory.NewCodeStore()
|
||||
authKeyID := [8]byte{8, 8, 8}
|
||||
legacy := store.PhoneCode{
|
||||
Version: 0, Phone: "15550014002", Code: "12345", Channel: codeChannelPhone,
|
||||
Purpose: store.PhoneCodePurposeChangePhone, UserID: 43, AuthKeyID: authKeyID,
|
||||
}
|
||||
svc := NewService(memory.NewUserStore(), memory.NewAuthorizationStore(), codes, nil, nil, "12345", WithCodeTTL(time.Minute))
|
||||
|
||||
if err := codes.Set(ctx, "legacy-resend", legacy, time.Minute); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if _, err := svc.ResendCodeForAuthKey(ctx, authKeyID, legacy.Phone, "legacy-resend"); err != ErrCodeExpired {
|
||||
t.Fatalf("legacy resend err=%v, want ErrCodeExpired", err)
|
||||
}
|
||||
if _, found, _ := codes.Get(ctx, "legacy-resend"); found {
|
||||
t.Fatal("legacy resend left code active")
|
||||
}
|
||||
|
||||
if err := codes.Set(ctx, "legacy-cancel", legacy, time.Minute); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if err := svc.CancelCodeForAuthKey(ctx, authKeyID, legacy.Phone, "legacy-cancel"); err != ErrCodeExpired {
|
||||
t.Fatalf("legacy cancel err=%v, want ErrCodeExpired", err)
|
||||
}
|
||||
if _, found, _ := codes.Get(ctx, "legacy-cancel"); found {
|
||||
t.Fatal("legacy cancel left code active")
|
||||
}
|
||||
}
|
||||
|
|
|
|||
366
internal/app/auth/login_code_delivery_test.go
Normal file
366
internal/app/auth/login_code_delivery_test.go
Normal file
|
|
@ -0,0 +1,366 @@
|
|||
package auth
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"telesrv/internal/domain"
|
||||
"telesrv/internal/store"
|
||||
"telesrv/internal/store/memory"
|
||||
)
|
||||
|
||||
type captureLoginCodeDelivery struct {
|
||||
requests []domain.LoginCodeDeliveryRequest
|
||||
result domain.LoginCodeDeliveryResult
|
||||
err error
|
||||
failAt int
|
||||
}
|
||||
|
||||
func (d *captureLoginCodeDelivery) DeliverLoginCodeMessage(_ context.Context, req domain.LoginCodeDeliveryRequest) (domain.LoginCodeDeliveryResult, error) {
|
||||
d.requests = append(d.requests, req)
|
||||
if d.err != nil && (d.failAt == 0 || len(d.requests) == d.failAt) {
|
||||
return domain.LoginCodeDeliveryResult{}, d.err
|
||||
}
|
||||
return d.result, nil
|
||||
}
|
||||
|
||||
type trackingCodeStore struct {
|
||||
store.CodeStore
|
||||
lastSetHash string
|
||||
deleted []string
|
||||
deleteCtx []error
|
||||
deleteErr error
|
||||
}
|
||||
|
||||
func (s *trackingCodeStore) Set(ctx context.Context, hash string, code store.PhoneCode, ttl time.Duration) error {
|
||||
s.lastSetHash = hash
|
||||
return s.CodeStore.Set(ctx, hash, code, ttl)
|
||||
}
|
||||
|
||||
func (s *trackingCodeStore) Del(ctx context.Context, hash string) error {
|
||||
s.deleted = append(s.deleted, hash)
|
||||
s.deleteCtx = append(s.deleteCtx, ctx.Err())
|
||||
if s.deleteErr != nil {
|
||||
return s.deleteErr
|
||||
}
|
||||
return s.CodeStore.Del(ctx, hash)
|
||||
}
|
||||
|
||||
func TestExistingAccountSendCodeDeliversBeforeSignInAndDoesNotRedeliver(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
users := memory.NewUserStore()
|
||||
authz := memory.NewAuthorizationStore()
|
||||
codes := memory.NewCodeStore()
|
||||
u, err := users.Create(ctx, domain.User{Phone: "15550009201", FirstName: "Existing"})
|
||||
if err != nil {
|
||||
t.Fatalf("create user: %v", err)
|
||||
}
|
||||
delivery := &captureLoginCodeDelivery{result: domain.LoginCodeDeliveryResult{Created: true}}
|
||||
svc := NewService(users, authz, codes, nil, nil, "12345", WithLoginCodeDelivery(delivery))
|
||||
|
||||
before := int(time.Now().Unix())
|
||||
hash, err := svc.SendCode(ctx, "+1 555 000 9201")
|
||||
if err != nil {
|
||||
t.Fatalf("SendCode: %v", err)
|
||||
}
|
||||
if hash == "" || len(delivery.requests) != 1 {
|
||||
t.Fatalf("SendCode hash=%q delivery calls=%d, want non-empty/1", hash, len(delivery.requests))
|
||||
}
|
||||
req := delivery.requests[0]
|
||||
if req.UserID != u.ID || req.PhoneCodeHash != hash || req.Code != "12345" || req.Date < before || req.ExpiresAt < int64(before)+int64((5*time.Minute)/time.Second)-1 {
|
||||
t.Fatalf("delivery request = %+v, want user=%d hash=%q code=12345 date>=%d", req, u.ID, hash, before)
|
||||
}
|
||||
if rec, found, err := codes.Get(ctx, hash); err != nil || !found || rec.Code != "12345" {
|
||||
t.Fatalf("code after synchronous delivery = %+v found=%v err=%v", rec, found, err)
|
||||
}
|
||||
|
||||
var key [8]byte
|
||||
key[0] = 0x92
|
||||
got, lateMessage, needSignUp, err := svc.SignIn(ctx, domain.Authorization{AuthKeyID: key}, "+15550009201", hash, "12345")
|
||||
if err != nil || needSignUp || got.ID != u.ID {
|
||||
t.Fatalf("SignIn user=%d needSignUp=%v err=%v, want %d/false", got.ID, needSignUp, err, u.ID)
|
||||
}
|
||||
if lateMessage.ID != 0 || len(delivery.requests) != 1 {
|
||||
t.Fatalf("SignIn lateMessage=%+v delivery calls=%d, want zero/unchanged", lateMessage, len(delivery.requests))
|
||||
}
|
||||
}
|
||||
|
||||
func TestDeliveredLoginCodeSurvivesWrongSignInAndCancelWithoutDuplicate(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
users := memory.NewUserStore()
|
||||
u, err := users.Create(ctx, domain.User{Phone: "15550009208", FirstName: "Cancel"})
|
||||
if err != nil {
|
||||
t.Fatalf("create user: %v", err)
|
||||
}
|
||||
codes := memory.NewCodeStore()
|
||||
dialogs := memory.NewDialogStore()
|
||||
messages := memory.NewMessageStore(dialogs)
|
||||
events := memory.NewUpdateEventStore()
|
||||
delivery := memory.NewLoginCodeDeliveryStore(messages, events)
|
||||
svc := NewService(users, memory.NewAuthorizationStore(), codes, nil, nil, "12345", WithLoginCodeDelivery(delivery))
|
||||
|
||||
hash, err := svc.SendCode(ctx, "15550009208")
|
||||
if err != nil {
|
||||
t.Fatalf("SendCode: %v", err)
|
||||
}
|
||||
assertFacts := func(stage string) {
|
||||
t.Helper()
|
||||
history, historyErr := messages.ListByUser(ctx, u.ID, domain.MessageFilter{
|
||||
HasPeer: true,
|
||||
Peer: domain.Peer{Type: domain.PeerTypeUser, ID: domain.OfficialSystemUserID},
|
||||
Limit: 10,
|
||||
})
|
||||
durable, eventErr := events.ListAfter(ctx, u.ID, 0, 10)
|
||||
if historyErr != nil || eventErr != nil || len(history.Messages) != 1 || len(durable) != 1 {
|
||||
t.Fatalf("%s messages=%d events=%d historyErr=%v eventErr=%v, want 1/1", stage, len(history.Messages), len(durable), historyErr, eventErr)
|
||||
}
|
||||
}
|
||||
assertFacts("after SendCode")
|
||||
|
||||
if _, late, _, err := svc.SignIn(ctx, domain.Authorization{}, "15550009208", hash, "00000"); !errors.Is(err, ErrCodeInvalid) || late.ID != 0 {
|
||||
t.Fatalf("wrong SignIn late=%+v err=%v, want ErrCodeInvalid/no message", late, err)
|
||||
}
|
||||
assertFacts("after wrong SignIn")
|
||||
|
||||
if err := svc.CancelCode(ctx, "15550009208", hash); err != nil {
|
||||
t.Fatalf("CancelCode: %v", err)
|
||||
}
|
||||
assertFacts("after CancelCode")
|
||||
if _, late, _, err := svc.SignIn(ctx, domain.Authorization{}, "15550009208", hash, "12345"); !errors.Is(err, ErrCodeExpired) || late.ID != 0 {
|
||||
t.Fatalf("SignIn after cancel late=%+v err=%v, want ErrCodeExpired/no message", late, err)
|
||||
}
|
||||
assertFacts("after canceled SignIn")
|
||||
}
|
||||
|
||||
func TestCodeIssuedBeforeConcurrentOwnerCreationIsRejected(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
users := memory.NewUserStore()
|
||||
authz := memory.NewAuthorizationStore()
|
||||
codes := memory.NewCodeStore()
|
||||
delivery := &captureLoginCodeDelivery{}
|
||||
svc := NewService(users, authz, codes, nil, nil, "12345", WithLoginCodeDelivery(delivery))
|
||||
|
||||
hash, err := svc.SendCode(ctx, "15550009209")
|
||||
if err != nil {
|
||||
t.Fatalf("SendCode before signup: %v", err)
|
||||
}
|
||||
rec, found, err := codes.Get(ctx, hash)
|
||||
if err != nil || !found || rec.Version != store.PhoneCodeVersionCurrent || rec.IssuedUserID != 0 || rec.SignUpVerified || len(delivery.requests) != 0 {
|
||||
t.Fatalf("pre-signup code=%+v found=%v err=%v deliveries=%d", rec, found, err, len(delivery.requests))
|
||||
}
|
||||
u, err := users.Create(ctx, domain.User{Phone: "15550009209", FirstName: "Concurrent"})
|
||||
if err != nil {
|
||||
t.Fatalf("concurrent create user: %v", err)
|
||||
}
|
||||
var key [8]byte
|
||||
key[0] = 0x93
|
||||
got, lateMessage, needSignUp, err := svc.SignIn(ctx, domain.Authorization{AuthKeyID: key}, "15550009209", hash, "12345")
|
||||
if !errors.Is(err, ErrCodeInvalid) || needSignUp || got.ID != 0 || lateMessage.ID != 0 {
|
||||
t.Fatalf("SignIn after owner creation got=%+v late=%+v needSignUp=%v err=%v, want invalid", got, lateMessage, needSignUp, err)
|
||||
}
|
||||
if len(delivery.requests) != 0 {
|
||||
t.Fatalf("owner-transfer code was delivered to new owner: %+v", delivery.requests)
|
||||
}
|
||||
if bound, ok, err := svc.UserID(ctx, key); err != nil || ok || bound != 0 {
|
||||
t.Fatalf("bound user=%d ok=%v err=%v, want no authorization (created uid=%d)", bound, ok, err, u.ID)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExistingAccountRepeatedSendCodeDeliversEachIssuedHashWithoutSignIn(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
users := memory.NewUserStore()
|
||||
if _, err := users.Create(ctx, domain.User{Phone: "15550009202"}); err != nil {
|
||||
t.Fatalf("create user: %v", err)
|
||||
}
|
||||
delivery := &captureLoginCodeDelivery{}
|
||||
svc := NewService(users, memory.NewAuthorizationStore(), memory.NewCodeStore(), nil, nil, "12345", WithLoginCodeDelivery(delivery))
|
||||
|
||||
first, err := svc.SendCode(ctx, "15550009202")
|
||||
if err != nil {
|
||||
t.Fatalf("first SendCode: %v", err)
|
||||
}
|
||||
second, err := svc.SendCode(ctx, "15550009202")
|
||||
if err != nil {
|
||||
t.Fatalf("second SendCode: %v", err)
|
||||
}
|
||||
if first == second || len(delivery.requests) != 2 {
|
||||
t.Fatalf("hashes=%q/%q delivery calls=%d, want distinct/2", first, second, len(delivery.requests))
|
||||
}
|
||||
if delivery.requests[0].PhoneCodeHash != first || delivery.requests[1].PhoneCodeHash != second {
|
||||
t.Fatalf("delivery hashes = %q/%q, want %q/%q", delivery.requests[0].PhoneCodeHash, delivery.requests[1].PhoneCodeHash, first, second)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExistingAccountSendCodeDeliveryFailureRevokesCode(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
users := memory.NewUserStore()
|
||||
if _, err := users.Create(ctx, domain.User{Phone: "15550009203"}); err != nil {
|
||||
t.Fatalf("create user: %v", err)
|
||||
}
|
||||
baseCodes := memory.NewCodeStore()
|
||||
codes := &trackingCodeStore{CodeStore: baseCodes}
|
||||
deliveryCause := errors.New("durable write failed")
|
||||
delivery := &captureLoginCodeDelivery{err: deliveryCause}
|
||||
svc := NewService(users, memory.NewAuthorizationStore(), codes, nil, nil, "12345", WithLoginCodeDelivery(delivery))
|
||||
|
||||
hash, err := svc.SendCode(ctx, "15550009203")
|
||||
if hash != "" || !errors.Is(err, ErrLoginCodeDeliveryFailed) || !errors.Is(err, deliveryCause) {
|
||||
t.Fatalf("SendCode hash=%q err=%v, want empty ErrLoginCodeDeliveryFailed+cause", hash, err)
|
||||
}
|
||||
if codes.lastSetHash == "" || len(codes.deleted) != 1 || codes.deleted[0] != codes.lastSetHash {
|
||||
t.Fatalf("set hash=%q deleted=%v, want exact rollback", codes.lastSetHash, codes.deleted)
|
||||
}
|
||||
if _, found, getErr := baseCodes.Get(ctx, codes.lastSetHash); getErr != nil || found {
|
||||
t.Fatalf("rolled-back hash found=%v err=%v", found, getErr)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExistingAccountAmbiguousDeliveryPreservesCodeForIdempotentRetry(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
users := memory.NewUserStore()
|
||||
if _, err := users.Create(ctx, domain.User{Phone: "15550009213"}); err != nil {
|
||||
t.Fatalf("create user: %v", err)
|
||||
}
|
||||
baseCodes := memory.NewCodeStore()
|
||||
codes := &trackingCodeStore{CodeStore: baseCodes}
|
||||
delivery := &captureLoginCodeDelivery{err: domain.ErrLoginCodeDeliveryCommitAmbiguous}
|
||||
svc := NewService(users, memory.NewAuthorizationStore(), codes, nil, nil, "12345", WithLoginCodeDelivery(delivery))
|
||||
|
||||
hash, err := svc.SendCode(ctx, "15550009213")
|
||||
if hash != "" || !errors.Is(err, ErrLoginCodeDeliveryFailed) || !errors.Is(err, domain.ErrLoginCodeDeliveryCommitAmbiguous) {
|
||||
t.Fatalf("SendCode hash=%q err=%v, want ambiguous delivery failure", hash, err)
|
||||
}
|
||||
if codes.lastSetHash == "" || len(codes.deleted) != 0 {
|
||||
t.Fatalf("ambiguous delivery set=%q deleted=%v, want code preserved", codes.lastSetHash, codes.deleted)
|
||||
}
|
||||
if rec, found, getErr := baseCodes.Get(ctx, codes.lastSetHash); getErr != nil || !found || rec.Code != "12345" {
|
||||
t.Fatalf("ambiguous delivery code=%+v found=%v err=%v", rec, found, getErr)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExplicitDeliveryFailureRollsBackWithDetachedContext(t *testing.T) {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
users := memory.NewUserStore()
|
||||
if _, err := users.Create(context.Background(), domain.User{Phone: "15550009214"}); err != nil {
|
||||
t.Fatalf("create user: %v", err)
|
||||
}
|
||||
baseCodes := memory.NewCodeStore()
|
||||
codes := &trackingCodeStore{CodeStore: baseCodes}
|
||||
delivery := &captureLoginCodeDelivery{err: errors.New("definite rollback")}
|
||||
svc := NewService(users, memory.NewAuthorizationStore(), codes, nil, nil, "12345", WithLoginCodeDelivery(delivery))
|
||||
|
||||
if hash, err := svc.SendCode(ctx, "15550009214"); hash != "" || !errors.Is(err, ErrLoginCodeDeliveryFailed) {
|
||||
t.Fatalf("SendCode hash=%q err=%v, want definite failure", hash, err)
|
||||
}
|
||||
if len(codes.deleted) != 1 || len(codes.deleteCtx) != 1 || codes.deleteCtx[0] != nil {
|
||||
t.Fatalf("rollback deleted=%v ctxErr=%v, want one detached delete", codes.deleted, codes.deleteCtx)
|
||||
}
|
||||
if _, found, err := baseCodes.Get(context.Background(), codes.lastSetHash); err != nil || found {
|
||||
t.Fatalf("detached rollback found=%v err=%v", found, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExistingAccountMissingDeliveryFailsClosedAndRevokesCode(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
users := memory.NewUserStore()
|
||||
if _, err := users.Create(ctx, domain.User{Phone: "15550009204"}); err != nil {
|
||||
t.Fatalf("create user: %v", err)
|
||||
}
|
||||
baseCodes := memory.NewCodeStore()
|
||||
codes := &trackingCodeStore{CodeStore: baseCodes}
|
||||
svc := NewService(users, memory.NewAuthorizationStore(), codes, nil, nil, "12345")
|
||||
|
||||
hash, err := svc.SendCode(ctx, "15550009204")
|
||||
if hash != "" || !errors.Is(err, ErrLoginCodeDeliveryUnavailable) {
|
||||
t.Fatalf("SendCode hash=%q err=%v, want unavailable", hash, err)
|
||||
}
|
||||
if codes.lastSetHash == "" {
|
||||
t.Fatal("missing delivery was checked before code creation; want rollback path covered")
|
||||
}
|
||||
if _, found, getErr := baseCodes.Get(ctx, codes.lastSetHash); getErr != nil || found {
|
||||
t.Fatalf("unavailable delivery hash found=%v err=%v", found, getErr)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExistingAccountResendDeliversNewHashAndInvalidatesOld(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
users := memory.NewUserStore()
|
||||
if _, err := users.Create(ctx, domain.User{Phone: "15550009205"}); err != nil {
|
||||
t.Fatalf("create user: %v", err)
|
||||
}
|
||||
codes := memory.NewCodeStore()
|
||||
delivery := &captureLoginCodeDelivery{}
|
||||
svc := NewService(users, memory.NewAuthorizationStore(), codes, nil, nil, "12345", WithLoginCodeDelivery(delivery))
|
||||
|
||||
oldHash, err := svc.SendCode(ctx, "15550009205")
|
||||
if err != nil {
|
||||
t.Fatalf("SendCode: %v", err)
|
||||
}
|
||||
newHash, err := svc.ResendCode(ctx, "15550009205", oldHash)
|
||||
if err != nil {
|
||||
t.Fatalf("ResendCode: %v", err)
|
||||
}
|
||||
if oldHash == newHash || len(delivery.requests) != 2 || delivery.requests[1].PhoneCodeHash != newHash {
|
||||
t.Fatalf("old/new=%q/%q deliveries=%+v", oldHash, newHash, delivery.requests)
|
||||
}
|
||||
if _, found, err := codes.Get(ctx, oldHash); err != nil || found {
|
||||
t.Fatalf("old code found=%v err=%v", found, err)
|
||||
}
|
||||
if _, found, err := codes.Get(ctx, newHash); err != nil || !found {
|
||||
t.Fatalf("new code found=%v err=%v", found, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExistingAccountResendDeliveryFailureLeavesNoUsableCode(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
users := memory.NewUserStore()
|
||||
if _, err := users.Create(ctx, domain.User{Phone: "15550009206"}); err != nil {
|
||||
t.Fatalf("create user: %v", err)
|
||||
}
|
||||
codes := memory.NewCodeStore()
|
||||
delivery := &captureLoginCodeDelivery{err: errors.New("second delivery failed"), failAt: 2}
|
||||
svc := NewService(users, memory.NewAuthorizationStore(), codes, nil, nil, "12345", WithLoginCodeDelivery(delivery))
|
||||
|
||||
oldHash, err := svc.SendCode(ctx, "15550009206")
|
||||
if err != nil {
|
||||
t.Fatalf("SendCode: %v", err)
|
||||
}
|
||||
newHash, err := svc.ResendCode(ctx, "15550009206", oldHash)
|
||||
if newHash != "" || !errors.Is(err, ErrLoginCodeDeliveryFailed) || len(delivery.requests) != 2 {
|
||||
t.Fatalf("ResendCode hash=%q err=%v deliveries=%d", newHash, err, len(delivery.requests))
|
||||
}
|
||||
failedHash := delivery.requests[1].PhoneCodeHash
|
||||
for _, hash := range []string{oldHash, failedHash} {
|
||||
if _, found, getErr := codes.Get(ctx, hash); getErr != nil || found {
|
||||
t.Fatalf("failed resend hash %q found=%v err=%v", hash, found, getErr)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestConfiguredEmailLoginDoesNotLeakCodeThroughAppDelivery(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
users := memory.NewUserStore()
|
||||
if _, err := users.Create(ctx, domain.User{Phone: "15550009207"}); err != nil {
|
||||
t.Fatalf("create user: %v", err)
|
||||
}
|
||||
emails := &testLoginEmailStore{emails: map[string]string{"15550009207": "secure@example.test"}}
|
||||
mailSender := &testMailSender{}
|
||||
delivery := &captureLoginCodeDelivery{}
|
||||
svc := NewService(users, memory.NewAuthorizationStore(), memory.NewCodeStore(), nil, nil, "12345",
|
||||
WithLoginEmail(LoginEmailOptions{Enabled: true, CodeLength: 6, Store: emails, Sender: mailSender}),
|
||||
WithLoginCodeDelivery(delivery),
|
||||
)
|
||||
|
||||
if _, err := svc.SendCode(ctx, "15550009207"); err != nil {
|
||||
t.Fatalf("SendCode: %v", err)
|
||||
}
|
||||
if mailSender.to != "secure@example.test" || mailSender.code == "" {
|
||||
t.Fatalf("email delivery = %q/%q", mailSender.to, mailSender.code)
|
||||
}
|
||||
if len(delivery.requests) != 0 {
|
||||
t.Fatalf("email code leaked into app delivery: %+v", delivery.requests)
|
||||
}
|
||||
}
|
||||
|
|
@ -19,8 +19,7 @@ func (s *testLoginEmailStore) LoginEmailByPhone(_ context.Context, phone string)
|
|||
return email, ok, nil
|
||||
}
|
||||
|
||||
func (s *testLoginEmailStore) SetLoginEmailByPhone(_ context.Context, phone, email string) error {
|
||||
s.emails[domain.NormalizePhone(phone)] = email
|
||||
func (s *testLoginEmailStore) SetLoginEmail(_ context.Context, _ int64, _ string) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -9,13 +9,17 @@ import (
|
|||
"telesrv/internal/store/memory"
|
||||
)
|
||||
|
||||
// TestSignInWithEmailCompletesLogin 验证带 email_verification 的登录:注册账号→登出→
|
||||
// 重新 sendCode→用任意邮箱验证码经 SignInWithEmail 完成登录。
|
||||
// TestSignInWithEmailCompletesLogin 验证旧客户端把 phone channel 放进
|
||||
// email_verification 时仍可登录,但验证码必须精确匹配,不能用任意非空值绕过。
|
||||
func TestSignInWithEmailCompletesLogin(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
users := memory.NewUserStore()
|
||||
authz := memory.NewAuthorizationStore()
|
||||
svc := NewService(users, authz, memory.NewCodeStore(), nil, nil, "12345")
|
||||
dialogs := memory.NewDialogStore()
|
||||
messages := memory.NewMessageStore(dialogs)
|
||||
svc := NewService(users, authz, memory.NewCodeStore(), nil, nil, "12345",
|
||||
WithLoginCodeDelivery(memory.NewLoginCodeDeliveryStore(messages, memory.NewUpdateEventStore())),
|
||||
)
|
||||
var key [8]byte
|
||||
key[0] = 0x42
|
||||
|
||||
|
|
@ -23,6 +27,7 @@ func TestSignInWithEmailCompletesLogin(t *testing.T) {
|
|||
if err != nil {
|
||||
t.Fatalf("SendCode signup: %v", err)
|
||||
}
|
||||
verifyCodeForSignUp(t, svc, "+15550009001", hash, "12345")
|
||||
u, _, err := svc.SignUp(ctx, domain.Authorization{AuthKeyID: key}, "+15550009001", hash, "Email", "Login")
|
||||
if err != nil {
|
||||
t.Fatalf("SignUp: %v", err)
|
||||
|
|
@ -35,7 +40,10 @@ func TestSignInWithEmailCompletesLogin(t *testing.T) {
|
|||
if err != nil {
|
||||
t.Fatalf("SendCode signin: %v", err)
|
||||
}
|
||||
got, _, needSignUp, err := svc.SignInWithEmail(ctx, domain.Authorization{AuthKeyID: key}, "+15550009001", hash, "anything-goes")
|
||||
if _, _, _, err := svc.SignInWithEmail(ctx, domain.Authorization{AuthKeyID: key}, "+15550009001", hash, "anything-goes"); !errors.Is(err, ErrCodeInvalid) {
|
||||
t.Fatalf("SignInWithEmail arbitrary nonempty code err=%v, want ErrCodeInvalid", err)
|
||||
}
|
||||
got, _, needSignUp, err := svc.SignInWithEmail(ctx, domain.Authorization{AuthKeyID: key}, "+15550009001", hash, "12345")
|
||||
if err != nil {
|
||||
t.Fatalf("SignInWithEmail: %v", err)
|
||||
}
|
||||
|
|
@ -48,7 +56,7 @@ func TestSignInWithEmailCompletesLogin(t *testing.T) {
|
|||
}
|
||||
}
|
||||
|
||||
// TestSignInWithEmailRejectsEmptyCode 空邮箱验证码必须被拒(即使开发环境码任意,也不能空)。
|
||||
// TestSignInWithEmailRejectsEmptyCode 空邮箱验证码必须被拒。
|
||||
func TestSignInWithEmailRejectsEmptyCode(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
svc := NewService(memory.NewUserStore(), memory.NewAuthorizationStore(), memory.NewCodeStore(), nil, nil, "12345")
|
||||
|
|
@ -66,7 +74,12 @@ func TestSignInWithEmailRejectsEmptyCode(t *testing.T) {
|
|||
func TestSignInWithEmailStillHonorsTwoFactor(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
passwords := memory.NewPasswordStore()
|
||||
svc := NewService(memory.NewUserStore(), memory.NewAuthorizationStore(), memory.NewCodeStore(), nil, nil, "12345", WithPasswords(passwords))
|
||||
dialogs := memory.NewDialogStore()
|
||||
messages := memory.NewMessageStore(dialogs)
|
||||
svc := NewService(memory.NewUserStore(), memory.NewAuthorizationStore(), memory.NewCodeStore(), nil, nil, "12345",
|
||||
WithPasswords(passwords),
|
||||
WithLoginCodeDelivery(memory.NewLoginCodeDeliveryStore(messages, memory.NewUpdateEventStore())),
|
||||
)
|
||||
var key [8]byte
|
||||
key[0] = 0x43
|
||||
|
||||
|
|
@ -74,6 +87,7 @@ func TestSignInWithEmailStillHonorsTwoFactor(t *testing.T) {
|
|||
if err != nil {
|
||||
t.Fatalf("SendCode signup: %v", err)
|
||||
}
|
||||
verifyCodeForSignUp(t, svc, "+15550009003", hash, "12345")
|
||||
u, _, err := svc.SignUp(ctx, domain.Authorization{AuthKeyID: key}, "+15550009003", hash, "Two", "Factor")
|
||||
if err != nil {
|
||||
t.Fatalf("SignUp: %v", err)
|
||||
|
|
@ -89,7 +103,7 @@ func TestSignInWithEmailStillHonorsTwoFactor(t *testing.T) {
|
|||
if err != nil {
|
||||
t.Fatalf("SendCode signin: %v", err)
|
||||
}
|
||||
got, _, _, err := svc.SignInWithEmail(ctx, domain.Authorization{AuthKeyID: key}, "+15550009003", hash, "any-email-code")
|
||||
got, _, _, err := svc.SignInWithEmail(ctx, domain.Authorization{AuthKeyID: key}, "+15550009003", hash, "12345")
|
||||
if !errors.Is(err, domain.ErrSessionPasswordNeeded) {
|
||||
t.Fatalf("SignInWithEmail err = %v, want ErrSessionPasswordNeeded", err)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -19,6 +19,7 @@ func TestSignUpPremiumGrant(t *testing.T) {
|
|||
if err != nil {
|
||||
t.Fatalf("SendCode: %v", err)
|
||||
}
|
||||
verifyCodeForSignUp(t, svc, "+15550004401", hash, "12345")
|
||||
u, _, err := svc.SignUp(ctx, domain.Authorization{}, "+15550004401", hash, "Prem", "User")
|
||||
if err != nil {
|
||||
t.Fatalf("SignUp: %v", err)
|
||||
|
|
@ -41,6 +42,7 @@ func TestSignUpPremiumGrantDisabled(t *testing.T) {
|
|||
if err != nil {
|
||||
t.Fatalf("SendCode: %v", err)
|
||||
}
|
||||
verifyCodeForSignUp(t, svc, "+15550004402", hash, "12345")
|
||||
u, _, err := svc.SignUp(ctx, domain.Authorization{}, "+15550004402", hash, "Free", "User")
|
||||
if err != nil {
|
||||
t.Fatalf("SignUp: %v", err)
|
||||
|
|
|
|||
|
|
@ -27,6 +27,13 @@ var (
|
|||
ErrCodeExpired = errors.New("phone code expired or not found")
|
||||
ErrCodeInvalid = errors.New("phone code invalid")
|
||||
ErrEncryptedMessageInvalid = errors.New("encrypted message invalid")
|
||||
// ErrLoginCodeDeliveryUnavailable 表示已有账号的 app-code 没有可用的
|
||||
// durable message/event/outbox 投递边界。这是服务端配置错误,不能降级成
|
||||
// “继续返回 sentCode,等 signIn 后补发”。
|
||||
ErrLoginCodeDeliveryUnavailable = errors.New("login code durable delivery unavailable")
|
||||
// ErrLoginCodeDeliveryFailed 表示 durable 投递未成功。SendCode/ResendCode
|
||||
// 必须同时撤销刚写入的 CodeStore hash,防止客户拿到无法送达的码。
|
||||
ErrLoginCodeDeliveryFailed = errors.New("login code durable delivery failed")
|
||||
// ErrPhoneNumberInvalid 表示手机号为空或非纯数字/长度越界。
|
||||
// 0090 把 users.phone 唯一约束改为忽略空串的部分索引(bot 行 phone=''),
|
||||
// 因此 phone 校验必须前移到 auth 入口,否则 sendCode/signUp 可无限铸造
|
||||
|
|
@ -40,6 +47,7 @@ const (
|
|||
codeChannelPhone = "phone"
|
||||
codeChannelEmailLogin = "email_login"
|
||||
codeChannelEmailSetupRequired = "email_setup_required"
|
||||
loginCodeRollbackTimeout = 2 * time.Second
|
||||
)
|
||||
|
||||
// validPhone 校验规范化后的手机号:5-32 位纯数字(上限对齐 users.phone 列宽)。
|
||||
|
|
@ -68,6 +76,7 @@ type Service struct {
|
|||
passwords store.PasswordStore
|
||||
messages store.MessageStore
|
||||
dialogs store.DialogStore
|
||||
loginCodeDelivery store.LoginCodeDeliveryStore
|
||||
bots store.BotStore
|
||||
fixedCode string
|
||||
codeTTL time.Duration
|
||||
|
|
@ -83,7 +92,7 @@ type Service struct {
|
|||
|
||||
type loginEmailStore interface {
|
||||
LoginEmailByPhone(ctx context.Context, phone string) (string, bool, error)
|
||||
SetLoginEmailByPhone(ctx context.Context, phone, email string) error
|
||||
SetLoginEmail(ctx context.Context, userID int64, email string) error
|
||||
}
|
||||
|
||||
type LoginEmailOptions struct {
|
||||
|
|
@ -102,7 +111,9 @@ type authorizationRevoker interface {
|
|||
// Option 调整登录服务的可选依赖。
|
||||
type Option func(*Service)
|
||||
|
||||
// WithLoginMessages 在登录成功后写入官方系统账号的登录消息与会话摘要。
|
||||
// WithLoginMessages 在新用户注册成功后写入官方系统账号的首条登录消息与会话摘要。
|
||||
// 已有账号的 app 验证码必须在 auth.sendCode/resendCode 阶段通过
|
||||
// WithLoginCodeDelivery 持久化,禁止在 signIn 成功后补发。
|
||||
func WithLoginMessages(messages store.MessageStore, dialogs store.DialogStore) Option {
|
||||
return func(s *Service) {
|
||||
s.messages = messages
|
||||
|
|
@ -110,6 +121,15 @@ func WithLoginMessages(messages store.MessageStore, dialogs store.DialogStore) O
|
|||
}
|
||||
}
|
||||
|
||||
// WithLoginCodeDelivery 注入已有账号 app-code 的 durable 投递边界。
|
||||
// 实现必须以 user_id + phone_code_hash 幂等,并原子写入 777000
|
||||
// message/dialog/user update event/dispatch outbox。
|
||||
func WithLoginCodeDelivery(delivery store.LoginCodeDeliveryStore) Option {
|
||||
return func(s *Service) {
|
||||
s.loginCodeDelivery = delivery
|
||||
}
|
||||
}
|
||||
|
||||
// WithPasswords lets sign-in stop at SESSION_PASSWORD_NEEDED for 2FA accounts.
|
||||
func WithPasswords(passwords store.PasswordStore) Option {
|
||||
return func(s *Service) {
|
||||
|
|
@ -271,53 +291,157 @@ func (s *Service) SendCode(ctx context.Context, phone string) (string, error) {
|
|||
if systemLoginPhoneForbidden(phone) {
|
||||
return "", ErrSystemUserLoginForbidden
|
||||
}
|
||||
existing, found, err := s.currentPhoneOwner(ctx, phone)
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("lookup login-code recipient: %w", err)
|
||||
}
|
||||
if found && systemUserLoginForbidden(existing) {
|
||||
return "", ErrSystemUserLoginForbidden
|
||||
}
|
||||
issuedUserID := int64(0)
|
||||
if found {
|
||||
issuedUserID = existing.ID
|
||||
}
|
||||
if s.loginEmailEnabled && s.loginEmails != nil {
|
||||
email, found, err := s.loginEmails.LoginEmailByPhone(ctx, phone)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if found && strings.TrimSpace(email) != "" {
|
||||
return s.createEmailLoginCode(ctx, phone, email)
|
||||
return s.createEmailLoginCode(ctx, phone, email, issuedUserID)
|
||||
}
|
||||
if s.loginEmailRequireSetup {
|
||||
return s.createSetupRequiredCode(ctx, phone)
|
||||
return s.createSetupRequiredCode(ctx, phone, issuedUserID)
|
||||
}
|
||||
}
|
||||
return s.createPhoneCode(ctx, phone)
|
||||
return s.createPhoneCode(ctx, phone, issuedUserID)
|
||||
}
|
||||
|
||||
func (s *Service) createPhoneCode(ctx context.Context, phone string) (string, error) {
|
||||
func (s *Service) currentPhoneOwner(ctx context.Context, phone string) (domain.User, bool, error) {
|
||||
if s == nil || s.users == nil {
|
||||
return domain.User{}, false, fmt.Errorf("user store is not configured")
|
||||
}
|
||||
return s.users.ByPhone(ctx, phone)
|
||||
}
|
||||
|
||||
func (s *Service) issuedOwnerMatches(ctx context.Context, phone string, issuedUserID int64) (bool, error) {
|
||||
current, found, err := s.currentPhoneOwner(ctx, phone)
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
currentUserID := int64(0)
|
||||
if found {
|
||||
currentUserID = current.ID
|
||||
}
|
||||
return currentUserID == issuedUserID, nil
|
||||
}
|
||||
|
||||
func (s *Service) ensureIssuedOwnerAfterSet(ctx context.Context, hash string, rec store.PhoneCode) error {
|
||||
matches, err := s.issuedOwnerMatches(ctx, rec.Phone, rec.IssuedUserID)
|
||||
if err == nil && matches {
|
||||
return nil
|
||||
}
|
||||
cause := err
|
||||
if cause == nil {
|
||||
cause = ErrCodeInvalid
|
||||
}
|
||||
return s.rollbackUndeliveredCode(ctx, hash, cause)
|
||||
}
|
||||
|
||||
func (s *Service) invalidateLoginCodeDetached(ctx context.Context, hash, phone string) {
|
||||
cleanupCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), loginCodeRollbackTimeout)
|
||||
defer cancel()
|
||||
_, _ = s.codes.InvalidateLoginCode(cleanupCtx, hash, phone)
|
||||
}
|
||||
|
||||
func (s *Service) createPhoneCode(ctx context.Context, phone string, existingUserID int64) (string, error) {
|
||||
hash, err := randomHex(8)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if err := s.codes.Set(ctx, hash, store.PhoneCode{
|
||||
Phone: phone,
|
||||
Code: s.fixedCode,
|
||||
Channel: codeChannelPhone,
|
||||
MaxAttempts: s.codeMaxAttempts,
|
||||
Version: store.PhoneCodeVersionCurrent,
|
||||
IssuedUserID: existingUserID,
|
||||
Phone: phone,
|
||||
Code: s.fixedCode,
|
||||
Channel: codeChannelPhone,
|
||||
MaxAttempts: s.codeMaxAttempts,
|
||||
}, s.codeTTL); err != nil {
|
||||
return "", fmt.Errorf("store code: %w", err)
|
||||
}
|
||||
rec := store.PhoneCode{Phone: phone, IssuedUserID: existingUserID}
|
||||
if err := s.ensureIssuedOwnerAfterSet(ctx, hash, rec); err != nil {
|
||||
return "", err
|
||||
}
|
||||
// 新手机号还没有 owner/dialog,只能在 SignUp 创建用户后写第一条
|
||||
// 777000 消息。已有账号则必须在 sendCode RPC 返回前把 app-code
|
||||
// 作为普通 incoming message + durable update/outbox 提交;登录成功不再补发。
|
||||
if existingUserID == 0 {
|
||||
return hash, nil
|
||||
}
|
||||
if err := s.deliverLoginCode(ctx, existingUserID, hash, s.fixedCode); err != nil {
|
||||
return "", s.rollbackUndeliveredCode(ctx, hash, err)
|
||||
}
|
||||
if err := s.ensureIssuedOwnerAfterSet(ctx, hash, rec); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return hash, nil
|
||||
}
|
||||
|
||||
func (s *Service) createSetupRequiredCode(ctx context.Context, phone string) (string, error) {
|
||||
func (s *Service) deliverLoginCode(ctx context.Context, userID int64, phoneCodeHash, code string) error {
|
||||
if s.loginCodeDelivery == nil {
|
||||
return ErrLoginCodeDeliveryUnavailable
|
||||
}
|
||||
now := time.Now()
|
||||
if _, err := s.loginCodeDelivery.DeliverLoginCodeMessage(ctx, domain.LoginCodeDeliveryRequest{
|
||||
UserID: userID,
|
||||
PhoneCodeHash: phoneCodeHash,
|
||||
Code: code,
|
||||
Date: int(now.Unix()),
|
||||
ExpiresAt: now.Add(s.codeTTL).Unix(),
|
||||
}); err != nil {
|
||||
return errors.Join(ErrLoginCodeDeliveryFailed, err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *Service) rollbackUndeliveredCode(ctx context.Context, phoneCodeHash string, cause error) error {
|
||||
// lib/pq can report an I/O failure after COMMIT reached PostgreSQL. In that
|
||||
// state deleting the code could turn an already delivered 777000 message
|
||||
// into an unusable login attempt. Preserve it and let the delivery receipt
|
||||
// make the retry idempotent.
|
||||
if errors.Is(cause, domain.ErrLoginCodeDeliveryCommitAmbiguous) {
|
||||
return cause
|
||||
}
|
||||
cleanupCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), loginCodeRollbackTimeout)
|
||||
defer cancel()
|
||||
if err := s.codes.Del(cleanupCtx, phoneCodeHash); err != nil {
|
||||
return errors.Join(cause, fmt.Errorf("rollback undelivered login code: %w", err))
|
||||
}
|
||||
return cause
|
||||
}
|
||||
|
||||
func (s *Service) createSetupRequiredCode(ctx context.Context, phone string, issuedUserID int64) (string, error) {
|
||||
hash, err := randomHex(8)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if err := s.codes.Set(ctx, hash, store.PhoneCode{
|
||||
Phone: phone,
|
||||
Channel: codeChannelEmailSetupRequired,
|
||||
MaxAttempts: s.codeMaxAttempts,
|
||||
Version: store.PhoneCodeVersionCurrent,
|
||||
IssuedUserID: issuedUserID,
|
||||
Phone: phone,
|
||||
Channel: codeChannelEmailSetupRequired,
|
||||
MaxAttempts: s.codeMaxAttempts,
|
||||
}, s.codeTTL); err != nil {
|
||||
return "", fmt.Errorf("store code: %w", err)
|
||||
}
|
||||
if err := s.ensureIssuedOwnerAfterSet(ctx, hash, store.PhoneCode{Phone: phone, IssuedUserID: issuedUserID}); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return hash, nil
|
||||
}
|
||||
|
||||
func (s *Service) createEmailLoginCode(ctx context.Context, phone, email string) (string, error) {
|
||||
func (s *Service) createEmailLoginCode(ctx context.Context, phone, email string, issuedUserID int64) (string, error) {
|
||||
hash, err := randomHex(8)
|
||||
if err != nil {
|
||||
return "", err
|
||||
|
|
@ -327,22 +451,28 @@ func (s *Service) createEmailLoginCode(ctx context.Context, phone, email string)
|
|||
return "", err
|
||||
}
|
||||
rec := store.PhoneCode{
|
||||
Phone: phone,
|
||||
Code: code,
|
||||
Channel: codeChannelEmailLogin,
|
||||
Email: strings.TrimSpace(email),
|
||||
MaxAttempts: s.codeMaxAttempts,
|
||||
Version: store.PhoneCodeVersionCurrent,
|
||||
IssuedUserID: issuedUserID,
|
||||
Phone: phone,
|
||||
Code: code,
|
||||
Channel: codeChannelEmailLogin,
|
||||
Email: strings.TrimSpace(email),
|
||||
MaxAttempts: s.codeMaxAttempts,
|
||||
}
|
||||
if err := s.codes.Set(ctx, hash, rec, s.codeTTL); err != nil {
|
||||
return "", fmt.Errorf("store email code: %w", err)
|
||||
}
|
||||
if err := s.ensureIssuedOwnerAfterSet(ctx, hash, rec); err != nil {
|
||||
return "", err
|
||||
}
|
||||
if s.loginEmailSender == nil {
|
||||
_ = s.codes.Del(ctx, hash)
|
||||
return "", fmt.Errorf("login email sender is not configured")
|
||||
return "", s.rollbackUndeliveredCode(ctx, hash, fmt.Errorf("login email sender is not configured"))
|
||||
}
|
||||
if err := s.loginEmailSender.SendLoginCode(ctx, rec.Email, code, s.codeTTL); err != nil {
|
||||
_ = s.codes.Del(ctx, hash)
|
||||
return "", fmt.Errorf("send login email code: %w", err)
|
||||
return "", s.rollbackUndeliveredCode(ctx, hash, fmt.Errorf("send login email code: %w", err))
|
||||
}
|
||||
if err := s.ensureIssuedOwnerAfterSet(ctx, hash, rec); err != nil {
|
||||
return "", err
|
||||
}
|
||||
return hash, nil
|
||||
}
|
||||
|
|
@ -396,20 +526,53 @@ func (s *Service) resendCode(ctx context.Context, authKeyID [8]byte, phone, phon
|
|||
if rec.Phone != phone {
|
||||
return "", ErrCodeInvalid
|
||||
}
|
||||
if rec.Purpose == store.PhoneCodePurposeChangePhone && (authKeyID == ([8]byte{}) || rec.AuthKeyID != authKeyID) {
|
||||
if rec.Purpose == store.PhoneCodePurposeChangePhone {
|
||||
if authKeyID == ([8]byte{}) || rec.AuthKeyID != authKeyID {
|
||||
return "", ErrCodeInvalid
|
||||
}
|
||||
consumed, ok, err := s.codes.ConsumeScoped(ctx, phoneCodeHash, rec.Scope())
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if !ok {
|
||||
return "", ErrCodeExpired
|
||||
}
|
||||
return s.recreateChangePhoneCode(ctx, consumed)
|
||||
}
|
||||
if rec.Version != store.PhoneCodeVersionCurrent {
|
||||
_, _, _ = s.codes.TakeLoginCode(ctx, phoneCodeHash, phone)
|
||||
return "", ErrCodeExpired
|
||||
}
|
||||
if matches, err := s.issuedOwnerMatches(ctx, phone, rec.IssuedUserID); err != nil {
|
||||
return "", err
|
||||
} else if !matches {
|
||||
s.invalidateLoginCodeDetached(ctx, phoneCodeHash, phone)
|
||||
return "", ErrCodeInvalid
|
||||
}
|
||||
_ = s.codes.Del(ctx, phoneCodeHash)
|
||||
if rec.Purpose == store.PhoneCodePurposeChangePhone {
|
||||
return s.recreateChangePhoneCode(ctx, rec)
|
||||
consumed, ok, err := s.codes.TakeLoginCode(ctx, phoneCodeHash, phone)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if !ok {
|
||||
return "", ErrCodeExpired
|
||||
}
|
||||
rec = consumed
|
||||
if matches, err := s.issuedOwnerMatches(ctx, phone, rec.IssuedUserID); err != nil {
|
||||
return "", err
|
||||
} else if !matches {
|
||||
s.invalidateLoginCodeDetached(ctx, phoneCodeHash, phone)
|
||||
return "", ErrCodeInvalid
|
||||
}
|
||||
if rec.Channel == codeChannelEmailLogin && strings.TrimSpace(rec.Email) != "" {
|
||||
return s.createEmailLoginCode(ctx, phone, rec.Email)
|
||||
return s.createEmailLoginCode(ctx, phone, rec.Email, rec.IssuedUserID)
|
||||
}
|
||||
if rec.Channel == codeChannelEmailSetupRequired {
|
||||
return s.createSetupRequiredCode(ctx, phone)
|
||||
return s.createSetupRequiredCode(ctx, phone, rec.IssuedUserID)
|
||||
}
|
||||
return s.SendCode(ctx, phone)
|
||||
if rec.Channel != codeChannelPhone {
|
||||
return "", ErrCodeInvalid
|
||||
}
|
||||
return s.createPhoneCode(ctx, phone, rec.IssuedUserID)
|
||||
}
|
||||
|
||||
func (s *Service) recreateChangePhoneCode(ctx context.Context, rec store.PhoneCode) (string, error) {
|
||||
|
|
@ -439,6 +602,75 @@ func (s *Service) CancelCodeForAuthKey(ctx context.Context, authKeyID [8]byte, p
|
|||
return s.cancelCode(ctx, authKeyID, phone, phoneCodeHash)
|
||||
}
|
||||
|
||||
// ConsumeLoginEmailReset authorizes auth.resetLoginEmail with the exact
|
||||
// email-login hash previously issued for this phone owner. Possession of only
|
||||
// a phone number is never sufficient to remove an authentication factor.
|
||||
func (s *Service) ConsumeLoginEmailReset(ctx context.Context, phone, phoneCodeHash string) (int64, error) {
|
||||
phone = normalizePhone(phone)
|
||||
rec, found, err := s.codes.Get(ctx, phoneCodeHash)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if !found {
|
||||
return 0, ErrCodeExpired
|
||||
}
|
||||
if rec.Version != store.PhoneCodeVersionCurrent {
|
||||
_, _, _ = s.codes.TakeLoginCode(ctx, phoneCodeHash, phone)
|
||||
return 0, ErrCodeExpired
|
||||
}
|
||||
if rec.Purpose != "" || rec.Phone != phone || rec.Channel != codeChannelEmailLogin || rec.SignUpVerified {
|
||||
return 0, ErrCodeInvalid
|
||||
}
|
||||
before, beforeFound, err := s.currentPhoneOwner(ctx, phone)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if !beforeFound || systemUserLoginForbidden(before) || rec.IssuedUserID == 0 || rec.IssuedUserID != before.ID {
|
||||
s.invalidateLoginCodeDetached(ctx, phoneCodeHash, phone)
|
||||
return 0, ErrCodeInvalid
|
||||
}
|
||||
consumed, consumedOK, err := s.codes.TakeLoginCode(ctx, phoneCodeHash, phone)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if !consumedOK {
|
||||
return 0, ErrCodeExpired
|
||||
}
|
||||
if consumed.Channel != codeChannelEmailLogin || consumed.IssuedUserID != before.ID {
|
||||
return 0, ErrCodeInvalid
|
||||
}
|
||||
after, afterFound, err := s.currentPhoneOwner(ctx, phone)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if !afterFound || after.ID != before.ID {
|
||||
return 0, ErrCodeInvalid
|
||||
}
|
||||
return before.ID, nil
|
||||
}
|
||||
|
||||
// SendPhoneCodeAfterLoginEmailReset issues the replacement app code only for
|
||||
// the exact user selected by ConsumeLoginEmailReset. It deliberately bypasses
|
||||
// SendCode's phone→owner reclassification so an A→B transfer cannot send B a
|
||||
// code and return that hash to A's reset flow.
|
||||
func (s *Service) SendPhoneCodeAfterLoginEmailReset(ctx context.Context, phone string, expectedUserID int64) (string, error) {
|
||||
phone = normalizePhone(phone)
|
||||
if !validPhone(phone) {
|
||||
return "", ErrPhoneNumberInvalid
|
||||
}
|
||||
if expectedUserID == 0 || systemLoginPhoneForbidden(phone) {
|
||||
return "", ErrCodeInvalid
|
||||
}
|
||||
owner, found, err := s.currentPhoneOwner(ctx, phone)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
if !found || owner.ID != expectedUserID || systemUserLoginForbidden(owner) {
|
||||
return "", ErrCodeInvalid
|
||||
}
|
||||
return s.createPhoneCode(ctx, phone, expectedUserID)
|
||||
}
|
||||
|
||||
func (s *Service) cancelCode(ctx context.Context, authKeyID [8]byte, phone, phoneCodeHash string) error {
|
||||
phone = normalizePhone(phone)
|
||||
rec, found, err := s.codes.Get(ctx, phoneCodeHash)
|
||||
|
|
@ -451,10 +683,37 @@ func (s *Service) cancelCode(ctx context.Context, authKeyID [8]byte, phone, phon
|
|||
if rec.Phone != phone {
|
||||
return ErrCodeInvalid
|
||||
}
|
||||
if rec.Purpose == store.PhoneCodePurposeChangePhone && (authKeyID == ([8]byte{}) || rec.AuthKeyID != authKeyID) {
|
||||
if rec.Purpose == store.PhoneCodePurposeChangePhone {
|
||||
if authKeyID == ([8]byte{}) || rec.AuthKeyID != authKeyID {
|
||||
return ErrCodeInvalid
|
||||
}
|
||||
_, consumed, err := s.codes.ConsumeScoped(ctx, phoneCodeHash, rec.Scope())
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if !consumed {
|
||||
return ErrCodeExpired
|
||||
}
|
||||
return nil
|
||||
}
|
||||
if rec.Version != store.PhoneCodeVersionCurrent {
|
||||
_, _, _ = s.codes.TakeLoginCode(ctx, phoneCodeHash, phone)
|
||||
return ErrCodeExpired
|
||||
}
|
||||
if matches, err := s.issuedOwnerMatches(ctx, phone, rec.IssuedUserID); err != nil {
|
||||
return err
|
||||
} else if !matches {
|
||||
s.invalidateLoginCodeDetached(ctx, phoneCodeHash, phone)
|
||||
return ErrCodeInvalid
|
||||
}
|
||||
return s.codes.Del(ctx, phoneCodeHash)
|
||||
_, consumed, err := s.codes.TakeLoginCode(ctx, phoneCodeHash, phone)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if !consumed {
|
||||
return ErrCodeExpired
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// SignIn 校验验证码并尝试登录。
|
||||
|
|
@ -464,102 +723,164 @@ func (s *Service) SignIn(ctx context.Context, auth domain.Authorization, phone,
|
|||
if systemLoginPhoneForbidden(phone) {
|
||||
return domain.User{}, domain.Message{}, false, ErrSystemUserLoginForbidden
|
||||
}
|
||||
rec, found, err := s.codes.Get(ctx, phoneCodeHash)
|
||||
if err != nil {
|
||||
return domain.User{}, domain.Message{}, false, err
|
||||
}
|
||||
if !found {
|
||||
return domain.User{}, domain.Message{}, false, ErrCodeExpired
|
||||
}
|
||||
if rec.Phone != phone || rec.Channel == codeChannelEmailSetupRequired {
|
||||
return domain.User{}, domain.Message{}, false, ErrCodeInvalid
|
||||
}
|
||||
if rec.Channel == codeChannelEmailLogin {
|
||||
return domain.User{}, domain.Message{}, false, ErrCodeInvalid
|
||||
}
|
||||
if rec.Code != code {
|
||||
return domain.User{}, domain.Message{}, false, s.rejectCode(ctx, phoneCodeHash, rec, ErrCodeInvalid)
|
||||
}
|
||||
|
||||
existing, found, err := s.users.ByPhone(ctx, phone)
|
||||
if err != nil {
|
||||
return domain.User{}, domain.Message{}, false, err
|
||||
}
|
||||
if !found {
|
||||
return domain.User{}, domain.Message{}, true, nil // 验证码对、但需注册
|
||||
}
|
||||
return s.finishSignIn(ctx, auth, existing, phoneCodeHash, rec.Code)
|
||||
}
|
||||
|
||||
// SignInWithEmail 处理带 email_verification 的 auth.signIn:账号设置了登录邮箱后,新设备
|
||||
// 的验证码改投递到邮箱,客户端凭邮箱码(而非短信码)登录。开启真实登录邮箱后必须匹配
|
||||
// 随机邮箱码;未开启该特性时仅保留旧开发路径的任意非空兼容。仍校验 phone_code_hash
|
||||
// 有效、手机号匹配,并与短信登录共用 2FA 门控——即便走邮箱验证,开启了两步验证的账号
|
||||
// 同样会停在 SESSION_PASSWORD_NEEDED。
|
||||
func (s *Service) SignInWithEmail(ctx context.Context, auth domain.Authorization, phone, phoneCodeHash, code string) (domain.User, domain.Message, bool, error) {
|
||||
phone = normalizePhone(phone)
|
||||
if systemLoginPhoneForbidden(phone) {
|
||||
return domain.User{}, domain.Message{}, false, ErrSystemUserLoginForbidden
|
||||
}
|
||||
rec, found, err := s.codes.Get(ctx, phoneCodeHash)
|
||||
if err != nil {
|
||||
return domain.User{}, domain.Message{}, false, err
|
||||
}
|
||||
if !found {
|
||||
return domain.User{}, domain.Message{}, false, ErrCodeExpired
|
||||
}
|
||||
if rec.Phone != phone {
|
||||
return domain.User{}, domain.Message{}, false, ErrCodeInvalid
|
||||
}
|
||||
if rec.Channel != codeChannelEmailLogin {
|
||||
if s.loginEmailEnabled {
|
||||
return domain.User{}, domain.Message{}, false, ErrCodeInvalid
|
||||
}
|
||||
if strings.TrimSpace(code) == "" {
|
||||
return domain.User{}, domain.Message{}, false, ErrCodeInvalid
|
||||
}
|
||||
} else if rec.Code != strings.TrimSpace(code) {
|
||||
return domain.User{}, domain.Message{}, false, s.rejectCode(ctx, phoneCodeHash, rec, ErrCodeInvalid)
|
||||
}
|
||||
existing, found, err := s.users.ByPhone(ctx, phone)
|
||||
_, existing, found, err := s.verifyLoginCode(ctx, phone, phoneCodeHash, code, false)
|
||||
if err != nil {
|
||||
return domain.User{}, domain.Message{}, false, err
|
||||
}
|
||||
if !found {
|
||||
return domain.User{}, domain.Message{}, true, nil
|
||||
}
|
||||
return s.finishSignIn(ctx, auth, existing, phoneCodeHash, rec.Code)
|
||||
return s.finishSignIn(ctx, auth, existing)
|
||||
}
|
||||
|
||||
// SignInWithEmail 处理带 email_verification 的 auth.signIn:账号设置了登录邮箱后,新设备
|
||||
// 的验证码改投递到邮箱,客户端凭邮箱码(而非短信码)登录。开启真实登录邮箱后必须匹配
|
||||
// 随机邮箱码;未开启该特性时仍允许旧客户端把 phone channel 放进
|
||||
// email_verification,但必须精确匹配该 phone code,不能再接受任意非空值。
|
||||
// 两条路径共用 owner 绑定、原子尝试计数与 2FA 门控。
|
||||
func (s *Service) SignInWithEmail(ctx context.Context, auth domain.Authorization, phone, phoneCodeHash, code string) (domain.User, domain.Message, bool, error) {
|
||||
phone = normalizePhone(phone)
|
||||
if systemLoginPhoneForbidden(phone) {
|
||||
return domain.User{}, domain.Message{}, false, ErrSystemUserLoginForbidden
|
||||
}
|
||||
_, existing, found, err := s.verifyLoginCode(ctx, phone, phoneCodeHash, strings.TrimSpace(code), true)
|
||||
if err != nil {
|
||||
return domain.User{}, domain.Message{}, false, err
|
||||
}
|
||||
if !found {
|
||||
return domain.User{}, domain.Message{}, true, nil
|
||||
}
|
||||
return s.finishSignIn(ctx, auth, existing)
|
||||
}
|
||||
|
||||
// verifyLoginCode closes the login-code state transition around one atomic
|
||||
// CodeStore verification. The phone owner is read both before and after that
|
||||
// linearization point. A hash issued for an unregistered number therefore can
|
||||
// never authorize whichever account happens to acquire that number later.
|
||||
func (s *Service) verifyLoginCode(ctx context.Context, phone, phoneCodeHash, code string, emailPath bool) (store.PhoneCode, domain.User, bool, error) {
|
||||
rec, found, err := s.codes.Get(ctx, phoneCodeHash)
|
||||
if err != nil {
|
||||
return store.PhoneCode{}, domain.User{}, false, err
|
||||
}
|
||||
if !found {
|
||||
return store.PhoneCode{}, domain.User{}, false, ErrCodeExpired
|
||||
}
|
||||
if rec.Version != store.PhoneCodeVersionCurrent {
|
||||
_, _, _ = s.codes.TakeLoginCode(ctx, phoneCodeHash, phone)
|
||||
return store.PhoneCode{}, domain.User{}, false, ErrCodeExpired
|
||||
}
|
||||
if rec.Phone != phone || rec.Purpose != "" {
|
||||
return store.PhoneCode{}, domain.User{}, false, ErrCodeInvalid
|
||||
}
|
||||
channelAllowed := rec.Channel == codeChannelPhone && !emailPath
|
||||
if emailPath {
|
||||
channelAllowed = rec.Channel == codeChannelEmailLogin || (!s.loginEmailEnabled && rec.Channel == codeChannelPhone)
|
||||
}
|
||||
if !channelAllowed {
|
||||
return store.PhoneCode{}, domain.User{}, false, ErrCodeInvalid
|
||||
}
|
||||
|
||||
before, beforeFound, err := s.currentPhoneOwner(ctx, phone)
|
||||
if err != nil {
|
||||
return store.PhoneCode{}, domain.User{}, false, err
|
||||
}
|
||||
beforeUserID := int64(0)
|
||||
if beforeFound {
|
||||
if systemUserLoginForbidden(before) {
|
||||
return store.PhoneCode{}, domain.User{}, false, ErrSystemUserLoginForbidden
|
||||
}
|
||||
beforeUserID = before.ID
|
||||
}
|
||||
if rec.IssuedUserID != beforeUserID {
|
||||
s.invalidateLoginCodeDetached(ctx, phoneCodeHash, phone)
|
||||
return store.PhoneCode{}, domain.User{}, false, ErrCodeInvalid
|
||||
}
|
||||
// A verified sign-up marker may precede auth.signIn on the email-setup
|
||||
// path, and a normal signIn response can be lost and retried. The marker is
|
||||
// already the durable authorization fact; return signUpRequired
|
||||
// idempotently without asking CodeStore to verify it a second time.
|
||||
if rec.SignUpVerified {
|
||||
if beforeFound || rec.IssuedUserID != 0 || subtle.ConstantTimeCompare([]byte(rec.Code), []byte(code)) != 1 {
|
||||
return store.PhoneCode{}, domain.User{}, false, ErrCodeInvalid
|
||||
}
|
||||
after, afterFound, err := s.currentPhoneOwner(ctx, phone)
|
||||
if err != nil {
|
||||
return store.PhoneCode{}, domain.User{}, false, err
|
||||
}
|
||||
if afterFound || after.ID != 0 {
|
||||
s.invalidateLoginCodeDetached(ctx, phoneCodeHash, phone)
|
||||
return store.PhoneCode{}, domain.User{}, false, ErrCodeInvalid
|
||||
}
|
||||
return rec, domain.User{}, false, nil
|
||||
}
|
||||
|
||||
result, err := s.codes.VerifyLogin(ctx, phoneCodeHash, phone, code, !beforeFound, s.codeMaxAttempts)
|
||||
if err != nil {
|
||||
return store.PhoneCode{}, domain.User{}, false, err
|
||||
}
|
||||
after, afterFound, ownerErr := s.currentPhoneOwner(ctx, phone)
|
||||
if ownerErr != nil {
|
||||
return store.PhoneCode{}, domain.User{}, false, ownerErr
|
||||
}
|
||||
afterUserID := int64(0)
|
||||
if afterFound {
|
||||
afterUserID = after.ID
|
||||
}
|
||||
recordOwnerMismatch := result.Status != store.LoginCodeVerifyMissing && result.Record.IssuedUserID != rec.IssuedUserID
|
||||
if beforeUserID != afterUserID || recordOwnerMismatch {
|
||||
// keepForSignUp may have left a verified marker behind. Remove it on
|
||||
// owner drift so a later transfer-back cannot resurrect authorization.
|
||||
s.invalidateLoginCodeDetached(ctx, phoneCodeHash, phone)
|
||||
return store.PhoneCode{}, domain.User{}, false, ErrCodeInvalid
|
||||
}
|
||||
switch result.Status {
|
||||
case store.LoginCodeVerifyMissing:
|
||||
return store.PhoneCode{}, domain.User{}, false, ErrCodeExpired
|
||||
case store.LoginCodeVerifyInvalid:
|
||||
return store.PhoneCode{}, domain.User{}, false, ErrCodeInvalid
|
||||
case store.LoginCodeVerifyAccepted:
|
||||
if result.Record.Version != store.PhoneCodeVersionCurrent || result.Record.Phone != phone || result.Record.IssuedUserID != afterUserID {
|
||||
return store.PhoneCode{}, domain.User{}, false, ErrCodeInvalid
|
||||
}
|
||||
if afterFound && systemUserLoginForbidden(after) {
|
||||
return store.PhoneCode{}, domain.User{}, false, ErrSystemUserLoginForbidden
|
||||
}
|
||||
return result.Record, after, afterFound, nil
|
||||
default:
|
||||
return store.PhoneCode{}, domain.User{}, false, ErrCodeInvalid
|
||||
}
|
||||
}
|
||||
|
||||
// finishSignIn 是短信/邮箱两条登录路径在「验证码已通过、用户已存在」之后的共用收尾:
|
||||
// 处理 2FA password_pending 绑定、写登录消息、消费验证码。
|
||||
func (s *Service) finishSignIn(ctx context.Context, auth domain.Authorization, existing domain.User, phoneCodeHash, loginCode string) (domain.User, domain.Message, bool, error) {
|
||||
// 验证码已由 VerifyLogin 原子消费;这里只处理 2FA password_pending 绑定。已有账号的 app-code 消息已在
|
||||
// SendCode/ResendCode 返回前持久化与入 outbox,这里绝不能再创建或补发;
|
||||
// 否则未完成登录/2FA 的真实验证码反而不会及时到达旧设备。
|
||||
func (s *Service) finishSignIn(ctx context.Context, auth domain.Authorization, existing domain.User) (domain.User, domain.Message, bool, error) {
|
||||
if systemUserLoginForbidden(existing) {
|
||||
_ = s.codes.Del(ctx, phoneCodeHash)
|
||||
return domain.User{}, domain.Message{}, false, ErrSystemUserLoginForbidden
|
||||
}
|
||||
// 开启两步验证的账号:把授权标记为 password_pending 再写入,业务鉴权据此拒绝该 auth_key,
|
||||
// 直到 auth.checkPassword 通过。绝不能先以完全授权写入再返回 SESSION_PASSWORD_NEEDED,
|
||||
// 否则客户端忽略该错误即可直接调用业务 RPC 绕过两步验证。
|
||||
passwordNeeded := s.passwordNeeded(ctx, existing.ID)
|
||||
passwordNeeded, err := s.passwordNeeded(ctx, existing.ID)
|
||||
if err != nil {
|
||||
// Password state is part of the authentication decision. Treat store
|
||||
// failures as fail-closed and leave the auth key entirely unbound.
|
||||
return domain.User{}, domain.Message{}, false, err
|
||||
}
|
||||
auth.PasswordPending = passwordNeeded
|
||||
if err := s.bind(ctx, auth, existing.ID); err != nil {
|
||||
return domain.User{}, domain.Message{}, false, err
|
||||
}
|
||||
if passwordNeeded {
|
||||
_ = s.codes.Del(ctx, phoneCodeHash)
|
||||
return existing, domain.Message{}, false, domain.ErrSessionPasswordNeeded
|
||||
}
|
||||
loginMessage, err := s.recordLoginMessage(ctx, existing.ID, loginCode)
|
||||
if err != nil {
|
||||
return domain.User{}, domain.Message{}, false, err
|
||||
}
|
||||
_ = s.codes.Del(ctx, phoneCodeHash)
|
||||
return existing, loginMessage, false, nil
|
||||
return existing, domain.Message{}, false, nil
|
||||
}
|
||||
|
||||
// SignUp 在 SignIn 判定需注册后创建用户并绑定授权。
|
||||
// signUp 的 TL 请求不带验证码,这里校验 phone_code_hash 仍有效且手机号匹配。
|
||||
// signUp 的 TL 请求不带验证码,因此只消费由正确 SignIn/email setup 原子
|
||||
// 标记过的 hash。直接 SendCode→SignUp 永远不能创建账号。
|
||||
func (s *Service) SignUp(ctx context.Context, auth domain.Authorization, phone, phoneCodeHash, firstName, lastName string) (domain.User, domain.Message, error) {
|
||||
phone = normalizePhone(phone)
|
||||
if !validPhone(phone) {
|
||||
|
|
@ -580,15 +901,48 @@ func (s *Service) SignUp(ctx context.Context, auth domain.Authorization, phone,
|
|||
if !found {
|
||||
return domain.User{}, domain.Message{}, ErrCodeExpired
|
||||
}
|
||||
if rec.Phone != phone {
|
||||
if rec.Version != store.PhoneCodeVersionCurrent {
|
||||
_, _, _ = s.codes.ConsumeSignUpVerified(ctx, phoneCodeHash, phone)
|
||||
return domain.User{}, domain.Message{}, ErrCodeExpired
|
||||
}
|
||||
if rec.Phone != phone || rec.Purpose != "" {
|
||||
return domain.User{}, domain.Message{}, ErrCodeInvalid
|
||||
}
|
||||
if rec.Channel == codeChannelEmailSetupRequired {
|
||||
if !rec.SignUpVerified {
|
||||
return domain.User{}, domain.Message{}, ErrCodeInvalid
|
||||
}
|
||||
if rec.IssuedUserID != 0 {
|
||||
s.invalidateLoginCodeDetached(ctx, phoneCodeHash, phone)
|
||||
return domain.User{}, domain.Message{}, ErrCodeInvalid
|
||||
}
|
||||
if rec.Channel != codeChannelPhone && rec.Channel != codeChannelEmailLogin {
|
||||
return domain.User{}, domain.Message{}, ErrCodeInvalid
|
||||
}
|
||||
if s.loginEmailRequireSetup && !rec.VerifiedEmail && strings.TrimSpace(rec.PendingEmail) == "" {
|
||||
return domain.User{}, domain.Message{}, ErrCodeInvalid
|
||||
}
|
||||
if current, currentFound, err := s.currentPhoneOwner(ctx, phone); err != nil {
|
||||
return domain.User{}, domain.Message{}, err
|
||||
} else if currentFound || current.ID != 0 {
|
||||
s.invalidateLoginCodeDetached(ctx, phoneCodeHash, phone)
|
||||
return domain.User{}, domain.Message{}, ErrCodeInvalid
|
||||
}
|
||||
consumed, consumedOK, err := s.codes.ConsumeSignUpVerified(ctx, phoneCodeHash, phone)
|
||||
if err != nil {
|
||||
return domain.User{}, domain.Message{}, err
|
||||
}
|
||||
if !consumedOK {
|
||||
return domain.User{}, domain.Message{}, ErrCodeExpired
|
||||
}
|
||||
rec = consumed
|
||||
if rec.IssuedUserID != 0 || !rec.SignUpVerified || (rec.Channel != codeChannelPhone && rec.Channel != codeChannelEmailLogin) {
|
||||
return domain.User{}, domain.Message{}, ErrCodeInvalid
|
||||
}
|
||||
if current, currentFound, err := s.currentPhoneOwner(ctx, phone); err != nil {
|
||||
return domain.User{}, domain.Message{}, err
|
||||
} else if currentFound || current.ID != 0 {
|
||||
return domain.User{}, domain.Message{}, ErrCodeInvalid
|
||||
}
|
||||
|
||||
accessHash, err := randomInt64()
|
||||
if err != nil {
|
||||
|
|
@ -610,18 +964,22 @@ func (s *Service) SignUp(ctx context.Context, auth domain.Authorization, phone,
|
|||
return domain.User{}, domain.Message{}, err
|
||||
}
|
||||
if rec.VerifiedEmail && strings.TrimSpace(rec.PendingEmail) != "" && s.loginEmails != nil {
|
||||
if err := s.loginEmails.SetLoginEmailByPhone(ctx, phone, rec.PendingEmail); err != nil {
|
||||
if err := s.loginEmails.SetLoginEmail(ctx, u.ID, rec.PendingEmail); err != nil {
|
||||
return domain.User{}, domain.Message{}, err
|
||||
}
|
||||
}
|
||||
if err := s.bind(ctx, auth, u.ID); err != nil {
|
||||
return domain.User{}, domain.Message{}, err
|
||||
}
|
||||
loginMessage, err := s.recordLoginMessage(ctx, u.ID, rec.Code)
|
||||
if err != nil {
|
||||
return domain.User{}, domain.Message{}, err
|
||||
loginMessage := domain.Message{}
|
||||
// SMTP setup/login codes are secret factors, not 777000 app messages. Only
|
||||
// the normal phone/app-code registration path creates the bootstrap dialog.
|
||||
if rec.Channel == codeChannelPhone {
|
||||
loginMessage, err = s.recordLoginMessage(ctx, u.ID, rec.Code)
|
||||
if err != nil {
|
||||
return domain.User{}, domain.Message{}, err
|
||||
}
|
||||
}
|
||||
_ = s.codes.Del(ctx, phoneCodeHash)
|
||||
return u, loginMessage, nil
|
||||
}
|
||||
|
||||
|
|
@ -865,15 +1223,21 @@ func (s *Service) authorizationsByUserExcept(ctx context.Context, userID int64,
|
|||
|
||||
func (s *Service) bind(ctx context.Context, auth domain.Authorization, userID int64) error {
|
||||
auth.UserID = userID
|
||||
// Bind 是授权切换的持久化状态边界:生产 store 会先清同 auth key 的旧用户
|
||||
// update state,再原子建立新用户 baseline。RPC 层不得在 Bind 成功后清整个 key,
|
||||
// 否则会把刚建立的 retained-floor checkpoint 一并删除。
|
||||
return s.auths.Bind(ctx, auth)
|
||||
}
|
||||
|
||||
func (s *Service) passwordNeeded(ctx context.Context, userID int64) bool {
|
||||
func (s *Service) passwordNeeded(ctx context.Context, userID int64) (bool, error) {
|
||||
if s.passwords == nil {
|
||||
return false
|
||||
return false, nil
|
||||
}
|
||||
settings, found, err := s.passwords.GetByUser(ctx, userID)
|
||||
return err == nil && found && settings.HasPassword
|
||||
if err != nil {
|
||||
return false, err
|
||||
}
|
||||
return found && settings.HasPassword, nil
|
||||
}
|
||||
|
||||
const loginMessageTpl = `Login code: %s. Do not give this code to anyone, even if they say they are from Telegram!
|
||||
|
|
@ -1005,20 +1369,6 @@ func authKeyIDInt64(id [8]byte) int64 {
|
|||
return int64(binary.LittleEndian.Uint64(id[:]))
|
||||
}
|
||||
|
||||
func (s *Service) rejectCode(ctx context.Context, hash string, rec store.PhoneCode, ret error) error {
|
||||
rec.Attempts++
|
||||
max := rec.MaxAttempts
|
||||
if max <= 0 {
|
||||
max = s.codeMaxAttempts
|
||||
}
|
||||
if max > 0 && rec.Attempts >= max {
|
||||
_ = s.codes.Del(ctx, hash)
|
||||
return ret
|
||||
}
|
||||
_ = s.codes.Update(ctx, hash, rec)
|
||||
return ret
|
||||
}
|
||||
|
||||
func normalizePhone(phone string) string {
|
||||
return domain.NormalizePhone(phone)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -178,6 +178,14 @@ func TestPhoneCodeAcceptsTDesktopDigitsOnlySignIn(t *testing.T) {
|
|||
}
|
||||
}
|
||||
|
||||
func verifyCodeForSignUp(t *testing.T, svc *Service, phone, hash, code string) {
|
||||
t.Helper()
|
||||
got, msg, needSignUp, err := svc.SignIn(context.Background(), domain.Authorization{}, phone, hash, code)
|
||||
if err != nil || !needSignUp || got.ID != 0 || msg.ID != 0 {
|
||||
t.Fatalf("SignIn before SignUp user=%+v message=%+v needSignUp=%v err=%v, want empty/empty/true/nil", got, msg, needSignUp, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSystemUserPhoneCannotLoginOrSignUp(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
codes := memory.NewCodeStore()
|
||||
|
|
@ -257,6 +265,7 @@ func TestMultipleAuthKeysKeepSeparateUsers(t *testing.T) {
|
|||
if err != nil {
|
||||
t.Fatalf("SendCode user1: %v", err)
|
||||
}
|
||||
verifyCodeForSignUp(t, svc, "+15550005001", hash1, "12345")
|
||||
user1, _, err := svc.SignUp(ctx, domain.Authorization{AuthKeyID: key1}, "+15550005001", hash1, "One", "")
|
||||
if err != nil {
|
||||
t.Fatalf("SignUp user1: %v", err)
|
||||
|
|
@ -265,6 +274,7 @@ func TestMultipleAuthKeysKeepSeparateUsers(t *testing.T) {
|
|||
if err != nil {
|
||||
t.Fatalf("SendCode user2: %v", err)
|
||||
}
|
||||
verifyCodeForSignUp(t, svc, "+15550005002", hash2, "12345")
|
||||
user2, _, err := svc.SignUp(ctx, domain.Authorization{AuthKeyID: key2}, "+15550005002", hash2, "Two", "")
|
||||
if err != nil {
|
||||
t.Fatalf("SignUp user2: %v", err)
|
||||
|
|
@ -294,6 +304,7 @@ func TestLogOutThenSignInSameAuthKeySwitchesUser(t *testing.T) {
|
|||
if err != nil {
|
||||
t.Fatalf("SendCode user1: %v", err)
|
||||
}
|
||||
verifyCodeForSignUp(t, svc, "+15550006001", hash1, "12345")
|
||||
user1, _, err := svc.SignUp(ctx, domain.Authorization{AuthKeyID: key}, "+15550006001", hash1, "One", "")
|
||||
if err != nil {
|
||||
t.Fatalf("SignUp user1: %v", err)
|
||||
|
|
@ -312,6 +323,7 @@ func TestLogOutThenSignInSameAuthKeySwitchesUser(t *testing.T) {
|
|||
if err != nil {
|
||||
t.Fatalf("SendCode user2: %v", err)
|
||||
}
|
||||
verifyCodeForSignUp(t, svc, "+15550006002", hash2, "12345")
|
||||
user2, _, err := svc.SignUp(ctx, domain.Authorization{AuthKeyID: key}, "+15550006002", hash2, "Two", "")
|
||||
if err != nil {
|
||||
t.Fatalf("SignUp user2: %v", err)
|
||||
|
|
@ -337,6 +349,7 @@ func TestResetAuthorizationDeletesProtocolAuthKey(t *testing.T) {
|
|||
if err != nil {
|
||||
t.Fatalf("SendCode: %v", err)
|
||||
}
|
||||
verifyCodeForSignUp(t, svc, "+15550007001", hash, "12345")
|
||||
u, _, err := svc.SignUp(ctx, domain.Authorization{AuthKeyID: key}, "+15550007001", hash, "One", "")
|
||||
if err != nil {
|
||||
t.Fatalf("SignUp: %v", err)
|
||||
|
|
@ -375,6 +388,7 @@ func TestResetAuthorizationsDeletesOnlyRevokedProtocolAuthKeys(t *testing.T) {
|
|||
if err != nil {
|
||||
t.Fatalf("SendCode: %v", err)
|
||||
}
|
||||
verifyCodeForSignUp(t, svc, "+15550007002", hash, "12345")
|
||||
u, _, err := svc.SignUp(ctx, domain.Authorization{AuthKeyID: keep}, "+15550007002", hash, "Two", "")
|
||||
if err != nil {
|
||||
t.Fatalf("SignUp: %v", err)
|
||||
|
|
@ -399,16 +413,27 @@ func TestSignUpWritesOfficialLoginMessage(t *testing.T) {
|
|||
ctx := context.Background()
|
||||
dialogs := memory.NewDialogStore()
|
||||
messages := memory.NewMessageStore(dialogs)
|
||||
svc := NewService(memory.NewUserStore(), memory.NewAuthorizationStore(), memory.NewCodeStore(), nil, nil, "12345", WithLoginMessages(messages, dialogs))
|
||||
delivery := &captureLoginCodeDelivery{}
|
||||
svc := NewService(memory.NewUserStore(), memory.NewAuthorizationStore(), memory.NewCodeStore(), nil, nil, "12345",
|
||||
WithLoginMessages(messages, dialogs),
|
||||
WithLoginCodeDelivery(delivery),
|
||||
)
|
||||
|
||||
hash, err := svc.SendCode(ctx, "+15550004311")
|
||||
if err != nil {
|
||||
t.Fatalf("SendCode: %v", err)
|
||||
}
|
||||
if len(delivery.requests) != 0 {
|
||||
t.Fatalf("unregistered SendCode delivered before user exists: %+v", delivery.requests)
|
||||
}
|
||||
verifyCodeForSignUp(t, svc, "+15550004311", hash, "12345")
|
||||
u, msg, err := svc.SignUp(ctx, domain.Authorization{}, "+15550004311", hash, "Test", "User")
|
||||
if err != nil {
|
||||
t.Fatalf("SignUp: %v", err)
|
||||
}
|
||||
if len(delivery.requests) != 0 {
|
||||
t.Fatalf("SignUp unexpectedly used existing-account delivery: %+v", delivery.requests)
|
||||
}
|
||||
|
||||
list, err := dialogs.ListByUser(ctx, u.ID, domain.DialogFilter{Limit: 10})
|
||||
if err != nil {
|
||||
|
|
@ -431,17 +456,23 @@ func TestSignUpWritesOfficialLoginMessage(t *testing.T) {
|
|||
}
|
||||
}
|
||||
|
||||
func TestSignInLoginMessagePreservesOfficialDialogReadWatermark(t *testing.T) {
|
||||
func TestSendCodeLoginMessagePreservesOfficialDialogReadWatermark(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
dialogs := memory.NewDialogStore()
|
||||
messages := memory.NewMessageStore(dialogs)
|
||||
svc := NewService(memory.NewUserStore(), memory.NewAuthorizationStore(), memory.NewCodeStore(), nil, nil, "12345", WithLoginMessages(messages, dialogs))
|
||||
events := memory.NewUpdateEventStore()
|
||||
delivery := memory.NewLoginCodeDeliveryStore(messages, events)
|
||||
svc := NewService(memory.NewUserStore(), memory.NewAuthorizationStore(), memory.NewCodeStore(), nil, nil, "12345",
|
||||
WithLoginMessages(messages, dialogs),
|
||||
WithLoginCodeDelivery(delivery),
|
||||
)
|
||||
phone := "+15550004312"
|
||||
|
||||
hash, err := svc.SendCode(ctx, phone)
|
||||
if err != nil {
|
||||
t.Fatalf("SendCode signup: %v", err)
|
||||
}
|
||||
verifyCodeForSignUp(t, svc, phone, hash, "12345")
|
||||
u, first, err := svc.SignUp(ctx, domain.Authorization{}, phone, hash, "Test", "User")
|
||||
if err != nil {
|
||||
t.Fatalf("SignUp: %v", err)
|
||||
|
|
@ -452,15 +483,6 @@ func TestSignInLoginMessagePreservesOfficialDialogReadWatermark(t *testing.T) {
|
|||
} else if read.MaxID != first.ID || read.StillUnreadCount != 0 {
|
||||
t.Fatalf("read first login message = %+v, want max_id %d unread 0", read, first.ID)
|
||||
}
|
||||
|
||||
hash, err = svc.SendCode(ctx, phone)
|
||||
if err != nil {
|
||||
t.Fatalf("SendCode signin second: %v", err)
|
||||
}
|
||||
_, second, needSignUp, err := svc.SignIn(ctx, domain.Authorization{}, phone, hash, "12345")
|
||||
if err != nil || needSignUp {
|
||||
t.Fatalf("SignIn second needSignUp=%v err=%v", needSignUp, err)
|
||||
}
|
||||
assertOfficialDialog := func(wantTop, wantRead, wantUnread int) {
|
||||
t.Helper()
|
||||
list, err := dialogs.ListByUser(ctx, u.ID, domain.DialogFilter{Limit: 10})
|
||||
|
|
@ -475,23 +497,67 @@ func TestSignInLoginMessagePreservesOfficialDialogReadWatermark(t *testing.T) {
|
|||
t.Fatalf("dialog = %+v, want top=%d read=%d unread=%d", got, wantTop, wantRead, wantUnread)
|
||||
}
|
||||
}
|
||||
latestLoginMessage := func(wantCount int) domain.Message {
|
||||
t.Helper()
|
||||
history, err := messages.ListByUser(ctx, u.ID, domain.MessageFilter{
|
||||
HasPeer: true,
|
||||
Peer: peer,
|
||||
Limit: 10,
|
||||
})
|
||||
if err != nil || len(history.Messages) != wantCount {
|
||||
t.Fatalf("official history count=%d err=%v, want %d", len(history.Messages), err, wantCount)
|
||||
}
|
||||
latest := history.Messages[0]
|
||||
for _, msg := range history.Messages[1:] {
|
||||
if msg.ID > latest.ID {
|
||||
latest = msg
|
||||
}
|
||||
}
|
||||
return latest
|
||||
}
|
||||
|
||||
hash, err = svc.SendCode(ctx, phone)
|
||||
if err != nil {
|
||||
t.Fatalf("SendCode signin second: %v", err)
|
||||
}
|
||||
second := latestLoginMessage(2)
|
||||
// 核心时序:SendCode 返回时 message/dialog/unread 已提交,尚未 SignIn。
|
||||
assertOfficialDialog(second.ID, first.ID, 1)
|
||||
_, signInMessage, needSignUp, err := svc.SignIn(ctx, domain.Authorization{}, phone, hash, "12345")
|
||||
if err != nil || needSignUp {
|
||||
t.Fatalf("SignIn second needSignUp=%v err=%v", needSignUp, err)
|
||||
}
|
||||
if signInMessage.ID != 0 {
|
||||
t.Fatalf("SignIn second returned a late login message %+v", signInMessage)
|
||||
}
|
||||
assertOfficialDialog(second.ID, first.ID, 1)
|
||||
|
||||
hash, err = svc.SendCode(ctx, phone)
|
||||
if err != nil {
|
||||
t.Fatalf("SendCode signin third: %v", err)
|
||||
}
|
||||
_, third, needSignUp, err := svc.SignIn(ctx, domain.Authorization{}, phone, hash, "12345")
|
||||
third := latestLoginMessage(3)
|
||||
assertOfficialDialog(third.ID, first.ID, 2)
|
||||
_, signInMessage, needSignUp, err = svc.SignIn(ctx, domain.Authorization{}, phone, hash, "12345")
|
||||
if err != nil || needSignUp {
|
||||
t.Fatalf("SignIn third needSignUp=%v err=%v", needSignUp, err)
|
||||
}
|
||||
if signInMessage.ID != 0 {
|
||||
t.Fatalf("SignIn third returned a late login message %+v", signInMessage)
|
||||
}
|
||||
assertOfficialDialog(third.ID, first.ID, 2)
|
||||
}
|
||||
|
||||
func TestSignInExistingTwoFactorAccountNeedsPassword(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
passwords := memory.NewPasswordStore()
|
||||
svc := NewService(memory.NewUserStore(), memory.NewAuthorizationStore(), memory.NewCodeStore(), nil, nil, "12345", WithPasswords(passwords))
|
||||
dialogs := memory.NewDialogStore()
|
||||
messages := memory.NewMessageStore(dialogs)
|
||||
delivery := memory.NewLoginCodeDeliveryStore(messages, memory.NewUpdateEventStore())
|
||||
svc := NewService(memory.NewUserStore(), memory.NewAuthorizationStore(), memory.NewCodeStore(), nil, nil, "12345",
|
||||
WithPasswords(passwords),
|
||||
WithLoginCodeDelivery(delivery),
|
||||
)
|
||||
var key [8]byte
|
||||
key[0] = 7
|
||||
|
||||
|
|
@ -499,6 +565,7 @@ func TestSignInExistingTwoFactorAccountNeedsPassword(t *testing.T) {
|
|||
if err != nil {
|
||||
t.Fatalf("SendCode signup: %v", err)
|
||||
}
|
||||
verifyCodeForSignUp(t, svc, "+15550004312", hash, "12345")
|
||||
u, _, err := svc.SignUp(ctx, domain.Authorization{AuthKeyID: key}, "+15550004312", hash, "Two", "Factor")
|
||||
if err != nil {
|
||||
t.Fatalf("SignUp: %v", err)
|
||||
|
|
@ -514,13 +581,16 @@ func TestSignInExistingTwoFactorAccountNeedsPassword(t *testing.T) {
|
|||
if err != nil {
|
||||
t.Fatalf("SendCode signin: %v", err)
|
||||
}
|
||||
got, _, needSignUp, err := svc.SignIn(ctx, domain.Authorization{AuthKeyID: key}, "+15550004312", hash, "12345")
|
||||
got, signInMessage, needSignUp, err := svc.SignIn(ctx, domain.Authorization{AuthKeyID: key}, "+15550004312", hash, "12345")
|
||||
if !errors.Is(err, domain.ErrSessionPasswordNeeded) {
|
||||
t.Fatalf("SignIn err = %v, want ErrSessionPasswordNeeded", err)
|
||||
}
|
||||
if needSignUp || got.ID != u.ID {
|
||||
t.Fatalf("SignIn user=%+v needSignUp=%v, want existing 2FA user", got, needSignUp)
|
||||
}
|
||||
if signInMessage.ID != 0 {
|
||||
t.Fatalf("2FA SignIn returned a late login message %+v", signInMessage)
|
||||
}
|
||||
// 两步验证未完成:业务鉴权(UserID)必须视为未登录,避免绕过 2FA。
|
||||
bound, found, err := svc.UserID(ctx, key)
|
||||
if err != nil || found || bound != 0 {
|
||||
|
|
@ -539,6 +609,10 @@ func TestSignInExistingTwoFactorAccountNeedsPassword(t *testing.T) {
|
|||
if err != nil || !found || bound != u.ID {
|
||||
t.Fatalf("UserID after 2FA passed = %d found=%v err=%v, want %d", bound, found, err, u.ID)
|
||||
}
|
||||
list, err := dialogs.ListByUser(ctx, u.ID, domain.DialogFilter{Limit: 10})
|
||||
if err != nil || len(list.Messages) != 1 {
|
||||
t.Fatalf("2FA login-code messages after password = %+v err=%v, want exactly the SendCode message", list.Messages, err)
|
||||
}
|
||||
}
|
||||
|
||||
func testAuthKey(seed byte) mtcrypto.AuthKey {
|
||||
|
|
|
|||
554
internal/app/auth/signup_state_test.go
Normal file
554
internal/app/auth/signup_state_test.go
Normal file
|
|
@ -0,0 +1,554 @@
|
|||
package auth
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
accountapp "telesrv/internal/app/account"
|
||||
"telesrv/internal/domain"
|
||||
"telesrv/internal/store"
|
||||
"telesrv/internal/store/memory"
|
||||
)
|
||||
|
||||
func TestSignUpRequiresCorrectSignInAndConsumesMarkerOnce(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
users := memory.NewUserStore()
|
||||
codes := memory.NewCodeStore()
|
||||
svc := NewService(users, memory.NewAuthorizationStore(), codes, nil, nil, "12345")
|
||||
phone := "15550009301"
|
||||
|
||||
hash, err := svc.SendCode(ctx, phone)
|
||||
if err != nil {
|
||||
t.Fatalf("SendCode: %v", err)
|
||||
}
|
||||
if _, _, err := svc.SignUp(ctx, domain.Authorization{}, phone, hash, "Direct", "Bypass"); !errors.Is(err, ErrCodeInvalid) {
|
||||
t.Fatalf("direct SignUp err=%v, want ErrCodeInvalid", err)
|
||||
}
|
||||
if _, _, _, err := svc.SignIn(ctx, domain.Authorization{}, phone, hash, "00000"); !errors.Is(err, ErrCodeInvalid) {
|
||||
t.Fatalf("wrong SignIn err=%v, want ErrCodeInvalid", err)
|
||||
}
|
||||
if rec, found, err := codes.Get(ctx, hash); err != nil || !found || rec.SignUpVerified {
|
||||
t.Fatalf("wrong code marker=%v found=%v err=%v, want live/unverified", rec.SignUpVerified, found, err)
|
||||
}
|
||||
if _, _, err := svc.SignUp(ctx, domain.Authorization{}, phone, hash, "Wrong", "Code"); !errors.Is(err, ErrCodeInvalid) {
|
||||
t.Fatalf("SignUp after wrong code err=%v, want ErrCodeInvalid", err)
|
||||
}
|
||||
|
||||
verifyCodeForSignUp(t, svc, phone, hash, "12345")
|
||||
if rec, found, err := codes.Get(ctx, hash); err != nil || !found || !rec.SignUpVerified || rec.IssuedUserID != 0 {
|
||||
t.Fatalf("verified record=%+v found=%v err=%v", rec, found, err)
|
||||
}
|
||||
if _, msg, needSignUp, err := svc.SignIn(ctx, domain.Authorization{}, phone, hash, "12345"); err != nil || !needSignUp || msg.ID != 0 {
|
||||
t.Fatalf("idempotent SignIn needSignUp=%v message=%+v err=%v", needSignUp, msg, err)
|
||||
}
|
||||
u, _, err := svc.SignUp(ctx, domain.Authorization{}, phone, hash, "Verified", "User")
|
||||
if err != nil || u.Phone != phone {
|
||||
t.Fatalf("verified SignUp user=%+v err=%v", u, err)
|
||||
}
|
||||
if _, _, err := svc.SignUp(ctx, domain.Authorization{}, phone, hash, "Replay", "User"); !errors.Is(err, ErrCodeExpired) {
|
||||
t.Fatalf("replayed SignUp err=%v, want ErrCodeExpired", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestConcurrentSignUpConsumesVerifiedHashExactlyOnce(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
users := memory.NewUserStore()
|
||||
svc := NewService(users, memory.NewAuthorizationStore(), memory.NewCodeStore(), nil, nil, "12345")
|
||||
phone := "15550009302"
|
||||
hash, err := svc.SendCode(ctx, phone)
|
||||
if err != nil {
|
||||
t.Fatalf("SendCode: %v", err)
|
||||
}
|
||||
verifyCodeForSignUp(t, svc, phone, hash, "12345")
|
||||
|
||||
const workers = 16
|
||||
start := make(chan struct{})
|
||||
errs := make(chan error, workers)
|
||||
for i := 0; i < workers; i++ {
|
||||
go func(i int) {
|
||||
<-start
|
||||
var key [8]byte
|
||||
key[0] = byte(i + 1)
|
||||
_, _, err := svc.SignUp(ctx, domain.Authorization{AuthKeyID: key}, phone, hash, "Concurrent", "User")
|
||||
errs <- err
|
||||
}(i)
|
||||
}
|
||||
close(start)
|
||||
successes := 0
|
||||
for i := 0; i < workers; i++ {
|
||||
err := <-errs
|
||||
switch {
|
||||
case err == nil:
|
||||
successes++
|
||||
case errors.Is(err, ErrCodeExpired), errors.Is(err, ErrCodeInvalid):
|
||||
default:
|
||||
t.Fatalf("concurrent SignUp err=%v", err)
|
||||
}
|
||||
}
|
||||
if successes != 1 {
|
||||
t.Fatalf("successful SignUp calls=%d, want 1", successes)
|
||||
}
|
||||
}
|
||||
|
||||
type afterVerifyCodeStore struct {
|
||||
store.CodeStore
|
||||
once sync.Once
|
||||
afterVerify func()
|
||||
}
|
||||
|
||||
type failingPasswordStore struct {
|
||||
store.PasswordStore
|
||||
err error
|
||||
}
|
||||
|
||||
func (s *failingPasswordStore) GetByUser(context.Context, int64) (domain.PasswordSettings, bool, error) {
|
||||
return domain.PasswordSettings{}, false, s.err
|
||||
}
|
||||
|
||||
type switchablePhoneOwnerStore struct {
|
||||
store.UserStore
|
||||
mu sync.RWMutex
|
||||
phone string
|
||||
override bool
|
||||
owner domain.User
|
||||
found bool
|
||||
}
|
||||
|
||||
func (s *switchablePhoneOwnerStore) ByPhone(ctx context.Context, phone string) (domain.User, bool, error) {
|
||||
s.mu.RLock()
|
||||
if s.override && domain.NormalizePhone(phone) == s.phone {
|
||||
owner, found := s.owner, s.found
|
||||
s.mu.RUnlock()
|
||||
return owner, found, nil
|
||||
}
|
||||
s.mu.RUnlock()
|
||||
return s.UserStore.ByPhone(ctx, phone)
|
||||
}
|
||||
|
||||
func (s *switchablePhoneOwnerStore) setOwnerView(phone string, owner domain.User, found bool) {
|
||||
s.mu.Lock()
|
||||
s.phone = domain.NormalizePhone(phone)
|
||||
s.owner = owner
|
||||
s.found = found
|
||||
s.override = true
|
||||
s.mu.Unlock()
|
||||
}
|
||||
|
||||
func (s *switchablePhoneOwnerStore) resetOwnerView() {
|
||||
s.mu.Lock()
|
||||
s.override = false
|
||||
s.mu.Unlock()
|
||||
}
|
||||
|
||||
func (s *afterVerifyCodeStore) VerifyLogin(ctx context.Context, hash, phone, code string, keep bool, maxAttempts int) (store.LoginCodeVerifyResult, error) {
|
||||
result, err := s.CodeStore.VerifyLogin(ctx, hash, phone, code, keep, maxAttempts)
|
||||
if err == nil && result.Status == store.LoginCodeVerifyAccepted && s.afterVerify != nil {
|
||||
s.once.Do(s.afterVerify)
|
||||
}
|
||||
return result, err
|
||||
}
|
||||
|
||||
func TestOwnerTransferAcrossVerifyInvalidatesHashPermanently(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
users := memory.NewUserStore()
|
||||
baseCodes := memory.NewCodeStore()
|
||||
var createErr error
|
||||
codes := &afterVerifyCodeStore{CodeStore: baseCodes}
|
||||
svc := NewService(users, memory.NewAuthorizationStore(), codes, nil, nil, "12345")
|
||||
phone := "15550009303"
|
||||
hash, err := svc.SendCode(ctx, phone)
|
||||
if err != nil {
|
||||
t.Fatalf("SendCode: %v", err)
|
||||
}
|
||||
codes.afterVerify = func() {
|
||||
_, createErr = users.Create(ctx, domain.User{Phone: phone, FirstName: "NewOwner"})
|
||||
}
|
||||
if _, _, needSignUp, err := svc.SignIn(ctx, domain.Authorization{}, phone, hash, "12345"); !errors.Is(err, ErrCodeInvalid) || needSignUp {
|
||||
t.Fatalf("SignIn across owner transfer needSignUp=%v err=%v, want invalid", needSignUp, err)
|
||||
}
|
||||
if createErr != nil {
|
||||
t.Fatalf("create concurrent owner: %v", createErr)
|
||||
}
|
||||
if _, found, err := baseCodes.Get(ctx, hash); err != nil || found {
|
||||
t.Fatalf("owner-drift hash found=%v err=%v, want invalidated", found, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPasswordLookupFailureNeverCreatesOrChangesAuthorization(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
users := memory.NewUserStore()
|
||||
target, err := users.Create(ctx, domain.User{Phone: "15550009320", FirstName: "Target"})
|
||||
if err != nil {
|
||||
t.Fatalf("create target: %v", err)
|
||||
}
|
||||
previous, err := users.Create(ctx, domain.User{Phone: "15550009321", FirstName: "Previous"})
|
||||
if err != nil {
|
||||
t.Fatalf("create previous: %v", err)
|
||||
}
|
||||
authz := memory.NewAuthorizationStore()
|
||||
lookupErr := errors.New("password store unavailable")
|
||||
passwords := &failingPasswordStore{PasswordStore: memory.NewPasswordStore(), err: lookupErr}
|
||||
svc := NewService(users, authz, memory.NewCodeStore(), nil, nil, "12345",
|
||||
WithPasswords(passwords),
|
||||
WithLoginCodeDelivery(&captureLoginCodeDelivery{}),
|
||||
)
|
||||
|
||||
t.Run("unbound-key-remains-unbound", func(t *testing.T) {
|
||||
key := [8]byte{0xC1}
|
||||
hash, err := svc.SendCode(ctx, target.Phone)
|
||||
if err != nil {
|
||||
t.Fatalf("SendCode: %v", err)
|
||||
}
|
||||
if _, _, _, err := svc.SignIn(ctx, domain.Authorization{AuthKeyID: key}, target.Phone, hash, "12345"); !errors.Is(err, lookupErr) {
|
||||
t.Fatalf("SignIn err=%v, want password lookup failure", err)
|
||||
}
|
||||
if got, found, err := authz.ByAuthKey(ctx, key); err != nil || found {
|
||||
t.Fatalf("authorization=%+v found=%v err=%v, want absent", got, found, err)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("previous-binding-remains-unchanged", func(t *testing.T) {
|
||||
key := [8]byte{0xC2}
|
||||
original := domain.Authorization{AuthKeyID: key, UserID: previous.ID, Hash: 987654321}
|
||||
if err := authz.Bind(ctx, original); err != nil {
|
||||
t.Fatalf("bind previous authorization: %v", err)
|
||||
}
|
||||
hash, err := svc.SendCode(ctx, target.Phone)
|
||||
if err != nil {
|
||||
t.Fatalf("SendCode: %v", err)
|
||||
}
|
||||
if _, _, _, err := svc.SignIn(ctx, domain.Authorization{AuthKeyID: key}, target.Phone, hash, "12345"); !errors.Is(err, lookupErr) {
|
||||
t.Fatalf("SignIn err=%v, want password lookup failure", err)
|
||||
}
|
||||
got, found, err := authz.ByAuthKey(ctx, key)
|
||||
if err != nil || !found || got.UserID != previous.ID || got.Hash != original.Hash || got.PasswordPending != original.PasswordPending {
|
||||
t.Fatalf("authorization after failure=%+v found=%v err=%v, want unchanged %+v", got, found, err, original)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestOwnerTransferAwayAndBackCannotReviveLoginHash(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
t.Run("unregistered-signin", func(t *testing.T) {
|
||||
baseUsers := memory.NewUserStore()
|
||||
other, err := baseUsers.Create(ctx, domain.User{Phone: "15550009311", FirstName: "Other"})
|
||||
if err != nil {
|
||||
t.Fatalf("create other owner: %v", err)
|
||||
}
|
||||
users := &switchablePhoneOwnerStore{UserStore: baseUsers}
|
||||
codes := memory.NewCodeStore()
|
||||
svc := NewService(users, memory.NewAuthorizationStore(), codes, nil, nil, "12345")
|
||||
phone := "15550009310"
|
||||
hash, err := svc.SendCode(ctx, phone)
|
||||
if err != nil {
|
||||
t.Fatalf("SendCode: %v", err)
|
||||
}
|
||||
users.setOwnerView(phone, other, true)
|
||||
if _, _, _, err := svc.SignIn(ctx, domain.Authorization{}, phone, hash, "12345"); !errors.Is(err, ErrCodeInvalid) {
|
||||
t.Fatalf("SignIn after 0->B owner transfer err=%v, want invalid", err)
|
||||
}
|
||||
users.resetOwnerView()
|
||||
if _, _, _, err := svc.SignIn(ctx, domain.Authorization{}, phone, hash, "12345"); !errors.Is(err, ErrCodeExpired) {
|
||||
t.Fatalf("SignIn after 0->B->0 err=%v, want expired", err)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("existing-resend", func(t *testing.T) {
|
||||
baseUsers := memory.NewUserStore()
|
||||
ownerA, err := baseUsers.Create(ctx, domain.User{Phone: "15550009312", FirstName: "A"})
|
||||
if err != nil {
|
||||
t.Fatalf("create owner A: %v", err)
|
||||
}
|
||||
ownerB, err := baseUsers.Create(ctx, domain.User{Phone: "15550009313", FirstName: "B"})
|
||||
if err != nil {
|
||||
t.Fatalf("create owner B: %v", err)
|
||||
}
|
||||
users := &switchablePhoneOwnerStore{UserStore: baseUsers}
|
||||
codes := memory.NewCodeStore()
|
||||
svc := NewService(users, memory.NewAuthorizationStore(), codes, nil, nil, "12345", WithLoginCodeDelivery(&captureLoginCodeDelivery{}))
|
||||
hash, err := svc.SendCode(ctx, ownerA.Phone)
|
||||
if err != nil {
|
||||
t.Fatalf("SendCode: %v", err)
|
||||
}
|
||||
users.setOwnerView(ownerA.Phone, ownerB, true)
|
||||
if _, err := svc.ResendCode(ctx, ownerA.Phone, hash); !errors.Is(err, ErrCodeInvalid) {
|
||||
t.Fatalf("ResendCode after A->B err=%v, want invalid", err)
|
||||
}
|
||||
users.resetOwnerView()
|
||||
if _, _, _, err := svc.SignIn(ctx, domain.Authorization{}, ownerA.Phone, hash, "12345"); !errors.Is(err, ErrCodeExpired) {
|
||||
t.Fatalf("SignIn after A->B->A err=%v, want expired", err)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("existing-cancel", func(t *testing.T) {
|
||||
baseUsers := memory.NewUserStore()
|
||||
ownerA, err := baseUsers.Create(ctx, domain.User{Phone: "15550009314", FirstName: "A"})
|
||||
if err != nil {
|
||||
t.Fatalf("create owner A: %v", err)
|
||||
}
|
||||
ownerB, err := baseUsers.Create(ctx, domain.User{Phone: "15550009315", FirstName: "B"})
|
||||
if err != nil {
|
||||
t.Fatalf("create owner B: %v", err)
|
||||
}
|
||||
users := &switchablePhoneOwnerStore{UserStore: baseUsers}
|
||||
codes := memory.NewCodeStore()
|
||||
svc := NewService(users, memory.NewAuthorizationStore(), codes, nil, nil, "12345", WithLoginCodeDelivery(&captureLoginCodeDelivery{}))
|
||||
hash, err := svc.SendCode(ctx, ownerA.Phone)
|
||||
if err != nil {
|
||||
t.Fatalf("SendCode: %v", err)
|
||||
}
|
||||
users.setOwnerView(ownerA.Phone, ownerB, true)
|
||||
if err := svc.CancelCode(ctx, ownerA.Phone, hash); !errors.Is(err, ErrCodeInvalid) {
|
||||
t.Fatalf("CancelCode after A->B err=%v, want invalid", err)
|
||||
}
|
||||
users.resetOwnerView()
|
||||
if _, _, _, err := svc.SignIn(ctx, domain.Authorization{}, ownerA.Phone, hash, "12345"); !errors.Is(err, ErrCodeExpired) {
|
||||
t.Fatalf("SignIn after canceled A->B->A err=%v, want expired", err)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestEmailSetupVerificationAuthorizesSignUpWithout777000Message(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
users := memory.NewUserStore()
|
||||
codes := memory.NewCodeStore()
|
||||
passwords := memory.NewPasswordStore()
|
||||
sender := &testMailSender{}
|
||||
accountSvc := accountapp.NewService(passwords,
|
||||
accountapp.WithUsers(users),
|
||||
accountapp.WithLoginEmailVerification(codes, sender, time.Minute, 3, 6),
|
||||
)
|
||||
dialogs := memory.NewDialogStore()
|
||||
messages := memory.NewMessageStore(dialogs)
|
||||
authSvc := NewService(users, memory.NewAuthorizationStore(), codes, nil, nil, "12345",
|
||||
WithLoginMessages(messages, dialogs),
|
||||
WithLoginEmail(LoginEmailOptions{
|
||||
Enabled: true,
|
||||
RequireSetup: true,
|
||||
CodeLength: 6,
|
||||
Store: accountSvc,
|
||||
Sender: sender,
|
||||
}),
|
||||
)
|
||||
phone := "15550009304"
|
||||
hash, err := authSvc.SendCode(ctx, phone)
|
||||
if err != nil {
|
||||
t.Fatalf("SendCode: %v", err)
|
||||
}
|
||||
if _, _, err := authSvc.SignUp(ctx, domain.Authorization{}, phone, hash, "Direct", "Email"); !errors.Is(err, ErrCodeInvalid) {
|
||||
t.Fatalf("SignUp before email setup err=%v, want ErrCodeInvalid", err)
|
||||
}
|
||||
if _, _, err := accountSvc.SendLoginEmailCode(ctx, 0, phone, hash, "new@example.test", true); err != nil {
|
||||
t.Fatalf("SendLoginEmailCode: %v", err)
|
||||
}
|
||||
bad := wrongCode(sender.code, '0')
|
||||
if _, err := accountSvc.VerifyLoginEmail(ctx, 0, phone, hash, bad, true); !errors.Is(err, domain.ErrEmailCodeInvalid) {
|
||||
t.Fatalf("wrong VerifyLoginEmail err=%v, want ErrEmailCodeInvalid", err)
|
||||
}
|
||||
if rec, found, err := codes.Get(ctx, hash); err != nil || !found || rec.SignUpVerified {
|
||||
t.Fatalf("wrong SMTP code marker=%v found=%v err=%v", rec.SignUpVerified, found, err)
|
||||
}
|
||||
if _, err := accountSvc.VerifyLoginEmail(ctx, 0, phone, hash, sender.code, true); err != nil {
|
||||
t.Fatalf("VerifyLoginEmail: %v", err)
|
||||
}
|
||||
if rec, found, err := codes.Get(ctx, hash); err != nil || !found || !rec.SignUpVerified || rec.Channel != codeChannelEmailLogin {
|
||||
t.Fatalf("email-verified phone code=%+v found=%v err=%v", rec, found, err)
|
||||
}
|
||||
if _, msg, needSignUp, err := authSvc.SignInWithEmail(ctx, domain.Authorization{}, phone, hash, sender.code); err != nil || !needSignUp || msg.ID != 0 {
|
||||
t.Fatalf("SignInWithEmail after setup needSignUp=%v message=%+v err=%v", needSignUp, msg, err)
|
||||
}
|
||||
u, msg, err := authSvc.SignUp(ctx, domain.Authorization{}, phone, hash, "Email", "User")
|
||||
if err != nil {
|
||||
t.Fatalf("SignUp after email setup: %v", err)
|
||||
}
|
||||
if msg.ID != 0 || msg.Body != "" {
|
||||
t.Fatalf("email SignUp returned SMTP code message: %+v", msg)
|
||||
}
|
||||
list, err := dialogs.ListByUser(ctx, u.ID, domain.DialogFilter{Limit: 10})
|
||||
if err != nil {
|
||||
t.Fatalf("ListByUser: %v", err)
|
||||
}
|
||||
if len(list.Dialogs) != 0 || len(list.Messages) != 0 {
|
||||
t.Fatalf("email SignUp created 777000 bootstrap state: dialogs=%+v messages=%+v", list.Dialogs, list.Messages)
|
||||
}
|
||||
if email, found, err := accountSvc.LoginEmailByPhone(ctx, phone); err != nil || !found || email != "new@example.test" {
|
||||
t.Fatalf("LoginEmailByPhone email=%q found=%v err=%v", email, found, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestConsumeLoginEmailResetRequiresExactIssuedHash(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
baseUsers := memory.NewUserStore()
|
||||
owner, err := baseUsers.Create(ctx, domain.User{Phone: "15550009330", FirstName: "Owner"})
|
||||
if err != nil {
|
||||
t.Fatalf("create owner: %v", err)
|
||||
}
|
||||
other, err := baseUsers.Create(ctx, domain.User{Phone: "15550009331", FirstName: "Other"})
|
||||
if err != nil {
|
||||
t.Fatalf("create other: %v", err)
|
||||
}
|
||||
users := &switchablePhoneOwnerStore{UserStore: baseUsers}
|
||||
codes := memory.NewCodeStore()
|
||||
delivery := &captureLoginCodeDelivery{}
|
||||
svc := NewService(users, memory.NewAuthorizationStore(), codes, nil, nil, "12345", WithLoginCodeDelivery(delivery))
|
||||
seed := func(hash, channel string) {
|
||||
t.Helper()
|
||||
if err := codes.Set(ctx, hash, store.PhoneCode{
|
||||
Version: store.PhoneCodeVersionCurrent,
|
||||
IssuedUserID: owner.ID,
|
||||
Phone: owner.Phone,
|
||||
Code: "654321",
|
||||
Channel: channel,
|
||||
MaxAttempts: 5,
|
||||
}, time.Minute); err != nil {
|
||||
t.Fatalf("seed %s: %v", hash, err)
|
||||
}
|
||||
}
|
||||
|
||||
if _, err := svc.ConsumeLoginEmailReset(ctx, owner.Phone, "arbitrary-missing"); !errors.Is(err, ErrCodeExpired) {
|
||||
t.Fatalf("arbitrary hash err=%v, want expired", err)
|
||||
}
|
||||
seed("wrong-phone", codeChannelEmailLogin)
|
||||
if _, err := svc.ConsumeLoginEmailReset(ctx, other.Phone, "wrong-phone"); !errors.Is(err, ErrCodeInvalid) {
|
||||
t.Fatalf("wrong phone err=%v, want invalid", err)
|
||||
}
|
||||
if _, found, err := codes.Get(ctx, "wrong-phone"); err != nil || !found {
|
||||
t.Fatalf("wrong-phone probe destroyed valid hash found=%v err=%v", found, err)
|
||||
}
|
||||
seed("wrong-channel", codeChannelPhone)
|
||||
if _, err := svc.ConsumeLoginEmailReset(ctx, owner.Phone, "wrong-channel"); !errors.Is(err, ErrCodeInvalid) {
|
||||
t.Fatalf("wrong channel err=%v, want invalid", err)
|
||||
}
|
||||
|
||||
seed("owner-drift", codeChannelEmailLogin)
|
||||
users.setOwnerView(owner.Phone, other, true)
|
||||
if _, err := svc.ConsumeLoginEmailReset(ctx, owner.Phone, "owner-drift"); !errors.Is(err, ErrCodeInvalid) {
|
||||
t.Fatalf("A->B reset err=%v, want invalid", err)
|
||||
}
|
||||
users.resetOwnerView()
|
||||
if _, err := svc.ConsumeLoginEmailReset(ctx, owner.Phone, "owner-drift"); !errors.Is(err, ErrCodeExpired) {
|
||||
t.Fatalf("A->B->A reset err=%v, want expired", err)
|
||||
}
|
||||
|
||||
seed("successful-reset", codeChannelEmailLogin)
|
||||
resetUserID, err := svc.ConsumeLoginEmailReset(ctx, owner.Phone, "successful-reset")
|
||||
if err != nil || resetUserID != owner.ID {
|
||||
t.Fatalf("successful reset consume uid=%d err=%v", resetUserID, err)
|
||||
}
|
||||
replacementHash, err := svc.SendPhoneCodeAfterLoginEmailReset(ctx, owner.Phone, resetUserID)
|
||||
if err != nil || replacementHash == "" {
|
||||
t.Fatalf("replacement hash=%q err=%v", replacementHash, err)
|
||||
}
|
||||
if len(delivery.requests) != 1 || delivery.requests[0].UserID != owner.ID || delivery.requests[0].PhoneCodeHash != replacementHash {
|
||||
t.Fatalf("replacement delivery=%+v", delivery.requests)
|
||||
}
|
||||
if rec, found, err := codes.Get(ctx, replacementHash); err != nil || !found || rec.Version != store.PhoneCodeVersionCurrent || rec.IssuedUserID != owner.ID || rec.Channel != codeChannelPhone {
|
||||
t.Fatalf("replacement code=%+v found=%v err=%v", rec, found, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestConcurrentLoginEmailResetHasSingleConsumer(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
users := memory.NewUserStore()
|
||||
owner, err := users.Create(ctx, domain.User{Phone: "15550009332", FirstName: "Owner"})
|
||||
if err != nil {
|
||||
t.Fatalf("create owner: %v", err)
|
||||
}
|
||||
codes := memory.NewCodeStore()
|
||||
hash := "concurrent-email-reset"
|
||||
if err := codes.Set(ctx, hash, store.PhoneCode{
|
||||
Version: store.PhoneCodeVersionCurrent,
|
||||
IssuedUserID: owner.ID,
|
||||
Phone: owner.Phone,
|
||||
Code: "654321",
|
||||
Channel: codeChannelEmailLogin,
|
||||
MaxAttempts: 5,
|
||||
}, time.Minute); err != nil {
|
||||
t.Fatalf("seed code: %v", err)
|
||||
}
|
||||
svc := NewService(users, memory.NewAuthorizationStore(), codes, nil, nil, "12345")
|
||||
const workers = 24
|
||||
start := make(chan struct{})
|
||||
errs := make(chan error, workers)
|
||||
for i := 0; i < workers; i++ {
|
||||
go func() {
|
||||
<-start
|
||||
_, err := svc.ConsumeLoginEmailReset(ctx, owner.Phone, hash)
|
||||
errs <- err
|
||||
}()
|
||||
}
|
||||
close(start)
|
||||
successes := 0
|
||||
for i := 0; i < workers; i++ {
|
||||
err := <-errs
|
||||
if err == nil {
|
||||
successes++
|
||||
continue
|
||||
}
|
||||
if !errors.Is(err, ErrCodeExpired) {
|
||||
t.Fatalf("concurrent reset err=%v", err)
|
||||
}
|
||||
}
|
||||
if successes != 1 {
|
||||
t.Fatalf("successful reset consumers=%d, want 1", successes)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoginEmailResetLocksUserAcrossOwnerTransfer(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
baseUsers := memory.NewUserStore()
|
||||
ownerA, err := baseUsers.Create(ctx, domain.User{Phone: "15550009340", FirstName: "A"})
|
||||
if err != nil {
|
||||
t.Fatalf("create A: %v", err)
|
||||
}
|
||||
ownerB, err := baseUsers.Create(ctx, domain.User{Phone: "15550009341", FirstName: "B"})
|
||||
if err != nil {
|
||||
t.Fatalf("create B: %v", err)
|
||||
}
|
||||
users := &switchablePhoneOwnerStore{UserStore: baseUsers}
|
||||
passwords := memory.NewPasswordStore()
|
||||
accountSvc := accountapp.NewService(passwords, accountapp.WithUsers(users))
|
||||
if err := accountSvc.SetLoginEmail(ctx, ownerA.ID, "a@example.test"); err != nil {
|
||||
t.Fatalf("SetLoginEmail A: %v", err)
|
||||
}
|
||||
if err := accountSvc.SetLoginEmail(ctx, ownerB.ID, "b@example.test"); err != nil {
|
||||
t.Fatalf("SetLoginEmail B: %v", err)
|
||||
}
|
||||
codes := memory.NewCodeStore()
|
||||
hash := "locked-reset-user"
|
||||
if err := codes.Set(ctx, hash, store.PhoneCode{
|
||||
Version: store.PhoneCodeVersionCurrent,
|
||||
IssuedUserID: ownerA.ID,
|
||||
Phone: ownerA.Phone,
|
||||
Code: "654321",
|
||||
Channel: codeChannelEmailLogin,
|
||||
MaxAttempts: 5,
|
||||
}, time.Minute); err != nil {
|
||||
t.Fatalf("seed reset code: %v", err)
|
||||
}
|
||||
delivery := &captureLoginCodeDelivery{}
|
||||
authSvc := NewService(users, memory.NewAuthorizationStore(), codes, nil, nil, "12345", WithLoginCodeDelivery(delivery))
|
||||
resetUserID, err := authSvc.ConsumeLoginEmailReset(ctx, ownerA.Phone, hash)
|
||||
if err != nil || resetUserID != ownerA.ID {
|
||||
t.Fatalf("ConsumeLoginEmailReset uid=%d err=%v", resetUserID, err)
|
||||
}
|
||||
users.setOwnerView(ownerA.Phone, ownerB, true)
|
||||
if err := accountSvc.ClearLoginEmail(ctx, resetUserID); err != nil {
|
||||
t.Fatalf("ClearLoginEmail exact A: %v", err)
|
||||
}
|
||||
if _, err := authSvc.SendPhoneCodeAfterLoginEmailReset(ctx, ownerA.Phone, resetUserID); !errors.Is(err, ErrCodeInvalid) {
|
||||
t.Fatalf("SendPhoneCodeAfterLoginEmailReset across A->B err=%v, want invalid", err)
|
||||
}
|
||||
if _, found, err := accountSvc.LoginEmail(ctx, ownerA.ID); err != nil || found {
|
||||
t.Fatalf("A login email found=%v err=%v, want cleared", found, err)
|
||||
}
|
||||
if email, found, err := accountSvc.LoginEmail(ctx, ownerB.ID); err != nil || !found || email != "b@example.test" {
|
||||
t.Fatalf("B login email=%q found=%v err=%v, want unchanged", email, found, err)
|
||||
}
|
||||
if len(delivery.requests) != 0 {
|
||||
t.Fatalf("owner B received reset replacement code: %+v", delivery.requests)
|
||||
}
|
||||
}
|
||||
|
|
@ -1242,6 +1242,25 @@ func (s *Service) SendMessage(ctx context.Context, userID int64, req domain.Send
|
|||
if req.UserID != userID {
|
||||
return domain.SendChannelMessageResult{}, domain.ErrChannelInvalid
|
||||
}
|
||||
if req.RandomID != 0 && !req.IdempotencyPreflighted {
|
||||
fingerprint, err := store.ChannelSendFingerprint(req)
|
||||
if err != nil {
|
||||
return domain.SendChannelMessageResult{}, err
|
||||
}
|
||||
req.IdempotencyFingerprint = fingerprint
|
||||
if replayStore, ok := s.channels.(store.ChannelSendReplayStore); ok {
|
||||
replay, found, err := replayStore.LookupChannelSendReplay(ctx, domain.ChannelSendReplayRequest{
|
||||
ChannelID: req.ChannelID,
|
||||
SenderUserID: req.UserID,
|
||||
RandomID: req.RandomID,
|
||||
IdempotencyFingerprint: fingerprint,
|
||||
})
|
||||
if err != nil || found {
|
||||
return replay, err
|
||||
}
|
||||
req.IdempotencyPreflighted = true
|
||||
}
|
||||
}
|
||||
if err := s.ensureCanSend(ctx, req.UserID); err != nil {
|
||||
return domain.SendChannelMessageResult{}, err
|
||||
}
|
||||
|
|
@ -1253,6 +1272,25 @@ func (s *Service) SendMessage(ctx context.Context, userID int64, req domain.Send
|
|||
return s.channels.SendChannelMessage(ctx, req)
|
||||
}
|
||||
|
||||
// LookupChannelSendReplay reads a regular-channel or monoforum receipt without current
|
||||
// membership/send-gate checks. The authenticated caller remains bound to SenderUserID.
|
||||
func (s *Service) LookupChannelSendReplay(ctx context.Context, userID int64, req domain.ChannelSendReplayRequest) (domain.SendChannelMessageResult, bool, error) {
|
||||
if s == nil || s.channels == nil || userID == 0 {
|
||||
return domain.SendChannelMessageResult{}, false, nil
|
||||
}
|
||||
if req.SenderUserID == 0 {
|
||||
req.SenderUserID = userID
|
||||
}
|
||||
if req.SenderUserID != userID || req.ChannelID == 0 || req.RandomID == 0 {
|
||||
return domain.SendChannelMessageResult{}, false, domain.ErrChannelInvalid
|
||||
}
|
||||
replayStore, ok := s.channels.(store.ChannelSendReplayStore)
|
||||
if !ok {
|
||||
return domain.SendChannelMessageResult{}, false, nil
|
||||
}
|
||||
return replayStore.LookupChannelSendReplay(ctx, req)
|
||||
}
|
||||
|
||||
func (s *Service) ensureCanSend(ctx context.Context, userID int64) error {
|
||||
if s == nil || s.sendGate == nil || userID == 0 {
|
||||
return nil
|
||||
|
|
@ -1754,6 +1792,26 @@ func (s *Service) SendMonoforumMessage(ctx context.Context, req domain.SendMonof
|
|||
if s == nil || s.channels == nil || req.MonoforumID == 0 || req.SenderUserID == 0 || req.SavedPeer.ID == 0 {
|
||||
return domain.SendChannelMessageResult{}, domain.ErrChannelInvalid
|
||||
}
|
||||
if req.RandomID != 0 && !req.IdempotencyPreflighted {
|
||||
fingerprint, err := store.MonoforumSendFingerprint(req)
|
||||
if err != nil {
|
||||
return domain.SendChannelMessageResult{}, err
|
||||
}
|
||||
req.IdempotencyFingerprint = fingerprint
|
||||
if replayStore, ok := s.channels.(store.ChannelSendReplayStore); ok {
|
||||
replay, found, err := replayStore.LookupChannelSendReplay(ctx, domain.ChannelSendReplayRequest{
|
||||
ChannelID: req.MonoforumID,
|
||||
SenderUserID: req.SenderUserID,
|
||||
SavedPeer: req.SavedPeer,
|
||||
RandomID: req.RandomID,
|
||||
IdempotencyFingerprint: fingerprint,
|
||||
})
|
||||
if err != nil || found {
|
||||
return replay, err
|
||||
}
|
||||
req.IdempotencyPreflighted = true
|
||||
}
|
||||
}
|
||||
if err := s.ensureCanSend(ctx, req.SenderUserID); err != nil {
|
||||
return domain.SendChannelMessageResult{}, err
|
||||
}
|
||||
|
|
@ -2021,6 +2079,26 @@ func (s *Service) DirtyActiveChannelsForUser(ctx context.Context, userID int64,
|
|||
return s.channels.ListDirtyActiveChannelsForUser(ctx, userID, sinceDate, afterChannelID, limit)
|
||||
}
|
||||
|
||||
// MaxChannelPts returns the durable channel watermark used by the fan-out saturation recovery
|
||||
// sweep. It intentionally performs no viewer access check: target visibility is derived from the
|
||||
// process-local joined-membership index, while getChannelDifference performs authoritative access
|
||||
// validation when a client consumes the nudge.
|
||||
func (s *Service) MaxChannelPts(ctx context.Context, channelID int64) (int, error) {
|
||||
if s == nil || s.channels == nil || channelID == 0 {
|
||||
return 0, domain.ErrChannelInvalid
|
||||
}
|
||||
return s.channels.MaxChannelPts(ctx, channelID)
|
||||
}
|
||||
|
||||
// MaxChannelPtsBatch reloads a bounded recovery page in one store call. Missing ids are omitted:
|
||||
// they represent channels deleted after the process-local online-membership snapshot was taken.
|
||||
func (s *Service) MaxChannelPtsBatch(ctx context.Context, channelIDs []int64) (map[int64]int, error) {
|
||||
if s == nil || s.channels == nil {
|
||||
return nil, domain.ErrChannelInvalid
|
||||
}
|
||||
return s.channels.MaxChannelPtsBatch(ctx, channelIDs)
|
||||
}
|
||||
|
||||
// ActiveMemberIDs returns a bounded list for transient online fanout such as typing.
|
||||
func (s *Service) ActiveMemberIDs(ctx context.Context, userID, channelID int64, limit int) ([]int64, error) {
|
||||
if s == nil || s.channels == nil || userID == 0 || channelID == 0 {
|
||||
|
|
|
|||
|
|
@ -30,6 +30,46 @@ func TestServiceSendMessageHonorsSendPermissionGate(t *testing.T) {
|
|||
}
|
||||
}
|
||||
|
||||
func TestServiceChannelReplayPrecedesCurrentSendPermissionGate(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
channels := memory.NewChannelStore()
|
||||
created, err := channels.CreateChannel(ctx, domain.CreateChannelRequest{
|
||||
CreatorUserID: 1001,
|
||||
Title: "replay gate",
|
||||
Megagroup: true,
|
||||
Date: 1_700_000_000,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("CreateChannel: %v", err)
|
||||
}
|
||||
req := domain.SendChannelMessageRequest{
|
||||
ChannelID: created.Channel.ID,
|
||||
RandomID: 92,
|
||||
Message: "committed before restriction",
|
||||
Date: 1_700_000_001,
|
||||
}
|
||||
allowed := NewService(channels)
|
||||
first, err := allowed.SendMessage(ctx, 1001, req)
|
||||
if err != nil {
|
||||
t.Fatalf("first SendMessage: %v", err)
|
||||
}
|
||||
|
||||
denied := NewService(channels, WithSendPermissionChecker(channelDenySendChecker{}))
|
||||
req.Date++
|
||||
replay, err := denied.SendMessage(ctx, 1001, req)
|
||||
if err != nil {
|
||||
t.Fatalf("replay through denied gate: %v", err)
|
||||
}
|
||||
if !replay.Duplicate || replay.Message.ID != first.Message.ID {
|
||||
t.Fatalf("replay = %+v, want committed duplicate %d", replay, first.Message.ID)
|
||||
}
|
||||
|
||||
req.Message = "different intent"
|
||||
if _, err := denied.SendMessage(ctx, 1001, req); !errors.Is(err, domain.ErrMessageRandomIDDuplicate) {
|
||||
t.Fatalf("conflicting replay err=%v, want ErrMessageRandomIDDuplicate before send gate", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestServiceSendMonoforumMessageHonorsSendPermissionGate(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
svc := NewService(memory.NewChannelStore(), WithSendPermissionChecker(channelDenySendChecker{}))
|
||||
|
|
@ -44,6 +84,52 @@ func TestServiceSendMonoforumMessageHonorsSendPermissionGate(t *testing.T) {
|
|||
}
|
||||
}
|
||||
|
||||
func TestServiceMonoforumReplayPrecedesCurrentSendPermissionGate(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
channels := memory.NewChannelStore()
|
||||
parent, err := channels.CreateChannel(ctx, domain.CreateChannelRequest{
|
||||
CreatorUserID: 1001,
|
||||
Title: "direct messages",
|
||||
Broadcast: true,
|
||||
Date: 1_700_000_010,
|
||||
})
|
||||
if err != nil {
|
||||
t.Fatalf("CreateChannel: %v", err)
|
||||
}
|
||||
enabled, err := channels.SetPaidMessagesPrice(ctx, 1001, parent.Channel.ID, 0, true)
|
||||
if err != nil {
|
||||
t.Fatalf("SetPaidMessagesPrice: %v", err)
|
||||
}
|
||||
req := domain.SendMonoforumMessageRequest{
|
||||
MonoforumID: enabled.Channel.LinkedMonoforumID,
|
||||
SenderUserID: 1002,
|
||||
SavedPeer: domain.Peer{Type: domain.PeerTypeUser, ID: 1002},
|
||||
RandomID: 93,
|
||||
Message: "committed direct message",
|
||||
Date: 1_700_000_011,
|
||||
}
|
||||
allowed := NewService(channels)
|
||||
first, err := allowed.SendMonoforumMessage(ctx, req)
|
||||
if err != nil {
|
||||
t.Fatalf("first SendMonoforumMessage: %v", err)
|
||||
}
|
||||
|
||||
denied := NewService(channels, WithSendPermissionChecker(channelDenySendChecker{}))
|
||||
req.Date++
|
||||
replay, err := denied.SendMonoforumMessage(ctx, req)
|
||||
if err != nil {
|
||||
t.Fatalf("monoforum replay through denied gate: %v", err)
|
||||
}
|
||||
if !replay.Duplicate || replay.Message.ID != first.Message.ID {
|
||||
t.Fatalf("monoforum replay = %+v, want committed duplicate %d", replay, first.Message.ID)
|
||||
}
|
||||
|
||||
req.Message = "different intent"
|
||||
if _, err := denied.SendMonoforumMessage(ctx, req); !errors.Is(err, domain.ErrMessageRandomIDDuplicate) {
|
||||
t.Fatalf("conflicting monoforum replay err=%v, want ErrMessageRandomIDDuplicate before send gate", err)
|
||||
}
|
||||
}
|
||||
|
||||
type channelDenySendChecker struct{}
|
||||
|
||||
func (channelDenySendChecker) CanSendMessages(context.Context, int64) error {
|
||||
|
|
@ -813,7 +899,8 @@ func TestCreateChatCreatesMegagroupWithChannelPts(t *testing.T) {
|
|||
duplicate, err := service.SendMessage(ctx, 1001, domain.SendChannelMessageRequest{
|
||||
ChannelID: created.Channel.ID,
|
||||
RandomID: 99,
|
||||
Message: "hello again",
|
||||
Message: "hello",
|
||||
ViaBotID: 1003,
|
||||
Date: 12,
|
||||
})
|
||||
if err != nil {
|
||||
|
|
@ -2032,12 +2119,12 @@ func TestChannelEditDeleteAndLocalClearUseChannelPts(t *testing.T) {
|
|||
if edited.Event.Type != domain.ChannelUpdateEditMessage || edited.Event.Pts != 4 || edited.Event.PtsCount != 1 {
|
||||
t.Fatalf("edit event = %+v, want channel edit pts=4 count=1", edited.Event)
|
||||
}
|
||||
duplicate, err := service.SendMessage(ctx, 1002, domain.SendChannelMessageRequest{ChannelID: created.Channel.ID, RandomID: 2, Message: "two retry", Date: 13})
|
||||
duplicate, err := service.SendMessage(ctx, 1002, domain.SendChannelMessageRequest{ChannelID: created.Channel.ID, RandomID: 2, Message: "two", Date: 13})
|
||||
if err != nil {
|
||||
t.Fatalf("duplicate SendMessage after edit: %v", err)
|
||||
}
|
||||
if !duplicate.Duplicate || duplicate.Event.Type != domain.ChannelUpdateNewMessage || duplicate.Message.Body != "two" || duplicate.Event.Message.Body != "two" {
|
||||
t.Fatalf("duplicate after edit = %+v, want original new-message snapshot", duplicate)
|
||||
if !duplicate.Duplicate || duplicate.Event.Type != domain.ChannelUpdateNewMessage || duplicate.Message.Body != "two edited" || duplicate.Event.Message.Body != "two edited" {
|
||||
t.Fatalf("duplicate after edit = %+v, want current message in new-message replay", duplicate)
|
||||
}
|
||||
|
||||
deleted, err := service.DeleteMessages(ctx, 1001, domain.DeleteChannelMessagesRequest{
|
||||
|
|
|
|||
|
|
@ -8,10 +8,10 @@ import (
|
|||
"fmt"
|
||||
"image"
|
||||
"image/color"
|
||||
"io"
|
||||
stddraw "image/draw"
|
||||
_ "image/jpeg" // 注册 jpeg DecodeConfig,用于读取上传头像/图片尺寸
|
||||
"image/png"
|
||||
"io"
|
||||
"math"
|
||||
"strings"
|
||||
"time"
|
||||
|
|
@ -48,14 +48,40 @@ func (s *Service) UploadProfilePhotoKind(ctx context.Context, ownerType domain.P
|
|||
|
||||
// CreatePhotoFromUpload 把已上传文件组装成 Photo(不绑定 profile_photos),用于频道头像 / 图片消息。
|
||||
func (s *Service) CreatePhotoFromUpload(ctx context.Context, file domain.UploadedFileRef) (domain.Photo, error) {
|
||||
data, err := s.assembleUpload(ctx, file.OwnerUserID, file.FileID, file.Parts)
|
||||
intentHash, err := uploadedMediaIntentHash(domain.UploadedMediaPhoto, file, nil)
|
||||
if err != nil {
|
||||
return domain.Photo{}, err
|
||||
}
|
||||
if photo, found, err := s.replayUploadedPhoto(ctx, file, intentHash); err != nil || found {
|
||||
return photo, err
|
||||
}
|
||||
data, err := s.readUploadBytes(ctx, file.OwnerUserID, file.FileID, file.Parts)
|
||||
if err != nil {
|
||||
return domain.Photo{}, err
|
||||
}
|
||||
if len(data) == 0 {
|
||||
return domain.Photo{}, domain.ErrPhotoInvalid
|
||||
}
|
||||
return s.createPhoto(ctx, data, photoSizeSpecsForMessage(data))
|
||||
photo, err := s.createPhoto(ctx, data, photoSizeSpecsForMessage(data))
|
||||
if err != nil {
|
||||
return domain.Photo{}, err
|
||||
}
|
||||
receipt, err := s.commitUploadedMediaReceipt(ctx, file, domain.UploadedMediaPhoto, intentHash, photo.ID)
|
||||
if err != nil {
|
||||
return domain.Photo{}, err
|
||||
}
|
||||
if receipt.MediaID != photo.ID {
|
||||
winner, found, err := s.media.GetPhoto(ctx, receipt.MediaID)
|
||||
if err != nil {
|
||||
return domain.Photo{}, err
|
||||
}
|
||||
if !found {
|
||||
return domain.Photo{}, fmt.Errorf("concurrent upload receipt references missing photo %d", receipt.MediaID)
|
||||
}
|
||||
photo = winner
|
||||
}
|
||||
s.cleanupMaterializedUpload(ctx, file, "photo materialized")
|
||||
return photo, nil
|
||||
}
|
||||
|
||||
// CreatePhotoFromBytes stores already-fetched image bytes as a message Photo.
|
||||
|
|
@ -202,6 +228,13 @@ func validateAvatarMarkupSize(size domain.PhotoSize) error {
|
|||
|
||||
// CreateDocumentFromUpload 把已上传文件组装成 Document(文件/视频/音频/gif/贴纸消息),落 blob + documents。
|
||||
func (s *Service) CreateDocumentFromUpload(ctx context.Context, file domain.UploadedFileRef, spec domain.DocumentSpec) (domain.Document, error) {
|
||||
intentHash, err := uploadedMediaIntentHash(domain.UploadedMediaDocument, file, &spec)
|
||||
if err != nil {
|
||||
return domain.Document{}, err
|
||||
}
|
||||
if doc, found, err := s.replayUploadedDocument(ctx, file, intentHash); err != nil || found {
|
||||
return doc, err
|
||||
}
|
||||
body, err := s.assembleUploadBlob(ctx, file.OwnerUserID, file.FileID, file.Parts)
|
||||
if err != nil {
|
||||
return domain.Document{}, err
|
||||
|
|
@ -234,11 +267,13 @@ func (s *Service) CreateDocumentFromUpload(ctx context.Context, file domain.Uplo
|
|||
DCID: s.dc,
|
||||
Attributes: spec.Attributes,
|
||||
}
|
||||
thumbMaterialized := false
|
||||
if spec.Thumb != nil {
|
||||
thumbData, err := s.assembleUpload(ctx, spec.Thumb.OwnerUserID, spec.Thumb.FileID, spec.Thumb.Parts)
|
||||
thumbData, err := s.readUploadBytes(ctx, spec.Thumb.OwnerUserID, spec.Thumb.FileID, spec.Thumb.Parts)
|
||||
if err == nil && len(thumbData) > 0 {
|
||||
if thumb, err := s.putDocumentThumb(ctx, docID, thumbData); err == nil {
|
||||
doc.Thumbs = []domain.PhotoSize{thumb}
|
||||
thumbMaterialized = true
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -250,12 +285,23 @@ func (s *Service) CreateDocumentFromUpload(ctx context.Context, file domain.Uplo
|
|||
if err := s.media.PutDocument(ctx, doc); err != nil {
|
||||
return domain.Document{}, err
|
||||
}
|
||||
if err := s.cleanupUploadParts(ctx, file.OwnerUserID, file.FileID); err != nil {
|
||||
s.log.Warn("cleanup assembled document upload parts failed",
|
||||
zap.Int64("owner_user_id", file.OwnerUserID),
|
||||
zap.Int64("file_id", file.FileID),
|
||||
zap.Int64("document_id", docID),
|
||||
zap.Error(err))
|
||||
receipt, err := s.commitUploadedMediaReceipt(ctx, file, domain.UploadedMediaDocument, intentHash, doc.ID)
|
||||
if err != nil {
|
||||
return domain.Document{}, err
|
||||
}
|
||||
if receipt.MediaID != doc.ID {
|
||||
winner, found, err := s.media.GetDocument(ctx, receipt.MediaID)
|
||||
if err != nil {
|
||||
return domain.Document{}, err
|
||||
}
|
||||
if !found {
|
||||
return domain.Document{}, fmt.Errorf("concurrent upload receipt references missing document %d", receipt.MediaID)
|
||||
}
|
||||
doc = winner
|
||||
}
|
||||
s.cleanupMaterializedUpload(ctx, file, "document materialized")
|
||||
if spec.Thumb != nil && thumbMaterialized {
|
||||
s.cleanupMaterializedUpload(ctx, *spec.Thumb, "document thumbnail materialized")
|
||||
}
|
||||
return doc, nil
|
||||
}
|
||||
|
|
@ -277,6 +323,7 @@ var faststartVideoMimes = map[string]bool{
|
|||
// 此时只发生几次 16 字节读,不读整段媒体。
|
||||
// 2. 仅 moov 在末尾时才重写;且优先走流式(仅 ftyp+moov 进内存,mdat 大块分块流式拼接),
|
||||
// 不把整段视频 2× 驻留内存。moov 非末尾的罕见排布回退到全量重排。
|
||||
//
|
||||
// 任何不适用/失败都返回原 body,绝不让上传失败或损坏数据。
|
||||
func (s *Service) maybeFaststartVideoBlob(ctx context.Context, mimeType string, body assembledUploadBlob) assembledUploadBlob {
|
||||
if !faststartVideoMimes[strings.ToLower(strings.TrimSpace(mimeType))] {
|
||||
|
|
|
|||
|
|
@ -26,6 +26,7 @@ type fakeMediaStore struct {
|
|||
parts map[string][]domain.UploadPart
|
||||
webPages map[int64]domain.MessageWebPage
|
||||
seedState map[string]string
|
||||
receipts map[string]domain.UploadedMediaReceipt
|
||||
}
|
||||
|
||||
func newFakeMediaStore() *fakeMediaStore {
|
||||
|
|
@ -36,9 +37,35 @@ func newFakeMediaStore() *fakeMediaStore {
|
|||
sets: map[int64]domain.StickerSet{},
|
||||
parts: map[string][]domain.UploadPart{},
|
||||
seedState: map[string]string{},
|
||||
receipts: map[string]domain.UploadedMediaReceipt{},
|
||||
}
|
||||
}
|
||||
|
||||
func fakeUploadReceiptKey(ownerUserID, fileID int64) string {
|
||||
return fmt.Sprintf("%d/%d", ownerUserID, fileID)
|
||||
}
|
||||
|
||||
func (f *fakeMediaStore) GetUploadedMediaReceipt(_ context.Context, ownerUserID, fileID int64) (domain.UploadedMediaReceipt, bool, error) {
|
||||
f.mu.Lock()
|
||||
defer f.mu.Unlock()
|
||||
receipt, ok := f.receipts[fakeUploadReceiptKey(ownerUserID, fileID)]
|
||||
receipt.IntentHash = append([]byte(nil), receipt.IntentHash...)
|
||||
return receipt, ok, nil
|
||||
}
|
||||
|
||||
func (f *fakeMediaStore) PutUploadedMediaReceipt(_ context.Context, receipt domain.UploadedMediaReceipt) (domain.UploadedMediaReceipt, bool, error) {
|
||||
f.mu.Lock()
|
||||
defer f.mu.Unlock()
|
||||
key := fakeUploadReceiptKey(receipt.OwnerUserID, receipt.FileID)
|
||||
if stored, ok := f.receipts[key]; ok {
|
||||
stored.IntentHash = append([]byte(nil), stored.IntentHash...)
|
||||
return stored, false, nil
|
||||
}
|
||||
receipt.IntentHash = append([]byte(nil), receipt.IntentHash...)
|
||||
f.receipts[key] = receipt
|
||||
return receipt, true, nil
|
||||
}
|
||||
|
||||
func (f *fakeMediaStore) SaveFilePart(_ context.Context, part domain.UploadPart) error {
|
||||
f.mu.Lock()
|
||||
defer f.mu.Unlock()
|
||||
|
|
|
|||
|
|
@ -514,6 +514,20 @@ func orderDocuments(docs []domain.Document, ids []int64) []domain.Document {
|
|||
// assembleUpload 把已上传分片按 part 顺序拼成完整字节,并清理分片。
|
||||
// expectedParts>0 时校验分片连续且齐全。
|
||||
func (s *Service) assembleUpload(ctx context.Context, ownerUserID, fileID int64, expectedParts int) ([]byte, error) {
|
||||
buf, err := s.readUploadBytes(ctx, ownerUserID, fileID, expectedParts)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := s.cleanupUploadParts(ctx, ownerUserID, fileID); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return buf, nil
|
||||
}
|
||||
|
||||
// readUploadBytes validates and reads all parts without consuming them. Message-media
|
||||
// materialization persists an upload receipt before cleanup; callers that do not need replayability
|
||||
// continue to use assembleUpload.
|
||||
func (s *Service) readUploadBytes(ctx context.Context, ownerUserID, fileID int64, expectedParts int) ([]byte, error) {
|
||||
parts, _, err := s.loadAndValidateUploadParts(ctx, ownerUserID, fileID, expectedParts)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
|
|
@ -532,9 +546,6 @@ func (s *Service) assembleUpload(ctx context.Context, ownerUserID, fileID int64,
|
|||
}
|
||||
buf = append(buf, data...)
|
||||
}
|
||||
if err := s.cleanupUploadParts(ctx, ownerUserID, fileID); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return buf, nil
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -124,6 +124,51 @@ func TestCreateDocumentFromUploadStreamsBodyAndCleansParts(t *testing.T) {
|
|||
if string(body) != strings.Join(parts, "") {
|
||||
t.Fatalf("body blob mismatch")
|
||||
}
|
||||
replayed, err := svc.CreateDocumentFromUpload(ctx,
|
||||
domain.UploadedFileRef{OwnerUserID: 10, FileID: 200, Parts: len(parts), Name: "large.bin", Big: true},
|
||||
domain.DocumentSpec{MimeType: "application/octet-stream"},
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatalf("replay CreateDocumentFromUpload after part cleanup: %v", err)
|
||||
}
|
||||
if replayed.ID != doc.ID || replayed.AccessHash != doc.AccessHash {
|
||||
t.Fatalf("replayed document = %d/%d, want original %d/%d", replayed.ID, replayed.AccessHash, doc.ID, doc.AccessHash)
|
||||
}
|
||||
if _, err := svc.CreateDocumentFromUpload(ctx,
|
||||
domain.UploadedFileRef{OwnerUserID: 10, FileID: 200, Parts: len(parts), Name: "large.bin", Big: true},
|
||||
domain.DocumentSpec{MimeType: "text/plain"},
|
||||
); !errors.Is(err, domain.ErrFilePartsInvalid) {
|
||||
t.Fatalf("changed materialization intent err = %v, want ErrFilePartsInvalid", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCreatePhotoFromUploadReceiptReplaysAfterPartCleanup(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
media := newFakeMediaStore()
|
||||
svc, _ := newUploadPartTestService(t, media, domain.UploadPartQuota{})
|
||||
file := domain.UploadedFileRef{OwnerUserID: 10, FileID: 201, Parts: 1, Name: "photo.jpg"}
|
||||
if _, err := svc.SaveFilePart(ctx, file.OwnerUserID, file.FileID, 0, []byte("image-bytes")); err != nil {
|
||||
t.Fatalf("SaveFilePart: %v", err)
|
||||
}
|
||||
first, err := svc.CreatePhotoFromUpload(ctx, file)
|
||||
if err != nil {
|
||||
t.Fatalf("CreatePhotoFromUpload: %v", err)
|
||||
}
|
||||
if remaining, err := media.LoadFileParts(ctx, file.OwnerUserID, file.FileID); err != nil || len(remaining) != 0 {
|
||||
t.Fatalf("upload parts after photo materialization = %+v err=%v", remaining, err)
|
||||
}
|
||||
replayed, err := svc.CreatePhotoFromUpload(ctx, file)
|
||||
if err != nil {
|
||||
t.Fatalf("replay CreatePhotoFromUpload: %v", err)
|
||||
}
|
||||
if replayed.ID != first.ID || replayed.AccessHash != first.AccessHash {
|
||||
t.Fatalf("replayed photo = %d/%d, want original %d/%d", replayed.ID, replayed.AccessHash, first.ID, first.AccessHash)
|
||||
}
|
||||
changed := file
|
||||
changed.Name = "different.jpg"
|
||||
if _, err := svc.CreatePhotoFromUpload(ctx, changed); !errors.Is(err, domain.ErrFilePartsInvalid) {
|
||||
t.Fatalf("changed photo intent err = %v, want ErrFilePartsInvalid", err)
|
||||
}
|
||||
}
|
||||
|
||||
type countingUploadPartBackend struct {
|
||||
|
|
|
|||
106
internal/app/files/upload_receipt.go
Normal file
106
internal/app/files/upload_receipt.go
Normal file
|
|
@ -0,0 +1,106 @@
|
|||
package files
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
|
||||
"go.uber.org/zap"
|
||||
|
||||
"telesrv/internal/domain"
|
||||
)
|
||||
|
||||
const uploadedMediaIntentVersion = 1
|
||||
|
||||
type uploadedMediaIntent struct {
|
||||
Version int `json:"version"`
|
||||
Kind domain.UploadedMediaKind `json:"kind"`
|
||||
File domain.UploadedFileRef `json:"file"`
|
||||
Spec *domain.DocumentSpec `json:"spec,omitempty"`
|
||||
}
|
||||
|
||||
func uploadedMediaIntentHash(kind domain.UploadedMediaKind, file domain.UploadedFileRef, spec *domain.DocumentSpec) ([]byte, error) {
|
||||
payload, err := json.Marshal(uploadedMediaIntent{
|
||||
Version: uploadedMediaIntentVersion,
|
||||
Kind: kind,
|
||||
File: file,
|
||||
Spec: spec,
|
||||
})
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("marshal uploaded media intent: %w", err)
|
||||
}
|
||||
sum := sha256.Sum256(payload)
|
||||
return sum[:], nil
|
||||
}
|
||||
|
||||
func sameUploadedMediaReceipt(receipt domain.UploadedMediaReceipt, kind domain.UploadedMediaKind, intentHash []byte) bool {
|
||||
return receipt.Kind == kind && len(intentHash) == sha256.Size && bytes.Equal(receipt.IntentHash, intentHash)
|
||||
}
|
||||
|
||||
func (s *Service) replayUploadedPhoto(ctx context.Context, file domain.UploadedFileRef, intentHash []byte) (domain.Photo, bool, error) {
|
||||
receipt, found, err := s.media.GetUploadedMediaReceipt(ctx, file.OwnerUserID, file.FileID)
|
||||
if err != nil || !found {
|
||||
return domain.Photo{}, false, err
|
||||
}
|
||||
if !sameUploadedMediaReceipt(receipt, domain.UploadedMediaPhoto, intentHash) {
|
||||
return domain.Photo{}, false, domain.ErrFilePartsInvalid
|
||||
}
|
||||
photo, found, err := s.media.GetPhoto(ctx, receipt.MediaID)
|
||||
if err != nil {
|
||||
return domain.Photo{}, false, err
|
||||
}
|
||||
if !found {
|
||||
return domain.Photo{}, false, fmt.Errorf("uploaded photo receipt %d/%d references missing photo %d", file.OwnerUserID, file.FileID, receipt.MediaID)
|
||||
}
|
||||
s.cleanupMaterializedUpload(ctx, file, "photo replay")
|
||||
return photo, true, nil
|
||||
}
|
||||
|
||||
func (s *Service) replayUploadedDocument(ctx context.Context, file domain.UploadedFileRef, intentHash []byte) (domain.Document, bool, error) {
|
||||
receipt, found, err := s.media.GetUploadedMediaReceipt(ctx, file.OwnerUserID, file.FileID)
|
||||
if err != nil || !found {
|
||||
return domain.Document{}, false, err
|
||||
}
|
||||
if !sameUploadedMediaReceipt(receipt, domain.UploadedMediaDocument, intentHash) {
|
||||
return domain.Document{}, false, domain.ErrFilePartsInvalid
|
||||
}
|
||||
doc, found, err := s.media.GetDocument(ctx, receipt.MediaID)
|
||||
if err != nil {
|
||||
return domain.Document{}, false, err
|
||||
}
|
||||
if !found {
|
||||
return domain.Document{}, false, fmt.Errorf("uploaded document receipt %d/%d references missing document %d", file.OwnerUserID, file.FileID, receipt.MediaID)
|
||||
}
|
||||
s.cleanupMaterializedUpload(ctx, file, "document replay")
|
||||
return doc, true, nil
|
||||
}
|
||||
|
||||
func (s *Service) commitUploadedMediaReceipt(ctx context.Context, file domain.UploadedFileRef, kind domain.UploadedMediaKind, intentHash []byte, mediaID int64) (domain.UploadedMediaReceipt, error) {
|
||||
receipt, _, err := s.media.PutUploadedMediaReceipt(ctx, domain.UploadedMediaReceipt{
|
||||
OwnerUserID: file.OwnerUserID,
|
||||
FileID: file.FileID,
|
||||
IntentHash: intentHash,
|
||||
Kind: kind,
|
||||
MediaID: mediaID,
|
||||
})
|
||||
if err != nil {
|
||||
return domain.UploadedMediaReceipt{}, err
|
||||
}
|
||||
if !sameUploadedMediaReceipt(receipt, kind, intentHash) {
|
||||
return domain.UploadedMediaReceipt{}, domain.ErrFilePartsInvalid
|
||||
}
|
||||
return receipt, nil
|
||||
}
|
||||
|
||||
func (s *Service) cleanupMaterializedUpload(ctx context.Context, file domain.UploadedFileRef, reason string) {
|
||||
if err := s.cleanupUploadParts(ctx, file.OwnerUserID, file.FileID); err != nil {
|
||||
s.log.Warn("cleanup materialized upload parts failed",
|
||||
zap.String("reason", reason),
|
||||
zap.Int64("owner_user_id", file.OwnerUserID),
|
||||
zap.Int64("file_id", file.FileID),
|
||||
zap.Error(err),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
|
@ -17,12 +17,49 @@ type TempAuthKeyRetentionStore interface {
|
|||
DeleteExpired(ctx context.Context, expiredBefore int64, limit int) (int, error)
|
||||
}
|
||||
|
||||
// OrphanAuthKeyRetentionStore 回收从未形成授权/temp binding 的旧握手 key。
|
||||
// protected 是当前连接注册表实际使用的 raw auth_key_id 快照。
|
||||
type OrphanAuthKeyRetentionStore interface {
|
||||
DeleteOrphaned(ctx context.Context, olderThan time.Duration, limit int, protected [][8]byte) (int, error)
|
||||
}
|
||||
|
||||
type ActiveRawAuthKeyProvider interface {
|
||||
ActiveRawAuthKeyIDs() [][8]byte
|
||||
}
|
||||
|
||||
// ActiveAuthKeyHeartbeatStore 把本实例仍在使用的 raw auth key 活性持久化。多实例下
|
||||
// orphan GC 不能只看当前进程的 active 快照;其它实例的 heartbeat 会推进数据库
|
||||
// last_used_at,使它们不会被误判为孤儿。
|
||||
type ActiveAuthKeyHeartbeatStore interface {
|
||||
TouchActiveRawAuthKeys(ctx context.Context, ids [][8]byte) error
|
||||
}
|
||||
|
||||
// BotAPIUpdateRetentionStore 回收 Bot API getUpdates 投递队列的死行(性能审计 H1):
|
||||
// 已确认且超过宽限期的行 + 按消息 date 超过保留期的行(官方 Bot API updates 最多保留 24h)。
|
||||
type BotAPIUpdateRetentionStore interface {
|
||||
DeleteDeliveredOrExpired(ctx context.Context, confirmedGrace, maxAge time.Duration, limit int) (int, error)
|
||||
}
|
||||
|
||||
// UserUpdateEventRetentionStore 只回收所有当前授权设备都明确确认过的账号事件前缀。
|
||||
// 它不是普通 TTL:任一授权缺 state 时确认水位为 0,不得删除该设备可能仍需的事件。
|
||||
type UserUpdateEventRetentionStore interface {
|
||||
DeleteConfirmedPrefix(ctx context.Context, olderThan time.Duration, limit int) (int, error)
|
||||
}
|
||||
|
||||
// ChannelUpdateEventRetentionStore 回收超过保留期的 channel durable update 连续前缀。
|
||||
// 具体 store 必须在同一事务内删除事件并推进 retained floor,低于 floor 的客户端由
|
||||
// updates.getChannelDifference 走 channelDifferenceTooLong 快照恢复。
|
||||
type ChannelUpdateEventRetentionStore interface {
|
||||
DeleteExpiredChannelUpdateEvents(ctx context.Context, olderThan time.Duration, limit int) (int, error)
|
||||
}
|
||||
|
||||
// LoginCodeDeliveryRetentionStore reclaims only compact idempotency receipts
|
||||
// after their associated opaque code lifetime. It must not delete the message,
|
||||
// durable update event, or outbox facts created by the delivery transaction.
|
||||
type LoginCodeDeliveryRetentionStore interface {
|
||||
DeleteExpiredLoginCodeDeliveries(ctx context.Context, expiredBefore time.Time, limit int) (int, error)
|
||||
}
|
||||
|
||||
// botAPIConfirmedGrace 是已确认 Bot API update 行的删除宽限:确认水位之下的行不会再被
|
||||
// getUpdates 读取(fromID 恒 > confirmed),宽限仅防御 offset 回拨调试;回收目标是清堆积。
|
||||
const botAPIConfirmedGrace = 15 * time.Minute
|
||||
|
|
@ -32,23 +69,39 @@ const botAPIConfirmedGrace = 15 * time.Minute
|
|||
// 连接;回收目标是清堆积,晚一天无妨。
|
||||
const tempAuthKeyExpiryGrace = 24 * time.Hour
|
||||
|
||||
const (
|
||||
// terminal failed outbox 只承担短期诊断隔离;它不是 durable update log。
|
||||
// 删除该任务会由 head trigger 立即放行同账号下一 pts,而 user_update_events
|
||||
// 继续保留,在线漏推由正常 difference 路径补偿。
|
||||
defaultOutboxPoisonRetention = time.Minute
|
||||
defaultOutboxPoisonInterval = 15 * time.Second
|
||||
)
|
||||
|
||||
// RetentionWorker 周期性回收存储中的死数据。
|
||||
//
|
||||
// 注意:本 worker 刻意不清理 user_update_events —— pts log 永久保留。原因:TDesktop 不支持
|
||||
// 账号级 updates.differenceTooLong(api_updates.cpp 收到该响应只打一行日志,且漏掉
|
||||
// setRequesting(false),会永久锁死整个 update 引擎),服务端因此无法让"落后超过保留期"的
|
||||
// 客户端整库重置;一旦裁剪 events,落后客户端的 getDifference 会拿到不完整的事件链而静默
|
||||
// 丢消息。详见 docs/performance-audit.md 与 docs/compatibility-matrix.md。user_update_events
|
||||
// 长期膨胀作为已知 todo。
|
||||
// 注意:TDesktop 不支持账号级 updates.differenceTooLong(api_updates.cpp 收到该响应只
|
||||
// 记录日志且不清 requesting,会永久锁死 update 引擎),因此绝不能按普通 TTL 硬裁剪
|
||||
// user_update_events。本 worker 只允许 store 删除“所有当前授权设备都明确确认”的连续安全
|
||||
// 前缀;落后或缺 state 的任一设备都会把 floor 压回 0。客户端偶然带回已确认前的旧 pts 时,
|
||||
// updates 服务通过普通 differenceSlice checkpoint 推进,不发送 differenceTooLong。
|
||||
type RetentionWorker struct {
|
||||
outbox DispatchOutboxRetentionStore
|
||||
tempKeys TempAuthKeyRetentionStore // 可为 nil(不回收 temp key 绑定)
|
||||
botAPIUpdates BotAPIUpdateRetentionStore // 可为 nil(不回收 Bot API 队列)
|
||||
logger *zap.Logger
|
||||
retention time.Duration
|
||||
botAPIRetention time.Duration
|
||||
interval time.Duration
|
||||
batch int
|
||||
outbox DispatchOutboxRetentionStore
|
||||
tempKeys TempAuthKeyRetentionStore // 可为 nil(不回收 temp key 绑定)
|
||||
botAPIUpdates BotAPIUpdateRetentionStore // 可为 nil(不回收 Bot API 队列)
|
||||
userUpdates UserUpdateEventRetentionStore
|
||||
channelUpdates ChannelUpdateEventRetentionStore
|
||||
loginCodeDeliveries LoginCodeDeliveryRetentionStore
|
||||
orphanAuthKeys OrphanAuthKeyRetentionStore
|
||||
activeAuthKeys ActiveRawAuthKeyProvider
|
||||
activeAuthKeyHeartbeat ActiveAuthKeyHeartbeatStore
|
||||
logger *zap.Logger
|
||||
retention time.Duration
|
||||
botAPIRetention time.Duration
|
||||
orphanRetention time.Duration
|
||||
outboxPoisonRetention time.Duration
|
||||
outboxPoisonInterval time.Duration
|
||||
interval time.Duration
|
||||
batch int
|
||||
}
|
||||
|
||||
func NewRetentionWorker(outbox DispatchOutboxRetentionStore, tempKeys TempAuthKeyRetentionStore, logger *zap.Logger, retention, interval time.Duration, batch int) *RetentionWorker {
|
||||
|
|
@ -65,15 +118,32 @@ func NewRetentionWorker(outbox DispatchOutboxRetentionStore, tempKeys TempAuthKe
|
|||
batch = 10000
|
||||
}
|
||||
return &RetentionWorker{
|
||||
outbox: outbox,
|
||||
tempKeys: tempKeys,
|
||||
logger: logger,
|
||||
retention: retention,
|
||||
interval: interval,
|
||||
batch: batch,
|
||||
outbox: outbox,
|
||||
tempKeys: tempKeys,
|
||||
logger: logger,
|
||||
retention: retention,
|
||||
outboxPoisonRetention: defaultOutboxPoisonRetention,
|
||||
outboxPoisonInterval: defaultOutboxPoisonInterval,
|
||||
interval: interval,
|
||||
batch: batch,
|
||||
}
|
||||
}
|
||||
|
||||
// WithDispatchOutboxPoisonPolicy 配置 terminal failed head 的独立短隔离与清理周期。
|
||||
// 该周期不能复用 durable update 的周级保留期,否则一条确定性构造错误会冻结该
|
||||
// 用户整条在线 pts lane。<=0 分别回退到 1m/15s 的安全默认值。
|
||||
func (w *RetentionWorker) WithDispatchOutboxPoisonPolicy(retention, interval time.Duration) *RetentionWorker {
|
||||
if retention <= 0 {
|
||||
retention = defaultOutboxPoisonRetention
|
||||
}
|
||||
if interval <= 0 {
|
||||
interval = defaultOutboxPoisonInterval
|
||||
}
|
||||
w.outboxPoisonRetention = retention
|
||||
w.outboxPoisonInterval = interval
|
||||
return w
|
||||
}
|
||||
|
||||
// WithBotAPIUpdateRetention 启用 bot_api_updates 队列回收;retention <=0 时用官方语义默认 24h。
|
||||
func (w *RetentionWorker) WithBotAPIUpdateRetention(store BotAPIUpdateRetentionStore, retention time.Duration) *RetentionWorker {
|
||||
if retention <= 0 {
|
||||
|
|
@ -84,26 +154,102 @@ func (w *RetentionWorker) WithBotAPIUpdateRetention(store BotAPIUpdateRetentionS
|
|||
return w
|
||||
}
|
||||
|
||||
// WithUserUpdateRetention 启用账号 update 的共同确认安全前缀回收。TDesktop 不支持
|
||||
// account differenceTooLong,具体 store 必须保证未确认前缀永不删除。
|
||||
func (w *RetentionWorker) WithUserUpdateRetention(store UserUpdateEventRetentionStore) *RetentionWorker {
|
||||
w.userUpdates = store
|
||||
return w
|
||||
}
|
||||
|
||||
// WithChannelUpdateRetention 启用 channel durable update 的有界 TTL 回收;复用 worker 的
|
||||
// retention/interval/batch,并由 store 的 retained floor 保证旧 pts 不会读到静默空洞。
|
||||
func (w *RetentionWorker) WithChannelUpdateRetention(store ChannelUpdateEventRetentionStore) *RetentionWorker {
|
||||
w.channelUpdates = store
|
||||
return w
|
||||
}
|
||||
|
||||
// WithLoginCodeDeliveryRetention enables bounded seek cleanup for compact
|
||||
// phone_code_hash receipts. Each row carries its own expiry derived from the
|
||||
// code TTL, so this cleanup intentionally does not reuse update-log retention.
|
||||
func (w *RetentionWorker) WithLoginCodeDeliveryRetention(store LoginCodeDeliveryRetentionStore) *RetentionWorker {
|
||||
w.loginCodeDeliveries = store
|
||||
return w
|
||||
}
|
||||
|
||||
// WithOrphanAuthKeyRetention 启用未授权握手 key 的有界回收。active 必须提供 raw key,
|
||||
// 不能提供 temp→perm business key;否则未登录或 PFS 连接会被误判为 orphan。
|
||||
func (w *RetentionWorker) WithOrphanAuthKeyRetention(store OrphanAuthKeyRetentionStore, active ActiveRawAuthKeyProvider, retention time.Duration) *RetentionWorker {
|
||||
w.orphanAuthKeys = store
|
||||
w.activeAuthKeys = active
|
||||
w.activeAuthKeyHeartbeat, _ = store.(ActiveAuthKeyHeartbeatStore)
|
||||
w.orphanRetention = retention
|
||||
return w
|
||||
}
|
||||
|
||||
func (w *RetentionWorker) Run(ctx context.Context) {
|
||||
w.runOnce(ctx)
|
||||
ticker := time.NewTicker(w.interval)
|
||||
defer ticker.Stop()
|
||||
retentionTicker := time.NewTicker(w.interval)
|
||||
defer retentionTicker.Stop()
|
||||
poisonTicker := time.NewTicker(w.outboxPoisonInterval)
|
||||
defer poisonTicker.Stop()
|
||||
var (
|
||||
heartbeatTicker *time.Ticker
|
||||
heartbeatC <-chan time.Time
|
||||
)
|
||||
if interval := w.orphanHeartbeatInterval(); interval > 0 {
|
||||
heartbeatTicker = time.NewTicker(interval)
|
||||
heartbeatC = heartbeatTicker.C
|
||||
defer heartbeatTicker.Stop()
|
||||
}
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-ticker.C:
|
||||
w.runOnce(ctx)
|
||||
case <-retentionTicker.C:
|
||||
w.runRetentionOnce(ctx)
|
||||
case <-poisonTicker.C:
|
||||
w.runOutboxPoisonOnce(ctx)
|
||||
case <-heartbeatC:
|
||||
w.heartbeatActiveAuthKeys(ctx)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (w *RetentionWorker) runOnce(ctx context.Context) {
|
||||
outboxDeleted, err := w.outbox.DeleteFailed(ctx, w.retention, w.batch)
|
||||
w.runOutboxPoisonOnce(ctx)
|
||||
w.runRetentionOnce(ctx)
|
||||
}
|
||||
|
||||
func (w *RetentionWorker) runOutboxPoisonOnce(ctx context.Context) {
|
||||
if w.outbox == nil {
|
||||
return
|
||||
}
|
||||
outboxDeleted, err := w.outbox.DeleteFailed(ctx, w.outboxPoisonRetention, w.batch)
|
||||
if err != nil {
|
||||
w.logger.Warn("清理 failed dispatch_outbox 失败", zap.Error(err))
|
||||
w.logger.Error("清理 terminal failed dispatch_outbox 失败",
|
||||
zap.String("signal", "dispatch_outbox_poison_cleanup_failed"),
|
||||
zap.Duration("quarantine", w.outboxPoisonRetention),
|
||||
zap.Error(err),
|
||||
)
|
||||
} else if outboxDeleted > 0 {
|
||||
w.logger.Info("清理 failed dispatch_outbox 完成", zap.Int("deleted", outboxDeleted))
|
||||
// Error 级结构化信号刻意保留:发生 terminal failed 代表确定性编码、事件缺失
|
||||
// 或其它不可自动重试故障。任务删除只解冻在线 lane,不会删除 durable event。
|
||||
w.logger.Error("terminal failed dispatch_outbox 已结束隔离并释放用户 lane",
|
||||
zap.String("signal", "dispatch_outbox_poison_released"),
|
||||
zap.Int("deleted", outboxDeleted),
|
||||
zap.Duration("quarantine", w.outboxPoisonRetention),
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
func (w *RetentionWorker) runRetentionOnce(ctx context.Context) {
|
||||
if w.loginCodeDeliveries != nil {
|
||||
deleted, err := w.loginCodeDeliveries.DeleteExpiredLoginCodeDeliveries(ctx, time.Now(), w.batch)
|
||||
if err != nil {
|
||||
w.logger.Warn("回收过期 login-code delivery 回执失败", zap.Error(err))
|
||||
} else if deleted > 0 {
|
||||
w.logger.Info("回收过期 login-code delivery 回执完成", zap.Int("deleted", deleted))
|
||||
}
|
||||
}
|
||||
if w.tempKeys != nil {
|
||||
expiredBefore := time.Now().Add(-tempAuthKeyExpiryGrace).Unix()
|
||||
|
|
@ -114,6 +260,24 @@ func (w *RetentionWorker) runOnce(ctx context.Context) {
|
|||
w.logger.Info("回收过期 temp auth key 绑定完成", zap.Int("deleted", tempDeleted))
|
||||
}
|
||||
}
|
||||
if w.orphanAuthKeys != nil && w.orphanRetention > 0 {
|
||||
var protected [][8]byte
|
||||
if w.activeAuthKeys != nil {
|
||||
protected = w.activeAuthKeys.ActiveRawAuthKeyIDs()
|
||||
}
|
||||
if !w.touchActiveAuthKeys(ctx, protected) {
|
||||
// Fail safe: if this instance cannot publish its own active set, deleting against a
|
||||
// stale database heartbeat could evict keys used by another instance too. Keep all
|
||||
// candidates for this pass and retry after the next heartbeat.
|
||||
} else {
|
||||
orphanDeleted, err := w.orphanAuthKeys.DeleteOrphaned(ctx, w.orphanRetention, w.batch, protected)
|
||||
if err != nil {
|
||||
w.logger.Warn("回收未授权 orphan auth key 失败", zap.Error(err))
|
||||
} else if orphanDeleted > 0 {
|
||||
w.logger.Info("回收未授权 orphan auth key 完成", zap.Int("deleted", orphanDeleted))
|
||||
}
|
||||
}
|
||||
}
|
||||
if w.botAPIUpdates != nil {
|
||||
botAPIDeleted, err := w.botAPIUpdates.DeleteDeliveredOrExpired(ctx, botAPIConfirmedGrace, w.botAPIRetention, w.batch)
|
||||
if err != nil {
|
||||
|
|
@ -122,4 +286,63 @@ func (w *RetentionWorker) runOnce(ctx context.Context) {
|
|||
w.logger.Info("回收 bot_api_updates 队列完成", zap.Int("deleted", botAPIDeleted))
|
||||
}
|
||||
}
|
||||
if w.userUpdates != nil {
|
||||
userDeleted, err := w.userUpdates.DeleteConfirmedPrefix(ctx, w.retention, w.batch)
|
||||
if err != nil {
|
||||
w.logger.Warn("回收已共同确认的 user_update_events 前缀失败", zap.Error(err))
|
||||
} else if userDeleted > 0 {
|
||||
w.logger.Info("回收已共同确认的 user_update_events 前缀完成", zap.Int("deleted", userDeleted))
|
||||
}
|
||||
}
|
||||
if w.channelUpdates != nil {
|
||||
channelDeleted, err := w.channelUpdates.DeleteExpiredChannelUpdateEvents(ctx, w.retention, w.batch)
|
||||
if err != nil {
|
||||
// store 会逐频道隔离坏 gap 后继续本轮;deleted 可能非零,必须同时记录,
|
||||
// 既不能把全局 pass 伪装成完全失败,也不能吞掉不变量错误。
|
||||
w.logger.Warn("回收过期 channel_update_events 存在隔离频道",
|
||||
zap.Int("deleted", channelDeleted),
|
||||
zap.Error(err),
|
||||
)
|
||||
} else if channelDeleted > 0 {
|
||||
w.logger.Info("回收过期 channel_update_events 连续前缀完成", zap.Int("deleted", channelDeleted))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (w *RetentionWorker) orphanHeartbeatInterval() time.Duration {
|
||||
if w.activeAuthKeyHeartbeat == nil || w.activeAuthKeys == nil || w.orphanRetention <= 0 {
|
||||
return 0
|
||||
}
|
||||
interval := w.orphanRetention / 3
|
||||
if interval <= 0 {
|
||||
interval = time.Nanosecond
|
||||
}
|
||||
if w.interval > 0 && w.interval < interval {
|
||||
interval = w.interval
|
||||
}
|
||||
return interval
|
||||
}
|
||||
|
||||
func (w *RetentionWorker) heartbeatActiveAuthKeys(ctx context.Context) {
|
||||
if w.activeAuthKeys == nil {
|
||||
return
|
||||
}
|
||||
w.touchActiveAuthKeys(ctx, w.activeAuthKeys.ActiveRawAuthKeyIDs())
|
||||
}
|
||||
|
||||
// touchActiveAuthKeys returns false only when a configured durable heartbeat failed. A store that
|
||||
// predates the optional heartbeat interface keeps single-instance behavior.
|
||||
func (w *RetentionWorker) touchActiveAuthKeys(ctx context.Context, protected [][8]byte) bool {
|
||||
if w.activeAuthKeyHeartbeat == nil {
|
||||
return true
|
||||
}
|
||||
if err := w.activeAuthKeyHeartbeat.TouchActiveRawAuthKeys(ctx, protected); err != nil {
|
||||
w.logger.Error("刷新 active raw auth key heartbeat 失败,本轮跳过 orphan GC",
|
||||
zap.String("signal", "auth_key_heartbeat_failed"),
|
||||
zap.Int("active_keys", len(protected)),
|
||||
zap.Error(err),
|
||||
)
|
||||
return false
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
|
|
|||
|
|
@ -2,19 +2,57 @@ package maintenance
|
|||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"go.uber.org/zap"
|
||||
"go.uber.org/zap/zapcore"
|
||||
"go.uber.org/zap/zaptest/observer"
|
||||
)
|
||||
|
||||
type fakeOutboxRetention struct {
|
||||
calls int
|
||||
calls int
|
||||
olderThan time.Duration
|
||||
limit int
|
||||
deleted int
|
||||
}
|
||||
|
||||
func (f *fakeOutboxRetention) DeleteFailed(context.Context, time.Duration, int) (int, error) {
|
||||
func (f *fakeOutboxRetention) DeleteFailed(_ context.Context, olderThan time.Duration, limit int) (int, error) {
|
||||
f.calls++
|
||||
return 0, nil
|
||||
f.olderThan = olderThan
|
||||
f.limit = limit
|
||||
return f.deleted, nil
|
||||
}
|
||||
|
||||
func TestRetentionWorkerUsesIndependentOutboxPoisonPolicyAndSignalsRelease(t *testing.T) {
|
||||
core, logs := observer.New(zapcore.ErrorLevel)
|
||||
outbox := &fakeOutboxRetention{deleted: 2}
|
||||
w := NewRetentionWorker(outbox, nil, zap.New(core), 7*24*time.Hour, time.Hour, 73).
|
||||
WithDispatchOutboxPoisonPolicy(2*time.Minute, 7*time.Second)
|
||||
|
||||
w.runOnce(context.Background())
|
||||
|
||||
if outbox.calls != 1 || outbox.olderThan != 2*time.Minute || outbox.limit != 73 {
|
||||
t.Fatalf("outbox poison calls/args = %d/%v/%d, want 1/2m/73", outbox.calls, outbox.olderThan, outbox.limit)
|
||||
}
|
||||
entries := logs.FilterMessage("terminal failed dispatch_outbox 已结束隔离并释放用户 lane").All()
|
||||
if len(entries) != 1 {
|
||||
t.Fatalf("poison release error signals = %d, want 1", len(entries))
|
||||
}
|
||||
if got := entries[0].ContextMap()["signal"]; got != "dispatch_outbox_poison_released" {
|
||||
t.Fatalf("poison signal = %v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRetentionWorkerOutboxPoisonPolicyDefaultsAreShort(t *testing.T) {
|
||||
outbox := &fakeOutboxRetention{}
|
||||
w := NewRetentionWorker(outbox, nil, zap.NewNop(), 168*time.Hour, time.Hour, 100).
|
||||
WithDispatchOutboxPoisonPolicy(0, 0)
|
||||
w.runOutboxPoisonOnce(context.Background())
|
||||
if outbox.olderThan != defaultOutboxPoisonRetention || w.outboxPoisonInterval != defaultOutboxPoisonInterval {
|
||||
t.Fatalf("default poison policy = %v/%v, want %v/%v", outbox.olderThan, w.outboxPoisonInterval, defaultOutboxPoisonRetention, defaultOutboxPoisonInterval)
|
||||
}
|
||||
}
|
||||
|
||||
type fakeTempKeyRetention struct {
|
||||
|
|
@ -65,6 +103,34 @@ type fakeBotAPIRetention struct {
|
|||
limit int
|
||||
}
|
||||
|
||||
type fakeLoginCodeDeliveryRetention struct {
|
||||
calls int
|
||||
expiredBefore time.Time
|
||||
limit int
|
||||
}
|
||||
|
||||
func (f *fakeLoginCodeDeliveryRetention) DeleteExpiredLoginCodeDeliveries(_ context.Context, expiredBefore time.Time, limit int) (int, error) {
|
||||
f.calls++
|
||||
f.expiredBefore = expiredBefore
|
||||
f.limit = limit
|
||||
return 4, nil
|
||||
}
|
||||
|
||||
func TestRetentionWorkerReclaimsExpiredLoginCodeDeliveryReceipts(t *testing.T) {
|
||||
loginCodes := &fakeLoginCodeDeliveryRetention{}
|
||||
w := NewRetentionWorker(&fakeOutboxRetention{}, nil, zap.NewNop(), 168*time.Hour, time.Hour, 83).
|
||||
WithLoginCodeDeliveryRetention(loginCodes)
|
||||
before := time.Now()
|
||||
w.runRetentionOnce(context.Background())
|
||||
after := time.Now()
|
||||
if loginCodes.calls != 1 || loginCodes.limit != 83 {
|
||||
t.Fatalf("login-code retention calls/limit = %d/%d, want 1/83", loginCodes.calls, loginCodes.limit)
|
||||
}
|
||||
if loginCodes.expiredBefore.Before(before) || loginCodes.expiredBefore.After(after) {
|
||||
t.Fatalf("login-code expiry boundary = %v, want within [%v,%v]", loginCodes.expiredBefore, before, after)
|
||||
}
|
||||
}
|
||||
|
||||
func (f *fakeBotAPIRetention) DeleteDeliveredOrExpired(_ context.Context, confirmedGrace, maxAge time.Duration, limit int) (int, error) {
|
||||
f.calls++
|
||||
f.confirmedGrace = confirmedGrace
|
||||
|
|
@ -99,3 +165,151 @@ func TestRetentionWorkerBotAPIRetentionDefaultsTo24h(t *testing.T) {
|
|||
t.Fatalf("default bot api retention = %v, want 24h", botAPI.maxAge)
|
||||
}
|
||||
}
|
||||
|
||||
type fakeUserUpdateRetention struct {
|
||||
calls int
|
||||
olderThan time.Duration
|
||||
limit int
|
||||
}
|
||||
|
||||
func (f *fakeUserUpdateRetention) DeleteConfirmedPrefix(_ context.Context, olderThan time.Duration, limit int) (int, error) {
|
||||
f.calls++
|
||||
f.olderThan = olderThan
|
||||
f.limit = limit
|
||||
return 9, nil
|
||||
}
|
||||
|
||||
func TestRetentionWorkerReclaimsOnlyConfirmedUserUpdatePrefix(t *testing.T) {
|
||||
const retention = 7 * 24 * time.Hour
|
||||
store := &fakeUserUpdateRetention{}
|
||||
w := NewRetentionWorker(&fakeOutboxRetention{}, nil, zap.NewNop(), retention, time.Hour, 91).
|
||||
WithUserUpdateRetention(store)
|
||||
w.runOnce(context.Background())
|
||||
if store.calls != 1 || store.olderThan != retention || store.limit != 91 {
|
||||
t.Fatalf("user update retention calls/args = %d/%v/%d, want 1/%v/91", store.calls, store.olderThan, store.limit, retention)
|
||||
}
|
||||
}
|
||||
|
||||
type fakeChannelUpdateRetention struct {
|
||||
calls int
|
||||
olderThan time.Duration
|
||||
limit int
|
||||
}
|
||||
|
||||
func (f *fakeChannelUpdateRetention) DeleteExpiredChannelUpdateEvents(_ context.Context, olderThan time.Duration, limit int) (int, error) {
|
||||
f.calls++
|
||||
f.olderThan = olderThan
|
||||
f.limit = limit
|
||||
return 7, nil
|
||||
}
|
||||
|
||||
func TestRetentionWorkerReclaimsChannelUpdates(t *testing.T) {
|
||||
const (
|
||||
retention = 14 * 24 * time.Hour
|
||||
batch = 321
|
||||
)
|
||||
channelUpdates := &fakeChannelUpdateRetention{}
|
||||
w := NewRetentionWorker(&fakeOutboxRetention{}, nil, zap.NewNop(), retention, time.Hour, batch).
|
||||
WithChannelUpdateRetention(channelUpdates)
|
||||
|
||||
w.runOnce(context.Background())
|
||||
|
||||
if channelUpdates.calls != 1 {
|
||||
t.Fatalf("channel update retention calls = %d, want 1", channelUpdates.calls)
|
||||
}
|
||||
if channelUpdates.olderThan != retention || channelUpdates.limit != batch {
|
||||
t.Fatalf("channel update retention args = (%v, %d), want (%v, %d)",
|
||||
channelUpdates.olderThan, channelUpdates.limit, retention, batch)
|
||||
}
|
||||
}
|
||||
|
||||
type fakeOrphanAuthKeyRetention struct {
|
||||
calls int
|
||||
olderThan time.Duration
|
||||
limit int
|
||||
protected [][8]byte
|
||||
}
|
||||
|
||||
func (f *fakeOrphanAuthKeyRetention) DeleteOrphaned(_ context.Context, olderThan time.Duration, limit int, protected [][8]byte) (int, error) {
|
||||
f.calls++
|
||||
f.olderThan = olderThan
|
||||
f.limit = limit
|
||||
f.protected = append([][8]byte(nil), protected...)
|
||||
return 2, nil
|
||||
}
|
||||
|
||||
type fakeActiveRawAuthKeys struct{ ids [][8]byte }
|
||||
|
||||
func (f fakeActiveRawAuthKeys) ActiveRawAuthKeyIDs() [][8]byte {
|
||||
return append([][8]byte(nil), f.ids...)
|
||||
}
|
||||
|
||||
func TestRetentionWorkerProtectsActiveRawAuthKeysFromOrphanGC(t *testing.T) {
|
||||
store := &fakeOrphanAuthKeyRetention{}
|
||||
active := fakeActiveRawAuthKeys{ids: [][8]byte{{1}, {2}}}
|
||||
w := NewRetentionWorker(&fakeOutboxRetention{}, nil, zap.NewNop(), time.Hour, time.Hour, 73).
|
||||
WithOrphanAuthKeyRetention(store, active, 24*time.Hour)
|
||||
|
||||
w.runOnce(context.Background())
|
||||
|
||||
if store.calls != 1 || store.olderThan != 24*time.Hour || store.limit != 73 {
|
||||
t.Fatalf("orphan retention calls/args = %d/%v/%d, want 1/24h/73", store.calls, store.olderThan, store.limit)
|
||||
}
|
||||
if len(store.protected) != 2 || store.protected[0] != ([8]byte{1}) || store.protected[1] != ([8]byte{2}) {
|
||||
t.Fatalf("protected raw auth keys = %v, want {1},{2}", store.protected)
|
||||
}
|
||||
}
|
||||
|
||||
type fakeHeartbeatOrphanRetention struct {
|
||||
fakeOrphanAuthKeyRetention
|
||||
heartbeatCalls int
|
||||
heartbeatIDs [][8]byte
|
||||
heartbeatErr error
|
||||
}
|
||||
|
||||
func (f *fakeHeartbeatOrphanRetention) TouchActiveRawAuthKeys(_ context.Context, ids [][8]byte) error {
|
||||
f.heartbeatCalls++
|
||||
f.heartbeatIDs = append([][8]byte(nil), ids...)
|
||||
return f.heartbeatErr
|
||||
}
|
||||
|
||||
func TestRetentionWorkerHeartbeatsActiveKeysBeforeOrphanDelete(t *testing.T) {
|
||||
store := &fakeHeartbeatOrphanRetention{}
|
||||
active := fakeActiveRawAuthKeys{ids: [][8]byte{{3}, {4}}}
|
||||
w := NewRetentionWorker(&fakeOutboxRetention{}, nil, zap.NewNop(), time.Hour, 2*time.Hour, 19).
|
||||
WithOrphanAuthKeyRetention(store, active, 3*time.Hour)
|
||||
|
||||
w.runRetentionOnce(context.Background())
|
||||
|
||||
if store.heartbeatCalls != 1 || store.calls != 1 {
|
||||
t.Fatalf("heartbeat/delete calls = %d/%d, want 1/1", store.heartbeatCalls, store.calls)
|
||||
}
|
||||
if len(store.heartbeatIDs) != 2 || store.heartbeatIDs[0] != ([8]byte{3}) || store.heartbeatIDs[1] != ([8]byte{4}) {
|
||||
t.Fatalf("heartbeat ids = %v, want {3},{4}", store.heartbeatIDs)
|
||||
}
|
||||
// min(retention worker interval=2h, orphan retention/3=1h)
|
||||
if got := w.orphanHeartbeatInterval(); got != time.Hour {
|
||||
t.Fatalf("heartbeat interval = %v, want 1h", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRetentionWorkerSkipsOrphanDeleteWhenHeartbeatFails(t *testing.T) {
|
||||
core, logs := observer.New(zapcore.ErrorLevel)
|
||||
store := &fakeHeartbeatOrphanRetention{heartbeatErr: errors.New("db unavailable")}
|
||||
active := fakeActiveRawAuthKeys{ids: [][8]byte{{5}}}
|
||||
w := NewRetentionWorker(&fakeOutboxRetention{}, nil, zap.New(core), time.Hour, 30*time.Minute, 11).
|
||||
WithOrphanAuthKeyRetention(store, active, 24*time.Hour)
|
||||
if got := w.orphanHeartbeatInterval(); got != 30*time.Minute {
|
||||
t.Fatalf("heartbeat interval = %v, want worker interval 30m", got)
|
||||
}
|
||||
|
||||
w.runRetentionOnce(context.Background())
|
||||
|
||||
if store.heartbeatCalls != 1 || store.calls != 0 {
|
||||
t.Fatalf("heartbeat/delete calls = %d/%d, want 1/0", store.heartbeatCalls, store.calls)
|
||||
}
|
||||
entries := logs.FilterMessage("刷新 active raw auth key heartbeat 失败,本轮跳过 orphan GC").All()
|
||||
if len(entries) != 1 || entries[0].ContextMap()["signal"] != "auth_key_heartbeat_failed" {
|
||||
t.Fatalf("heartbeat failure signals = %+v", entries)
|
||||
}
|
||||
}
|
||||
|
|
|
|||
28
internal/app/messages/album_group.go
Normal file
28
internal/app/messages/album_group.go
Normal file
|
|
@ -0,0 +1,28 @@
|
|||
package messages
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
|
||||
"telesrv/internal/domain"
|
||||
"telesrv/internal/store"
|
||||
)
|
||||
|
||||
// ReserveAlbumGroup 把 RPC 已验证的一批 album item 交给持久层原子预留。
|
||||
// 该能力只传 domain DTO;上传媒体解析与 tg 类型仍停留在 RPC edge。
|
||||
func (s *Service) ReserveAlbumGroup(ctx context.Context, userID int64, req domain.AlbumGroupReservationRequest) (int64, error) {
|
||||
if s == nil || s.messages == nil || userID <= 0 {
|
||||
return 0, domain.ErrAlbumGroupReservationInvalid
|
||||
}
|
||||
if req.SenderUserID == 0 {
|
||||
req.SenderUserID = userID
|
||||
}
|
||||
if req.SenderUserID != userID {
|
||||
return 0, domain.ErrAlbumGroupReservationInvalid
|
||||
}
|
||||
reservations, ok := s.messages.(store.AlbumGroupStore)
|
||||
if !ok {
|
||||
return 0, errors.New("message store does not support album group reservations")
|
||||
}
|
||||
return reservations.ReserveAlbumGroup(ctx, req)
|
||||
}
|
||||
|
|
@ -2,6 +2,7 @@ package messages
|
|||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
|
||||
"telesrv/internal/app/userprojection"
|
||||
"telesrv/internal/domain"
|
||||
|
|
@ -96,6 +97,28 @@ func (s *Service) SendPrivateText(ctx context.Context, userID int64, req domain.
|
|||
if req.SenderUserID == 0 {
|
||||
req.SenderUserID = userID
|
||||
}
|
||||
if req.SenderUserID != userID {
|
||||
return domain.SendPrivateTextResult{}, domain.ErrUserSendRestricted
|
||||
}
|
||||
if req.RandomID != 0 && !req.IdempotencyPreflighted {
|
||||
fingerprint, err := store.PrivateSendFingerprint(req)
|
||||
if err != nil {
|
||||
return domain.SendPrivateTextResult{}, err
|
||||
}
|
||||
req.IdempotencyFingerprint = fingerprint
|
||||
if replayStore, ok := s.messages.(store.PrivateSendReplayStore); ok {
|
||||
replay, found, err := replayStore.LookupPrivateSendReplay(ctx, domain.PrivateSendReplayRequest{
|
||||
SenderUserID: req.SenderUserID,
|
||||
RecipientUserID: req.RecipientUserID,
|
||||
RandomID: req.RandomID,
|
||||
IdempotencyFingerprint: fingerprint,
|
||||
})
|
||||
if err != nil || found {
|
||||
return replay, err
|
||||
}
|
||||
req.IdempotencyPreflighted = true
|
||||
}
|
||||
}
|
||||
if err := s.ensureCanSend(ctx, req.SenderUserID); err != nil {
|
||||
return domain.SendPrivateTextResult{}, err
|
||||
}
|
||||
|
|
@ -113,6 +136,26 @@ func (s *Service) SendPrivateText(ctx context.Context, userID int64, req domain.
|
|||
return res, err
|
||||
}
|
||||
|
||||
// LookupPrivateSendReplay exposes the immutable receipt to the RPC boundary without executing
|
||||
// send permission checks, business automation or bot responders. Sender identity is still bound
|
||||
// to the authenticated app-service caller.
|
||||
func (s *Service) LookupPrivateSendReplay(ctx context.Context, userID int64, req domain.PrivateSendReplayRequest) (domain.SendPrivateTextResult, bool, error) {
|
||||
if s == nil || s.messages == nil || userID == 0 {
|
||||
return domain.SendPrivateTextResult{}, false, nil
|
||||
}
|
||||
if req.SenderUserID == 0 {
|
||||
req.SenderUserID = userID
|
||||
}
|
||||
if req.SenderUserID != userID || req.RecipientUserID == 0 || req.RandomID == 0 {
|
||||
return domain.SendPrivateTextResult{}, false, fmt.Errorf("private send replay: invalid authenticated scope")
|
||||
}
|
||||
replayStore, ok := s.messages.(store.PrivateSendReplayStore)
|
||||
if !ok {
|
||||
return domain.SendPrivateTextResult{}, false, nil
|
||||
}
|
||||
return replayStore.LookupPrivateSendReplay(ctx, req)
|
||||
}
|
||||
|
||||
func (s *Service) ensureCanSend(ctx context.Context, userID int64) error {
|
||||
if s == nil || s.sendGate == nil || userID == 0 {
|
||||
return nil
|
||||
|
|
|
|||
|
|
@ -29,6 +29,38 @@ func TestServiceSendPrivateTextHonorsSendPermissionGate(t *testing.T) {
|
|||
}
|
||||
}
|
||||
|
||||
func TestServicePrivateReplayPrecedesCurrentSendPermissionGate(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
messages := memory.NewMessageStore()
|
||||
allowed := NewService(messages, nil)
|
||||
req := domain.SendPrivateTextRequest{
|
||||
SenderUserID: 1001,
|
||||
RecipientUserID: 1002,
|
||||
RandomID: 91,
|
||||
Message: "committed before restriction",
|
||||
Date: 1_700_000_000,
|
||||
}
|
||||
first, err := allowed.SendPrivateText(ctx, 1001, req)
|
||||
if err != nil {
|
||||
t.Fatalf("first SendPrivateText: %v", err)
|
||||
}
|
||||
|
||||
denied := NewService(messages, nil, WithSendPermissionChecker(denySendChecker{}))
|
||||
req.Date++ // execution time is not part of the immutable send intent.
|
||||
replay, err := denied.SendPrivateText(ctx, 1001, req)
|
||||
if err != nil {
|
||||
t.Fatalf("replay through denied gate: %v", err)
|
||||
}
|
||||
if !replay.Duplicate || replay.SenderMessage.ID != first.SenderMessage.ID {
|
||||
t.Fatalf("replay = %+v, want committed duplicate %d", replay, first.SenderMessage.ID)
|
||||
}
|
||||
|
||||
req.Message = "different intent"
|
||||
if _, err := denied.SendPrivateText(ctx, 1001, req); !errors.Is(err, domain.ErrMessageRandomIDDuplicate) {
|
||||
t.Fatalf("conflicting replay err=%v, want ErrMessageRandomIDDuplicate before send gate", err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestServiceForwardPrivateMessagesHonorsSendPermissionGate(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
store := &gateMessageStore{}
|
||||
|
|
|
|||
|
|
@ -27,6 +27,10 @@ type newMessageEventFinder interface {
|
|||
FindNewMessageEvent(ctx context.Context, userID int64, messageBoxID int) (domain.UpdateEvent, bool, error)
|
||||
}
|
||||
|
||||
type userUpdateRetentionCheckpointStore interface {
|
||||
UserUpdateRetentionCheckpoint(ctx context.Context, authKeyID [8]byte, userID int64) (pts, date int, ok bool, err error)
|
||||
}
|
||||
|
||||
// ServiceOption 调整 updates 服务的运行时依赖。
|
||||
type ServiceOption func(*Service)
|
||||
|
||||
|
|
@ -146,6 +150,11 @@ func (s *Service) AcknowledgeCurrentState(ctx context.Context, authKeyID [8]byte
|
|||
if err := s.saveConfirmedState(ctx, authKeyID, userID, st); err != nil {
|
||||
return domain.UpdateState{}, err
|
||||
}
|
||||
// getState 明确建立“从当前快照开始同步”的 baseline;即使响应丢失,客户端也会
|
||||
// 重试 getState/重新拉 snapshot,而不会依赖 baseline 之前的 durable event。
|
||||
if err := s.observeClientState(ctx, authKeyID, userID, st); err != nil {
|
||||
return domain.UpdateState{}, err
|
||||
}
|
||||
return st, nil
|
||||
}
|
||||
|
||||
|
|
@ -162,6 +171,27 @@ func (s *Service) GetDifference(ctx context.Context, authKeyID [8]byte, userID i
|
|||
if err != nil {
|
||||
return domain.UpdateDifference{}, err
|
||||
}
|
||||
// 只把客户端在本次请求中实际带回的 cursor 记为 observed。绝不能把本次将要
|
||||
// 返回的 State 当确认:响应可能在 socket/进程故障中丢失。恶意/损坏客户端带来的
|
||||
// 超前 pts 钳到账号当前连续水位,避免把 retention 安全边界推过 durable truth。
|
||||
observed := from
|
||||
if observed.Pts < 0 {
|
||||
observed.Pts = 0
|
||||
}
|
||||
if observed.Pts > st.Pts {
|
||||
observed.Pts = st.Pts
|
||||
}
|
||||
if err := s.observeClientState(ctx, authKeyID, userID, observed); err != nil {
|
||||
return domain.UpdateDifference{}, err
|
||||
}
|
||||
// TDesktop 不支持账号级 updates.differenceTooLong。retention 只能删除所有授权
|
||||
// 设备都已确认的共同前缀;当前设备若仍带更旧 pts,用一个空的普通
|
||||
// differenceSlice 把 IntermediateState 推进到已确认 checkpoint,再从 live tail 续拉。
|
||||
if checkpoint, found, err := s.retainedPrefixCheckpoint(ctx, authKeyID, userID, from, st); err != nil {
|
||||
return domain.UpdateDifference{}, err
|
||||
} else if found {
|
||||
return checkpoint, nil
|
||||
}
|
||||
if s.events == nil || from.Pts >= st.Pts {
|
||||
if from.Date != 0 {
|
||||
st.Date = from.Date
|
||||
|
|
@ -176,6 +206,17 @@ func (s *Service) GetDifference(ctx context.Context, authKeyID [8]byte, userID i
|
|||
return domain.UpdateDifference{}, err
|
||||
}
|
||||
contiguous, gapEvent, expectedPts := contiguousPrefixAndGap(events, from.Pts)
|
||||
// Retention may advance after the pre-read checkpoint probe and before ListAfter obtains its
|
||||
// statement snapshot. If it removed the whole requested prefix, the read is empty or starts at a
|
||||
// gap. Re-read the checkpoint before returning a non-advancing empty difference; otherwise a
|
||||
// client can believe synchronization completed while retaining a cursor below deleted history.
|
||||
if len(contiguous) == 0 && from.Pts < st.Pts {
|
||||
if checkpoint, found, err := s.retainedPrefixCheckpoint(ctx, authKeyID, userID, from, st); err != nil {
|
||||
return domain.UpdateDifference{}, err
|
||||
} else if found {
|
||||
return checkpoint, nil
|
||||
}
|
||||
}
|
||||
last := from.Pts
|
||||
if len(contiguous) > 0 {
|
||||
last = contiguous[len(contiguous)-1].Pts
|
||||
|
|
@ -215,6 +256,32 @@ func (s *Service) GetDifference(ctx context.Context, authKeyID [8]byte, userID i
|
|||
}, nil
|
||||
}
|
||||
|
||||
func (s *Service) retainedPrefixCheckpoint(ctx context.Context, authKeyID [8]byte, userID int64, from, current domain.UpdateState) (domain.UpdateDifference, bool, error) {
|
||||
checkpoints, ok := s.events.(userUpdateRetentionCheckpointStore)
|
||||
if !ok {
|
||||
return domain.UpdateDifference{}, false, nil
|
||||
}
|
||||
pts, date, found, err := checkpoints.UserUpdateRetentionCheckpoint(ctx, authKeyID, userID)
|
||||
if err != nil {
|
||||
return domain.UpdateDifference{}, false, err
|
||||
}
|
||||
if !found || from.Pts >= pts {
|
||||
return domain.UpdateDifference{}, false, nil
|
||||
}
|
||||
checkpoint := from
|
||||
checkpoint.Pts = pts
|
||||
checkpoint.Seq = 0
|
||||
if date > 0 {
|
||||
checkpoint.Date = date
|
||||
} else if checkpoint.Date == 0 {
|
||||
checkpoint.Date = current.Date
|
||||
}
|
||||
if err := s.saveConfirmedState(ctx, authKeyID, userID, checkpoint); err != nil {
|
||||
return domain.UpdateDifference{}, false, err
|
||||
}
|
||||
return domain.UpdateDifference{State: checkpoint, Partial: true}, true, nil
|
||||
}
|
||||
|
||||
func (s *Service) currentState(ctx context.Context, userID int64) (domain.UpdateState, error) {
|
||||
current, err := s.currentPts(ctx, userID)
|
||||
if err != nil {
|
||||
|
|
@ -235,6 +302,14 @@ func (s *Service) saveConfirmedState(ctx context.Context, authKeyID [8]byte, use
|
|||
return s.states.Save(ctx, authKeyID, userID, st)
|
||||
}
|
||||
|
||||
func (s *Service) observeClientState(ctx context.Context, authKeyID [8]byte, userID int64, st domain.UpdateState) error {
|
||||
if s.states == nil {
|
||||
return nil
|
||||
}
|
||||
st.Seq = 0
|
||||
return s.states.ObserveClientState(ctx, authKeyID, userID, st)
|
||||
}
|
||||
|
||||
// contiguousPrefix 返回从 from 起 pts 严格连续(from+1, from+2, ...)的事件前缀。
|
||||
// 先按 pts 升序排序以兼容存储返回顺序,遇到空洞即停。
|
||||
func contiguousPrefix(events []domain.UpdateEvent, from int) []domain.UpdateEvent {
|
||||
|
|
@ -287,7 +362,7 @@ func (s *Service) RecordNewMessage(ctx context.Context, authKeyID [8]byte, userI
|
|||
if date == 0 {
|
||||
date = int(time.Now().Unix())
|
||||
}
|
||||
return s.recordEvent(ctx, authKeyID, userID, domain.UpdateEvent{
|
||||
return s.recordEvent(ctx, authKeyID, [8]byte{}, userID, domain.UpdateEvent{
|
||||
Type: domain.UpdateEventNewMessage,
|
||||
Date: date,
|
||||
Message: msg,
|
||||
|
|
@ -318,7 +393,7 @@ func (s *Service) PublishNewMessage(ctx context.Context, userID int64, msg domai
|
|||
if date == 0 {
|
||||
date = int(time.Now().Unix())
|
||||
}
|
||||
return s.recordEventCore(ctx, [8]byte{}, userID, domain.UpdateEvent{
|
||||
return s.recordEventCore(ctx, [8]byte{}, [8]byte{}, userID, domain.UpdateEvent{
|
||||
Type: domain.UpdateEventNewMessage,
|
||||
Date: date,
|
||||
Message: msg,
|
||||
|
|
@ -369,11 +444,11 @@ func (s *Service) RecordMessagePoll(ctx context.Context, authKeyID [8]byte, user
|
|||
}
|
||||
|
||||
// RecordStory records a story snapshot change for offline difference replay.
|
||||
func (s *Service) RecordStory(ctx context.Context, authKeyID [8]byte, userID int64, story domain.Story, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
|
||||
func (s *Service) RecordStory(ctx context.Context, stateAuthKeyID [8]byte, userID int64, story domain.Story, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
|
||||
if userID == 0 && story.Owner.Type == domain.PeerTypeUser {
|
||||
userID = story.Owner.ID
|
||||
}
|
||||
return s.recordEvent(ctx, authKeyID, userID, domain.UpdateEvent{
|
||||
return s.recordEvent(ctx, stateAuthKeyID, excludeAuthKeyID, userID, domain.UpdateEvent{
|
||||
Type: domain.UpdateEventStory,
|
||||
Date: story.Date,
|
||||
Peer: story.Owner,
|
||||
|
|
@ -389,7 +464,7 @@ func (s *Service) RecordStoryFanout(ctx context.Context, userID int64, story dom
|
|||
if userID == 0 {
|
||||
return domain.UpdateEvent{}, domain.UpdateState{}, domain.ErrStoryPeerInvalid
|
||||
}
|
||||
return s.recordEventCore(ctx, [8]byte{}, userID, domain.UpdateEvent{
|
||||
return s.recordEventCore(ctx, [8]byte{}, [8]byte{}, userID, domain.UpdateEvent{
|
||||
Type: domain.UpdateEventStory,
|
||||
Date: story.Date,
|
||||
Peer: story.Owner,
|
||||
|
|
@ -399,11 +474,11 @@ func (s *Service) RecordStoryFanout(ctx context.Context, userID int64, story dom
|
|||
}
|
||||
|
||||
// RecordReadStories records a read boundary update for multi-device sync.
|
||||
func (s *Service) RecordReadStories(ctx context.Context, authKeyID [8]byte, userID int64, read domain.StoryReadResult, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
|
||||
func (s *Service) RecordReadStories(ctx context.Context, stateAuthKeyID [8]byte, userID int64, read domain.StoryReadResult, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
|
||||
if userID == 0 {
|
||||
userID = read.ViewerID
|
||||
}
|
||||
return s.recordEvent(ctx, authKeyID, userID, domain.UpdateEvent{
|
||||
return s.recordEvent(ctx, stateAuthKeyID, excludeAuthKeyID, userID, domain.UpdateEvent{
|
||||
Type: domain.UpdateEventReadStories,
|
||||
Date: read.Date,
|
||||
Peer: read.Peer,
|
||||
|
|
@ -413,11 +488,11 @@ func (s *Service) RecordReadStories(ctx context.Context, authKeyID [8]byte, user
|
|||
}
|
||||
|
||||
// RecordSentStoryReaction records the current user's story reaction for multi-device sync.
|
||||
func (s *Service) RecordSentStoryReaction(ctx context.Context, authKeyID [8]byte, userID int64, reaction domain.StoryReactionResult, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
|
||||
func (s *Service) RecordSentStoryReaction(ctx context.Context, stateAuthKeyID [8]byte, userID int64, reaction domain.StoryReactionResult, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
|
||||
if userID == 0 {
|
||||
userID = reaction.ViewerID
|
||||
}
|
||||
return s.recordEvent(ctx, authKeyID, userID, domain.UpdateEvent{
|
||||
return s.recordEvent(ctx, stateAuthKeyID, excludeAuthKeyID, userID, domain.UpdateEvent{
|
||||
Type: domain.UpdateEventSentStoryReaction,
|
||||
Date: reaction.Date,
|
||||
Peer: reaction.Peer,
|
||||
|
|
@ -432,7 +507,7 @@ func (s *Service) RecordSentStoryReaction(ctx context.Context, authKeyID [8]byte
|
|||
// sent by another user. It does not advance any owner device confirmation state:
|
||||
// the owner did not initiate the RPC, but online outbox and offline difference
|
||||
// must still see the durable event.
|
||||
func (s *Service) RecordNewStoryReaction(ctx context.Context, authKeyID [8]byte, ownerUserID int64, reaction domain.StoryReactionResult, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
|
||||
func (s *Service) RecordNewStoryReaction(ctx context.Context, stateAuthKeyID [8]byte, ownerUserID int64, reaction domain.StoryReactionResult, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
|
||||
if ownerUserID == 0 && reaction.Story.Owner.Type == domain.PeerTypeUser {
|
||||
ownerUserID = reaction.Story.Owner.ID
|
||||
}
|
||||
|
|
@ -442,7 +517,7 @@ func (s *Service) RecordNewStoryReaction(ctx context.Context, authKeyID [8]byte,
|
|||
if ownerUserID == 0 || reaction.ViewerID == 0 || reaction.Reaction == nil {
|
||||
return domain.UpdateEvent{}, domain.UpdateState{}, domain.ErrStoryPeerInvalid
|
||||
}
|
||||
return s.recordEventCore(ctx, authKeyID, ownerUserID, domain.UpdateEvent{
|
||||
return s.recordEventCore(ctx, stateAuthKeyID, excludeAuthKeyID, ownerUserID, domain.UpdateEvent{
|
||||
Type: domain.UpdateEventNewStoryReaction,
|
||||
Date: reaction.Date,
|
||||
Peer: domain.Peer{Type: domain.PeerTypeUser, ID: reaction.ViewerID},
|
||||
|
|
@ -456,7 +531,7 @@ func (s *Service) RecordNewStoryReaction(ctx context.Context, authKeyID [8]byte,
|
|||
// RecordQuickReplyMutation records account-local quick reply state changes for
|
||||
// multi-device sync. Quick-reply TL updates do not carry pts, so outbox appends
|
||||
// auxiliary pts bookkeeping just like other account settings events.
|
||||
func (s *Service) RecordQuickReplyMutation(ctx context.Context, authKeyID [8]byte, userID int64, mutation domain.QuickReplyMutation, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
|
||||
func (s *Service) RecordQuickReplyMutation(ctx context.Context, stateAuthKeyID [8]byte, userID int64, mutation domain.QuickReplyMutation, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
|
||||
if userID == 0 {
|
||||
userID = mutation.List.OwnerUserID
|
||||
}
|
||||
|
|
@ -481,16 +556,16 @@ func (s *Service) RecordQuickReplyMutation(ctx context.Context, authKeyID [8]byt
|
|||
default:
|
||||
event.Type = domain.UpdateEventQuickReplies
|
||||
}
|
||||
return s.recordEvent(ctx, authKeyID, userID, event, true, excludeSessionID)
|
||||
return s.recordEvent(ctx, stateAuthKeyID, excludeAuthKeyID, userID, event, true, excludeSessionID)
|
||||
}
|
||||
|
||||
// RecordReadHistory 推进 update 状态并追加一条 read_history_inbox 事件。
|
||||
func (s *Service) RecordReadHistory(ctx context.Context, authKeyID [8]byte, userID int64, read domain.ReadHistoryResult, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
|
||||
func (s *Service) RecordReadHistory(ctx context.Context, stateAuthKeyID [8]byte, userID int64, read domain.ReadHistoryResult, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
|
||||
if userID == 0 {
|
||||
userID = read.OwnerUserID
|
||||
}
|
||||
date := int(time.Now().Unix())
|
||||
return s.recordEvent(ctx, authKeyID, userID, domain.UpdateEvent{
|
||||
return s.recordEvent(ctx, stateAuthKeyID, excludeAuthKeyID, userID, domain.UpdateEvent{
|
||||
Type: domain.UpdateEventReadHistoryInbox,
|
||||
Date: date,
|
||||
Peer: read.Peer,
|
||||
|
|
@ -503,8 +578,8 @@ func (s *Service) RecordReadHistory(ctx context.Context, authKeyID [8]byte, user
|
|||
|
||||
// RecordChannelState 记录当前账号与某频道成员关系变化(leave/kick),
|
||||
// 离线设备经 difference 收到 updateChannel 后重拉 channel 状态。
|
||||
func (s *Service) RecordChannelState(ctx context.Context, authKeyID [8]byte, userID, channelID int64, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
|
||||
return s.recordEvent(ctx, authKeyID, userID, domain.UpdateEvent{
|
||||
func (s *Service) RecordChannelState(ctx context.Context, stateAuthKeyID [8]byte, userID, channelID int64, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
|
||||
return s.recordEvent(ctx, stateAuthKeyID, excludeAuthKeyID, userID, domain.UpdateEvent{
|
||||
Type: domain.UpdateEventChannelState,
|
||||
Peer: domain.Peer{Type: domain.PeerTypeChannel, ID: channelID},
|
||||
PtsCount: 1,
|
||||
|
|
@ -512,8 +587,8 @@ func (s *Service) RecordChannelState(ctx context.Context, authKeyID [8]byte, use
|
|||
}
|
||||
|
||||
// RecordContactsReset 记录通讯录视角变化,供离线设备通过 updates.getDifference 触发重拉。
|
||||
func (s *Service) RecordContactsReset(ctx context.Context, authKeyID [8]byte, userID int64, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
|
||||
return s.recordEvent(ctx, authKeyID, userID, domain.UpdateEvent{
|
||||
func (s *Service) RecordContactsReset(ctx context.Context, stateAuthKeyID [8]byte, userID int64, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
|
||||
return s.recordEvent(ctx, stateAuthKeyID, excludeAuthKeyID, userID, domain.UpdateEvent{
|
||||
Type: domain.UpdateEventContactsReset,
|
||||
PtsCount: 1,
|
||||
}, true, excludeSessionID)
|
||||
|
|
@ -522,8 +597,8 @@ func (s *Service) RecordContactsReset(ctx context.Context, authKeyID [8]byte, us
|
|||
// RecordDraftMessage 记录某会话云草稿变化(保存/清空都是同一事件——草稿是绝对
|
||||
// 状态,重放时按 peer 重载当前值)。updateDraftMessage 无 pts 字段,走 LacksWirePts
|
||||
// aux 簿记;topMsgID 是 forum 话题草稿键(复用 MaxID 列持久化)。
|
||||
func (s *Service) RecordDraftMessage(ctx context.Context, authKeyID [8]byte, userID int64, peer domain.Peer, topMsgID int, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
|
||||
return s.recordEvent(ctx, authKeyID, userID, domain.UpdateEvent{
|
||||
func (s *Service) RecordDraftMessage(ctx context.Context, stateAuthKeyID [8]byte, userID int64, peer domain.Peer, topMsgID int, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
|
||||
return s.recordEvent(ctx, stateAuthKeyID, excludeAuthKeyID, userID, domain.UpdateEvent{
|
||||
Type: domain.UpdateEventDraftMessage,
|
||||
Peer: peer,
|
||||
MaxID: topMsgID,
|
||||
|
|
@ -533,8 +608,8 @@ func (s *Service) RecordDraftMessage(ctx context.Context, authKeyID [8]byte, use
|
|||
|
||||
// RecordDialogPinned 记录单个会话置顶状态变化;folderID 是会话所在 folder
|
||||
// (0 主列表/1 归档),缺失会让离线设备把归档内置顶重放到主列表。
|
||||
func (s *Service) RecordDialogPinned(ctx context.Context, authKeyID [8]byte, userID int64, peer domain.Peer, pinned bool, folderID int, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
|
||||
return s.recordEvent(ctx, authKeyID, userID, domain.UpdateEvent{
|
||||
func (s *Service) RecordDialogPinned(ctx context.Context, stateAuthKeyID [8]byte, userID int64, peer domain.Peer, pinned bool, folderID int, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
|
||||
return s.recordEvent(ctx, stateAuthKeyID, excludeAuthKeyID, userID, domain.UpdateEvent{
|
||||
Type: domain.UpdateEventDialogPinned,
|
||||
Peer: peer,
|
||||
Bool: pinned,
|
||||
|
|
@ -544,8 +619,8 @@ func (s *Service) RecordDialogPinned(ctx context.Context, authKeyID [8]byte, use
|
|||
}
|
||||
|
||||
// RecordPinnedDialogs 记录指定 folder 内置顶顺序变化,并把新顺序持久化给 getDifference/outbox。
|
||||
func (s *Service) RecordPinnedDialogs(ctx context.Context, authKeyID [8]byte, userID int64, folderID int, order []domain.Peer, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
|
||||
return s.recordEvent(ctx, authKeyID, userID, domain.UpdateEvent{
|
||||
func (s *Service) RecordPinnedDialogs(ctx context.Context, stateAuthKeyID [8]byte, userID int64, folderID int, order []domain.Peer, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
|
||||
return s.recordEvent(ctx, stateAuthKeyID, excludeAuthKeyID, userID, domain.UpdateEvent{
|
||||
Type: domain.UpdateEventPinnedDialogs,
|
||||
Peers: append([]domain.Peer(nil), order...),
|
||||
FolderID: folderID,
|
||||
|
|
@ -554,8 +629,8 @@ func (s *Service) RecordPinnedDialogs(ctx context.Context, authKeyID [8]byte, us
|
|||
}
|
||||
|
||||
// RecordSavedDialogPinned 记录收藏夹单个子会话置顶状态变化。
|
||||
func (s *Service) RecordSavedDialogPinned(ctx context.Context, authKeyID [8]byte, userID int64, peer domain.Peer, pinned bool, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
|
||||
return s.recordEvent(ctx, authKeyID, userID, domain.UpdateEvent{
|
||||
func (s *Service) RecordSavedDialogPinned(ctx context.Context, stateAuthKeyID [8]byte, userID int64, peer domain.Peer, pinned bool, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
|
||||
return s.recordEvent(ctx, stateAuthKeyID, excludeAuthKeyID, userID, domain.UpdateEvent{
|
||||
Type: domain.UpdateEventSavedDialogPinned,
|
||||
Peer: peer,
|
||||
Bool: pinned,
|
||||
|
|
@ -564,8 +639,8 @@ func (s *Service) RecordSavedDialogPinned(ctx context.Context, authKeyID [8]byte
|
|||
}
|
||||
|
||||
// RecordPinnedSavedDialogs 记录收藏夹置顶顺序变化,新顺序持久化给 getDifference/outbox。
|
||||
func (s *Service) RecordPinnedSavedDialogs(ctx context.Context, authKeyID [8]byte, userID int64, order []domain.Peer, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
|
||||
return s.recordEvent(ctx, authKeyID, userID, domain.UpdateEvent{
|
||||
func (s *Service) RecordPinnedSavedDialogs(ctx context.Context, stateAuthKeyID [8]byte, userID int64, order []domain.Peer, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
|
||||
return s.recordEvent(ctx, stateAuthKeyID, excludeAuthKeyID, userID, domain.UpdateEvent{
|
||||
Type: domain.UpdateEventPinnedSavedDialogs,
|
||||
Peers: append([]domain.Peer(nil), order...),
|
||||
PtsCount: 1,
|
||||
|
|
@ -573,8 +648,8 @@ func (s *Service) RecordPinnedSavedDialogs(ctx context.Context, authKeyID [8]byt
|
|||
}
|
||||
|
||||
// RecordDialogUnreadMark 记录手动未读标记变化。
|
||||
func (s *Service) RecordDialogUnreadMark(ctx context.Context, authKeyID [8]byte, userID int64, peer domain.Peer, unread bool, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
|
||||
return s.recordEvent(ctx, authKeyID, userID, domain.UpdateEvent{
|
||||
func (s *Service) RecordDialogUnreadMark(ctx context.Context, stateAuthKeyID [8]byte, userID int64, peer domain.Peer, unread bool, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
|
||||
return s.recordEvent(ctx, stateAuthKeyID, excludeAuthKeyID, userID, domain.UpdateEvent{
|
||||
Type: domain.UpdateEventDialogUnreadMark,
|
||||
Peer: peer,
|
||||
Bool: unread,
|
||||
|
|
@ -583,8 +658,8 @@ func (s *Service) RecordDialogUnreadMark(ctx context.Context, authKeyID [8]byte,
|
|||
}
|
||||
|
||||
// RecordChannelViewForumAsMessages records a per-account forum presentation state change.
|
||||
func (s *Service) RecordChannelViewForumAsMessages(ctx context.Context, authKeyID [8]byte, userID, channelID int64, enabled bool, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
|
||||
return s.recordEvent(ctx, authKeyID, userID, domain.UpdateEvent{
|
||||
func (s *Service) RecordChannelViewForumAsMessages(ctx context.Context, stateAuthKeyID [8]byte, userID, channelID int64, enabled bool, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
|
||||
return s.recordEvent(ctx, stateAuthKeyID, excludeAuthKeyID, userID, domain.UpdateEvent{
|
||||
Type: domain.UpdateEventChannelViewForum,
|
||||
Peer: domain.Peer{Type: domain.PeerTypeChannel, ID: channelID},
|
||||
Bool: enabled,
|
||||
|
|
@ -594,8 +669,8 @@ func (s *Service) RecordChannelViewForumAsMessages(ctx context.Context, authKeyI
|
|||
|
||||
// RecordChannelDiscussionInbox 记录 forum 话题级已读(updateReadChannelDiscussionInbox),
|
||||
// 占一个账号 pts(LacksWirePts),供自己其它设备在线同步与离线差分恢复。
|
||||
func (s *Service) RecordChannelDiscussionInbox(ctx context.Context, authKeyID [8]byte, userID, channelID int64, topicID, maxID int, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
|
||||
return s.recordEvent(ctx, authKeyID, userID, domain.UpdateEvent{
|
||||
func (s *Service) RecordChannelDiscussionInbox(ctx context.Context, stateAuthKeyID [8]byte, userID, channelID int64, topicID, maxID int, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
|
||||
return s.recordEvent(ctx, stateAuthKeyID, excludeAuthKeyID, userID, domain.UpdateEvent{
|
||||
Type: domain.UpdateEventReadChannelDiscussionInbox,
|
||||
Peer: domain.Peer{Type: domain.PeerTypeChannel, ID: channelID},
|
||||
TopMsgID: topicID,
|
||||
|
|
@ -605,8 +680,8 @@ func (s *Service) RecordChannelDiscussionInbox(ctx context.Context, authKeyID [8
|
|||
}
|
||||
|
||||
// RecordPeerSettings 记录 peer settings 变化。
|
||||
func (s *Service) RecordPeerSettings(ctx context.Context, authKeyID [8]byte, userID int64, peer domain.Peer, settings domain.PeerSettings, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
|
||||
return s.recordEvent(ctx, authKeyID, userID, domain.UpdateEvent{
|
||||
func (s *Service) RecordPeerSettings(ctx context.Context, stateAuthKeyID [8]byte, userID int64, peer domain.Peer, settings domain.PeerSettings, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
|
||||
return s.recordEvent(ctx, stateAuthKeyID, excludeAuthKeyID, userID, domain.UpdateEvent{
|
||||
Type: domain.UpdateEventPeerSettings,
|
||||
Peer: peer,
|
||||
Settings: settings,
|
||||
|
|
@ -615,8 +690,8 @@ func (s *Service) RecordPeerSettings(ctx context.Context, authKeyID [8]byte, use
|
|||
}
|
||||
|
||||
// RecordPeerStoryBlocked 记录当前账号 story blocklist 对某个 peer 的可见状态变化。
|
||||
func (s *Service) RecordPeerStoryBlocked(ctx context.Context, authKeyID [8]byte, userID int64, peer domain.Peer, blocked bool, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
|
||||
return s.recordEvent(ctx, authKeyID, userID, domain.UpdateEvent{
|
||||
func (s *Service) RecordPeerStoryBlocked(ctx context.Context, stateAuthKeyID [8]byte, userID int64, peer domain.Peer, blocked bool, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
|
||||
return s.recordEvent(ctx, stateAuthKeyID, excludeAuthKeyID, userID, domain.UpdateEvent{
|
||||
Type: domain.UpdateEventPeerStoryBlocked,
|
||||
Peer: peer,
|
||||
Bool: blocked,
|
||||
|
|
@ -625,13 +700,13 @@ func (s *Service) RecordPeerStoryBlocked(ctx context.Context, authKeyID [8]byte,
|
|||
}
|
||||
|
||||
// RecordDialogFilter 记录单个 filter 的创建、更新或删除;folder 为 nil 表示删除。
|
||||
func (s *Service) RecordDialogFilter(ctx context.Context, authKeyID [8]byte, userID int64, folderID int, folder *domain.DialogFolder, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
|
||||
func (s *Service) RecordDialogFilter(ctx context.Context, stateAuthKeyID [8]byte, userID int64, folderID int, folder *domain.DialogFolder, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
|
||||
var copyFolder *domain.DialogFolder
|
||||
if folder != nil {
|
||||
f := *folder
|
||||
copyFolder = &f
|
||||
}
|
||||
return s.recordEvent(ctx, authKeyID, userID, domain.UpdateEvent{
|
||||
return s.recordEvent(ctx, stateAuthKeyID, excludeAuthKeyID, userID, domain.UpdateEvent{
|
||||
Type: domain.UpdateEventDialogFilter,
|
||||
FilterID: folderID,
|
||||
DialogFilter: copyFolder,
|
||||
|
|
@ -640,8 +715,8 @@ func (s *Service) RecordDialogFilter(ctx context.Context, authKeyID [8]byte, use
|
|||
}
|
||||
|
||||
// RecordDialogFilterOrder 记录 filter 顺序变化。
|
||||
func (s *Service) RecordDialogFilterOrder(ctx context.Context, authKeyID [8]byte, userID int64, order []int, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
|
||||
return s.recordEvent(ctx, authKeyID, userID, domain.UpdateEvent{
|
||||
func (s *Service) RecordDialogFilterOrder(ctx context.Context, stateAuthKeyID [8]byte, userID int64, order []int, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
|
||||
return s.recordEvent(ctx, stateAuthKeyID, excludeAuthKeyID, userID, domain.UpdateEvent{
|
||||
Type: domain.UpdateEventDialogFilterOrder,
|
||||
FilterOrder: append([]int(nil), order...),
|
||||
PtsCount: 1,
|
||||
|
|
@ -649,16 +724,16 @@ func (s *Service) RecordDialogFilterOrder(ctx context.Context, authKeyID [8]byte
|
|||
}
|
||||
|
||||
// RecordDialogFiltersReload 通知其他设备重新拉取 filter 列表。
|
||||
func (s *Service) RecordDialogFiltersReload(ctx context.Context, authKeyID [8]byte, userID int64, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
|
||||
return s.recordEvent(ctx, authKeyID, userID, domain.UpdateEvent{
|
||||
func (s *Service) RecordDialogFiltersReload(ctx context.Context, stateAuthKeyID [8]byte, userID int64, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
|
||||
return s.recordEvent(ctx, stateAuthKeyID, excludeAuthKeyID, userID, domain.UpdateEvent{
|
||||
Type: domain.UpdateEventDialogFilters,
|
||||
PtsCount: 1,
|
||||
}, true, excludeSessionID)
|
||||
}
|
||||
|
||||
// RecordFolderPeers 记录归档/还原会话的 folder_id 变化。
|
||||
func (s *Service) RecordFolderPeers(ctx context.Context, authKeyID [8]byte, userID int64, peers []domain.FolderPeerUpdate, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
|
||||
return s.recordEvent(ctx, authKeyID, userID, domain.UpdateEvent{
|
||||
func (s *Service) RecordFolderPeers(ctx context.Context, stateAuthKeyID [8]byte, userID int64, peers []domain.FolderPeerUpdate, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
|
||||
return s.recordEvent(ctx, stateAuthKeyID, excludeAuthKeyID, userID, domain.UpdateEvent{
|
||||
Type: domain.UpdateEventFolderPeers,
|
||||
FolderPeers: append([]domain.FolderPeerUpdate(nil), peers...),
|
||||
PtsCount: 1,
|
||||
|
|
@ -666,8 +741,8 @@ func (s *Service) RecordFolderPeers(ctx context.Context, authKeyID [8]byte, user
|
|||
}
|
||||
|
||||
// RecordChannelAvailableMessages records a local channel history clear for multi-device sync.
|
||||
func (s *Service) RecordChannelAvailableMessages(ctx context.Context, authKeyID [8]byte, userID, channelID int64, availableMinID int, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
|
||||
return s.recordEvent(ctx, authKeyID, userID, domain.UpdateEvent{
|
||||
func (s *Service) RecordChannelAvailableMessages(ctx context.Context, stateAuthKeyID [8]byte, userID, channelID int64, availableMinID int, excludeAuthKeyID [8]byte, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
|
||||
return s.recordEvent(ctx, stateAuthKeyID, excludeAuthKeyID, userID, domain.UpdateEvent{
|
||||
Type: domain.UpdateEventChannelAvailable,
|
||||
Peer: domain.Peer{Type: domain.PeerTypeChannel, ID: channelID},
|
||||
MaxID: availableMinID,
|
||||
|
|
@ -675,15 +750,15 @@ func (s *Service) RecordChannelAvailableMessages(ctx context.Context, authKeyID
|
|||
}, true, excludeSessionID)
|
||||
}
|
||||
|
||||
func (s *Service) recordEvent(ctx context.Context, authKeyID [8]byte, userID int64, event domain.UpdateEvent, dispatch bool, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
|
||||
return s.recordEventCore(ctx, authKeyID, userID, event, dispatch, excludeSessionID, true)
|
||||
func (s *Service) recordEvent(ctx context.Context, stateAuthKeyID, excludeAuthKeyID [8]byte, userID int64, event domain.UpdateEvent, dispatch bool, excludeSessionID int64) (domain.UpdateEvent, domain.UpdateState, error) {
|
||||
return s.recordEventCore(ctx, stateAuthKeyID, excludeAuthKeyID, userID, event, dispatch, excludeSessionID, true)
|
||||
}
|
||||
|
||||
func (s *Service) recordEventWithoutState(ctx context.Context, userID int64, event domain.UpdateEvent) (domain.UpdateEvent, domain.UpdateState, error) {
|
||||
return s.recordEventCore(ctx, [8]byte{}, userID, event, false, 0, false)
|
||||
return s.recordEventCore(ctx, [8]byte{}, [8]byte{}, userID, event, false, 0, false)
|
||||
}
|
||||
|
||||
func (s *Service) recordEventCore(ctx context.Context, authKeyID [8]byte, userID int64, event domain.UpdateEvent, dispatch bool, excludeSessionID int64, saveState bool) (domain.UpdateEvent, domain.UpdateState, error) {
|
||||
func (s *Service) recordEventCore(ctx context.Context, stateAuthKeyID, excludeAuthKeyID [8]byte, userID int64, event domain.UpdateEvent, dispatch bool, excludeSessionID int64, saveState bool) (domain.UpdateEvent, domain.UpdateState, error) {
|
||||
date := event.Date
|
||||
if date == 0 {
|
||||
date = int(time.Now().Unix())
|
||||
|
|
@ -698,7 +773,7 @@ func (s *Service) recordEventCore(ctx context.Context, authKeyID [8]byte, userID
|
|||
var err error
|
||||
if dispatch {
|
||||
if appender, ok := s.events.(dispatchingEventAppender); ok {
|
||||
event, err = appender.AppendAllocatedWithDispatch(ctx, userID, event, authKeyID, excludeSessionID)
|
||||
event, err = appender.AppendAllocatedWithDispatch(ctx, userID, event, excludeAuthKeyID, excludeSessionID)
|
||||
} else {
|
||||
event, err = s.events.AppendAllocated(ctx, userID, event)
|
||||
}
|
||||
|
|
@ -735,7 +810,7 @@ func (s *Service) recordEventCore(ctx context.Context, authKeyID [8]byte, userID
|
|||
st.Pts = event.Pts
|
||||
}
|
||||
if saveState && s.states != nil {
|
||||
if err := s.states.Save(ctx, authKeyID, userID, st); err != nil {
|
||||
if err := s.states.Save(ctx, stateAuthKeyID, userID, st); err != nil {
|
||||
return domain.UpdateEvent{}, domain.UpdateState{}, err
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -100,7 +100,7 @@ func TestRecordReadHistoryFeedsGetDifference(t *testing.T) {
|
|||
Peer: peer,
|
||||
MaxID: 10,
|
||||
Changed: true,
|
||||
}, 0)
|
||||
}, [8]byte{}, 0)
|
||||
if err != nil {
|
||||
t.Fatalf("RecordReadHistory: %v", err)
|
||||
}
|
||||
|
|
@ -132,7 +132,7 @@ func TestRecordChannelReadHistoryKeepsChannelPtsPayload(t *testing.T) {
|
|||
StillUnreadCount: 3,
|
||||
ChannelPts: 77,
|
||||
Changed: true,
|
||||
}, 0)
|
||||
}, [8]byte{}, 0)
|
||||
if err != nil {
|
||||
t.Fatalf("RecordReadHistory: %v", err)
|
||||
}
|
||||
|
|
@ -157,24 +157,24 @@ func TestRecordSettingsEventsFeedGetDifference(t *testing.T) {
|
|||
ownerUserID := int64(1000000001)
|
||||
peer := domain.Peer{Type: domain.PeerTypeUser, ID: 1000000002}
|
||||
|
||||
if _, _, err := svc.RecordContactsReset(ctx, authKeyID, ownerUserID, 0); err != nil {
|
||||
if _, _, err := svc.RecordContactsReset(ctx, authKeyID, ownerUserID, [8]byte{}, 0); err != nil {
|
||||
t.Fatalf("RecordContactsReset: %v", err)
|
||||
}
|
||||
if _, _, err := svc.RecordDialogPinned(ctx, authKeyID, ownerUserID, peer, true, 0, 0); err != nil {
|
||||
if _, _, err := svc.RecordDialogPinned(ctx, authKeyID, ownerUserID, peer, true, 0, [8]byte{}, 0); err != nil {
|
||||
t.Fatalf("RecordDialogPinned: %v", err)
|
||||
}
|
||||
order := []domain.Peer{peer}
|
||||
if _, _, err := svc.RecordPinnedDialogs(ctx, authKeyID, ownerUserID, 0, order, 0); err != nil {
|
||||
if _, _, err := svc.RecordPinnedDialogs(ctx, authKeyID, ownerUserID, 0, order, [8]byte{}, 0); err != nil {
|
||||
t.Fatalf("RecordPinnedDialogs: %v", err)
|
||||
}
|
||||
if _, _, err := svc.RecordDialogUnreadMark(ctx, authKeyID, ownerUserID, peer, false, 0); err != nil {
|
||||
if _, _, err := svc.RecordDialogUnreadMark(ctx, authKeyID, ownerUserID, peer, false, [8]byte{}, 0); err != nil {
|
||||
t.Fatalf("RecordDialogUnreadMark: %v", err)
|
||||
}
|
||||
settings := domain.PeerSettings{ShareContact: true}
|
||||
if _, _, err := svc.RecordPeerSettings(ctx, authKeyID, ownerUserID, peer, settings, 0); err != nil {
|
||||
if _, _, err := svc.RecordPeerSettings(ctx, authKeyID, ownerUserID, peer, settings, [8]byte{}, 0); err != nil {
|
||||
t.Fatalf("RecordPeerSettings: %v", err)
|
||||
}
|
||||
stateEvent, state, err := svc.RecordPeerStoryBlocked(ctx, authKeyID, ownerUserID, peer, true, 0)
|
||||
stateEvent, state, err := svc.RecordPeerStoryBlocked(ctx, authKeyID, ownerUserID, peer, true, [8]byte{}, 0)
|
||||
if err != nil {
|
||||
t.Fatalf("RecordPeerStoryBlocked: %v", err)
|
||||
}
|
||||
|
|
@ -221,22 +221,29 @@ func TestRecordSettingsEventsFeedGetDifference(t *testing.T) {
|
|||
|
||||
func TestRecordSettingsEventUsesDispatchAppender(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
var authKeyID [8]byte
|
||||
authKeyID[0] = 4
|
||||
authKeyID := [8]byte{4}
|
||||
rawAuthKeyID := [8]byte{4, 9}
|
||||
events := &captureDispatchAppender{UpdateEventStore: memory.NewUpdateEventStore()}
|
||||
svc := NewService(memory.NewUpdateStateStore(), events)
|
||||
states := &captureStateStore{}
|
||||
svc := NewService(states, events)
|
||||
peer := domain.Peer{Type: domain.PeerTypeUser, ID: 1000000002}
|
||||
|
||||
event, state, err := svc.RecordDialogPinned(ctx, authKeyID, 1000000001, peer, true, 0, 42)
|
||||
event, state, err := svc.RecordDialogPinned(ctx, authKeyID, 1000000001, peer, true, 0, rawAuthKeyID, 42)
|
||||
if err != nil {
|
||||
t.Fatalf("RecordDialogPinned: %v", err)
|
||||
}
|
||||
if event.Pts != 1 || state.Pts != 1 {
|
||||
t.Fatalf("event/state = %+v / %+v, want first pts", event, state)
|
||||
}
|
||||
if !events.dispatched || events.excludeAuthKeyID != authKeyID || events.excludeSessionID != 42 || events.event.Type != domain.UpdateEventDialogPinned || events.event.Peer != peer {
|
||||
if !events.dispatched || events.excludeAuthKeyID != rawAuthKeyID || events.excludeSessionID != 42 || events.event.Type != domain.UpdateEventDialogPinned || events.event.Peer != peer {
|
||||
t.Fatalf("dispatch capture = %+v exclude_auth=%v exclude_session=%d dispatched=%v, want dialog_pinned outbox", events.event, events.excludeAuthKeyID, events.excludeSessionID, events.dispatched)
|
||||
}
|
||||
if states.lastSaveAuthKeyID != authKeyID {
|
||||
t.Fatalf("device state auth key = %x, want business/perm %x", states.lastSaveAuthKeyID, authKeyID)
|
||||
}
|
||||
if _, found, err := states.Get(ctx, rawAuthKeyID, 1000000001); err != nil || found {
|
||||
t.Fatalf("raw temp key unexpectedly owns device state: found=%v err=%v", found, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRecordSettingsEventDispatchFailureDoesNotRecordEvent(t *testing.T) {
|
||||
|
|
@ -246,7 +253,7 @@ func TestRecordSettingsEventDispatchFailureDoesNotRecordEvent(t *testing.T) {
|
|||
events := &failingDispatchAppender{UpdateEventStore: memory.NewUpdateEventStore()}
|
||||
svc := NewService(memory.NewUpdateStateStore(), events)
|
||||
|
||||
_, _, err := svc.RecordDialogPinned(ctx, authKeyID, 1000000001, domain.Peer{Type: domain.PeerTypeUser, ID: 1000000002}, true, 0, 42)
|
||||
_, _, err := svc.RecordDialogPinned(ctx, authKeyID, 1000000001, domain.Peer{Type: domain.PeerTypeUser, ID: 1000000002}, true, 0, authKeyID, 42)
|
||||
if !errors.Is(err, errDispatchFailed) {
|
||||
t.Fatalf("RecordDialogPinned err = %v, want dispatch failure", err)
|
||||
}
|
||||
|
|
@ -267,7 +274,7 @@ func TestRecordPeerStoryBlockedUsesDispatchAppender(t *testing.T) {
|
|||
svc := NewService(memory.NewUpdateStateStore(), events)
|
||||
peer := domain.Peer{Type: domain.PeerTypeUser, ID: 1000000002}
|
||||
|
||||
event, state, err := svc.RecordPeerStoryBlocked(ctx, authKeyID, 1000000001, peer, true, 91)
|
||||
event, state, err := svc.RecordPeerStoryBlocked(ctx, authKeyID, 1000000001, peer, true, authKeyID, 91)
|
||||
if err != nil {
|
||||
t.Fatalf("RecordPeerStoryBlocked: %v", err)
|
||||
}
|
||||
|
|
@ -294,7 +301,7 @@ func TestRecordStoryUsesDispatchAppenderExcludeCurrentSession(t *testing.T) {
|
|||
Caption: "owner story",
|
||||
}
|
||||
|
||||
event, state, err := svc.RecordStory(ctx, authKeyID, owner.ID, story, 1234)
|
||||
event, state, err := svc.RecordStory(ctx, authKeyID, owner.ID, story, authKeyID, 1234)
|
||||
if err != nil {
|
||||
t.Fatalf("RecordStory: %v", err)
|
||||
}
|
||||
|
|
@ -330,7 +337,7 @@ func TestRecordStoryReadAndSentReactionExcludeCurrentSession(t *testing.T) {
|
|||
MaxReadID: story.ID,
|
||||
Advanced: true,
|
||||
Date: 1700000201,
|
||||
}, 2233)
|
||||
}, authKeyID, 2233)
|
||||
if err != nil {
|
||||
t.Fatalf("RecordReadStories: %v", err)
|
||||
}
|
||||
|
|
@ -350,7 +357,7 @@ func TestRecordStoryReadAndSentReactionExcludeCurrentSession(t *testing.T) {
|
|||
Reaction: reaction,
|
||||
Changed: true,
|
||||
Date: 1700000202,
|
||||
}, 2233)
|
||||
}, authKeyID, 2233)
|
||||
if err != nil {
|
||||
t.Fatalf("RecordSentStoryReaction: %v", err)
|
||||
}
|
||||
|
|
@ -384,7 +391,7 @@ func TestRecordNewStoryReactionDispatchesWithoutSavingDeviceState(t *testing.T)
|
|||
},
|
||||
Reaction: reaction,
|
||||
Date: 1700000101,
|
||||
}, 0)
|
||||
}, [8]byte{}, 0)
|
||||
if err != nil {
|
||||
t.Fatalf("RecordNewStoryReaction: %v", err)
|
||||
}
|
||||
|
|
@ -493,7 +500,8 @@ func TestAcknowledgeCurrentStateAdvancesConfirmedWatermark(t *testing.T) {
|
|||
authKeyID[0] = 11
|
||||
userID := int64(1000000001)
|
||||
events := memory.NewUpdateEventStore()
|
||||
svc := NewService(memory.NewUpdateStateStore(), events)
|
||||
states := memory.NewUpdateStateStore()
|
||||
svc := NewService(states, events)
|
||||
if err := events.Append(ctx, userID, domain.UpdateEvent{
|
||||
UserID: userID, Type: domain.UpdateEventNewMessage, Pts: 1, PtsCount: 1,
|
||||
Date: 1700000001, Message: domain.Message{ID: 1, OwnerUserID: userID},
|
||||
|
|
@ -527,6 +535,138 @@ func TestAcknowledgeCurrentStateAdvancesConfirmedWatermark(t *testing.T) {
|
|||
if confirmed.Pts != 3 {
|
||||
t.Fatalf("confirmed watermark = %d, want advanced to 3", confirmed.Pts)
|
||||
}
|
||||
observed, ok := states.ObservedClientState(authKeyID, userID)
|
||||
if !ok || observed.Pts != 3 {
|
||||
t.Fatalf("getState observed watermark = %+v/%v, want pts=3", observed, ok)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetDifferenceRetainsOnlyClientObservedInputCursor(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
authKeyID := [8]byte{12}
|
||||
const userID int64 = 1000000012
|
||||
events := memory.NewUpdateEventStore()
|
||||
states := memory.NewUpdateStateStore()
|
||||
svc := NewService(states, events)
|
||||
for pts := 1; pts <= 2; pts++ {
|
||||
if err := events.Append(ctx, userID, domain.UpdateEvent{
|
||||
UserID: userID, Type: domain.UpdateEventNewMessage, Pts: pts, PtsCount: 1,
|
||||
Date: 1700000100 + pts, Message: domain.Message{ID: pts, OwnerUserID: userID},
|
||||
}); err != nil {
|
||||
t.Fatalf("append pts=%d: %v", pts, err)
|
||||
}
|
||||
}
|
||||
|
||||
// 服务端把 pts=1..2 放进 response,并不证明客户端收到了 response;observed 只能
|
||||
// 保持在本次 request 实际携带的 pts=0。
|
||||
diff, err := svc.GetDifference(ctx, authKeyID, userID, domain.UpdateState{Pts: 0, Date: 1700000100})
|
||||
if err != nil {
|
||||
t.Fatalf("first difference: %v", err)
|
||||
}
|
||||
if diff.State.Pts != 2 || len(diff.Events) != 2 {
|
||||
t.Fatalf("first difference = %+v, want response through pts=2", diff)
|
||||
}
|
||||
observed, ok := states.ObservedClientState(authKeyID, userID)
|
||||
if !ok || observed.Pts != 0 {
|
||||
t.Fatalf("observed after merely sending response = %+v/%v, want pts=0", observed, ok)
|
||||
}
|
||||
|
||||
// 客户端下一次明确带回 pts=2 后,才允许 retention 把共同安全水位推进到 2。
|
||||
if _, err := svc.GetDifference(ctx, authKeyID, userID, domain.UpdateState{Pts: 2, Date: 1700000102}); err != nil {
|
||||
t.Fatalf("confirming difference: %v", err)
|
||||
}
|
||||
observed, ok = states.ObservedClientState(authKeyID, userID)
|
||||
if !ok || observed.Pts != 2 {
|
||||
t.Fatalf("observed after client carried cursor = %+v/%v, want pts=2", observed, ok)
|
||||
}
|
||||
}
|
||||
|
||||
type retentionCheckpointEvents struct {
|
||||
*memory.UpdateEventStore
|
||||
pts int
|
||||
date int
|
||||
current int
|
||||
missFirst bool
|
||||
calls int
|
||||
}
|
||||
|
||||
func (s *retentionCheckpointEvents) UserUpdateRetentionCheckpoint(_ context.Context, _ [8]byte, _ int64) (int, int, bool, error) {
|
||||
s.calls++
|
||||
if s.missFirst && s.calls == 1 {
|
||||
return 0, 0, false, nil
|
||||
}
|
||||
return s.pts, s.date, s.pts > 0, nil
|
||||
}
|
||||
|
||||
func (s *retentionCheckpointEvents) MaxContiguousPts(_ context.Context, _ int64) (int, error) {
|
||||
return s.current, nil
|
||||
}
|
||||
|
||||
func TestGetDifferenceBelowRetainedFloorUsesEmptySliceCheckpoint(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
authKeyID := [8]byte{13}
|
||||
const userID int64 = 1000000013
|
||||
base := memory.NewUpdateEventStore()
|
||||
events := &retentionCheckpointEvents{UpdateEventStore: base, pts: 2, date: 1700000202, current: 3}
|
||||
states := memory.NewUpdateStateStore()
|
||||
svc := NewService(states, events)
|
||||
// Retention already removed pts 1..2; only the live tail remains.
|
||||
if err := base.Append(ctx, userID, domain.UpdateEvent{
|
||||
UserID: userID, Type: domain.UpdateEventNoop, Pts: 3, PtsCount: 1, Date: 1700000203,
|
||||
}); err != nil {
|
||||
t.Fatalf("append live tail: %v", err)
|
||||
}
|
||||
|
||||
checkpoint, err := svc.GetDifference(ctx, authKeyID, userID, domain.UpdateState{Pts: 0, Date: 1700000200})
|
||||
if err != nil {
|
||||
t.Fatalf("difference below retained floor: %v", err)
|
||||
}
|
||||
if !checkpoint.Partial || len(checkpoint.Events) != 0 || checkpoint.State.Pts != 2 || checkpoint.State.Date != 1700000202 {
|
||||
t.Fatalf("checkpoint difference = %+v, want empty differenceSlice at pts/date 2/1700000202", checkpoint)
|
||||
}
|
||||
|
||||
tail, err := svc.GetDifference(ctx, authKeyID, userID, checkpoint.State)
|
||||
if err != nil {
|
||||
t.Fatalf("difference from retained floor: %v", err)
|
||||
}
|
||||
if tail.Partial || len(tail.Events) != 1 || tail.Events[0].Pts != 3 || tail.State.Pts != 3 {
|
||||
t.Fatalf("tail difference = %+v, want normal event pts=3", tail)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGetDifferenceRechecksCheckpointWhenRetentionRacesEventRead(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
authKeyID := [8]byte{14}
|
||||
const userID int64 = 1000000014
|
||||
base := memory.NewUpdateEventStore()
|
||||
events := &retentionCheckpointEvents{
|
||||
UpdateEventStore: base,
|
||||
pts: 2,
|
||||
date: 1700000302,
|
||||
current: 3,
|
||||
missFirst: true,
|
||||
}
|
||||
if err := base.Append(ctx, userID, domain.UpdateEvent{
|
||||
UserID: userID, Type: domain.UpdateEventNoop, Pts: 3, PtsCount: 1, Date: 1700000303,
|
||||
}); err != nil {
|
||||
t.Fatalf("append live tail: %v", err)
|
||||
}
|
||||
|
||||
diff, err := NewService(memory.NewUpdateStateStore(), events).GetDifference(
|
||||
ctx,
|
||||
authKeyID,
|
||||
userID,
|
||||
domain.UpdateState{Pts: 0, Date: 1700000300},
|
||||
)
|
||||
if err != nil {
|
||||
t.Fatalf("difference across retention race: %v", err)
|
||||
}
|
||||
if events.calls != 2 {
|
||||
t.Fatalf("checkpoint probes = %d, want pre-read plus post-gap recheck", events.calls)
|
||||
}
|
||||
if !diff.Partial || len(diff.Events) != 0 || diff.State.Pts != 2 || diff.State.Date != 1700000302 {
|
||||
t.Fatalf("race checkpoint difference = %+v, want empty differenceSlice at retained floor", diff)
|
||||
}
|
||||
}
|
||||
|
||||
type captureDispatchAppender struct {
|
||||
|
|
@ -559,8 +699,9 @@ func (s *failingDispatchAppender) AppendAllocatedWithDispatch(context.Context, i
|
|||
}
|
||||
|
||||
type captureStateStore struct {
|
||||
saveCount int
|
||||
states map[[16]byte]domain.UpdateState
|
||||
saveCount int
|
||||
lastSaveAuthKeyID [8]byte
|
||||
states map[[16]byte]domain.UpdateState
|
||||
}
|
||||
|
||||
func (s *captureStateStore) Get(_ context.Context, authKeyID [8]byte, userID int64) (domain.UpdateState, bool, error) {
|
||||
|
|
@ -576,10 +717,15 @@ func (s *captureStateStore) Save(_ context.Context, authKeyID [8]byte, userID in
|
|||
s.states = make(map[[16]byte]domain.UpdateState)
|
||||
}
|
||||
s.saveCount++
|
||||
s.lastSaveAuthKeyID = authKeyID
|
||||
s.states[captureStateKey(authKeyID, userID)] = state
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *captureStateStore) ObserveClientState(_ context.Context, _ [8]byte, _ int64, _ domain.UpdateState) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *captureStateStore) Delete(_ context.Context, authKeyID [8]byte, userID int64) error {
|
||||
if s.states != nil {
|
||||
delete(s.states, captureStateKey(authKeyID, userID))
|
||||
|
|
|
|||
|
|
@ -20,6 +20,7 @@
|
|||
| `schema/canonical-227.tl` | **embed**,运行期 walker 的 227 字段布局(= gotd `td/_schema/tdesktop.tl` 的副本) | gotd 升级时 re-sync |
|
||||
| `_schema/layer-2NN.tl` | 历史层官方 schema(从 TDesktop git 抽,**仅生成期用**,下划线=不编译/不 embed) | 升级/下探 floor 时抽取 |
|
||||
| `schema/client-drift.tl` | **声明式**:客户端发的旧构造器老布局(body 与 227 不同的) | 发现客户端漂移时 +1 行 |
|
||||
| `schema/routable-compat.tl` | **仅结构预检**:已有 RPC fallback adapter 的非 canonical wire 布局(当前只含 4 个 DrKLO theme 构造器);与 canonical 图合并后完整 walk,但不自动升级 | 收敛既有手写 adapter 时维护,禁止借此新增业务 fallback |
|
||||
| `client_aliases.go` | 客户端漂移里 **body 与 227 字节一致**的,纯 `老CRC→227CRC` | 发现纯换 CRC 漂移时 +1 条 |
|
||||
| `tables_gen.go` | **生成产物**(勿手改):官方层降级表 + 入站升级表 + 新类型集 | 跑 `gen` 重生成 |
|
||||
| `gen/main.go` | 生成器:对拍 schema、证明机械性、产 `tables_gen.go` | 升级逻辑变更时 |
|
||||
|
|
@ -95,7 +96,7 @@ gofmt -w internal/compat/layerwire/ && go build ./... && go vet ./internal/...
|
|||
- 绿 = 通用引擎已能自动升级(复制共享字段 + 插 flags=0 + 按 kind 补默认)。**完事**。
|
||||
- `TestInboundDriftCoverage` 报 `needs converter A->B` = 有字段类型变更 → 往 `inbound.go fieldConverters` 加一条 `"A->B"`(可复用,参照 `Vector<int>->Vector<InputMessage>`)。
|
||||
- 报 `field X not defaultable` 或字段**改名** → 往 `inbound.go driftFieldRenames` 加 `"<method>\x00<227字段>": "<老字段>"`(参照 `bots.exportBotToken\x00bot`)。
|
||||
4. **绝不**为此写一个新的 `handleLegacyXxx` 解码 handler——那是旧做法,已全删。统一走数据 + 通用引擎。
|
||||
4. **绝不**为此写一个新的 `handleLegacyXxx` 解码 handler——统一走数据 + 通用引擎。`routable-compat.tl` 只给既存 DrKLO theme fallback 补 dispatcher 前结构门禁,不是新增 adapter 的入口。
|
||||
|
||||
## 操作 4:出站 `TestCoverageGate` 失败
|
||||
|
||||
|
|
|
|||
|
|
@ -30,8 +30,11 @@ func init() {
|
|||
// replaceWithBare consumes the canonical (227-only) object and emits a
|
||||
// bodyless constructor id the target layer understands.
|
||||
func replaceWithBare(id uint32) fallbackFunc {
|
||||
return func(cl *ctorLayout, in, out *bin.Buffer, layer int) error {
|
||||
if err := canonical.skipObject(in); err != nil {
|
||||
return func(cl *ctorLayout, in, out *bin.Buffer, layer, depth int, walk *walkState) error {
|
||||
if err := in.ConsumeID(cl.crc); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := walk.skipCtorBody(canonical, in, cl, depth); err != nil {
|
||||
return err
|
||||
}
|
||||
out.PutID(id)
|
||||
|
|
@ -47,6 +50,8 @@ var peerVectorField = fieldLayout{
|
|||
elem: &fieldLayout{kind: kindObject, typeName: "Peer", flagBit: -1},
|
||||
}
|
||||
|
||||
var pollOptionBytesField = fieldLayout{kind: kindBytes, flagBit: -1}
|
||||
|
||||
// transcodePollAnswerVoters downgrades pollAnswerVoters: canonical (227) made
|
||||
// voters conditional (flags.2?int) and added recent_voters (flags.2?Vector<Peer>);
|
||||
// older layers carry voters as a plain int. The leading CRC is already consumed.
|
||||
|
|
@ -54,34 +59,35 @@ var peerVectorField = fieldLayout{
|
|||
// 227: flags:# chosen:flags.0?true correct:flags.1?true option:bytes
|
||||
// voters:flags.2?int recent_voters:flags.2?Vector<Peer>
|
||||
// <=226: flags:# chosen:flags.0?true correct:flags.1?true option:bytes voters:int
|
||||
func transcodePollAnswerVoters(cl *ctorLayout, target uint32, in, out *bin.Buffer, layer int) error {
|
||||
func transcodePollAnswerVoters(cl *ctorLayout, target uint32, in, out *bin.Buffer, layer, depth int, walk *walkState) error {
|
||||
flags, err := in.Uint32()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
option, err := in.Bytes()
|
||||
if err != nil {
|
||||
optionStart := in.Buf
|
||||
if err := walk.skipValue(canonical, in, &pollOptionBytesField, cl, depth); err != nil {
|
||||
return err
|
||||
}
|
||||
optionRaw := optionStart[:len(optionStart)-len(in.Buf)]
|
||||
var voters int
|
||||
if flags&(1<<2) != 0 {
|
||||
if voters, err = in.Int(); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := canonical.skipValue(in, &peerVectorField); err != nil {
|
||||
if err := walk.skipValue(canonical, in, &peerVectorField, cl, depth); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
out.PutID(target)
|
||||
out.PutUint32(flags & 0b11) // retain chosen/correct, clear the moved bit 2
|
||||
out.PutBytes(option)
|
||||
out.Put(optionRaw)
|
||||
out.PutInt(voters)
|
||||
return nil
|
||||
}
|
||||
|
||||
// fallbackMessageEntity replaces any 227-only MessageEntity with
|
||||
// messageEntityUnknown, preserving offset/length so text positions stay valid.
|
||||
func fallbackMessageEntity(cl *ctorLayout, in, out *bin.Buffer, layer int) error {
|
||||
func fallbackMessageEntity(cl *ctorLayout, in, out *bin.Buffer, layer, depth int, walk *walkState) error {
|
||||
id, err := in.PeekID()
|
||||
if err != nil {
|
||||
return err
|
||||
|
|
@ -89,7 +95,7 @@ func fallbackMessageEntity(cl *ctorLayout, in, out *bin.Buffer, layer int) error
|
|||
if err := in.ConsumeID(id); err != nil {
|
||||
return err
|
||||
}
|
||||
offset, length, err := canonical.decodeOffsetLength(in, cl)
|
||||
offset, length, err := canonical.decodeOffsetLength(in, cl, depth, walk)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
|
|
@ -101,7 +107,7 @@ func fallbackMessageEntity(cl *ctorLayout, in, out *bin.Buffer, layer int) error
|
|||
|
||||
// decodeOffsetLength walks a constructor body (no leading CRC) per the canonical
|
||||
// layout, returning its offset/length int fields and discarding the rest.
|
||||
func (m *schemaModel) decodeOffsetLength(in *bin.Buffer, cl *ctorLayout) (offset, length int, err error) {
|
||||
func (m *schemaModel) decodeOffsetLength(in *bin.Buffer, cl *ctorLayout, depth int, walk *walkState) (offset, length int, err error) {
|
||||
var flags map[string]uint32
|
||||
for i := range cl.fields {
|
||||
f := &cl.fields[i]
|
||||
|
|
@ -129,7 +135,7 @@ func (m *schemaModel) decodeOffsetLength(in *bin.Buffer, cl *ctorLayout) (offset
|
|||
return
|
||||
}
|
||||
default:
|
||||
if err = m.skipValue(in, f); err != nil {
|
||||
if err = walk.skipValue(m, in, f, cl, depth); err != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -47,16 +47,19 @@ var driftFieldRenames = map[string]string{
|
|||
// fieldConverter rewrites one field whose wire type changed between the old and
|
||||
// canonical layout. Keyed by "<oldTypeSig>-><newTypeSig>"; raw is the old field's
|
||||
// encoded bytes. Reusable across any method with the same type change.
|
||||
type fieldConverter func(raw []byte, out *bin.Buffer) error
|
||||
type fieldConverter func(raw []byte, out *bin.Buffer, walk *walkState, owner *ctorLayout, field *fieldLayout) error
|
||||
|
||||
var fieldConverters = map[string]fieldConverter{
|
||||
// id:Vector<int> -> id:Vector<InputMessage> (wrap each int in inputMessageID).
|
||||
"Vector<int>->Vector<InputMessage>": func(raw []byte, out *bin.Buffer) error {
|
||||
"Vector<int>->Vector<InputMessage>": func(raw []byte, out *bin.Buffer, walk *walkState, owner *ctorLayout, field *fieldLayout) error {
|
||||
in := &bin.Buffer{Buf: raw}
|
||||
n, err := in.VectorHeader()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if max := walk.vectorLimit(owner, field); n > max {
|
||||
return limitf("vector %s.%s length %d exceeds limit %d", ownerName(owner), fieldName(field), n, max)
|
||||
}
|
||||
out.PutVectorHeader(n)
|
||||
for i := 0; i < n; i++ {
|
||||
v, err := in.Int()
|
||||
|
|
@ -66,10 +69,13 @@ var fieldConverters = map[string]fieldConverter{
|
|||
out.PutID(inputMessageID)
|
||||
out.PutInt(v)
|
||||
}
|
||||
if in.Len() != 0 {
|
||||
return malformedf("%d trailing bytes in Vector<int> converter", in.Len())
|
||||
}
|
||||
return nil
|
||||
},
|
||||
// bot_id:long -> bot:InputUser{user_id, access_hash=0}.
|
||||
"long->InputUser": func(raw []byte, out *bin.Buffer) error {
|
||||
"long->InputUser": func(raw []byte, out *bin.Buffer, walk *walkState, owner *ctorLayout, field *fieldLayout) error {
|
||||
in := &bin.Buffer{Buf: raw}
|
||||
id, err := in.Long()
|
||||
if err != nil {
|
||||
|
|
@ -78,11 +84,14 @@ var fieldConverters = map[string]fieldConverter{
|
|||
out.PutID(inputUserID)
|
||||
out.PutLong(id)
|
||||
out.PutLong(0)
|
||||
if in.Len() != 0 {
|
||||
return malformedf("%d trailing bytes in long converter", in.Len())
|
||||
}
|
||||
return nil
|
||||
},
|
||||
// channel:InputChannel -> peer:InputPeer for the old channels.editCreator
|
||||
// Android constructor. Concrete layouts are otherwise byte-compatible.
|
||||
"InputChannel->InputPeer": func(raw []byte, out *bin.Buffer) error {
|
||||
"InputChannel->InputPeer": func(raw []byte, out *bin.Buffer, walk *walkState, owner *ctorLayout, field *fieldLayout) error {
|
||||
in := &bin.Buffer{Buf: raw}
|
||||
id, err := in.ID()
|
||||
if err != nil {
|
||||
|
|
@ -116,8 +125,12 @@ var fieldConverters = map[string]fieldConverter{
|
|||
// id + body) is what to dispatch.
|
||||
func UpgradeInbound(id uint32, in *bin.Buffer) (*bin.Buffer, bool, error) {
|
||||
if newID, ok := UpgradeMethodCRC(id); ok {
|
||||
if len(in.Buf) < 4 {
|
||||
return nil, false, fmt.Errorf("layerwire: short inbound buffer for %#08x", id)
|
||||
target := canonical.byCRC[newID]
|
||||
if target == nil || !target.isFunc {
|
||||
return nil, true, malformedf("alias %#08x targets unknown canonical method %#08x", id, newID)
|
||||
}
|
||||
if err := validateAliasedMethod(id, target, in.Buf); err != nil {
|
||||
return nil, true, err
|
||||
}
|
||||
// Copy rather than rewrite in place: never mutate the caller's buffer
|
||||
// (matches the body-transform path, which also returns a fresh buffer).
|
||||
|
|
@ -126,15 +139,36 @@ func UpgradeInbound(id uint32, in *bin.Buffer) (*bin.Buffer, bool, error) {
|
|||
return out, true, nil
|
||||
}
|
||||
if old := driftModel.byCRC[id]; old != nil {
|
||||
out, err := upgradeFromDrift(old, in)
|
||||
out, err := upgradeFromDrift(old, in, newWalkState())
|
||||
if err != nil {
|
||||
return nil, false, fmt.Errorf("layerwire: upgrade %s (%#08x): %w", old.name, id, err)
|
||||
return nil, true, classifyWalkError(fmt.Errorf("layerwire: upgrade %s (%#08x): %w", old.name, id, err))
|
||||
}
|
||||
return out, true, nil
|
||||
}
|
||||
return nil, false, nil
|
||||
}
|
||||
|
||||
// validateAliasedMethod validates the old-id/canonical-body shape before
|
||||
// allocating the replacement buffer. The body is walked against the canonical
|
||||
// target layout while the original constructor id remains untouched.
|
||||
func validateAliasedMethod(oldID uint32, target *ctorLayout, raw []byte) error {
|
||||
walk := newWalkState()
|
||||
if err := walk.enter(1, "constructor"); err != nil {
|
||||
return err
|
||||
}
|
||||
b := &bin.Buffer{Buf: raw}
|
||||
if err := b.ConsumeID(oldID); err != nil {
|
||||
return classifyWalkError(err)
|
||||
}
|
||||
if err := walk.skipCtorBody(canonical, b, target, 1); err != nil {
|
||||
return classifyWalkError(err)
|
||||
}
|
||||
if b.Len() != 0 {
|
||||
return malformedf("%d trailing bytes after aliased method %s", b.Len(), target.name)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// IsClientDrift reports whether id is a client-private constructor (DrKLO
|
||||
// constructor drift), as opposed to official layer drift from api.tl.
|
||||
func IsClientDrift(id uint32) bool {
|
||||
|
|
@ -146,11 +180,14 @@ func IsClientDrift(id uint32) bool {
|
|||
|
||||
// upgradeFromDrift rebuilds a canonical (227) request from an old client-drift
|
||||
// body, comparing the declared old layout to the canonical layout field by field.
|
||||
func upgradeFromDrift(old *ctorLayout, in *bin.Buffer) (*bin.Buffer, error) {
|
||||
func upgradeFromDrift(old *ctorLayout, in *bin.Buffer, walk *walkState) (*bin.Buffer, error) {
|
||||
target := canonical.byName[old.name]
|
||||
if target == nil {
|
||||
return nil, fmt.Errorf("no canonical method %q", old.name)
|
||||
}
|
||||
if err := walk.enter(1, "constructor"); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := in.ConsumeID(old.crc); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
|
@ -179,7 +216,7 @@ func upgradeFromDrift(old *ctorLayout, in *bin.Buffer) (*bin.Buffer, error) {
|
|||
continue
|
||||
}
|
||||
pre := in.Buf
|
||||
if err := canonical.skipValue(in, f); err != nil {
|
||||
if err := walk.skipValue(canonical, in, f, old, 1); err != nil {
|
||||
return nil, fmt.Errorf("decode old field %q: %w", f.name, err)
|
||||
}
|
||||
vals[f.name] = pre[:len(pre)-len(in.Buf)]
|
||||
|
|
@ -208,7 +245,7 @@ func upgradeFromDrift(old *ctorLayout, in *bin.Buffer) (*bin.Buffer, error) {
|
|||
if conv == nil {
|
||||
return nil, fmt.Errorf("field %q: no converter %s->%s", nf.name, typeSig(of), typeSig(nf))
|
||||
}
|
||||
if err := conv(vals[oldName], out); err != nil {
|
||||
if err := conv(vals[oldName], out, walk, old, of); err != nil {
|
||||
return nil, fmt.Errorf("field %q convert: %w", nf.name, err)
|
||||
}
|
||||
} else {
|
||||
|
|
|
|||
80
internal/compat/layerwire/routable.go
Normal file
80
internal/compat/layerwire/routable.go
Normal file
|
|
@ -0,0 +1,80 @@
|
|||
package layerwire
|
||||
|
||||
import (
|
||||
_ "embed"
|
||||
"fmt"
|
||||
|
||||
"github.com/gotd/td/bin"
|
||||
)
|
||||
|
||||
const maxOpaqueRequestBytes = 16 << 20
|
||||
|
||||
//go:embed schema/routable-compat.tl
|
||||
var routableCompatSchema string
|
||||
|
||||
// routable combines the canonical Layer 227 model with the small set of
|
||||
// explicitly declared compatibility-only methods. Nested objects in those
|
||||
// methods are canonical Input* constructors, so one combined graph is needed
|
||||
// for the same depth/vector/bytes walker to validate the complete request.
|
||||
var routable = mustLoadRoutable()
|
||||
|
||||
func mustLoadRoutable() *schemaModel {
|
||||
compat, err := parseSchemaModel(routableCompatSchema)
|
||||
if err != nil {
|
||||
panic("layerwire: parse routable compat schema: " + err.Error())
|
||||
}
|
||||
m := &schemaModel{
|
||||
byCRC: make(map[uint32]*ctorLayout, len(canonical.byCRC)+len(compat.byCRC)),
|
||||
byName: make(map[string]*ctorLayout, len(canonical.byName)+len(compat.byName)),
|
||||
bareByT: make(map[string]*ctorLayout, len(canonical.bareByT)),
|
||||
ctorsOfT: make(map[string][]*ctorLayout, len(canonical.ctorsOfT)),
|
||||
}
|
||||
for id, cl := range canonical.byCRC {
|
||||
m.byCRC[id] = cl
|
||||
}
|
||||
for name, cl := range canonical.byName {
|
||||
m.byName[name] = cl
|
||||
}
|
||||
for name, cl := range canonical.bareByT {
|
||||
m.bareByT[name] = cl
|
||||
}
|
||||
for name, ctors := range canonical.ctorsOfT {
|
||||
m.ctorsOfT[name] = ctors
|
||||
}
|
||||
for id, cl := range compat.byCRC {
|
||||
if existing := m.byCRC[id]; existing != nil {
|
||||
panic(fmt.Sprintf("layerwire: routable compat crc %#08x collides with %s", id, existing.name))
|
||||
}
|
||||
m.byCRC[id] = cl
|
||||
m.byName[cl.name] = cl
|
||||
}
|
||||
return m
|
||||
}
|
||||
|
||||
// ValidateRoutableRequest validates every request shape the router knows how to
|
||||
// decode, including compatibility-only fallback methods. known=false denotes
|
||||
// a genuinely unknown top-level constructor. Such a request is never decoded:
|
||||
// it is treated as opaque, word-aligned TL data, bounded by both this total-size
|
||||
// cap and mtprotoedge's transport/RPC budgets, and must continue to the router's
|
||||
// compatibility trace rather than being mislabeled as malformed input.
|
||||
func ValidateRoutableRequest(body []byte) (known bool, err error) {
|
||||
b := &bin.Buffer{Buf: body}
|
||||
id, err := b.PeekID()
|
||||
if err != nil {
|
||||
return false, classifyWalkError(err)
|
||||
}
|
||||
cl := routable.byCRC[id]
|
||||
if cl == nil {
|
||||
if len(body) > maxOpaqueRequestBytes {
|
||||
return false, limitf("opaque request length %d exceeds limit %d", len(body), maxOpaqueRequestBytes)
|
||||
}
|
||||
if len(body)%bin.Word != 0 {
|
||||
return false, malformedf("opaque request length %d is not word aligned", len(body))
|
||||
}
|
||||
return false, nil
|
||||
}
|
||||
if !cl.isFunc {
|
||||
return true, malformedf("constructor %s (%#08x) is not a method", cl.name, id)
|
||||
}
|
||||
return true, validateRequestLayout(routable, cl, body)
|
||||
}
|
||||
43
internal/compat/layerwire/routable_test.go
Normal file
43
internal/compat/layerwire/routable_test.go
Normal file
|
|
@ -0,0 +1,43 @@
|
|||
package layerwire
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"testing"
|
||||
|
||||
"github.com/gotd/td/bin"
|
||||
"github.com/gotd/td/tg"
|
||||
)
|
||||
|
||||
func TestValidateRoutableRequestCompatibilityAndUnknown(t *testing.T) {
|
||||
t.Run("legacy theme is fully walked", func(t *testing.T) {
|
||||
var b bin.Buffer
|
||||
b.PutID(0x8d9d742b)
|
||||
b.PutString("android")
|
||||
(&tg.InputThemeSlug{Slug: "night"}).Encode(&b)
|
||||
b.PutLong(42)
|
||||
known, err := ValidateRoutableRequest(b.Buf)
|
||||
if err != nil || !known {
|
||||
t.Fatalf("legacy theme known=%v err=%v, want true/nil", known, err)
|
||||
}
|
||||
|
||||
b.Buf = b.Buf[:len(b.Buf)-4]
|
||||
known, err = ValidateRoutableRequest(b.Buf)
|
||||
if !known || !errors.Is(err, ErrMalformed) {
|
||||
t.Fatalf("truncated legacy theme known=%v err=%v, want true/malformed", known, err)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("unknown stays opaque and bounded", func(t *testing.T) {
|
||||
var b bin.Buffer
|
||||
b.PutID(0x12345678)
|
||||
b.PutUint32(0xffffffff)
|
||||
known, err := ValidateRoutableRequest(b.Buf)
|
||||
if err != nil || known {
|
||||
t.Fatalf("opaque unknown known=%v err=%v, want false/nil", known, err)
|
||||
}
|
||||
known, err = ValidateRoutableRequest(append(b.Buf, 1))
|
||||
if known || !errors.Is(err, ErrMalformed) {
|
||||
t.Fatalf("unaligned unknown known=%v err=%v, want false/malformed", known, err)
|
||||
}
|
||||
})
|
||||
}
|
||||
11
internal/compat/layerwire/schema/routable-compat.tl
Normal file
11
internal/compat/layerwire/schema/routable-compat.tl
Normal file
|
|
@ -0,0 +1,11 @@
|
|||
// Hand-maintained request layouts that are intentionally handled by the RPC
|
||||
// fallback instead of gotd's canonical ServerDispatcher. They still belong in
|
||||
// the structural preflight model: fallback handlers must never become a way to
|
||||
// bypass the canonical vector/depth/bytes budgets.
|
||||
|
||||
---functions---
|
||||
|
||||
compat.legacyCreateTheme#8432c21f flags:# slug:string title:string document:flags.2?InputDocument settings:flags.3?InputThemeSettings = Object;
|
||||
compat.legacyUpdateTheme#5cb367d5 flags:# format:string theme:InputTheme slug:flags.0?string title:flags.1?string document:flags.2?InputDocument settings:flags.3?InputThemeSettings = Object;
|
||||
compat.legacyInstallTheme#7ae43737 flags:# dark:flags.0?true format:flags.1?string theme:flags.1?InputTheme = Object;
|
||||
compat.legacyGetTheme#8d9d742b format:string theme:InputTheme document_id:long = Object;
|
||||
|
|
@ -139,11 +139,11 @@ func (lt *layerTables) fieldDirty(f *fieldLayout) bool {
|
|||
// field drop. The leading CRC has already been consumed from in; the transform
|
||||
// reads the canonical body from in and writes the target-layer object (whose
|
||||
// constructor id is target) to out.
|
||||
type structuralFunc func(cl *ctorLayout, target uint32, in, out *bin.Buffer, layer int) error
|
||||
type structuralFunc func(cl *ctorLayout, target uint32, in, out *bin.Buffer, layer, depth int, walk *walkState) error
|
||||
|
||||
// fallbackFunc replaces a layer-absent (227-only) constructor with an
|
||||
// equivalent the target layer understands. The leading CRC is NOT yet consumed.
|
||||
type fallbackFunc func(cl *ctorLayout, in, out *bin.Buffer, layer int) error
|
||||
type fallbackFunc func(cl *ctorLayout, in, out *bin.Buffer, layer, depth int, walk *walkState) error
|
||||
|
||||
// structuralTransforms and the newType fallback registries are populated in
|
||||
// fallback.go. newTypeFallbacks is keyed by canonical CRC (specific override);
|
||||
|
|
@ -175,11 +175,12 @@ func Transcode(canonicalBytes []byte, layer int) ([]byte, error) {
|
|||
}
|
||||
in := &bin.Buffer{Buf: canonicalBytes}
|
||||
out := &bin.Buffer{}
|
||||
if err := lt.transcodeObject(in, out, layer); err != nil {
|
||||
return nil, err
|
||||
walk := newWalkState()
|
||||
if err := lt.transcodeObject(in, out, layer, 1, walk); err != nil {
|
||||
return nil, classifyWalkError(err)
|
||||
}
|
||||
if in.Len() != 0 {
|
||||
return nil, fmt.Errorf("layerwire: %d trailing bytes after transcode to layer %d", in.Len(), layer)
|
||||
return nil, malformedf("%d trailing bytes after transcode to layer %d", in.Len(), layer)
|
||||
}
|
||||
return out.Buf, nil
|
||||
}
|
||||
|
|
@ -198,7 +199,10 @@ func UpgradeMethodCRC(oldID uint32) (uint32, bool) {
|
|||
return newID, ok
|
||||
}
|
||||
|
||||
func (lt *layerTables) transcodeObject(in, out *bin.Buffer, layer int) error {
|
||||
func (lt *layerTables) transcodeObject(in, out *bin.Buffer, layer, depth int, walk *walkState) error {
|
||||
if err := walk.enter(depth, "constructor"); err != nil {
|
||||
return err
|
||||
}
|
||||
id, err := in.PeekID()
|
||||
if err != nil {
|
||||
return err
|
||||
|
|
@ -216,10 +220,10 @@ func (lt *layerTables) transcodeObject(in, out *bin.Buffer, layer int) error {
|
|||
if fn == nil {
|
||||
return fmt.Errorf("layerwire: no structural transform %q for %s@%d", rule.structural, cl.name, layer)
|
||||
}
|
||||
return fn(cl, rule.target, in, out, layer)
|
||||
return fn(cl, rule.target, in, out, layer, depth, walk)
|
||||
}
|
||||
out.PutID(rule.target)
|
||||
return lt.transcodeBody(in, out, cl, rule.keep, layer)
|
||||
return lt.transcodeBody(in, out, cl, rule.keep, layer, depth, walk)
|
||||
}
|
||||
if lt.newTypes[id] {
|
||||
fn := newTypeFallbacks[id]
|
||||
|
|
@ -229,12 +233,15 @@ func (lt *layerTables) transcodeObject(in, out *bin.Buffer, layer int) error {
|
|||
if fn == nil {
|
||||
return fmt.Errorf("layerwire: %s (%#08x) absent at layer %d and no fallback", cl.name, id, layer)
|
||||
}
|
||||
return fn(cl, in, out, layer)
|
||||
return fn(cl, in, out, layer, depth, walk)
|
||||
}
|
||||
if !lt.dirty[id] {
|
||||
// Unaffected subtree: byte-for-byte copy.
|
||||
pre := in.Buf
|
||||
if err := canonical.skipObject(in); err != nil {
|
||||
if err := in.ConsumeID(id); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := walk.skipCtorBody(canonical, in, cl, depth); err != nil {
|
||||
return err
|
||||
}
|
||||
out.Put(pre[:len(pre)-len(in.Buf)])
|
||||
|
|
@ -245,14 +252,14 @@ func (lt *layerTables) transcodeObject(in, out *bin.Buffer, layer int) error {
|
|||
return err
|
||||
}
|
||||
out.PutID(id)
|
||||
return lt.transcodeBody(in, out, cl, nil, layer)
|
||||
return lt.transcodeBody(in, out, cl, nil, layer, depth, walk)
|
||||
}
|
||||
|
||||
// transcodeBody re-encodes a constructor body. keep==nil means retain every
|
||||
// field (recursing into dirty descendants); otherwise only the named canonical
|
||||
// fields are written, flag integers are remasked to the retained bits, and
|
||||
// dropped fields are read-and-discarded.
|
||||
func (lt *layerTables) transcodeBody(in, out *bin.Buffer, cl *ctorLayout, keep map[string]bool, layer int) error {
|
||||
func (lt *layerTables) transcodeBody(in, out *bin.Buffer, cl *ctorLayout, keep map[string]bool, layer, depth int, walk *walkState) error {
|
||||
kept := func(name string) bool { return keep == nil || keep[name] }
|
||||
var flags map[string]uint32
|
||||
for i := range cl.fields {
|
||||
|
|
@ -276,10 +283,10 @@ func (lt *layerTables) transcodeBody(in, out *bin.Buffer, cl *ctorLayout, keep m
|
|||
continue
|
||||
}
|
||||
if kept(f.name) {
|
||||
if err := lt.transcodeValue(in, out, f, layer); err != nil {
|
||||
if err := lt.transcodeValue(in, out, f, cl, layer, depth, walk); err != nil {
|
||||
return fmt.Errorf("%s.%s: %w", cl.name, f.name, err)
|
||||
}
|
||||
} else if err := canonical.skipValue(in, f); err != nil {
|
||||
} else if err := walk.skipValue(canonical, in, f, cl, depth); err != nil {
|
||||
return fmt.Errorf("%s.%s (drop): %w", cl.name, f.name, err)
|
||||
}
|
||||
}
|
||||
|
|
@ -301,10 +308,10 @@ func (lt *layerTables) keptMask(cl *ctorLayout, flagName string, kept func(strin
|
|||
|
||||
// transcodeValue writes one present field value, recursing only into dirty
|
||||
// subtrees and byte-copying everything else.
|
||||
func (lt *layerTables) transcodeValue(in, out *bin.Buffer, f *fieldLayout, layer int) error {
|
||||
func (lt *layerTables) transcodeValue(in, out *bin.Buffer, f *fieldLayout, owner *ctorLayout, layer, depth int, walk *walkState) error {
|
||||
if !lt.fieldDirty(f) {
|
||||
pre := in.Buf
|
||||
if err := canonical.skipValue(in, f); err != nil {
|
||||
if err := walk.skipValue(canonical, in, f, owner, depth); err != nil {
|
||||
return err
|
||||
}
|
||||
out.Put(pre[:len(pre)-len(in.Buf)])
|
||||
|
|
@ -312,6 +319,10 @@ func (lt *layerTables) transcodeValue(in, out *bin.Buffer, f *fieldLayout, layer
|
|||
}
|
||||
switch f.kind {
|
||||
case kindVector, kindVectorBare:
|
||||
vectorDepth := depth + 1
|
||||
if vectorDepth <= 0 || vectorDepth > walk.limits.maxDepth {
|
||||
return limitf("vector nesting depth %d exceeds limit %d", vectorDepth, walk.limits.maxDepth)
|
||||
}
|
||||
if f.kind == kindVector {
|
||||
id, err := in.Uint32()
|
||||
if err != nil {
|
||||
|
|
@ -326,23 +337,36 @@ func (lt *layerTables) transcodeValue(in, out *bin.Buffer, f *fieldLayout, layer
|
|||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if n < 0 {
|
||||
return malformedf("negative vector length %d", n)
|
||||
}
|
||||
if max := walk.vectorLimit(owner, f); n > max {
|
||||
return limitf("vector %s.%s length %d exceeds limit %d", ownerName(owner), fieldName(f), n, max)
|
||||
}
|
||||
if err := walk.addUnits(n, "vector "+ownerName(owner)+"."+fieldName(f)); err != nil {
|
||||
return err
|
||||
}
|
||||
out.PutInt(n)
|
||||
for i := 0; i < n; i++ {
|
||||
if err := lt.transcodeValue(in, out, f.elem, layer); err != nil {
|
||||
if err := lt.transcodeValue(in, out, f.elem, nil, layer, vectorDepth, walk); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
return nil
|
||||
case kindObject:
|
||||
return lt.transcodeObject(in, out, layer)
|
||||
return lt.transcodeObject(in, out, layer, depth+1, walk)
|
||||
case kindBareObject:
|
||||
bareDepth := depth + 1
|
||||
if err := walk.enter(bareDepth, "bare constructor"); err != nil {
|
||||
return err
|
||||
}
|
||||
cl, ok := canonical.bareByT[f.typeName]
|
||||
if !ok {
|
||||
return fmt.Errorf("unknown bare type %q", f.typeName)
|
||||
}
|
||||
// Bare objects have no CRC and (within 220..227) no changed bare ctor;
|
||||
// recurse all-kept to reach any dirty descendants.
|
||||
return lt.transcodeBody(in, out, cl, nil, layer)
|
||||
return lt.transcodeBody(in, out, cl, nil, layer, bareDepth, walk)
|
||||
default:
|
||||
// Primitive marked dirty should be impossible.
|
||||
return fmt.Errorf("unexpected dirty primitive kind %d", f.kind)
|
||||
|
|
|
|||
|
|
@ -1,32 +1,205 @@
|
|||
package layerwire
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"math"
|
||||
|
||||
"github.com/gotd/td/bin"
|
||||
)
|
||||
|
||||
// ErrMalformed identifies invalid or truncated TL wire data. Callers may use
|
||||
// errors.Is to distinguish it from an otherwise well-formed request which was
|
||||
// rejected by a walker resource limit.
|
||||
var ErrMalformed = errors.New("layerwire: malformed TL")
|
||||
|
||||
// ErrResourceLimit identifies structurally valid-looking TL input which would
|
||||
// exceed a walker resource budget.
|
||||
var ErrResourceLimit = errors.New("layerwire: resource limit")
|
||||
|
||||
const (
|
||||
defaultMaxVectorElements = 4096
|
||||
defaultMaxWalkDepth = 32
|
||||
defaultMaxWalkUnits = 131072 // constructors + declared vector elements
|
||||
defaultMaxFieldBytes = 16 << 20
|
||||
defaultMaxTotalBytes = 32 << 20
|
||||
)
|
||||
|
||||
// A very small number of API methods have a documented limit above the
|
||||
// package-wide default. Keeping overrides keyed by constructor and field makes
|
||||
// every exception explicit and prevents a large vector in an unrelated method
|
||||
// from inheriting the larger allowance.
|
||||
type vectorLimitKey struct {
|
||||
owner string
|
||||
field string
|
||||
}
|
||||
|
||||
var vectorElementLimitOverrides = map[vectorLimitKey]int{
|
||||
{owner: "contacts.editCloseFriends", field: "id"}: 5000,
|
||||
{owner: "contacts.setBlocked", field: "id"}: 5000,
|
||||
}
|
||||
|
||||
type walkLimits struct {
|
||||
maxVectorElements int
|
||||
maxDepth int
|
||||
maxUnits uint64
|
||||
maxFieldBytes uint64
|
||||
maxTotalBytes uint64
|
||||
}
|
||||
|
||||
var defaultWalkLimits = walkLimits{
|
||||
maxVectorElements: defaultMaxVectorElements,
|
||||
maxDepth: defaultMaxWalkDepth,
|
||||
maxUnits: defaultMaxWalkUnits,
|
||||
maxFieldBytes: defaultMaxFieldBytes,
|
||||
maxTotalBytes: defaultMaxTotalBytes,
|
||||
}
|
||||
|
||||
// walkState is deliberately request-scoped. Every branch of one transform
|
||||
// shares it, so splitting a large value across nested constructors or vectors
|
||||
// cannot reset the aggregate budgets.
|
||||
type walkState struct {
|
||||
limits walkLimits
|
||||
units uint64
|
||||
bytes uint64
|
||||
}
|
||||
|
||||
func newWalkState() *walkState {
|
||||
return &walkState{limits: defaultWalkLimits}
|
||||
}
|
||||
|
||||
func malformedf(format string, args ...any) error {
|
||||
return fmt.Errorf("%w: %s", ErrMalformed, fmt.Sprintf(format, args...))
|
||||
}
|
||||
|
||||
func limitf(format string, args ...any) error {
|
||||
return fmt.Errorf("%w: %s", ErrResourceLimit, fmt.Sprintf(format, args...))
|
||||
}
|
||||
|
||||
// classifyWalkError makes all public walker/transform failures classifiable,
|
||||
// including errors returned by the low-level gotd bin decoder.
|
||||
func classifyWalkError(err error) error {
|
||||
if err == nil || errors.Is(err, ErrMalformed) || errors.Is(err, ErrResourceLimit) {
|
||||
return err
|
||||
}
|
||||
return fmt.Errorf("%w: %v", ErrMalformed, err)
|
||||
}
|
||||
|
||||
func (s *walkState) enter(depth int, what string) error {
|
||||
if depth <= 0 || depth > s.limits.maxDepth {
|
||||
return limitf("%s nesting depth %d exceeds limit %d", what, depth, s.limits.maxDepth)
|
||||
}
|
||||
return s.addUnits(1, what)
|
||||
}
|
||||
|
||||
func (s *walkState) addUnits(n int, what string) error {
|
||||
if n < 0 {
|
||||
return malformedf("negative %s count %d", what, n)
|
||||
}
|
||||
u := uint64(n)
|
||||
// Subtraction form avoids overflow even if limits are changed later.
|
||||
if s.units > s.limits.maxUnits || u > s.limits.maxUnits-s.units {
|
||||
return limitf("constructor/vector element budget exceeds %d at %s", s.limits.maxUnits, what)
|
||||
}
|
||||
s.units += u
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *walkState) addBytes(n uint64, what string) error {
|
||||
if n > s.limits.maxFieldBytes {
|
||||
return limitf("%s payload length %d exceeds per-field limit %d", what, n, s.limits.maxFieldBytes)
|
||||
}
|
||||
if s.bytes > s.limits.maxTotalBytes || n > s.limits.maxTotalBytes-s.bytes {
|
||||
return limitf("string/bytes payload budget exceeds %d at %s", s.limits.maxTotalBytes, what)
|
||||
}
|
||||
s.bytes += n
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *walkState) vectorLimit(owner *ctorLayout, f *fieldLayout) int {
|
||||
if owner != nil && f != nil {
|
||||
if n := vectorElementLimitOverrides[vectorLimitKey{owner: owner.name, field: f.name}]; n > 0 {
|
||||
return n
|
||||
}
|
||||
}
|
||||
return s.limits.maxVectorElements
|
||||
}
|
||||
|
||||
const maxConstructorFlagWords = 8
|
||||
|
||||
type constructorFlagWord struct {
|
||||
name string
|
||||
value uint32
|
||||
}
|
||||
|
||||
// ValidateCanonicalRequest performs a complete, allocation-free structural
|
||||
// preflight of one canonical Layer 227 method request. It is intended for the
|
||||
// router seam immediately before typed dispatch. A successful result means the
|
||||
// walker consumed exactly one known function constructor and all of its body.
|
||||
func ValidateCanonicalRequest(body []byte) error {
|
||||
b := &bin.Buffer{Buf: body}
|
||||
id, err := b.PeekID()
|
||||
if err != nil {
|
||||
return classifyWalkError(err)
|
||||
}
|
||||
cl := canonical.byCRC[id]
|
||||
if cl == nil {
|
||||
return malformedf("unknown canonical request constructor %#08x", id)
|
||||
}
|
||||
if !cl.isFunc {
|
||||
return malformedf("constructor %s (%#08x) is not a method", cl.name, id)
|
||||
}
|
||||
return validateRequestLayout(canonical, cl, body)
|
||||
}
|
||||
|
||||
func validateRequestLayout(m *schemaModel, cl *ctorLayout, body []byte) error {
|
||||
b := &bin.Buffer{Buf: body}
|
||||
s := newWalkState()
|
||||
if err := s.skipObject(m, b, 1); err != nil {
|
||||
return classifyWalkError(err)
|
||||
}
|
||||
if b.Len() != 0 {
|
||||
return malformedf("%d trailing bytes after canonical request %s", b.Len(), cl.name)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// skipObject advances b past one boxed object (CRC + body), resolving the
|
||||
// constructor from the canonical schema.
|
||||
// constructor from m. This compatibility wrapper creates a fresh budget; all
|
||||
// production transforms call the stateful variant directly.
|
||||
func (m *schemaModel) skipObject(b *bin.Buffer) error {
|
||||
return classifyWalkError(newWalkState().skipObject(m, b, 1))
|
||||
}
|
||||
|
||||
func (s *walkState) skipObject(m *schemaModel, b *bin.Buffer, depth int) error {
|
||||
if err := s.enter(depth, "constructor"); err != nil {
|
||||
return err
|
||||
}
|
||||
id, err := b.PeekID()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
cl, ok := m.byCRC[id]
|
||||
if !ok {
|
||||
return fmt.Errorf("layerwire: unknown constructor %#08x", id)
|
||||
return malformedf("unknown constructor %#08x", id)
|
||||
}
|
||||
if err := b.ConsumeID(id); err != nil {
|
||||
return err
|
||||
}
|
||||
return m.skipCtorBody(b, cl)
|
||||
return s.skipCtorBody(m, b, cl, depth)
|
||||
}
|
||||
|
||||
// skipCtorBody advances b past a constructor body (no leading CRC), evaluating
|
||||
// flag integers so conditional fields are read iff present.
|
||||
func (m *schemaModel) skipCtorBody(b *bin.Buffer, cl *ctorLayout) error {
|
||||
var flags map[string]uint32
|
||||
// flag integers so conditional fields are read iff present. The constructor's
|
||||
// unit and depth have already been charged by the caller.
|
||||
func (s *walkState) skipCtorBody(m *schemaModel, b *bin.Buffer, cl *ctorLayout, depth int) error {
|
||||
// Layer 227 constructors currently use at most flags + flags2. Keep generous fixed stack
|
||||
// storage so the allocation-free preflight remains allocation-free on the hottest flagged
|
||||
// methods; the explicit bound also prevents a future malformed/generated layout from turning
|
||||
// every request into an attacker-amplified map allocation.
|
||||
var flags [maxConstructorFlagWords]constructorFlagWord
|
||||
flagCount := 0
|
||||
for i := range cl.fields {
|
||||
f := &cl.fields[i]
|
||||
if f.isFlags {
|
||||
|
|
@ -34,59 +207,80 @@ func (m *schemaModel) skipCtorBody(b *bin.Buffer, cl *ctorLayout) error {
|
|||
if err != nil {
|
||||
return fmt.Errorf("%s.%s: %w", cl.name, f.name, err)
|
||||
}
|
||||
if flags == nil {
|
||||
flags = make(map[string]uint32, 2)
|
||||
if flagCount >= len(flags) {
|
||||
return limitf("constructor %s has more than %d flags words", cl.name, len(flags))
|
||||
}
|
||||
flags[f.name] = v
|
||||
flags[flagCount] = constructorFlagWord{name: f.name, value: v}
|
||||
flagCount++
|
||||
continue
|
||||
}
|
||||
if f.conditional() && flags[f.flagName]&(1<<uint(f.flagBit)) == 0 {
|
||||
continue
|
||||
if f.conditional() {
|
||||
var (
|
||||
flagValue uint32
|
||||
found bool
|
||||
)
|
||||
for j := 0; j < flagCount; j++ {
|
||||
if flags[j].name == f.flagName {
|
||||
flagValue = flags[j].value
|
||||
found = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !found {
|
||||
return malformedf("constructor %s conditional field %s references missing flags word %s", cl.name, f.name, f.flagName)
|
||||
}
|
||||
if flagValue&(1<<uint(f.flagBit)) == 0 {
|
||||
continue
|
||||
}
|
||||
}
|
||||
if err := m.skipValue(b, f); err != nil {
|
||||
if err := s.skipValue(m, b, f, cl, depth); err != nil {
|
||||
return fmt.Errorf("%s.%s: %w", cl.name, f.name, err)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// skipValue advances b past one (already known-present) field value.
|
||||
func (m *schemaModel) skipValue(b *bin.Buffer, f *fieldLayout) error {
|
||||
// skipValue advances b past one already-known-present field value.
|
||||
func (s *walkState) skipValue(m *schemaModel, b *bin.Buffer, f *fieldLayout, owner *ctorLayout, depth int) error {
|
||||
switch f.kind {
|
||||
case kindInt:
|
||||
_, err := b.Int()
|
||||
return err
|
||||
case kindLong:
|
||||
_, err := b.Long()
|
||||
return err
|
||||
case kindDouble:
|
||||
_, err := b.Double()
|
||||
return err
|
||||
return skipFixed(b, 4)
|
||||
case kindLong, kindDouble:
|
||||
return skipFixed(b, 8)
|
||||
case kindInt128:
|
||||
_, err := b.Int128()
|
||||
return err
|
||||
return skipFixed(b, 16)
|
||||
case kindInt256:
|
||||
_, err := b.Int256()
|
||||
return err
|
||||
return skipFixed(b, 32)
|
||||
case kindBytes:
|
||||
_, err := b.Bytes()
|
||||
return err
|
||||
return s.skipTLBytes(b, "bytes")
|
||||
case kindString:
|
||||
_, err := b.String()
|
||||
return err
|
||||
return s.skipTLBytes(b, "string")
|
||||
case kindBool:
|
||||
_, err := b.Bool()
|
||||
return err
|
||||
if err := s.addUnits(1, "Bool constructor"); err != nil {
|
||||
return err
|
||||
}
|
||||
id, err := b.Uint32()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if id != bin.TypeTrue && id != bin.TypeFalse {
|
||||
return malformedf("invalid Bool constructor %#08x", id)
|
||||
}
|
||||
return nil
|
||||
case kindTrue:
|
||||
return nil
|
||||
case kindVector, kindVectorBare:
|
||||
vectorDepth := depth + 1
|
||||
if vectorDepth <= 0 || vectorDepth > s.limits.maxDepth {
|
||||
return limitf("vector nesting depth %d exceeds limit %d", vectorDepth, s.limits.maxDepth)
|
||||
}
|
||||
if f.kind == kindVector {
|
||||
id, err := b.Uint32()
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if id != vectorTypeID {
|
||||
return fmt.Errorf("expected vector id, got %#08x", id)
|
||||
return malformedf("expected vector id, got %#08x", id)
|
||||
}
|
||||
}
|
||||
n, err := b.Int()
|
||||
|
|
@ -94,23 +288,146 @@ func (m *schemaModel) skipValue(b *bin.Buffer, f *fieldLayout) error {
|
|||
return err
|
||||
}
|
||||
if n < 0 {
|
||||
return fmt.Errorf("negative vector length %d", n)
|
||||
return malformedf("negative vector length %d", n)
|
||||
}
|
||||
if max := s.vectorLimit(owner, f); n > max {
|
||||
return limitf("vector %s.%s length %d exceeds limit %d", ownerName(owner), fieldName(f), n, max)
|
||||
}
|
||||
if err := s.addUnits(n, "vector "+ownerName(owner)+"."+fieldName(f)); err != nil {
|
||||
return err
|
||||
}
|
||||
if width, ok := fixedWireWidth(f.elem); ok {
|
||||
total, ok := checkedMulInt(n, width)
|
||||
if !ok {
|
||||
return malformedf("vector byte length overflow: %d * %d", n, width)
|
||||
}
|
||||
return skipFixed(b, total)
|
||||
}
|
||||
for i := 0; i < n; i++ {
|
||||
if err := m.skipValue(b, f.elem); err != nil {
|
||||
return err
|
||||
if err := s.skipValue(m, b, f.elem, nil, vectorDepth); err != nil {
|
||||
return fmt.Errorf("vector element %d: %w", i, err)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
case kindObject:
|
||||
return m.skipObject(b)
|
||||
return s.skipObject(m, b, depth+1)
|
||||
case kindBareObject:
|
||||
bareDepth := depth + 1
|
||||
if err := s.enter(bareDepth, "bare constructor"); err != nil {
|
||||
return err
|
||||
}
|
||||
cl, ok := m.bareByT[f.typeName]
|
||||
if !ok {
|
||||
return fmt.Errorf("unknown bare type %q", f.typeName)
|
||||
return malformedf("unknown bare type %q", f.typeName)
|
||||
}
|
||||
return m.skipCtorBody(b, cl)
|
||||
return s.skipCtorBody(m, b, cl, bareDepth)
|
||||
default:
|
||||
return fmt.Errorf("bad wire kind %d", f.kind)
|
||||
return malformedf("bad wire kind %d", f.kind)
|
||||
}
|
||||
}
|
||||
|
||||
// skipTLBytes parses TL's 1/4-byte length prefix directly and advances the
|
||||
// input slice. Unlike bin.Buffer.Bytes it never copies payload data.
|
||||
func (s *walkState) skipTLBytes(b *bin.Buffer, what string) error {
|
||||
if len(b.Buf) == 0 {
|
||||
return io.ErrUnexpectedEOF
|
||||
}
|
||||
var header, payload uint64
|
||||
switch b.Buf[0] {
|
||||
case 254:
|
||||
if len(b.Buf) < 4 {
|
||||
return io.ErrUnexpectedEOF
|
||||
}
|
||||
header = 4
|
||||
payload = uint64(b.Buf[1]) | uint64(b.Buf[2])<<8 | uint64(b.Buf[3])<<16
|
||||
case 255:
|
||||
return malformedf("invalid %s length prefix 255", what)
|
||||
default:
|
||||
header = 1
|
||||
payload = uint64(b.Buf[0])
|
||||
}
|
||||
if err := s.addBytes(payload, what); err != nil {
|
||||
return err
|
||||
}
|
||||
encoded, ok := checkedAddUint64(header, payload)
|
||||
if !ok {
|
||||
return malformedf("%s encoded length overflow", what)
|
||||
}
|
||||
withPadding, ok := checkedAddUint64(encoded, 3)
|
||||
if !ok {
|
||||
return malformedf("%s padded length overflow", what)
|
||||
}
|
||||
padded := withPadding &^ uint64(3)
|
||||
if padded > uint64(math.MaxInt) {
|
||||
return malformedf("%s padded length %d overflows int", what, padded)
|
||||
}
|
||||
if uint64(len(b.Buf)) < padded {
|
||||
return io.ErrUnexpectedEOF
|
||||
}
|
||||
b.Buf = b.Buf[int(padded):]
|
||||
return nil
|
||||
}
|
||||
|
||||
func skipFixed(b *bin.Buffer, n int) error {
|
||||
if n < 0 {
|
||||
return malformedf("negative fixed-width skip %d", n)
|
||||
}
|
||||
if len(b.Buf) < n {
|
||||
return io.ErrUnexpectedEOF
|
||||
}
|
||||
b.Buf = b.Buf[n:]
|
||||
return nil
|
||||
}
|
||||
|
||||
func fixedWireWidth(f *fieldLayout) (int, bool) {
|
||||
if f == nil {
|
||||
return 0, false
|
||||
}
|
||||
switch f.kind {
|
||||
case kindInt:
|
||||
return 4, true
|
||||
case kindLong, kindDouble:
|
||||
return 8, true
|
||||
case kindInt128:
|
||||
return 16, true
|
||||
case kindInt256:
|
||||
return 32, true
|
||||
case kindTrue:
|
||||
return 0, true
|
||||
default:
|
||||
// Bool deliberately stays on the element loop so constructor ids are
|
||||
// validated and charged to the aggregate constructor budget.
|
||||
return 0, false
|
||||
}
|
||||
}
|
||||
|
||||
func checkedMulInt(a, b int) (int, bool) {
|
||||
if a < 0 || b < 0 {
|
||||
return 0, false
|
||||
}
|
||||
if a != 0 && b > math.MaxInt/a {
|
||||
return 0, false
|
||||
}
|
||||
return a * b, true
|
||||
}
|
||||
|
||||
func checkedAddUint64(a, b uint64) (uint64, bool) {
|
||||
if b > math.MaxUint64-a {
|
||||
return 0, false
|
||||
}
|
||||
return a + b, true
|
||||
}
|
||||
|
||||
func ownerName(cl *ctorLayout) string {
|
||||
if cl == nil || cl.name == "" {
|
||||
return "<nested>"
|
||||
}
|
||||
return cl.name
|
||||
}
|
||||
|
||||
func fieldName(f *fieldLayout) string {
|
||||
if f == nil || f.name == "" {
|
||||
return "<element>"
|
||||
}
|
||||
return f.name
|
||||
}
|
||||
|
|
|
|||
263
internal/compat/layerwire/walk_limits_test.go
Normal file
263
internal/compat/layerwire/walk_limits_test.go
Normal file
|
|
@ -0,0 +1,263 @@
|
|||
package layerwire
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"math"
|
||||
"testing"
|
||||
|
||||
"github.com/gotd/td/bin"
|
||||
"github.com/gotd/td/tg"
|
||||
)
|
||||
|
||||
func TestValidateCanonicalRequestFlaggedHotPathAllocatesNothing(t *testing.T) {
|
||||
var body bin.Buffer
|
||||
req := &tg.MessagesSendMessageRequest{
|
||||
Peer: &tg.InputPeerSelf{},
|
||||
Message: "hello",
|
||||
RandomID: 7,
|
||||
}
|
||||
if err := req.Encode(&body); err != nil {
|
||||
t.Fatalf("encode request: %v", err)
|
||||
}
|
||||
if err := ValidateCanonicalRequest(body.Buf); err != nil {
|
||||
t.Fatalf("validate request: %v", err)
|
||||
}
|
||||
if allocs := testing.AllocsPerRun(1000, func() {
|
||||
if err := ValidateCanonicalRequest(body.Buf); err != nil {
|
||||
panic(err)
|
||||
}
|
||||
}); allocs != 0 {
|
||||
t.Fatalf("canonical request preflight allocations = %.2f, want 0", allocs)
|
||||
}
|
||||
}
|
||||
|
||||
func TestValidateCanonicalRequestVectorLimits(t *testing.T) {
|
||||
editCloseFriends := canonical.byName["contacts.editCloseFriends"]
|
||||
if editCloseFriends == nil {
|
||||
t.Fatal("contacts.editCloseFriends missing from canonical schema")
|
||||
}
|
||||
|
||||
t.Run("explicit_5000_override", func(t *testing.T) {
|
||||
var body bin.Buffer
|
||||
body.PutID(editCloseFriends.crc)
|
||||
body.PutVectorHeader(5000)
|
||||
for i := 0; i < 5000; i++ {
|
||||
body.PutLong(int64(i))
|
||||
}
|
||||
if err := ValidateCanonicalRequest(body.Buf); err != nil {
|
||||
t.Fatalf("validate legal 5000-element close-friends request: %v", err)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("override_stops_at_5000", func(t *testing.T) {
|
||||
var body bin.Buffer
|
||||
body.PutID(editCloseFriends.crc)
|
||||
body.PutVectorHeader(5001)
|
||||
err := ValidateCanonicalRequest(body.Buf)
|
||||
if !errors.Is(err, ErrResourceLimit) {
|
||||
t.Fatalf("error = %v, want ErrResourceLimit", err)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("default_4096", func(t *testing.T) {
|
||||
getMessages := canonical.byName["messages.getMessages"]
|
||||
var body bin.Buffer
|
||||
body.PutID(getMessages.crc)
|
||||
body.PutVectorHeader(defaultMaxVectorElements + 1)
|
||||
err := ValidateCanonicalRequest(body.Buf)
|
||||
if !errors.Is(err, ErrResourceLimit) {
|
||||
t.Fatalf("error = %v, want ErrResourceLimit", err)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("max_int32_count_rejected_before_iteration", func(t *testing.T) {
|
||||
var body bin.Buffer
|
||||
body.PutID(editCloseFriends.crc)
|
||||
body.PutID(vectorTypeID)
|
||||
body.PutInt32(math.MaxInt32)
|
||||
err := ValidateCanonicalRequest(body.Buf)
|
||||
if !errors.Is(err, ErrResourceLimit) {
|
||||
t.Fatalf("error = %v, want ErrResourceLimit", err)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestValidateCanonicalRequestDepthLimit(t *testing.T) {
|
||||
invoke := canonical.byName["invokeWithoutUpdates"]
|
||||
leaf := canonical.byName["help.getConfig"]
|
||||
if invoke == nil || leaf == nil {
|
||||
t.Fatal("generic wrapper methods missing from canonical schema")
|
||||
}
|
||||
request := func(wrappers int) []byte {
|
||||
var body bin.Buffer
|
||||
for i := 0; i < wrappers; i++ {
|
||||
body.PutID(invoke.crc)
|
||||
}
|
||||
body.PutID(leaf.crc)
|
||||
return body.Buf
|
||||
}
|
||||
if err := ValidateCanonicalRequest(request(defaultMaxWalkDepth - 1)); err != nil {
|
||||
t.Fatalf("depth exactly %d rejected: %v", defaultMaxWalkDepth, err)
|
||||
}
|
||||
err := ValidateCanonicalRequest(request(defaultMaxWalkDepth))
|
||||
if !errors.Is(err, ErrResourceLimit) {
|
||||
t.Fatalf("depth %d error = %v, want ErrResourceLimit", defaultMaxWalkDepth+1, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTLBytesSkipIsZeroCopyAndBounded(t *testing.T) {
|
||||
var encoded bin.Buffer
|
||||
encoded.PutBytes([]byte("payload"))
|
||||
fieldLen := len(encoded.Buf)
|
||||
raw := append(encoded.Copy(), 0xaa, 0xbb, 0xcc, 0xdd)
|
||||
b := &bin.Buffer{Buf: raw}
|
||||
walk := newWalkState()
|
||||
if err := walk.skipTLBytes(b, "bytes"); err != nil {
|
||||
t.Fatalf("skip bytes: %v", err)
|
||||
}
|
||||
if len(b.Buf) != 4 || &b.Buf[0] != &raw[fieldLen] {
|
||||
t.Fatalf("walker did not retain the original backing buffer")
|
||||
}
|
||||
|
||||
t.Run("per_field_budget", func(t *testing.T) {
|
||||
limited := newWalkState()
|
||||
limited.limits.maxFieldBytes = 3
|
||||
probe := &bin.Buffer{Buf: encoded.Copy()}
|
||||
err := limited.skipTLBytes(probe, "bytes")
|
||||
if !errors.Is(err, ErrResourceLimit) {
|
||||
t.Fatalf("error = %v, want ErrResourceLimit", err)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("aggregate_budget", func(t *testing.T) {
|
||||
limited := newWalkState()
|
||||
limited.limits.maxTotalBytes = 10
|
||||
first := &bin.Buffer{Buf: encoded.Copy()}
|
||||
if err := limited.skipTLBytes(first, "bytes"); err != nil {
|
||||
t.Fatalf("first field: %v", err)
|
||||
}
|
||||
second := &bin.Buffer{Buf: encoded.Copy()}
|
||||
err := limited.skipTLBytes(second, "bytes")
|
||||
if !errors.Is(err, ErrResourceLimit) {
|
||||
t.Fatalf("second error = %v, want ErrResourceLimit", err)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("truncated_payload_is_malformed", func(t *testing.T) {
|
||||
importAuth := canonical.byName["auth.importAuthorization"]
|
||||
var body bin.Buffer
|
||||
body.PutID(importAuth.crc)
|
||||
body.PutLong(1)
|
||||
body.Put([]byte{5, 'a', 'b'}) // declares five bytes, lacks payload/padding
|
||||
err := ValidateCanonicalRequest(body.Buf)
|
||||
if !errors.Is(err, ErrMalformed) || errors.Is(err, ErrResourceLimit) {
|
||||
t.Fatalf("error = %v, want only ErrMalformed", err)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestInboundTransformsShareWalkerBudgets(t *testing.T) {
|
||||
t.Run("canonical_alias", func(t *testing.T) {
|
||||
var body bin.Buffer
|
||||
body.PutID(0x41d41ade) // DrKLO messages.forwardMessages alias
|
||||
body.PutUint32(0)
|
||||
body.PutID(canonical.byName["inputPeerEmpty"].crc)
|
||||
body.PutID(vectorTypeID)
|
||||
body.PutInt32(math.MaxInt32)
|
||||
_, ok, err := UpgradeInbound(0x41d41ade, &body)
|
||||
if !ok || !errors.Is(err, ErrResourceLimit) {
|
||||
t.Fatalf("ok=%v error=%v, want matched ErrResourceLimit", ok, err)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("drift_body_transform", func(t *testing.T) {
|
||||
var body bin.Buffer
|
||||
body.PutID(0x2e1ee318) // DrKLO langpack.getStrings body transform
|
||||
body.PutString("en")
|
||||
body.PutID(vectorTypeID)
|
||||
body.PutInt32(math.MaxInt32)
|
||||
_, ok, err := UpgradeInbound(0x2e1ee318, &body)
|
||||
if !ok || !errors.Is(err, ErrResourceLimit) {
|
||||
t.Fatalf("ok=%v error=%v, want matched ErrResourceLimit", ok, err)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("outbound_structural_transform", func(t *testing.T) {
|
||||
poll := canonical.byName["pollAnswerVoters"]
|
||||
var body bin.Buffer
|
||||
body.PutID(poll.crc)
|
||||
body.PutUint32(1 << 2)
|
||||
body.PutBytes(nil)
|
||||
body.PutInt(1)
|
||||
body.PutID(vectorTypeID)
|
||||
body.PutInt32(math.MaxInt32)
|
||||
_, err := Transcode(body.Buf, CanonicalLayer-1)
|
||||
if !errors.Is(err, ErrResourceLimit) {
|
||||
t.Fatalf("error = %v, want ErrResourceLimit", err)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestWalkerArithmeticAndMalformedClassification(t *testing.T) {
|
||||
if defaultMaxWalkUnits != 131072 {
|
||||
t.Fatalf("default constructor/vector budget = %d, want 131072", defaultMaxWalkUnits)
|
||||
}
|
||||
if _, ok := checkedMulInt(math.MaxInt, 2); ok {
|
||||
t.Fatal("checkedMulInt accepted overflow")
|
||||
}
|
||||
if _, ok := checkedAddUint64(math.MaxUint64, 1); ok {
|
||||
t.Fatal("checkedAddUint64 accepted overflow")
|
||||
}
|
||||
|
||||
t.Run("aggregate_constructor_and_vector_units", func(t *testing.T) {
|
||||
editCloseFriends := canonical.byName["contacts.editCloseFriends"]
|
||||
var body bin.Buffer
|
||||
body.PutID(editCloseFriends.crc)
|
||||
body.PutVectorHeader(4)
|
||||
for i := 0; i < 4; i++ {
|
||||
body.PutLong(int64(i))
|
||||
}
|
||||
walk := newWalkState()
|
||||
walk.limits.maxUnits = 4 // top constructor + four elements needs five
|
||||
probe := &bin.Buffer{Buf: body.Buf}
|
||||
err := walk.skipObject(canonical, probe, 1)
|
||||
if !errors.Is(err, ErrResourceLimit) {
|
||||
t.Fatalf("error = %v, want ErrResourceLimit", err)
|
||||
}
|
||||
})
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
body []byte
|
||||
}{
|
||||
{name: "empty"},
|
||||
{name: "unknown_constructor", body: []byte{1, 2, 3, 4}},
|
||||
{name: "trailing_bytes", body: append(methodIDBytes(canonical.byName["help.getConfig"].crc), 0, 0, 0, 0)},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
err := ValidateCanonicalRequest(tt.body)
|
||||
if !errors.Is(err, ErrMalformed) || errors.Is(err, ErrResourceLimit) {
|
||||
t.Fatalf("error = %v, want only ErrMalformed", err)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func methodIDBytes(id uint32) []byte {
|
||||
var b bin.Buffer
|
||||
b.PutID(id)
|
||||
return b.Buf
|
||||
}
|
||||
|
||||
func FuzzValidateCanonicalRequest(f *testing.F) {
|
||||
f.Add(methodIDBytes(canonical.byName["help.getConfig"].crc))
|
||||
f.Add([]byte{})
|
||||
f.Add([]byte{1, 2, 3, 4})
|
||||
f.Fuzz(func(t *testing.T, body []byte) {
|
||||
err := ValidateCanonicalRequest(body)
|
||||
if err != nil && !errors.Is(err, ErrMalformed) && !errors.Is(err, ErrResourceLimit) {
|
||||
t.Fatalf("unclassified walker error: %v", err)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
|
@ -29,6 +29,28 @@ type Config struct {
|
|||
RSAKeyPath string
|
||||
// DC 是本 server 的 DC ID。
|
||||
DC int
|
||||
// MTProtoMaxConnections / PerIP 覆盖 raw Accept、codec sniff、握手到认证 session
|
||||
// 的完整物理连接生命周期;负数关闭对应 admission 上限。
|
||||
MTProtoMaxConnections int
|
||||
MTProtoMaxConnectionsPerIP int
|
||||
// MTProtoMaxConcurrentHandshakes 限制昂贵 RSA/DH exchange 并发;负数关闭。
|
||||
MTProtoMaxConcurrentHandshakes int
|
||||
// MTProto RPC 使用 Server 共享公平调度器;per-connection 与 global 预算共同限制
|
||||
// goroutine、排队任务和 request body 内存。
|
||||
MTProtoRPCMaxInflight int
|
||||
MTProtoRPCQueueSize int
|
||||
MTProtoRPCTimeout time.Duration
|
||||
MTProtoRPCGlobalWorkers int
|
||||
MTProtoRPCGlobalMaxTasks int
|
||||
MTProtoRPCGlobalMaxBytes int64
|
||||
// MTProtoInboundFrameGlobalMaxBytes 是 transport wire + 最大解密 plaintext 的
|
||||
// 进程级在途预算;frame 长度读出后、payload 分配前预留。
|
||||
MTProtoInboundFrameGlobalMaxBytes int64
|
||||
// MTProto outbound mailbox 按连接有界;resend pending body 另受 Server 全局预算约束。
|
||||
MTProtoOutboundQueueSize int
|
||||
MTProtoOutboundControlQueueSize int
|
||||
MTProtoOutboundTrackedGlobalMaxBytes int64
|
||||
MTProtoOutboundWriteGlobalMaxBytes int64
|
||||
|
||||
// DebugAddr 是 net/http/pprof 调试端点监听地址(CPU/heap/goroutine/mutex/block 剖析)。
|
||||
// telesrv 是宿主进程、不在 docker 内,docker stats 看不到它,性能定位主要靠此端点。
|
||||
|
|
@ -76,6 +98,12 @@ type Config struct {
|
|||
// AuthCodeMaxAttempts 是同一 phone_code_hash / email verification code 的最大错误次数。
|
||||
// 达到上限后验证码立即失效,用户必须重发。
|
||||
AuthCodeMaxAttempts int
|
||||
// AuthCodePhoneRateLimit / AuthCodeAuthKeyRateLimit 对未授权验证码签发按规范化手机号摘要
|
||||
// 与连接实际 raw auth_key 分别限流。两个维度共用 AuthCodeRateWindow;<=0 关闭对应维度。
|
||||
// 手机号只以 SHA-256 摘要进入限流 key,禁止把原文写入 Redis key 或日志。
|
||||
AuthCodePhoneRateLimit int
|
||||
AuthCodeAuthKeyRateLimit int
|
||||
AuthCodeRateWindow time.Duration
|
||||
// LoginEmailEnable 启用手机号登录流程中的邮箱验证码投递。
|
||||
LoginEmailEnable bool
|
||||
// LoginEmailRequireSetup 为 true 时,没有登录邮箱的账号/新手机号会要求先设置邮箱。
|
||||
|
|
@ -136,6 +164,9 @@ type Config struct {
|
|||
AIPrivacyLogContent bool
|
||||
// TempKeyResolveCacheMaxEntries 是 Router temp→perm 解析缓存容量。
|
||||
TempKeyResolveCacheMaxEntries int
|
||||
// TempKeyResolveCacheTTL 是 temp→perm 绑定的进程内复核周期。绑定/revoke 有精确
|
||||
// 失效,TTL 作为跨进程或异常路径兜底;默认 30m 避免大连接数下每 5s 全量打 PG。
|
||||
TempKeyResolveCacheTTL time.Duration
|
||||
|
||||
// ChannelRowCacheMaxEntries 是「共享频道行」进程内缓存容量(channelID→domain.Channel)。
|
||||
// 由 channels 表 LISTEN/NOTIFY 触发器实时失效(强一致、零 TTL)。<=0 禁用缓存与监听。
|
||||
|
|
@ -153,8 +184,8 @@ type Config struct {
|
|||
// ChannelBoostCacheTTL 是 boost 读投影在未收到写侧通知时的最大陈旧窗口。
|
||||
ChannelBoostCacheTTL time.Duration
|
||||
|
||||
// OutboxWorkers 是并发 claim 的 outbox worker 数。默认 1,保证同一用户 pts update
|
||||
// 在线投递顺序与持久化顺序一致;后续需要吞吐时应改成按 target_user_id 分片的串行 worker。
|
||||
// OutboxWorkers 是并发 outbox worker 数。用户先稳定哈希到固定 logical shard,
|
||||
// 每个 shard 只归一个 worker,故提高 worker 数不会破坏同一用户 pts 顺序。
|
||||
OutboxWorkers int
|
||||
// OutboxBatch 是 transactional outbox worker 每次 claim 的最大条数。
|
||||
// 调大提升吞吐、增大单批 PG/推送压力;调小降低延迟抖动。配套压测见 docs/message-module.md。
|
||||
|
|
@ -164,6 +195,13 @@ type Config struct {
|
|||
// OutboxLeaseTimeout 是 'dispatching' 行被判定为租约过期、允许其它 worker 重新 claim 的时长。
|
||||
// 取值需大于单批投递耗时,否则会重复推送;过大则 worker 崩溃后积压恢复变慢。
|
||||
OutboxLeaseTimeout time.Duration
|
||||
// OutboxPoisonRetention 是 terminal failed outbox head 的隔离窗口。隔离期内保留
|
||||
// last_error 供排障,期满只删除在线投递任务;durable user_update_events 仍保留,
|
||||
// 客户端可经 updates.getDifference 恢复。
|
||||
OutboxPoisonRetention time.Duration
|
||||
// OutboxPoisonCleanupInterval 独立于大表 retention 周期清理 terminal failed head,
|
||||
// 避免一条确定性坏事件长期冻结同账号更高 pts 的在线投递 lane。
|
||||
OutboxPoisonCleanupInterval time.Duration
|
||||
// OutboundPushTimeout 是 best-effort updates 推送等待 outbound 队列接受的最长时间。
|
||||
OutboundPushTimeout time.Duration
|
||||
// SendRateLimit 是账号级发送窗口内允许的消息条数;<=0 表示关闭发送限流。
|
||||
|
|
@ -182,6 +220,9 @@ type Config struct {
|
|||
// BotAPIUpdateRetention 是 bot_api_updates 投递队列的最大保留期(官方 Bot API 语义 24h);
|
||||
// 已确认的行另按固定短宽限提前回收(性能审计 H1)。
|
||||
BotAPIUpdateRetention time.Duration
|
||||
// OrphanAuthKeyRetention 是握手已创建、但没有 authorization/temp binding/活跃连接的
|
||||
// auth key 最短保留期。过期后由有界 GC 回收;客户端收到 -404 会重建 key。
|
||||
OrphanAuthKeyRetention time.Duration
|
||||
// RetentionInterval 是 retention worker 的运行间隔。
|
||||
RetentionInterval time.Duration
|
||||
// RetentionBatch 是单次 retention 最多删除的行数。
|
||||
|
|
@ -323,19 +364,33 @@ func Load() (Config, error) {
|
|||
// AdvertiseIP 当前不影响 help.getConfig——getConfig 返回空 DCOptions,
|
||||
// 客户端使用其写死的 static DC 地址(见 compat/tdesktop/config.go)。
|
||||
// 字段与默认值保留,供未来需要显式下发 DC 地址时使用。
|
||||
AdvertiseIP: envOr("TELESRV_ADVERTISE_IP", "127.0.0.1"),
|
||||
RSAKeyPath: envOr("TELESRV_RSA_KEY", "data/server_rsa.pem"),
|
||||
DC: envIntOr("TELESRV_DC", 2),
|
||||
DebugAddr: envAllowEmptyOr("TELESRV_DEBUG_ADDR", "127.0.0.1:6060"),
|
||||
BotAPIAddr: envAllowEmptyOr("TELESRV_BOT_API_ADDR", ""),
|
||||
AdminAPIAddr: envAllowEmptyOr("TELESRV_ADMIN_API_ADDR", ""),
|
||||
AdminAPIToken: envOr("TELESRV_ADMIN_API_TOKEN", ""),
|
||||
PublicBaseURL: publicBaseURL,
|
||||
PublicLinkWebAddr: envAllowEmptyOr("TELESRV_PUBLIC_LINK_WEB_ADDR", ""),
|
||||
AdminUIAddr: envOr("TELESRV_ADMIN_UI_ADDR", "127.0.0.1:2600"),
|
||||
AdminUIPassword: envOr("TELESRV_ADMIN_UI_PASSWORD", ""),
|
||||
AdminUIToken: envOr("TELESRV_ADMIN_UI_TOKEN", ""),
|
||||
AdminSessionKey: envOr("TELESRV_ADMIN_SESSION_KEY", ""),
|
||||
AdvertiseIP: envOr("TELESRV_ADVERTISE_IP", "127.0.0.1"),
|
||||
RSAKeyPath: envOr("TELESRV_RSA_KEY", "data/server_rsa.pem"),
|
||||
DC: envIntOr("TELESRV_DC", 2),
|
||||
MTProtoMaxConnections: envIntOr("TELESRV_MTPROTO_MAX_CONNECTIONS", 200000),
|
||||
MTProtoMaxConnectionsPerIP: envIntOr("TELESRV_MTPROTO_MAX_CONNECTIONS_PER_IP", 4096),
|
||||
MTProtoMaxConcurrentHandshakes: envIntOr("TELESRV_MTPROTO_MAX_CONCURRENT_HANDSHAKES", 256),
|
||||
MTProtoRPCMaxInflight: envIntOr("TELESRV_MTPROTO_RPC_MAX_INFLIGHT", 32),
|
||||
MTProtoRPCQueueSize: envIntOr("TELESRV_MTPROTO_RPC_QUEUE_SIZE", 64),
|
||||
MTProtoRPCTimeout: envDurationOr("TELESRV_MTPROTO_RPC_TIMEOUT", 30*time.Second),
|
||||
MTProtoRPCGlobalWorkers: envIntOr("TELESRV_MTPROTO_RPC_GLOBAL_WORKERS", 256),
|
||||
MTProtoRPCGlobalMaxTasks: envIntOr("TELESRV_MTPROTO_RPC_GLOBAL_MAX_TASKS", 8192),
|
||||
MTProtoRPCGlobalMaxBytes: envInt64Or("TELESRV_MTPROTO_RPC_GLOBAL_MAX_BYTES", 512<<20),
|
||||
MTProtoInboundFrameGlobalMaxBytes: envInt64Or("TELESRV_MTPROTO_INBOUND_FRAME_GLOBAL_MAX_BYTES", 512<<20),
|
||||
MTProtoOutboundQueueSize: envIntOr("TELESRV_MTPROTO_OUTBOUND_QUEUE_SIZE", 128),
|
||||
MTProtoOutboundControlQueueSize: envIntOr("TELESRV_MTPROTO_OUTBOUND_CONTROL_QUEUE_SIZE", 32),
|
||||
MTProtoOutboundTrackedGlobalMaxBytes: envInt64Or("TELESRV_MTPROTO_OUTBOUND_TRACKED_GLOBAL_MAX_BYTES", 512<<20),
|
||||
MTProtoOutboundWriteGlobalMaxBytes: envInt64Or("TELESRV_MTPROTO_OUTBOUND_WRITE_GLOBAL_MAX_BYTES", 512<<20),
|
||||
DebugAddr: envAllowEmptyOr("TELESRV_DEBUG_ADDR", "127.0.0.1:6060"),
|
||||
BotAPIAddr: envAllowEmptyOr("TELESRV_BOT_API_ADDR", ""),
|
||||
AdminAPIAddr: envAllowEmptyOr("TELESRV_ADMIN_API_ADDR", ""),
|
||||
AdminAPIToken: envOr("TELESRV_ADMIN_API_TOKEN", ""),
|
||||
PublicBaseURL: publicBaseURL,
|
||||
PublicLinkWebAddr: envAllowEmptyOr("TELESRV_PUBLIC_LINK_WEB_ADDR", ""),
|
||||
AdminUIAddr: envOr("TELESRV_ADMIN_UI_ADDR", "127.0.0.1:2600"),
|
||||
AdminUIPassword: envOr("TELESRV_ADMIN_UI_PASSWORD", ""),
|
||||
AdminUIToken: envOr("TELESRV_ADMIN_UI_TOKEN", ""),
|
||||
AdminSessionKey: envOr("TELESRV_ADMIN_SESSION_KEY", ""),
|
||||
|
||||
// 用 127.0.0.1 而非 localhost:localhost 在 Windows 上会先解析到 IPv6 ::1,而 Docker
|
||||
// Desktop 的端口转发只在 IPv4 监听,IPv6 连接要等 ~1s 超时才回退 IPv4(实测 localhost
|
||||
|
|
@ -351,6 +406,9 @@ func Load() (Config, error) {
|
|||
DevAuthCode: envOr("TELESRV_DEV_AUTH_CODE", "12345"),
|
||||
AuthCodeTTL: envDurationOr("TELESRV_AUTH_CODE_TTL", 5*time.Minute),
|
||||
AuthCodeMaxAttempts: envIntOr("TELESRV_AUTH_CODE_MAX_ATTEMPTS", 5),
|
||||
AuthCodePhoneRateLimit: envIntOr("TELESRV_AUTH_CODE_PHONE_RATE_LIMIT", 5),
|
||||
AuthCodeAuthKeyRateLimit: envIntOr("TELESRV_AUTH_CODE_AUTH_KEY_RATE_LIMIT", 20),
|
||||
AuthCodeRateWindow: envDurationOr("TELESRV_AUTH_CODE_RATE_WINDOW", 10*time.Minute),
|
||||
LoginEmailEnable: envBoolOr("TELESRV_LOGIN_EMAIL_ENABLE", false),
|
||||
LoginEmailRequireSetup: envBoolOr("TELESRV_LOGIN_EMAIL_REQUIRE_SETUP", false),
|
||||
LoginEmailCodeLength: envIntOr("TELESRV_LOGIN_EMAIL_CODE_LENGTH", 6),
|
||||
|
|
@ -381,17 +439,22 @@ func Load() (Config, error) {
|
|||
AIRateLimit: envIntOr("TELESRV_AI_RATE_LIMIT", 20),
|
||||
AIRateWindow: envDurationOr("TELESRV_AI_RATE_WINDOW", time.Minute),
|
||||
AIPrivacyLogContent: envBoolOr("TELESRV_AI_LOG_CONTENT", false),
|
||||
TempKeyResolveCacheMaxEntries: envIntOr("TELESRV_TEMP_KEY_CACHE_MAX_ENTRIES", 4096),
|
||||
TempKeyResolveCacheMaxEntries: envIntOr("TELESRV_TEMP_KEY_CACHE_MAX_ENTRIES", 262144),
|
||||
TempKeyResolveCacheTTL: envDurationOr("TELESRV_TEMP_KEY_CACHE_TTL", 30*time.Minute),
|
||||
ChannelRowCacheMaxEntries: envIntOr("TELESRV_CHANNEL_ROW_CACHE_MAX", 50000),
|
||||
ChannelMemberCacheMaxEntries: envIntOr("TELESRV_CHANNEL_MEMBER_CACHE_MAX", 100000),
|
||||
ChannelDialogCacheMaxEntries: envIntOr("TELESRV_CHANNEL_DIALOG_CACHE_MAX", 100000),
|
||||
ChannelBoostCacheMaxEntries: envIntOr("TELESRV_CHANNEL_BOOST_CACHE_MAX", 100000),
|
||||
ChannelBoostCacheTTL: envDurationOr("TELESRV_CHANNEL_BOOST_CACHE_TTL", 10*time.Second),
|
||||
|
||||
OutboxWorkers: envIntOr("TELESRV_OUTBOX_WORKERS", 1),
|
||||
OutboxBatch: envIntOr("TELESRV_OUTBOX_BATCH", 100),
|
||||
OutboxInterval: envDurationOr("TELESRV_OUTBOX_INTERVAL", 200*time.Millisecond),
|
||||
OutboxLeaseTimeout: envDurationOr("TELESRV_OUTBOX_LEASE_TIMEOUT", 30*time.Second),
|
||||
OutboxWorkers: envIntOr("TELESRV_OUTBOX_WORKERS", 4),
|
||||
OutboxBatch: envIntOr("TELESRV_OUTBOX_BATCH", 100),
|
||||
OutboxInterval: envDurationOr("TELESRV_OUTBOX_INTERVAL", 200*time.Millisecond),
|
||||
OutboxLeaseTimeout: envDurationOr("TELESRV_OUTBOX_LEASE_TIMEOUT", 30*time.Second),
|
||||
OutboxPoisonRetention: envDurationOr("TELESRV_OUTBOX_POISON_RETENTION", time.Minute),
|
||||
OutboxPoisonCleanupInterval: envDurationOr(
|
||||
"TELESRV_OUTBOX_POISON_CLEANUP_INTERVAL", 15*time.Second,
|
||||
),
|
||||
OutboundPushTimeout: envDurationOr("TELESRV_OUTBOUND_PUSH_TIMEOUT", 200*time.Millisecond),
|
||||
SendRateLimit: envIntOr("TELESRV_SEND_RATE_LIMIT", 30),
|
||||
SendRateWindow: envDurationOr("TELESRV_SEND_RATE_WINDOW", time.Minute),
|
||||
|
|
@ -400,6 +463,7 @@ func Load() (Config, error) {
|
|||
ChannelNudgeMaxTargets: envIntOr("TELESRV_CHANNEL_NUDGE_MAX_TARGETS", 0),
|
||||
UpdateEventRetention: envDurationOr("TELESRV_UPDATE_EVENT_RETENTION", 168*time.Hour),
|
||||
BotAPIUpdateRetention: envDurationOr("TELESRV_BOT_API_UPDATE_RETENTION", 24*time.Hour),
|
||||
OrphanAuthKeyRetention: envDurationOr("TELESRV_ORPHAN_AUTH_KEY_RETENTION", 24*time.Hour),
|
||||
RetentionInterval: envDurationOr("TELESRV_RETENTION_INTERVAL", time.Hour),
|
||||
RetentionBatch: envIntOr("TELESRV_RETENTION_BATCH", 10000),
|
||||
UploadPartTTL: envDurationOr("TELESRV_UPLOAD_PART_TTL", 24*time.Hour),
|
||||
|
|
|
|||
|
|
@ -37,6 +37,62 @@ func TestLoadUsesExplicitAdvertiseIP(t *testing.T) {
|
|||
}
|
||||
}
|
||||
|
||||
func TestLoadMTProtoAdmissionAndRPCBudgets(t *testing.T) {
|
||||
disableDefaultConfigFile(t)
|
||||
t.Setenv("TELESRV_MTPROTO_MAX_CONNECTIONS", "12345")
|
||||
t.Setenv("TELESRV_MTPROTO_MAX_CONNECTIONS_PER_IP", "234")
|
||||
t.Setenv("TELESRV_MTPROTO_MAX_CONCURRENT_HANDSHAKES", "45")
|
||||
t.Setenv("TELESRV_MTPROTO_RPC_MAX_INFLIGHT", "7")
|
||||
t.Setenv("TELESRV_MTPROTO_RPC_QUEUE_SIZE", "19")
|
||||
t.Setenv("TELESRV_MTPROTO_RPC_TIMEOUT", "9s")
|
||||
t.Setenv("TELESRV_MTPROTO_RPC_GLOBAL_WORKERS", "33")
|
||||
t.Setenv("TELESRV_MTPROTO_RPC_GLOBAL_MAX_TASKS", "444")
|
||||
t.Setenv("TELESRV_MTPROTO_RPC_GLOBAL_MAX_BYTES", "555555")
|
||||
t.Setenv("TELESRV_MTPROTO_INBOUND_FRAME_GLOBAL_MAX_BYTES", "777777")
|
||||
t.Setenv("TELESRV_MTPROTO_OUTBOUND_QUEUE_SIZE", "88")
|
||||
t.Setenv("TELESRV_MTPROTO_OUTBOUND_CONTROL_QUEUE_SIZE", "22")
|
||||
t.Setenv("TELESRV_MTPROTO_OUTBOUND_TRACKED_GLOBAL_MAX_BYTES", "888888")
|
||||
t.Setenv("TELESRV_MTPROTO_OUTBOUND_WRITE_GLOBAL_MAX_BYTES", "999999")
|
||||
t.Setenv("TELESRV_TEMP_KEY_CACHE_MAX_ENTRIES", "666")
|
||||
t.Setenv("TELESRV_TEMP_KEY_CACHE_TTL", "17m")
|
||||
t.Setenv("TELESRV_ORPHAN_AUTH_KEY_RETENTION", "36h")
|
||||
|
||||
cfg, err := Load()
|
||||
if err != nil {
|
||||
t.Fatalf("Load: %v", err)
|
||||
}
|
||||
if cfg.MTProtoMaxConnections != 12345 || cfg.MTProtoMaxConnectionsPerIP != 234 || cfg.MTProtoMaxConcurrentHandshakes != 45 {
|
||||
t.Fatalf("admission config = %d/%d/%d", cfg.MTProtoMaxConnections, cfg.MTProtoMaxConnectionsPerIP, cfg.MTProtoMaxConcurrentHandshakes)
|
||||
}
|
||||
if cfg.MTProtoRPCMaxInflight != 7 || cfg.MTProtoRPCQueueSize != 19 || cfg.MTProtoRPCTimeout != 9*time.Second ||
|
||||
cfg.MTProtoRPCGlobalWorkers != 33 || cfg.MTProtoRPCGlobalMaxTasks != 444 || cfg.MTProtoRPCGlobalMaxBytes != 555555 {
|
||||
t.Fatalf("rpc budget config = %d/%d/%v/%d/%d/%d", cfg.MTProtoRPCMaxInflight, cfg.MTProtoRPCQueueSize, cfg.MTProtoRPCTimeout, cfg.MTProtoRPCGlobalWorkers, cfg.MTProtoRPCGlobalMaxTasks, cfg.MTProtoRPCGlobalMaxBytes)
|
||||
}
|
||||
if cfg.MTProtoInboundFrameGlobalMaxBytes != 777777 {
|
||||
t.Fatalf("inbound frame budget config = %d", cfg.MTProtoInboundFrameGlobalMaxBytes)
|
||||
}
|
||||
if cfg.MTProtoOutboundQueueSize != 88 || cfg.MTProtoOutboundControlQueueSize != 22 || cfg.MTProtoOutboundTrackedGlobalMaxBytes != 888888 || cfg.MTProtoOutboundWriteGlobalMaxBytes != 999999 {
|
||||
t.Fatalf("outbound config = %d/%d/%d/%d", cfg.MTProtoOutboundQueueSize, cfg.MTProtoOutboundControlQueueSize, cfg.MTProtoOutboundTrackedGlobalMaxBytes, cfg.MTProtoOutboundWriteGlobalMaxBytes)
|
||||
}
|
||||
if cfg.TempKeyResolveCacheMaxEntries != 666 || cfg.TempKeyResolveCacheTTL != 17*time.Minute || cfg.OrphanAuthKeyRetention != 36*time.Hour {
|
||||
t.Fatalf("auth key resource config = %d/%v/%v", cfg.TempKeyResolveCacheMaxEntries, cfg.TempKeyResolveCacheTTL, cfg.OrphanAuthKeyRetention)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoadOutboxPoisonPolicy(t *testing.T) {
|
||||
disableDefaultConfigFile(t)
|
||||
t.Setenv("TELESRV_OUTBOX_POISON_RETENTION", "2m")
|
||||
t.Setenv("TELESRV_OUTBOX_POISON_CLEANUP_INTERVAL", "7s")
|
||||
|
||||
cfg, err := Load()
|
||||
if err != nil {
|
||||
t.Fatalf("Load: %v", err)
|
||||
}
|
||||
if cfg.OutboxPoisonRetention != 2*time.Minute || cfg.OutboxPoisonCleanupInterval != 7*time.Second {
|
||||
t.Fatalf("outbox poison policy = %v/%v, want 2m/7s", cfg.OutboxPoisonRetention, cfg.OutboxPoisonCleanupInterval)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoadBusinessAIProvider(t *testing.T) {
|
||||
disableDefaultConfigFile(t)
|
||||
t.Setenv("TELESRV_BUSINESS_AI_PROVIDER", "echo")
|
||||
|
|
@ -76,8 +132,11 @@ func TestLoadLoginEmailDefaultsDisabled(t *testing.T) {
|
|||
if cfg.LoginEmailRequireSetup {
|
||||
t.Fatal("LoginEmailRequireSetup = true, want false")
|
||||
}
|
||||
if cfg.AuthCodeTTL != 5*time.Minute || cfg.AuthCodeMaxAttempts != 5 || cfg.LoginEmailCodeLength != 6 {
|
||||
t.Fatalf("auth/login email defaults = %v/%d/%d", cfg.AuthCodeTTL, cfg.AuthCodeMaxAttempts, cfg.LoginEmailCodeLength)
|
||||
if cfg.AuthCodeTTL != 5*time.Minute || cfg.AuthCodeMaxAttempts != 5 || cfg.LoginEmailCodeLength != 6 ||
|
||||
cfg.AuthCodePhoneRateLimit != 5 || cfg.AuthCodeAuthKeyRateLimit != 20 || cfg.AuthCodeRateWindow != 10*time.Minute {
|
||||
t.Fatalf("auth/login email defaults = ttl=%v attempts=%d length=%d phone_limit=%d key_limit=%d window=%v",
|
||||
cfg.AuthCodeTTL, cfg.AuthCodeMaxAttempts, cfg.LoginEmailCodeLength,
|
||||
cfg.AuthCodePhoneRateLimit, cfg.AuthCodeAuthKeyRateLimit, cfg.AuthCodeRateWindow)
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -87,6 +146,9 @@ func TestLoadLoginEmailSMTPConfig(t *testing.T) {
|
|||
t.Setenv("TELESRV_LOGIN_EMAIL_REQUIRE_SETUP", "true")
|
||||
t.Setenv("TELESRV_AUTH_CODE_TTL", "3m")
|
||||
t.Setenv("TELESRV_AUTH_CODE_MAX_ATTEMPTS", "4")
|
||||
t.Setenv("TELESRV_AUTH_CODE_PHONE_RATE_LIMIT", "3")
|
||||
t.Setenv("TELESRV_AUTH_CODE_AUTH_KEY_RATE_LIMIT", "9")
|
||||
t.Setenv("TELESRV_AUTH_CODE_RATE_WINDOW", "2m")
|
||||
t.Setenv("TELESRV_LOGIN_EMAIL_CODE_LENGTH", "7")
|
||||
t.Setenv("TELESRV_SMTP_HOST", "smtp.example.test")
|
||||
t.Setenv("TELESRV_SMTP_PORT", "2525")
|
||||
|
|
@ -103,8 +165,11 @@ func TestLoadLoginEmailSMTPConfig(t *testing.T) {
|
|||
if !cfg.LoginEmailEnable || !cfg.LoginEmailRequireSetup {
|
||||
t.Fatalf("login email flags = %v/%v, want true/true", cfg.LoginEmailEnable, cfg.LoginEmailRequireSetup)
|
||||
}
|
||||
if cfg.AuthCodeTTL != 3*time.Minute || cfg.AuthCodeMaxAttempts != 4 || cfg.LoginEmailCodeLength != 7 {
|
||||
t.Fatalf("auth/login email config = %v/%d/%d", cfg.AuthCodeTTL, cfg.AuthCodeMaxAttempts, cfg.LoginEmailCodeLength)
|
||||
if cfg.AuthCodeTTL != 3*time.Minute || cfg.AuthCodeMaxAttempts != 4 || cfg.LoginEmailCodeLength != 7 ||
|
||||
cfg.AuthCodePhoneRateLimit != 3 || cfg.AuthCodeAuthKeyRateLimit != 9 || cfg.AuthCodeRateWindow != 2*time.Minute {
|
||||
t.Fatalf("auth/login email config = ttl=%v attempts=%d length=%d phone_limit=%d key_limit=%d window=%v",
|
||||
cfg.AuthCodeTTL, cfg.AuthCodeMaxAttempts, cfg.LoginEmailCodeLength,
|
||||
cfg.AuthCodePhoneRateLimit, cfg.AuthCodeAuthKeyRateLimit, cfg.AuthCodeRateWindow)
|
||||
}
|
||||
if cfg.SMTPHost != "smtp.example.test" || cfg.SMTPPort != 2525 || cfg.SMTPUsername != "smtp-user" || cfg.SMTPPassword != "smtp-pass" || cfg.SMTPFrom != "noreply@example.test" || cfg.SMTPTLSMode != "none" || cfg.SMTPTimeout != 2*time.Second {
|
||||
t.Fatalf("smtp config = %#v", cfg)
|
||||
|
|
|
|||
51
internal/domain/album_group.go
Normal file
51
internal/domain/album_group.go
Normal file
|
|
@ -0,0 +1,51 @@
|
|||
package domain
|
||||
|
||||
import (
|
||||
"crypto/sha256"
|
||||
"errors"
|
||||
)
|
||||
|
||||
// ErrAlbumGroupReservationInvalid 表示相册分组预留缺少发送者、目标、
|
||||
// random_id 或 proposed grouped_id。RPC 边界通常会更早拦截这些输入;
|
||||
// domain/store 仍 fail-fast,避免坏绑定进入持久层。
|
||||
var ErrAlbumGroupReservationInvalid = errors.New("album group reservation invalid")
|
||||
|
||||
// AlbumGroupReservationRequest 在任何相册 item 落库或上传媒体解析前,原子地把
|
||||
// 一组 random_id 绑定到同一个 grouped_id。Peer 是幂等作用域的一部分:同一发送者
|
||||
// 可以在不同会话中复用 random_id,而不会互相污染相册分组。
|
||||
type AlbumGroupReservationRequest struct {
|
||||
SenderUserID int64
|
||||
Peer Peer
|
||||
Items []AlbumGroupReservationItem
|
||||
ProposedGroupedID int64
|
||||
}
|
||||
|
||||
// AlbumGroupReservationItem 把 random_id 与该 item 的不可变客户端意图绑定。
|
||||
// IntentHash 是在媒体解析/服务端派生字段产生前计算的 SHA-256;相同 random_id
|
||||
// 若携带不同意图必须报冲突,不能借旧 album reservation 绕过发送幂等校验。
|
||||
type AlbumGroupReservationItem struct {
|
||||
RandomID int64
|
||||
IntentHash []byte
|
||||
}
|
||||
|
||||
// Validate 校验持久层必须依赖的最小不变量。同一批内重复 random_id 与发送幂等
|
||||
// 冲突同义,必须显式失败,不能静默去重后改变客户端请求的消息条数。
|
||||
func (r AlbumGroupReservationRequest) Validate() error {
|
||||
if r.SenderUserID <= 0 || r.Peer.ID <= 0 || r.ProposedGroupedID == 0 || len(r.Items) == 0 {
|
||||
return ErrAlbumGroupReservationInvalid
|
||||
}
|
||||
if r.Peer.Type != PeerTypeUser && r.Peer.Type != PeerTypeChannel {
|
||||
return ErrAlbumGroupReservationInvalid
|
||||
}
|
||||
seen := make(map[int64]struct{}, len(r.Items))
|
||||
for _, item := range r.Items {
|
||||
if item.RandomID == 0 || len(item.IntentHash) != sha256.Size {
|
||||
return ErrAlbumGroupReservationInvalid
|
||||
}
|
||||
if _, exists := seen[item.RandomID]; exists {
|
||||
return ErrMessageRandomIDDuplicate
|
||||
}
|
||||
seen[item.RandomID] = struct{}{}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
|
@ -11,6 +11,9 @@ const (
|
|||
MaxChannelDifferenceLimit = 100
|
||||
// MaxChannelDifferenceTooLongMessages limits the latest message snapshot returned by channelDifferenceTooLong.
|
||||
MaxChannelDifferenceTooLongMessages = 100
|
||||
// MaxChannelUpdateRetentionBatch bounds one channel durable-log pruning transaction.
|
||||
// Retention advances the recoverable floor only through rows actually deleted in that transaction.
|
||||
MaxChannelUpdateRetentionBatch = 10000
|
||||
// MaxChannelParticipantsLimit limits a single participants page.
|
||||
MaxChannelParticipantsLimit = 200
|
||||
// MaxChannelParticipantsOffset bounds channels.getParticipants deep OFFSET work.
|
||||
|
|
@ -1184,6 +1187,23 @@ type DirtyChannel struct {
|
|||
Pts int
|
||||
}
|
||||
|
||||
// ChannelUpdateRetentionCheckpoint is the durable recovery boundary for one channel.
|
||||
// Events with pts <= RetainedThroughPts may be absent; callers below that floor must receive
|
||||
// channelDifferenceTooLong. LatestEventDate/LatestPts survive event pruning and keep account-level
|
||||
// dirty-channel nudges reconstructable after the hot event rows have been removed.
|
||||
type ChannelUpdateRetentionCheckpoint struct {
|
||||
ChannelID int64
|
||||
RetainedThroughPts int
|
||||
LatestEventDate int
|
||||
LatestPts int
|
||||
}
|
||||
|
||||
// ChannelUpdateRetentionResult describes one bounded, atomic prune operation.
|
||||
type ChannelUpdateRetentionResult struct {
|
||||
Checkpoint ChannelUpdateRetentionCheckpoint
|
||||
Deleted int
|
||||
}
|
||||
|
||||
// CreateChannelRequest creates a broadcast channel or megagroup.
|
||||
type CreateChannelRequest struct {
|
||||
CreatorUserID int64
|
||||
|
|
@ -1372,14 +1392,20 @@ type DeleteChannelResult struct {
|
|||
|
||||
// SendChannelMessageRequest sends one channel/supergroup message.
|
||||
type SendChannelMessageRequest struct {
|
||||
UserID int64
|
||||
ChannelID int64
|
||||
RandomID int64
|
||||
Message string
|
||||
Entities []MessageEntity
|
||||
Media *MessageMedia
|
||||
MentionUserIDs []int64
|
||||
SkipDeliveryUserIDs []int64
|
||||
UserID int64
|
||||
ChannelID int64
|
||||
RandomID int64
|
||||
// IdempotencyFingerprint is the SHA-256 of the immutable client send intent. RPC callers
|
||||
// provide a raw-TL per-item value; internal callers leave it empty for the store fallback.
|
||||
IdempotencyFingerprint []byte
|
||||
// IdempotencyPreflighted is trusted internal execution metadata; see the private-send
|
||||
// equivalent. It is deliberately excluded from the durable fingerprint.
|
||||
IdempotencyPreflighted bool
|
||||
Message string
|
||||
Entities []MessageEntity
|
||||
Media *MessageMedia
|
||||
MentionUserIDs []int64
|
||||
SkipDeliveryUserIDs []int64
|
||||
// SkipRecipientLookup lets high-level realtime fan-out use the online member
|
||||
// read model instead of forcing store.SendChannelMessage to synchronously
|
||||
// return an active-member recipient list after commit.
|
||||
|
|
@ -1405,13 +1431,26 @@ type SendChannelMessageRequest struct {
|
|||
// 虚拟频道 id;SavedPeer 是订阅者子会话分组键(订阅者发=自己,管理员回复=目标订阅者);
|
||||
// SenderUserID 是实际发件人。发件权限(订阅者身份/管理员)在 RPC 层校验,store 只校验 monoforum 存在。
|
||||
type SendMonoforumMessageRequest struct {
|
||||
MonoforumID int64
|
||||
SenderUserID int64
|
||||
SavedPeer Peer
|
||||
RandomID int64
|
||||
Message string
|
||||
Entities []MessageEntity
|
||||
Date int
|
||||
MonoforumID int64
|
||||
SenderUserID int64
|
||||
SavedPeer Peer
|
||||
RandomID int64
|
||||
IdempotencyFingerprint []byte
|
||||
IdempotencyPreflighted bool
|
||||
Message string
|
||||
Entities []MessageEntity
|
||||
Date int
|
||||
}
|
||||
|
||||
// ChannelSendReplayRequest addresses either a regular channel send (SavedPeer is zero) or one
|
||||
// monoforum sub-dialog send (SavedPeer is the subscriber scope). Lookup is read-only and must
|
||||
// never re-run membership/permission checks or allocate pts/message ids.
|
||||
type ChannelSendReplayRequest struct {
|
||||
ChannelID int64
|
||||
SenderUserID int64
|
||||
SavedPeer Peer
|
||||
RandomID int64
|
||||
IdempotencyFingerprint []byte
|
||||
}
|
||||
|
||||
// MonoforumHistoryFilter 按订阅者子会话拉取 monoforum 私信历史。
|
||||
|
|
@ -1520,7 +1559,11 @@ type SendChannelMessageResult struct {
|
|||
Event ChannelUpdateEvent
|
||||
Recipients []int64
|
||||
Duplicate bool
|
||||
Discussion *SendChannelDiscussionResult
|
||||
// ReplayDeleteEvent is the existing durable channel delete event paired
|
||||
// with a deleted exact-random_id replay. It must be returned only to the
|
||||
// caller echo and must never be fanned out as a fresh event.
|
||||
ReplayDeleteEvent *ChannelUpdateEvent
|
||||
Discussion *SendChannelDiscussionResult
|
||||
// MentionUserIDs 是本条消息解析出的被 @ 成员;在线 fanout 按它为
|
||||
// 每个接收者投影 message.mentioned/media_unread。
|
||||
MentionUserIDs []int64
|
||||
|
|
|
|||
55
internal/domain/login_code_delivery.go
Normal file
55
internal/domain/login_code_delivery.go
Normal file
|
|
@ -0,0 +1,55 @@
|
|||
package domain
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"math"
|
||||
"strings"
|
||||
)
|
||||
|
||||
const officialLoginCodeMessageTemplate = `Login code: %s. Do not give this code to anyone, even if they say they are from Telegram!
|
||||
|
||||
This code can be used to log in to your Telegram account. We never ask it for anything else.
|
||||
|
||||
If you didn't request this code by trying to log in on another device, simply ignore this message.`
|
||||
|
||||
// LoginCodeDeliveryRequest describes one durable 777000 login-code delivery.
|
||||
// PhoneCodeHash is an opaque idempotency token and must never be persisted in
|
||||
// plaintext; store implementations persist only its SHA-256 digest.
|
||||
type LoginCodeDeliveryRequest struct {
|
||||
UserID int64
|
||||
PhoneCodeHash string
|
||||
Code string
|
||||
Date int
|
||||
// ExpiresAt is the unix second after which the compact idempotency receipt
|
||||
// may be reclaimed. It must cover the corresponding code's usable lifetime.
|
||||
ExpiresAt int64
|
||||
}
|
||||
|
||||
// LoginCodeDeliveryResult returns the immutable first delivery. Created is
|
||||
// false when the same phone_code_hash was already committed and replayed.
|
||||
type LoginCodeDeliveryResult struct {
|
||||
Message Message
|
||||
Created bool
|
||||
}
|
||||
|
||||
// OfficialLoginCodeMessage builds the account-visible incoming message from
|
||||
// Telegram's official notification account. Persistence assigns ID, UID and
|
||||
// Pts atomically.
|
||||
func OfficialLoginCodeMessage(userID int64, code string, date int) (Message, error) {
|
||||
if userID <= 0 || IsSystemUserID(userID) || strings.TrimSpace(code) == "" || len(code) > 64 || date < 0 || date > math.MaxInt32 {
|
||||
return Message{}, fmt.Errorf("%w: user=%d code_length=%d date=%d", ErrLoginCodeDeliveryInvalid, userID, len(code), date)
|
||||
}
|
||||
body := fmt.Sprintf(officialLoginCodeMessageTemplate, code)
|
||||
codeOffset := len("Login code: ")
|
||||
return Message{
|
||||
OwnerUserID: userID,
|
||||
Peer: Peer{Type: PeerTypeUser, ID: OfficialSystemUserID},
|
||||
From: Peer{Type: PeerTypeUser, ID: OfficialSystemUserID},
|
||||
Date: date,
|
||||
Body: body,
|
||||
Entities: []MessageEntity{
|
||||
{Type: MessageEntityBold, Offset: 0, Length: len("Login code:")},
|
||||
{Type: MessageEntityBold, Offset: codeOffset, Length: len(code)},
|
||||
},
|
||||
}, nil
|
||||
}
|
||||
|
|
@ -3,6 +3,7 @@ package domain
|
|||
import (
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
// 本文件定义媒体相关的业务值对象(文档、照片、贴纸集、可用 reaction、消息媒体)。
|
||||
|
|
@ -81,6 +82,27 @@ type UploadedFileRef struct {
|
|||
MD5 string // small file 客户端 md5_checksum(hex),可校验;big file 为空
|
||||
}
|
||||
|
||||
// UploadedMediaKind identifies the durable object materialized from a one-shot upload file id.
|
||||
type UploadedMediaKind string
|
||||
|
||||
const (
|
||||
UploadedMediaPhoto UploadedMediaKind = "photo"
|
||||
UploadedMediaDocument UploadedMediaKind = "document"
|
||||
)
|
||||
|
||||
// UploadedMediaReceipt makes InputMediaUploaded* replayable after transient upload parts have
|
||||
// been consumed. IntentHash binds the file id to the complete materialization intent (kind, part
|
||||
// metadata and document spec); MediaID points at the immutable Photo/Document returned on every
|
||||
// exact replay.
|
||||
type UploadedMediaReceipt struct {
|
||||
OwnerUserID int64
|
||||
FileID int64
|
||||
IntentHash []byte
|
||||
Kind UploadedMediaKind
|
||||
MediaID int64
|
||||
CreatedAt time.Time
|
||||
}
|
||||
|
||||
// DocumentSpec 描述从上传文件创建 Document 的元数据(来自 InputMediaUploadedDocument)。
|
||||
type DocumentSpec struct {
|
||||
MimeType string
|
||||
|
|
|
|||
|
|
@ -260,8 +260,18 @@ type SendPrivateTextRequest struct {
|
|||
OriginAuthKeyID [8]byte
|
||||
OriginSessionID int64
|
||||
RecipientBlocked bool
|
||||
TTLPeriod int
|
||||
ViaBotID int64
|
||||
// IdempotencyFingerprint 是调用边界对原始、不可变发送请求计算的 SHA-256。
|
||||
// RPC 层应优先填入原始 TL 请求指纹,避免链接预览、骰子结果、上传媒体
|
||||
// 等服务端派生字段让合法重放看起来不同;内部调用留空时 store 会基于
|
||||
// domain command 的不可变字段生成等价指纹。
|
||||
IdempotencyFingerprint []byte
|
||||
// IdempotencyPreflighted is internal execution metadata. A trusted caller sets it only
|
||||
// after a read-only replay lookup returned absent, allowing the app/store layers to avoid
|
||||
// repeating the same indexed lookup. The transactional unique-key path still fences a
|
||||
// concurrent first writer; this flag is never part of the durable request fingerprint.
|
||||
IdempotencyPreflighted bool
|
||||
TTLPeriod int
|
||||
ViaBotID int64
|
||||
// GroupedID 相册分组 id(sendMultiMedia 同组共享非零值,非相册恒 0)。
|
||||
GroupedID int64
|
||||
// Effect 消息特效 id(私聊专属,0 表无特效;调用方已对 catalog 校验过合法性)。
|
||||
|
|
@ -275,6 +285,16 @@ type SendPrivateTextRequest struct {
|
|||
RichMessage *MessageRichMessage
|
||||
}
|
||||
|
||||
// PrivateSendReplayRequest identifies one already-committed private send without carrying any
|
||||
// mutable or resolver-derived message fields. The fingerprint is computed at the original
|
||||
// request boundary and must be a complete SHA-256 value.
|
||||
type PrivateSendReplayRequest struct {
|
||||
SenderUserID int64
|
||||
RecipientUserID int64
|
||||
RandomID int64
|
||||
IdempotencyFingerprint []byte
|
||||
}
|
||||
|
||||
// SendPrivateTextResult 描述一次私聊文本发送的双端结果。
|
||||
type SendPrivateTextResult struct {
|
||||
SenderMessage Message
|
||||
|
|
@ -282,6 +302,10 @@ type SendPrivateTextResult struct {
|
|||
SenderEvent UpdateEvent
|
||||
RecipientEvent UpdateEvent
|
||||
Duplicate bool
|
||||
// ReplayDeleteEvent is the already-durable sender-side deletion that must
|
||||
// follow the first-send snapshot in an exact random_id replay. It never
|
||||
// represents a newly allocated event.
|
||||
ReplayDeleteEvent *UpdateEvent
|
||||
}
|
||||
|
||||
// SetPrivateChatThemeRequest changes the shared theme token for a private dialog.
|
||||
|
|
@ -351,12 +375,13 @@ type ForwardPrivateMessagesRequest struct {
|
|||
|
||||
// ForwardPrivateMessagesResult 描述一次私聊转发的 owner 维度结果。
|
||||
type ForwardPrivateMessagesResult struct {
|
||||
OwnerUserID int64
|
||||
SenderMessages []Message
|
||||
RecipientMessages []Message
|
||||
SenderEvents []UpdateEvent
|
||||
RecipientEvents []UpdateEvent
|
||||
Duplicates []bool
|
||||
OwnerUserID int64
|
||||
SenderMessages []Message
|
||||
RecipientMessages []Message
|
||||
SenderEvents []UpdateEvent
|
||||
RecipientEvents []UpdateEvent
|
||||
Duplicates []bool
|
||||
ReplayDeleteEvents []*UpdateEvent
|
||||
}
|
||||
|
||||
// ReadHistoryRequest 是账号视角的 messages.readHistory 命令。
|
||||
|
|
|
|||
|
|
@ -3,13 +3,30 @@ package domain
|
|||
import "errors"
|
||||
|
||||
var (
|
||||
ErrMessageIDInvalid = errors.New("message id invalid")
|
||||
ErrMessageEmpty = errors.New("message empty")
|
||||
ErrMessageAuthorRequired = errors.New("message author required")
|
||||
ErrMessageNotModified = errors.New("message not modified")
|
||||
ErrMessageNotReadYet = errors.New("message not read yet")
|
||||
ErrReplyMessageIDInvalid = errors.New("reply message id invalid")
|
||||
ErrChatForwardsRestricted = errors.New("chat forwards restricted")
|
||||
ErrMessageIDInvalid = errors.New("message id invalid")
|
||||
ErrMessageEmpty = errors.New("message empty")
|
||||
ErrMessageAuthorRequired = errors.New("message author required")
|
||||
ErrMessageNotModified = errors.New("message not modified")
|
||||
ErrMessageNotReadYet = errors.New("message not read yet")
|
||||
// ErrMessageRandomIDDuplicate 表示同一发送者重复使用 random_id,且本次
|
||||
// 不可变请求载荷与首次成功发送不一致。完全相同的重放不返回此错误,
|
||||
// 而是复用首次发送结果。
|
||||
ErrMessageRandomIDDuplicate = errors.New("message random id duplicate")
|
||||
// ErrLoginCodeDeliveryInvalid rejects malformed durable 777000 delivery
|
||||
// commands before allocating message/pts facts.
|
||||
ErrLoginCodeDeliveryInvalid = errors.New("login code delivery invalid")
|
||||
// ErrLoginCodeDeliveryConflict means one phone_code_hash digest was reused
|
||||
// for a different account or code. It must fail closed rather than expose or
|
||||
// overwrite the first account's immutable receipt.
|
||||
ErrLoginCodeDeliveryConflict = errors.New("login code delivery conflict")
|
||||
// ErrLoginCodeDeliveryCommitAmbiguous means PostgreSQL lost the commit
|
||||
// acknowledgement and an independent receipt probe could not prove whether
|
||||
// the durable 777000 transaction committed. Callers must retain the opaque
|
||||
// code record until TTL expiry; deleting it could invalidate a committed but
|
||||
// undisclosed delivery and make a retry impossible to reconcile.
|
||||
ErrLoginCodeDeliveryCommitAmbiguous = errors.New("login code delivery commit ambiguous")
|
||||
ErrReplyMessageIDInvalid = errors.New("reply message id invalid")
|
||||
ErrChatForwardsRestricted = errors.New("chat forwards restricted")
|
||||
// ErrPinnedSavedDialogsTooMuch 映射 PINNED_TOO_MUCH:收藏夹子会话置顶
|
||||
// 数量达到 MaxPinnedSavedDialogs 上限。
|
||||
ErrPinnedSavedDialogsTooMuch = errors.New("pinned saved dialogs too much")
|
||||
|
|
|
|||
166
internal/mtprotoedge/admission.go
Normal file
166
internal/mtprotoedge/admission.go
Normal file
|
|
@ -0,0 +1,166 @@
|
|||
package mtprotoedge
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
const (
|
||||
defaultMaxConnections = 200_000
|
||||
defaultMaxConnectionsPerIP = 4_096
|
||||
defaultMaxConcurrentHandshakes = 256
|
||||
acceptRetryInitialDelay = 5 * time.Millisecond
|
||||
acceptRetryMaxDelay = time.Second
|
||||
)
|
||||
|
||||
// admissionController 把 raw socket 与昂贵的 RSA/DH exchange 分开限流。
|
||||
// raw 配额覆盖连接从 Accept 到物理 Close 的完整生命周期;handshake 配额只覆盖
|
||||
// auth_key_id=0 的 exchange(包括 TDesktop 每条候选连接的 fake req_pq 探活)。
|
||||
type admissionController struct {
|
||||
mu sync.Mutex
|
||||
maxConnections int
|
||||
maxPerIP int
|
||||
connections int
|
||||
byIP map[string]int
|
||||
handshakes chan struct{}
|
||||
}
|
||||
|
||||
func newAdmissionController(maxConnections, maxPerIP, maxHandshakes int) *admissionController {
|
||||
a := &admissionController{
|
||||
maxConnections: maxConnections,
|
||||
maxPerIP: maxPerIP,
|
||||
byIP: make(map[string]int),
|
||||
}
|
||||
if maxHandshakes > 0 {
|
||||
a.handshakes = make(chan struct{}, maxHandshakes)
|
||||
}
|
||||
return a
|
||||
}
|
||||
|
||||
func (a *admissionController) wrapListener(ln net.Listener) net.Listener {
|
||||
if a == nil {
|
||||
return ln
|
||||
}
|
||||
return &admissionListener{Listener: ln, admission: a}
|
||||
}
|
||||
|
||||
func (a *admissionController) acquireConnection(addr net.Addr) (func(), bool) {
|
||||
if a == nil {
|
||||
return func() {}, true
|
||||
}
|
||||
ip := remoteAdmissionKey(addr)
|
||||
a.mu.Lock()
|
||||
if (a.maxConnections > 0 && a.connections >= a.maxConnections) ||
|
||||
(a.maxPerIP > 0 && a.byIP[ip] >= a.maxPerIP) {
|
||||
a.mu.Unlock()
|
||||
return nil, false
|
||||
}
|
||||
a.connections++
|
||||
a.byIP[ip]++
|
||||
a.mu.Unlock()
|
||||
|
||||
var once sync.Once
|
||||
return func() {
|
||||
once.Do(func() {
|
||||
a.mu.Lock()
|
||||
a.connections--
|
||||
a.byIP[ip]--
|
||||
if a.byIP[ip] == 0 {
|
||||
delete(a.byIP, ip)
|
||||
}
|
||||
a.mu.Unlock()
|
||||
})
|
||||
}, true
|
||||
}
|
||||
|
||||
func (a *admissionController) tryAcquireHandshake() (func(), bool) {
|
||||
if a == nil || a.handshakes == nil {
|
||||
return func() {}, true
|
||||
}
|
||||
select {
|
||||
case a.handshakes <- struct{}{}:
|
||||
var once sync.Once
|
||||
return func() {
|
||||
once.Do(func() { <-a.handshakes })
|
||||
}, true
|
||||
default:
|
||||
return nil, false
|
||||
}
|
||||
}
|
||||
|
||||
func remoteAdmissionKey(addr net.Addr) string {
|
||||
if addr == nil {
|
||||
return "<unknown>"
|
||||
}
|
||||
if host, _, err := net.SplitHostPort(addr.String()); err == nil {
|
||||
if ip := net.ParseIP(host); ip != nil {
|
||||
return ip.String()
|
||||
}
|
||||
return host
|
||||
}
|
||||
return addr.Network() + ":" + addr.String()
|
||||
}
|
||||
|
||||
// admissionListener 在最早的原始 Accept 边界记账,因此 mixed TCP/WebSocket 的
|
||||
// sniff/upgrade 阶段也受 raw cap 保护。admittedConn.Close 负责幂等归还配额。
|
||||
type admissionListener struct {
|
||||
net.Listener
|
||||
admission *admissionController
|
||||
}
|
||||
|
||||
func (l *admissionListener) Accept() (net.Conn, error) {
|
||||
for {
|
||||
conn, err := l.Listener.Accept()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
release, ok := l.admission.acquireConnection(conn.RemoteAddr())
|
||||
if !ok {
|
||||
_ = conn.Close()
|
||||
continue
|
||||
}
|
||||
return &admittedConn{Conn: conn, release: release}, nil
|
||||
}
|
||||
}
|
||||
|
||||
type admittedConn struct {
|
||||
net.Conn
|
||||
release func()
|
||||
once sync.Once
|
||||
}
|
||||
|
||||
func (c *admittedConn) Close() error {
|
||||
err := c.Conn.Close()
|
||||
c.once.Do(c.release)
|
||||
return err
|
||||
}
|
||||
|
||||
func isTemporaryAcceptError(err error) bool {
|
||||
type temporary interface{ Temporary() bool }
|
||||
e, ok := err.(temporary)
|
||||
return ok && e.Temporary()
|
||||
}
|
||||
|
||||
func nextAcceptRetryDelay(previous time.Duration) time.Duration {
|
||||
if previous <= 0 {
|
||||
return acceptRetryInitialDelay
|
||||
}
|
||||
next := previous * 2
|
||||
if next > acceptRetryMaxDelay {
|
||||
return acceptRetryMaxDelay
|
||||
}
|
||||
return next
|
||||
}
|
||||
|
||||
func waitAcceptRetry(ctx context.Context, delay time.Duration) bool {
|
||||
timer := time.NewTimer(delay)
|
||||
defer timer.Stop()
|
||||
select {
|
||||
case <-timer.C:
|
||||
return true
|
||||
case <-ctx.Done():
|
||||
return false
|
||||
}
|
||||
}
|
||||
313
internal/mtprotoedge/admission_test.go
Normal file
313
internal/mtprotoedge/admission_test.go
Normal file
|
|
@ -0,0 +1,313 @@
|
|||
package mtprotoedge
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"net"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/gotd/td/bin"
|
||||
"github.com/gotd/td/proto/codec"
|
||||
|
||||
"telesrv/internal/store"
|
||||
"telesrv/internal/store/memory"
|
||||
)
|
||||
|
||||
func TestAdmissionConnectionLimitsAndIdempotentRelease(t *testing.T) {
|
||||
a := newAdmissionController(2, 1, 1)
|
||||
ip1a := &net.TCPAddr{IP: net.ParseIP("203.0.113.1"), Port: 1000}
|
||||
ip1b := &net.TCPAddr{IP: net.ParseIP("203.0.113.1"), Port: 1001}
|
||||
ip2 := &net.TCPAddr{IP: net.ParseIP("203.0.113.2"), Port: 1000}
|
||||
ip3 := &net.TCPAddr{IP: net.ParseIP("203.0.113.3"), Port: 1000}
|
||||
|
||||
release1, ok := a.acquireConnection(ip1a)
|
||||
if !ok {
|
||||
t.Fatal("first connection rejected")
|
||||
}
|
||||
if _, ok := a.acquireConnection(ip1b); ok {
|
||||
t.Fatal("second connection from same IP bypassed per-IP cap")
|
||||
}
|
||||
release2, ok := a.acquireConnection(ip2)
|
||||
if !ok {
|
||||
t.Fatal("second IP connection rejected below global cap")
|
||||
}
|
||||
if _, ok := a.acquireConnection(ip3); ok {
|
||||
t.Fatal("third connection bypassed global cap")
|
||||
}
|
||||
|
||||
release1()
|
||||
release1() // 幂等归还不得把计数减成负数。
|
||||
releaseAgain, ok := a.acquireConnection(ip1b)
|
||||
if !ok {
|
||||
t.Fatal("released per-IP/global slot was not reusable")
|
||||
}
|
||||
releaseAgain()
|
||||
release2()
|
||||
|
||||
a.mu.Lock()
|
||||
defer a.mu.Unlock()
|
||||
if a.connections != 0 || len(a.byIP) != 0 {
|
||||
t.Fatalf("admission counters after release = %d/%v, want 0/empty", a.connections, a.byIP)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAdmissionHandshakeLimitAndRelease(t *testing.T) {
|
||||
a := newAdmissionController(-1, -1, 1)
|
||||
release, ok := a.tryAcquireHandshake()
|
||||
if !ok {
|
||||
t.Fatal("first handshake rejected")
|
||||
}
|
||||
if _, ok := a.tryAcquireHandshake(); ok {
|
||||
t.Fatal("second handshake bypassed semaphore")
|
||||
}
|
||||
release()
|
||||
release() // 幂等
|
||||
release2, ok := a.tryAcquireHandshake()
|
||||
if !ok {
|
||||
t.Fatal("released handshake slot was not reusable")
|
||||
}
|
||||
release2()
|
||||
}
|
||||
|
||||
type oneConnListener struct {
|
||||
conn net.Conn
|
||||
once sync.Once
|
||||
}
|
||||
|
||||
func (l *oneConnListener) Accept() (net.Conn, error) {
|
||||
var conn net.Conn
|
||||
l.once.Do(func() {
|
||||
conn = l.conn
|
||||
})
|
||||
if conn == nil {
|
||||
return nil, net.ErrClosed
|
||||
}
|
||||
return conn, nil
|
||||
}
|
||||
func (l *oneConnListener) Close() error { return l.conn.Close() }
|
||||
func (l *oneConnListener) Addr() net.Addr { return l.conn.LocalAddr() }
|
||||
|
||||
func TestAdmissionListenerTracksUntilPhysicalClose(t *testing.T) {
|
||||
serverSide, clientSide := net.Pipe()
|
||||
defer clientSide.Close()
|
||||
a := newAdmissionController(1, 1, 1)
|
||||
ln := a.wrapListener(&oneConnListener{conn: serverSide})
|
||||
conn, err := ln.Accept()
|
||||
if err != nil {
|
||||
t.Fatalf("Accept: %v", err)
|
||||
}
|
||||
a.mu.Lock()
|
||||
active := a.connections
|
||||
a.mu.Unlock()
|
||||
if active != 1 {
|
||||
t.Fatalf("active after Accept = %d, want 1", active)
|
||||
}
|
||||
_ = conn.Close()
|
||||
_ = conn.Close()
|
||||
a.mu.Lock()
|
||||
active = a.connections
|
||||
a.mu.Unlock()
|
||||
if active != 0 {
|
||||
t.Fatalf("active after physical Close = %d, want 0", active)
|
||||
}
|
||||
}
|
||||
|
||||
type temporaryAcceptTestError struct{}
|
||||
|
||||
func (temporaryAcceptTestError) Error() string { return "temporary accept failure" }
|
||||
func (temporaryAcceptTestError) Timeout() bool { return false }
|
||||
func (temporaryAcceptTestError) Temporary() bool { return true }
|
||||
|
||||
type temporaryThenConnListener struct {
|
||||
conn net.Conn
|
||||
closed chan struct{}
|
||||
closeOnce sync.Once
|
||||
calls atomic.Int32
|
||||
}
|
||||
|
||||
type connThenErrorListener struct {
|
||||
conn net.Conn
|
||||
err error
|
||||
closeOnce sync.Once
|
||||
calls atomic.Int32
|
||||
}
|
||||
|
||||
func (l *connThenErrorListener) Accept() (net.Conn, error) {
|
||||
if l.calls.Add(1) == 1 {
|
||||
return l.conn, nil
|
||||
}
|
||||
return nil, l.err
|
||||
}
|
||||
|
||||
func (l *connThenErrorListener) Close() error {
|
||||
var err error
|
||||
l.closeOnce.Do(func() {
|
||||
if l.conn != nil {
|
||||
err = l.conn.Close()
|
||||
}
|
||||
})
|
||||
return err
|
||||
}
|
||||
|
||||
func (l *connThenErrorListener) Addr() net.Addr {
|
||||
if l.conn != nil {
|
||||
return l.conn.LocalAddr()
|
||||
}
|
||||
return &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 12345}
|
||||
}
|
||||
|
||||
type fixedErrorListener struct {
|
||||
err error
|
||||
addr net.Addr
|
||||
}
|
||||
|
||||
func (l *fixedErrorListener) Accept() (net.Conn, error) { return nil, l.err }
|
||||
func (*fixedErrorListener) Close() error { return nil }
|
||||
func (l *fixedErrorListener) Addr() net.Addr { return l.addr }
|
||||
|
||||
func (l *temporaryThenConnListener) Accept() (net.Conn, error) {
|
||||
call := l.calls.Add(1)
|
||||
if call == 1 {
|
||||
return nil, temporaryAcceptTestError{}
|
||||
}
|
||||
if call == 2 {
|
||||
return l.conn, nil
|
||||
}
|
||||
<-l.closed
|
||||
return nil, net.ErrClosed
|
||||
}
|
||||
|
||||
func (l *temporaryThenConnListener) Close() error {
|
||||
l.closeOnce.Do(func() {
|
||||
close(l.closed)
|
||||
_ = l.conn.Close()
|
||||
})
|
||||
return nil
|
||||
}
|
||||
|
||||
func (l *temporaryThenConnListener) Addr() net.Addr {
|
||||
return &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 12345}
|
||||
}
|
||||
|
||||
func TestAcceptLoopRetriesTemporaryError(t *testing.T) {
|
||||
serverSide, clientSide := net.Pipe()
|
||||
defer clientSide.Close()
|
||||
ln := &temporaryThenConnListener{conn: serverSide, closed: make(chan struct{})}
|
||||
srv := New(Options{HandshakeIdleTimeout: 100 * time.Millisecond})
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
done := make(chan error, 1)
|
||||
go func() { done <- srv.acceptLoop(ctx, ln, false) }()
|
||||
|
||||
deadline := time.Now().Add(time.Second)
|
||||
for ln.calls.Load() < 3 && time.Now().Before(deadline) {
|
||||
time.Sleep(time.Millisecond)
|
||||
}
|
||||
if ln.calls.Load() < 3 {
|
||||
cancel()
|
||||
<-done
|
||||
t.Fatalf("accept calls = %d, want temporary retry then next accept", ln.calls.Load())
|
||||
}
|
||||
cancel()
|
||||
select {
|
||||
case err := <-done:
|
||||
if err != nil {
|
||||
t.Fatalf("acceptLoop after temporary error: %v", err)
|
||||
}
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("acceptLoop did not stop after cancel")
|
||||
}
|
||||
}
|
||||
|
||||
func TestAcceptLoopPermanentErrorCancelsAcceptedConnectionsBeforeWait(t *testing.T) {
|
||||
serverSide, clientSide := net.Pipe()
|
||||
defer clientSide.Close()
|
||||
wantErr := errors.New("permanent accept failure")
|
||||
ln := &connThenErrorListener{conn: serverSide, err: wantErr}
|
||||
srv := New(Options{HandshakeIdleTimeout: time.Hour})
|
||||
|
||||
done := make(chan error, 1)
|
||||
go func() {
|
||||
done <- srv.acceptLoop(context.Background(), ln, false)
|
||||
}()
|
||||
|
||||
select {
|
||||
case err := <-done:
|
||||
if !errors.Is(err, wantErr) {
|
||||
t.Fatalf("acceptLoop error = %v, want %v", err, wantErr)
|
||||
}
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("acceptLoop waited for an accepted connection before canceling it")
|
||||
}
|
||||
|
||||
_ = clientSide.SetReadDeadline(time.Now().Add(time.Second))
|
||||
var one [1]byte
|
||||
if _, err := clientSide.Read(one[:]); err == nil {
|
||||
t.Fatal("accepted connection remained open after permanent accept failure")
|
||||
}
|
||||
}
|
||||
|
||||
func TestServeMixedStopsAllComponentsWhenOneReturnsCleanly(t *testing.T) {
|
||||
srv := New(Options{WebSocket: true})
|
||||
ln := &fixedErrorListener{
|
||||
err: net.ErrClosed,
|
||||
addr: &net.TCPAddr{IP: net.IPv4(127, 0, 0, 1), Port: 2398},
|
||||
}
|
||||
|
||||
done := make(chan error, 1)
|
||||
go func() {
|
||||
done <- srv.serveMixed(context.Background(), ln)
|
||||
}()
|
||||
select {
|
||||
case err := <-done:
|
||||
if err != nil {
|
||||
t.Fatalf("serveMixed error = %v, want nil closed-listener shutdown", err)
|
||||
}
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("serveMixed did not stop remaining components after one clean exit")
|
||||
}
|
||||
}
|
||||
|
||||
type countingAuthKeyStore struct {
|
||||
store.AuthKeyStore
|
||||
gets atomic.Int32
|
||||
}
|
||||
|
||||
func (s *countingAuthKeyStore) Get(ctx context.Context, id [8]byte) (store.AuthKeyData, bool, error) {
|
||||
s.gets.Add(1)
|
||||
return s.AuthKeyStore.Get(ctx, id)
|
||||
}
|
||||
|
||||
func TestUnknownAuthKeyRespondsOnceThenCloses(t *testing.T) {
|
||||
keys := &countingAuthKeyStore{AuthKeyStore: memory.NewAuthKeyStore()}
|
||||
addr, _, _ := startTestServer(t, Options{AuthKeys: keys})
|
||||
conn := dialTransportOnly(t, addr)
|
||||
|
||||
var request bin.Buffer
|
||||
request.PutLong(0x0102030405060708)
|
||||
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
|
||||
defer cancel()
|
||||
if err := conn.Send(ctx, &request); err != nil {
|
||||
t.Fatalf("send unknown auth key: %v", err)
|
||||
}
|
||||
var response bin.Buffer
|
||||
err := conn.Recv(ctx, &response)
|
||||
var protocolErr *codec.ProtocolErr
|
||||
if !errors.As(err, &protocolErr) || protocolErr.Code != codec.CodeAuthKeyNotFound {
|
||||
t.Fatalf("first recv err = %T %v, want protocol -404", err, err)
|
||||
}
|
||||
if got := keys.gets.Load(); got != 1 {
|
||||
t.Fatalf("AuthKeyStore.Get calls = %d, want 1", got)
|
||||
}
|
||||
|
||||
response.Reset()
|
||||
err = conn.Recv(ctx, &response)
|
||||
if err == nil {
|
||||
t.Fatal("connection remained readable after terminal -404")
|
||||
}
|
||||
if got := keys.gets.Load(); got != 1 {
|
||||
t.Fatalf("AuthKeyStore.Get calls after close = %d, want 1", got)
|
||||
}
|
||||
}
|
||||
56
internal/mtprotoedge/auth_key_switch_test.go
Normal file
56
internal/mtprotoedge/auth_key_switch_test.go
Normal file
|
|
@ -0,0 +1,56 @@
|
|||
package mtprotoedge
|
||||
|
||||
import (
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/gotd/td/mt"
|
||||
"github.com/gotd/td/proto"
|
||||
)
|
||||
|
||||
func TestEncryptedConnectionSwitchesAuthKeyEvenWhenSessionIDIsReused(t *testing.T) {
|
||||
const dc = 2
|
||||
addr, pub, srv := startTestServer(t, Options{DC: dc})
|
||||
connA, authA, cipherA := dialHandshake(t, addr, dc, pub)
|
||||
_, authB, cipherB := dialHandshake(t, addr, dc, pub)
|
||||
msgID := proto.NewMessageIDGen(time.Now)
|
||||
|
||||
sendEncrypted(t, connA, cipherA, authA, msgID.New(proto.MessageFromClient), &mt.PingRequest{PingID: 1})
|
||||
for range 3 { // new_session_created + pong + msgs_ack; leave no A-key frame on the socket.
|
||||
readServerMessage(t, connA, cipherA, authA.AuthKey)
|
||||
}
|
||||
|
||||
// Reuse A's session id on the same physical TCP socket, but encrypt with the independently
|
||||
// established key B and B's salt. Session identity is (raw auth_key_id, session_id): comparing
|
||||
// session_id alone would keep A's cached key/user identity and encrypt the reply with A.
|
||||
body := encodeClientMessageBodyForTest(t, &mt.PingRequest{PingID: 2})
|
||||
sendEncryptedWithSessionSaltAndSeq(
|
||||
t,
|
||||
connA,
|
||||
cipherB,
|
||||
authB,
|
||||
authA.SessionID,
|
||||
authB.ServerSalt,
|
||||
msgID.New(proto.MessageFromClient),
|
||||
1,
|
||||
body,
|
||||
)
|
||||
seenPong := false
|
||||
for range 3 {
|
||||
_, typeID, _ := readServerMessage(t, connA, cipherB, authB.AuthKey)
|
||||
seenPong = seenPong || typeID == mt.PongTypeID
|
||||
}
|
||||
if !seenPong {
|
||||
t.Fatal("new auth key did not receive pong")
|
||||
}
|
||||
|
||||
oldKey := sessionKey{authKeyID: authA.AuthKey.ID, sessionID: authA.SessionID}
|
||||
newKey := sessionKey{authKeyID: authB.AuthKey.ID, sessionID: authA.SessionID}
|
||||
srv.conns.mu.RLock()
|
||||
_, oldAlive := srv.conns.bySession[oldKey]
|
||||
current := srv.conns.bySession[newKey]
|
||||
srv.conns.mu.RUnlock()
|
||||
if oldAlive || current == nil || current.authKeyID != authB.AuthKey.ID {
|
||||
t.Fatalf("registry after key switch: old_alive=%v current=%v", oldAlive, current != nil)
|
||||
}
|
||||
}
|
||||
|
|
@ -46,27 +46,53 @@ type Conn struct {
|
|||
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
|
||||
// terminal 表示该 logical Conn 已停止接受新的出站操作。写失败时由
|
||||
// outbound actor 置位并只发停止信号,不能在 actor 内等待自身退出。
|
||||
terminal atomic.Bool
|
||||
transportClose sync.Once
|
||||
|
||||
rpcQueue chan inboundRPC
|
||||
rpcStop chan struct{}
|
||||
rpcCancel context.CancelFunc
|
||||
rpcClose sync.Once
|
||||
rpcWG sync.WaitGroup
|
||||
rpcTimeout time.Duration
|
||||
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
|
||||
// inflightRPCBytes 跟踪已入队未完成的 inbound RPC body 总字节,配合 maxInflightRPCBytes
|
||||
// 给 RPC 队列设字节预算(不止限条数),防对抗客户端发大请求撑内存。
|
||||
inflightRPCBytes atomic.Int64
|
||||
// RPC worker 懒启动:首个 RPC 入队时才起 worker(ensureInboundRPCWorkers),
|
||||
// 避免握手后静默 / 纯推送目标连接白白钉住 rpcMaxInflight 个 goroutine。
|
||||
// 单连接只保留并发配额;实际 worker 来自 Server 共享池,避免每连接预留 goroutine。
|
||||
rpcRootCtx context.Context
|
||||
rpcMaxInflight int
|
||||
rpcWorkersOnce sync.Once
|
||||
|
||||
// sentContentMessages 只由 outbound actor 访问,用于生成 MTProto seq_no。
|
||||
sentContentMessages int32
|
||||
// outboundPlain/outboundWire 只由 outbound actor 访问,用于复用出站加密缓冲。
|
||||
outboundPlain bin.Buffer
|
||||
outboundWire bin.Buffer
|
||||
// outboundRand 只由 outbound actor 访问:对 cipher 随机源的缓冲预读,
|
||||
// 把每帧 padding 的 getrandom syscall 摊薄成 ~1KiB 一次。
|
||||
outboundRand *bufio.Reader
|
||||
|
|
|
|||
|
|
@ -1,6 +1,8 @@
|
|||
package mtprotoedge
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"compress/gzip"
|
||||
"context"
|
||||
"crypto/sha256"
|
||||
"encoding/binary"
|
||||
|
|
@ -8,6 +10,7 @@ import (
|
|||
"fmt"
|
||||
"io"
|
||||
"math"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"go.uber.org/zap"
|
||||
|
|
@ -59,6 +62,25 @@ func (cs *connState) reset() {
|
|||
|
||||
const (
|
||||
maxTrackedClientMsgIDs = 400
|
||||
// maxContainerMessages bounds per-frame recursive work and ack growth. Official clients batch
|
||||
// far fewer messages; 1024 leaves ample headroom while preventing a 16 MiB frame of zero-body
|
||||
// container entries from expanding into tens of MiB of Go objects.
|
||||
maxContainerMessages = 1024
|
||||
// maxDispatchDepth bounds gzip/container wrapper recursion. Normal shapes are RPC, gzip(RPC),
|
||||
// container(RPC...) and gzip(container(...)); deeper nesting has no compatibility value.
|
||||
maxDispatchDepth = 4
|
||||
// gotd already caps each gzip expansion at 10 MiB. This cumulative cap prevents several nested
|
||||
// gzip layers in one transport frame from repeatedly allocating/decompressing that allowance.
|
||||
maxDispatchExpandedBytes = 32 << 20
|
||||
maxSingleGZIPExpandedBytes = 10 << 20
|
||||
// MTProto service vectors operate on bounded connection tracking tables. Accepting more IDs
|
||||
// only burns decode/CPU and cannot improve the result.
|
||||
maxServiceMessageIDs = 4096
|
||||
// A decoded container descriptor is 48 bytes on 64-bit Go today. Charge 64 bytes per entry
|
||||
// before allocating the exact-size slice so allocator rounding and future field growth remain
|
||||
// inside the process-wide inbound budget. Message bodies stay as zero-copy views of the already
|
||||
// charged plaintext frame/gzip expansion.
|
||||
containerDescriptorBudgetBytes = 64
|
||||
|
||||
msgStateUnknown byte = 1
|
||||
msgStateNotReceived byte = 2
|
||||
|
|
@ -101,7 +123,7 @@ func (s *Server) handleEncrypted(ctx context.Context, tc transport.Conn, cs *con
|
|||
if frame.salt != serverSalt {
|
||||
c := current
|
||||
temp := false
|
||||
if c == nil || c.sessionID != frame.sessionID {
|
||||
if c == nil || c.sessionID != frame.sessionID || c.authKeyID != key.ID {
|
||||
c = s.newConn(tc, key, frame.sessionID, serverSalt)
|
||||
temp = true
|
||||
}
|
||||
|
|
@ -113,7 +135,7 @@ func (s *Server) handleEncrypted(ctx context.Context, tc transport.Conn, cs *con
|
|||
}
|
||||
|
||||
// 首个加密消息或 session 变化时(重新)注册连接到 SessionManager。
|
||||
if current == nil || current.sessionID != frame.sessionID {
|
||||
if current == nil || current.sessionID != frame.sessionID || current.authKeyID != key.ID {
|
||||
if current != nil {
|
||||
cs.reset()
|
||||
}
|
||||
|
|
@ -150,7 +172,7 @@ func (s *Server) handleEncrypted(ctx context.Context, tc transport.Conn, cs *con
|
|||
)
|
||||
return current, s.sendBadMsg(ctx, current, frame.messageID, frame.seqNo, code)
|
||||
}
|
||||
if err := sendQuickAckIfRequested(ctx, tc, key, frame.plaintext); err != nil {
|
||||
if err := sendQuickAckIfRequested(ctx, tc, key, frame.plaintext, s.writeTimeout); err != nil {
|
||||
return current, err
|
||||
}
|
||||
|
||||
|
|
@ -224,12 +246,28 @@ func (s *Server) maybePersistSession(ctx context.Context, c *Conn, sessionID int
|
|||
}
|
||||
}
|
||||
|
||||
func sendQuickAckIfRequested(ctx context.Context, tc transport.Conn, key crypto.AuthKey, plaintext []byte) error {
|
||||
func sendQuickAckIfRequested(ctx context.Context, tc transport.Conn, key crypto.AuthKey, plaintext []byte, writeTimeout time.Duration) error {
|
||||
q, ok := tc.(quickAckTransport)
|
||||
if !ok || !q.ConsumeQuickAckRequested() {
|
||||
return nil
|
||||
}
|
||||
return q.SendQuickAck(ctx, clientQuickAckToken(key, plaintext))
|
||||
token := clientQuickAckToken(key, plaintext)
|
||||
deadline := time.Time{}
|
||||
if writeTimeout > 0 {
|
||||
deadline = time.Now().Add(writeTimeout)
|
||||
}
|
||||
if d, ok := ctx.Deadline(); ok && (deadline.IsZero() || d.Before(deadline)) {
|
||||
deadline = d
|
||||
}
|
||||
if dq, ok := tc.(deadlineQuickAckTransport); ok {
|
||||
return dq.SendQuickAckDeadline(deadline, token)
|
||||
}
|
||||
if deadline.IsZero() {
|
||||
return q.SendQuickAck(ctx, token)
|
||||
}
|
||||
sendCtx, cancel := context.WithDeadline(ctx, deadline)
|
||||
defer cancel()
|
||||
return q.SendQuickAck(sendCtx, token)
|
||||
}
|
||||
|
||||
// clientQuickAckToken 按 Android MTProto v2 公式计算 quick ack:SHA256(auth_key[88:120] +
|
||||
|
|
@ -246,6 +284,20 @@ func clientQuickAckToken(key crypto.AuthKey, plaintext []byte) uint32 {
|
|||
// dispatch 处理一条明文消息:解包 container/gzip,处理服务消息,其余转 RPC 路由。
|
||||
// content-related 消息(ping、RPC)的 msg_id 会收集到 acks 以便统一确认。
|
||||
func (s *Server) dispatch(ctx context.Context, cs *connState, c *Conn, msgID int64, seqNo int32, b *bin.Buffer, acks *[]int64) error {
|
||||
expanded := 0
|
||||
return s.dispatchWithBudget(ctx, cs, c, msgID, seqNo, b, acks, dispatchBudget{expanded: &expanded})
|
||||
}
|
||||
|
||||
type dispatchBudget struct {
|
||||
depth int
|
||||
containerDepth int
|
||||
expanded *int
|
||||
}
|
||||
|
||||
func (s *Server) dispatchWithBudget(ctx context.Context, cs *connState, c *Conn, msgID int64, seqNo int32, b *bin.Buffer, acks *[]int64, budget dispatchBudget) error {
|
||||
if budget.depth > maxDispatchDepth {
|
||||
return fmt.Errorf("mtproto wrapper depth %d exceeds %d", budget.depth, maxDispatchDepth)
|
||||
}
|
||||
id, err := b.PeekID()
|
||||
if err != nil {
|
||||
return fmt.Errorf("peek type id: %w", err)
|
||||
|
|
@ -258,20 +310,39 @@ func (s *Server) dispatch(ctx context.Context, cs *connState, c *Conn, msgID int
|
|||
|
||||
switch id {
|
||||
case proto.GZIPTypeID:
|
||||
var gz proto.GZIP
|
||||
if err := gz.Decode(b); err != nil {
|
||||
data, releaseExpansion, err := s.decodeGZIPWithGlobalBudget(b)
|
||||
if err != nil {
|
||||
return fmt.Errorf("decode gzip: %w", err)
|
||||
}
|
||||
return s.dispatch(ctx, cs, c, msgID, seqNo, &bin.Buffer{Buf: gz.Data}, acks)
|
||||
defer releaseExpansion()
|
||||
*budget.expanded += len(data)
|
||||
if *budget.expanded > maxDispatchExpandedBytes {
|
||||
return fmt.Errorf("cumulative gzip expansion %d exceeds %d", *budget.expanded, maxDispatchExpandedBytes)
|
||||
}
|
||||
budget.depth++
|
||||
return s.dispatchWithBudget(ctx, cs, c, msgID, seqNo, &bin.Buffer{Buf: data}, acks, budget)
|
||||
|
||||
case proto.MessageContainerTypeID:
|
||||
var container proto.MessageContainer
|
||||
if err := container.Decode(b); err != nil {
|
||||
if budget.containerDepth != 0 {
|
||||
return s.sendBadMsg(ctx, c, msgID, seqNo, badMsgContainer)
|
||||
}
|
||||
count, err := containerMessageCount(b)
|
||||
if err != nil {
|
||||
return fmt.Errorf("decode container count: %w", err)
|
||||
}
|
||||
if count > maxContainerMessages {
|
||||
return s.sendBadMsg(ctx, c, msgID, seqNo, badMsgContainer)
|
||||
}
|
||||
container, releaseContainer, err := s.decodeMessageContainerViews(b, count)
|
||||
if err != nil {
|
||||
return fmt.Errorf("decode container: %w", err)
|
||||
}
|
||||
defer releaseContainer()
|
||||
if code := validateClientContainer(msgID, seqNo, container); code != 0 {
|
||||
return s.sendBadMsg(ctx, c, msgID, seqNo, code)
|
||||
}
|
||||
budget.depth++
|
||||
budget.containerDepth++
|
||||
for i := range container.Messages {
|
||||
m := container.Messages[i]
|
||||
typeID, err := (&bin.Buffer{Buf: m.Body}).PeekID()
|
||||
|
|
@ -292,7 +363,7 @@ func (s *Server) dispatch(ctx context.Context, cs *connState, c *Conn, msgID int
|
|||
return s.sendBadMsg(ctx, c, m.ID, int32(m.SeqNo), code)
|
||||
}
|
||||
cs.track(m.ID, int32(m.SeqNo), content, msgStateReceived)
|
||||
if err := s.dispatch(ctx, cs, c, m.ID, int32(m.SeqNo), &bin.Buffer{Buf: m.Body}, acks); err != nil {
|
||||
if err := s.dispatchWithBudget(ctx, cs, c, m.ID, int32(m.SeqNo), &bin.Buffer{Buf: m.Body}, acks, budget); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
|
@ -323,6 +394,9 @@ func (s *Server) dispatch(ctx context.Context, cs *connState, c *Conn, msgID int
|
|||
return s.sendFutureSalts(ctx, c, msgID, req.Num)
|
||||
|
||||
case mt.MsgsAckTypeID:
|
||||
if err := validateFirstVectorCount(b, maxServiceMessageIDs); err != nil {
|
||||
return fmt.Errorf("msgs_ack vector: %w", err)
|
||||
}
|
||||
var ack mt.MsgsAck
|
||||
if err := ack.Decode(b); err != nil {
|
||||
return fmt.Errorf("decode msgs_ack: %w", err)
|
||||
|
|
@ -332,6 +406,9 @@ func (s *Server) dispatch(ctx context.Context, cs *connState, c *Conn, msgID int
|
|||
return nil
|
||||
|
||||
case mt.MsgsStateReqTypeID:
|
||||
if err := validateFirstVectorCount(b, maxServiceMessageIDs); err != nil {
|
||||
return fmt.Errorf("msgs_state_req vector: %w", err)
|
||||
}
|
||||
var req mt.MsgsStateReq
|
||||
if err := req.Decode(b); err != nil {
|
||||
return fmt.Errorf("decode msgs_state_req: %w", err)
|
||||
|
|
@ -344,6 +421,9 @@ func (s *Server) dispatch(ctx context.Context, cs *connState, c *Conn, msgID int
|
|||
return s.sendMsgsStateInfo(ctx, c, msgID, mergeStateInfo(outgoing, cs.stateInfo(req.MsgIDs)))
|
||||
|
||||
case mt.MsgResendReqTypeID:
|
||||
if err := validateFirstVectorCount(b, maxServiceMessageIDs); err != nil {
|
||||
return fmt.Errorf("msg_resend_req vector: %w", err)
|
||||
}
|
||||
var req mt.MsgResendReq
|
||||
if err := req.Decode(b); err != nil {
|
||||
return fmt.Errorf("decode msg_resend_req: %w", err)
|
||||
|
|
@ -356,19 +436,22 @@ func (s *Server) dispatch(ctx context.Context, cs *connState, c *Conn, msgID int
|
|||
return s.sendMsgsStateInfo(ctx, c, msgID, mergeStateInfo(outgoing, cs.stateInfo(req.MsgIDs)))
|
||||
|
||||
case mt.MsgsStateInfoTypeID:
|
||||
var info mt.MsgsStateInfo
|
||||
if err := info.Decode(b); err != nil {
|
||||
reqMsgID, info, err := msgsStateInfoView(b)
|
||||
if err != nil {
|
||||
return fmt.Errorf("decode msgs_state_info: %w", err)
|
||||
}
|
||||
s.log.Debug("Received msgs_state_info", zap.Int64("req_msg_id", info.ReqMsgID), zap.Int("len", len(info.Info)))
|
||||
s.log.Debug("Received msgs_state_info", zap.Int64("req_msg_id", reqMsgID), zap.Int("len", len(info)))
|
||||
return nil
|
||||
|
||||
case mt.MsgsAllInfoTypeID:
|
||||
var info mt.MsgsAllInfo
|
||||
if err := info.Decode(b); err != nil {
|
||||
count, info, err := msgsAllInfoView(b)
|
||||
if err != nil {
|
||||
return fmt.Errorf("decode msgs_all_info: %w", err)
|
||||
}
|
||||
s.log.Debug("Received msgs_all_info", zap.Int("msg_ids", len(info.MsgIDs)), zap.Int("len", len(info.Info)))
|
||||
if len(info) != count {
|
||||
return fmt.Errorf("decode msgs_all_info: info length %d does not match msg_ids %d", len(info), count)
|
||||
}
|
||||
s.log.Debug("Received msgs_all_info", zap.Int("msg_ids", count), zap.Int("len", len(info)))
|
||||
return nil
|
||||
|
||||
case mt.DestroySessionRequestTypeID:
|
||||
|
|
@ -424,11 +507,228 @@ func (s *Server) dispatch(ctx context.Context, cs *connState, c *Conn, msgID int
|
|||
|
||||
default:
|
||||
ackContent()
|
||||
body := b.Copy()
|
||||
return s.enqueueRPC(ctx, c, msgID, id, body)
|
||||
return s.enqueueRPC(ctx, c, msgID, id, b)
|
||||
}
|
||||
}
|
||||
|
||||
// decodeGZIPWithGlobalBudget reserves the maximum single-wrapper output before
|
||||
// decompression starts. Once the actual size is known the excess reservation is
|
||||
// returned, while the actual output remains charged through recursive dispatch.
|
||||
// This closes the gap where every connection read goroutine could otherwise hold
|
||||
// an unaccounted 10 MiB expansion before the shared RPC scheduler saw the body.
|
||||
func (s *Server) decodeGZIPWithGlobalBudget(b *bin.Buffer) ([]byte, func(), error) {
|
||||
compressed, err := gzipPackedBytesView(b)
|
||||
if err != nil {
|
||||
return nil, func() {}, err
|
||||
}
|
||||
reserved := int64(0)
|
||||
release := func() {
|
||||
if reserved > 0 && s.frameBudget != nil {
|
||||
s.frameBudget.release(reserved)
|
||||
reserved = 0
|
||||
}
|
||||
}
|
||||
if s.frameBudget != nil {
|
||||
reserved, err = s.frameBudget.reserve(maxSingleGZIPExpandedBytes, 0)
|
||||
if err != nil {
|
||||
return nil, func() {}, err
|
||||
}
|
||||
}
|
||||
|
||||
r, err := gzip.NewReader(bytes.NewReader(compressed))
|
||||
if err != nil {
|
||||
release()
|
||||
return nil, func() {}, err
|
||||
}
|
||||
data, readErr := io.ReadAll(io.LimitReader(r, maxSingleGZIPExpandedBytes+1))
|
||||
closeErr := r.Close()
|
||||
if readErr != nil {
|
||||
release()
|
||||
return nil, func() {}, readErr
|
||||
}
|
||||
if closeErr != nil {
|
||||
release()
|
||||
return nil, func() {}, closeErr
|
||||
}
|
||||
if len(data) > maxSingleGZIPExpandedBytes {
|
||||
release()
|
||||
return nil, func() {}, fmt.Errorf("gzip expansion %d exceeds %d", len(data), maxSingleGZIPExpandedBytes)
|
||||
}
|
||||
if reserved > int64(len(data)) {
|
||||
s.frameBudget.release(reserved - int64(len(data)))
|
||||
reserved = int64(len(data))
|
||||
}
|
||||
return data, release, nil
|
||||
}
|
||||
|
||||
// gzipPackedBytesView parses the TL bytes envelope without copying the compressed
|
||||
// payload. proto.GZIP.Decode calls bin.Buffer.Bytes, which duplicates the compressed
|
||||
// frame before allocating the decompressed result.
|
||||
func gzipPackedBytesView(b *bin.Buffer) ([]byte, error) {
|
||||
if b == nil || len(b.Buf) < 5 {
|
||||
return nil, io.ErrUnexpectedEOF
|
||||
}
|
||||
if binary.LittleEndian.Uint32(b.Buf[:4]) != proto.GZIPTypeID {
|
||||
return nil, fmt.Errorf("unexpected gzip constructor %#x", binary.LittleEndian.Uint32(b.Buf[:4]))
|
||||
}
|
||||
payload, _, err := tlBytesView(b.Buf[4:], -1)
|
||||
return payload, err
|
||||
}
|
||||
|
||||
// tlBytesView validates one TL bytes envelope and returns a view into the caller-owned buffer.
|
||||
// maxPayload < 0 means that the enclosing frame budget is the only size limit. The limit is
|
||||
// checked from the encoded length before touching the payload, so service messages cannot make
|
||||
// generated decoders allocate an attacker-selected []byte first and validate it afterwards.
|
||||
func tlBytesView(raw []byte, maxPayload int) ([]byte, int, error) {
|
||||
if len(raw) < 1 {
|
||||
return nil, 0, io.ErrUnexpectedEOF
|
||||
}
|
||||
header, size := 1, int(raw[0])
|
||||
if size == 254 {
|
||||
if len(raw) < 4 {
|
||||
return nil, 0, io.ErrUnexpectedEOF
|
||||
}
|
||||
header = 4
|
||||
size = int(raw[1]) | int(raw[2])<<8 | int(raw[3])<<16
|
||||
} else if size == 255 {
|
||||
return nil, 0, errors.New("invalid TL bytes length marker 255")
|
||||
}
|
||||
if maxPayload >= 0 && size > maxPayload {
|
||||
return nil, 0, fmt.Errorf("TL bytes length %d exceeds %d", size, maxPayload)
|
||||
}
|
||||
padded := (header + size + 3) &^ 3
|
||||
if size < 0 || padded < header || len(raw) < padded {
|
||||
return nil, 0, io.ErrUnexpectedEOF
|
||||
}
|
||||
return raw[header : header+size : header+size], padded, nil
|
||||
}
|
||||
|
||||
// decodeMessageContainerViews parses the container without proto.Message.Decode's per-body
|
||||
// copies. Bodies are immutable views of b and stay alive only for this synchronous dispatch;
|
||||
// enqueueRPC takes its own budgeted copy before returning. Only the exact-size descriptor slice
|
||||
// is new memory, and that allocation is reserved globally first.
|
||||
func (s *Server) decodeMessageContainerViews(b *bin.Buffer, count int) (proto.MessageContainer, func(), error) {
|
||||
release := func() {}
|
||||
if b == nil || len(b.Buf) < 8 {
|
||||
return proto.MessageContainer{}, release, io.ErrUnexpectedEOF
|
||||
}
|
||||
if got := binary.LittleEndian.Uint32(b.Buf[:4]); got != proto.MessageContainerTypeID {
|
||||
return proto.MessageContainer{}, release, fmt.Errorf("unexpected constructor %#x", got)
|
||||
}
|
||||
declared := int(int32(binary.LittleEndian.Uint32(b.Buf[4:8])))
|
||||
if declared != count || count < 0 || count > maxContainerMessages {
|
||||
return proto.MessageContainer{}, release, fmt.Errorf("invalid message count %d", declared)
|
||||
}
|
||||
|
||||
reserved := int64(0)
|
||||
if count > 0 && s.frameBudget != nil {
|
||||
var err error
|
||||
reserved, err = s.frameBudget.reserve(int64(count*containerDescriptorBudgetBytes), 0)
|
||||
if err != nil {
|
||||
return proto.MessageContainer{}, release, err
|
||||
}
|
||||
release = func() {
|
||||
if reserved > 0 {
|
||||
s.frameBudget.release(reserved)
|
||||
reserved = 0
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
messages := make([]proto.Message, count)
|
||||
offset := 8
|
||||
for i := range messages {
|
||||
if len(b.Buf)-offset < 16 {
|
||||
release()
|
||||
return proto.MessageContainer{}, func() {}, io.ErrUnexpectedEOF
|
||||
}
|
||||
id := int64(binary.LittleEndian.Uint64(b.Buf[offset : offset+8]))
|
||||
seqNo := int32(binary.LittleEndian.Uint32(b.Buf[offset+8 : offset+12]))
|
||||
bodyLen := int(int32(binary.LittleEndian.Uint32(b.Buf[offset+12 : offset+16])))
|
||||
offset += 16
|
||||
if bodyLen < 0 || bodyLen > 1024*1024 {
|
||||
release()
|
||||
return proto.MessageContainer{}, func() {}, fmt.Errorf("message length %d is invalid", bodyLen)
|
||||
}
|
||||
if bodyLen > len(b.Buf)-offset {
|
||||
release()
|
||||
return proto.MessageContainer{}, func() {}, io.ErrUnexpectedEOF
|
||||
}
|
||||
bodyEnd := offset + bodyLen
|
||||
messages[i] = proto.Message{
|
||||
ID: id,
|
||||
SeqNo: int(seqNo),
|
||||
Bytes: bodyLen,
|
||||
Body: b.Buf[offset:bodyEnd:bodyEnd],
|
||||
}
|
||||
offset = bodyEnd
|
||||
}
|
||||
return proto.MessageContainer{Messages: messages}, release, nil
|
||||
}
|
||||
|
||||
func msgsStateInfoView(b *bin.Buffer) (int64, []byte, error) {
|
||||
if b == nil || len(b.Buf) < 12 {
|
||||
return 0, nil, io.ErrUnexpectedEOF
|
||||
}
|
||||
if got := binary.LittleEndian.Uint32(b.Buf[:4]); got != mt.MsgsStateInfoTypeID {
|
||||
return 0, nil, fmt.Errorf("unexpected constructor %#x", got)
|
||||
}
|
||||
info, _, err := tlBytesView(b.Buf[12:], maxServiceMessageIDs)
|
||||
if err != nil {
|
||||
return 0, nil, err
|
||||
}
|
||||
return int64(binary.LittleEndian.Uint64(b.Buf[4:12])), info, nil
|
||||
}
|
||||
|
||||
func msgsAllInfoView(b *bin.Buffer) (int, []byte, error) {
|
||||
if err := validateFirstVectorCount(b, maxServiceMessageIDs); err != nil {
|
||||
return 0, nil, fmt.Errorf("vector: %w", err)
|
||||
}
|
||||
count := int(int32(binary.LittleEndian.Uint32(b.Buf[8:12])))
|
||||
// count is already non-negative and capped, but check remaining bytes before multiplying into
|
||||
// an offset so malformed frames cannot produce an out-of-bounds slice.
|
||||
if count > (len(b.Buf)-12)/8 {
|
||||
return 0, nil, io.ErrUnexpectedEOF
|
||||
}
|
||||
offset := 12 + count*8
|
||||
info, _, err := tlBytesView(b.Buf[offset:], maxServiceMessageIDs)
|
||||
if err != nil {
|
||||
return 0, nil, err
|
||||
}
|
||||
return count, info, nil
|
||||
}
|
||||
|
||||
func containerMessageCount(b *bin.Buffer) (int, error) {
|
||||
if b == nil || len(b.Buf) < 8 {
|
||||
return 0, io.ErrUnexpectedEOF
|
||||
}
|
||||
if binary.LittleEndian.Uint32(b.Buf[:4]) != proto.MessageContainerTypeID {
|
||||
return 0, fmt.Errorf("unexpected constructor %#x", binary.LittleEndian.Uint32(b.Buf[:4]))
|
||||
}
|
||||
count := int(int32(binary.LittleEndian.Uint32(b.Buf[4:8])))
|
||||
if count < 0 {
|
||||
return 0, fmt.Errorf("negative message count %d", count)
|
||||
}
|
||||
return count, nil
|
||||
}
|
||||
|
||||
func validateFirstVectorCount(b *bin.Buffer, max int) error {
|
||||
if b == nil || len(b.Buf) < 12 {
|
||||
return io.ErrUnexpectedEOF
|
||||
}
|
||||
if got := binary.LittleEndian.Uint32(b.Buf[4:8]); got != bin.TypeVector {
|
||||
return fmt.Errorf("unexpected vector constructor %#x", got)
|
||||
}
|
||||
count := int(int32(binary.LittleEndian.Uint32(b.Buf[8:12])))
|
||||
if count < 0 {
|
||||
return fmt.Errorf("negative vector count %d", count)
|
||||
}
|
||||
if count > max {
|
||||
return fmt.Errorf("vector count %d exceeds %d", count, max)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func mergeStateInfo(primary, fallback []byte) []byte {
|
||||
if len(primary) == 0 {
|
||||
return fallback
|
||||
|
|
@ -448,7 +748,7 @@ func mergeStateInfo(primary, fallback []byte) []byte {
|
|||
|
||||
// enqueueRPC 把一条 RPC 请求交给连接的 inbound 调度器。typeID 由 dispatch 传入
|
||||
// (已 PeekID 过一次),method 只解析一次并随任务透传,避免同一请求三处重复 PeekID/typeName。
|
||||
func (s *Server) enqueueRPC(ctx context.Context, c *Conn, msgID int64, typeID uint32, body []byte) error {
|
||||
func (s *Server) enqueueRPC(ctx context.Context, c *Conn, msgID int64, typeID uint32, request *bin.Buffer) error {
|
||||
method := s.typeName(typeID)
|
||||
if cached, ok := s.cachedRPCResult(c, msgID); ok {
|
||||
s.log.Info("RPC duplicate replay from session cache",
|
||||
|
|
@ -459,13 +759,48 @@ func (s *Server) enqueueRPC(ctx context.Context, c *Conn, msgID int64, typeID ui
|
|||
)
|
||||
return c.SendEncoded(ctx, proto.MessageServerResponse, cached)
|
||||
}
|
||||
err := c.enqueueInboundRPC(ctx, inboundRPC{
|
||||
method: method,
|
||||
size: len(body),
|
||||
// 两级条数/字节预算必须先于 Copy:对抗客户端不能用大量满尺寸请求在“判断队列满”
|
||||
// 之前制造一轮无上限的临时 body 分配。reservation 在 commit/abort 间唯一持有预算。
|
||||
reservation, err := c.reserveInboundRPC(ctx, method, request.Len())
|
||||
if err != nil {
|
||||
return s.handleInboundRPCAdmissionError(ctx, c, msgID, method, err)
|
||||
}
|
||||
defer reservation.abort()
|
||||
body := request.Copy()
|
||||
responseGate := &rpcResponseGate{}
|
||||
timeoutResponse := func() {
|
||||
if !responseGate.tryTimeout() {
|
||||
return
|
||||
}
|
||||
// 原 task context 已到期,使用有界的新 context 回显明确的可重试超时;
|
||||
// 500 保持 TDesktop 默认重试语义,错误名区分于容量型 FLOOD_WAIT。
|
||||
writeTimeout := c.writeTimeout
|
||||
if writeTimeout <= 0 || writeTimeout > 5*time.Second {
|
||||
writeTimeout = 5 * time.Second
|
||||
}
|
||||
responseCtx, cancel := context.WithTimeout(context.Background(), writeTimeout)
|
||||
defer cancel()
|
||||
if sendErr := s.sendResult(responseCtx, c, msgID, &mt.RPCError{
|
||||
ErrorCode: 500,
|
||||
ErrorMessage: "RPC_TIMEOUT",
|
||||
}); sendErr != nil && !isClientDisconnect(sendErr) {
|
||||
s.log.Debug("Send RPC timeout failed",
|
||||
zap.String("method", method),
|
||||
zap.Int64("msg_id", msgID),
|
||||
zap.String("auth_key_id", c.authKeyHex),
|
||||
zap.Int64("session_id", c.sessionID),
|
||||
zap.Error(sendErr),
|
||||
)
|
||||
}
|
||||
}
|
||||
err = reservation.commit(inboundRPC{
|
||||
method: method,
|
||||
size: len(body),
|
||||
onTimeout: timeoutResponse,
|
||||
run: func(taskCtx context.Context) error {
|
||||
// body 已是 enqueueRPC 入参的独立副本(dispatch 里 b.Copy()),且每个任务只 run 一次,
|
||||
// body 是预算成功后生成的独立副本,且每个任务只 run 一次,
|
||||
// 无需再 append 拷贝;直接复用,省掉一份 inbound 在途内存。
|
||||
if err := s.handleRPC(taskCtx, c, msgID, method, &bin.Buffer{Buf: body}); err != nil {
|
||||
if err := s.handleRPC(taskCtx, c, msgID, method, &bin.Buffer{Buf: body}, responseGate); err != nil {
|
||||
fields := []zap.Field{
|
||||
zap.Int64("msg_id", msgID),
|
||||
zap.String("auth_key_id", c.authKeyHex),
|
||||
|
|
@ -482,8 +817,12 @@ func (s *Server) enqueueRPC(ctx context.Context, c *Conn, msgID int64, typeID ui
|
|||
return nil
|
||||
},
|
||||
})
|
||||
return s.handleInboundRPCAdmissionError(ctx, c, msgID, method, err)
|
||||
}
|
||||
|
||||
func (s *Server) handleInboundRPCAdmissionError(ctx context.Context, c *Conn, msgID int64, method string, err error) error {
|
||||
if errors.Is(err, ErrInboundRPCQueueFull) {
|
||||
s.log.Debug("Inbound RPC queue full",
|
||||
s.log.Debug("Inbound RPC capacity exhausted",
|
||||
zap.String("method", method),
|
||||
zap.Int64("msg_id", msgID),
|
||||
zap.String("auth_key_id", c.authKeyHex),
|
||||
|
|
@ -498,7 +837,7 @@ func (s *Server) enqueueRPC(ctx context.Context, c *Conn, msgID int64, typeID ui
|
|||
}
|
||||
|
||||
// handleRPC 把明文 RPC 请求交给 RPC 路由,并将结果或错误包成 rpc_result 回发。
|
||||
func (s *Server) handleRPC(ctx context.Context, c *Conn, msgID int64, method string, b *bin.Buffer) error {
|
||||
func (s *Server) handleRPC(ctx context.Context, c *Conn, msgID int64, method string, b *bin.Buffer, responseGate *rpcResponseGate) error {
|
||||
if s.rpc == nil {
|
||||
s.log.Warn("No RPC handler configured; dropping request", zap.String("method", method))
|
||||
return nil
|
||||
|
|
@ -534,12 +873,24 @@ func (s *Server) handleRPC(ctx context.Context, c *Conn, msgID int64, method str
|
|||
}
|
||||
fields = dbtrace.AppendZapFields(fields, "", dbStats.Snapshot())
|
||||
|
||||
if ctxErr := ctx.Err(); ctxErr != nil && err != nil {
|
||||
// A canceled request context means the result cannot be delivered. Do not
|
||||
// turn cancellation-derived handler errors into cacheable rpc_error replies.
|
||||
s.log.Info("RPC canceled", append(fields, zap.NamedError("dispatch_error", err), zap.NamedError("context_error", ctxErr))...)
|
||||
if ctxErr := ctx.Err(); ctxErr != nil {
|
||||
// A canceled request context means neither a success nor an error can be delivered
|
||||
// with this expired context. In particular, do not cache a late successful result and
|
||||
// hand it to outbound: a past write deadline would correctly poison that transport and
|
||||
// could prevent the scheduler's fresh-context RPC_TIMEOUT response from being sent.
|
||||
cancelFields := append(fields, zap.NamedError("context_error", ctxErr))
|
||||
if err != nil {
|
||||
cancelFields = append(cancelFields, zap.NamedError("dispatch_error", err))
|
||||
}
|
||||
s.log.Info("RPC canceled", cancelFields...)
|
||||
return ctxErr
|
||||
}
|
||||
// A deadline callback may have already emitted RPC_TIMEOUT while Dispatch was returning.
|
||||
// Claim the single normal-response slot before serializing any success/error rpc_result.
|
||||
if responseGate != nil && !responseGate.tryNormal() {
|
||||
s.log.Info("RPC result suppressed after timeout", fields...)
|
||||
return context.DeadlineExceeded
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
var rpcErr *tgerr.Error
|
||||
|
|
@ -565,6 +916,21 @@ func (s *Server) handleRPC(ctx context.Context, c *Conn, msgID int64, method str
|
|||
return nil
|
||||
}
|
||||
|
||||
// rpcResponseGate guarantees exactly one terminal rpc_result per request. A running deadline
|
||||
// races legitimately with a handler completing at the boundary; whichever path claims state
|
||||
// first owns the response, and the other path becomes a no-op.
|
||||
type rpcResponseGate struct {
|
||||
state atomic.Uint32
|
||||
}
|
||||
|
||||
func (g *rpcResponseGate) tryNormal() bool {
|
||||
return g == nil || g.state.CompareAndSwap(0, 1)
|
||||
}
|
||||
|
||||
func (g *rpcResponseGate) tryTimeout() bool {
|
||||
return g != nil && g.state.CompareAndSwap(0, 2)
|
||||
}
|
||||
|
||||
// sendResult 把 RPC 结果包成 rpc_result 并加密回发。
|
||||
func (s *Server) sendResult(ctx context.Context, c *Conn, reqMsgID int64, result bin.Encoder) error {
|
||||
encoded, err := s.encodeRPCResult(c, reqMsgID, result)
|
||||
|
|
|
|||
|
|
@ -2,9 +2,9 @@ package mtprotoedge
|
|||
|
||||
import (
|
||||
"context"
|
||||
"encoding/binary"
|
||||
"errors"
|
||||
"fmt"
|
||||
"sync"
|
||||
|
||||
"go.uber.org/zap"
|
||||
|
||||
|
|
@ -12,7 +12,6 @@ import (
|
|||
"github.com/gotd/td/crypto"
|
||||
"github.com/gotd/td/exchange"
|
||||
"github.com/gotd/td/mt"
|
||||
"github.com/gotd/td/proto"
|
||||
"github.com/gotd/td/proto/codec"
|
||||
"github.com/gotd/td/transport"
|
||||
|
||||
|
|
@ -31,7 +30,8 @@ func peekAuthKeyID(b *bin.Buffer) (id [8]byte, err error) {
|
|||
// handleExchange 在收到 auth_key_id==0 的首帧后执行服务端 MTProto 密钥交换。
|
||||
//
|
||||
// first 是已读取的首帧(req_pq*),通过 bufferedConn 交还给 exchange 流程,
|
||||
// 使其能从头读取握手消息。成功后将 auth key + server salt 落入 AuthKeyStore。
|
||||
// 使其能从头读取握手消息。auth key + server salt 会在 DhGenOk 发出前落入
|
||||
// AuthKeyStore;持久化失败时不向客户端确认握手成功。
|
||||
func (s *Server) handleExchange(ctx context.Context, conn transport.Conn, first *bin.Buffer) (*bin.Buffer, error) {
|
||||
if s.key.Zero() {
|
||||
s.log.Error("Key exchange requested but server RSA key is not configured")
|
||||
|
|
@ -62,11 +62,6 @@ func (s *Server) handleExchange(ctx context.Context, conn transport.Conn, first
|
|||
var encErr *exchange.UnexpectedEncryptedError
|
||||
if errors.As(err, &encErr) {
|
||||
replay := encErr.Frame
|
||||
if len(replay) == 0 {
|
||||
if lf := buffered.lastFrame(); lf != nil {
|
||||
replay = lf.Buf
|
||||
}
|
||||
}
|
||||
if len(replay) > 0 {
|
||||
s.log.Debug("Key exchange interrupted by encrypted frame; replaying as existing session")
|
||||
return &bin.Buffer{Buf: replay}, nil
|
||||
|
|
@ -99,7 +94,7 @@ func (s *Server) handleExchange(ctx context.Context, conn transport.Conn, first
|
|||
zap.Duration("dur", s.clock.Now().Sub(start)),
|
||||
)
|
||||
|
||||
return nil, s.authKeys.Save(ctx, authKeyData(res.Key, res.ServerSalt, s.clock.Now().Unix()))
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
// authKeyData 把握手结果转换为 store 记录。
|
||||
|
|
@ -142,9 +137,7 @@ var errTooManyHandshakeReqPQ = errors.New("too many req_pq frames in one handsha
|
|||
// 用于密钥交换:serveConn 已读首帧用于 peek auth_key_id,再 push 回来交给 exchange。
|
||||
type bufferedConn struct {
|
||||
transport.Conn
|
||||
mu sync.Mutex
|
||||
pending []bin.Buffer
|
||||
last bin.Buffer
|
||||
reqPQCount int // 本次握手已见 req_pq(_multi) 帧数;只在握手期访问(Recv 单 goroutine)
|
||||
}
|
||||
|
||||
|
|
@ -153,32 +146,38 @@ func newBufferedConn(conn transport.Conn) *bufferedConn {
|
|||
}
|
||||
|
||||
func (c *bufferedConn) push(b *bin.Buffer) {
|
||||
c.mu.Lock()
|
||||
c.pending = append(c.pending, bin.Buffer{Buf: b.Copy()})
|
||||
c.mu.Unlock()
|
||||
if b == nil {
|
||||
return
|
||||
}
|
||||
// serveConn is synchronously blocked in handleExchange, so the first frame's
|
||||
// backing remains stable until the exchange returns. Keep a slice view instead
|
||||
// of copying an attacker-sized transport frame.
|
||||
buf := b.Buf
|
||||
b.Buf = nil // transfer ownership; serveConn must not pin the frame after next Recv releases it
|
||||
c.pending = append(c.pending, bin.Buffer{Buf: buf})
|
||||
}
|
||||
|
||||
// Recv 优先返回已 push 的帧(FIFO),耗尽后读取底层连接。
|
||||
func (c *bufferedConn) Recv(ctx context.Context, b *bin.Buffer) error {
|
||||
for {
|
||||
c.mu.Lock()
|
||||
if len(c.pending) > 0 {
|
||||
e := c.pending[0]
|
||||
c.pending[0] = bin.Buffer{}
|
||||
c.pending = c.pending[1:]
|
||||
c.last.ResetTo(e.Copy())
|
||||
c.mu.Unlock()
|
||||
b.ResetTo(e.Buf)
|
||||
} else {
|
||||
c.mu.Unlock()
|
||||
if err := c.Conn.Recv(ctx, b); err != nil {
|
||||
return err
|
||||
}
|
||||
c.mu.Lock()
|
||||
c.last.ResetTo(b.Copy())
|
||||
c.mu.Unlock()
|
||||
}
|
||||
|
||||
if isUnencryptedMsgsAckFrame(b) {
|
||||
// The ack is intentionally ignored during exchange. Drop its transport
|
||||
// backing and shrink the retained high-water charge before the next Recv;
|
||||
// otherwise a large trailing frame can consume global admission budget for
|
||||
// the rest of a CPU-heavy key exchange even though no backing remains live.
|
||||
b.Buf = nil
|
||||
retainInboundFrameBackings(c.Conn, b)
|
||||
continue
|
||||
}
|
||||
// req_pq 计数上界:仅在握手期生效(bufferedConn 只用于密钥交换),且 payload id 探测
|
||||
|
|
@ -196,21 +195,17 @@ func (c *bufferedConn) Recv(ctx context.Context, b *bin.Buffer) error {
|
|||
// unencryptedPayloadID 返回未加密消息(auth_key_id==0)内层 TL payload 的 type id。
|
||||
// 非未加密消息 / 解码失败时 ok=false。
|
||||
func unencryptedPayloadID(frame *bin.Buffer) (uint32, bool) {
|
||||
authKeyID, err := peekAuthKeyID(frame)
|
||||
if err != nil || authKeyID != emptyAuthKeyID {
|
||||
if frame == nil || len(frame.Buf) < 24 {
|
||||
return 0, false
|
||||
}
|
||||
var msg proto.UnencryptedMessage
|
||||
cp := &bin.Buffer{Buf: frame.Copy()}
|
||||
if err := msg.Decode(cp); err != nil {
|
||||
if binary.LittleEndian.Uint64(frame.Buf[:8]) != 0 {
|
||||
return 0, false
|
||||
}
|
||||
payload := &bin.Buffer{Buf: msg.MessageData}
|
||||
id, err := payload.PeekID()
|
||||
if err != nil {
|
||||
dataLen := int64(int32(binary.LittleEndian.Uint32(frame.Buf[16:20])))
|
||||
if dataLen < 4 || dataLen > int64(len(frame.Buf)-20) {
|
||||
return 0, false
|
||||
}
|
||||
return id, true
|
||||
return binary.LittleEndian.Uint32(frame.Buf[20:24]), true
|
||||
}
|
||||
|
||||
func isUnencryptedMsgsAckFrame(frame *bin.Buffer) bool {
|
||||
|
|
@ -222,12 +217,3 @@ func isUnencryptedReqPQFrame(frame *bin.Buffer) bool {
|
|||
id, ok := unencryptedPayloadID(frame)
|
||||
return ok && (id == mt.ReqPqRequestTypeID || id == mt.ReqPqMultiRequestTypeID)
|
||||
}
|
||||
|
||||
func (c *bufferedConn) lastFrame() *bin.Buffer {
|
||||
c.mu.Lock()
|
||||
defer c.mu.Unlock()
|
||||
if c.last.Len() == 0 {
|
||||
return nil
|
||||
}
|
||||
return &bin.Buffer{Buf: c.last.Copy()}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -31,27 +31,42 @@ import (
|
|||
// matches this server DC.
|
||||
func (s *Server) runServerExchange(ctx context.Context, conn transport.Conn) (exchange.ServerExchangeResult, error) {
|
||||
ex := serverExchangeCompat{
|
||||
conn: conn,
|
||||
clock: s.clock,
|
||||
rand: s.rand,
|
||||
timeout: exchange.DefaultTimeout,
|
||||
key: s.key,
|
||||
dc: s.dc,
|
||||
log: s.log.Named("exchange"),
|
||||
rng: compatServerRNG{rand: s.rand},
|
||||
conn: conn,
|
||||
clock: s.clock,
|
||||
rand: s.rand,
|
||||
timeout: exchange.DefaultTimeout,
|
||||
key: s.key,
|
||||
dc: s.dc,
|
||||
log: s.log.Named("exchange"),
|
||||
rng: compatServerRNG{rand: s.rand},
|
||||
commitKey: s.commitExchangeAuthKey,
|
||||
}
|
||||
return ex.run(ctx)
|
||||
}
|
||||
|
||||
// commitExchangeAuthKey is the durable commit point of the server exchange.
|
||||
// It must complete before DhGenOk is put on the wire: after that response the
|
||||
// client is allowed to immediately use the new key, possibly on another TCP
|
||||
// connection. Persisting after the response creates a split-brain window when
|
||||
// storage fails or the process exits between those two operations.
|
||||
func (s *Server) commitExchangeAuthKey(ctx context.Context, result exchange.ServerExchangeResult) error {
|
||||
createdAt := s.clock.Now().Unix()
|
||||
if err := s.authKeys.Save(ctx, authKeyData(result.Key, result.ServerSalt, createdAt)); err != nil {
|
||||
return fmt.Errorf("persist auth key before DhGenOk: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
type serverExchangeCompat struct {
|
||||
conn transport.Conn
|
||||
clock clock.Clock
|
||||
rand io.Reader
|
||||
timeout time.Duration
|
||||
key exchange.PrivateKey
|
||||
dc int
|
||||
log *zap.Logger
|
||||
rng compatServerRNG
|
||||
conn transport.Conn
|
||||
clock clock.Clock
|
||||
rand io.Reader
|
||||
timeout time.Duration
|
||||
key exchange.PrivateKey
|
||||
dc int
|
||||
log *zap.Logger
|
||||
rng compatServerRNG
|
||||
commitKey func(context.Context, exchange.ServerExchangeResult) error
|
||||
}
|
||||
|
||||
func (s serverExchangeCompat) run(ctx context.Context) (exchange.ServerExchangeResult, error) {
|
||||
|
|
@ -193,6 +208,21 @@ SendResPQ:
|
|||
return exchange.ServerExchangeResult{}, wrapKeyNotFound(err)
|
||||
}
|
||||
|
||||
serverResult := exchange.ServerExchangeResult{
|
||||
Key: authKey.WithID(),
|
||||
ServerSalt: crypto.ServerSalt(innerData.NewNonce, serverNonce),
|
||||
}
|
||||
// DhGenOk is the externally visible commit acknowledgement. Require a
|
||||
// durable key commit before sending it, rather than allowing callers to
|
||||
// persist after run returns. A nil hook is rejected so a future call site
|
||||
// cannot accidentally reintroduce the unsafe ordering.
|
||||
if s.commitKey == nil {
|
||||
return exchange.ServerExchangeResult{}, gofaster.New("auth key commit hook is required before DhGenOk")
|
||||
}
|
||||
if err := s.commitKey(ctx, serverResult); err != nil {
|
||||
return exchange.ServerExchangeResult{}, err
|
||||
}
|
||||
|
||||
s.log.Debug("Sending DhGenOk")
|
||||
if err := s.writeUnencrypted(ctx, b, &mt.DhGenOk{
|
||||
Nonce: req.Nonce,
|
||||
|
|
@ -202,11 +232,7 @@ SendResPQ:
|
|||
return exchange.ServerExchangeResult{}, err
|
||||
}
|
||||
|
||||
serverSalt := crypto.ServerSalt(innerData.NewNonce, serverNonce)
|
||||
return exchange.ServerExchangeResult{
|
||||
Key: authKey.WithID(),
|
||||
ServerSalt: serverSalt,
|
||||
}, nil
|
||||
return serverResult, nil
|
||||
}
|
||||
|
||||
func (s serverExchangeCompat) validatePQInnerDataDC(d mt.PQInnerDataClass) error {
|
||||
|
|
@ -278,9 +304,14 @@ func (s serverExchangeCompat) readUnencrypted(ctx context.Context, b *bin.Buffer
|
|||
|
||||
var keyID [8]byte
|
||||
if err := b.PeekN(keyID[:], len(keyID)); err == nil && keyID != ([8]byte{}) {
|
||||
// The exchange aborts immediately on an encrypted frame, so transfer the received backing
|
||||
// to the replay error instead of making an unbudgeted near-transport-limit copy. serveConn
|
||||
// keeps the existing frame reservation until replay dispatch has finished.
|
||||
frame := b.Buf
|
||||
b.Buf = nil
|
||||
return &exchange.UnexpectedEncryptedError{
|
||||
AuthKeyID: keyID,
|
||||
Frame: append([]byte(nil), b.Buf...),
|
||||
Frame: frame,
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ import (
|
|||
"context"
|
||||
"crypto/rand"
|
||||
"crypto/rsa"
|
||||
"encoding/binary"
|
||||
"errors"
|
||||
"net"
|
||||
"testing"
|
||||
|
|
@ -106,6 +107,214 @@ func TestKeyExchange(t *testing.T) {
|
|||
}
|
||||
}
|
||||
|
||||
type authKeySaveContextObservation struct {
|
||||
hasDeadline bool
|
||||
deadline time.Time
|
||||
}
|
||||
|
||||
type observingAuthKeyStore struct {
|
||||
store.AuthKeyStore
|
||||
saveContext chan authKeySaveContextObservation
|
||||
}
|
||||
|
||||
func (s *observingAuthKeyStore) Save(ctx context.Context, key store.AuthKeyData) error {
|
||||
deadline, hasDeadline := ctx.Deadline()
|
||||
select {
|
||||
case s.saveContext <- authKeySaveContextObservation{hasDeadline: hasDeadline, deadline: deadline}:
|
||||
default:
|
||||
}
|
||||
return s.AuthKeyStore.Save(ctx, key)
|
||||
}
|
||||
|
||||
type gatedAuthKeyStore struct {
|
||||
store.AuthKeyStore
|
||||
entered chan store.AuthKeyData
|
||||
release chan struct{}
|
||||
saveErr error
|
||||
}
|
||||
|
||||
type ownershipFrameConn struct {
|
||||
transport.Conn
|
||||
frame []byte
|
||||
}
|
||||
|
||||
func (c *ownershipFrameConn) Recv(_ context.Context, b *bin.Buffer) error {
|
||||
b.ResetTo(c.frame)
|
||||
return nil
|
||||
}
|
||||
|
||||
func TestExchangeEncryptedReplayTransfersFrameOwnership(t *testing.T) {
|
||||
backing := make([]byte, 64)
|
||||
copy(backing[:8], []byte{1, 2, 3, 4, 5, 6, 7, 8})
|
||||
conn := &ownershipFrameConn{frame: backing}
|
||||
ex := serverExchangeCompat{conn: conn, timeout: time.Second}
|
||||
var b bin.Buffer
|
||||
err := ex.readUnencrypted(context.Background(), &b, &compatReqPQ{})
|
||||
var encrypted *exchange.UnexpectedEncryptedError
|
||||
if !errors.As(err, &encrypted) {
|
||||
t.Fatalf("read encrypted frame err = %v, want UnexpectedEncryptedError", err)
|
||||
}
|
||||
if len(encrypted.Frame) != len(backing) || &encrypted.Frame[0] != &backing[0] {
|
||||
t.Fatal("encrypted replay copied the received frame instead of transferring ownership")
|
||||
}
|
||||
if b.Buf != nil {
|
||||
t.Fatal("exchange buffer retained transferred encrypted frame backing")
|
||||
}
|
||||
}
|
||||
|
||||
func (s *gatedAuthKeyStore) Save(ctx context.Context, key store.AuthKeyData) error {
|
||||
select {
|
||||
case s.entered <- key:
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
}
|
||||
select {
|
||||
case <-s.release:
|
||||
case <-ctx.Done():
|
||||
return ctx.Err()
|
||||
}
|
||||
if s.saveErr != nil {
|
||||
return s.saveErr
|
||||
}
|
||||
return s.AuthKeyStore.Save(ctx, key)
|
||||
}
|
||||
|
||||
// TestKeyExchangeDoesNotAcknowledgeBeforeAuthKeyCommit pins the protocol commit
|
||||
// boundary: while durable Save is blocked, the client must not receive DhGenOk
|
||||
// and therefore must not report a successful exchange.
|
||||
func TestKeyExchangeDoesNotAcknowledgeBeforeAuthKeyCommit(t *testing.T) {
|
||||
base := memory.NewAuthKeyStore()
|
||||
keys := &gatedAuthKeyStore{
|
||||
AuthKeyStore: base,
|
||||
entered: make(chan store.AuthKeyData, 1),
|
||||
release: make(chan struct{}, 1),
|
||||
}
|
||||
addr, pub, _ := startTestServer(t, Options{DC: 2, AuthKeys: keys})
|
||||
conn := dialTransportOnly(t, addr)
|
||||
|
||||
type exchangeOutcome struct {
|
||||
result exchange.ClientExchangeResult
|
||||
err error
|
||||
}
|
||||
outcome := make(chan exchangeOutcome, 1)
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
defer cancel()
|
||||
go func() {
|
||||
result, err := exchange.NewExchanger(conn, 2).
|
||||
WithRand(rand.Reader).
|
||||
Client([]exchange.PublicKey{pub}).
|
||||
Run(ctx)
|
||||
outcome <- exchangeOutcome{result: result, err: err}
|
||||
}()
|
||||
|
||||
var pending store.AuthKeyData
|
||||
select {
|
||||
case pending = <-keys.entered:
|
||||
case <-time.After(5 * time.Second):
|
||||
t.Fatal("AuthKeyStore.Save was not reached")
|
||||
}
|
||||
defer func() {
|
||||
select {
|
||||
case keys.release <- struct{}{}:
|
||||
default:
|
||||
}
|
||||
}()
|
||||
|
||||
select {
|
||||
case got := <-outcome:
|
||||
t.Fatalf("client exchange completed before auth key commit: err=%v", got.err)
|
||||
case <-time.After(150 * time.Millisecond):
|
||||
}
|
||||
if _, found, err := base.Get(context.Background(), pending.ID); err != nil {
|
||||
t.Fatalf("Get before commit: %v", err)
|
||||
} else if found {
|
||||
t.Fatal("auth key became visible while durable Save was blocked")
|
||||
}
|
||||
|
||||
keys.release <- struct{}{}
|
||||
select {
|
||||
case got := <-outcome:
|
||||
if got.err != nil {
|
||||
t.Fatalf("client exchange after commit: %v", got.err)
|
||||
}
|
||||
if got.result.AuthKey.ID != pending.ID {
|
||||
t.Fatalf("committed auth key id = %x, client got %x", pending.ID, got.result.AuthKey.ID)
|
||||
}
|
||||
case <-time.After(5 * time.Second):
|
||||
t.Fatal("client exchange did not finish after auth key commit")
|
||||
}
|
||||
if _, found, err := base.Get(context.Background(), pending.ID); err != nil {
|
||||
t.Fatalf("Get after commit: %v", err)
|
||||
} else if !found {
|
||||
t.Fatal("auth key is not durable after successful client exchange")
|
||||
}
|
||||
}
|
||||
|
||||
// TestKeyExchangeAuthKeyCommitFailureWithholdsDhGenOk proves the failure side
|
||||
// of the same invariant. The client must not observe success if storage rejects
|
||||
// the key; the server closes this exchange and lets the client retry cleanly.
|
||||
func TestKeyExchangeAuthKeyCommitFailureWithholdsDhGenOk(t *testing.T) {
|
||||
base := memory.NewAuthKeyStore()
|
||||
keys := &gatedAuthKeyStore{
|
||||
AuthKeyStore: base,
|
||||
entered: make(chan store.AuthKeyData, 1),
|
||||
release: make(chan struct{}, 1),
|
||||
saveErr: errors.New("injected auth key persistence failure"),
|
||||
}
|
||||
keys.release <- struct{}{}
|
||||
addr, pub, _ := startTestServer(t, Options{DC: 2, AuthKeys: keys})
|
||||
conn := dialTransportOnly(t, addr)
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
_, err := exchange.NewExchanger(conn, 2).
|
||||
WithRand(rand.Reader).
|
||||
Client([]exchange.PublicKey{pub}).
|
||||
Run(ctx)
|
||||
if err == nil {
|
||||
t.Fatal("client exchange succeeded even though auth key commit failed")
|
||||
}
|
||||
|
||||
select {
|
||||
case attempted := <-keys.entered:
|
||||
if _, found, getErr := base.Get(context.Background(), attempted.ID); getErr != nil {
|
||||
t.Fatalf("Get failed key: %v", getErr)
|
||||
} else if found {
|
||||
t.Fatal("failed auth key commit became visible")
|
||||
}
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("AuthKeyStore.Save was not attempted")
|
||||
}
|
||||
}
|
||||
|
||||
func TestKeyExchangeAuthKeySaveUsesHandshakeDeadline(t *testing.T) {
|
||||
const handshakeMax = 10 * time.Second
|
||||
observed := make(chan authKeySaveContextObservation, 1)
|
||||
keys := &observingAuthKeyStore{
|
||||
AuthKeyStore: memory.NewAuthKeyStore(),
|
||||
saveContext: observed,
|
||||
}
|
||||
addr, pub, _ := startTestServer(t, Options{
|
||||
DC: 2,
|
||||
AuthKeys: keys,
|
||||
HandshakeMaxDuration: handshakeMax,
|
||||
})
|
||||
|
||||
_, _, _ = dialHandshake(t, addr, 2, pub)
|
||||
select {
|
||||
case got := <-observed:
|
||||
if !got.hasDeadline {
|
||||
t.Fatal("AuthKeyStore.Save context has no handshake deadline")
|
||||
}
|
||||
remaining := time.Until(got.deadline)
|
||||
if remaining <= 0 || remaining > handshakeMax {
|
||||
t.Fatalf("AuthKeyStore.Save deadline remaining = %v, want (0, %v]", remaining, handshakeMax)
|
||||
}
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("AuthKeyStore.Save was not called")
|
||||
}
|
||||
}
|
||||
|
||||
func TestKeyExchangeAcceptsAndroidMediaTempNegativeDC(t *testing.T) {
|
||||
const dc = 2
|
||||
addr, pub, srv := startTestServer(t, Options{DC: dc})
|
||||
|
|
@ -227,6 +436,82 @@ func TestKeyExchangeIgnoresUnencryptedMsgsAck(t *testing.T) {
|
|||
}
|
||||
}
|
||||
|
||||
func TestBufferedExchangePushTransfersFrameOwnershipWithoutCopy(t *testing.T) {
|
||||
backing := make([]byte, 64)
|
||||
for i := range backing {
|
||||
backing[i] = byte(i)
|
||||
}
|
||||
source := &bin.Buffer{Buf: backing}
|
||||
buffered := newBufferedConn(nil)
|
||||
buffered.push(source)
|
||||
if source.Buf != nil {
|
||||
t.Fatal("push retained ownership in the source buffer")
|
||||
}
|
||||
|
||||
var got bin.Buffer
|
||||
if err := buffered.Recv(context.Background(), &got); err != nil {
|
||||
t.Fatalf("Recv pending frame: %v", err)
|
||||
}
|
||||
if len(got.Buf) != len(backing) || &got.Buf[0] != &backing[0] {
|
||||
t.Fatal("pending frame was copied instead of transferring its backing")
|
||||
}
|
||||
if len(buffered.pending) != 0 || cap(buffered.pending) != 0 {
|
||||
t.Fatalf("consumed pending ownership retained: len=%d cap=%d", len(buffered.pending), cap(buffered.pending))
|
||||
}
|
||||
}
|
||||
|
||||
func TestBufferedExchangeLargeTrailingMsgsAckReleasesFrameBeforeNextRecv(t *testing.T) {
|
||||
encodeUnencrypted := func(msg bin.Encoder, msgID int64) []byte {
|
||||
var payload bin.Buffer
|
||||
if err := msg.Encode(&payload); err != nil {
|
||||
t.Fatalf("encode payload: %v", err)
|
||||
}
|
||||
var frame bin.Buffer
|
||||
if err := (tgproto.UnencryptedMessage{MessageID: msgID, MessageData: payload.Raw()}).Encode(&frame); err != nil {
|
||||
t.Fatalf("encode unencrypted frame: %v", err)
|
||||
}
|
||||
return frame.Copy()
|
||||
}
|
||||
intermediate := func(frame []byte) []byte {
|
||||
packet := make([]byte, bin.Word+len(frame))
|
||||
binary.LittleEndian.PutUint32(packet, uint32(len(frame)))
|
||||
copy(packet[bin.Word:], frame)
|
||||
return packet
|
||||
}
|
||||
|
||||
// Make the ignored ack larger than the per-codec retained-buffer threshold. The following
|
||||
// small req_pq frame forces bufferedConn to cross the next-Recv ownership boundary while the
|
||||
// same destination bin.Buffer is reused.
|
||||
ids := make([]int64, 300_000)
|
||||
for i := range ids {
|
||||
ids[i] = int64(i + 1)
|
||||
}
|
||||
ackFrame := encodeUnencrypted(&mt.MsgsAck{MsgIDs: ids}, 4)
|
||||
reqFrame := encodeUnencrypted(&mt.ReqPqMultiRequest{}, 8)
|
||||
packet := append(intermediate(ackFrame), intermediate(reqFrame)...)
|
||||
budget := newInboundFrameBudget(2 * int64(len(ackFrame)))
|
||||
conn, _ := newFrameBudgetTestTransport(packet, &quickAckIntermediateCodec{}, budget)
|
||||
buffered := newBufferedConn(conn)
|
||||
|
||||
var got bin.Buffer
|
||||
if err := buffered.Recv(context.Background(), &got); err != nil {
|
||||
t.Fatalf("Recv after large msgs_ack: %v", err)
|
||||
}
|
||||
if id, ok := unencryptedPayloadID(&got); !ok || id != mt.ReqPqMultiRequestTypeID {
|
||||
t.Fatalf("returned frame type = 0x%x ok=%v, want req_pq_multi", id, ok)
|
||||
}
|
||||
if used, want := budget.usedBytes(), 2*int64(len(reqFrame)); used != want {
|
||||
t.Fatalf("inbound budget after skipped ack = %d, want only next frame %d", used, want)
|
||||
}
|
||||
if cap(got.Buf) >= len(ackFrame)/2 {
|
||||
t.Fatalf("large ignored ack backing retained by next frame: cap=%d ack=%d", cap(got.Buf), len(ackFrame))
|
||||
}
|
||||
conn.releaseInboundFrame()
|
||||
if used := budget.usedBytes(); used != 0 {
|
||||
t.Fatalf("inbound budget after final ownership release = %d, want 0", used)
|
||||
}
|
||||
}
|
||||
|
||||
type ackingExchangeConn struct {
|
||||
transport.Conn
|
||||
t *testing.T
|
||||
|
|
|
|||
251
internal/mtprotoedge/frame_budget.go
Normal file
251
internal/mtprotoedge/frame_budget.go
Normal file
|
|
@ -0,0 +1,251 @@
|
|||
package mtprotoedge
|
||||
|
||||
import (
|
||||
"encoding/binary"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"sync/atomic"
|
||||
|
||||
"github.com/gotd/td/bin"
|
||||
"github.com/gotd/td/proto/codec"
|
||||
"github.com/gotd/td/transport"
|
||||
)
|
||||
|
||||
const defaultInboundFrameGlobalMaxBytes int64 = 512 << 20
|
||||
|
||||
var (
|
||||
// ErrInboundFrameBudgetExceeded means the process-wide wire+plaintext reservation for a
|
||||
// newly announced transport frame could not be acquired. The length prefix has been read,
|
||||
// but the payload buffer has not been allocated and the connection must be closed.
|
||||
ErrInboundFrameBudgetExceeded = errors.New("inbound frame global byte budget exceeded")
|
||||
|
||||
errInboundFrameCodecUnsupported = errors.New("transport codec cannot preflight inbound frame length")
|
||||
errInboundFrameNotReserved = errors.New("transport codec returned a frame without reserving inbound bytes")
|
||||
)
|
||||
|
||||
// InboundFrameBudgetedCodec is the fail-safe extension point for a custom Options.Codec.
|
||||
// Implementations must parse and validate the frame length, call reserve exactly once before
|
||||
// allocating or growing the payload buffer, and keep the reservation valid until Read returns.
|
||||
// Built-in abridged/intermediate/padded-intermediate/full codecs are recognized directly.
|
||||
type InboundFrameBudgetedCodec interface {
|
||||
transport.Codec
|
||||
ReadWithInboundFrameBudget(r io.Reader, b *bin.Buffer, reserve func(wireBytes, plaintextBytes int64) error) error
|
||||
}
|
||||
|
||||
// inboundFrameBudget accounts the two per-frame buffers that can coexist while an encrypted
|
||||
// request is handled: transport/wire bytes and decrypted plaintext. It deliberately charges the
|
||||
// maximum plaintext size announced by framing even for an unencrypted handshake frame; that
|
||||
// conservative rule makes admission independent of auth state and prevents allocation before
|
||||
// auth_key_id can be inspected.
|
||||
type inboundFrameBudget struct {
|
||||
max int64
|
||||
used atomic.Int64
|
||||
}
|
||||
|
||||
func newInboundFrameBudget(max int64) *inboundFrameBudget {
|
||||
if max <= 0 {
|
||||
max = defaultInboundFrameGlobalMaxBytes
|
||||
}
|
||||
return &inboundFrameBudget{max: max}
|
||||
}
|
||||
|
||||
func (b *inboundFrameBudget) reserve(wireBytes, plaintextBytes int64) (int64, error) {
|
||||
return b.growReservation(0, wireBytes, plaintextBytes)
|
||||
}
|
||||
|
||||
// growReservation atomically raises one connection's existing retained/frame reservation to
|
||||
// cover a newly announced frame. Keeping the old charge until this transition is what makes a
|
||||
// reused transport/plaintext backing remain accounted between frames; a small next frame cannot
|
||||
// release a previously large allocation while still retaining its capacity.
|
||||
func (b *inboundFrameBudget) growReservation(current, wireBytes, plaintextBytes int64) (int64, error) {
|
||||
if current < 0 || wireBytes <= 0 || plaintextBytes < 0 || wireBytes > b.max || plaintextBytes > b.max-wireBytes {
|
||||
return 0, fmt.Errorf("%w: wire=%d plaintext=%d limit=%d", ErrInboundFrameBudgetExceeded, wireBytes, plaintextBytes, b.max)
|
||||
}
|
||||
target := wireBytes + plaintextBytes
|
||||
if target <= current {
|
||||
return current, nil
|
||||
}
|
||||
n := target - current
|
||||
|
||||
for {
|
||||
used := b.used.Load()
|
||||
if n > b.max-used {
|
||||
return 0, fmt.Errorf("%w: requested=%d used=%d limit=%d", ErrInboundFrameBudgetExceeded, n, used, b.max)
|
||||
}
|
||||
if b.used.CompareAndSwap(used, used+n) {
|
||||
return target, nil
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (b *inboundFrameBudget) release(n int64) {
|
||||
if n == 0 {
|
||||
return
|
||||
}
|
||||
used := b.used.Add(-n)
|
||||
if used < 0 {
|
||||
// This is an internal ownership invariant, not recoverable input. A negative value would
|
||||
// silently disable admission for subsequent frames, so fail loudly during development.
|
||||
panic("mtprotoedge: inbound frame budget released more than reserved")
|
||||
}
|
||||
}
|
||||
|
||||
func (b *inboundFrameBudget) usedBytes() int64 {
|
||||
return b.used.Load()
|
||||
}
|
||||
|
||||
type inboundFrameCodecKind uint8
|
||||
|
||||
const (
|
||||
inboundFrameCodecUnknown inboundFrameCodecKind = iota
|
||||
inboundFrameCodecQuickAckAbridged
|
||||
inboundFrameCodecAbridged
|
||||
inboundFrameCodecIntermediate
|
||||
inboundFrameCodecPaddedIntermediate
|
||||
inboundFrameCodecFull
|
||||
inboundFrameCodecCustom
|
||||
)
|
||||
|
||||
func classifyInboundFrameCodec(c transport.Codec) inboundFrameCodecKind {
|
||||
switch v := c.(type) {
|
||||
case *quickAckAbridgedCodec:
|
||||
return inboundFrameCodecQuickAckAbridged
|
||||
case codec.Abridged, *codec.Abridged:
|
||||
return inboundFrameCodecAbridged
|
||||
case *quickAckIntermediateCodec, codec.Intermediate, *codec.Intermediate:
|
||||
return inboundFrameCodecIntermediate
|
||||
case *quickAckPaddedIntermediateCodec, codec.PaddedIntermediate, *codec.PaddedIntermediate:
|
||||
return inboundFrameCodecPaddedIntermediate
|
||||
case *codec.Full:
|
||||
return inboundFrameCodecFull
|
||||
case codec.NoHeader:
|
||||
return classifyInboundFrameCodec(v.Codec)
|
||||
case *codec.NoHeader:
|
||||
if v == nil {
|
||||
return inboundFrameCodecUnknown
|
||||
}
|
||||
return classifyInboundFrameCodec(v.Codec)
|
||||
case InboundFrameBudgetedCodec:
|
||||
return inboundFrameCodecCustom
|
||||
default:
|
||||
return inboundFrameCodecUnknown
|
||||
}
|
||||
}
|
||||
|
||||
func unwrapInboundFrameBudgetedCodec(c transport.Codec) InboundFrameBudgetedCodec {
|
||||
switch v := c.(type) {
|
||||
case InboundFrameBudgetedCodec:
|
||||
return v
|
||||
case codec.NoHeader:
|
||||
return unwrapInboundFrameBudgetedCodec(v.Codec)
|
||||
case *codec.NoHeader:
|
||||
if v != nil {
|
||||
return unwrapInboundFrameBudgetedCodec(v.Codec)
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// inboundFramePreflightReader consumes only the framing length prefix, reserves the announced
|
||||
// wire+plaintext bytes, and only then exposes the final prefix bytes to the codec. Consequently a
|
||||
// budget error is observed by codec.Read before it can ResetN/Expand the payload buffer.
|
||||
type inboundFramePreflightReader struct {
|
||||
r io.Reader
|
||||
kind inboundFrameCodecKind
|
||||
reserve func(wireBytes, plaintextBytes int64) error
|
||||
|
||||
abridgedFirstDelivered bool
|
||||
done bool
|
||||
}
|
||||
|
||||
func (r *inboundFramePreflightReader) Read(p []byte) (int, error) {
|
||||
if len(p) == 0 {
|
||||
return 0, nil
|
||||
}
|
||||
if r.done {
|
||||
return r.r.Read(p)
|
||||
}
|
||||
|
||||
switch r.kind {
|
||||
case inboundFrameCodecQuickAckAbridged:
|
||||
return r.readAbridgedPrefix(p, true)
|
||||
case inboundFrameCodecAbridged:
|
||||
return r.readAbridgedPrefix(p, false)
|
||||
case inboundFrameCodecIntermediate, inboundFrameCodecPaddedIntermediate:
|
||||
return r.readWordPrefix(p, false)
|
||||
case inboundFrameCodecFull:
|
||||
return r.readWordPrefix(p, true)
|
||||
default:
|
||||
return 0, errInboundFrameCodecUnsupported
|
||||
}
|
||||
}
|
||||
|
||||
func (r *inboundFramePreflightReader) readAbridgedPrefix(p []byte, quickAck bool) (int, error) {
|
||||
if !r.abridgedFirstDelivered {
|
||||
var first [1]byte
|
||||
if _, err := io.ReadFull(r.r, first[:]); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
lengthByte := first[0]
|
||||
extended := lengthByte >= 0x7f
|
||||
if quickAck {
|
||||
lengthByte &= 0x7f
|
||||
extended = lengthByte == 0x7f
|
||||
}
|
||||
if !extended {
|
||||
n := int64(lengthByte) * bin.Word
|
||||
if err := reserveCompatFrame(r.reserve, n, n); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
r.done = true
|
||||
}
|
||||
r.abridgedFirstDelivered = true
|
||||
p[0] = first[0]
|
||||
return 1, nil
|
||||
}
|
||||
|
||||
var tail [3]byte
|
||||
if _, err := io.ReadFull(r.r, tail[:]); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
words := uint32(tail[0]) | uint32(tail[1])<<8 | uint32(tail[2])<<16
|
||||
n := int64(words) * bin.Word
|
||||
if err := reserveCompatFrame(r.reserve, n, n); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
r.done = true
|
||||
return copy(p, tail[:]), nil
|
||||
}
|
||||
|
||||
func (r *inboundFramePreflightReader) readWordPrefix(p []byte, full bool) (int, error) {
|
||||
var header [bin.Word]byte
|
||||
if _, err := io.ReadFull(r.r, header[:]); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
raw := int64(binary.LittleEndian.Uint32(header[:]))
|
||||
var wireBytes, plaintextBytes int64
|
||||
if full {
|
||||
// Full transport length includes length + sequence + payload + CRC.
|
||||
if raw < 3*bin.Word || raw > maxTransportMessageSize {
|
||||
return 0, fmt.Errorf("invalid full transport message length %d", raw)
|
||||
}
|
||||
wireBytes = raw
|
||||
plaintextBytes = raw - 3*bin.Word
|
||||
} else {
|
||||
wireBytes = raw &^ int64(quickAckResponseFlag)
|
||||
plaintextBytes = wireBytes
|
||||
}
|
||||
if err := reserveCompatFrame(r.reserve, wireBytes, plaintextBytes); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
r.done = true
|
||||
return copy(p, header[:]), nil
|
||||
}
|
||||
|
||||
func reserveCompatFrame(reserve func(wireBytes, plaintextBytes int64) error, wireBytes, plaintextBytes int64) error {
|
||||
if wireBytes <= 0 || wireBytes > maxTransportMessageSize {
|
||||
return fmt.Errorf("invalid transport message length %d", wireBytes)
|
||||
}
|
||||
return reserve(wireBytes, plaintextBytes)
|
||||
}
|
||||
337
internal/mtprotoedge/frame_budget_test.go
Normal file
337
internal/mtprotoedge/frame_budget_test.go
Normal file
|
|
@ -0,0 +1,337 @@
|
|||
package mtprotoedge
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/binary"
|
||||
"errors"
|
||||
"io"
|
||||
"net"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/gotd/td/bin"
|
||||
"github.com/gotd/td/proto/codec"
|
||||
"github.com/gotd/td/transport"
|
||||
)
|
||||
|
||||
type frameBudgetTestConn struct {
|
||||
reader bytes.Reader
|
||||
read int
|
||||
closed bool
|
||||
}
|
||||
|
||||
func newFrameBudgetTestConn(packet []byte) *frameBudgetTestConn {
|
||||
c := &frameBudgetTestConn{}
|
||||
c.reader.Reset(packet)
|
||||
return c
|
||||
}
|
||||
|
||||
func (c *frameBudgetTestConn) Read(p []byte) (int, error) {
|
||||
n, err := c.reader.Read(p)
|
||||
c.read += n
|
||||
return n, err
|
||||
}
|
||||
|
||||
func (*frameBudgetTestConn) Write(p []byte) (int, error) { return len(p), nil }
|
||||
func (c *frameBudgetTestConn) Close() error {
|
||||
c.closed = true
|
||||
return nil
|
||||
}
|
||||
func (*frameBudgetTestConn) LocalAddr() net.Addr { return frameBudgetTestAddr("local") }
|
||||
func (*frameBudgetTestConn) RemoteAddr() net.Addr { return frameBudgetTestAddr("remote") }
|
||||
func (*frameBudgetTestConn) SetDeadline(time.Time) error { return nil }
|
||||
func (*frameBudgetTestConn) SetReadDeadline(time.Time) error { return nil }
|
||||
func (*frameBudgetTestConn) SetWriteDeadline(time.Time) error { return nil }
|
||||
|
||||
type frameBudgetTestAddr string
|
||||
|
||||
func (a frameBudgetTestAddr) Network() string { return "frame-budget-test" }
|
||||
func (a frameBudgetTestAddr) String() string { return string(a) }
|
||||
|
||||
func newFrameBudgetTestTransport(packet []byte, c transport.Codec, budget *inboundFrameBudget) (*compatTransportConn, *frameBudgetTestConn) {
|
||||
raw := newFrameBudgetTestConn(packet)
|
||||
return &compatTransportConn{conn: raw, codec: c, budget: budget}, raw
|
||||
}
|
||||
|
||||
func TestInboundFrameBudgetSupportsBuiltInCodecs(t *testing.T) {
|
||||
payload := []byte{1, 2, 3, 4, 5, 6, 7, 8}
|
||||
|
||||
abridged := append([]byte{byte(len(payload) / bin.Word)}, payload...)
|
||||
intermediate := make([]byte, bin.Word+len(payload))
|
||||
binary.LittleEndian.PutUint32(intermediate, uint32(len(payload)))
|
||||
copy(intermediate[bin.Word:], payload)
|
||||
padded := make([]byte, bin.Word+len(payload)+1)
|
||||
binary.LittleEndian.PutUint32(padded, uint32(len(payload)+1))
|
||||
copy(padded[bin.Word:], payload)
|
||||
padded[len(padded)-1] = 0xa5
|
||||
|
||||
var full bytes.Buffer
|
||||
fullCodec := &codec.Full{}
|
||||
fullPayload := &bin.Buffer{Buf: append([]byte(nil), payload...)}
|
||||
if err := fullCodec.Write(&full, fullPayload); err != nil {
|
||||
t.Fatalf("encode full frame: %v", err)
|
||||
}
|
||||
|
||||
tests := []struct {
|
||||
name string
|
||||
packet []byte
|
||||
codec transport.Codec
|
||||
reservation int64
|
||||
}{
|
||||
{name: "abridged", packet: abridged, codec: &quickAckAbridgedCodec{}, reservation: 2 * int64(len(payload))},
|
||||
{name: "intermediate", packet: intermediate, codec: &quickAckIntermediateCodec{}, reservation: 2 * int64(len(payload))},
|
||||
{name: "padded_intermediate", packet: padded, codec: &quickAckPaddedIntermediateCodec{}, reservation: 2 * int64(len(payload)+1)},
|
||||
{name: "full", packet: full.Bytes(), codec: &codec.Full{}, reservation: int64(full.Len() + len(payload))},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
budget := newInboundFrameBudget(tt.reservation)
|
||||
conn, _ := newFrameBudgetTestTransport(tt.packet, tt.codec, budget)
|
||||
var got bin.Buffer
|
||||
if err := conn.Recv(context.Background(), &got); err != nil {
|
||||
t.Fatalf("Recv: %v", err)
|
||||
}
|
||||
if !bytes.Equal(got.Raw(), payload) {
|
||||
t.Fatalf("payload = %x, want %x", got.Raw(), payload)
|
||||
}
|
||||
if used := budget.usedBytes(); used != tt.reservation {
|
||||
t.Fatalf("held budget = %d, want %d", used, tt.reservation)
|
||||
}
|
||||
if err := conn.Close(); err != nil {
|
||||
t.Fatalf("Close: %v", err)
|
||||
}
|
||||
if used := budget.usedBytes(); used != tt.reservation {
|
||||
t.Fatalf("budget after concurrent Close = %d, want delivered ownership %d", used, tt.reservation)
|
||||
}
|
||||
conn.releaseInboundFrame()
|
||||
if used := budget.usedBytes(); used != 0 {
|
||||
t.Fatalf("budget after ownership release = %d, want 0", used)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestInboundFrameBudgetRejectsBeforePayloadAllocation(t *testing.T) {
|
||||
const payloadBytes = 1 << 20
|
||||
var header [bin.Word]byte
|
||||
binary.LittleEndian.PutUint32(header[:], payloadBytes)
|
||||
budget := newInboundFrameBudget(2*payloadBytes - 1)
|
||||
conn, raw := newFrameBudgetTestTransport(header[:], &quickAckIntermediateCodec{}, budget)
|
||||
var got bin.Buffer
|
||||
|
||||
err := conn.Recv(context.Background(), &got)
|
||||
if !errors.Is(err, ErrInboundFrameBudgetExceeded) {
|
||||
t.Fatalf("Recv error = %v, want ErrInboundFrameBudgetExceeded", err)
|
||||
}
|
||||
if raw.read != bin.Word {
|
||||
t.Fatalf("wire bytes read = %d, want only %d-byte length prefix", raw.read, bin.Word)
|
||||
}
|
||||
if cap(got.Buf) != 0 {
|
||||
t.Fatalf("payload buffer capacity = %d, want 0 before admission", cap(got.Buf))
|
||||
}
|
||||
if used := budget.usedBytes(); used != 0 {
|
||||
t.Fatalf("budget after rejected preflight = %d, want 0", used)
|
||||
}
|
||||
}
|
||||
|
||||
func TestInboundFrameBudgetAbridgedPreflightMatchesCodecSemantics(t *testing.T) {
|
||||
payload := []byte{1, 2, 3, 4, 5, 6, 7, 8}
|
||||
quickPacket := append([]byte{0x80 | byte(len(payload)/bin.Word)}, payload...)
|
||||
quickBudget := newInboundFrameBudget(int64(2 * len(payload)))
|
||||
quick, _ := newFrameBudgetTestTransport(quickPacket, &quickAckAbridgedCodec{}, quickBudget)
|
||||
var got bin.Buffer
|
||||
if err := quick.Recv(context.Background(), &got); err != nil {
|
||||
t.Fatalf("quick-ack abridged Recv: %v", err)
|
||||
}
|
||||
requested := quick.ConsumeQuickAckRequested()
|
||||
if !bytes.Equal(got.Raw(), payload) || !requested {
|
||||
t.Fatalf("quick-ack frame = %x requested=%v", got.Raw(), requested)
|
||||
}
|
||||
quick.releaseInboundFrame()
|
||||
_ = quick.Close()
|
||||
|
||||
// gotd's plain codec treats every first byte >= 0x7f as the extended form (it does not
|
||||
// implement the quick-ack high bit). The preflight parser must mirror that behavior; treating
|
||||
// 0x82 as a short two-word frame would let the codec allocate from the following three bytes.
|
||||
malicious := []byte{0x82, 0xff, 0xff, 0xff}
|
||||
plainBudget := newInboundFrameBudget(defaultInboundFrameGlobalMaxBytes)
|
||||
plain, raw := newFrameBudgetTestTransport(malicious, codec.Abridged{}, plainBudget)
|
||||
got.Reset()
|
||||
err := plain.Recv(context.Background(), &got)
|
||||
if err == nil {
|
||||
t.Fatal("plain abridged accepted oversized extended length")
|
||||
}
|
||||
if raw.read != 4 || cap(got.Buf) > 2*bin.Word {
|
||||
t.Fatalf("plain abridged read=%d buffer_cap=%d, want prefix-only allocation", raw.read, cap(got.Buf))
|
||||
}
|
||||
_ = plain.Close()
|
||||
}
|
||||
|
||||
func TestInboundFrameBudgetReleasedAtNextRecvAndReusable(t *testing.T) {
|
||||
payload := []byte{1, 2, 3, 4, 5, 6, 7, 8}
|
||||
frame := make([]byte, bin.Word+len(payload))
|
||||
binary.LittleEndian.PutUint32(frame, uint32(len(payload)))
|
||||
copy(frame[bin.Word:], payload)
|
||||
packet := append(append([]byte(nil), frame...), frame...)
|
||||
reservation := int64(2 * len(payload))
|
||||
budget := newInboundFrameBudget(reservation)
|
||||
conn, _ := newFrameBudgetTestTransport(packet, &quickAckIntermediateCodec{}, budget)
|
||||
|
||||
for i := 0; i < 2; i++ {
|
||||
var got bin.Buffer
|
||||
if err := conn.Recv(context.Background(), &got); err != nil {
|
||||
t.Fatalf("Recv %d: %v", i+1, err)
|
||||
}
|
||||
if used := budget.usedBytes(); used != reservation {
|
||||
t.Fatalf("held budget after frame %d = %d, want %d", i+1, used, reservation)
|
||||
}
|
||||
}
|
||||
conn.releaseInboundFrame()
|
||||
_ = conn.Close()
|
||||
}
|
||||
|
||||
func TestInboundFrameRetainedBackingStaysChargedAcrossSmallFrame(t *testing.T) {
|
||||
const largeBytes = 1 << 20
|
||||
large := make([]byte, bin.Word+largeBytes)
|
||||
binary.LittleEndian.PutUint32(large, largeBytes)
|
||||
smallPayload := []byte{1, 2, 3, 4, 5, 6, 7, 8}
|
||||
small := make([]byte, bin.Word+len(smallPayload))
|
||||
binary.LittleEndian.PutUint32(small, uint32(len(smallPayload)))
|
||||
copy(small[bin.Word:], smallPayload)
|
||||
packet := append(large, small...)
|
||||
budget := newInboundFrameBudget(2 * largeBytes)
|
||||
conn, _ := newFrameBudgetTestTransport(packet, &quickAckIntermediateCodec{}, budget)
|
||||
var wire bin.Buffer
|
||||
if err := conn.Recv(context.Background(), &wire); err != nil {
|
||||
t.Fatalf("large Recv: %v", err)
|
||||
}
|
||||
// Model decryptClientFrame's exact-size plaintext reuse buffer.
|
||||
plain := bin.Buffer{Buf: make([]byte, largeBytes)}
|
||||
retainInboundFrameBackings(conn, &wire, &plain)
|
||||
retained := int64(cap(wire.Buf) + cap(plain.Buf))
|
||||
if got := budget.usedBytes(); got != retained {
|
||||
t.Fatalf("retained budget after large frame = %d, want capacities %d", got, retained)
|
||||
}
|
||||
|
||||
wire.Reset()
|
||||
if err := conn.Recv(context.Background(), &wire); err != nil {
|
||||
t.Fatalf("small Recv: %v", err)
|
||||
}
|
||||
// The small announcement must not release the large backing's charge. This was the
|
||||
// warm-many-connections bypass: each socket retained MiBs while the global budget saw bytes.
|
||||
if got := budget.usedBytes(); got != retained {
|
||||
t.Fatalf("budget after small frame = %d, want retained high-water %d", got, retained)
|
||||
}
|
||||
|
||||
wire.Buf = nil
|
||||
plain.Buf = nil
|
||||
retainInboundFrameBackings(conn, &wire, &plain)
|
||||
if got := budget.usedBytes(); got != 0 {
|
||||
t.Fatalf("budget after dropping reusable backings = %d, want 0", got)
|
||||
}
|
||||
_ = conn.Close()
|
||||
}
|
||||
|
||||
func TestInboundFrameBudgetClosePreservesDeliveredOwnershipUntilRelease(t *testing.T) {
|
||||
payload := []byte{1, 2, 3, 4, 5, 6, 7, 8}
|
||||
frame := make([]byte, bin.Word+len(payload))
|
||||
binary.LittleEndian.PutUint32(frame, uint32(len(payload)))
|
||||
copy(frame[bin.Word:], payload)
|
||||
budget := newInboundFrameBudget(int64(2 * len(payload)))
|
||||
|
||||
first, _ := newFrameBudgetTestTransport(frame, &quickAckIntermediateCodec{}, budget)
|
||||
var got bin.Buffer
|
||||
if err := first.Recv(context.Background(), &got); err != nil {
|
||||
t.Fatalf("first Recv: %v", err)
|
||||
}
|
||||
blocked, _ := newFrameBudgetTestTransport(frame, &quickAckIntermediateCodec{}, budget)
|
||||
var blockedPayload bin.Buffer
|
||||
if err := blocked.Recv(context.Background(), &blockedPayload); !errors.Is(err, ErrInboundFrameBudgetExceeded) {
|
||||
t.Fatalf("concurrent Recv error = %v, want global budget rejection", err)
|
||||
}
|
||||
if cap(blockedPayload.Buf) != 0 {
|
||||
t.Fatalf("blocked connection allocated payload capacity %d", cap(blockedPayload.Buf))
|
||||
}
|
||||
_ = blocked.Close()
|
||||
if err := first.Close(); err != nil {
|
||||
t.Fatalf("first Close: %v", err)
|
||||
}
|
||||
if used := budget.usedBytes(); used != int64(2*len(payload)) {
|
||||
t.Fatalf("budget after concurrent Close = %d, want delivered frame still charged", used)
|
||||
}
|
||||
|
||||
second, _ := newFrameBudgetTestTransport(frame, &quickAckIntermediateCodec{}, budget)
|
||||
got.Reset()
|
||||
if err := second.Recv(context.Background(), &got); !errors.Is(err, ErrInboundFrameBudgetExceeded) {
|
||||
t.Fatalf("second Recv before ownership release = %v, want budget rejection", err)
|
||||
}
|
||||
_ = second.Close()
|
||||
|
||||
first.releaseInboundFrame()
|
||||
third, _ := newFrameBudgetTestTransport(frame, &quickAckIntermediateCodec{}, budget)
|
||||
got.Reset()
|
||||
if err := third.Recv(context.Background(), &got); err != nil {
|
||||
t.Fatalf("third Recv after ownership release: %v", err)
|
||||
}
|
||||
third.releaseInboundFrame()
|
||||
_ = third.Close()
|
||||
}
|
||||
|
||||
type unsafeFrameBudgetCodec struct {
|
||||
readCalled bool
|
||||
}
|
||||
|
||||
func (*unsafeFrameBudgetCodec) WriteHeader(io.Writer) error { return nil }
|
||||
func (*unsafeFrameBudgetCodec) ReadHeader(io.Reader) error { return nil }
|
||||
func (*unsafeFrameBudgetCodec) Write(io.Writer, *bin.Buffer) error { return nil }
|
||||
func (c *unsafeFrameBudgetCodec) Read(io.Reader, *bin.Buffer) error { c.readCalled = true; return nil }
|
||||
|
||||
func TestCustomCodecWithoutPreflightFailsClosed(t *testing.T) {
|
||||
raw := newFrameBudgetTestConn([]byte{1, 2, 3, 4})
|
||||
listener := newSingleConnListener(raw)
|
||||
custom := &unsafeFrameBudgetCodec{}
|
||||
budgeted := newCompatTransportListener(func() transport.Codec { return custom }, listener, newInboundFrameBudget(1024))
|
||||
|
||||
conn, err := budgeted.Accept()
|
||||
if !errors.Is(err, errInboundFrameCodecUnsupported) {
|
||||
t.Fatalf("Accept error = %v, want unsupported preflight codec", err)
|
||||
}
|
||||
if conn != nil {
|
||||
t.Fatal("unsupported custom codec unexpectedly accepted")
|
||||
}
|
||||
if custom.readCalled || raw.read != 0 {
|
||||
t.Fatalf("custom codec touched frame before rejection: read_called=%v wire_read=%d", custom.readCalled, raw.read)
|
||||
}
|
||||
}
|
||||
|
||||
func TestExplicitBuiltInCodecUsesBudgetedTransport(t *testing.T) {
|
||||
payload := []byte{1, 2, 3, 4, 5, 6, 7, 8}
|
||||
packet := append([]byte(nil), codec.IntermediateClientStart[:]...)
|
||||
var header [bin.Word]byte
|
||||
binary.LittleEndian.PutUint32(header[:], uint32(len(payload)))
|
||||
packet = append(packet, header[:]...)
|
||||
packet = append(packet, payload...)
|
||||
|
||||
raw := newFrameBudgetTestConn(packet)
|
||||
budget := newInboundFrameBudget(int64(2 * len(payload)))
|
||||
listener := newCompatTransportListener(
|
||||
func() transport.Codec { return codec.Intermediate{} },
|
||||
newSingleConnListener(raw),
|
||||
budget,
|
||||
)
|
||||
conn, err := listener.Accept()
|
||||
if err != nil {
|
||||
t.Fatalf("Accept: %v", err)
|
||||
}
|
||||
var got bin.Buffer
|
||||
if err := conn.Recv(context.Background(), &got); err != nil {
|
||||
t.Fatalf("Recv: %v", err)
|
||||
}
|
||||
if !bytes.Equal(got.Raw(), payload) || budget.usedBytes() != int64(2*len(payload)) {
|
||||
t.Fatalf("payload=%x budget=%d", got.Raw(), budget.usedBytes())
|
||||
}
|
||||
conn.(*compatTransportConn).releaseInboundFrame()
|
||||
_ = conn.Close()
|
||||
}
|
||||
|
|
@ -1,30 +1,337 @@
|
|||
package mtprotoedge
|
||||
|
||||
import (
|
||||
"container/list"
|
||||
"context"
|
||||
"errors"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
)
|
||||
|
||||
// ErrInboundRPCQueueFull 表示单连接 RPC 队列已满。
|
||||
// ErrInboundRPCQueueFull 表示 inbound RPC 已触达单连接或进程级预算。
|
||||
var ErrInboundRPCQueueFull = errors.New("inbound rpc queue full")
|
||||
|
||||
// maxInflightRPCBytes 是单连接已入队未完成 inbound RPC body 的总字节上限。
|
||||
// 队列除按条数(queueSize)限制外,再按字节预算兜底:对抗客户端发满大请求时按字节先拒绝。
|
||||
// maxInflightRPCBytes 是单连接所有已预留、排队和执行中 RPC body 的总字节上限。
|
||||
// 进程级预算在 Copy 前先兜底;这里再隔离单个连接,避免一个客户端独占全局内存。
|
||||
const maxInflightRPCBytes = 32 << 20 // 32 MiB
|
||||
|
||||
// rpcCloseWaitTimeout 是连接关闭时等待 inbound RPC worker 退出的上限。
|
||||
// rpcCloseWaitTimeout 是连接/Server 关闭时等待在途 RPC 或共享 worker 退出的上限。
|
||||
const rpcCloseWaitTimeout = 5 * time.Second
|
||||
|
||||
type inboundRPC struct {
|
||||
ctx context.Context
|
||||
method string
|
||||
enqueuedAt time.Time
|
||||
size int
|
||||
run func(context.Context) error
|
||||
ctx context.Context
|
||||
cancel context.CancelFunc
|
||||
stopRoot func() bool
|
||||
stopTimeout func() bool
|
||||
method string
|
||||
enqueuedAt time.Time
|
||||
deadline time.Time
|
||||
size int
|
||||
run func(context.Context) error
|
||||
onTimeout func()
|
||||
budget *inboundRPCGlobalReservation
|
||||
ticket *inboundRPCTicket
|
||||
}
|
||||
|
||||
func (c *Conn) startInboundRPCScheduler(maxInflight, queueSize int, timeout time.Duration) {
|
||||
const (
|
||||
inboundRPCTicketQueued int32 = iota
|
||||
inboundRPCTicketRunning
|
||||
inboundRPCTicketDone
|
||||
)
|
||||
|
||||
type inboundRPCTicket struct {
|
||||
state atomic.Int32
|
||||
onTimeout func()
|
||||
}
|
||||
|
||||
// inboundRPCScheduler 是 Server 级共享调度器。ready 中每个 Conn 最多只有一个有效令牌;
|
||||
// worker 每次只从该连接取一条,再把仍可运行的连接放回队尾,因此单个热点连接不能长期
|
||||
// 占住共享池。worker 在首条任务到达后才创建,空闲 Server 不预起 256 个 goroutine。
|
||||
type inboundRPCScheduler struct {
|
||||
workers int
|
||||
maxTasks int
|
||||
maxBytes int64
|
||||
|
||||
// ready is an intrusive scheduler-owned queue rather than a bounded channel. A connection
|
||||
// has at most one element, and close removes that element in O(1). This prevents closed-Conn
|
||||
// stale tokens from filling a channel and making every worker block while trying to reschedule.
|
||||
readyMu sync.Mutex
|
||||
ready *list.List
|
||||
readyIndex map[*Conn]*list.Element
|
||||
readyWake chan struct{}
|
||||
stopCh chan struct{}
|
||||
|
||||
lifecycleMu sync.Mutex
|
||||
started bool
|
||||
stopped bool
|
||||
workersStarted bool
|
||||
workerWG sync.WaitGroup
|
||||
|
||||
budgetMu sync.Mutex
|
||||
tasks int
|
||||
bytes int64
|
||||
}
|
||||
|
||||
type inboundRPCGlobalReservation struct {
|
||||
scheduler *inboundRPCScheduler
|
||||
size int64
|
||||
once sync.Once
|
||||
}
|
||||
|
||||
// inboundRPCReservation 同时持有全局和单连接的“Copy 前”预算。commit/abort 只能成功一次;
|
||||
// 无论 Copy 后连接关闭、入队成功还是调用方提前返回,预算都有唯一归还路径。
|
||||
type inboundRPCReservation struct {
|
||||
conn *Conn
|
||||
global *inboundRPCGlobalReservation
|
||||
ctx context.Context
|
||||
method string
|
||||
size int
|
||||
enqueuedAt time.Time
|
||||
deadline time.Time
|
||||
once sync.Once
|
||||
}
|
||||
|
||||
func newInboundRPCScheduler(workers, maxTasks int, maxBytes int64) *inboundRPCScheduler {
|
||||
if workers <= 0 {
|
||||
workers = 1
|
||||
}
|
||||
if maxTasks <= 0 {
|
||||
maxTasks = 1
|
||||
}
|
||||
if maxBytes <= 0 {
|
||||
maxBytes = 1
|
||||
}
|
||||
return &inboundRPCScheduler{
|
||||
workers: workers,
|
||||
maxTasks: maxTasks,
|
||||
maxBytes: maxBytes,
|
||||
ready: list.New(),
|
||||
readyIndex: make(map[*Conn]*list.Element),
|
||||
readyWake: make(chan struct{}, 1),
|
||||
stopCh: make(chan struct{}),
|
||||
}
|
||||
}
|
||||
|
||||
// start 允许共享池开始消费。已在 start 前进入 ready 的任务会保留顺序,便于启动突发,
|
||||
// 也使测试能够确定性验证轮转公平性。
|
||||
func (s *inboundRPCScheduler) start() {
|
||||
s.lifecycleMu.Lock()
|
||||
if s.stopped {
|
||||
s.lifecycleMu.Unlock()
|
||||
return
|
||||
}
|
||||
s.started = true
|
||||
shouldStart := s.readyLen() > 0
|
||||
s.lifecycleMu.Unlock()
|
||||
if shouldStart {
|
||||
s.ensureWorkers()
|
||||
}
|
||||
}
|
||||
|
||||
func (s *inboundRPCScheduler) ensureWorkers() {
|
||||
s.lifecycleMu.Lock()
|
||||
defer s.lifecycleMu.Unlock()
|
||||
if !s.started || s.stopped || s.workersStarted {
|
||||
return
|
||||
}
|
||||
s.workersStarted = true
|
||||
s.workerWG.Add(s.workers)
|
||||
for i := 0; i < s.workers; i++ {
|
||||
go s.worker()
|
||||
}
|
||||
}
|
||||
|
||||
func (s *inboundRPCScheduler) stop(timeout time.Duration) {
|
||||
s.lifecycleMu.Lock()
|
||||
if !s.stopped {
|
||||
s.stopped = true
|
||||
s.budgetMu.Lock()
|
||||
// 与 reserveGlobal 在同一把锁下切断新任务;已持有 reservation 的任务仍由
|
||||
// 对应 Conn 的 commit/abort/close 路径精确归还。
|
||||
close(s.stopCh)
|
||||
s.budgetMu.Unlock()
|
||||
}
|
||||
s.lifecycleMu.Unlock()
|
||||
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
s.workerWG.Wait()
|
||||
close(done)
|
||||
}()
|
||||
if timeout <= 0 {
|
||||
<-done
|
||||
return
|
||||
}
|
||||
timer := time.NewTimer(timeout)
|
||||
defer timer.Stop()
|
||||
select {
|
||||
case <-done:
|
||||
case <-timer.C:
|
||||
}
|
||||
}
|
||||
|
||||
func (s *inboundRPCScheduler) reserveGlobal(size int) (*inboundRPCGlobalReservation, string, error) {
|
||||
if size < 0 {
|
||||
size = 0
|
||||
}
|
||||
size64 := int64(size)
|
||||
s.budgetMu.Lock()
|
||||
defer s.budgetMu.Unlock()
|
||||
|
||||
select {
|
||||
case <-s.stopCh:
|
||||
return nil, "scheduler_closed", ErrConnClosed
|
||||
default:
|
||||
}
|
||||
if s.tasks >= s.maxTasks {
|
||||
return nil, "global_task_budget", ErrInboundRPCQueueFull
|
||||
}
|
||||
// 用减法比较避免 s.bytes+size64 溢出。
|
||||
if size64 > s.maxBytes-s.bytes {
|
||||
return nil, "global_byte_budget", ErrInboundRPCQueueFull
|
||||
}
|
||||
s.tasks++
|
||||
s.bytes += size64
|
||||
return &inboundRPCGlobalReservation{scheduler: s, size: size64}, "", nil
|
||||
}
|
||||
|
||||
func (r *inboundRPCGlobalReservation) release() {
|
||||
if r == nil || r.scheduler == nil {
|
||||
return
|
||||
}
|
||||
r.once.Do(func() {
|
||||
s := r.scheduler
|
||||
s.budgetMu.Lock()
|
||||
s.tasks--
|
||||
s.bytes -= r.size
|
||||
s.budgetMu.Unlock()
|
||||
})
|
||||
}
|
||||
|
||||
func (s *inboundRPCScheduler) budgetSnapshot() (tasks int, bytes int64) {
|
||||
s.budgetMu.Lock()
|
||||
defer s.budgetMu.Unlock()
|
||||
return s.tasks, s.bytes
|
||||
}
|
||||
|
||||
func (s *inboundRPCScheduler) schedule(c *Conn) {
|
||||
if s == nil || c == nil {
|
||||
return
|
||||
}
|
||||
// rpcReady/rpcClosed and queue membership must be tested/installed while holding rpcMu.
|
||||
// Otherwise close can remove the old token between the test and enqueue, leaving a new stale
|
||||
// token behind after the connection is already terminal.
|
||||
c.rpcMu.Lock()
|
||||
eligible := c.rpcReady && !c.rpcClosed
|
||||
added := false
|
||||
if eligible {
|
||||
added = s.enqueueReady(c)
|
||||
}
|
||||
c.rpcMu.Unlock()
|
||||
if !added {
|
||||
return
|
||||
}
|
||||
s.signalReady()
|
||||
s.ensureWorkers()
|
||||
}
|
||||
|
||||
func (s *inboundRPCScheduler) worker() {
|
||||
defer s.workerWG.Done()
|
||||
for {
|
||||
select {
|
||||
case <-s.stopCh:
|
||||
return
|
||||
default:
|
||||
}
|
||||
if c := s.popReady(); c != nil {
|
||||
task, ok, reschedule := c.takeInboundRPC()
|
||||
if reschedule {
|
||||
s.schedule(c)
|
||||
}
|
||||
if ok {
|
||||
c.runInboundRPC(task)
|
||||
}
|
||||
continue
|
||||
}
|
||||
select {
|
||||
case <-s.readyWake:
|
||||
case <-s.stopCh:
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (s *inboundRPCScheduler) enqueueReady(c *Conn) bool {
|
||||
select {
|
||||
case <-s.stopCh:
|
||||
return false
|
||||
default:
|
||||
}
|
||||
s.readyMu.Lock()
|
||||
defer s.readyMu.Unlock()
|
||||
select {
|
||||
case <-s.stopCh:
|
||||
return false
|
||||
default:
|
||||
}
|
||||
if _, exists := s.readyIndex[c]; exists {
|
||||
return false
|
||||
}
|
||||
s.readyIndex[c] = s.ready.PushBack(c)
|
||||
return true
|
||||
}
|
||||
|
||||
func (s *inboundRPCScheduler) popReady() *Conn {
|
||||
s.readyMu.Lock()
|
||||
front := s.ready.Front()
|
||||
if front == nil {
|
||||
s.readyMu.Unlock()
|
||||
return nil
|
||||
}
|
||||
c, _ := front.Value.(*Conn)
|
||||
s.ready.Remove(front)
|
||||
delete(s.readyIndex, c)
|
||||
hasMore := s.ready.Len() > 0
|
||||
s.readyMu.Unlock()
|
||||
if hasMore {
|
||||
// Wake another worker while this worker begins the task. A capacity-one wake channel is
|
||||
// sufficient: every pop cascades another wake until the queue is drained.
|
||||
s.signalReady()
|
||||
}
|
||||
return c
|
||||
}
|
||||
|
||||
func (s *inboundRPCScheduler) unschedule(c *Conn) {
|
||||
if s == nil || c == nil {
|
||||
return
|
||||
}
|
||||
s.readyMu.Lock()
|
||||
if el := s.readyIndex[c]; el != nil {
|
||||
s.ready.Remove(el)
|
||||
delete(s.readyIndex, c)
|
||||
}
|
||||
hasMore := s.ready.Len() > 0
|
||||
s.readyMu.Unlock()
|
||||
if hasMore {
|
||||
s.signalReady()
|
||||
}
|
||||
}
|
||||
|
||||
func (s *inboundRPCScheduler) readyLen() int {
|
||||
s.readyMu.Lock()
|
||||
defer s.readyMu.Unlock()
|
||||
return s.ready.Len()
|
||||
}
|
||||
|
||||
func (s *inboundRPCScheduler) signalReady() {
|
||||
select {
|
||||
case s.readyWake <- struct{}{}:
|
||||
default:
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Conn) startInboundRPCScheduler(scheduler *inboundRPCScheduler, maxInflight, queueSize int, timeout time.Duration) {
|
||||
if c.metrics == nil {
|
||||
c.metrics = NopMetrics{}
|
||||
}
|
||||
|
|
@ -35,163 +342,413 @@ func (c *Conn) startInboundRPCScheduler(maxInflight, queueSize int, timeout time
|
|||
queueSize = 1
|
||||
}
|
||||
rootCtx, cancel := context.WithCancel(context.Background())
|
||||
c.rpcQueue = make(chan inboundRPC, queueSize)
|
||||
c.rpcStop = make(chan struct{})
|
||||
c.rpcScheduler = scheduler
|
||||
c.rpcCancel = cancel
|
||||
c.rpcTimeout = timeout
|
||||
c.rpcRootCtx = rootCtx
|
||||
c.rpcMaxInflight = maxInflight
|
||||
// worker 懒启动:不在此处起 worker;首个 RPC 入队时由 ensureInboundRPCWorkers 起,
|
||||
// 避免握手后静默 / 纯推送目标连接白白钉住 maxInflight 个 goroutine。
|
||||
c.rpcQueueSize = queueSize
|
||||
// rpcQueue 保持 nil;首个成功 commit 才由 append 分配,静默连接零队列内存。
|
||||
}
|
||||
|
||||
// ensureInboundRPCWorkers 懒启动 maxInflight 个 RPC worker(仅一次),在 enqueueInboundRPC
|
||||
// 入队成功后调用。从不发 RPC 的连接(半开 / 纯推送)由此完全不起 worker。
|
||||
func (c *Conn) ensureInboundRPCWorkers() {
|
||||
c.rpcWorkersOnce.Do(func() {
|
||||
c.rpcWG.Add(c.rpcMaxInflight)
|
||||
for i := 0; i < c.rpcMaxInflight; i++ {
|
||||
go c.inboundRPCWorker(c.rpcRootCtx)
|
||||
// reserveInboundRPC 必须在 request body Copy 前调用。它先拿进程级条数/字节预算,
|
||||
// 再预占单连接队列槽和字节预算;commit 或 abort 负责唯一释放。
|
||||
func (c *Conn) reserveInboundRPC(ctx context.Context, method string, size int) (*inboundRPCReservation, error) {
|
||||
if ctx == nil {
|
||||
ctx = context.Background()
|
||||
}
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
c.metrics.InboundRPCDropped(method, "context_done")
|
||||
return nil, ctx.Err()
|
||||
default:
|
||||
}
|
||||
if c.rpcScheduler == nil {
|
||||
c.metrics.InboundRPCDropped(method, "scheduler_closed")
|
||||
return nil, ErrConnClosed
|
||||
}
|
||||
global, reason, err := c.rpcScheduler.reserveGlobal(size)
|
||||
if err != nil {
|
||||
c.metrics.InboundRPCDropped(method, reason)
|
||||
return nil, err
|
||||
}
|
||||
|
||||
now := time.Now()
|
||||
deadline := time.Time{}
|
||||
if c.rpcTimeout > 0 {
|
||||
deadline = now.Add(c.rpcTimeout)
|
||||
}
|
||||
if ctxDeadline, ok := ctx.Deadline(); ok && (deadline.IsZero() || ctxDeadline.Before(deadline)) {
|
||||
deadline = ctxDeadline
|
||||
}
|
||||
if size < 0 {
|
||||
size = 0
|
||||
}
|
||||
|
||||
c.rpcMu.Lock()
|
||||
if err := ctx.Err(); err != nil {
|
||||
c.rpcMu.Unlock()
|
||||
global.release()
|
||||
c.metrics.InboundRPCDropped(method, "context_done")
|
||||
return nil, err
|
||||
}
|
||||
if c.rpcClosed {
|
||||
c.rpcMu.Unlock()
|
||||
global.release()
|
||||
c.metrics.InboundRPCDropped(method, "scheduler_closed")
|
||||
return nil, ErrConnClosed
|
||||
}
|
||||
if c.rpcReserved+len(c.rpcQueue) >= c.rpcQueueSize {
|
||||
c.rpcMu.Unlock()
|
||||
global.release()
|
||||
c.metrics.InboundRPCDropped(method, "queue_full")
|
||||
return nil, ErrInboundRPCQueueFull
|
||||
}
|
||||
if int64(size) > maxInflightRPCBytes-c.inflightRPCBytes.Load() {
|
||||
c.rpcMu.Unlock()
|
||||
global.release()
|
||||
c.metrics.InboundRPCDropped(method, "byte_budget")
|
||||
return nil, ErrInboundRPCQueueFull
|
||||
}
|
||||
c.rpcReserved++
|
||||
c.inflightRPCBytes.Add(int64(size))
|
||||
// Add 与 close 的 Wait 由 rpcMu 排序:close 置 rpcClosed 后不会再发生 Add。
|
||||
c.rpcReservationWG.Add(1)
|
||||
c.rpcMu.Unlock()
|
||||
|
||||
return &inboundRPCReservation{
|
||||
conn: c,
|
||||
global: global,
|
||||
ctx: ctx,
|
||||
method: method,
|
||||
size: size,
|
||||
enqueuedAt: now,
|
||||
deadline: deadline,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// enqueueInboundRPC 是测试和已持有独立 body 的便捷入口。生产收包路径使用
|
||||
// reserveInboundRPC -> Copy -> commit,保证真正的 Copy 前预算。
|
||||
func (c *Conn) enqueueInboundRPC(ctx context.Context, task inboundRPC) error {
|
||||
reservation, err := c.reserveInboundRPC(ctx, task.method, task.size)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer reservation.abort()
|
||||
return reservation.commit(task)
|
||||
}
|
||||
|
||||
func (r *inboundRPCReservation) commit(task inboundRPC) error {
|
||||
result := ErrConnClosed
|
||||
var (
|
||||
committed bool
|
||||
reschedule bool
|
||||
queueLen int
|
||||
queueCap int
|
||||
)
|
||||
r.once.Do(func() {
|
||||
c := r.conn
|
||||
c.rpcMu.Lock()
|
||||
c.rpcReserved--
|
||||
if c.rpcClosed {
|
||||
c.inflightRPCBytes.Add(-int64(r.size))
|
||||
} else {
|
||||
// The request deadline starts when admission succeeds, not when a worker
|
||||
// eventually dequeues the request. This bounds total queue + execution
|
||||
// latency and lets a queued request emit its explicit timeout on time.
|
||||
if r.deadline.IsZero() {
|
||||
task.ctx, task.cancel = context.WithCancel(r.ctx)
|
||||
} else {
|
||||
task.ctx, task.cancel = context.WithDeadline(r.ctx, r.deadline)
|
||||
}
|
||||
task.stopRoot = context.AfterFunc(c.rpcRootCtx, task.cancel)
|
||||
task.method = r.method
|
||||
task.enqueuedAt = r.enqueuedAt
|
||||
task.deadline = r.deadline
|
||||
task.size = r.size
|
||||
task.budget = r.global
|
||||
ticket := &inboundRPCTicket{}
|
||||
if task.onTimeout != nil {
|
||||
onTimeout := task.onTimeout
|
||||
var timeoutOnce sync.Once
|
||||
ticket.onTimeout = func() {
|
||||
timeoutOnce.Do(onTimeout)
|
||||
}
|
||||
task.onTimeout = ticket.onTimeout
|
||||
}
|
||||
task.ticket = ticket
|
||||
if task.onTimeout != nil && !task.deadline.IsZero() {
|
||||
taskCtx := task.ctx
|
||||
task.stopTimeout = context.AfterFunc(taskCtx, func() {
|
||||
if errors.Is(taskCtx.Err(), context.DeadlineExceeded) {
|
||||
c.expireInboundRPCTicket(ticket)
|
||||
}
|
||||
})
|
||||
}
|
||||
c.rpcQueue = append(c.rpcQueue, task)
|
||||
queueLen = len(c.rpcQueue)
|
||||
queueCap = c.rpcQueueSize
|
||||
if c.rpcRunning < c.rpcMaxInflight && !c.rpcReady {
|
||||
c.rpcReady = true
|
||||
reschedule = true
|
||||
}
|
||||
committed = true
|
||||
result = nil
|
||||
}
|
||||
c.rpcMu.Unlock()
|
||||
c.rpcReservationWG.Done()
|
||||
if !committed {
|
||||
r.global.release()
|
||||
}
|
||||
})
|
||||
if committed {
|
||||
r.conn.metrics.InboundRPCQueued(r.method, queueLen, queueCap)
|
||||
if reschedule {
|
||||
r.conn.rpcScheduler.schedule(r.conn)
|
||||
}
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
func (r *inboundRPCReservation) abort() {
|
||||
if r == nil {
|
||||
return
|
||||
}
|
||||
r.once.Do(func() {
|
||||
c := r.conn
|
||||
c.rpcMu.Lock()
|
||||
c.rpcReserved--
|
||||
c.inflightRPCBytes.Add(-int64(r.size))
|
||||
c.rpcMu.Unlock()
|
||||
c.rpcReservationWG.Done()
|
||||
r.global.release()
|
||||
})
|
||||
}
|
||||
|
||||
func (c *Conn) enqueueInboundRPC(ctx context.Context, task inboundRPC) error {
|
||||
if ctx == nil {
|
||||
ctx = context.Background()
|
||||
func (c *Conn) takeInboundRPC() (task inboundRPC, ok, reschedule bool) {
|
||||
c.rpcMu.Lock()
|
||||
defer c.rpcMu.Unlock()
|
||||
// ready token 是可替代的:收到一个 token 就消费当前“已调度”状态。关闭后或
|
||||
// 已被另一 token 抢先处理时,这只是一个无害 stale token。
|
||||
if !c.rpcReady {
|
||||
return inboundRPC{}, false, false
|
||||
}
|
||||
if c.rpcQueue == nil || c.rpcStop == nil {
|
||||
c.metrics.InboundRPCDropped(task.method, "scheduler_closed")
|
||||
return ErrConnClosed
|
||||
c.rpcReady = false
|
||||
if c.rpcClosed || len(c.rpcQueue) == 0 || c.rpcRunning >= c.rpcMaxInflight {
|
||||
return inboundRPC{}, false, false
|
||||
}
|
||||
task.ctx = ctx
|
||||
task.enqueuedAt = time.Now()
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
task = c.rpcQueue[0]
|
||||
c.rpcQueue[0] = inboundRPC{}
|
||||
c.rpcQueue = c.rpcQueue[1:]
|
||||
if len(c.rpcQueue) == 0 {
|
||||
c.rpcQueue = nil
|
||||
}
|
||||
c.rpcRunning++
|
||||
if task.ticket != nil {
|
||||
task.ticket.state.Store(inboundRPCTicketRunning)
|
||||
}
|
||||
c.rpcWG.Add(1)
|
||||
if len(c.rpcQueue) > 0 && c.rpcRunning < c.rpcMaxInflight {
|
||||
c.rpcReady = true
|
||||
reschedule = true
|
||||
}
|
||||
return task, true, reschedule
|
||||
}
|
||||
|
||||
func (c *Conn) runInboundRPC(task inboundRPC) {
|
||||
defer c.finishInboundRPC(task)
|
||||
|
||||
now := time.Now()
|
||||
ctxErr := task.ctx.Err()
|
||||
if (!task.deadline.IsZero() && !now.Before(task.deadline)) || errors.Is(ctxErr, context.DeadlineExceeded) {
|
||||
c.metrics.InboundRPCDropped(task.method, "queue_timeout")
|
||||
if task.onTimeout != nil {
|
||||
task.onTimeout()
|
||||
}
|
||||
return
|
||||
}
|
||||
if ctxErr != nil {
|
||||
c.metrics.InboundRPCDropped(task.method, "context_done")
|
||||
return ctx.Err()
|
||||
case <-c.rpcStop:
|
||||
c.metrics.InboundRPCDropped(task.method, "scheduler_closed")
|
||||
return ErrConnClosed
|
||||
default:
|
||||
return
|
||||
}
|
||||
// 字节预算:先预扣 size,超 maxInflightRPCBytes 则回滚并拒绝(与条数上限并列的第二道闸)。
|
||||
if task.size > 0 {
|
||||
if c.inflightRPCBytes.Add(int64(task.size)) > maxInflightRPCBytes {
|
||||
c.inflightRPCBytes.Add(-int64(task.size))
|
||||
c.metrics.InboundRPCDropped(task.method, "byte_budget")
|
||||
return ErrInboundRPCQueueFull
|
||||
}
|
||||
}
|
||||
select {
|
||||
case c.rpcQueue <- task:
|
||||
c.ensureInboundRPCWorkers()
|
||||
c.metrics.InboundRPCQueued(task.method, len(c.rpcQueue), cap(c.rpcQueue))
|
||||
return nil
|
||||
case <-ctx.Done():
|
||||
c.releaseInflightRPCBytes(task.size)
|
||||
c.metrics.InboundRPCDropped(task.method, "context_done")
|
||||
return ctx.Err()
|
||||
case <-c.rpcStop:
|
||||
c.releaseInflightRPCBytes(task.size)
|
||||
c.metrics.InboundRPCDropped(task.method, "scheduler_closed")
|
||||
return ErrConnClosed
|
||||
default:
|
||||
c.releaseInflightRPCBytes(task.size)
|
||||
c.metrics.InboundRPCDropped(task.method, "queue_full")
|
||||
return ErrInboundRPCQueueFull
|
||||
}
|
||||
}
|
||||
|
||||
// releaseInflightRPCBytes 归还字节预算。与 enqueueInboundRPC 的预扣严格配对:
|
||||
// 入队失败时回滚、worker 执行完(runInboundRPC)或排空丢弃(drainInboundRPCQueue)时释放。
|
||||
func (c *Conn) releaseInflightRPCBytes(size int) {
|
||||
if size > 0 {
|
||||
c.inflightRPCBytes.Add(-int64(size))
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Conn) inboundRPCWorker(rootCtx context.Context) {
|
||||
defer c.rpcWG.Done()
|
||||
for {
|
||||
select {
|
||||
case <-c.rpcStop:
|
||||
return
|
||||
default:
|
||||
}
|
||||
select {
|
||||
case task := <-c.rpcQueue:
|
||||
c.runInboundRPC(rootCtx, task)
|
||||
case <-c.rpcStop:
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Conn) runInboundRPC(rootCtx context.Context, task inboundRPC) {
|
||||
defer c.releaseInflightRPCBytes(task.size)
|
||||
queueWait := time.Since(task.enqueuedAt)
|
||||
c.metrics.InboundRPCStarted(task.method, queueWait)
|
||||
c.metrics.InboundRPCStarted(task.method, now.Sub(task.enqueuedAt))
|
||||
ctx := task.ctx
|
||||
if ctx == nil {
|
||||
ctx = context.Background()
|
||||
if task.run != nil {
|
||||
_ = task.run(ctx)
|
||||
}
|
||||
// 合并两个取消源(task.ctx 与 rootCtx)+ 超时为最少的 context 层数:
|
||||
// WithTimeout/WithCancel 的 cancel 直接作为 AfterFunc 回调,省掉单独的中间层。
|
||||
var cancel context.CancelFunc
|
||||
if c.rpcTimeout > 0 {
|
||||
ctx, cancel = context.WithTimeout(ctx, c.rpcTimeout)
|
||||
} else {
|
||||
ctx, cancel = context.WithCancel(ctx)
|
||||
}
|
||||
|
||||
func (c *Conn) finishInboundRPC(task inboundRPC) {
|
||||
if task.ticket != nil {
|
||||
task.ticket.state.Store(inboundRPCTicketDone)
|
||||
}
|
||||
stopInboundRPCTask(task)
|
||||
var reschedule bool
|
||||
c.rpcMu.Lock()
|
||||
c.rpcRunning--
|
||||
c.inflightRPCBytes.Add(-int64(task.size))
|
||||
if !c.rpcClosed && len(c.rpcQueue) > 0 && c.rpcRunning < c.rpcMaxInflight && !c.rpcReady {
|
||||
c.rpcReady = true
|
||||
reschedule = true
|
||||
}
|
||||
c.rpcMu.Unlock()
|
||||
reservation := task.budget
|
||||
// The scheduler budget may be reused immediately after release. Clear request-owned
|
||||
// closures/context references first so slow metrics/rescheduling cannot overlap the old body
|
||||
// with a newly admitted body under the same byte accounting.
|
||||
task = inboundRPC{}
|
||||
reservation.release()
|
||||
c.rpcWG.Done()
|
||||
if reschedule {
|
||||
c.rpcScheduler.schedule(c)
|
||||
}
|
||||
}
|
||||
|
||||
// expireInboundRPCTicket removes a request that is still queued and returns its
|
||||
// memory/task reservations immediately. If the worker won the dequeue race, the
|
||||
// same callback only signals the running request's response gate; its body remains
|
||||
// owned until the handler exits.
|
||||
func (c *Conn) expireInboundRPCTicket(ticket *inboundRPCTicket) {
|
||||
if ticket == nil {
|
||||
return
|
||||
}
|
||||
var (
|
||||
task inboundRPC
|
||||
found bool
|
||||
unschedule bool
|
||||
)
|
||||
c.rpcMu.Lock()
|
||||
for i := range c.rpcQueue {
|
||||
if c.rpcQueue[i].ticket != ticket {
|
||||
continue
|
||||
}
|
||||
task = c.rpcQueue[i]
|
||||
copy(c.rpcQueue[i:], c.rpcQueue[i+1:])
|
||||
last := len(c.rpcQueue) - 1
|
||||
c.rpcQueue[last] = inboundRPC{}
|
||||
c.rpcQueue = c.rpcQueue[:last]
|
||||
if len(c.rpcQueue) == 0 {
|
||||
c.rpcQueue = nil
|
||||
if c.rpcReady {
|
||||
c.rpcReady = false
|
||||
unschedule = true
|
||||
}
|
||||
}
|
||||
c.inflightRPCBytes.Add(-int64(task.size))
|
||||
ticket.state.Store(inboundRPCTicketDone)
|
||||
found = true
|
||||
break
|
||||
}
|
||||
c.rpcMu.Unlock()
|
||||
|
||||
if unschedule {
|
||||
c.rpcScheduler.unschedule(c)
|
||||
}
|
||||
if found {
|
||||
method := task.method
|
||||
reservation := task.budget
|
||||
stopInboundRPCTask(task)
|
||||
// Drop the run/context closures before returning the byte reservation. Otherwise an
|
||||
// onTimeout callback that blocks or performs a slow write can keep the copied request body
|
||||
// reachable after the global scheduler has advertised those bytes as available again.
|
||||
task = inboundRPC{}
|
||||
reservation.release()
|
||||
c.metrics.InboundRPCDropped(method, "queue_timeout")
|
||||
if ticket.onTimeout != nil {
|
||||
ticket.onTimeout()
|
||||
}
|
||||
return
|
||||
}
|
||||
if ticket.state.Load() == inboundRPCTicketRunning && ticket.onTimeout != nil {
|
||||
ticket.onTimeout()
|
||||
}
|
||||
}
|
||||
|
||||
// stopInboundRPCTask disarms callbacks before canceling the context so a normal
|
||||
// completion or connection close cannot manufacture an RPC_TIMEOUT response.
|
||||
// A deadline callback already in flight is harmless because enqueueRPC's response
|
||||
// gate makes timeout and normal rpc_result mutually exclusive.
|
||||
func stopInboundRPCTask(task inboundRPC) {
|
||||
if task.stopTimeout != nil {
|
||||
task.stopTimeout()
|
||||
}
|
||||
if task.stopRoot != nil {
|
||||
task.stopRoot()
|
||||
}
|
||||
if task.cancel != nil {
|
||||
task.cancel()
|
||||
}
|
||||
defer cancel()
|
||||
stopRoot := context.AfterFunc(rootCtx, cancel)
|
||||
defer stopRoot()
|
||||
_ = task.run(ctx)
|
||||
}
|
||||
|
||||
func (c *Conn) closeInboundRPCScheduler() {
|
||||
if c.rpcStop == nil {
|
||||
c.beginCloseInboundRPCScheduler()
|
||||
if c.rpcScheduler == nil {
|
||||
return
|
||||
}
|
||||
c.waitInboundShutdown(rpcCloseWaitTimeout)
|
||||
}
|
||||
|
||||
// beginCloseInboundRPCScheduler publishes closure, cancels running work and releases queued
|
||||
// requests without waiting for handlers. ForceClose uses this phase before transport.Close so a
|
||||
// pathological/blocking transport implementation cannot leave the RPC admission gate open.
|
||||
func (c *Conn) beginCloseInboundRPCScheduler() {
|
||||
if c.rpcScheduler == nil {
|
||||
return
|
||||
}
|
||||
c.rpcClose.Do(func() {
|
||||
c.rpcMu.Lock()
|
||||
c.rpcClosed = true
|
||||
c.rpcReady = false
|
||||
queued := c.rpcQueue
|
||||
c.rpcQueue = nil
|
||||
for i := range queued {
|
||||
c.inflightRPCBytes.Add(-int64(queued[i].size))
|
||||
}
|
||||
c.rpcMu.Unlock()
|
||||
// Remove the scheduler-owned token after rpcClosed/rpcReady become visible. schedule()
|
||||
// takes rpcMu while installing a token, so either it finishes first and is removed here,
|
||||
// or it observes the closed state and cannot enqueue a new stale token afterward.
|
||||
c.rpcScheduler.unschedule(c)
|
||||
|
||||
if c.rpcCancel != nil {
|
||||
c.rpcCancel()
|
||||
}
|
||||
close(c.rpcStop)
|
||||
// 抢占懒启动 Once:若 worker 尚未起,封住其启动,避免后续 ensureInboundRPCWorkers 的
|
||||
// rpcWG.Add 与下面的 rpcWG.Wait 并发(WaitGroup 误用)。Once 互斥保证 Add happens-before Wait。
|
||||
c.rpcWorkersOnce.Do(func() {})
|
||||
c.drainInboundRPCQueue()
|
||||
// 等 worker 退出,使关闭对 inbound 与 outbound(<-outboundDone)收敛对称;带超时防慢 handler 卡死。
|
||||
c.waitInboundWorkers(rpcCloseWaitTimeout)
|
||||
for i := range queued {
|
||||
task := queued[i]
|
||||
queued[i] = inboundRPC{}
|
||||
if task.ticket != nil {
|
||||
task.ticket.state.Store(inboundRPCTicketDone)
|
||||
}
|
||||
method := task.method
|
||||
reservation := task.budget
|
||||
stopInboundRPCTask(task)
|
||||
task = inboundRPC{}
|
||||
reservation.release()
|
||||
c.metrics.InboundRPCDropped(method, "connection_closed")
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
// waitInboundWorkers 等所有 inbound RPC worker 退出,最长 timeout。超时则放弃等待,
|
||||
// worker 在其阻塞的底层调用返回后自行退出(rpcCancel 已发,最终收敛)。
|
||||
func (c *Conn) waitInboundWorkers(timeout time.Duration) {
|
||||
// waitInboundShutdown 等 Copy 前 reservation 完成 commit/abort,以及本连接已经出队的 RPC
|
||||
// 完成,二者共用一个 timeout。超时后 reservation/共享 worker 会在底层调用最终返回时自行
|
||||
// 收敛;连接 root context 已取消。
|
||||
func (c *Conn) waitInboundShutdown(timeout time.Duration) bool {
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
c.rpcReservationWG.Wait()
|
||||
c.rpcWG.Wait()
|
||||
close(done)
|
||||
}()
|
||||
if timeout <= 0 {
|
||||
return false
|
||||
}
|
||||
timer := time.NewTimer(timeout)
|
||||
defer timer.Stop()
|
||||
select {
|
||||
case <-done:
|
||||
return true
|
||||
case <-timer.C:
|
||||
}
|
||||
}
|
||||
|
||||
func (c *Conn) drainInboundRPCQueue() {
|
||||
for {
|
||||
select {
|
||||
case task := <-c.rpcQueue:
|
||||
c.releaseInflightRPCBytes(task.size)
|
||||
c.metrics.InboundRPCDropped(task.method, "connection_closed")
|
||||
default:
|
||||
return
|
||||
}
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -8,10 +8,40 @@ import (
|
|||
"time"
|
||||
)
|
||||
|
||||
func TestInboundRPCSchedulerBoundsConcurrentWork(t *testing.T) {
|
||||
func newInboundTestConn(s *inboundRPCScheduler, maxInflight, queueSize int, timeout time.Duration) *Conn {
|
||||
c := &Conn{metrics: NopMetrics{}}
|
||||
c.startInboundRPCScheduler(2, 4, time.Second)
|
||||
defer c.closeInboundRPCScheduler()
|
||||
c.startInboundRPCScheduler(s, maxInflight, queueSize, timeout)
|
||||
return c
|
||||
}
|
||||
|
||||
func TestInboundRPCSchedulerIsLazyPerConnectionAndServer(t *testing.T) {
|
||||
scheduler := newInboundRPCScheduler(4, 16, 1<<20)
|
||||
scheduler.start()
|
||||
c := newInboundTestConn(scheduler, 2, 4, time.Second)
|
||||
defer func() {
|
||||
c.closeInboundRPCScheduler()
|
||||
scheduler.stop(time.Second)
|
||||
}()
|
||||
|
||||
if c.rpcQueue != nil {
|
||||
t.Fatal("new connection eagerly allocated an inbound queue")
|
||||
}
|
||||
scheduler.lifecycleMu.Lock()
|
||||
workersStarted := scheduler.workersStarted
|
||||
scheduler.lifecycleMu.Unlock()
|
||||
if workersStarted {
|
||||
t.Fatal("empty server eagerly started inbound RPC workers")
|
||||
}
|
||||
}
|
||||
|
||||
func TestInboundRPCSchedulerBoundsConcurrentWork(t *testing.T) {
|
||||
scheduler := newInboundRPCScheduler(2, 32, 1<<20)
|
||||
scheduler.start()
|
||||
c := newInboundTestConn(scheduler, 2, 4, time.Second)
|
||||
defer func() {
|
||||
c.closeInboundRPCScheduler()
|
||||
scheduler.stop(time.Second)
|
||||
}()
|
||||
|
||||
var active atomic.Int64
|
||||
var maxActive atomic.Int64
|
||||
|
|
@ -73,4 +103,390 @@ func TestInboundRPCSchedulerBoundsConcurrentWork(t *testing.T) {
|
|||
time.Sleep(10 * time.Millisecond)
|
||||
}
|
||||
}
|
||||
if tasks, bytes := scheduler.budgetSnapshot(); tasks != 0 || bytes != 0 {
|
||||
t.Fatalf("global budget after completion = (%d tasks, %d bytes), want zero", tasks, bytes)
|
||||
}
|
||||
}
|
||||
|
||||
func TestInboundRPCSchedulerFairAcrossConnections(t *testing.T) {
|
||||
scheduler := newInboundRPCScheduler(1, 16, 1<<20)
|
||||
c1 := newInboundTestConn(scheduler, 1, 4, time.Second)
|
||||
c2 := newInboundTestConn(scheduler, 1, 4, time.Second)
|
||||
defer func() {
|
||||
c1.closeInboundRPCScheduler()
|
||||
c2.closeInboundRPCScheduler()
|
||||
scheduler.stop(time.Second)
|
||||
}()
|
||||
|
||||
order := make(chan string, 3)
|
||||
enqueue := func(c *Conn, label string) {
|
||||
t.Helper()
|
||||
if err := c.enqueueInboundRPC(context.Background(), inboundRPC{
|
||||
method: label,
|
||||
run: func(context.Context) error {
|
||||
order <- label
|
||||
return nil
|
||||
},
|
||||
}); err != nil {
|
||||
t.Fatalf("enqueue %s: %v", label, err)
|
||||
}
|
||||
}
|
||||
|
||||
// 先在 worker 启动前形成 [c1, c2] ready 顺序。c1 每次只执行一条后回到队尾,
|
||||
// 因此 c2 必须在 c1 的第二条之前获得执行机会。
|
||||
enqueue(c1, "c1-first")
|
||||
enqueue(c1, "c1-second")
|
||||
enqueue(c2, "c2-first")
|
||||
scheduler.start()
|
||||
|
||||
want := []string{"c1-first", "c2-first", "c1-second"}
|
||||
for i := range want {
|
||||
select {
|
||||
case got := <-order:
|
||||
if got != want[i] {
|
||||
t.Fatalf("execution[%d] = %q, want %q", i, got, want[i])
|
||||
}
|
||||
case <-time.After(time.Second):
|
||||
t.Fatalf("timed out waiting for execution[%d]", i)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestInboundRPCBudgetReservedBeforeCommitAndFullyReturned(t *testing.T) {
|
||||
scheduler := newInboundRPCScheduler(1, 2, 10)
|
||||
c1 := newInboundTestConn(scheduler, 1, 4, time.Second)
|
||||
c2 := newInboundTestConn(scheduler, 1, 4, time.Second)
|
||||
defer func() {
|
||||
c1.closeInboundRPCScheduler()
|
||||
c2.closeInboundRPCScheduler()
|
||||
scheduler.stop(time.Second)
|
||||
}()
|
||||
|
||||
r1, err := c1.reserveInboundRPC(context.Background(), "one", 6)
|
||||
if err != nil {
|
||||
t.Fatalf("reserve first body: %v", err)
|
||||
}
|
||||
if tasks, bytes := scheduler.budgetSnapshot(); tasks != 1 || bytes != 6 {
|
||||
t.Fatalf("budget after first pre-Copy reservation = (%d, %d), want (1, 6)", tasks, bytes)
|
||||
}
|
||||
if _, err := c2.reserveInboundRPC(context.Background(), "too-large", 5); !errors.Is(err, ErrInboundRPCQueueFull) {
|
||||
t.Fatalf("reserve over byte budget err = %v, want queue full", err)
|
||||
}
|
||||
r2, err := c2.reserveInboundRPC(context.Background(), "two", 4)
|
||||
if err != nil {
|
||||
t.Fatalf("reserve second body: %v", err)
|
||||
}
|
||||
if _, err := c1.reserveInboundRPC(context.Background(), "too-many", 0); !errors.Is(err, ErrInboundRPCQueueFull) {
|
||||
t.Fatalf("reserve over task budget err = %v, want queue full", err)
|
||||
}
|
||||
|
||||
r1.abort()
|
||||
r2.abort()
|
||||
if tasks, bytes := scheduler.budgetSnapshot(); tasks != 0 || bytes != 0 {
|
||||
t.Fatalf("budget after aborts = (%d, %d), want zero", tasks, bytes)
|
||||
}
|
||||
if got := c1.inflightRPCBytes.Load(); got != 0 {
|
||||
t.Fatalf("c1 inflight bytes = %d, want zero", got)
|
||||
}
|
||||
if got := c2.inflightRPCBytes.Load(); got != 0 {
|
||||
t.Fatalf("c2 inflight bytes = %d, want zero", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestInboundRPCPerConnectionByteBudgetRejectedBeforeCommit(t *testing.T) {
|
||||
scheduler := newInboundRPCScheduler(1, 2, int64(maxInflightRPCBytes)+1)
|
||||
c := newInboundTestConn(scheduler, 1, 2, time.Second)
|
||||
defer func() {
|
||||
c.closeInboundRPCScheduler()
|
||||
scheduler.stop(time.Second)
|
||||
}()
|
||||
|
||||
if _, err := c.reserveInboundRPC(context.Background(), "oversized", maxInflightRPCBytes+1); !errors.Is(err, ErrInboundRPCQueueFull) {
|
||||
t.Fatalf("reserve over per-connection byte budget err = %v, want queue full", err)
|
||||
}
|
||||
if tasks, bytes := scheduler.budgetSnapshot(); tasks != 0 || bytes != 0 {
|
||||
t.Fatalf("global budget after per-connection rejection = (%d, %d), want zero", tasks, bytes)
|
||||
}
|
||||
if got := c.inflightRPCBytes.Load(); got != 0 {
|
||||
t.Fatalf("connection bytes after rejection = %d, want zero", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestInboundRPCCommitRacingCloseReturnsReservation(t *testing.T) {
|
||||
scheduler := newInboundRPCScheduler(1, 4, 1<<20)
|
||||
c := newInboundTestConn(scheduler, 1, 2, time.Second)
|
||||
defer scheduler.stop(time.Second)
|
||||
|
||||
reservation, err := c.reserveInboundRPC(context.Background(), "closing", 13)
|
||||
if err != nil {
|
||||
t.Fatalf("reserve: %v", err)
|
||||
}
|
||||
closed := make(chan struct{})
|
||||
go func() {
|
||||
c.closeInboundRPCScheduler()
|
||||
close(closed)
|
||||
}()
|
||||
deadline := time.Now().Add(time.Second)
|
||||
for {
|
||||
c.rpcMu.Lock()
|
||||
isClosed := c.rpcClosed
|
||||
c.rpcMu.Unlock()
|
||||
if isClosed {
|
||||
break
|
||||
}
|
||||
if time.Now().After(deadline) {
|
||||
t.Fatal("connection scheduler was not marked closed")
|
||||
}
|
||||
time.Sleep(time.Millisecond)
|
||||
}
|
||||
|
||||
if err := reservation.commit(inboundRPC{run: func(context.Context) error { return nil }}); !errors.Is(err, ErrConnClosed) {
|
||||
t.Fatalf("commit after close err = %v, want ErrConnClosed", err)
|
||||
}
|
||||
select {
|
||||
case <-closed:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("close did not finish after reservation commit")
|
||||
}
|
||||
if tasks, bytes := scheduler.budgetSnapshot(); tasks != 0 || bytes != 0 {
|
||||
t.Fatalf("global budget after close/commit race = (%d, %d), want zero", tasks, bytes)
|
||||
}
|
||||
if got := c.inflightRPCBytes.Load(); got != 0 {
|
||||
t.Fatalf("connection bytes after close/commit race = %d, want zero", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestInboundRPCSchedulerCloseRemovesReadyTokenBeforeStart(t *testing.T) {
|
||||
scheduler := newInboundRPCScheduler(1, 1, 1<<20)
|
||||
defer scheduler.stop(time.Second)
|
||||
|
||||
// A bounded ready channel used to retain one stale token per closed connection. With workers
|
||||
// not started yet, the second connection then blocked forever trying to publish its token even
|
||||
// though the first connection had returned every task/byte budget.
|
||||
for i := 0; i < 32; i++ {
|
||||
c := newInboundTestConn(scheduler, 1, 1, time.Second)
|
||||
done := make(chan error, 1)
|
||||
go func() {
|
||||
done <- c.enqueueInboundRPC(context.Background(), inboundRPC{
|
||||
method: "close-before-start",
|
||||
run: func(context.Context) error { return nil },
|
||||
})
|
||||
}()
|
||||
select {
|
||||
case err := <-done:
|
||||
if err != nil {
|
||||
t.Fatalf("enqueue iteration %d: %v", i, err)
|
||||
}
|
||||
case <-time.After(time.Second):
|
||||
t.Fatalf("enqueue iteration %d blocked behind a stale ready token", i)
|
||||
}
|
||||
c.closeInboundRPCScheduler()
|
||||
if got := scheduler.readyLen(); got != 0 {
|
||||
t.Fatalf("ready tokens after close iteration %d = %d, want zero", i, got)
|
||||
}
|
||||
}
|
||||
if tasks, bytes := scheduler.budgetSnapshot(); tasks != 0 || bytes != 0 {
|
||||
t.Fatalf("budget after close churn = (%d, %d), want zero", tasks, bytes)
|
||||
}
|
||||
}
|
||||
|
||||
func TestInboundRPCExpiredInQueueNeverRunsAndSignalsTimeout(t *testing.T) {
|
||||
scheduler := newInboundRPCScheduler(1, 8, 1<<20)
|
||||
scheduler.start()
|
||||
c := newInboundTestConn(scheduler, 1, 4, 40*time.Millisecond)
|
||||
defer func() {
|
||||
c.closeInboundRPCScheduler()
|
||||
scheduler.stop(time.Second)
|
||||
}()
|
||||
|
||||
started := make(chan struct{})
|
||||
release := make(chan struct{})
|
||||
if err := c.enqueueInboundRPC(context.Background(), inboundRPC{
|
||||
method: "blocker",
|
||||
size: 7,
|
||||
run: func(context.Context) error {
|
||||
close(started)
|
||||
<-release // 刻意忽略 deadline,确保下一条在队列中到期。
|
||||
return nil
|
||||
},
|
||||
}); err != nil {
|
||||
t.Fatalf("enqueue blocker: %v", err)
|
||||
}
|
||||
select {
|
||||
case <-started:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("blocker did not start")
|
||||
}
|
||||
|
||||
var ran atomic.Bool
|
||||
timedOut := make(chan struct{})
|
||||
if err := c.enqueueInboundRPC(context.Background(), inboundRPC{
|
||||
method: "expires",
|
||||
size: 11,
|
||||
onTimeout: func() {
|
||||
close(timedOut)
|
||||
},
|
||||
run: func(context.Context) error {
|
||||
ran.Store(true)
|
||||
return nil
|
||||
},
|
||||
}); err != nil {
|
||||
t.Fatalf("enqueue expiring task: %v", err)
|
||||
}
|
||||
select {
|
||||
case <-timedOut:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("queued task did not signal timeout while the worker was still blocked")
|
||||
}
|
||||
deadline := time.Now().Add(time.Second)
|
||||
for {
|
||||
tasks, bytes := scheduler.budgetSnapshot()
|
||||
if tasks == 1 && bytes == 7 {
|
||||
break
|
||||
}
|
||||
if time.Now().After(deadline) {
|
||||
t.Fatalf("budget while blocker still runs = (%d, %d), want only blocker (1, 7)", tasks, bytes)
|
||||
}
|
||||
time.Sleep(time.Millisecond)
|
||||
}
|
||||
close(release)
|
||||
if ran.Load() {
|
||||
t.Fatal("expired queued task entered business handler")
|
||||
}
|
||||
|
||||
deadline = time.Now().Add(time.Second)
|
||||
for {
|
||||
tasks, bytes := scheduler.budgetSnapshot()
|
||||
if tasks == 0 && bytes == 0 {
|
||||
break
|
||||
}
|
||||
if time.Now().After(deadline) {
|
||||
t.Fatalf("budget after timeout = (%d, %d), want zero", tasks, bytes)
|
||||
}
|
||||
time.Sleep(time.Millisecond)
|
||||
}
|
||||
}
|
||||
|
||||
func TestInboundRPCCloseDisarmsQueuedTimeout(t *testing.T) {
|
||||
scheduler := newInboundRPCScheduler(1, 8, 1<<20)
|
||||
c := newInboundTestConn(scheduler, 1, 4, 30*time.Millisecond)
|
||||
defer scheduler.stop(time.Second)
|
||||
|
||||
timedOut := make(chan struct{}, 1)
|
||||
if err := c.enqueueInboundRPC(context.Background(), inboundRPC{
|
||||
method: "queued",
|
||||
size: 11,
|
||||
onTimeout: func() {
|
||||
timedOut <- struct{}{}
|
||||
},
|
||||
}); err != nil {
|
||||
t.Fatalf("enqueue queued task: %v", err)
|
||||
}
|
||||
c.closeInboundRPCScheduler()
|
||||
time.Sleep(60 * time.Millisecond)
|
||||
select {
|
||||
case <-timedOut:
|
||||
t.Fatal("connection close emitted a queued RPC timeout")
|
||||
default:
|
||||
}
|
||||
}
|
||||
|
||||
func TestInboundRPCRunningTimeoutSignalsWithoutReleasingBodyEarly(t *testing.T) {
|
||||
scheduler := newInboundRPCScheduler(1, 8, 1<<20)
|
||||
scheduler.start()
|
||||
c := newInboundTestConn(scheduler, 1, 4, 30*time.Millisecond)
|
||||
defer func() {
|
||||
c.closeInboundRPCScheduler()
|
||||
scheduler.stop(time.Second)
|
||||
}()
|
||||
|
||||
started := make(chan struct{})
|
||||
release := make(chan struct{})
|
||||
timedOut := make(chan struct{}, 1)
|
||||
if err := c.enqueueInboundRPC(context.Background(), inboundRPC{
|
||||
method: "running",
|
||||
size: 7,
|
||||
onTimeout: func() {
|
||||
timedOut <- struct{}{}
|
||||
},
|
||||
run: func(context.Context) error {
|
||||
close(started)
|
||||
<-release
|
||||
return nil
|
||||
},
|
||||
}); err != nil {
|
||||
t.Fatalf("enqueue running task: %v", err)
|
||||
}
|
||||
<-started
|
||||
select {
|
||||
case <-timedOut:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("running task did not signal timeout while handler ignored cancellation")
|
||||
}
|
||||
if tasks, bytes := scheduler.budgetSnapshot(); tasks != 1 || bytes != 7 {
|
||||
t.Fatalf("running body budget after timeout = (%d, %d), want retained (1, 7)", tasks, bytes)
|
||||
}
|
||||
close(release)
|
||||
deadline := time.Now().Add(time.Second)
|
||||
for {
|
||||
tasks, bytes := scheduler.budgetSnapshot()
|
||||
if tasks == 0 && bytes == 0 {
|
||||
break
|
||||
}
|
||||
if time.Now().After(deadline) {
|
||||
t.Fatalf("running body budget after completion = (%d, %d), want zero", tasks, bytes)
|
||||
}
|
||||
time.Sleep(time.Millisecond)
|
||||
}
|
||||
}
|
||||
|
||||
func TestInboundRPCCloseDrainsQueueAndReturnsBudgets(t *testing.T) {
|
||||
scheduler := newInboundRPCScheduler(1, 8, 1<<20)
|
||||
scheduler.start()
|
||||
c := newInboundTestConn(scheduler, 1, 4, time.Second)
|
||||
defer scheduler.stop(time.Second)
|
||||
|
||||
started := make(chan struct{})
|
||||
if err := c.enqueueInboundRPC(context.Background(), inboundRPC{
|
||||
method: "running",
|
||||
size: 7,
|
||||
run: func(ctx context.Context) error {
|
||||
close(started)
|
||||
<-ctx.Done()
|
||||
return ctx.Err()
|
||||
},
|
||||
}); err != nil {
|
||||
t.Fatalf("enqueue running task: %v", err)
|
||||
}
|
||||
select {
|
||||
case <-started:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("running task did not start")
|
||||
}
|
||||
|
||||
var queuedRan atomic.Bool
|
||||
if err := c.enqueueInboundRPC(context.Background(), inboundRPC{
|
||||
method: "queued",
|
||||
size: 11,
|
||||
run: func(context.Context) error {
|
||||
queuedRan.Store(true)
|
||||
return nil
|
||||
},
|
||||
}); err != nil {
|
||||
t.Fatalf("enqueue queued task: %v", err)
|
||||
}
|
||||
if tasks, bytes := scheduler.budgetSnapshot(); tasks != 2 || bytes != 18 {
|
||||
t.Fatalf("budget before close = (%d, %d), want (2, 18)", tasks, bytes)
|
||||
}
|
||||
|
||||
c.closeInboundRPCScheduler()
|
||||
if queuedRan.Load() {
|
||||
t.Fatal("queued task ran during connection close")
|
||||
}
|
||||
if tasks, bytes := scheduler.budgetSnapshot(); tasks != 0 || bytes != 0 {
|
||||
t.Fatalf("budget after close = (%d, %d), want zero", tasks, bytes)
|
||||
}
|
||||
if got := c.inflightRPCBytes.Load(); got != 0 {
|
||||
t.Fatalf("connection inflight bytes after close = %d, want zero", got)
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -71,11 +71,16 @@ func TestLoginEmailEndToEnd(t *testing.T) {
|
|||
passwordStore := memory.NewPasswordStore()
|
||||
helpStore := memory.NewHelpStore()
|
||||
codeStore := memory.NewCodeStore()
|
||||
dialogStore := memory.NewDialogStore()
|
||||
messageStore := memory.NewMessageStore(dialogStore)
|
||||
updateEventStore := memory.NewUpdateEventStore()
|
||||
emailSender := &loginEmailTestSender{}
|
||||
accountService := account.NewService(passwordStore,
|
||||
account.WithUsers(userStore),
|
||||
account.WithLoginEmailVerification(codeStore, emailSender, 5*time.Minute, 5, 6))
|
||||
authService := auth.NewService(userStore, authzStore, codeStore, authKeyStore, memory.NewTempAuthKeyBindingStore(), code,
|
||||
auth.WithLoginMessages(messageStore, dialogStore),
|
||||
auth.WithLoginCodeDelivery(memory.NewLoginCodeDeliveryStore(messageStore, updateEventStore)),
|
||||
auth.WithPasswords(passwordStore),
|
||||
auth.WithLoginEmail(auth.LoginEmailOptions{
|
||||
Enabled: true,
|
||||
|
|
@ -89,10 +94,10 @@ func TestLoginEmailEndToEnd(t *testing.T) {
|
|||
Account: accountService,
|
||||
Help: help.NewService(helpStore, helpStore),
|
||||
Users: users.NewService(userStore),
|
||||
Updates: updates.NewService(memory.NewUpdateStateStore(), memory.NewUpdateEventStore()),
|
||||
Updates: updates.NewService(memory.NewUpdateStateStore(), updateEventStore),
|
||||
|
||||
Contacts: contacts.NewService(memory.NewContactStore()),
|
||||
Dialogs: dialogs.NewService(memory.NewDialogStore()),
|
||||
Dialogs: dialogs.NewService(dialogStore),
|
||||
}
|
||||
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, RPC: router})
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
126
internal/mtprotoedge/outbound_scratch.go
Normal file
126
internal/mtprotoedge/outbound_scratch.go
Normal file
|
|
@ -0,0 +1,126 @@
|
|||
package mtprotoedge
|
||||
|
||||
import (
|
||||
"context"
|
||||
"time"
|
||||
|
||||
"github.com/gotd/td/bin"
|
||||
)
|
||||
|
||||
const (
|
||||
defaultOutboundWriteMaxBytes = int64(512 << 20)
|
||||
defaultOutboundScratchPool = 256
|
||||
)
|
||||
|
||||
// outboundScratchPool bounds and reuses the encrypted wire buffer across connections. A lease
|
||||
// reserves a conservative 3x wire size while writing (wire + codec/obfuscation copies), then
|
||||
// shrinks to the actual retained capacity while idle in the bounded pool. Large one-off frames are
|
||||
// dropped on return. This removes attacker-warmable per-Conn MiB buffers without returning to an
|
||||
// unbounded allocation-per-message design.
|
||||
type outboundScratchPool struct {
|
||||
budget *outboundTrackedBudget
|
||||
idle chan *outboundScratch
|
||||
}
|
||||
|
||||
type outboundScratch struct {
|
||||
wire bin.Buffer
|
||||
reserved int
|
||||
}
|
||||
|
||||
func newOutboundScratchPool(maxBytes int64) *outboundScratchPool {
|
||||
if maxBytes <= 0 {
|
||||
maxBytes = defaultOutboundWriteMaxBytes
|
||||
}
|
||||
return &outboundScratchPool{
|
||||
budget: newOutboundTrackedBudget(maxBytes),
|
||||
idle: make(chan *outboundScratch, defaultOutboundScratchPool),
|
||||
}
|
||||
}
|
||||
|
||||
func (p *outboundScratchPool) acquire(ctx context.Context, stop <-chan struct{}, wireBytes int) (*outboundScratch, error) {
|
||||
return p.acquireUntil(ctx, stop, wireBytes, time.Time{})
|
||||
}
|
||||
|
||||
func (p *outboundScratchPool) acquireUntil(ctx context.Context, stop <-chan struct{}, wireBytes int, deadline time.Time) (*outboundScratch, error) {
|
||||
if p == nil || wireBytes <= 0 {
|
||||
return nil, ErrOutboundMessageTooLarge
|
||||
}
|
||||
peak := wireBytes * 3
|
||||
if peak < wireBytes { // int overflow
|
||||
return nil, ErrOutboundMessageTooLarge
|
||||
}
|
||||
|
||||
var scratch *outboundScratch
|
||||
select {
|
||||
case scratch = <-p.idle:
|
||||
default:
|
||||
}
|
||||
if scratch == nil {
|
||||
if err := p.budget.waitReserveUntil(ctx, stop, peak, deadline); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &outboundScratch{wire: bin.Buffer{Buf: make([]byte, wireBytes)}, reserved: peak}, nil
|
||||
}
|
||||
|
||||
if cap(scratch.wire.Buf) >= wireBytes {
|
||||
if extra := peak - scratch.reserved; extra > 0 {
|
||||
if err := p.budget.waitReserveUntil(ctx, stop, extra, deadline); err != nil {
|
||||
p.putIdle(scratch)
|
||||
return nil, err
|
||||
}
|
||||
scratch.reserved += extra
|
||||
}
|
||||
scratch.wire.Buf = scratch.wire.Buf[:wireBytes]
|
||||
return scratch, nil
|
||||
}
|
||||
|
||||
// The old slice is no longer reachable after clearing it; return that retained charge before
|
||||
// waiting for a larger lease, otherwise old+peak may exceed the budget and deadlock a resize
|
||||
// that would fit after replacement.
|
||||
old := scratch.reserved
|
||||
scratch.wire.Buf = nil
|
||||
scratch.reserved = 0
|
||||
p.budget.release(old)
|
||||
if err := p.budget.waitReserveUntil(ctx, stop, peak, deadline); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
scratch.wire.Buf = make([]byte, wireBytes)
|
||||
scratch.reserved = peak
|
||||
return scratch, nil
|
||||
}
|
||||
|
||||
func (p *outboundScratchPool) release(scratch *outboundScratch) {
|
||||
if p == nil || scratch == nil {
|
||||
return
|
||||
}
|
||||
retained := cap(scratch.wire.Buf)
|
||||
if retained > maxRetainedConnBuffer {
|
||||
p.budget.release(scratch.reserved)
|
||||
scratch.wire.Buf = nil
|
||||
scratch.reserved = 0
|
||||
return
|
||||
}
|
||||
if scratch.reserved > retained {
|
||||
p.budget.release(scratch.reserved - retained)
|
||||
scratch.reserved = retained
|
||||
}
|
||||
scratch.wire.Buf = scratch.wire.Buf[:0]
|
||||
p.putIdle(scratch)
|
||||
}
|
||||
|
||||
func (p *outboundScratchPool) putIdle(scratch *outboundScratch) {
|
||||
select {
|
||||
case p.idle <- scratch:
|
||||
default:
|
||||
p.budget.release(scratch.reserved)
|
||||
scratch.wire.Buf = nil
|
||||
scratch.reserved = 0
|
||||
}
|
||||
}
|
||||
|
||||
func (p *outboundScratchPool) snapshot() int64 {
|
||||
if p == nil {
|
||||
return 0
|
||||
}
|
||||
return p.budget.snapshot()
|
||||
}
|
||||
|
|
@ -4,7 +4,10 @@ import (
|
|||
"bytes"
|
||||
"context"
|
||||
"crypto/rand"
|
||||
"errors"
|
||||
"io"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
|
|
@ -13,8 +16,426 @@ import (
|
|||
"github.com/gotd/td/mt"
|
||||
"github.com/gotd/td/proto"
|
||||
"github.com/gotd/td/tg"
|
||||
"github.com/gotd/td/transport"
|
||||
)
|
||||
|
||||
type failAfterTransport struct {
|
||||
failAt atomic.Int32
|
||||
sends atomic.Int32
|
||||
stored atomic.Int32
|
||||
closes atomic.Int32
|
||||
mu sync.Mutex
|
||||
last []byte
|
||||
}
|
||||
|
||||
type blockingOutboundTransport struct {
|
||||
started chan struct{}
|
||||
release chan struct{}
|
||||
once sync.Once
|
||||
sends atomic.Int32
|
||||
}
|
||||
|
||||
type blockingEncodeProbe struct {
|
||||
started chan struct{}
|
||||
release <-chan struct{}
|
||||
active atomic.Int32
|
||||
max atomic.Int32
|
||||
}
|
||||
|
||||
func (e *blockingEncodeProbe) Encode(b *bin.Buffer) error {
|
||||
active := e.active.Add(1)
|
||||
for {
|
||||
max := e.max.Load()
|
||||
if active <= max || e.max.CompareAndSwap(max, active) {
|
||||
break
|
||||
}
|
||||
}
|
||||
e.started <- struct{}{}
|
||||
<-e.release
|
||||
e.active.Add(-1)
|
||||
b.PutID(tg.UpdatesTooLongTypeID)
|
||||
return nil
|
||||
}
|
||||
|
||||
func newBlockingOutboundTransport() *blockingOutboundTransport {
|
||||
return &blockingOutboundTransport{started: make(chan struct{}), release: make(chan struct{})}
|
||||
}
|
||||
|
||||
func TestOutboundEncodingHasProcessWideConcurrencyBudget(t *testing.T) {
|
||||
const extra = 8
|
||||
total := defaultOutboundEncodeConcurrency + extra
|
||||
release := make(chan struct{})
|
||||
probe := &blockingEncodeProbe{
|
||||
started: make(chan struct{}, total),
|
||||
release: release,
|
||||
}
|
||||
errs := make(chan error, total)
|
||||
for range total {
|
||||
go func() {
|
||||
_, err := encodeOutboundMessage(probe)
|
||||
errs <- err
|
||||
}()
|
||||
}
|
||||
|
||||
for range defaultOutboundEncodeConcurrency {
|
||||
select {
|
||||
case <-probe.started:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("encode workers did not fill concurrency budget")
|
||||
}
|
||||
}
|
||||
select {
|
||||
case <-probe.started:
|
||||
t.Fatalf("more than %d outbound encodes ran concurrently", defaultOutboundEncodeConcurrency)
|
||||
case <-time.After(50 * time.Millisecond):
|
||||
}
|
||||
|
||||
close(release)
|
||||
for range total {
|
||||
if err := <-errs; err != nil {
|
||||
t.Fatalf("encode: %v", err)
|
||||
}
|
||||
}
|
||||
if got := probe.max.Load(); got != defaultOutboundEncodeConcurrency {
|
||||
t.Fatalf("peak concurrent encodes = %d, want %d", got, defaultOutboundEncodeConcurrency)
|
||||
}
|
||||
}
|
||||
|
||||
func TestConnectionCloseDoesNotWaitForRunningEncoder(t *testing.T) {
|
||||
release := make(chan struct{})
|
||||
probe := &blockingEncodeProbe{started: make(chan struct{}, 1), release: release}
|
||||
c := &Conn{metrics: NopMetrics{}}
|
||||
c.startOutbound()
|
||||
sendDone := make(chan error, 1)
|
||||
go func() {
|
||||
sendDone <- c.Send(context.Background(), proto.MessageFromServer, probe)
|
||||
}()
|
||||
select {
|
||||
case <-probe.started:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("encoder did not start")
|
||||
}
|
||||
|
||||
closeDone := make(chan struct{})
|
||||
go func() {
|
||||
c.Close()
|
||||
close(closeDone)
|
||||
}()
|
||||
select {
|
||||
case <-closeDone:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("Conn.Close waited for external Encoder")
|
||||
}
|
||||
close(release)
|
||||
select {
|
||||
case err := <-sendDone:
|
||||
if !errors.Is(err, ErrConnClosed) {
|
||||
t.Fatalf("send after concurrent close = %v, want ErrConnClosed", err)
|
||||
}
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("send did not return after encoder release")
|
||||
}
|
||||
}
|
||||
|
||||
func TestOutboundControlVectorsUseGlobalByteBudget(t *testing.T) {
|
||||
budget := newOutboundTrackedBudget(16)
|
||||
c := &Conn{outboundControlTrackedBudget: budget}
|
||||
op, err := c.newOutboundVectorOp(outboundAck, []int64{1, 2})
|
||||
if err != nil {
|
||||
t.Fatalf("reserve first vector: %v", err)
|
||||
}
|
||||
if got := budget.snapshot(); got != 16 {
|
||||
t.Fatalf("tracked bytes after reserve = %d, want 16", got)
|
||||
}
|
||||
if _, err := c.newOutboundVectorOp(outboundResend, []int64{3}); !errors.Is(err, ErrOutboundTrackedBudget) {
|
||||
t.Fatalf("reserve over budget error = %v, want %v", err, ErrOutboundTrackedBudget)
|
||||
}
|
||||
op.releaseReservation(budget)
|
||||
if got := budget.snapshot(); got != 0 {
|
||||
t.Fatalf("tracked bytes after release = %d, want 0", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEncodedControlFramesUseIndependentBudgetForQueuedAndPendingLifetime(t *testing.T) {
|
||||
bodyBudget := newOutboundTrackedBudget(4)
|
||||
controlBudget := newOutboundTrackedBudget(256)
|
||||
tr := &failAfterTransport{}
|
||||
c := newOutboundTestConn(t, tr, bodyBudget)
|
||||
c.outboundControlTrackedBudget = controlBudget
|
||||
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
|
||||
defer cancel()
|
||||
|
||||
// One content frame fills the ordinary body budget and remains pending.
|
||||
if err := c.Send(ctx, proto.MessageFromServer, &tg.UpdatesTooLong{}); err != nil {
|
||||
t.Fatalf("fill body budget: %v", err)
|
||||
}
|
||||
first, err := crypto.NewClientCipher(rand.Reader).DecryptFromBuffer(c.key, &bin.Buffer{Buf: tr.lastFrame()})
|
||||
if err != nil {
|
||||
t.Fatalf("decrypt ordinary frame: %v", err)
|
||||
}
|
||||
if got := bodyBudget.snapshot(); got != 4 {
|
||||
t.Fatalf("body budget = %d, want saturated 4", got)
|
||||
}
|
||||
|
||||
created := &mt.NewSessionCreated{FirstMsgID: 1, UniqueID: 2, ServerSalt: 3}
|
||||
encodedCreated, err := encodeOutboundMessageWithoutSlot(created)
|
||||
if err != nil {
|
||||
t.Fatalf("encode new_session_created: %v", err)
|
||||
}
|
||||
if err := c.SendAsync(ctx, proto.MessageFromServer, created); err != nil {
|
||||
t.Fatalf("new_session_created under saturated body budget: %v", err)
|
||||
}
|
||||
deadline := time.Now().Add(time.Second)
|
||||
for tr.stored.Load() < 2 && time.Now().Before(deadline) {
|
||||
time.Sleep(time.Millisecond)
|
||||
}
|
||||
if got := tr.stored.Load(); got != 2 {
|
||||
t.Fatalf("completed physical sends = %d, want 2", got)
|
||||
}
|
||||
second, err := crypto.NewClientCipher(rand.Reader).DecryptFromBuffer(c.key, &bin.Buffer{Buf: tr.lastFrame()})
|
||||
if err != nil {
|
||||
t.Fatalf("decrypt control frame: %v", err)
|
||||
}
|
||||
if got := bodyBudget.snapshot(); got != 4 {
|
||||
t.Fatalf("body budget after control send = %d, want unchanged 4", got)
|
||||
}
|
||||
if got := controlBudget.snapshot(); got != int64(len(encodedCreated.body)) {
|
||||
t.Fatalf("control pending budget = %d, want new_session_created body %d", got, len(encodedCreated.body))
|
||||
}
|
||||
select {
|
||||
case <-c.outboundDone:
|
||||
t.Fatal("ordinary body pressure closed a healthy connection")
|
||||
default:
|
||||
}
|
||||
|
||||
// Pong is non-pending, but must also remain admissible and return its control bytes after write.
|
||||
if err := c.SendAsync(ctx, proto.MessageServerResponse, &mt.Pong{MsgID: 4, PingID: 5}); err != nil {
|
||||
t.Fatalf("pong under saturated body budget: %v", err)
|
||||
}
|
||||
deadline = time.Now().Add(time.Second)
|
||||
for tr.stored.Load() < 3 && time.Now().Before(deadline) {
|
||||
time.Sleep(time.Millisecond)
|
||||
}
|
||||
if got := tr.stored.Load(); got != 3 {
|
||||
t.Fatalf("completed physical sends after pong = %d, want 3", got)
|
||||
}
|
||||
if got := controlBudget.snapshot(); got != int64(len(encodedCreated.body)) {
|
||||
t.Fatalf("control budget after non-pending pong = %d, want pending %d", got, len(encodedCreated.body))
|
||||
}
|
||||
|
||||
c.AckServerMessages([]int64{first.MessageID, second.MessageID})
|
||||
deadline = time.Now().Add(time.Second)
|
||||
for (bodyBudget.snapshot() != 0 || controlBudget.snapshot() != 0) && time.Now().Before(deadline) {
|
||||
time.Sleep(time.Millisecond)
|
||||
}
|
||||
if got := bodyBudget.snapshot(); got != 0 {
|
||||
t.Fatalf("body budget after ACK = %d, want 0", got)
|
||||
}
|
||||
if got := controlBudget.snapshot(); got != 0 {
|
||||
t.Fatalf("control budget after ACK = %d, want 0", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOutboundScratchPoolBoundsConcurrentWireCopies(t *testing.T) {
|
||||
pool := newOutboundScratchPool(300)
|
||||
first, err := pool.acquire(context.Background(), nil, 100) // 3x peak = full budget.
|
||||
if err != nil {
|
||||
t.Fatalf("acquire first scratch: %v", err)
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 20*time.Millisecond)
|
||||
defer cancel()
|
||||
if _, err := pool.acquire(ctx, nil, 100); !errors.Is(err, context.DeadlineExceeded) {
|
||||
t.Fatalf("second concurrent acquire = %v, want deadline backpressure", err)
|
||||
}
|
||||
pool.release(first)
|
||||
if got := pool.snapshot(); got != 100 {
|
||||
t.Fatalf("idle retained scratch = %d, want 100", got)
|
||||
}
|
||||
second, err := pool.acquire(context.Background(), nil, 100)
|
||||
if err != nil {
|
||||
t.Fatalf("reuse retained scratch: %v", err)
|
||||
}
|
||||
pool.release(second)
|
||||
if got := pool.snapshot(); got != 100 {
|
||||
t.Fatalf("scratch after reuse = %d, want one bounded idle buffer", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOutboundScratchAdmissionUsesWriteTimeoutWithoutClosingHealthyConnection(t *testing.T) {
|
||||
wireBytes := encryptedOutboundWireLen(4)
|
||||
pool := newOutboundScratchPool(int64(wireBytes * 3))
|
||||
blocker, err := pool.acquire(context.Background(), nil, wireBytes)
|
||||
if err != nil {
|
||||
t.Fatalf("occupy shared scratch budget: %v", err)
|
||||
}
|
||||
|
||||
tr := &failAfterTransport{}
|
||||
c := newOutboundTestConn(t, tr, newOutboundTrackedBudget(1<<20))
|
||||
c.outboundScratchPool = pool
|
||||
c.writeTimeout = 25 * time.Millisecond
|
||||
|
||||
start := time.Now()
|
||||
err = c.Send(context.Background(), proto.MessageFromServer, &tg.UpdatesTooLong{})
|
||||
elapsed := time.Since(start)
|
||||
if !errors.Is(err, context.DeadlineExceeded) {
|
||||
t.Fatalf("scratch admission err = %v, want deadline exceeded", err)
|
||||
}
|
||||
if elapsed > 250*time.Millisecond {
|
||||
t.Fatalf("scratch admission waited %v, want writeTimeout-bounded wait", elapsed)
|
||||
}
|
||||
if got := tr.sends.Load(); got != 0 {
|
||||
t.Fatalf("writer called %d times without scratch, want 0", got)
|
||||
}
|
||||
if c.terminal.Load() {
|
||||
t.Fatal("scratch admission timeout terminally closed a healthy connection")
|
||||
}
|
||||
select {
|
||||
case <-c.outboundDone:
|
||||
t.Fatal("outbound actor exited after pre-write scratch timeout")
|
||||
default:
|
||||
}
|
||||
|
||||
pool.release(blocker)
|
||||
c.writeTimeout = time.Second
|
||||
if err := c.Send(context.Background(), proto.MessageFromServer, &tg.UpdatesTooLong{}); err != nil {
|
||||
t.Fatalf("send after scratch capacity returned: %v", err)
|
||||
}
|
||||
if got := tr.sends.Load(); got != 1 {
|
||||
t.Fatalf("writer calls after recovery = %d, want 1", got)
|
||||
}
|
||||
}
|
||||
|
||||
func (t *blockingOutboundTransport) Send(context.Context, *bin.Buffer) error {
|
||||
if t.sends.Add(1) == 1 {
|
||||
close(t.started)
|
||||
}
|
||||
<-t.release
|
||||
return io.ErrClosedPipe
|
||||
}
|
||||
|
||||
func (t *blockingOutboundTransport) Recv(context.Context, *bin.Buffer) error { return io.EOF }
|
||||
func (t *blockingOutboundTransport) Close() error {
|
||||
t.once.Do(func() { close(t.release) })
|
||||
return nil
|
||||
}
|
||||
|
||||
func (t *failAfterTransport) Send(_ context.Context, b *bin.Buffer) error {
|
||||
n := t.sends.Add(1)
|
||||
if failAt := t.failAt.Load(); failAt > 0 && n >= failAt {
|
||||
return io.ErrClosedPipe
|
||||
}
|
||||
t.mu.Lock()
|
||||
t.last = append(t.last[:0], b.Raw()...)
|
||||
t.mu.Unlock()
|
||||
t.stored.Add(1)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (t *failAfterTransport) Recv(context.Context, *bin.Buffer) error { return io.EOF }
|
||||
func (t *failAfterTransport) Close() error {
|
||||
t.closes.Add(1)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (t *failAfterTransport) lastFrame() []byte {
|
||||
t.mu.Lock()
|
||||
defer t.mu.Unlock()
|
||||
return append([]byte(nil), t.last...)
|
||||
}
|
||||
|
||||
func newOutboundFailureTestConn(t *testing.T, tr transport.Conn) *Conn {
|
||||
return newOutboundTestConn(t, tr, nil)
|
||||
}
|
||||
|
||||
func newOutboundTestConn(t *testing.T, tr transport.Conn, budget *outboundTrackedBudget) *Conn {
|
||||
t.Helper()
|
||||
var key crypto.Key
|
||||
if _, err := rand.Read(key[:]); err != nil {
|
||||
t.Fatalf("rand key: %v", err)
|
||||
}
|
||||
c := &Conn{
|
||||
transport: tr,
|
||||
writer: tr,
|
||||
cipher: crypto.NewServerCipher(rand.Reader),
|
||||
msgID: proto.NewMessageIDGen(time.Now),
|
||||
writeTimeout: time.Second,
|
||||
metrics: NopMetrics{},
|
||||
key: key.WithID(),
|
||||
salt: 123,
|
||||
sessionID: 456,
|
||||
outboundTrackedBudget: budget,
|
||||
}
|
||||
c.startOutbound()
|
||||
t.Cleanup(c.Close)
|
||||
return c
|
||||
}
|
||||
|
||||
func TestOutboundQueueBackingUsesSmallConfigurableBounds(t *testing.T) {
|
||||
t.Run("defaults", func(t *testing.T) {
|
||||
c := &Conn{metrics: NopMetrics{}}
|
||||
c.startOutbound()
|
||||
defer c.Close()
|
||||
if got := cap(c.outbound); got != defaultOutboundQueueSize {
|
||||
t.Fatalf("normal queue cap = %d, want %d", got, defaultOutboundQueueSize)
|
||||
}
|
||||
if got := cap(c.outboundControl); got != defaultOutboundControlQueueSize {
|
||||
t.Fatalf("control queue cap = %d, want %d", got, defaultOutboundControlQueueSize)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("configured", func(t *testing.T) {
|
||||
c := &Conn{
|
||||
metrics: NopMetrics{},
|
||||
outboundQueueSize: 7,
|
||||
outboundControlQueueSize: 3,
|
||||
}
|
||||
c.startOutbound()
|
||||
defer c.Close()
|
||||
if got := cap(c.outbound); got != 7 {
|
||||
t.Fatalf("normal queue cap = %d, want 7", got)
|
||||
}
|
||||
if got := cap(c.outboundControl); got != 3 {
|
||||
t.Fatalf("control queue cap = %d, want 3", got)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestOutboundOptionsDefaults(t *testing.T) {
|
||||
opts := Options{}
|
||||
opts.setDefaults()
|
||||
if opts.OutboundQueueSize != 128 || opts.OutboundControlQueueSize != 32 {
|
||||
t.Fatalf("outbound queue defaults = %d/%d, want 128/32", opts.OutboundQueueSize, opts.OutboundControlQueueSize)
|
||||
}
|
||||
if opts.OutboundTrackedGlobalMaxBytes != 512<<20 {
|
||||
t.Fatalf("outbound tracked default = %d, want %d", opts.OutboundTrackedGlobalMaxBytes, 512<<20)
|
||||
}
|
||||
}
|
||||
|
||||
func TestServerNewConnectionsShareOutboundBudgetAndQueueLimits(t *testing.T) {
|
||||
srv := New(Options{
|
||||
OutboundQueueSize: 7,
|
||||
OutboundControlQueueSize: 3,
|
||||
OutboundTrackedGlobalMaxBytes: 20,
|
||||
})
|
||||
var rawKey crypto.Key
|
||||
key := rawKey.WithID()
|
||||
c1 := srv.newConn(nil, key, 1, 1)
|
||||
c2 := srv.newConn(nil, key, 2, 1)
|
||||
defer c1.Close()
|
||||
defer c2.Close()
|
||||
|
||||
if cap(c1.outbound) != 7 || cap(c1.outboundControl) != 3 || cap(c2.outbound) != 7 || cap(c2.outboundControl) != 3 {
|
||||
t.Fatalf("server queue caps = %d/%d and %d/%d, want 7/3",
|
||||
cap(c1.outbound), cap(c1.outboundControl), cap(c2.outbound), cap(c2.outboundControl))
|
||||
}
|
||||
if c1.outboundTrackedBudget != srv.outboundTrackedBudget || c2.outboundTrackedBudget != srv.outboundTrackedBudget {
|
||||
t.Fatal("server connections did not receive the shared outbound tracking budget")
|
||||
}
|
||||
if got := srv.outboundTrackedBudget.maxBytes; got != 20 {
|
||||
t.Fatalf("server outbound tracked max = %d, want 20", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestEncryptOutboundFrameDecryptsWithGotdCipher(t *testing.T) {
|
||||
var key crypto.Key
|
||||
if _, err := rand.Read(key[:]); err != nil {
|
||||
|
|
@ -104,8 +525,375 @@ func TestOutboundActorSerializesConcurrentSends(t *testing.T) {
|
|||
}
|
||||
}
|
||||
|
||||
func TestOutboundWriteErrorTerminallyClosesWithoutActorDeadlock(t *testing.T) {
|
||||
tr := &failAfterTransport{}
|
||||
tr.failAt.Store(1)
|
||||
c := newOutboundFailureTestConn(t, tr)
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
|
||||
defer cancel()
|
||||
if err := c.Send(ctx, proto.MessageFromServer, &tg.UpdatesTooLong{}); err == nil {
|
||||
t.Fatal("Send unexpectedly succeeded")
|
||||
}
|
||||
select {
|
||||
case <-c.outboundDone:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("outbound actor deadlocked while terminalizing its own write error")
|
||||
}
|
||||
if got := tr.closes.Load(); got != 1 {
|
||||
t.Fatalf("transport closes = %d, want 1", got)
|
||||
}
|
||||
if err := c.Send(ctx, proto.MessageFromServer, &tg.UpdatesTooLong{}); !errors.Is(err, ErrConnClosed) {
|
||||
t.Fatalf("second Send err = %v, want ErrConnClosed", err)
|
||||
}
|
||||
if got := tr.sends.Load(); got != 1 {
|
||||
t.Fatalf("physical sends after terminal error = %d, want 1", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOutboundResendWriteErrorTerminallyCloses(t *testing.T) {
|
||||
tr := &failAfterTransport{}
|
||||
c := newOutboundFailureTestConn(t, tr)
|
||||
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
|
||||
defer cancel()
|
||||
|
||||
if err := c.Send(ctx, proto.MessageFromServer, &tg.UpdatesTooLong{}); err != nil {
|
||||
t.Fatalf("initial Send: %v", err)
|
||||
}
|
||||
data, err := crypto.NewClientCipher(rand.Reader).DecryptFromBuffer(c.key, &bin.Buffer{Buf: tr.lastFrame()})
|
||||
if err != nil {
|
||||
t.Fatalf("decrypt initial frame: %v", err)
|
||||
}
|
||||
tr.failAt.Store(2)
|
||||
if _, err := c.ResendMessages(ctx, []int64{data.MessageID}); err == nil {
|
||||
t.Fatal("ResendMessages unexpectedly succeeded")
|
||||
}
|
||||
select {
|
||||
case <-c.outboundDone:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("outbound actor did not exit after resend write error")
|
||||
}
|
||||
if got := tr.closes.Load(); got != 1 {
|
||||
t.Fatalf("transport closes = %d, want 1", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOutboundTrackedBudgetSharedAcrossConnections(t *testing.T) {
|
||||
budget := newOutboundTrackedBudget(12)
|
||||
tr1 := &failAfterTransport{}
|
||||
tr2 := &failAfterTransport{}
|
||||
c1 := newOutboundTestConn(t, tr1, budget)
|
||||
c2 := newOutboundTestConn(t, tr2, budget)
|
||||
body := &encodedOutboundMessage{body: make([]byte, 8), typeID: tg.UpdatesTooLongTypeID}
|
||||
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
|
||||
defer cancel()
|
||||
|
||||
if err := c1.SendEncoded(ctx, proto.MessageFromServer, body); err != nil {
|
||||
t.Fatalf("first connection send: %v", err)
|
||||
}
|
||||
if got := budget.snapshot(); got != 8 {
|
||||
t.Fatalf("tracked bytes after first connection = %d, want 8", got)
|
||||
}
|
||||
if err := c2.SendEncoded(ctx, proto.MessageFromServer, body); !errors.Is(err, ErrOutboundTrackedBudget) && !errors.Is(err, ErrConnClosed) {
|
||||
t.Fatalf("second connection send err = %v, want tracked budget/closed", err)
|
||||
}
|
||||
select {
|
||||
case <-c2.outboundDone:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("budget-exhausted connection did not terminate")
|
||||
}
|
||||
if got := tr2.sends.Load(); got != 0 {
|
||||
t.Fatalf("budget-exhausted connection wrote %d frames, want 0", got)
|
||||
}
|
||||
if got := budget.snapshot(); got != 8 {
|
||||
t.Fatalf("tracked bytes after second rejection = %d, want first connection's 8", got)
|
||||
}
|
||||
|
||||
c1.Close()
|
||||
if got := budget.snapshot(); got != 0 {
|
||||
t.Fatalf("tracked bytes after first connection close = %d, want 0", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOutboundTrackedBudgetReleaseBroadcastsToAllWaiters(t *testing.T) {
|
||||
const waiters = 8
|
||||
budget := newOutboundTrackedBudget(waiters)
|
||||
if !budget.reserve(waiters) {
|
||||
t.Fatal("reserve initial saturated budget")
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
|
||||
defer cancel()
|
||||
results := make(chan error, waiters)
|
||||
for i := 0; i < waiters; i++ {
|
||||
go func() {
|
||||
results <- budget.waitReserve(ctx, nil, 1)
|
||||
}()
|
||||
}
|
||||
|
||||
deadline := time.Now().Add(time.Second)
|
||||
for {
|
||||
budget.wakeMu.Lock()
|
||||
got := budget.wake.waiters
|
||||
budget.wakeMu.Unlock()
|
||||
if got == waiters {
|
||||
break
|
||||
}
|
||||
if time.Now().After(deadline) {
|
||||
t.Fatalf("subscribed waiters = %d, want %d", got, waiters)
|
||||
}
|
||||
time.Sleep(time.Millisecond)
|
||||
}
|
||||
|
||||
// One batch release creates capacity for every waiter. A single-token notification strands
|
||||
// seven of them forever because successful reservations do not produce another wake-up.
|
||||
budget.release(waiters)
|
||||
for i := 0; i < waiters; i++ {
|
||||
if err := <-results; err != nil {
|
||||
t.Fatalf("waiter %d: %v", i, err)
|
||||
}
|
||||
}
|
||||
if got := budget.snapshot(); got != waiters {
|
||||
t.Fatalf("reserved bytes after broadcast = %d, want %d", got, waiters)
|
||||
}
|
||||
budget.release(waiters)
|
||||
}
|
||||
|
||||
func TestOutboundGlobalBudgetIncludesQueuedBodies(t *testing.T) {
|
||||
budget := newOutboundTrackedBudget(24)
|
||||
tr := newBlockingOutboundTransport()
|
||||
c := newOutboundTestConn(t, tr, budget)
|
||||
body := &encodedOutboundMessage{body: make([]byte, 8), typeID: tg.UpdatesTooLongTypeID}
|
||||
|
||||
if err := c.SendBestEffortEncoded(context.Background(), proto.MessageFromServer, body, 0); err != nil {
|
||||
t.Fatalf("enqueue writing body: %v", err)
|
||||
}
|
||||
select {
|
||||
case <-tr.started:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("outbound actor did not start blocked write")
|
||||
}
|
||||
for i := 0; i < 2; i++ {
|
||||
if err := c.SendBestEffortEncoded(context.Background(), proto.MessageFromServer, body, 0); err != nil {
|
||||
t.Fatalf("enqueue queued body %d: %v", i, err)
|
||||
}
|
||||
}
|
||||
if got := budget.snapshot(); got != 24 {
|
||||
t.Fatalf("writing + queued budget = %d, want 24", got)
|
||||
}
|
||||
if err := c.SendBestEffortEncoded(context.Background(), proto.MessageFromServer, body, 0); !errors.Is(err, ErrOutboundTrackedBudget) {
|
||||
t.Fatalf("over-budget enqueue err = %v, want ErrOutboundTrackedBudget", err)
|
||||
}
|
||||
select {
|
||||
case <-c.outboundDone:
|
||||
t.Fatal("best-effort global pressure terminated a healthy connection")
|
||||
case <-time.After(50 * time.Millisecond):
|
||||
}
|
||||
if got := budget.snapshot(); got != 24 {
|
||||
t.Fatalf("budget after non-terminal rejection = %d, want existing 24", got)
|
||||
}
|
||||
if err := tr.Close(); err != nil {
|
||||
t.Fatalf("close blocking transport: %v", err)
|
||||
}
|
||||
select {
|
||||
case <-c.outboundDone:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("outbound actor did not stop after transport failure")
|
||||
}
|
||||
if got := budget.snapshot(); got != 0 {
|
||||
t.Fatalf("budget after transport close = %d, want zero", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOutboundOversizedBodyRejectedBeforeEncryption(t *testing.T) {
|
||||
budget := newOutboundTrackedBudget(64 << 20)
|
||||
tr := &failAfterTransport{}
|
||||
c := newOutboundTestConn(t, tr, budget)
|
||||
body := &encodedOutboundMessage{body: make([]byte, maxOutboundBodyBytes+1), typeID: tg.UpdatesTooLongTypeID}
|
||||
err := c.SendEncoded(context.Background(), proto.MessageFromServer, body)
|
||||
if !errors.Is(err, ErrOutboundMessageTooLarge) {
|
||||
t.Fatalf("oversized outbound err = %v, want ErrOutboundMessageTooLarge", err)
|
||||
}
|
||||
if got := tr.sends.Load(); got != 0 {
|
||||
t.Fatalf("oversized outbound wrote %d frames, want zero", got)
|
||||
}
|
||||
if got := budget.snapshot(); got != 0 {
|
||||
t.Fatalf("oversized outbound reserved %d bytes, want zero", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOutboundCloseRaceDrainsEveryProducerReservation(t *testing.T) {
|
||||
budget := newOutboundTrackedBudget(1 << 20)
|
||||
c := newOutboundTestConn(t, &failAfterTransport{}, budget)
|
||||
body := &encodedOutboundMessage{body: make([]byte, 128), typeID: tg.UpdatesTooLongTypeID}
|
||||
start := make(chan struct{})
|
||||
var wg sync.WaitGroup
|
||||
for i := 0; i < 128; i++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
<-start
|
||||
_ = c.SendBestEffortEncoded(context.Background(), proto.MessageFromServer, body, 0)
|
||||
}()
|
||||
}
|
||||
close(start)
|
||||
c.Close()
|
||||
wg.Wait()
|
||||
if got := budget.snapshot(); got != 0 {
|
||||
t.Fatalf("outbound budget after close/enqueue race = %d, want zero", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOutboundTrackedBudgetAckAndCloseReturnExactly(t *testing.T) {
|
||||
t.Run("ack", func(t *testing.T) {
|
||||
budget := newOutboundTrackedBudget(64)
|
||||
tr := &failAfterTransport{}
|
||||
c := newOutboundTestConn(t, tr, budget)
|
||||
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
|
||||
defer cancel()
|
||||
body := &encodedOutboundMessage{body: make([]byte, 12), typeID: tg.UpdatesTooLongTypeID}
|
||||
if err := c.SendEncoded(ctx, proto.MessageFromServer, body); err != nil {
|
||||
t.Fatalf("send: %v", err)
|
||||
}
|
||||
if got := budget.snapshot(); got != 12 {
|
||||
t.Fatalf("tracked bytes after send = %d, want 12", got)
|
||||
}
|
||||
data, err := crypto.NewClientCipher(rand.Reader).DecryptFromBuffer(c.key, &bin.Buffer{Buf: tr.lastFrame()})
|
||||
if err != nil {
|
||||
t.Fatalf("decrypt frame: %v", err)
|
||||
}
|
||||
c.AckServerMessages([]int64{data.MessageID})
|
||||
deadline := time.Now().Add(time.Second)
|
||||
for budget.snapshot() != 0 && time.Now().Before(deadline) {
|
||||
time.Sleep(time.Millisecond)
|
||||
}
|
||||
if got := budget.snapshot(); got != 0 {
|
||||
t.Fatalf("tracked bytes after ack = %d, want 0", got)
|
||||
}
|
||||
})
|
||||
|
||||
t.Run("close", func(t *testing.T) {
|
||||
budget := newOutboundTrackedBudget(64)
|
||||
c := newOutboundTestConn(t, &failAfterTransport{}, budget)
|
||||
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
|
||||
defer cancel()
|
||||
body := &encodedOutboundMessage{body: make([]byte, 12), typeID: tg.UpdatesTooLongTypeID}
|
||||
if err := c.SendEncoded(ctx, proto.MessageFromServer, body); err != nil {
|
||||
t.Fatalf("send: %v", err)
|
||||
}
|
||||
if got := budget.snapshot(); got != 12 {
|
||||
t.Fatalf("tracked bytes after send = %d, want 12", got)
|
||||
}
|
||||
c.Close()
|
||||
if got := budget.snapshot(); got != 0 {
|
||||
t.Fatalf("tracked bytes after close = %d, want 0", got)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func TestOutboundTrackedBudgetWriteFailureReturnsReservation(t *testing.T) {
|
||||
budget := newOutboundTrackedBudget(64)
|
||||
tr := &failAfterTransport{}
|
||||
tr.failAt.Store(1)
|
||||
c := newOutboundTestConn(t, tr, budget)
|
||||
ctx, cancel := context.WithTimeout(context.Background(), time.Second)
|
||||
defer cancel()
|
||||
body := &encodedOutboundMessage{body: make([]byte, 12), typeID: tg.UpdatesTooLongTypeID}
|
||||
if err := c.SendEncoded(ctx, proto.MessageFromServer, body); err == nil {
|
||||
t.Fatal("send unexpectedly succeeded")
|
||||
}
|
||||
select {
|
||||
case <-c.outboundDone:
|
||||
case <-time.After(time.Second):
|
||||
t.Fatal("write-failed connection did not terminate")
|
||||
}
|
||||
if got := budget.snapshot(); got != 0 {
|
||||
t.Fatalf("tracked bytes after write failure = %d, want 0", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOutboundStateEvictionReturnsTrackedBudget(t *testing.T) {
|
||||
budget := newOutboundTrackedBudget(64)
|
||||
state := newOutboundStateWithLimits(budget, 2, 8)
|
||||
defer state.releaseAll()
|
||||
frames := make([]*outboundFrame, 0, 3)
|
||||
for id := int64(1); id <= 3; id++ {
|
||||
frame := &outboundFrame{msgID: id, body: make([]byte, 4), reservedBytes: 4}
|
||||
frames = append(frames, frame)
|
||||
if !budget.reserve(len(frame.body)) {
|
||||
t.Fatalf("reserve frame %d", id)
|
||||
}
|
||||
dropped := state.addReserved(frame)
|
||||
if id < 3 && dropped != 0 {
|
||||
t.Fatalf("frame %d dropped %d, want 0", id, dropped)
|
||||
}
|
||||
if id == 3 && dropped != 1 {
|
||||
t.Fatalf("third frame dropped %d, want 1", dropped)
|
||||
}
|
||||
}
|
||||
if got := budget.snapshot(); got != 8 {
|
||||
t.Fatalf("tracked bytes after eviction = %d, want 8", got)
|
||||
}
|
||||
if frames[0].body != nil {
|
||||
t.Fatal("evicted frame retained its body reference")
|
||||
}
|
||||
state.releaseAll()
|
||||
if got := budget.snapshot(); got != 0 {
|
||||
t.Fatalf("tracked bytes after state close = %d, want 0", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOutboundStateReleasesMixedBodyAndControlBudgets(t *testing.T) {
|
||||
bodyBudget := newOutboundTrackedBudget(16)
|
||||
controlBudget := newOutboundTrackedBudget(16)
|
||||
state := newOutboundStateWithLimits(bodyBudget, 1, 16)
|
||||
|
||||
if !controlBudget.reserve(4) {
|
||||
t.Fatal("reserve control frame")
|
||||
}
|
||||
controlFrame := &outboundFrame{
|
||||
msgID: 1,
|
||||
body: make([]byte, 4),
|
||||
reservedBytes: 4,
|
||||
reservationBudget: controlBudget,
|
||||
}
|
||||
if dropped := state.addReserved(controlFrame); dropped != 0 {
|
||||
t.Fatalf("first add dropped %d, want 0", dropped)
|
||||
}
|
||||
|
||||
if !bodyBudget.reserve(4) {
|
||||
t.Fatal("reserve body frame")
|
||||
}
|
||||
bodyFrame := &outboundFrame{
|
||||
msgID: 2,
|
||||
body: make([]byte, 4),
|
||||
reservedBytes: 4,
|
||||
reservationBudget: bodyBudget,
|
||||
}
|
||||
if dropped := state.addReserved(bodyFrame); dropped != 1 {
|
||||
t.Fatalf("second add dropped %d, want control frame eviction", dropped)
|
||||
}
|
||||
if got := controlBudget.snapshot(); got != 0 {
|
||||
t.Fatalf("control budget after eviction = %d, want 0", got)
|
||||
}
|
||||
if got := bodyBudget.snapshot(); got != 4 {
|
||||
t.Fatalf("body budget after eviction = %d, want 4", got)
|
||||
}
|
||||
if controlFrame.body != nil || controlFrame.reservationBudget != nil {
|
||||
t.Fatal("evicted control frame retained body or budget ownership")
|
||||
}
|
||||
|
||||
state.releaseAll()
|
||||
if got := bodyBudget.snapshot(); got != 0 {
|
||||
t.Fatalf("body budget after state close = %d, want 0", got)
|
||||
}
|
||||
if bodyFrame.body != nil || bodyFrame.reservationBudget != nil {
|
||||
t.Fatal("closed body frame retained body or budget ownership")
|
||||
}
|
||||
}
|
||||
|
||||
func TestSendBestEffortQueueFullBehavior(t *testing.T) {
|
||||
c := &Conn{metrics: NopMetrics{}}
|
||||
c := &Conn{metrics: NopMetrics{}, outboundTrackedBudget: newOutboundTrackedBudget(1 << 20)}
|
||||
c.outbound = make(chan outboundOp, 1)
|
||||
c.outboundControl = make(chan outboundOp, 1)
|
||||
c.outboundStop = make(chan struct{})
|
||||
|
|
@ -140,6 +928,21 @@ 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.outboundStop = make(chan struct{})
|
||||
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)
|
||||
}
|
||||
if got := len(c.outboundControl); got != 1 {
|
||||
t.Fatalf("control queue len = %d, want bounded at 1", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFrameNeedsAckServiceExceptions(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
|
|
|
|||
|
|
@ -96,17 +96,22 @@ func TestPasskeyEndToEnd(t *testing.T) {
|
|||
userStore := memory.NewUserStore()
|
||||
authKeyStore := memory.NewAuthKeyStore()
|
||||
helpStore := memory.NewHelpStore()
|
||||
dialogStore := memory.NewDialogStore()
|
||||
messageStore := memory.NewMessageStore(dialogStore)
|
||||
updateEventStore := memory.NewUpdateEventStore()
|
||||
passkeyService := passkeyapp.NewService(memory.NewPasskeyStore(), memory.NewPasskeyChallengeStore(), rpID, dc)
|
||||
|
||||
deps := rpc.Deps{
|
||||
Auth: auth.NewService(userStore, memory.NewAuthorizationStore(), memory.NewCodeStore(), authKeyStore, memory.NewTempAuthKeyBindingStore(), code),
|
||||
Auth: auth.NewService(userStore, memory.NewAuthorizationStore(), memory.NewCodeStore(), authKeyStore, memory.NewTempAuthKeyBindingStore(), code,
|
||||
auth.WithLoginMessages(messageStore, dialogStore),
|
||||
auth.WithLoginCodeDelivery(memory.NewLoginCodeDeliveryStore(messageStore, updateEventStore))),
|
||||
Account: account.NewService(memory.NewPasswordStore(), account.WithUsers(userStore)),
|
||||
Help: help.NewService(helpStore, helpStore),
|
||||
Users: users.NewService(userStore),
|
||||
Updates: updates.NewService(memory.NewUpdateStateStore(), memory.NewUpdateEventStore()),
|
||||
Updates: updates.NewService(memory.NewUpdateStateStore(), updateEventStore),
|
||||
|
||||
Contacts: contacts.NewService(memory.NewContactStore()),
|
||||
Dialogs: dialogs.NewService(memory.NewDialogStore()),
|
||||
Dialogs: dialogs.NewService(dialogStore),
|
||||
Passkey: passkeyService,
|
||||
}
|
||||
router := rpc.New(rpc.Config{DC: dc, IP: tcpAddr.IP.String(), Port: tcpAddr.Port}, deps, zaptest.NewLogger(t), clock.System)
|
||||
|
|
|
|||
Some files were not shown because too many files have changed in this diff Show more
Loading…
Add table
Add a link
Reference in a new issue